Medical QA를 위한 Llama 3.2 3B 파인튜닝(Fine-Tuning): 3주 차 – 첫 번째 트레이닝 실행

이번 주에 일어난 일

2주 차는 포맷팅을 마친 정제된 데이터셋을 Hugging Face Hub에 푸시하며 마무리되었습니다. 3주 차는 실제 파인튜닝 (Fine-Tuning)이 이루어지는 단계입니다. LoRA 어댑터 (adapters)를 구성하고, 트레이닝 루프 (training loop)를 실행하며, 파인튜닝된 모델의 출력을 1주 차 베이스라인 (baseline)과 비교합니다.

과정이 순탄하지는 않았습니다. 이 포스트는 무엇이 고장 났는지, 왜 고장 났는지, 그리고 실제 결과가 무엇을 보여주었는지를 기록합니다.

이번 주의 기술 스택 (Stack)

파인튜닝 스택은 1주 차에 설치한 것들을 기반으로 구축됩니다:

• peft – LoRA를 구현합니다. 동결된 베이스 모델 (base model)의 특정 레이어 (layers)에 학습 가능한 작은 어댑터 행렬 (adapter matrices)을 추가합니다.
• trl – 지도 학습 기반 파인튜닝 (SFT, Supervised Fine-Tuning) 트레이닝 루프인 SFTTrainer를 제공합니다. 배치 (batching), 그래디언트 누적 (gradient accumulation), 평가 (evaluation) 및 체크포인팅 (checkpointing)을 처리합니다.
• bitsandbytes – 3B 모델이 GPU 메모리에 들어갈 수 있도록 여전히 4비트 양자화 (4-bit quantization)를 처리합니다.

LoRA (Low-Rank Adaptation)가 실제로 하는 일

트레이닝 실행에 들어가기에 앞서, 이 프로젝트의 핵심 기술인 LoRA가 무엇을 하는지 설명할 가치가 있습니다.

베이스 모델은 32억 개의 파라미터 (parameters)를 가지고 있습니다. 파인튜닝 중에 이 모든 것을 업데이트하려면 약 24GB의 VRAM과 수 시간의 연산이 필요합니다. 이는 무료 GPU에서는 실행 불가능한 수준입니다.

LoRA는 원래의 가중치 (weights)를 전혀 업데이트하지 않습니다. 대신, 특정 레이어 옆에 두 개의 작은 학습 가능한 행렬을 추가합니다. 모델의 모든 어텐션 레이어 (attention layer)에는 쿼리 (query), 키 (key), 값 (value), 그리고 출력 프로젝션 (output projections)을 위한 가중치 행렬이 있습니다. 각 항목에 대해 LoRA는 다음을 추가합니다:

Matrix A: 16 x 3072 = 49,152 parameters
Matrix B: 3072 x 16 = 49,152 parameters

순전파 (forward pass) 동안, 레이어는 다음을 계산합니다:

W: frozen base model weights.

scale: 조절 강도를 제어합니다 (scale = lora_alpha / r)
...

W는 고정(frozen)됩니다. 오직 A와 B만이 그래디언트 업데이트(gradient updates)를 받습니다. 28개의 트랜스포머 레이어(transformer layers)와 레이어당 4개의 타겟 모듈(target modules) 전체에 걸쳐, 이는 총 32억 개 중 900만 개의 학습 가능한 파라미터(trainable parameters)에 해당합니다. 이는 모델의 0.28%만이 학습된다는 것을 의미합니다. 원래의 가중치(weights)는 완전히 손대지 않은 상태로 유지됩니다.

랭크(rank) 값 16은 어댑터(adapter)의 용량을 제어합니다. 랭크가 높을수록 어댑터의 표현력(expressive)은 좋아지지만 학습해야 할 파라미터도 많아집니다. 16은 3B 모델을 위한 표준적인 시작점입니다.

prepare_model_for_kbit_training이 하는 일

이 함수는 LoRA 어댑터를 적용하기 전에 호출되며, 정확히 설명할 가치가 있습니다.

모델은 4비트 양자화(4-bit quantization)로 로드되는데, 이는 가중치가 압축되고 고정(frozen)되었음을 의미합니다. 학습을 시작하기 전에, PyTorch는 고정된 양자화 레이어(quantized layers)를 통해 LoRA 어댑터에 도달할 수 있도록 그래디언트(gradients)를 역전파(backward propagation)하는 방법을 알아야 합니다.

기본적으로 PyTorch는 고정된 레이어를 발견하면 해당 레이어를 통한 그래디언트 흐름 추적을 완전히 중단합니다. 손실(loss)로부터 발생하는 에러 신호가 LoRA 행렬(matrices)에 결코 도달하지 못하게 됩니다. 그래디언트가 없으면 학습도 없습니다.

prepare_model_for_kbit_training은 해당 레이어들이 업데이트되지 않더라도 PyTorch가 고정된 레이어를 통해 계속해서 그래디언트를 전달하도록 지시합니다. 즉, 고정된 레이어들은 그래디언트 계산 그래프(gradient computation graph)에서 통과 노드(passthrough nodes) 역할을 하게 됩니다. 체인의 끝에 있는 LoRA 어댑터가 그래디언트를 전달받아 그에 따라 업데이트를 수행합니다.

이 호출이 없다면, 학습은 에러 없이 완료되지만 LoRA 가중치는 전혀 변하지 않습니다. 학습 후에도 모델은 베이스 모델(base model)과 동일한 상태일 것입니다.

하드웨어 문제 발생

초기 계획은 fp16 혼합 정밀도(mixed precision)를 활성화하여 Google Colab의 무료 T4 GPU에서 학습하는 것이었습니다. 하지만 실패했습니다.

발생한 에러는 다음과 같습니다:

NotImplementedError: "_amp_foreach_non_finite_check_and_unscale_cuda"
not implemented for 'BFloat16'

fp16 그래디언트 스케일러(gradient scaler)가 역전파(backward pass) 과정 중 prepare_model_for_kbit_training에 의해 내부적으로 생성된 BFloat16 텐서를 발견했습니다. T4 GPU는 float16을 네이티브로 지원하지만, 양자화된 베이스 모델 레이어(quantized base model layers)와 최신 PEFT 버전의 LoRA 어댑터 초기화 사이의 상호작용으로 인해 fp16 스케일러가 처리할 수 없는 BFloat16 중간 텐서(intermediate tensors)가 생성됩니다.

해결 방법은 fp16=False, bf16=False로 설정하여 혼합 정밀도(mixed precision)를 완전히 비활성화하고, Colab의 2~4시간 제한 대신 주당 30시간의 세션 제한을 제공하는 Kaggle의 T4 GPU로 이동하는 것이었습니다.

트레이드오프(tradeoff): T4에서 float32 그래디언트 계산은 fp16보다 대략 2~3배 느립니다. fp16을 사용했을 때 45분이 소요되었을 트레이닝 실행이 이를 사용하지 않자 1시간 13분이 걸렸습니다.

이는 TRL 1.5.1, PEFT, 그리고 T4 하드웨어의 특정 조합에서 발생하는 라이브러리 버전 호환성 문제입니다. A100이나 V100을 사용한다면 각각 bf16=True 또는 fp16=True를 설정하여 이러한 충돌 없이 작동할 것입니다.

트레이닝 설정 (Training Configuration)

training_args = SFTConfig(
    output_dir="/kaggle/working/checkpoints",
    num_train_epochs=1,
...

설명이 필요한 몇 가지 결정 사항들입니다:
gradient_accumulation_steps=4: 배치 크기(batch size)가 4일 때, 각 가중치 업데이트는 16개 샘플(4 x 4)로부터 누적된 그래디언트를 사용합니다. 이는 16개 샘플을 동시에 로드하는 VRAM 비용 없이 더 큰 배치 크기를 시뮬레이션합니다.

gradient_checkpointing=True: 순전파(forward pass) 동안 모든 중간 레이어 활성화 값(intermediate layer activations)을 VRAM에 저장하는 대신, PyTorch는 역전파(backward pass) 중에 필요할 때 이를 다시 계산합니다. 약 20% 정도의 트레이닝 속도 저하를 대가로 상당한 VRAM 절약 효과를 얻을 수 있습니다. float32를 실행 중인 15.6GB T4에서는 OOM(OutOfMemoryError: CUDA out of memory)을 방지하기 위해 이 설정이 필수적이었습니다.

load_best_model_at_end=True: 트레이닝이 완료된 후, 마지막 체크포인트 대신 평가 손실(eval loss)이 가장 낮은 체크포인트를 로드합니다. 이는 마지막 몇 단계에서 발생할 수 있는 미세한 과적합(overfitting)을 방지합니다.

learning_rate=2e-4: LoRA 파인튜닝 (fine-tuning)의 표준 시작점입니다. 이는 스텝당 어댑터 가중치 (adapter weights)가 얼마나 공격적으로 업데이트되는지를 제어합니다.

훈련 결과 (Training Results)

Step 100:  train loss 2.570  |  eval loss 2.558
Step 200:  train loss 2.525  |  eval loss 2.511
Step 300:  train loss 2.482  |  eval loss 2.495
...

훈련 손실 (training loss)은 309개 스텝 전체에 걸쳐 꾸준히 감소했습니다. 평가 손실 (eval loss)은 마지막 스텝에서 단 0.013의 차이만을 보이며 훈련 손실을 밀접하게 추적했습니다. 훈련 손실과 평가 손실 사이의 간격이 크고 점점 벌어진다면 이는 과적합 (overfitting)을 나타냅니다. 일관되게 작은 간격은 모델이 단순히 훈련 데이터를 암기하는 것이 아니라, 보지 못한 예시들에 대해 일반화 (generalizing)하고 있음을 나타냅니다.

309번째 스텝에서도 손실이 여전히 감소하고 있었으며, 이는 완전한 수렴 (convergence) 이전에 훈련이 중단되었음을 의미합니다. 이는 1 에포크 (epoch) 기준으로는 예상된 결과입니다. 4주 차에는 9,000개의 샘플로 구성된 전체 데이터셋을 사용하여 2 에포크를 실행할 예정입니다.

시스템 프롬프트는 모델의 일부입니다

이 점 또한 주목할 가치가 있습니다: 추론 (inference) 중에 사용되는 시스템 프롬프트 (system prompt)는 훈련 중에 사용된 것과 정확히 일치해야 합니다.

모든 훈련 샘플은 다음과 같은 형식으로 구성되었습니다:

If you are a doctor, please answer the medical questions
based on the patient's description.

모델은 해당 특정 프레임워크에 대응하여 임상적인 산문 (clinical prose)을 생성하는 법을 배우는 데 309 스텝을 소비했습니다. 추론 시점에 다른 시스템 프롬프트를 사용하면 모델을 파인튜닝 (fine-tuned)된 적이 없는 컨텍스트 (context)에 놓이게 되어 출력 품질이 저하됩니다. 이는 언급이 필요한 미묘하지만 중요한 제약 사항입니다.

전과 후 (Before and After)

다음은 1주 차에서 베이스 모델 (base model)에 질문했던 것과 동일한 5가지 질문을 훈련 후의 파인튜닝된 모델에 던진 결과입니다.

핵심 비교 — 제2형 당뇨병 (type 2 diabetes):

베이스 모델:

"몸에서 인슐린을 더 많이 생성하면 몸이 더 많은 수분을 보유하게 되어 갈증이 증가할 수 있습니다."

파인튜닝된 모델:

"갈증 및 배뇨 증가: 높은 혈당 수치는 신체가 더 많은 소변을 생성하게 하여 탈수와 갈증 증가를 유발할 수 있습니다."

환각 (Hallucination) 현상이 사라졌습니다. 베이스 모델은 인슐린과 수분 보유 사이의 인과 관계를 허구로 만들어냈습니다. 반면, 파인튜닝 (Fine-tuned)된 모델은 갈증 증가의 원인을 높은 혈당 수치로 인한 삼투성 이뇨 (Osmotic diuresis)로 정확하게 설명합니다. 이것이 바로 이 프로젝트가 해결하고자 설계된 정확한 메커니즘입니다.

그 외 개선된 점:

다섯 가지 응답 모두에서 불필요한 서두 (Filler openers)가 완전히 사라졌습니다. 베이스 모델은 모든 답변을 "의료 보조원으로서, 기꺼이 도와드리겠습니다..."와 같이 시작했습니다. 파인튜닝된 응답 중에는 이러한 패턴이 전혀 포함되어 있지 않습니다. 응답이 임상적 내용으로 바로 들어갑니다.

고혈압에 대한 응답이 크게 개선되었습니다. 파인튜닝된 모델은 중증도에 따라 치료법을 계층화하고, 구체적인 약물군과 예시를 명시합니다. 베이스 모델은 훨씬 더 일반적인 수준에 머물렀습니다.

여전히 개선이 필요한 점:

당뇨병 응답에서 "느린 말투"를 제2형 당뇨병의 증상으로 나열하고 있습니다. 이는 사실이 아닙니다. 느린 말투는 저혈당 (Hypoglycemia) 또는 뇌졸중 (Stroke)과 관련이 있으며, 초기 제2형 당뇨병과는 관련이 없습니다. 4,937개의 샘플에 대해 1 에포크 (Epoch)를 수행했음에도 모든 환각을 제거하지는 못했습니다.

심장마비 응답에는 "피를 토함"과 "얼굴, 특히 뺨이나 이마의 통증 또는 불편함"이 경고 징후로 포함되어 있습니다. 둘 다 전형적인 심장 증상은 아닙니다. 한 응답은 "이 답변이 질문에 도움이 되기를 바랍니다. 추가로 도와드릴 일이 있으면 알려주세요"로 끝나는데, 이는 정제 파이프라인 (Cleaning pipeline)에서 살아남은 불필요한 패턴입니다.

이러한 잔여 오류들은 4주 차의 구체적인 목표 대상입니다.

모델 저장 위치

파인튜닝된 어댑터 가중치 (Adapter weights)는 Hugging Face Hub에 게시되어 있습니다:

nicholas-ugbala-hf/llama-3.2-3b-medical-finetuned

해당 저장소(repo)에는 LoRA 어댑터 가중치(adapter_model.safetensors), 어댑터 설정(adapter_config.json), 그리고 커스텀 패딩 토큰(pad token)이 포함된 수정된 토크나이저(tokenizer)가 포함되어 있습니다. 베이스 모델(base model)은 Meta의 저장소에 그대로 유지됩니다. 파인튜닝(Fine-tuned)된 모델을 로드하려면 두 가지가 모두 필요합니다.

파일 크기에 관한 참고 사항: adapter_model.safetensors는 LoRA 어댑터의 예상 크기인 30~60MB가 아니라 3.19GB입니다. 이는 훈련 과정에서 토크나이저 어휘(vocabulary)가 토큰 하나만큼 확장됨에 따라, PEFT가 어댑터와 함께 전체 임베딩 레이어(embedding layer)를 저장했기 때문입니다. 기능은 동일하지만 파일 크기는 이러한 트레이드오프(tradeoff)를 반영합니다.

다음 단계

3주 차를 통해 파이프라인(pipeline)이 작동하며 측정 가능한 개선을 만들어냈음을 확인했습니다. 잔차 오차(residual errors) 또한 명확하게 식별되었습니다.

4주 차 목표: 9,000개의 전체 샘플 데이터셋, 2 에포크(epochs), 남아있는 필러 패턴(filler patterns)을 제거하기 위한 더 엄격한 데이터 정제(data cleaning), 그리고 사실적 근거(factual grounding)를 개선하기 위한 두 번째 데이터셋 활용입니다.

모델: huggingface.co/nicholas-ugbala-hf/llama-3.2-3b-medical-finetuned
저장소: github.com/nicholas-ugbala-dev/healthcare-llm-finetune
데이터셋: huggingface.co/datasets/nicholas-ugbala-hf/chatdoctor-cleaned-10k