Add train_sft.py

train_sft.py

@@ -0,0 +1,67 @@
+"""C17d: 모든 풀이 + 길이 필터 (1500자 이하만) + NaN 방지"""
+import json, re, random, torch, numpy as np, os
+from collections import defaultdict
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from trl import SFTTrainer, SFTConfig
+from transformers import EarlyStoppingCallback
+from datasets import Dataset
+
+SEED = 42
+random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
+if torch.cuda.is_available():
+    torch.cuda.manual_seed_all(SEED)
+    if torch.cuda.get_device_capability()[0] >= 8: torch.set_float32_matmul_precision('high')
+
+SP = "주어진 수학 문제를 단계별로 풀고 답변을 작성하세요.\n반드시 최종 답변을 \\boxed{정수} 형식으로 마지막 줄에 출력하세요.\n예시: \\boxed{42}"
+
+print("=== C17d: All solutions, length-filtered (≤1500 chars) ===")
+
+with open("data/GSM8K_full_qwen3_30b.json") as f:
+    data = json.load(f)
+
+# 길이 필터: 1500자 이하만
+filtered = [d for d in data if len(d['answer']) <= 1500]
+print(f"원본: {len(data)}개 → 필터 후: {len(filtered)}개 (제거: {len(data)-len(filtered)})")
+
+random.shuffle(filtered)
+uq = len(set(d["question"] for d in filtered))
+print(f"Unique: {uq}, avg {len(filtered)/uq:.1f}/q")
+
+split = int(len(filtered) * 0.95)
+train, test = filtered[:split], filtered[split:]
+def to_sft(ex):
+    return {"prompt": [{"role":"user","content":SP+"\n\n"+ex["question"]}],
+            "completion": [{"role":"assistant","content":ex["answer"]}]}
+
+cols = [c for c in Dataset.from_list(train[:1]).column_names if c not in ["prompt","completion"]]
+train_ds = Dataset.from_list(train).map(to_sft, remove_columns=cols)
+test_ds = Dataset.from_list(test).map(to_sft, remove_columns=cols)
+print(f"학습: {len(train_ds)} / 검증: {len(test_ds)}")
+
+tokenizer = AutoTokenizer.from_pretrained("outputs/models/gemma-3-1b-it")
+model = AutoModelForCausalLM.from_pretrained("outputs/models/gemma-3-1b-it", dtype=torch.bfloat16, device_map="auto", attn_implementation='flash_attention_2')
+tokenizer.pad_token = tokenizer.eos_token
+model.gradient_checkpointing_enable(); model.config.use_cache = False
+
+cfg = SFTConfig(
+    report_to='none', seed=SEED, eval_strategy="steps", eval_steps=200,
+    save_total_limit=2, load_best_model_at_end=True, metric_for_best_model="eval_loss",
+    save_steps=200, num_train_epochs=3, warmup_ratio=0.05, weight_decay=0.01, max_grad_norm=1.0,
+    neftune_noise_alpha=5, per_device_train_batch_size=8, gradient_accumulation_steps=4,
+    per_device_eval_batch_size=2, max_length=2048, lr_scheduler_type='cosine',
+    learning_rate=2e-5, bf16=True, optim="paged_adamw_8bit",
+    output_dir="outputs/c17d_checkpoints", logging_steps=50, save_strategy="steps",
+)
+
+trainer = SFTTrainer(model=model, processing_class=tokenizer, train_dataset=train_ds, eval_dataset=test_ds, args=cfg,
+                     callbacks=[EarlyStoppingCallback(early_stopping_patience=3)])
+print("학습 시작 (3 epochs, 모든 풀이, ≤1500자)")
+r = trainer.train()
+print(f"완료! Loss: {r.training_loss:.4f}")
+
+SAVE = "outputs/models/c17d-gemma-3-1b-it-Math"
+os.makedirs(SAVE, exist_ok=True)
+model.eval(); model.save_pretrained(SAVE, safe_serialization=False); tokenizer.save_pretrained(SAVE)
+print(f"저장: {SAVE}")
+del model, trainer; torch.cuda.empty_cache()
+print("GPU 해제")