Add train_rs_sft.py

train_rs_sft.py

@@ -0,0 +1,143 @@
+"""C20: Variants of C18-2 (the 48.5% recipe) with different replay ratios"""
+import json, re, random, torch, numpy as np, os
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from trl import SFTTrainer, SFTConfig
+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}"
+BASE = "outputs/models/c17d-gemma-3-1b-it-Math"
+
+# Load RS1+RS2 (the winning combo)
+with open("outputs/c17d_rs/sft_dataset.json") as f:
+    rs1 = json.load(f)
+with open("outputs/c17d_rs2/sft_dataset.json") as f:
+    rs2 = json.load(f)
+
+seen = set()
+rs_combined = []
+for d in rs1 + rs2:
+    if d["n_correct"] < 4: continue
+    key = (d["question"], d["answer"])
+    if key not in seen:
+        seen.add(key)
+        rs_combined.append({"question": d["question"], "answer": d["answer"], "source": "gsm8k"})
+print(f"RS1+RS2 combined: {len(rs_combined)}")
+
+with open("data/GSM8K_full_qwen3_30b.json") as f:
+    orig_data = json.load(f)
+orig_filtered = [d for d in orig_data if len(d["answer"]) <= 1500]
+
+def to_sft(ex):
+    return {"prompt": [{"role":"user","content":SP+"\n\n"+ex["question"]}],
+            "completion": [{"role":"assistant","content":ex["answer"]}]}
+
+# === Condition 1: RS1+RS2 + 2x replay (more aggressive RS) ===
+print("\n=== C20-1: RS1+RS2 + 2x replay ===")
+random.seed(SEED)
+rs_qs = set(d["question"] for d in rs_combined)
+replay = [d for d in orig_filtered if d["question"] not in rs_qs]
+random.shuffle(replay)
+replay1 = replay[:int(len(rs_combined) * 2)]
+mixed1 = rs_combined + replay1
+random.shuffle(mixed1)
+print(f"  RS: {len(rs_combined)} + replay: {len(replay1)} = {len(mixed1)}")
+
+ds1 = Dataset.from_list(mixed1)
+cols = [c for c in ds1.column_names if c not in ["prompt","completion"]]
+ds1 = ds1.map(to_sft, remove_columns=cols)
+
+tokenizer = AutoTokenizer.from_pretrained(BASE)
+model = AutoModelForCausalLM.from_pretrained(BASE, 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
+
+cfg1 = SFTConfig(report_to='none', seed=SEED, num_train_epochs=1, warmup_ratio=0.05,
+    weight_decay=0.01, max_grad_norm=1.0, per_device_train_batch_size=8,
+    gradient_accumulation_steps=4, max_length=2048, lr_scheduler_type='cosine',
+    learning_rate=2e-6, bf16=True, optim="paged_adamw_8bit",
+    output_dir="outputs/c20_1_ckpt", logging_steps=25, save_strategy="no")
+trainer = SFTTrainer(model=model, processing_class=tokenizer, train_dataset=ds1, args=cfg1)
+r = trainer.train()
+print(f"  Loss: {r.training_loss:.4f}")
+
+SAVE1 = "outputs/models/c20-1-2x-replay"
+os.makedirs(SAVE1, exist_ok=True)
+model.eval(); model.save_pretrained(SAVE1, safe_serialization=False)
+tokenizer.save_pretrained(SAVE1)
+del model, trainer; torch.cuda.empty_cache()
+
+# === Condition 2: RS1+RS2 + 5x replay (more teacher data) ===
+print("\n=== C20-2: RS1+RS2 + 5x replay ===")
+random.seed(SEED)
+replay = [d for d in orig_filtered if d["question"] not in rs_qs]
+random.shuffle(replay)
+replay2 = replay[:int(len(rs_combined) * 5)]
+mixed2 = rs_combined + replay2
+random.shuffle(mixed2)
+print(f"  RS: {len(rs_combined)} + replay: {len(replay2)} = {len(mixed2)}")
+
+ds2 = Dataset.from_list(mixed2)
+cols = [c for c in ds2.column_names if c not in ["prompt","completion"]]
+ds2 = ds2.map(to_sft, remove_columns=cols)
+
+tokenizer = AutoTokenizer.from_pretrained(BASE)
+model = AutoModelForCausalLM.from_pretrained(BASE, 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
+
+cfg2 = SFTConfig(report_to='none', seed=SEED, num_train_epochs=1, warmup_ratio=0.05,
+    weight_decay=0.01, max_grad_norm=1.0, per_device_train_batch_size=8,
+    gradient_accumulation_steps=4, max_length=2048, lr_scheduler_type='cosine',
+    learning_rate=2e-6, bf16=True, optim="paged_adamw_8bit",
+    output_dir="outputs/c20_2_ckpt", logging_steps=25, save_strategy="no")
+trainer = SFTTrainer(model=model, processing_class=tokenizer, train_dataset=ds2, args=cfg2)
+r = trainer.train()
+print(f"  Loss: {r.training_loss:.4f}")
+
+SAVE2 = "outputs/models/c20-2-5x-replay"
+os.makedirs(SAVE2, exist_ok=True)
+model.eval(); model.save_pretrained(SAVE2, safe_serialization=False)
+tokenizer.save_pretrained(SAVE2)
+del model, trainer; torch.cuda.empty_cache()
+
+# === Condition 3: RS1+RS2 + 3x replay + lr=3e-6 (higher lr) ===
+print("\n=== C20-3: RS1+RS2 + 3x replay + lr=3e-6 ===")
+random.seed(SEED)
+replay = [d for d in orig_filtered if d["question"] not in rs_qs]
+random.shuffle(replay)
+replay3 = replay[:int(len(rs_combined) * 3)]
+mixed3 = rs_combined + replay3
+random.shuffle(mixed3)
+print(f"  RS: {len(rs_combined)} + replay: {len(replay3)} = {len(mixed3)}")
+
+ds3 = Dataset.from_list(mixed3)
+cols = [c for c in ds3.column_names if c not in ["prompt","completion"]]
+ds3 = ds3.map(to_sft, remove_columns=cols)
+
+tokenizer = AutoTokenizer.from_pretrained(BASE)
+model = AutoModelForCausalLM.from_pretrained(BASE, 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
+
+cfg3 = SFTConfig(report_to='none', seed=SEED, num_train_epochs=1, warmup_ratio=0.05,
+    weight_decay=0.01, max_grad_norm=1.0, per_device_train_batch_size=8,
+    gradient_accumulation_steps=4, max_length=2048, lr_scheduler_type='cosine',
+    learning_rate=3e-6, bf16=True, optim="paged_adamw_8bit",
+    output_dir="outputs/c20_3_ckpt", logging_steps=25, save_strategy="no")
+trainer = SFTTrainer(model=model, processing_class=tokenizer, train_dataset=ds3, args=cfg3)
+r = trainer.train()
+print(f"  Loss: {r.training_loss:.4f}")
+
+SAVE3 = "outputs/models/c20-3-lr3e-6"
+os.makedirs(SAVE3, exist_ok=True)
+model.eval(); model.save_pretrained(SAVE3, safe_serialization=False)
+tokenizer.save_pretrained(SAVE3)
+del model, trainer; torch.cuda.empty_cache()
+
+print("\n=== All conditions complete ===")