21. Lab: Thinking 모델

난이도: 고급 | GPU: RTX 3090~5090 (VRAM 24GB+) | 소요 시간: 4~8시간

목표

Llama-3.2-3B에 Mixture-of-LoRAs 전략으로 일반 지식과 CoT(Chain-of-Thought) 추론 능력을 동시에 부여합니다. 2단계 파이프라인 훈련을 통해 기존 능력을 보존하면서 수학 추론을 강화하는 방법을 학습합니다.

핵심 포인트

Stage 1: 일반 SFT (50K 영한 혼합) -- 기초 instruction-following
Stage 2: CoT SFT (OpenR1-Math 40K) -- 단계별 수학 추론
Mixture-of-LoRAs: 4개 Expert가 도메인별로 자연 분화

사전 준비

데이터

19_data_collection.md의 섹션 4 (Reasoning/CoT 데이터) 를 실행하여 다음을 준비합니다:

data/reasoning/openr1_math_40k.jsonl   # CoT 훈련 (40K)
data/reasoning/openr1_math_1k.jsonl    # CoT 벤치 (1K)
data/sft_50k_en_ko_raw.jsonl           # 일반 SFT (기본 제공)
data/sft_1k_en_ko_raw.jsonl            # 일반 SFT 벤치 (기본 제공)

스크립트

cp examples/run_llama32_pipeline_reasoning.sh ./
chmod +x run_llama32_pipeline_reasoning.sh

2단계 훈련 설계

Stage 1: 일반 SFT

Config: configs/presets/reasoning/llama3.2_3b_mixture_lora_sft.yml

injection:
  strategy: mixture_lora
  num_experts: 4
  top_k: 2
  lora_r: 32
  lora_dropout: 0.08           # 초기에는 dropout 높게
training:
  phases:
    - step: 0
      trainable: ["router"]             # Router 단독 워밍업
    - step: 1000
      trainable: ["router", "lora", "attn_lora"]  # 전체 LoRA 오픈
  lr: 6.0e-5
  max_train_steps: 5000       # 스크립트에서 override

Router를 먼저 1000 step 단독 학습하여 Expert 분배를 안정화
높은 router_z_loss(0.0003)와 aux_loss(0.003)로 부하 균형 강제

Stage 2: CoT SFT (핵심 변경)

Config: configs/presets/reasoning/llama3.2_3b_mixture_lora_sft_cot.yml

Stage 1과의 주요 차이점:

파라미터	Stage 1 (일반)	Stage 2 (CoT)	이유
Router 워밍업	step 0~1000	없음	기존 라우팅 규칙 보호
`lr`	6.0e-5	1.5e-5	기존 지식 파괴 방지
`lora_dropout`	0.08	0.05	기존 지식 보존
`router_z_loss_coef`	0.0003	0.0001	MoE 패널티 완화
`aux_loss_coef`	0.003	0.001	Expert 쏠림 허용
`max_train_steps`	5000	3000	과적합 방지

training:
  phases:
    - step: 0
      trainable: ["router", "lora", "attn_lora"]  # 전부 동시 오픈
  lr: 1.5e-5                                       # 4배 낮은 LR

Router 단독 훈련 삭제 이유: Stage 1에서 이미 라우터가 초기화되어 있습니다. Stage 2에서 다시 Router만 학습하면 기존 라우팅 규칙이 수학 데이터로 재설정되어 일반 지식 Expert 배분이 무너질 수 있습니다.

MoE 패널티 완화 이유: 수학 데이터만 들어오므로 한두 개의 Expert에 트래픽이 집중되는 것이 자연스럽습니다. 강한 균형 패널티는 오히려 학습을 방해합니다.

실행

./run_llama32_pipeline_reasoning.sh

또는 단계별 수동 실행:

# Stage 1: 일반 SFT (5000 steps)
eulerforge train \
  --preset configs/presets/reasoning/llama3.2_3b_mixture_lora_sft.yml \
  --set model_name=meta-llama/Llama-3.2-3B \
  --set data.format=raw --set data.task=sft \
  --set data.path=data/sft_50k_en_ko_raw.jsonl \
  --set data.max_length=512 \
  --set training.max_train_steps=5000 \
  --set training.batch_size=10 --set training.grad_accum_steps=2 \
  --output-dir outputs/reasoning/stage1_general_sft

# Stage 2: CoT SFT (Stage 1 체크포인트 기반, 3000 steps)
eulerforge train \
  --preset configs/presets/reasoning/llama3.2_3b_mixture_lora_sft_cot.yml \
  --set model_name=outputs/reasoning/stage1_general_sft/final \
  --set data.format=raw --set data.task=sft \
  --set data.path=data/reasoning/openr1_math_40k.jsonl \
  --set data.max_length=1024 \
  --set training.max_train_steps=3000 \
  --set training.batch_size=8 --set training.grad_accum_steps=2 \
  --output-dir outputs/reasoning/stage2_reasoning_sft

결과 해석

기대 패턴

Stage 1 벤치: 일반 instruction-following 능력 확보 (base 대비 향상)
Stage 2 벤치: CoT 추론 능력 추가 + 일반 능력 유지 (또는 소폭 저하)

CoT 과적합 주의

Stage 2에서 수학 데이터만 오래 훈련하면 과적합이 발생할 수 있습니다: - 수학 bench 점수는 오르지만 일반 bench 점수가 급락 - Stage 2 steps를 줄이거나 lr을 더 낮춰 실험

심화: 1단계 통합 훈련 (숙제)

현재 스크립트는 2단계로 나뉘어 있지만, Mixture-of-LoRAs는 4개 Expert 구조이므로 두 데이터를 섞어서 한 번에 훈련하는 것도 가능합니다:

# 일반 SFT + CoT 데이터를 하나로 합침
cat data/sft_50k_en_ko_raw.jsonl data/reasoning/openr1_math_40k.jsonl \
  | shuf --random-source=<(yes 42) > data/reasoning/mixed_90k.jsonl

echo "Mixed: $(wc -l < data/reasoning/mixed_90k.jsonl) rows"

과제: llama3.2_3b_mixture_lora_sft.yml 프리셋 하나로 mixed_90k.jsonl을 훈련하고, 2단계 파이프라인 결과와 비교해 보세요.

장점: 파이프라인 단순화, Expert가 자연스럽게 도메인별 분화
단점: CoT 데이터의 비율 조절이 어려움 (2단계에서는 lr/steps로 세밀 제어 가능)

MoE의 라우터가 일반/수학 토큰을 자동으로 다른 Expert에 배분하는지 --metrics-level advanced로 확인해 보세요.