Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/combined_preference.jsonl filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,162 +1,146 @@
 ---
 language:
-- ko
-- en
+  - ko
+  - en
 license: mit
 tags:
-- mamba2
-- hybrid
-- korean
-- causal-lm
+  - mamba2
+  - hybrid
+  - transformer
+  - korean
+  - from-scratch
+  - dpo
+  - slerp
+  - orpo
+library_name: pytorch
 pipeline_tag: text-generation
 ---
 
-# EVAFRILL-Mo-3B
+# EVAFRILL-Mo 3B — Hybrid Mamba-2 + Transformer
 
-EVAFRILL-Mo-3B is a 2.94B-parameter **hybrid Mamba-2 + Transformer** language
-model optimised for Korean, trained from scratch on 55 billion tokens of
-Korean-dominant multilingual text.
+**A 3-billion-parameter hybrid Mamba-2 + Transformer language model built from scratch.**
 
-> **EVAFRILL-Mo** stands for *Efficient Variably-Architected Fusion of
-> Recurrent and Integrated Linear Layers for Language Model-based Output* — a
-> custom architecture inspired by [Nemotron-H](https://arxiv.org/abs/2501.14587)
-> that replaces most self-attention layers with Mamba-2 SSM blocks, achieving
-> linear-time inference without sacrificing generation quality.
+Inspired by the NVIDIA [Nemotron-H](https://arxiv.org/abs/2504.03624) architecture. Pretrained on 55B tokens across Korean, English, code, and math using 7× NVIDIA B200 GPUs.
 
----
+## Model Variants
+
+This repository contains **7 model versions** representing each stage of the training pipeline, plus training data and scripts for full reproducibility.
+
+| Variant | Directory | Size | Description | Recommended |
+|---------|-----------|------|-------------|:-----------:|
+| **SLERP** | `slerp/` | 6.3GB | SFT + DPO merged (α=0.5) | ⭐ **Yes** |
+| Pretrain | `pretrain/` | 12.6GB | Base model (319K steps, 55B tokens) | |
+| SFT v2 | `sft-v2/` | 6.3GB | Instruction-tuned (65K steps) | |
+| DPO R1 | `dpo-r1/` | 6.3GB | Preference-aligned Round 1 | |
+| DPO R2 | `dpo-r2/` | 6.3GB | Conservative fine-tuning Round 2 | |
+| ORPO | `orpo/` | 6.3GB | SFT+alignment simultaneous (experimental) | |
+| DPO R3 | `dpo-r3/` | 6.3GB | Repetition-targeted (experimental) | |
+
+## Training Pipeline
+
+```
+Pretrain (55B tokens, 7×B200, 60h)
+    ↓
+SFT v2 (65K steps, H100 MIG, 5 days)
+    ↓
+DPO Round 1 (3K steps, LoRA, loss 0.693→0.565)
+    ↓
+DPO Round 2 (2K steps, conservative, loss 0.692→0.689)
+    ↓
+SLERP Merge (α=0.5, SFT 50% + DPO 50%)  ← RECOMMENDED
+    ↓
+ORPO Experiment (10K steps, alternative approach)
+    ↓
+DPO Round 3 (1K steps, repetition-targeted experiment)
+```
 
 ## Architecture
 
-| Property | Value |
-|---|---|
-| Total parameters | ~2.94 B |
-| Layers | 26 (24 × Mamba-2 + 2 × Attention) |
-| Hidden size | 3072 |
-| Attention heads | 24 (GQA, 8 KV heads) |
-| FFN size | 9 216 |
-| Mamba-2 state dim | 128 |
-| Mamba-2 head dim | 64 |
-| Vocab size | 64 000 |
-| Max sequence length | 4 096 |
-| RoPE theta | 500 000 |
-
-The layer pattern places attention blocks at positions 9 and 18 (zero-indexed),
-mirroring the Nemotron-H 8B dense design scaled to 3B parameters. All other
-layers use Mamba-2 with SwiGLU FFN (mamba_d_ffn = 4 608). Attention layers use
-full SwiGLU FFN (d_ffn = 9 216).
+```
+Type:           Hybrid Mamba-2 + Transformer
+Parameters:     2.94B (2,975,397,632)
+Layers:         26 (24× Mamba-2 SSM + 2× Attention GQA)
+d_model:        3,072
+Vocabulary:     64,000 (custom SentencePiece)
+Max seq length: 4,096
+```
 
----
+## Benchmark Results (SLERP, recommended model)
 
-## Training
-
-### Pretraining
-- **Tokens**: 55 B (319 772 steps, effective batch ≈ 172 K tokens)
-- **Hardware**: 8× NVIDIA B200 (183 GB each), ~62 hours
-- **Optimizer**: AdamW, lr=2e-4, cosine decay, warmup 2 000 steps
-- **Precision**: FP8 (TransformerEngine MXFP8) + BF16 embedding
-- **Data**: Korean web corpus, Wikipedia, books, code (Korean-dominant)
-
-### Supervised Fine-Tuning (SFT)
-- **Steps**: 65 000 (≈ 1 epoch on 2.44M instruction samples)
-- **Effective batch**: 56 (2 per GPU × 7 GPU × 4 grad_accum)
-- **LR**: 1e-5 (pretrain/30, catastrophic-forgetting guard)
-- **NEFTune alpha**: 5.0 (repetition degeneracy mitigation)
-- **Data**: Combined Korean instruction set (filtered, 2.44M samples)
-
-### Direct Preference Optimisation (DPO)
-- **Rounds**: 2-round DPO (Nemotron-H style)
-  - Round 1: 3 000 steps, beta=0.1, lr=5e-7, LoRA rank=32
-  - Round 2: 2 000 steps, beta=0.05, lr=1e-7, LoRA rank=32
-- **Hardware**: 1× NVIDIA H100 MIG 3g.40gb (~42 GB VRAM)
-- **Method**: Native LoRA DPO (no TRL dependency)
-
-### SLERP Merge
-The final checkpoint is produced by **spherical linear interpolation (SLERP)**
-between the SFT-v2 and DPO-round-2 checkpoints (ratio 0.5), combining the
-instruction-following strengths of both stages.
+| Metric | Value |
+|--------|-------|
+| Greedy 3-gram repetition | 74.5% (→ 5.5% with rep_penalty=1.2) |
+| hellaswag (0-shot) | 34.6% |
+| arc_easy (0-shot) | 32.0% |
+| belebele_kor (0-shot) | 23.6% |
+| global_mmlu_ko (0-shot) | 23.7% |
 
----
+**Recommended inference**: `temperature=0.7, repetition_penalty=1.2`
 
-## Evaluation (SLERP checkpoint, lm-eval-harness)
+## SFT→DPO→SLERP vs ORPO Comparison
 
-| Benchmark | Metric | Score |
-|---|---|---|
-| HellaSwag | acc_norm | 0.42 |
-| ARC-Challenge | acc_norm | 0.22 |
-| ARC-Easy | acc_norm | 0.28 |
-| Belebele (kor_Hang) | acc | 0.30 |
-| Global-MMLU-ko (full) | acc | 0.233 |
-|  — Humanities | acc | 0.242 |
-|  — STEM | acc | 0.237 |
-|  — Social Sciences | acc | 0.221 |
-|  — Other | acc | 0.229 |
+| Metric | SLERP | ORPO | Winner |
+|--------|:-----:|:----:|:------:|
+| Greedy repetition | **74.5%** | 87.1% | SLERP |
+| Chat quality | ✅ Fluent | ❌ Broken | SLERP |
+| hellaswag | **39.0%** | 35.0% | SLERP |
+| Training time | 5d+8h | **12.8h** | ORPO |
 
-*Evaluated on 100-sample subsets per task. Numbers reflect the final
-SLERP-merged checkpoint.*
+ORPO's weakness: insufficient SFT learning at 10K steps (vs SFT's 65K).
 
----
+## Repository Structure
+
+```
+├── slerp/                    # ⭐ Recommended final model
+├── pretrain/                 # Base pretrained model
+├── sft-v2/                   # SFT instruction-tuned
+├── dpo-r1/                   # DPO Round 1 + LoRA weights
+├── dpo-r2/                   # DPO Round 2 + LoRA weights
+├── orpo/                     # ORPO experiment + LoRA weights
+├── dpo-r3/                   # DPO Round 3
+├── data/                     # Preference datasets for reproducibility
+│   ├── combined_preference.jsonl    (684K pairs, 2.6GB)
+│   └── repetition_preference.jsonl  (105 pairs, self-generated)
+├── configs/                  # Training YAML configs
+│   ├── korean_3b_sft_1gpu.yaml
+│   ├── dpo_3b_1gpu.yaml
+│   └── orpo_3b_1gpu.yaml
+└── scripts/                  # Training & evaluation code
+    ├── dpo.py, orpo_native.py, sft.py
+    ├── lora.py, merge_checkpoints.py
+    ├── evafrill_eval.py
+    └── generate_repetition_preference.py
+```
 
 ## Usage
 
 ```python
-# Requires: pip install transformers tokenizers safetensors
-from transformers import AutoTokenizer, AutoModelForCausalLM
+# This is a custom architecture — use the project's native loading code
+# Clone: https://github.com/pathcosmos/EVAFRILL-Mo
+
 import torch
+from model.transformer import LLM
+from tokenizers import Tokenizer
 
-model_id = "pathcosmos/EVAFRILL-Mo-3B"
-
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-)
-
-# Chat-style prompt
-prompt = "<|user|>\n안녕하세요! 자기소개를 해 주세요.\n<|assistant|>\n"
-inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-with torch.no_grad():
-    output = model.generate(
-        **inputs,
-        max_new_tokens=256,
-        temperature=0.8,
-        top_p=0.9,
-        do_sample=True,
-        repetition_penalty=1.1,
-    )
-
-print(tokenizer.decode(output[0], skip_special_tokens=False))
-```
+model = LLM.from_pretrained("checkpoints/3b_dpo/checkpoint-slerp")
+model = model.to(device="cuda:0", dtype=torch.bfloat16)
+model.eval()
 
----
+tok = Tokenizer.from_file("tokenizer/korean_sp/tokenizer.json")
+```
 
 ## Limitations
 
-- This model is an **experimental research checkpoint**, not a production system.
-- Korean is the dominant language; English and other languages are secondary.
-- The custom architecture (`evafrill-mo`) requires either
-  (a) the original project code for full inference, or
-  (b) a compatible HuggingFace integration that understands Mamba-2 hybrid layers.
-  The exported `model.safetensors` preserves the native weight layout.
-- Benchmark numbers were evaluated on small (100-sample) subsets and should be
-  treated as rough estimates.
-
----
-
-## Citation
+- **3B scale**: Factual accuracy and complex reasoning are limited
+- **GGUF/Ollama**: Not possible due to custom hybrid Mamba-2 architecture
+- **vLLM**: Theoretically possible but requires custom weight key mapping
+- **Greedy repetition**: ~74.5% without rep_penalty (use rep_penalty=1.2)
 
-```bibtex
-@misc{evafrill-mo-3b-2026,
-  title  = {EVAFRILL-Mo-3B: A Hybrid Mamba-2 + Transformer LLM for Korean},
-  author = {pathcosmos},
-  year   = {2026},
-  url    = {https://huggingface.co/pathcosmos/EVAFRILL-Mo-3B},
-}
-```
+## Links
 
----
+- **GitHub**: [pathcosmos/EVAFRILL-Mo](https://github.com/pathcosmos/EVAFRILL-Mo)
+- **Paper reference**: [Nemotron-H](https://arxiv.org/abs/2504.03624)
 
 ## License
 
-[MIT](LICENSE)
+MIT

configs/dpo_3b_1gpu.yaml

@@ -0,0 +1,31 @@
+# EVAFRILL-Mo 3B DPO — Single GPU (H100 MIG 3g.40gb, 42.3GB VRAM)
+#
+# Base model: checkpoints/3b_sft_v2/checkpoint-best (SFT 완료 모델)
+# Method: LoRA DPO (native, TRL 미사용)
+#
+# [설계 근거]
+#   - GPU: H100 PCIe MIG 3g.40gb (42.3GB VRAM)
+#   - LoRA DPO VRAM 예산: base(6GB) + LoRA(0.3GB) + optim(0.2GB) + act(10GB) + ref_fwd(6GB) ≈ 22GB
+#   - BF16 + Gradient Checkpointing (FP8 미지원)
+#   - eff_batch: 1 × 16 grad_accum = 16
+#   - Nemotron-H 스타일 2-round DPO
+
+train:
+  # Round 1 설정 (Round 2는 max_steps=2000, beta=0.05, lr=1e-7로 변경)
+  max_steps: 3000
+  batch_size: 1
+  grad_accum_steps: 16     # eff_batch = 16
+  lr: 5.0e-7               # DPO는 SFT보다 훨씬 낮은 lr
+  weight_decay: 0.01
+  warmup_steps: 100
+  max_length: 1024          # VRAM 제약으로 seq_len 제한
+  beta: 0.1                 # DPO temperature
+
+  # LoRA 설정
+  use_lora: true
+  lora_rank: 32
+  lora_alpha: 64
+
+  # 저장/로깅
+  save_interval: 500
+  log_interval: 10

configs/korean_3b_sft_1gpu.yaml

@@ -0,0 +1,31 @@
+# EVAFRILL-Mo 3B SFT — Single GPU (H100 MIG 3g.40gb, 42.3GB VRAM)
+#
+# Base model: checkpoints/3b_final/checkpoint-0319772
+# Fresh start from pretrained checkpoint
+#
+# [설계 근거 — 2026-03-17, 최적화 2026-03-17]
+#   - GPU: H100 PCIe MIG 3g.40gb (42.3GB VRAM, 46 SMs)
+#   - CPU: 45 cores (cgroup), RAM: 200GB (cgroup)
+#   - BF16 + Gradient Checkpointing (no FP8, MIG NVML 제약)
+#   - 벤치마크 결과: bs=4 ga=7 @ 27.7GB VRAM (68.7%), 5,475 tok/s (+10% vs bs=1)
+#   - eff_batch: 4 × 1GPU × 7 grad_accum = 28
+#   - 1 epoch: 3,774,413 / 28 ≈ 134,800 steps → max_steps=135000
+#   - 예상 시간: 135,000 steps × ~10.5s/step ≈ ~391 hours ≈ ~16 days
+
+train:
+  max_steps: 135000      # ≈ 1 epoch on 3.77M samples, eff_batch=28
+  batch_size: 4          # 벤치마크 최적: bs=4 @ 27.7GB (68.7% VRAM)
+  grad_accum_steps: 7    # eff_batch=28 (4×7), bs=1→4 전환으로 tok/s +10%
+  lr: 7.0e-6             # sqrt(28/56) * 1e-5 ≈ 7e-6 (linear scaling rule)
+  weight_decay: 0.01
+  warmup_steps: 500
+  max_grad_norm: 1.0
+  log_interval: 10
+  save_interval: 2000    # 크래시 복구 위해 자주 저장
+  eval_interval: 5000    # validation 1회 ~5분, 부담 최소화
+  neftune_alpha: 5.0     # NEFTune noise injection (반복 퇴화 완화)
+  max_val_batches: 500   # validation 배치 수 제한 (속도 최적화)
+
+tokenizer:
+  vocab_size: 64000
+  type: sentencepiece_unigram

configs/orpo_3b_1gpu.yaml

@@ -0,0 +1,27 @@
+# EVAFRILL-Mo 3B ORPO — Single GPU (H100 MIG 3g.40gb, 42.3GB VRAM)
+#
+# Base model: checkpoints/3b_final/checkpoint-0319772 (Pretrained, NOT SFT)
+# Method: ORPO (SFT + Odds Ratio Preference) with LoRA
+#
+# [설계 근거]
+#   - ORPO는 SFT+정렬을 동시에 학습 → pretrained 모델에서 시작
+#   - Reference model 불필요 → DPO보다 VRAM 절약
+#   - LoRA rank=32: base(6GB) + LoRA(0.3GB) + optim(0.2GB) + act(~8GB) ≈ 15GB
+#   - eff_batch: 1 × 16 grad_accum = 16
+
+train:
+  max_steps: 10000
+  batch_size: 1
+  grad_accum_steps: 16
+  lr: 5.0e-6
+  weight_decay: 0.01
+  warmup_steps: 500
+  max_length: 1024
+  lambda_or: 1.0
+
+  use_lora: true
+  lora_rank: 32
+  lora_alpha: 64
+
+  save_interval: 1000
+  log_interval: 10

data/combined_preference.jsonl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32b0c2a5ca3a523a22882c2d828917a3eb543605f04b51db2bab16e6bd262f95
+size 2721356504

dpo-r1/config.json

@@ -0,0 +1,29 @@
+{
+  "vocab_size": 64000,
+  "d_model": 3072,
+  "n_layers": 26,
+  "n_heads": 24,
+  "n_kv_heads": 8,
+  "d_ffn": 9216,
+  "max_seq_len": 4096,
+  "rope_theta": 500000.0,
+  "dropout": 0.0,
+  "bias": false,
+  "use_flash_attn": true,
+  "use_fp8": false,
+  "use_hybrid": true,
+  "hybrid_pattern": "M M M M M M M M M M M M A M M M M M M M M M M M A M",
+  "mamba_d_state": 128,
+  "mamba_head_dim": 64,
+  "mamba_expand": 2,
+  "mamba_conv_kernel": 4,
+  "mamba_n_groups": 8,
+  "mamba_d_ffn": 4608,
+  "mamba_chunk_size": 256,
+  "model_type": "evafrill-mo",
+  "architectures": [
+    "EvafrillMoForCausalLM"
+  ],
+  "_variant": "dpo-r1",
+  "_description": "DPO Round 1 (3K steps, loss 0.693->0.565)"
+}

dpo-r1/lora_weights.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1f161ffc4138d61116dedb9e29176497ff85adbed7374a1b7fd35f9672d21245
+size 42909589

dpo-r1/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a07011b219031b7cab8985a4c0dc811aa4758f44c79694da12b744059f77cd99
+size 6301164272

dpo-r2/config.json

@@ -0,0 +1,29 @@
+{
+  "vocab_size": 64000,
+  "d_model": 3072,
+  "n_layers": 26,
+  "n_heads": 24,
+  "n_kv_heads": 8,
+  "d_ffn": 9216,
+  "max_seq_len": 4096,
+  "rope_theta": 500000.0,
+  "dropout": 0.0,
+  "bias": false,
+  "use_flash_attn": true,
+  "use_fp8": false,
+  "use_hybrid": true,
+  "hybrid_pattern": "M M M M M M M M M M M M A M M M M M M M M M M M A M",
+  "mamba_d_state": 128,
+  "mamba_head_dim": 64,
+  "mamba_expand": 2,
+  "mamba_conv_kernel": 4,
+  "mamba_n_groups": 8,
+  "mamba_d_ffn": 4608,
+  "mamba_chunk_size": 256,
+  "model_type": "evafrill-mo",
+  "architectures": [
+    "EvafrillMoForCausalLM"
+  ],
+  "_variant": "dpo-r2",
+  "_description": "DPO Round 2 (2K steps, conservative fine-tuning)"
+}

dpo-r2/lora_weights.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eeda7e71c6d2f69ea5ed8a02759ee2ac1f6feadcbf6e440fd3f9da919628f947
+size 42909589

dpo-r2/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c64ad3c2c26b3e706e58011a8a3eb8dba3f5cee2b0ea62eaa0083ad3c4b7685e
+size 6301164272

dpo-r3/config.json

@@ -0,0 +1,29 @@
+{
+  "vocab_size": 64000,
+  "d_model": 3072,
+  "n_layers": 26,
+  "n_heads": 24,
+  "n_kv_heads": 8,
+  "d_ffn": 9216,
+  "max_seq_len": 4096,
+  "rope_theta": 500000.0,
+  "dropout": 0.0,
+  "bias": false,
+  "use_flash_attn": true,
+  "use_fp8": false,
+  "use_hybrid": true,
+  "hybrid_pattern": "M M M M M M M M M M M M A M M M M M M M M M M M A M",
+  "mamba_d_state": 128,
+  "mamba_head_dim": 64,
+  "mamba_expand": 2,
+  "mamba_conv_kernel": 4,
+  "mamba_n_groups": 8,
+  "mamba_d_ffn": 4608,
+  "mamba_chunk_size": 256,
+  "model_type": "evafrill-mo",
+  "architectures": [
+    "EvafrillMoForCausalLM"
+  ],
+  "_variant": "dpo-r3",
+  "_description": "DPO Round 3 (1K steps, repetition-targeted experiment)"
+}

dpo-r3/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d49c4459ae2db0f5fb8ef761c94afd06b6f518e07e7208407effe1298fa9bc20
+size 6301164272

orpo/config.json

@@ -0,0 +1,29 @@
+{
+  "vocab_size": 64000,
+  "d_model": 3072,
+  "n_layers": 26,
+  "n_heads": 24,
+  "n_kv_heads": 8,
+  "d_ffn": 9216,
+  "max_seq_len": 4096,
+  "rope_theta": 500000.0,
+  "dropout": 0.0,
+  "bias": false,
+  "use_flash_attn": true,
+  "use_fp8": false,
+  "use_hybrid": true,
+  "hybrid_pattern": "M M M M M M M M M M M M A M M M M M M M M M M M A M",
+  "mamba_d_state": 128,
+  "mamba_head_dim": 64,
+  "mamba_expand": 2,
+  "mamba_conv_kernel": 4,
+  "mamba_n_groups": 8,
+  "mamba_d_ffn": 4608,
+  "mamba_chunk_size": 256,
+  "model_type": "evafrill-mo",
+  "architectures": [
+    "EvafrillMoForCausalLM"
+  ],
+  "_variant": "orpo",
+  "_description": "ORPO experiment (10K steps, SFT+alignment simultaneous)"
+}

orpo/lora_weights.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a06dd455a91aa303be57a004db0ef2889027cf9b2a7854137e7e7a0aefbfeeda
+size 42909589

orpo/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:46b8ee870223a391c08ff08a363ac68948df0b13971910fb9ba3c85678f3f6e5
+size 6301164272

pretrain/config.json

@@ -0,0 +1,29 @@
+{
+  "vocab_size": 64000,
+  "d_model": 3072,
+  "n_layers": 26,
+  "n_heads": 24,
+  "n_kv_heads": 8,
+  "d_ffn": 9216,
+  "max_seq_len": 4096,
+  "rope_theta": 500000.0,
+  "dropout": 0.0,
+  "bias": false,
+  "use_flash_attn": true,
+  "use_fp8": true,
+  "use_hybrid": true,
+  "hybrid_pattern": "M M M M M M M M M M M M A M M M M M M M M M M M A M",
+  "mamba_d_state": 128,
+  "mamba_head_dim": 64,
+  "mamba_expand": 2,
+  "mamba_conv_kernel": 4,
+  "mamba_n_groups": 8,
+  "mamba_d_ffn": 4608,
+  "mamba_chunk_size": 256,
+  "model_type": "evafrill-mo",
+  "architectures": [
+    "EvafrillMoForCausalLM"
+  ],
+  "_variant": "pretrain",
+  "_description": "Pretrained base model (319K steps, 55B tokens, Chinchilla 93%)"
+}

pretrain/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b4e8796eb71489901cded57f3981889b2cd57f06552c03cf814c15fc52ad69df
+size 12602306810

scripts/dpo.py

@@ -0,0 +1,477 @@
+"""
+train/dpo.py — Direct Preference Optimization (DPO) training.
+
+Native DPO implementation (no TRL dependency) for EVAFRILL-Mo hybrid models.
+Supports LoRA adapters for memory-efficient training on single GPU.
+
+Launch:
+    python train/dpo.py \
+        --sft_checkpoint checkpoints/3b_sft_v2/checkpoint-best \
+        --dpo_data data/preference/combined_preference.jsonl \
+        --config configs/h100_mig/dpo_3b_1gpu.yaml \
+        --device cuda:0
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import random
+import signal
+import shutil
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, RandomSampler
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+torch.set_float32_matmul_precision("high")
+
+_PROJECT_ROOT = Path(__file__).resolve().parent.parent
+if str(_PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(_PROJECT_ROOT))
+
+from model import LLM
+from model.lora import apply_lora, get_lora_params, merge_lora, save_lora
+from data.dpo_dataset import DPODataset, dpo_collate_fn
+from train.utils import (
+    get_cosine_schedule_with_warmup,
+    is_main_process,
+    save_checkpoint,
+    load_checkpoint,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="DPO Training for EVAFRILL-Mo")
+
+    # Paths
+    parser.add_argument("--sft_checkpoint", type=Path, required=True,
+                        help="Path to SFT checkpoint directory")
+    parser.add_argument("--dpo_data", type=Path, required=True,
+                        help="Path to preference JSONL data")
+    parser.add_argument("--checkpoint_dir", type=Path, default=Path("checkpoints/3b_dpo"),
+                        help="Output checkpoint directory")
+    parser.add_argument("--resume", type=Path, default=None)
+    parser.add_argument("--tokenizer", type=Path, default=None)
+    parser.add_argument("--log_file", type=Path, default=None)
+    parser.add_argument("--config", type=Path, default=None)
+
+    # DPO hyperparameters
+    parser.add_argument("--beta", type=float, default=0.1, help="DPO temperature")
+    parser.add_argument("--max_steps", type=int, default=3000)
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--grad_accum", type=int, default=16)
+    parser.add_argument("--lr", type=float, default=5e-7)
+    parser.add_argument("--weight_decay", type=float, default=0.01)
+    parser.add_argument("--warmup_steps", type=int, default=100)
+    parser.add_argument("--max_length", type=int, default=1024)
+    parser.add_argument("--seed", type=int, default=42)
+
+    # LoRA
+    parser.add_argument("--use_lora", action="store_true", default=True)
+    parser.add_argument("--lora_rank", type=int, default=32)
+    parser.add_argument("--lora_alpha", type=float, default=64.0)
+
+    # Infra
+    parser.add_argument("--device", type=str, default=None)
+    parser.add_argument("--save_interval", type=int, default=500)
+    parser.add_argument("--log_interval", type=int, default=10)
+    parser.add_argument("--num_workers", type=int, default=4)
+
+    args, _ = parser.parse_known_args()
+
+    # Load YAML config
+    if args.config is not None:
+        if not args.config.exists():
+            raise FileNotFoundError(f"Config not found: {args.config}")
+        import yaml
+        with open(args.config) as f:
+            cfg = yaml.safe_load(f)
+        train_cfg = cfg.get("train", {})
+        yaml_map = {
+            "max_steps": "max_steps", "batch_size": "batch_size",
+            "grad_accum_steps": "grad_accum", "lr": "lr",
+            "weight_decay": "weight_decay", "warmup_steps": "warmup_steps",
+            "beta": "beta", "max_length": "max_length",
+            "save_interval": "save_interval", "log_interval": "log_interval",
+            "use_lora": "use_lora", "lora_rank": "lora_rank", "lora_alpha": "lora_alpha",
+        }
+        defaults = {}
+        for yk, ak in yaml_map.items():
+            if yk in train_cfg:
+                defaults[ak] = train_cfg[yk]
+        if defaults:
+            parser.set_defaults(**defaults)
+
+    return parser.parse_args()
+
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def compute_log_probs(
+    model: nn.Module,
+    input_ids: torch.Tensor,
+    labels: torch.Tensor,
+) -> torch.Tensor:
+    """Compute sum of log probabilities over non-masked tokens.
+
+    Args:
+        model: The LLM model
+        input_ids: (B, T) token ids
+        labels: (B, T) target ids, -1 for masked positions
+
+    Returns:
+        (B,) sum of log probs per sample
+    """
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        logits, _ = model(input_ids)  # (B, T, V)
+
+    # Shift: predict next token
+    # logits[:, :-1] predicts labels[:, 1:]
+    # But our labels already have the shifted targets (same as SFT convention)
+    # labels[i] = token_id means input_ids[i] should predict labels[i]
+    log_probs = F.log_softmax(logits.float(), dim=-1)  # (B, T, V)
+
+    # Gather log probs for target tokens
+    # For each position, get log_prob of the label token
+    mask = labels != -1  # (B, T)
+    # Clamp labels for gather (replace -1 with 0, will be masked out)
+    safe_labels = labels.clamp(min=0)  # (B, T)
+    per_token_logps = log_probs.gather(-1, safe_labels.unsqueeze(-1)).squeeze(-1)  # (B, T)
+    per_token_logps = per_token_logps * mask.float()  # zero out masked positions
+
+    return per_token_logps.sum(dim=-1)  # (B,)
+
+
+def dpo_loss(
+    policy_chosen_logps: torch.Tensor,
+    policy_rejected_logps: torch.Tensor,
+    ref_chosen_logps: torch.Tensor,
+    ref_rejected_logps: torch.Tensor,
+    beta: float = 0.1,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Compute DPO loss.
+
+    Returns:
+        (loss, chosen_rewards, rejected_rewards)
+    """
+    chosen_rewards = beta * (policy_chosen_logps - ref_chosen_logps)
+    rejected_rewards = beta * (policy_rejected_logps - ref_rejected_logps)
+
+    logits = chosen_rewards - rejected_rewards  # (B,)
+    loss = -F.logsigmoid(logits).mean()
+
+    return loss, chosen_rewards.detach().mean(), rejected_rewards.detach().mean()
+
+
+def _resolve_tokenizer_path(args: argparse.Namespace) -> Path:
+    if args.tokenizer is not None:
+        return Path(args.tokenizer)
+    ckpt_tok = args.sft_checkpoint / "tokenizer.json"
+    if ckpt_tok.exists():
+        return ckpt_tok
+    default_tok = _PROJECT_ROOT / "tokenizer" / "korean_sp" / "tokenizer.json"
+    if default_tok.exists():
+        return default_tok
+    raise FileNotFoundError("Cannot find tokenizer.json")
+
+
+def main() -> None:
+    args = parse_args()
+    set_seed(args.seed)
+
+    # Device setup
+    if args.device:
+        device = torch.device(args.device)
+    elif torch.cuda.is_available():
+        device = torch.device("cuda:0")
+    else:
+        device = torch.device("cpu")
+
+    # Validate checkpoint
+    if not args.sft_checkpoint.exists():
+        raise FileNotFoundError(f"SFT checkpoint not found: {args.sft_checkpoint}")
+
+    # Load SFT model as policy
+    print(f"Loading SFT model from {args.sft_checkpoint}...")
+    model = LLM.from_pretrained(args.sft_checkpoint)
+    model.config.use_fp8 = False  # H100 MIG: BF16 only
+    model = model.to(device=device, dtype=torch.bfloat16)
+
+    # Enable gradient checkpointing
+    if hasattr(model, 'gradient_checkpointing_enable'):
+        model.gradient_checkpointing_enable()
+        print("[INFO] Gradient checkpointing enabled")
+
+    # Compute reference log probs BEFORE applying LoRA
+    # (reference model = SFT model without LoRA)
+    # We'll compute ref logps on-the-fly with LoRA disabled via a context manager
+    # Actually for simplicity: precompute nothing, just use model without LoRA adapters
+    # For LoRA DPO: ref_model is the base (original weights), policy is base + LoRA
+    # Since LoRA is initialized to zero, at start policy = ref
+
+    # Apply LoRA
+    if args.use_lora:
+        n_lora_params = apply_lora(model, rank=args.lora_rank, alpha=args.lora_alpha)
+        lora_params = get_lora_params(model)
+        print(f"[INFO] LoRA: {n_lora_params:,} trainable params")
+    else:
+        # Full fine-tuning (risky for VRAM)
+        lora_params = None
+
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"Total params: {total_params:,}, Trainable: {trainable_params:,}")
+
+    # Tokenizer
+    tokenizer_path = _resolve_tokenizer_path(args)
+    print(f"Loading tokenizer from {tokenizer_path}")
+    from tokenizers import Tokenizer
+    tokenizer = Tokenizer.from_file(str(tokenizer_path))
+
+    # Dataset
+    train_dataset = DPODataset(
+        data_path=args.dpo_data,
+        tokenizer=tokenizer,
+        max_seq_len=args.max_length,
+    )
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=args.batch_size,
+        sampler=RandomSampler(train_dataset),
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=True,
+        collate_fn=dpo_collate_fn,
+        prefetch_factor=2,
+        persistent_workers=True,
+    )
+
+    # Optimizer — only LoRA params if using LoRA
+    if lora_params is not None:
+        optimizer = torch.optim.AdamW(
+            lora_params,
+            lr=args.lr,
+            betas=(0.9, 0.95),
+            weight_decay=args.weight_decay,
+            fused=torch.cuda.is_available(),
+        )
+    else:
+        optimizer = torch.optim.AdamW(
+            [p for p in model.parameters() if p.requires_grad],
+            lr=args.lr,
+            betas=(0.9, 0.95),
+            weight_decay=args.weight_decay,
+            fused=torch.cuda.is_available(),
+        )
+
+    scheduler = get_cosine_schedule_with_warmup(
+        optimizer=optimizer,
+        warmup_steps=args.warmup_steps,
+        total_steps=args.max_steps,
+    )
+
+    # Resume
+    start_step = 0
+    if args.resume is not None:
+        start_step, _ = load_checkpoint(args.resume, model, optimizer, scheduler)
+        print(f"Resumed from step {start_step}")
+
+    # Checkpoint dir
+    args.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+
+    # Copy tokenizer
+    dest_tok = args.checkpoint_dir / "tokenizer.json"
+    if not dest_tok.exists():
+        shutil.copy2(str(tokenizer_path), str(dest_tok))
+
+    # Log file
+    log_fh = None
+    if args.log_file:
+        Path(args.log_file).parent.mkdir(parents=True, exist_ok=True)
+        log_fh = open(args.log_file, "a", encoding="utf-8", buffering=1)
+
+    def log(msg: str, level: str = "INFO"):
+        import datetime
+        ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        line = f"[{ts}] [{level}] {msg}"
+        print(line)
+        if log_fh:
+            log_fh.write(line + "\n")
+
+    # Banner
+    eff_batch = args.batch_size * args.grad_accum
+    log(f"{'='*60}")
+    log(f"DPO Training — EVAFRILL-Mo 3B")
+    log(f"  SFT ckpt: {args.sft_checkpoint}")
+    log(f"  DPO data: {args.dpo_data} ({len(train_dataset):,} samples)")
+    log(f"  LoRA: rank={args.lora_rank} alpha={args.lora_alpha}")
+    log(f"  beta={args.beta}, lr={args.lr:.2e}, eff_batch={eff_batch}")
+    log(f"  max_steps={args.max_steps}, max_length={args.max_length}")
+    log(f"  device={device}")
+    log(f"{'='*60}")
+
+    # Training loop
+    import time
+    model.train()
+    loader_iter = iter(train_loader)
+    epoch = 0
+
+    def next_batch():
+        nonlocal loader_iter, epoch
+        try:
+            return next(loader_iter)
+        except StopIteration:
+            epoch += 1
+            loader_iter = iter(train_loader)
+            return next(loader_iter)
+
+    shutdown_requested = False
+    def shutdown_handler(signum, frame):
+        nonlocal shutdown_requested
+        shutdown_requested = True
+        log(f"Shutdown signal received ({signum})", "WARN")
+
+    signal.signal(signal.SIGHUP, shutdown_handler)
+    signal.signal(signal.SIGTERM, shutdown_handler)
+
+    t0 = time.perf_counter()
+    running_loss = 0.0
+    running_chosen_reward = 0.0
+    running_rejected_reward = 0.0
+    log_step_count = 0
+
+    for step in range(start_step, args.max_steps):
+        optimizer.zero_grad(set_to_none=True)
+        accum_loss = 0.0
+
+        for micro in range(args.grad_accum):
+            batch = next_batch()
+            chosen_ids = batch[0].to(device, dtype=torch.long, non_blocking=True)
+            chosen_labels = batch[1].to(device, dtype=torch.long, non_blocking=True)
+            rejected_ids = batch[2].to(device, dtype=torch.long, non_blocking=True)
+            rejected_labels = batch[3].to(device, dtype=torch.long, non_blocking=True)
+
+            # Policy log probs (with LoRA active)
+            policy_chosen_logps = compute_log_probs(model, chosen_ids, chosen_labels)
+            policy_rejected_logps = compute_log_probs(model, rejected_ids, rejected_labels)
+
+            # Reference log probs (LoRA disabled)
+            # For LoRA: temporarily set lora scaling to 0
+            with torch.no_grad():
+                # Save and zero LoRA params
+                if args.use_lora:
+                    saved_B = []
+                    for m in model.modules():
+                        from model.lora import LoRALinear
+                        if isinstance(m, LoRALinear):
+                            saved_B.append(m.lora_B.data.clone())
+                            m.lora_B.data.zero_()
+
+                ref_chosen_logps = compute_log_probs(model, chosen_ids, chosen_labels)
+                ref_rejected_logps = compute_log_probs(model, rejected_ids, rejected_labels)
+
+                # Restore LoRA params
+                if args.use_lora:
+                    idx = 0
+                    for m in model.modules():
+                        from model.lora import LoRALinear
+                        if isinstance(m, LoRALinear):
+                            m.lora_B.data.copy_(saved_B[idx])
+                            idx += 1
+
+            # DPO loss
+            loss, chosen_reward, rejected_reward = dpo_loss(
+                policy_chosen_logps, policy_rejected_logps,
+                ref_chosen_logps, ref_rejected_logps,
+                beta=args.beta,
+            )
+
+            scaled_loss = loss / args.grad_accum
+            scaled_loss.backward()
+            accum_loss += loss.item()
+
+        # Gradient clipping
+        grad_norm = torch.nn.utils.clip_grad_norm_(
+            [p for p in model.parameters() if p.requires_grad], 1.0
+        ).item()
+
+        optimizer.step()
+        scheduler.step()
+
+        avg_loss = accum_loss / args.grad_accum
+        running_loss += avg_loss
+        running_chosen_reward += chosen_reward.item()
+        running_rejected_reward += rejected_reward.item()
+        log_step_count += 1
+
+        # Shutdown check
+        if shutdown_requested:
+            log(f"Graceful shutdown at step {step + 1}", "WARN")
+            save_checkpoint(model, optimizer, scheduler, step + 1, avg_loss, str(args.checkpoint_dir))
+            if args.use_lora:
+                save_lora(model, args.checkpoint_dir / f"lora-{step+1:07d}")
+            break
+
+        # Logging
+        if (step + 1) % args.log_interval == 0:
+            t1 = time.perf_counter()
+            elapsed = t1 - t0
+            avg_l = running_loss / log_step_count
+            avg_cr = running_chosen_reward / log_step_count
+            avg_rr = running_rejected_reward / log_step_count
+            margin = avg_cr - avg_rr
+            lr = scheduler.get_last_lr()[0]
+            mem_gb = torch.cuda.memory_allocated() / 1e9
+
+            log(f"step {step+1:>6d} | loss {avg_l:.4f} | "
+                f"margin {margin:.4f} (c={avg_cr:.3f} r={avg_rr:.3f}) | "
+                f"lr {lr:.2e} | gnorm {grad_norm:.3f} | mem {mem_gb:.1f}GB")
+
+            running_loss = 0.0
+            running_chosen_reward = 0.0
+            running_rejected_reward = 0.0
+            log_step_count = 0
+            t0 = t1
+
+        # Save checkpoint
+        if (step + 1) % args.save_interval == 0:
+            ckpt_path = save_checkpoint(
+                model, optimizer, scheduler, step + 1, avg_loss, str(args.checkpoint_dir)
+            )
+            if args.use_lora:
+                save_lora(model, args.checkpoint_dir / f"lora-{step+1:07d}")
+            log(f"Checkpoint saved -> {ckpt_path}")
+
+    # Final save
+    final_path = save_checkpoint(
+        model, optimizer, scheduler, args.max_steps, avg_loss, str(args.checkpoint_dir)
+    )
+    if args.use_lora:
+        save_lora(model, args.checkpoint_dir / "lora-final")
+        # Also merge and save merged model
+        log("Merging LoRA weights into base model...")
+        merge_lora(model)
+        model.save_pretrained(args.checkpoint_dir / "checkpoint-merged")
+        log(f"Merged model saved -> {args.checkpoint_dir / 'checkpoint-merged'}")
+
+    log(f"DPO training complete. Final checkpoint -> {final_path}")
+
+    if log_fh:
+        log_fh.close()
+
+
+if __name__ == "__main__":
+    main()

scripts/evafrill_eval.py

@@ -0,0 +1,749 @@
+"""
+EVAFRILL-Mo 3B — 종합 평가 파이프라인
+======================================
+
+Phase 1: PPL (1-GPU 순차, 16개 val 셋)
+Phase 2: 생성 품질 + 반복률 분석 (cuda:0)
+Phase 3: Calibration (cuda:0)
+Phase 4: lm-eval 벤치마크 — 커스텀 래퍼 사용
+          (belebele_kor_Hang, global_mmlu_full_ko, hellaswag, arc_easy, arc_challenge, kmmlu)
+
+Usage:
+    cd /home/ghong/project-ghong/taketimes/llm-star
+    python eval/evafrill_eval.py
+    python eval/evafrill_eval.py --skip-phase4
+    python eval/evafrill_eval.py --checkpoint checkpoints/3b_final/checkpoint-0319772
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import sys
+import time
+from collections import Counter
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+
+_PROJECT_ROOT = Path(__file__).resolve().parent.parent
+if str(_PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(_PROJECT_ROOT))
+
+from model.transformer import LLM  # noqa: E402
+from tokenizers import Tokenizer  # noqa: E402
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+DEFAULT_CHECKPOINT = str(_PROJECT_ROOT / "checkpoints" / "3b_final" / "checkpoint-0319772")
+TOKENIZER_PATH = str(_PROJECT_ROOT / "tokenizer" / "korean_sp" / "tokenizer.json")
+DATA_DIR = _PROJECT_ROOT / "data"
+OUTPUT_DIR = _PROJECT_ROOT / "eval" / "outputs"
+
+# GPUs available
+N_GPUS = 1
+GPU_IDS = [0]
+
+# 한국어 생성 프롬프트 (15개)
+PROMPTS = [
+    "대한민국의 수도는",
+    "인공지능이란",
+    "한국의 전통 음식 중에서",
+    "지구 온난화의 주요 원인은",
+    "프로그래밍을 배우려면",
+    "조선시대에는",
+    "물리학에서 에너지란",
+    "한국어는 세계에서",
+    "경제 성장을 위해서는",
+    "우주 탐사의 역사를 보면",
+    "머신러닝과 딥러닝의 차이는",
+    "한국 문학의 대표적인 작품으로는",
+    "양자 컴퓨터란",
+    "건강한 식습관을 위해서는",
+    "세계 2차 대전 이후",
+]
+
+# PPL 태스크: GPU → val 파일 리스트
+PPL_TASKS: Dict[int, List[str]] = {
+    0: [
+        "3b_val.bin",
+        "korean_c4_val.bin", "korean_val.bin",
+        "hplt_ko_val.bin", "cc100_ko_val.bin",
+        "korean_wiki_val.bin", "korean_namuwiki_val.bin",
+        "cosmo_auto_math_text_val.bin", "cosmo_stories_val.bin", "cosmo_web_v2_val.bin",
+        "cosmo_stanford_val.bin", "cosmo_khanacademy_val.bin", "cosmo_openstax_val.bin", "cosmo_wikihow_val.bin",
+        "mathpile_val.bin", "open_web_math_val.bin",
+    ],
+}
+
+
+# ===========================================================================
+# Argument parsing
+# ===========================================================================
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="EVAFRILL-Mo 종합 평가")
+    parser.add_argument("--checkpoint", default=DEFAULT_CHECKPOINT)
+    parser.add_argument("--output-dir", default=None)
+    parser.add_argument("--seq-len", type=int, default=2048)
+    parser.add_argument("--stride", type=int, default=512)
+    parser.add_argument("--batch-size", type=int, default=2)
+    parser.add_argument("--max-new-tokens", type=int, default=256)
+    parser.add_argument("--skip-phase1", action="store_true")
+    parser.add_argument("--skip-phase2", action="store_true")
+    parser.add_argument("--skip-phase3", action="store_true")
+    parser.add_argument("--skip-phase4", action="store_true")
+    parser.add_argument("--limit", type=int, default=None,
+                        help="Limit examples per lm-eval task (for fast testing)")
+    parser.add_argument("--exclude-tasks", type=str, default=None,
+                        help="Comma-separated lm-eval tasks to exclude (e.g. kmmlu)")
+    return parser.parse_args()
+
+
+# ===========================================================================
+# Sliding-window PPL dataset
+# ===========================================================================
+
+class BinDataset(Dataset):
+    def __init__(self, path: str, seq_len: int, stride: int):
+        data = np.fromfile(path, dtype=np.uint16)
+        self.data = torch.from_numpy(data.astype(np.int64))
+        self.seq_len = seq_len
+        self.stride = stride
+        self.indices = list(range(0, max(1, len(self.data) - seq_len), stride))
+
+    def __len__(self):
+        return len(self.indices)
+
+    def __getitem__(self, idx):
+        start = self.indices[idx]
+        chunk = self.data[start: start + self.seq_len + 1]
+        if len(chunk) < self.seq_len + 1:
+            chunk = F.pad(chunk, (0, self.seq_len + 1 - len(chunk)))
+        return chunk[:-1], chunk[1:]
+
+
+# ===========================================================================
+# PPL worker (runs in separate process)
+# ===========================================================================
+
+def _ppl_worker(
+    checkpoint: str,
+    gpu_id: int,
+    val_files: List[str],
+    data_dir: str,
+    seq_len: int,
+    stride: int,
+    batch_size: int,
+) -> Dict[str, float]:
+    """각 GPU에서 여러 val 파일의 PPL을 계산."""
+    import torch
+    import sys
+    from pathlib import Path
+    sys.path.insert(0, str(Path(checkpoint).parent.parent.parent))  # project root
+
+    from model.transformer import LLM  # noqa
+
+    device = f"cuda:{gpu_id}"
+    model = LLM.from_pretrained(checkpoint)
+    model = model.to(device=device, dtype=torch.bfloat16)
+    model.eval()
+
+    results = {}
+    for fname in val_files:
+        fpath = Path(data_dir) / fname
+        if not fpath.exists():
+            results[fname.replace("_val.bin", "")] = None
+            continue
+
+        ds = BinDataset(str(fpath), seq_len, stride)
+        loader = DataLoader(ds, batch_size=batch_size, num_workers=0, pin_memory=True)
+
+        total_nll = 0.0
+        total_tokens = 0
+        with torch.no_grad():
+            for x, y in loader:
+                x, y = x.to(device), y.to(device)
+                logits, _ = model(x)
+                loss = F.cross_entropy(
+                    logits.reshape(-1, logits.size(-1)),
+                    y.reshape(-1),
+                    reduction="sum",
+                    ignore_index=0,
+                )
+                valid = (y != 0).sum().item()
+                total_nll += loss.item()
+                total_tokens += valid
+
+        ppl = math.exp(total_nll / max(total_tokens, 1))
+        key = fname.replace("_val.bin", "")
+        results[key] = round(ppl, 4)
+        print(f"[GPU {gpu_id}] {key}: PPL={ppl:.4f}", flush=True)
+
+    return results
+
+
+# ===========================================================================
+# Phase 1: PPL (병렬)
+# ===========================================================================
+
+def run_phase1(checkpoint: str, seq_len: int, stride: int, batch_size: int) -> Dict[str, float]:
+    print("\n" + "=" * 60)
+    print("Phase 1: PPL 평가 (1-GPU 순차)")
+    print("=" * 60)
+    t0 = time.time()
+
+    existing = [f for f in PPL_TASKS[0] if (DATA_DIR / f).exists()]
+    if not existing:
+        print("  평가할 val 파일 없음")
+        return {}
+
+    all_results = _ppl_worker(
+        checkpoint=checkpoint,
+        gpu_id=0,
+        val_files=existing,
+        data_dir=str(DATA_DIR),
+        seq_len=seq_len,
+        stride=stride,
+        batch_size=batch_size,
+    )
+
+    elapsed = time.time() - t0
+    print(f"\n  Phase 1 완료 ({elapsed:.1f}s)")
+    return all_results
+
+
+# ===========================================================================
+# Phase 2: 생성 품질 + 반복률
+# ===========================================================================
+
+def _ngram_repetition(tokens: List[int], n: int) -> float:
+    if len(tokens) < n:
+        return 0.0
+    ngrams = [tuple(tokens[i:i+n]) for i in range(len(tokens) - n + 1)]
+    total = len(ngrams)
+    unique = len(set(ngrams))
+    return round(1.0 - unique / total, 4) if total > 0 else 0.0
+
+
+def run_phase2(checkpoint: str, max_new_tokens: int) -> List[Dict]:
+    print("\n" + "=" * 60)
+    print("Phase 2: 생성 품질 + 반복률")
+    print("=" * 60)
+
+    device = "cuda:0"
+    model = LLM.from_pretrained(checkpoint)
+    model = model.to(device=device, dtype=torch.bfloat16)
+    model.eval()
+
+    tok = Tokenizer.from_file(TOKENIZER_PATH)
+
+    results = []
+    configs = [
+        ("greedy",    0.0, 1.0),
+        ("t0.7",      0.7, 1.0),
+        ("t0.7_r1.2", 0.7, 1.2),
+        ("t0.9_r1.1", 0.9, 1.1),
+    ]
+
+    for prompt in PROMPTS:
+        ids = tok.encode(prompt).ids
+        x = torch.tensor([ids], dtype=torch.long, device=device)
+
+        row = {"prompt": prompt, "configs": {}}
+        for cfg_name, temp, rep_pen in configs:
+            with torch.no_grad():
+                generated = list(ids)
+                for _ in range(max_new_tokens):
+                    inp = torch.tensor([generated[-2048:]], dtype=torch.long, device=device)
+                    logits, _ = model(inp)
+                    logits = logits[:, -1, :]
+
+                    # Repetition penalty
+                    if rep_pen != 1.0:
+                        for tok_id in set(generated[-64:]):
+                            logits[0, tok_id] /= rep_pen
+
+                    if temp == 0.0:
+                        next_tok = logits.argmax(dim=-1).item()
+                    else:
+                        probs = torch.softmax(logits / temp, dim=-1)
+                        next_tok = torch.multinomial(probs[0], 1).item()
+
+                    generated.append(next_tok)
+                    if next_tok in (tok.token_to_id("</s>"), tok.token_to_id("<eos>"), 2):
+                        break
+
+            new_ids = generated[len(ids):]
+            text = tok.decode(new_ids)
+            rep3 = _ngram_repetition(new_ids, 3)
+            rep4 = _ngram_repetition(new_ids, 4)
+            eos_hit = new_ids[-1] in (2,) if new_ids else False
+
+            row["configs"][cfg_name] = {
+                "text": text,
+                "tokens": len(new_ids),
+                "3gram_rep": rep3,
+                "4gram_rep": rep4,
+                "eos": eos_hit,
+            }
+
+        results.append(row)
+        greedy = row["configs"]["greedy"]
+        print(f"\n[{prompt}]")
+        print(f"  greedy({greedy['tokens']}tok, rep3={greedy['3gram_rep']:.2%}): {greedy['text'][:120]}")
+
+    del model
+    torch.cuda.empty_cache()
+    return results
+
+
+# ===========================================================================
+# Phase 3: Calibration
+# ===========================================================================
+
+def run_phase3(checkpoint: str) -> Dict:
+    print("\n" + "=" * 60)
+    print("Phase 3: Calibration 체크")
+    print("=" * 60)
+
+    device = "cuda:0"
+    model = LLM.from_pretrained(checkpoint)
+    model = model.to(device=device, dtype=torch.bfloat16)
+    model.eval()
+
+    val_path = DATA_DIR / "3b_val.bin"
+    if not val_path.exists():
+        print("  3b_val.bin 없음 — 스킵")
+        return {}
+
+    ds = BinDataset(str(val_path), seq_len=512, stride=256)
+    loader = DataLoader(ds, batch_size=8, num_workers=0)
+
+    top1 = top5 = top10 = total = 0
+    mean_probs, mean_entropies = [], []
+
+    CALIB_TOKENS = 50_000
+    token_count = 0
+
+    with torch.no_grad():
+        for x, y in loader:
+            x, y = x.to(device), y.to(device)
+            logits, _ = model(x)
+            probs = torch.softmax(logits, dim=-1)
+
+            mask = (y != 0)
+            labels = y[mask]
+            p = probs[mask]
+
+            ranks = (p > p.gather(1, labels.unsqueeze(1))).sum(dim=1)
+            top1 += (ranks < 1).sum().item()
+            top5 += (ranks < 5).sum().item()
+            top10 += (ranks < 10).sum().item()
+
+            chosen_p = p.gather(1, labels.unsqueeze(1)).squeeze(1)
+            mean_probs.append(chosen_p.mean().item())
+
+            ent = -(p * (p + 1e-10).log()).sum(dim=-1)  # p already masked → 1D
+            mean_entropies.append(ent.mean().item())
+
+            total += labels.size(0)
+            token_count += labels.size(0)
+            if token_count >= CALIB_TOKENS:
+                break
+
+    result = {
+        "top1_acc":  round(top1 / total, 4),
+        "top5_acc":  round(top5 / total, 4),
+        "top10_acc": round(top10 / total, 4),
+        "mean_prob": round(float(np.mean(mean_probs)), 4),
+        "mean_entropy": round(float(np.mean(mean_entropies)), 4),
+        "total_tokens": total,
+    }
+    print(f"  Top-1: {result['top1_acc']:.2%}  Top-5: {result['top5_acc']:.2%}  Top-10: {result['top10_acc']:.2%}")
+    print(f"  Mean prob: {result['mean_prob']:.4f}  Entropy: {result['mean_entropy']:.4f}")
+
+    del model
+    torch.cuda.empty_cache()
+    return result
+
+
+# ===========================================================================
+# Phase 4: lm-eval 벤치마크 (커스텀 래퍼)
+# ===========================================================================
+
+def run_phase4(checkpoint: str, limit: int = None, exclude_tasks: str = None) -> Dict:
+    print("\n" + "=" * 60)
+    print("Phase 4: lm-eval 벤치마크")
+    print("=" * 60)
+
+    try:
+        import lm_eval
+        from lm_eval.api.model import LM as BaseLM
+        from lm_eval.api.instance import Instance
+        from lm_eval import evaluator
+    except ImportError:
+        print("  lm-eval 미설치 — 스킵 (pip install lm-eval)")
+        return {}
+
+    device = "cuda:0"
+
+    class EvafrillLM(BaseLM):
+        """EVAFRILL-Mo를 lm-eval-harness에 연결하는 래퍼."""
+
+        def __init__(self, checkpoint: str, device: str, batch_size: int = 8):
+            super().__init__()
+            self._model = LLM.from_pretrained(checkpoint)
+            self._model = self._model.to(device=device, dtype=torch.bfloat16)
+            self._model.eval()
+            self._tok = Tokenizer.from_file(TOKENIZER_PATH)
+            self._device = device
+            self._batch_size = batch_size
+            self._max_len = 4096
+
+        @property
+        def eot_token_id(self) -> int:
+            return 2  # </s>
+
+        @property
+        def max_length(self) -> int:
+            return self._max_len
+
+        @property
+        def max_gen_toks(self) -> int:
+            return 256
+
+        @property
+        def batch_size(self) -> int:
+            return self._batch_size
+
+        @property
+        def device(self):
+            return self._device
+
+        def tok_encode(self, string: str) -> List[int]:
+            return self._tok.encode(string).ids
+
+        def tok_decode(self, tokens) -> str:
+            return self._tok.decode(list(tokens))
+
+        def _model_call(self, inps: torch.Tensor) -> torch.Tensor:
+            with torch.no_grad():
+                logits, _ = self._model(inps.to(self._device))
+                return logits
+
+        def loglikelihood(self, requests) -> List[Tuple[float, bool]]:
+            results = []
+            for req in requests:
+                ctx, cont = req.args[0], req.args[1]
+                ctx_ids = self.tok_encode(ctx)
+                cont_ids = self.tok_encode(cont)
+
+                all_ids = ctx_ids + cont_ids
+                if len(all_ids) > self._max_len:
+                    all_ids = all_ids[-self._max_len:]
+                    # adjust cont boundary
+                    cont_start = len(all_ids) - len(cont_ids)
+                else:
+                    cont_start = len(ctx_ids)
+
+                inp = torch.tensor([all_ids[:-1]], dtype=torch.long)
+                tgt = torch.tensor([all_ids[1:]], dtype=torch.long)
+
+                logits = self._model_call(inp)
+                log_probs = F.log_softmax(logits, dim=-1)
+
+                # sum log-probs over continuation tokens
+                cont_log_prob = 0.0
+                is_greedy = True
+                for i, t in enumerate(cont_ids):
+                    pos = cont_start - 1 + i
+                    if pos >= log_probs.size(1):
+                        break
+                    cont_log_prob += log_probs[0, pos, t].item()
+                    pred = log_probs[0, pos].argmax().item()
+                    if pred != t:
+                        is_greedy = False
+
+                results.append((cont_log_prob, is_greedy))
+            return results
+
+        def loglikelihood_rolling(self, requests) -> List[float]:
+            results = []
+            for req in requests:
+                text = req.args[0]
+                ids = self.tok_encode(text)
+                total_nll = 0.0
+                for start in range(0, len(ids) - 1, self._max_len - 1):
+                    chunk = ids[start: start + self._max_len]
+                    if len(chunk) < 2:
+                        break
+                    inp = torch.tensor([chunk[:-1]], dtype=torch.long)
+                    tgt = torch.tensor([chunk[1:]], dtype=torch.long)
+                    logits = self._model_call(inp)
+                    nll = F.cross_entropy(
+                        logits[0], tgt[0].to(self._device), reduction="sum"
+                    ).item()
+                    total_nll += nll
+                results.append(-total_nll)
+            return results
+
+        def generate_until(self, requests) -> List[str]:
+            results = []
+            for req in requests:
+                ctx = req.args[0]
+                gen_kwargs = req.args[1] if len(req.args) > 1 else {}
+                until = gen_kwargs.get("until", [])
+                max_gen = gen_kwargs.get("max_gen_toks", self.max_gen_toks)
+                temp = gen_kwargs.get("temperature", 0.0)
+
+                ids = self.tok_encode(ctx)
+                generated = list(ids)
+
+                with torch.no_grad():
+                    for _ in range(max_gen):
+                        inp = torch.tensor(
+                            [generated[-self._max_len:]], dtype=torch.long
+                        )
+                        logits = self._model_call(inp)[:, -1:, :].squeeze(1)
+                        if temp == 0.0:
+                            next_tok = logits.argmax(dim=-1).item()
+                        else:
+                            probs = torch.softmax(logits / temp, dim=-1)
+                            next_tok = torch.multinomial(probs[0], 1).item()
+                        generated.append(next_tok)
+                        if next_tok == self.eot_token_id:
+                            break
+                        decoded_new = self.tok_decode(generated[len(ids):])
+                        if any(stop in decoded_new for stop in until):
+                            break
+
+                new_text = self.tok_decode(generated[len(ids):])
+                for stop in until:
+                    if stop in new_text:
+                        new_text = new_text[:new_text.index(stop)]
+                results.append(new_text)
+            return results
+
+    lm = EvafrillLM(checkpoint, device=device, batch_size=2)
+
+    tasks = [
+        "belebele_kor_Hang",
+        "global_mmlu_full_ko",
+        "hellaswag",
+        "arc_easy",
+        "arc_challenge",
+        "kmmlu",
+    ]
+
+    if exclude_tasks:
+        excluded = {t.strip() for t in exclude_tasks.split(",")}
+        tasks = [t for t in tasks if t not in excluded]
+        print(f"  제외: {', '.join(excluded)}")
+
+    print(f"  태스크: {', '.join(tasks)}")
+    print("  (belebele/mmlu: 한국어, hellaswag/arc: 영어)")
+    if limit:
+        print(f"  limit: {limit} examples/task")
+
+    try:
+        results = evaluator.simple_evaluate(
+            model=lm,
+            tasks=tasks,
+            num_fewshot=0,
+            batch_size=2,
+            log_samples=False,
+            limit=limit,
+        )
+        return results.get("results", {})
+    except Exception as e:
+        print(f"  lm-eval 오류: {e}")
+        import traceback; traceback.print_exc()
+        return {}
+
+
+# ===========================================================================
+# Report generation
+# ===========================================================================
+
+def generate_report(
+    checkpoint: str,
+    output_dir: Path,
+    ppl: Dict,
+    gen: List[Dict],
+    calib: Dict,
+    bench: Dict,
+    elapsed: float,
+) -> Path:
+    now = datetime.now().strftime("%Y-%m-%d %H:%M")
+    run_tag = datetime.now().strftime("%Y%m%d_%H%M")
+    report_path = _PROJECT_ROOT / "reports" / f"{run_tag}_EVAFRILL_EVAL_REPORT.md"
+    report_path.parent.mkdir(parents=True, exist_ok=True)
+
+    lines = [
+        "# EVAFRILL-Mo 3B — 종합 평가 보고서",
+        "",
+        f"- **평가 일시**: {now}",
+        f"- **체크포인트**: `{Path(checkpoint).name}`",
+        f"- **총 소요 시간**: {elapsed/60:.1f}분",
+        "",
+        "---",
+        "",
+        "## 1. Executive Summary",
+        "",
+    ]
+
+    # PPL summary
+    if ppl:
+        avg_ko = np.mean([v for k, v in ppl.items() if v and "korean" in k or "hplt" in k or "cc100" in k])
+        lines += [
+            "### PPL (주요 셋)",
+            "",
+            "| 데이터셋 | PPL |",
+            "|---------|-----|",
+        ]
+        for k, v in sorted(ppl.items()):
+            if v is not None:
+                lines.append(f"| {k} | {v:.4f} |")
+        lines.append("")
+
+    # Generation summary
+    if gen:
+        greedy_reps = [r["configs"]["greedy"]["3gram_rep"] for r in gen if "greedy" in r["configs"]]
+        greedy_eos = [r["configs"]["greedy"]["eos"] for r in gen if "greedy" in r["configs"]]
+        t07r12_reps = [r["configs"].get("t0.7_r1.2", {}).get("3gram_rep", None) for r in gen]
+        t07r12_reps = [x for x in t07r12_reps if x is not None]
+
+        lines += [
+            "### 생성 품질 요약",
+            "",
+            f"| 설정 | 평균 3-gram 반복률 | EOS 종료율 |",
+            f"|------|-------------------|-----------|",
+            f"| greedy | {np.mean(greedy_reps):.2%} | {np.mean(greedy_eos):.0%} |",
+        ]
+        if t07r12_reps:
+            t07r12_eos = [r["configs"].get("t0.7_r1.2", {}).get("eos", False) for r in gen]
+            lines.append(f"| temp=0.7 rep=1.2 | {np.mean(t07r12_reps):.2%} | {np.mean(t07r12_eos):.0%} |")
+        lines.append("")
+
+    # Calibration
+    if calib:
+        lines += [
+            "### Calibration",
+            "",
+            f"| Top-1 | Top-5 | Top-10 |",
+            f"|-------|-------|--------|",
+            f"| {calib['top1_acc']:.2%} | {calib['top5_acc']:.2%} | {calib['top10_acc']:.2%} |",
+            "",
+        ]
+
+    # Benchmarks
+    if bench:
+        lines += [
+            "### lm-eval 벤치마크",
+            "",
+            "| 태스크 | Accuracy | 랜덤 기준 |",
+            "|--------|----------|----------|",
+        ]
+        random_baseline = {
+            "belebele_kor_Hang": 0.25,
+            "global_mmlu_full_ko": 0.25,
+            "hellaswag": 0.25,
+            "arc_easy": 0.25,
+            "arc_challenge": 0.25,
+            "kmmlu": 0.25,
+        }
+        for task, res in bench.items():
+            acc = res.get("acc,none", res.get("acc", "N/A"))
+            rb = random_baseline.get(task, "?")
+            lines.append(f"| {task} | {acc:.4f} | {rb} |")
+        lines.append("")
+
+    # Generation samples
+    if gen:
+        lines += ["## 2. 생성 샘플 (Greedy)", ""]
+        for r in gen:
+            gcfg = r["configs"].get("greedy", {})
+            lines += [
+                f"**[{r['prompt']}]**",
+                f"> {gcfg.get('text', '')[:200]}",
+                f"> *EOS={gcfg.get('eos')}, 3gram_rep={gcfg.get('3gram_rep', 0):.2%}, tokens={gcfg.get('tokens')}*",
+                "",
+            ]
+
+    report_path.write_text("\n".join(lines), encoding="utf-8")
+    print(f"\n  보고서 저장: {report_path}")
+
+    # JSON 결과도 저장
+    json_path = output_dir / "evafrill_eval_results.json"
+    json_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(json_path, "w", encoding="utf-8") as f:
+        json.dump({"ppl": ppl, "calib": calib, "bench": bench}, f, ensure_ascii=False, indent=2)
+
+    return report_path
+
+
+# ===========================================================================
+# Main
+# ===========================================================================
+
+def main():
+    args = parse_args()
+    t_start = time.time()
+
+    run_tag = datetime.now().strftime("%Y%m%d_%H%M")
+    output_dir = Path(args.output_dir) if args.output_dir else (
+        _PROJECT_ROOT / "eval" / "outputs" / f"evafrill_eval_{run_tag}"
+    )
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    print("=" * 60)
+    print("EVAFRILL-Mo 3B 종합 평가 시작")
+    print(f"체크포인트: {args.checkpoint}")
+    print(f"출력 디렉토리: {output_dir}")
+    print("=" * 60)
+
+    ppl_results = {}
+    gen_results = []
+    calib_results = {}
+    bench_results = {}
+
+    if not args.skip_phase1:
+        ppl_results = run_phase1(
+            args.checkpoint, args.seq_len, args.stride, args.batch_size
+        )
+
+    if not args.skip_phase2:
+        gen_results = run_phase2(args.checkpoint, args.max_new_tokens)
+
+    if not args.skip_phase3:
+        calib_results = run_phase3(args.checkpoint)
+
+    if not args.skip_phase4:
+        bench_results = run_phase4(args.checkpoint, limit=args.limit,
+                                    exclude_tasks=args.exclude_tasks)
+
+    elapsed = time.time() - t_start
+    report_path = generate_report(
+        args.checkpoint, output_dir,
+        ppl_results, gen_results, calib_results, bench_results,
+        elapsed,
+    )
+
+    print("\n" + "=" * 60)
+    print(f"평가 완료! 총 {elapsed/60:.1f}분")
+    print(f"보고서: {report_path}")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()

scripts/generate_repetition_preference.py

@@ -0,0 +1,480 @@
+#!/usr/bin/env python3
+"""
+data/generate_repetition_preference.py — Self-play preference data targeting repetition.
+
+Generates (prompt, chosen, rejected) pairs by:
+  - rejected: greedy decoding (temp=0, rep_penalty=1.0) → tends to repeat
+  - chosen:   sampling with repetition penalty (temp=0.7, rep_penalty=1.2) → cleaner
+
+Only keeps pairs where rejected has strictly higher 3-gram repetition rate than chosen.
+
+Usage:
+    python3 data/generate_repetition_preference.py \
+        --checkpoint checkpoints/3b_dpo/checkpoint-slerp
+
+    python3 data/generate_repetition_preference.py \
+        --checkpoint checkpoints/3b_dpo/checkpoint-slerp \
+        --output data/preference/repetition_preference.jsonl \
+        --num_prompts 100 \
+        --max_tokens 256
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import sys
+import time
+from pathlib import Path
+from typing import List, Optional
+
+import torch
+import torch.nn.functional as F
+
+_PROJECT_ROOT = Path(__file__).resolve().parent.parent
+if str(_PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(_PROJECT_ROOT))
+
+from model import LLM  # noqa: E402
+from tokenizers import Tokenizer  # noqa: E402
+
+# ---------------------------------------------------------------------------
+# Korean prompt bank — 100+ diverse prompts
+# ---------------------------------------------------------------------------
+
+# 15 existing eval prompts (completion style → wrapped in chat template)
+_EVAL_PROMPTS = [
+    "대한민국의 수도는 어디인지 설명해주세요.",
+    "인공지능이란 무엇인지 자세히 설명해주세요.",
+    "한국의 전통 음식 중에서 대표적인 것들을 소개해주세요.",
+    "지구 온난화의 주요 원인은 무엇인가요?",
+    "프로그래밍을 배우려면 어떻게 시작해야 하나요?",
+    "조선시대에는 어떤 일들이 있었나요?",
+    "물리학에서 에너지란 무엇인지 설명해주세요.",
+    "한국어는 세계에서 어떤 특징을 가지고 있나요?",
+    "경제 성장을 위해서는 무엇이 필요한가요?",
+    "우주 탐사의 역사를 간단히 설명해주세요.",
+    "머신러닝과 딥러닝의 차이는 무엇인가요?",
+    "한국 문학의 대표적인 작품으로는 어떤 것들이 있나요?",
+    "양자 컴퓨터란 무엇인지 설명해주세요.",
+    "건강한 식습관을 위해서는 어떻게 해야 하나요?",
+    "세계 2차 대전 이후 세계는 어떻게 변했나요?",
+]
+
+# Additional diverse prompts (~85 more)
+_EXTRA_PROMPTS = [
+    # 일상 대화
+    "오늘 날씨가 좋은데 뭐 하면 좋을까요?",
+    "주말에 뭐 하면 좋을지 추천해주세요.",
+    "좋은 하루를 시작하는 방법을 알려주세요.",
+    "집에서 할 수 있는 취미 활동을 추천해주세요.",
+    "친구와 싸웠을 때 어떻게 화해하면 좋을까요?",
+    "외로움을 느낄 때 어떻게 극복할 수 있나요?",
+    "시간 관리를 잘 하는 방법을 알려주세요.",
+    "아침 일찍 일어나는 습관을 만들려면 어떻게 해야 하나요?",
+    "새로운 도시로 이사했을 때 적응하는 방법은?",
+    "카페에서 혼자 시간 보내는 것의 장점은 무엇인가요?",
+
+    # 지식 — 과학
+    "DNA가 무엇인지 설명해주세요.",
+    "블랙홀이란 무엇인가요?",
+    "진화론이란 무엇인지 간단히 설명해주세요.",
+    "기후 변화가 생태계에 미치는 영향은 무엇인가요?",
+    "인체의 면역 시스템은 어떻게 작동하나요?",
+    "빛의 속도는 왜 중요한가요?",
+    "원자와 분자의 차이점은 무엇인가요?",
+    "광합성이란 무엇인지 설명해주세요.",
+    "중력파란 무엇인가요?",
+    "줄기세포 치료란 무엇이며 어떻게 활용되나요?",
+
+    # 지식 — 역사·사회
+    "한국의 역사에서 가장 중요한 사건은 무엇인가요?",
+    "민주주의란 무엇인지 설명해주세요.",
+    "산업혁명이 세계에 미친 영향은 무엇인가요?",
+    "냉전이란 무엇이었나요?",
+    "한국 전쟁의 원인과 결과를 설명해주세요.",
+    "세계화란 무엇이며 어떤 영향을 미치나요?",
+    "인권이란 무엇이며 왜 중요한가요?",
+    "실크로드가 역사적으로 중요한 이유는 무엇인가요?",
+    "르네상스 시대는 어떤 시기였나요?",
+    "한국의 독립운동에 대해 설명해주세요.",
+
+    # 조언 — 직업·학습
+    "취업 면접 잘 보는 방법은 무엇인가요?",
+    "이력서를 잘 쓰는 방법을 알려주세요.",
+    "대학 생활을 알차게 보내는 방법은?",
+    "공부 집중력을 높이는 방법을 알려주세요.",
+    "외국어를 빠르게 배우는 방법은 무엇인가요?",
+    "직장에서 상사와 잘 지내는 방법은?",
+    "프리랜서로 일하면 어떤 장단점이 있나요?",
+    "자기소개서를 잘 ��는 팁을 알려주세요.",
+    "독서 습관을 기르는 방법은 무엇인가요?",
+    "수학을 잘하기 위한 공부 방법은?",
+
+    # 조언 — 건강·심리
+    "스트레스 해소 방법을 알려주세요.",
+    "우울감을 극복하는 방법은 무엇인가요?",
+    "규칙적인 운동 습관을 만드는 방법은?",
+    "수면의 질을 높이는 방법을 알려주세요.",
+    "번아웃을 예방하는 방법은 무엇인가요?",
+    "마음의 평화를 찾는 방법은?",
+    "자존감을 높이는 방법을 알려주세요.",
+    "명상을 시작하려면 어떻게 해야 하나요?",
+    "건강한 체중을 유지하는 방법은?",
+    "디지털 중독을 극복하는 방법을 알려주세요.",
+
+    # 창작
+    "짧은 동화를 하나 만들어주세요.",
+    "봄에 대한 시를 써주세요.",
+    "미래 도시를 배경으로 한 짧은 이야기를 써주세요.",
+    "바다에 관한 짧은 수필을 써주세요.",
+    "고양이를 주인공으로 한 짧은 이야기를 만들어주세요.",
+    "가을 풍경을 묘사하는 글을 써주세요.",
+    "우정에 관한 짧은 시를 써주세요.",
+    "엄마에게 보내는 편지를 써주세요.",
+    "미래의 나에게 쓰는 편지를 작성해주세요.",
+    "어린 시절 추억에 관한 짧은 글을 써주세요.",
+
+    # 기술·IT
+    "클라우드 컴퓨팅이란 무엇인가요?",
+    "블록체인이 무엇인지 설명해주세요.",
+    "사이버 보안이 왜 중요한가요?",
+    "빅데이터란 무엇이며 어떻게 활용되나요?",
+    "5G 기술이 가져올 변화는 무엇인가요?",
+    "인터넷 검색 엔진은 어떻게 작동하나요?",
+    "스마트폰이 생활에 미친 영향은 무엇인가요?",
+    "가상현실과 증강현실의 차이는 무엇인가요?",
+    "자율주행 자동차 기술은 어디까지 왔나요?",
+    "오픈소스 소프트웨어란 무엇인가요?",
+
+    # 문화·예술
+    "K-팝이 세계적으로 인기를 얻은 이유는 무엇인가요?",
+    "한국 영화가 세계 시장에서 주목받는 이유는?",
+    "전통 예술과 현대 예술의 차이는 무엇인가요?",
+    "음악이 감정에 미치는 영향은 무엇인가요?",
+    "독서가 삶에 미치는 긍정적인 영향은?",
+    "미술 감상을 잘 하는 방법을 알려주세요.",
+    "한국 전통 음악인 국악의 특징은 무엇인가요?",
+    "영화 비평을 잘 쓰는 방법은?",
+    "여행이 사람을 성장시키는 이유는 무엇인가요?",
+    "사진 찍기를 잘 하는 팁을 알려주세요.",
+
+    # 환경·사회
+    "환경 보호를 위해 개인이 할 수 있는 일은?",
+    "재활용의 중요성과 방법을 설명해주세요.",
+    "채식주의의 장단점은 무엇인가요?",
+    "동물 복지란 무엇이며 왜 중요한가요?",
+    "지속 가능한 발전이란 무엇인가요?",
+    "노령화 사회가 가져오는 문제점은 무엇인가요?",
+    "교육 불평등을 해소하는 방법은?",
+    "빈곤 문제를 해결하기 위한 방법은?",
+    "다문화 사회에서 공존하는 방법은?",
+    "봉사 활동이 사회에 미치는 영향은 무엇인가요?",
+]
+
+ALL_PROMPTS = _EVAL_PROMPTS + _EXTRA_PROMPTS  # 15 + 85 = 100
+
+CHAT_TEMPLATE = "<|user|>\n{prompt}\n<|assistant|>\n"
+
+EOS_TOKEN_ID = 2
+
+
+# ---------------------------------------------------------------------------
+# Repetition metric
+# ---------------------------------------------------------------------------
+
+def compute_ngram_repetition_rate(tokens: List[int], n: int = 3) -> float:
+    """Fraction of n-gram positions that are repeats of an earlier occurrence."""
+    if len(tokens) < n:
+        return 0.0
+    ngrams = [tuple(tokens[i: i + n]) for i in range(len(tokens) - n + 1)]
+    if not ngrams:
+        return 0.0
+    seen: set = set()
+    repeated = 0
+    for ng in ngrams:
+        if ng in seen:
+            repeated += 1
+        seen.add(ng)
+    return repeated / len(ngrams)
+
+
+# ---------------------------------------------------------------------------
+# Generation
+# ---------------------------------------------------------------------------
+
+@torch.inference_mode()
+def generate(
+    model: torch.nn.Module,
+    input_ids: torch.Tensor,
+    max_new_tokens: int,
+    temperature: float,
+    repetition_penalty: float,
+    eos_token_id: int,
+) -> List[int]:
+    """Auto-regressive generation with optional repetition penalty.
+
+    Args:
+        model: LLM instance already on device
+        input_ids: (1, T) prompt token ids
+        max_new_tokens: max tokens to generate
+        temperature: sampling temperature (0 = greedy)
+        repetition_penalty: penalty > 1 reduces prob of previously seen tokens
+        eos_token_id: stop generation when this token is produced
+
+    Returns:
+        List of generated token ids (not including the prompt).
+    """
+    device = input_ids.device
+    generated: List[int] = []
+    current_ids = input_ids.clone()  # (1, T)
+
+    for _ in range(max_new_tokens):
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+            logits, _ = model(current_ids)  # (1, T, V)
+
+        next_logits = logits[0, -1, :].float()  # (V,)
+
+        # Repetition penalty: discount logits for already-generated tokens
+        if repetition_penalty != 1.0:
+            all_seen_ids = current_ids[0].tolist() + generated
+            for token_id in set(all_seen_ids):
+                if token_id < next_logits.shape[0]:
+                    if next_logits[token_id] < 0:
+                        next_logits[token_id] *= repetition_penalty
+                    else:
+                        next_logits[token_id] /= repetition_penalty
+
+        # Sample / greedy
+        if temperature == 0.0:
+            next_token = int(next_logits.argmax())
+        else:
+            next_logits = next_logits / temperature
+            probs = F.softmax(next_logits, dim=-1)
+            next_token = int(torch.multinomial(probs, num_samples=1).item())
+
+        generated.append(next_token)
+
+        if next_token == eos_token_id:
+            break
+
+        # Append to context
+        next_tensor = torch.tensor([[next_token]], dtype=torch.long, device=device)
+        current_ids = torch.cat([current_ids, next_tensor], dim=1)
+
+    return generated
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Generate self-play repetition preference data"
+    )
+    parser.add_argument(
+        "--checkpoint",
+        type=Path,
+        default=Path("checkpoints/3b_dpo/checkpoint-slerp"),
+        help="Path to model checkpoint directory",
+    )
+    parser.add_argument(
+        "--output",
+        type=Path,
+        default=Path("data/preference/repetition_preference.jsonl"),
+        help="Output JSONL path",
+    )
+    parser.add_argument(
+        "--num_prompts",
+        type=int,
+        default=None,
+        help="How many prompts to use (default: all ~100)",
+    )
+    parser.add_argument(
+        "--max_tokens",
+        type=int,
+        default=256,
+        help="Max new tokens per generation",
+    )
+    parser.add_argument(
+        "--tokenizer",
+        type=Path,
+        default=None,
+        help="Path to tokenizer.json (default: auto-resolve)",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda:0",
+        help="Torch device string",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Random seed for reproducibility",
+    )
+    parser.add_argument(
+        "--min_rep_diff",
+        type=float,
+        default=0.0,
+        help="Minimum difference (rejected_rep - chosen_rep) to keep a pair (default: >0)",
+    )
+    return parser.parse_args()
+
+
+def _resolve_tokenizer(args: argparse.Namespace) -> Path:
+    if args.tokenizer is not None:
+        return Path(args.tokenizer)
+    # Try checkpoint dir first
+    ckpt_tok = args.checkpoint / "tokenizer.json"
+    if ckpt_tok.exists():
+        return ckpt_tok
+    # Fall back to project default
+    default_tok = _PROJECT_ROOT / "tokenizer" / "korean_sp" / "tokenizer.json"
+    if default_tok.exists():
+        return default_tok
+    raise FileNotFoundError(
+        "Cannot find tokenizer.json — specify with --tokenizer"
+    )
+
+
+def main() -> None:
+    args = parse_args()
+
+    # Reproducibility
+    torch.manual_seed(args.seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(args.seed)
+
+    # Prompts
+    prompts = ALL_PROMPTS
+    if args.num_prompts is not None:
+        prompts = prompts[: args.num_prompts]
+    print(f"[INFO] Using {len(prompts)} prompts")
+
+    # Device
+    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
+    print(f"[INFO] Device: {device}")
+
+    # Tokenizer
+    tokenizer_path = _resolve_tokenizer(args)
+    print(f"[INFO] Loading tokenizer from {tokenizer_path}")
+    tokenizer = Tokenizer.from_file(str(tokenizer_path))
+
+    # Model
+    checkpoint_path = _PROJECT_ROOT / args.checkpoint if not args.checkpoint.is_absolute() else args.checkpoint
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
+    print(f"[INFO] Loading model from {checkpoint_path} ...")
+    t0 = time.perf_counter()
+    model = LLM.from_pretrained(checkpoint_path)
+    model = model.to(device=device, dtype=torch.bfloat16)
+    model.eval()
+    print(f"[INFO] Model loaded in {time.perf_counter() - t0:.1f}s")
+
+    # Output dir
+    output_path = _PROJECT_ROOT / args.output if not args.output.is_absolute() else args.output
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Stats
+    valid_pairs = 0
+    skipped = 0
+    total_rejected_rep = 0.0
+    total_chosen_rep = 0.0
+
+    t_start = time.perf_counter()
+
+    with open(output_path, "w", encoding="utf-8") as fout:
+        for idx, prompt_text in enumerate(prompts):
+            prompt_str = CHAT_TEMPLATE.format(prompt=prompt_text)
+
+            # Tokenize prompt
+            encoding = tokenizer.encode(prompt_str)
+            prompt_ids = encoding.ids
+            if not prompt_ids:
+                print(f"  [{idx+1}/{len(prompts)}] SKIP: empty tokenization for prompt")
+                skipped += 1
+                continue
+
+            input_ids = torch.tensor([prompt_ids], dtype=torch.long, device=device)
+
+            # --- Generate REJECTED: greedy, no rep penalty ---
+            rej_tokens = generate(
+                model=model,
+                input_ids=input_ids,
+                max_new_tokens=args.max_tokens,
+                temperature=0.0,
+                repetition_penalty=1.0,
+                eos_token_id=EOS_TOKEN_ID,
+            )
+
+            # --- Generate CHOSEN: sampling + rep penalty ---
+            cho_tokens = generate(
+                model=model,
+                input_ids=input_ids,
+                max_new_tokens=args.max_tokens,
+                temperature=0.7,
+                repetition_penalty=1.2,
+                eos_token_id=EOS_TOKEN_ID,
+            )
+
+            # Decode (strip EOS)
+            rej_clean = [t for t in rej_tokens if t != EOS_TOKEN_ID]
+            cho_clean = [t for t in cho_tokens if t != EOS_TOKEN_ID]
+
+            rej_text = tokenizer.decode(rej_clean)
+            cho_text = tokenizer.decode(cho_clean)
+
+            # Compute 3-gram repetition rates on generated tokens
+            rej_rep = compute_ngram_repetition_rate(rej_clean, n=3)
+            cho_rep = compute_ngram_repetition_rate(cho_clean, n=3)
+
+            # Filter: only keep if rejected is more repetitive than chosen
+            diff = rej_rep - cho_rep
+            if diff <= args.min_rep_diff:
+                status = "SKIP"
+                skipped += 1
+            else:
+                status = "KEEP"
+                valid_pairs += 1
+                total_rejected_rep += rej_rep
+                total_chosen_rep += cho_rep
+                record = {
+                    "prompt": prompt_str,
+                    "chosen": cho_text,
+                    "rejected": rej_text,
+                }
+                fout.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+            elapsed = time.perf_counter() - t_start
+            print(
+                f"  [{idx+1:3d}/{len(prompts)}] {status:4s} "
+                f"rej_rep={rej_rep:.3f} cho_rep={cho_rep:.3f} diff={diff:+.3f} "
+                f"| rej_len={len(rej_clean)} cho_len={len(cho_clean)} "
+                f"| elapsed={elapsed:.1f}s"
+            )
+
+    # Summary
+    elapsed_total = time.perf_counter() - t_start
+    print()
+    print("=" * 60)
+    print(f"Generation complete in {elapsed_total:.1f}s")
+    print(f"  Total prompts processed : {len(prompts)}")
+    print(f"  Valid pairs kept        : {valid_pairs}")
+    print(f"  Skipped (rep filter)    : {skipped}")
+    if valid_pairs > 0:
+        avg_rej = total_rejected_rep / valid_pairs
+        avg_cho = total_chosen_rep / valid_pairs
+        print(f"  Avg rejected 3-gram rep : {avg_rej:.4f}")
+        print(f"  Avg chosen  3-gram rep  : {avg_cho:.4f}")
+        print(f"  Avg improvement         : {avg_rej - avg_cho:+.4f}")
+    print(f"  Output saved to         : {output_path}")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()

scripts/lora.py

@@ -0,0 +1,240 @@
+"""
+model/lora.py — LoRA (Low-Rank Adaptation) for EVAFRILL-Mo hybrid models.
+
+Injects trainable low-rank adapters into:
+  - Attention layers: qkv_proj, out_proj
+  - Mamba-2 layers: in_proj, out_proj
+
+Usage:
+    model = LLM.from_pretrained(checkpoint)
+    apply_lora(model, rank=32, alpha=64)
+    # Only LoRA params are trainable; base model is frozen
+
+    # After training, merge LoRA weights back:
+    merge_lora(model)
+
+    # Or save/load LoRA weights separately:
+    save_lora(model, path)
+    load_lora(model, path)
+"""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .attention import MultiHeadAttention
+from .mamba_block import Mamba2Block
+
+
+class LoRALinear(nn.Module):
+    """LoRA adapter wrapping an existing nn.Linear layer.
+
+    Computes: output = original_linear(x) + (alpha/rank) * x @ A^T @ B^T
+    where A: (rank, in_features), B: (out_features, rank)
+    """
+
+    def __init__(
+        self,
+        original: nn.Linear,
+        rank: int = 32,
+        alpha: float = 64.0,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.original = original
+        self.rank = rank
+        self.alpha = alpha
+        self.scaling = alpha / rank
+
+        in_features = original.in_features
+        out_features = original.out_features
+
+        # A: down-projection (in_features → rank)
+        # Create on same device/dtype as original weights
+        _dev = original.weight.device
+        _dt = original.weight.dtype
+        self.lora_A = nn.Parameter(torch.empty(rank, in_features, device=_dev, dtype=_dt))
+        # B: up-projection (rank → out_features)
+        self.lora_B = nn.Parameter(torch.zeros(out_features, rank, device=_dev, dtype=_dt))
+
+        # Initialize A with kaiming uniform, B with zeros
+        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+        # B is already zeros → initial LoRA output is zero
+
+        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+
+        # Freeze original weights
+        original.weight.requires_grad = False
+        if original.bias is not None:
+            original.bias.requires_grad = False
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Original forward
+        result = self.original(x)
+        # LoRA path: x → dropout → A → B → scale
+        lora_out = self.dropout(x)
+        lora_out = F.linear(lora_out, self.lora_A)  # (..., rank)
+        lora_out = F.linear(lora_out, self.lora_B)  # (..., out_features)
+        return result + lora_out * self.scaling
+
+    def merge_weights(self) -> None:
+        """Merge LoRA weights into the original linear layer permanently."""
+        with torch.no_grad():
+            # W' = W + scaling * B @ A
+            self.original.weight.add_(
+                (self.lora_B @ self.lora_A) * self.scaling
+            )
+
+    @property
+    def weight(self) -> torch.Tensor:
+        """Access original weight for compatibility."""
+        return self.original.weight
+
+    @property
+    def bias(self) -> Optional[torch.Tensor]:
+        return self.original.bias
+
+
+def apply_lora(
+    model: nn.Module,
+    rank: int = 32,
+    alpha: float = 64.0,
+    dropout: float = 0.0,
+    target_modules: Optional[list[str]] = None,
+) -> int:
+    """Apply LoRA adapters to a model, freeze base weights.
+
+    Args:
+        model: The LLM model (raw, not DDP-wrapped).
+        rank: LoRA rank (default 32).
+        alpha: LoRA scaling factor (default 64).
+        dropout: Dropout on LoRA path (default 0).
+        target_modules: List of module attribute names to adapt.
+            Default: ["qkv_proj", "out_proj", "in_proj"]
+            (covers both Attention and Mamba layers).
+
+    Returns:
+        Number of LoRA parameters added.
+    """
+    if target_modules is None:
+        target_modules = ["qkv_proj", "out_proj", "in_proj"]
+
+    # First, freeze ALL parameters
+    for param in model.parameters():
+        param.requires_grad = False
+
+    lora_count = 0
+    total_lora_params = 0
+
+    for name, module in model.named_modules():
+        # Check Attention layers
+        if isinstance(module, MultiHeadAttention):
+            for attr in target_modules:
+                if hasattr(module, attr):
+                    original = getattr(module, attr)
+                    if isinstance(original, nn.Linear):
+                        lora_layer = LoRALinear(original, rank=rank, alpha=alpha, dropout=dropout)
+                        setattr(module, attr, lora_layer)
+                        params = rank * original.in_features + original.out_features * rank
+                        total_lora_params += params
+                        lora_count += 1
+
+        # Check Mamba layers
+        elif isinstance(module, Mamba2Block):
+            for attr in target_modules:
+                if hasattr(module, attr):
+                    original = getattr(module, attr)
+                    if isinstance(original, nn.Linear):
+                        lora_layer = LoRALinear(original, rank=rank, alpha=alpha, dropout=dropout)
+                        setattr(module, attr, lora_layer)
+                        params = rank * original.in_features + original.out_features * rank
+                        total_lora_params += params
+                        lora_count += 1
+
+    print(f"[LoRA] Applied {lora_count} adapters, {total_lora_params:,} trainable params "
+          f"(rank={rank}, alpha={alpha})")
+    return total_lora_params
+
+
+def merge_lora(model: nn.Module) -> int:
+    """Merge all LoRA weights back into base model and remove LoRA layers.
+
+    Returns:
+        Number of LoRA layers merged.
+    """
+    merged = 0
+    for name, module in model.named_modules():
+        for attr_name in list(vars(module).keys()):
+            # Check nn.Module children
+            pass
+
+        if isinstance(module, (MultiHeadAttention, Mamba2Block)):
+            for attr in ["qkv_proj", "out_proj", "in_proj"]:
+                if hasattr(module, attr):
+                    layer = getattr(module, attr)
+                    if isinstance(layer, LoRALinear):
+                        layer.merge_weights()
+                        setattr(module, attr, layer.original)
+                        merged += 1
+
+    # Unfreeze all parameters after merging
+    for param in model.parameters():
+        param.requires_grad = True
+
+    print(f"[LoRA] Merged {merged} adapters back into base model")
+    return merged
+
+
+def get_lora_params(model: nn.Module) -> list[nn.Parameter]:
+    """Get all LoRA trainable parameters."""
+    params = []
+    for module in model.modules():
+        if isinstance(module, LoRALinear):
+            params.append(module.lora_A)
+            params.append(module.lora_B)
+    return params
+
+
+def save_lora(model: nn.Module, path: str | Path) -> Path:
+    """Save only the LoRA adapter weights."""
+    path = Path(path)
+    path.mkdir(parents=True, exist_ok=True)
+
+    lora_state = {}
+    for name, module in model.named_modules():
+        if isinstance(module, LoRALinear):
+            lora_state[f"{name}.lora_A"] = module.lora_A.data.cpu()
+            lora_state[f"{name}.lora_B"] = module.lora_B.data.cpu()
+
+    save_path = path / "lora_weights.pt"
+    torch.save(lora_state, save_path)
+    n_params = sum(v.numel() for v in lora_state.values())
+    size_mb = save_path.stat().st_size / 1e6
+    print(f"[LoRA] Saved {len(lora_state)} tensors ({n_params:,} params, {size_mb:.1f} MB) → {save_path}")
+    return save_path
+
+
+def load_lora(model: nn.Module, path: str | Path) -> int:
+    """Load LoRA adapter weights. LoRA layers must already be applied."""
+    path = Path(path)
+    lora_file = path / "lora_weights.pt" if path.is_dir() else path
+    lora_state = torch.load(lora_file, map_location="cpu", weights_only=True)
+
+    loaded = 0
+    for name, module in model.named_modules():
+        if isinstance(module, LoRALinear):
+            a_key = f"{name}.lora_A"
+            b_key = f"{name}.lora_B"
+            if a_key in lora_state and b_key in lora_state:
+                module.lora_A.data.copy_(lora_state[a_key])
+                module.lora_B.data.copy_(lora_state[b_key])
+                loaded += 1
+
+    print(f"[LoRA] Loaded {loaded} adapter weight pairs from {lora_file}")
+    return loaded

scripts/merge_checkpoints.py

@@ -0,0 +1,194 @@
+#!/usr/bin/env python3
+"""
+scripts/merge_checkpoints.py — Slerp (Spherical Linear Interpolation) checkpoint merge.
+
+Merges two model checkpoints (e.g., SFT + DPO) using SLERP interpolation
+to balance knowledge retention and alignment improvement.
+
+Reference: Nemotron-H paper — SLERP merging reduces alignment tax.
+
+Usage:
+    python scripts/merge_checkpoints.py \
+        --ckpt_a checkpoints/3b_sft_v2/checkpoint-best \
+        --ckpt_b checkpoints/3b_dpo/checkpoint-merged \
+        --output checkpoints/3b_dpo/checkpoint-slerp \
+        --alpha 0.5
+
+    alpha=0.0 → pure ckpt_a (SFT)
+    alpha=1.0 → pure ckpt_b (DPO)
+    alpha=0.5 → equal blend (recommended starting point)
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import shutil
+from pathlib import Path
+
+import torch
+import yaml
+
+
+def slerp(t: float, v0: torch.Tensor, v1: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
+    """Spherical linear interpolation between two tensors.
+
+    Args:
+        t: Interpolation factor in [0, 1]. 0 → v0, 1 → v1.
+        v0: First tensor (flattened internally).
+        v1: Second tensor (same shape as v0).
+        eps: Small value to avoid division by zero.
+
+    Returns:
+        Interpolated tensor with the same shape as v0.
+    """
+    original_shape = v0.shape
+    v0_flat = v0.flatten().float()
+    v1_flat = v1.flatten().float()
+
+    # Normalize
+    v0_norm = v0_flat / (v0_flat.norm() + eps)
+    v1_norm = v1_flat / (v1_flat.norm() + eps)
+
+    # Cosine of angle between vectors
+    cos_omega = torch.dot(v0_norm, v1_norm).clamp(-1.0, 1.0)
+
+    # If vectors are very similar, fall back to linear interpolation
+    if abs(cos_omega.item()) > 0.9995:
+        result = (1.0 - t) * v0_flat + t * v1_flat
+        return result.reshape(original_shape).to(v0.dtype)
+
+    omega = torch.acos(cos_omega)
+    sin_omega = torch.sin(omega)
+
+    s0 = torch.sin((1.0 - t) * omega) / sin_omega
+    s1 = torch.sin(t * omega) / sin_omega
+
+    # Interpolate using original (non-normalized) vectors to preserve scale
+    result = s0 * v0_flat + s1 * v1_flat
+    return result.reshape(original_shape).to(v0.dtype)
+
+
+def lerp(t: float, v0: torch.Tensor, v1: torch.Tensor) -> torch.Tensor:
+    """Simple linear interpolation."""
+    return ((1.0 - t) * v0.float() + t * v1.float()).to(v0.dtype)
+
+
+def merge_state_dicts(
+    sd_a: dict[str, torch.Tensor],
+    sd_b: dict[str, torch.Tensor],
+    alpha: float = 0.5,
+    method: str = "slerp",
+) -> dict[str, torch.Tensor]:
+    """Merge two state dicts using SLERP or LERP.
+
+    Args:
+        sd_a: State dict A (e.g., SFT model).
+        sd_b: State dict B (e.g., DPO model).
+        alpha: Interpolation factor. 0 → A, 1 → B.
+        method: "slerp" or "lerp".
+
+    Returns:
+        Merged state dict.
+    """
+    interp_fn = slerp if method == "slerp" else lerp
+
+    merged = {}
+    keys_a = set(sd_a.keys())
+    keys_b = set(sd_b.keys())
+
+    common = keys_a & keys_b
+    only_a = keys_a - keys_b
+    only_b = keys_b - keys_a
+
+    if only_a:
+        print(f"[WARN] {len(only_a)} keys only in ckpt_a (kept as-is)")
+    if only_b:
+        print(f"[WARN] {len(only_b)} keys only in ckpt_b (kept as-is)")
+
+    for key in sorted(common):
+        va = sd_a[key]
+        vb = sd_b[key]
+
+        if va.shape != vb.shape:
+            print(f"[WARN] Shape mismatch for {key}: {va.shape} vs {vb.shape}, keeping ckpt_a")
+            merged[key] = va
+            continue
+
+        # Only interpolate float parameters (skip int buffers, etc.)
+        if va.is_floating_point() and va.numel() > 1:
+            merged[key] = interp_fn(alpha, va, vb)
+        else:
+            merged[key] = va  # Keep from ckpt_a for non-float/scalar
+
+    # Include keys unique to each
+    for key in only_a:
+        merged[key] = sd_a[key]
+    for key in only_b:
+        merged[key] = sd_b[key]
+
+    return merged
+
+
+def main():
+    parser = argparse.ArgumentParser(description="SLERP checkpoint merge")
+    parser.add_argument("--ckpt_a", type=Path, required=True,
+                        help="Path to checkpoint A (e.g., SFT)")
+    parser.add_argument("--ckpt_b", type=Path, required=True,
+                        help="Path to checkpoint B (e.g., DPO)")
+    parser.add_argument("--output", type=Path, required=True,
+                        help="Output checkpoint directory")
+    parser.add_argument("--alpha", type=float, default=0.5,
+                        help="Interpolation factor (0=A, 1=B, default 0.5)")
+    parser.add_argument("--method", choices=["slerp", "lerp"], default="slerp",
+                        help="Interpolation method (default: slerp)")
+    args = parser.parse_args()
+
+    print(f"Merge: {args.ckpt_a.name} ←({1-args.alpha:.1%})— ({args.alpha:.1%})→ {args.ckpt_b.name}")
+    print(f"Method: {args.method}, alpha={args.alpha}")
+
+    # Load state dicts
+    print("Loading checkpoint A...")
+    sd_a = torch.load(args.ckpt_a / "model.pt", map_location="cpu", weights_only=True)
+    print(f"  {len(sd_a)} keys loaded")
+
+    print("Loading checkpoint B...")
+    sd_b = torch.load(args.ckpt_b / "model.pt", map_location="cpu", weights_only=True)
+    print(f"  {len(sd_b)} keys loaded")
+
+    # Merge
+    print("Merging...")
+    merged_sd = merge_state_dicts(sd_a, sd_b, alpha=args.alpha, method=args.method)
+    print(f"  {len(merged_sd)} keys in merged state dict")
+
+    # Save
+    args.output.mkdir(parents=True, exist_ok=True)
+    torch.save(merged_sd, args.output / "model.pt")
+
+    # Copy config from ckpt_a
+    config_src = args.ckpt_a / "config.yaml"
+    if config_src.exists():
+        shutil.copy2(str(config_src), str(args.output / "config.yaml"))
+
+    # Copy tokenizer if available
+    for tok_name in ["tokenizer.json", "tokenizer.model"]:
+        tok_src = args.ckpt_a / tok_name
+        if tok_src.exists():
+            shutil.copy2(str(tok_src), str(args.output / tok_name))
+
+    # Write merge metadata
+    meta = {
+        "ckpt_a": str(args.ckpt_a),
+        "ckpt_b": str(args.ckpt_b),
+        "alpha": args.alpha,
+        "method": args.method,
+    }
+    with open(args.output / "merge_info.yaml", "w") as f:
+        yaml.safe_dump(meta, f)
+
+    size_mb = (args.output / "model.pt").stat().st_size / 1e6
+    print(f"\nMerged checkpoint saved → {args.output} ({size_mb:.0f} MB)")
+
+
+if __name__ == "__main__":
+    main()

scripts/orpo_native.py

@@ -0,0 +1,636 @@
+"""
+train/orpo_native.py — ORPO (Odds Ratio Preference Optimization) training.
+
+Native ORPO implementation (no TRL, no HuggingFace Trainer) for EVAFRILL-Mo
+hybrid Mamba-2+Transformer models. Unlike DPO, ORPO requires NO reference model
+and performs SFT + alignment in a single training stage, making it ideal for
+starting from a raw pretrained checkpoint.
+
+Reference: Hong et al., "ORPO: Monolithic Preference Optimization without
+Reference Model" (2024), https://arxiv.org/abs/2403.07691
+
+Loss:
+    L_ORPO = L_SFT + λ * L_OR
+    L_SFT  = CrossEntropy(chosen_logits, chosen_labels)
+    L_OR   = -E[log σ(log(odds_chosen / odds_rejected))]
+    odds(x) = P(x) / (1 - P(x)),  P(x) = exp(avg_log_prob(x))
+
+Launch:
+    python train/orpo_native.py \
+        --pretrained_checkpoint checkpoints/3b_final/checkpoint-0319772 \
+        --preference_data data/preference/combined_preference.jsonl \
+        --config configs/h100_mig/dpo_3b_1gpu.yaml \
+        --device cuda:0
+"""
+
+from __future__ import annotations
+
+import argparse
+import datetime
+import os
+import random
+import signal
+import shutil
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, RandomSampler
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+torch.set_float32_matmul_precision("high")
+
+_PROJECT_ROOT = Path(__file__).resolve().parent.parent
+if str(_PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(_PROJECT_ROOT))
+
+from model import LLM
+from model.lora import apply_lora, get_lora_params, merge_lora, save_lora
+from data.dpo_dataset import DPODataset, dpo_collate_fn
+from train.utils import (
+    get_cosine_schedule_with_warmup,
+    is_main_process,
+    save_checkpoint,
+    load_checkpoint,
+)
+
+
+# ---------------------------------------------------------------------------
+# Argument parsing
+# ---------------------------------------------------------------------------
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="ORPO Training for EVAFRILL-Mo")
+
+    # Paths
+    parser.add_argument("--pretrained_checkpoint", type=Path, required=True,
+                        help="Path to pretrained model checkpoint directory "
+                             "(e.g. checkpoints/3b_final/checkpoint-0319772)")
+    parser.add_argument("--preference_data", type=Path, required=True,
+                        help="Path to preference JSONL data (prompt/chosen/rejected)")
+    parser.add_argument("--checkpoint_dir", type=Path, default=Path("checkpoints/3b_orpo"),
+                        help="Output checkpoint directory (default: checkpoints/3b_orpo)")
+    parser.add_argument("--resume", type=Path, default=None,
+                        help="Resume training from an existing ORPO checkpoint directory")
+    parser.add_argument("--tokenizer", type=Path, default=None,
+                        help="Path to tokenizer.json (auto-detected if omitted)")
+    parser.add_argument("--log_file", type=Path, default=None,
+                        help="Append logs to this file in addition to stdout")
+    parser.add_argument("--config", type=Path, default=None,
+                        help="YAML config to load defaults from (train: section)")
+
+    # ORPO hyperparameters
+    parser.add_argument("--lambda_or", type=float, default=1.0,
+                        help="ORPO odds-ratio loss weight λ (default: 1.0)")
+    parser.add_argument("--max_steps", type=int, default=3000,
+                        help="Total optimisation steps (default: 3000)")
+    parser.add_argument("--batch_size", type=int, default=1,
+                        help="Per-step micro-batch size (default: 1)")
+    parser.add_argument("--grad_accum", type=int, default=16,
+                        help="Gradient accumulation steps (default: 16)")
+    parser.add_argument("--lr", type=float, default=5e-6,
+                        help="Peak learning rate (default: 5e-6; higher than DPO because "
+                             "ORPO starts from pretrained, not SFT)")
+    parser.add_argument("--weight_decay", type=float, default=0.01)
+    parser.add_argument("--warmup_steps", type=int, default=100)
+    parser.add_argument("--max_length", type=int, default=1024)
+    parser.add_argument("--seed", type=int, default=42)
+
+    # LoRA
+    parser.add_argument("--use_lora", action="store_true", default=True,
+                        help="Use LoRA adapters for memory-efficient training (default: on)")
+    parser.add_argument("--lora_rank", type=int, default=32)
+    parser.add_argument("--lora_alpha", type=float, default=64.0)
+
+    # Infrastructure
+    parser.add_argument("--device", type=str, default=None,
+                        help="Device string, e.g. cuda:0 (auto-detected if omitted)")
+    parser.add_argument("--save_interval", type=int, default=500)
+    parser.add_argument("--log_interval", type=int, default=10)
+    parser.add_argument("--num_workers", type=int, default=4)
+
+    args, _ = parser.parse_known_args()
+
+    # Load YAML config and apply as defaults (CLI flags override YAML)
+    if args.config is not None:
+        if not args.config.exists():
+            raise FileNotFoundError(f"Config not found: {args.config}")
+        import yaml
+        with open(args.config) as f:
+            cfg = yaml.safe_load(f)
+        train_cfg = cfg.get("train", {})
+        yaml_map = {
+            "max_steps":       "max_steps",
+            "batch_size":      "batch_size",
+            "grad_accum_steps": "grad_accum",
+            "lr":              "lr",
+            "weight_decay":    "weight_decay",
+            "warmup_steps":    "warmup_steps",
+            "lambda_or":       "lambda_or",
+            "max_length":      "max_length",
+            "save_interval":   "save_interval",
+            "log_interval":    "log_interval",
+            "use_lora":        "use_lora",
+            "lora_rank":       "lora_rank",
+            "lora_alpha":      "lora_alpha",
+        }
+        defaults: dict = {}
+        for yk, ak in yaml_map.items():
+            if yk in train_cfg:
+                defaults[ak] = train_cfg[yk]
+        if defaults:
+            parser.set_defaults(**defaults)
+
+    return parser.parse_args()
+
+
+# ---------------------------------------------------------------------------
+# Utilities
+# ---------------------------------------------------------------------------
+
+def set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def _resolve_tokenizer_path(args: argparse.Namespace) -> Path:
+    """Find tokenizer.json: explicit flag > checkpoint dir > project default."""
+    if args.tokenizer is not None:
+        return Path(args.tokenizer)
+    ckpt_tok = args.pretrained_checkpoint / "tokenizer.json"
+    if ckpt_tok.exists():
+        return ckpt_tok
+    default_tok = _PROJECT_ROOT / "tokenizer" / "korean_sp" / "tokenizer.json"
+    if default_tok.exists():
+        return default_tok
+    raise FileNotFoundError(
+        "Cannot find tokenizer.json. Provide --tokenizer or place it in the checkpoint dir."
+    )
+
+
+# ---------------------------------------------------------------------------
+# ORPO loss
+# ---------------------------------------------------------------------------
+
+def get_avg_log_prob(
+    logits: torch.Tensor,
+    labels: torch.Tensor,
+) -> torch.Tensor:
+    """Compute average log probability over non-masked (response) tokens.
+
+    Args:
+        logits: (B, T, V) raw model logits — already in float32.
+        labels: (B, T) token ids; -1 marks prompt/padding positions to ignore.
+
+    Returns:
+        (B,) mean log probability over response tokens per sample.
+        Returns 0 for samples where no response token is present (shouldn't
+        happen with well-formed data, but guarded for safety).
+    """
+    log_probs = F.log_softmax(logits.float(), dim=-1)  # (B, T, V)
+
+    mask = labels != -1                                # (B, T)  True = response token
+    safe_labels = labels.clamp(min=0)                  # replace -1 with 0 for gather
+    per_token_logps = log_probs.gather(
+        -1, safe_labels.unsqueeze(-1)
+    ).squeeze(-1)                                      # (B, T)
+
+    # Zero out masked positions
+    per_token_logps = per_token_logps * mask.float()   # (B, T)
+
+    # Average over response tokens; clamp denominator to avoid div-by-zero
+    n_tokens = mask.float().sum(dim=-1).clamp(min=1.0)  # (B,)
+    return per_token_logps.sum(dim=-1) / n_tokens        # (B,)
+
+
+def compute_orpo_loss(
+    model: nn.Module,
+    chosen_ids: torch.Tensor,
+    chosen_labels: torch.Tensor,
+    rejected_ids: torch.Tensor,
+    rejected_labels: torch.Tensor,
+    lambda_or: float = 1.0,
+    vocab_size: int | None = None,
+) -> tuple[torch.Tensor, float, float]:
+    """Compute ORPO loss = SFT loss + λ * OR loss.
+
+    No reference model is needed.  The SFT loss trains the model to generate
+    chosen responses; the OR loss simultaneously teaches the model to prefer
+    chosen over rejected by maximising the log odds ratio.
+
+    Args:
+        model:            The policy model (frozen base + trainable LoRA).
+        chosen_ids:       (B, T) token ids for chosen sequences.
+        chosen_labels:    (B, T) labels for chosen; -1 on prompt tokens.
+        rejected_ids:     (B, T) token ids for rejected sequences.
+        rejected_labels:  (B, T) labels for rejected; -1 on prompt tokens.
+        lambda_or:        Weight of the OR loss term (paper default = 1.0).
+        vocab_size:       Vocabulary size for reshape; inferred from logits if None.
+
+    Returns:
+        (total_loss, sft_loss_scalar, or_loss_scalar)
+    """
+    # -----------------------------------------------------------------------
+    # 1. Forward pass — chosen
+    # -----------------------------------------------------------------------
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        chosen_logits, _ = model(chosen_ids)   # (B, T, V)
+
+    # Infer vocab size from logits if not given
+    V = chosen_logits.size(-1) if vocab_size is None else vocab_size
+
+    # SFT loss: next-token prediction on response positions only.
+    # logits[:, :-1] predicts labels[:, 1:] (standard causal shift).
+    sft_logits = chosen_logits[:, :-1].contiguous().reshape(-1, V).float()
+    sft_targets = chosen_labels[:, 1:].contiguous().reshape(-1)
+
+    # F.cross_entropy ignores index -1 via ignore_index; -1 covers prompt tokens
+    # AND the last padding position shifted out of the window.
+    sft_loss: torch.Tensor = F.cross_entropy(sft_logits, sft_targets, ignore_index=-1)
+
+    # Average log-prob over response tokens (used for OR computation)
+    # Labels are NOT shifted here — get_avg_log_prob handles the alignment
+    # by using labels directly as targets at each position.
+    chosen_avg_logp: torch.Tensor = get_avg_log_prob(chosen_logits.float(), chosen_labels)
+
+    # -----------------------------------------------------------------------
+    # 2. Forward pass — rejected
+    # -----------------------------------------------------------------------
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+        rejected_logits, _ = model(rejected_ids)  # (B, T, V)
+
+    rejected_avg_logp: torch.Tensor = get_avg_log_prob(rejected_logits.float(), rejected_labels)
+
+    # -----------------------------------------------------------------------
+    # 3. Odds ratio loss
+    #
+    # odds(x)  = P(x) / (1 - P(x))
+    # log odds = log P(x) - log(1 - P(x)) = log P(x) - log1p(-exp(log P(x)))
+    #
+    # We use log1p(-exp(·)) with clamping to keep values numerically stable:
+    #   - avg_log_prob is always ≤ 0
+    #   - exp(avg_log_prob) ∈ (0, 1] → 1 - exp ∈ [0, 1)
+    #   - clamp to avoid log(0) when avg_log_prob ≈ 0 (very high confidence)
+    # -----------------------------------------------------------------------
+    # Clamp to (-33, -1e-6): upper bound avoids 1-exp≈0 → log(0); lower keeps
+    # values finite (exp(-33) ≈ 5e-15, no underflow in float32).
+    eps_low, eps_high = -33.0, -1e-6
+
+    chosen_avg_logp_clamped   = chosen_avg_logp.clamp(eps_low, eps_high)
+    rejected_avg_logp_clamped = rejected_avg_logp.clamp(eps_low, eps_high)
+
+    log_odds_chosen   = chosen_avg_logp_clamped   - torch.log1p(-chosen_avg_logp_clamped.exp())
+    log_odds_rejected = rejected_avg_logp_clamped - torch.log1p(-rejected_avg_logp_clamped.exp())
+
+    log_odds_ratio = log_odds_chosen - log_odds_rejected  # (B,)
+    or_loss: torch.Tensor = -F.logsigmoid(log_odds_ratio).mean()
+
+    # -----------------------------------------------------------------------
+    # 4. Combined loss
+    # -----------------------------------------------------------------------
+    total_loss = sft_loss + lambda_or * or_loss
+
+    return total_loss, sft_loss.item(), or_loss.item()
+
+
+# ---------------------------------------------------------------------------
+# Main training loop
+# ---------------------------------------------------------------------------
+
+def main() -> None:
+    args = parse_args()
+    set_seed(args.seed)
+
+    # ------------------------------------------------------------------
+    # Device
+    # ------------------------------------------------------------------
+    if args.device:
+        device = torch.device(args.device)
+    elif torch.cuda.is_available():
+        device = torch.device("cuda:0")
+    else:
+        device = torch.device("cpu")
+
+    # ------------------------------------------------------------------
+    # Load pretrained model
+    # ------------------------------------------------------------------
+    if not args.pretrained_checkpoint.exists():
+        raise FileNotFoundError(
+            f"Pretrained checkpoint not found: {args.pretrained_checkpoint}"
+        )
+
+    print(f"Loading pretrained model from {args.pretrained_checkpoint} ...")
+    model: nn.Module = LLM.from_pretrained(args.pretrained_checkpoint)
+    model.config.use_fp8 = False   # H100 MIG: BF16 only; B200 may set fp8 via config
+    model = model.to(device=device, dtype=torch.bfloat16)
+
+    # Gradient checkpointing — reduces VRAM at cost of ~20% speed
+    if hasattr(model, "gradient_checkpointing_enable"):
+        model.gradient_checkpointing_enable()
+        print("[INFO] Gradient checkpointing enabled")
+
+    # ------------------------------------------------------------------
+    # LoRA
+    # ------------------------------------------------------------------
+    if args.use_lora:
+        n_lora = apply_lora(model, rank=args.lora_rank, alpha=args.lora_alpha)
+        lora_params = get_lora_params(model)
+        print(f"[INFO] LoRA: {n_lora:,} trainable params "
+              f"(rank={args.lora_rank}, alpha={args.lora_alpha})")
+    else:
+        lora_params = None
+        print("[INFO] Full fine-tuning (all parameters trainable)")
+
+    total_params     = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"Total params: {total_params:,}  |  Trainable: {trainable_params:,}")
+
+    # ------------------------------------------------------------------
+    # Tokenizer
+    # ------------------------------------------------------------------
+    tokenizer_path = _resolve_tokenizer_path(args)
+    print(f"Loading tokenizer from {tokenizer_path}")
+    from tokenizers import Tokenizer
+    tokenizer = Tokenizer.from_file(str(tokenizer_path))
+
+    # ------------------------------------------------------------------
+    # Dataset & DataLoader
+    # ------------------------------------------------------------------
+    train_dataset = DPODataset(
+        data_path=args.preference_data,
+        tokenizer=tokenizer,
+        max_seq_len=args.max_length,
+    )
+    if len(train_dataset) == 0:
+        raise ValueError(f"Preference dataset is empty: {args.preference_data}")
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=args.batch_size,
+        sampler=RandomSampler(train_dataset),
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=True,
+        collate_fn=dpo_collate_fn,
+        prefetch_factor=2,
+        persistent_workers=(args.num_workers > 0),
+    )
+
+    # ------------------------------------------------------------------
+    # Optimizer
+    # ------------------------------------------------------------------
+    if lora_params is not None:
+        opt_params = lora_params
+    else:
+        opt_params = [p for p in model.parameters() if p.requires_grad]
+
+    optimizer = torch.optim.AdamW(
+        opt_params,
+        lr=args.lr,
+        betas=(0.9, 0.95),
+        weight_decay=args.weight_decay,
+        fused=torch.cuda.is_available(),
+    )
+
+    scheduler = get_cosine_schedule_with_warmup(
+        optimizer=optimizer,
+        warmup_steps=args.warmup_steps,
+        total_steps=args.max_steps,
+    )
+
+    # ------------------------------------------------------------------
+    # Resume
+    # ------------------------------------------------------------------
+    start_step = 0
+    if args.resume is not None:
+        if not args.resume.exists():
+            raise FileNotFoundError(f"Resume checkpoint not found: {args.resume}")
+        start_step, _ = load_checkpoint(args.resume, model, optimizer, scheduler)
+        print(f"Resumed from step {start_step}")
+
+    # ------------------------------------------------------------------
+    # Output directory & tokenizer copy
+    # ------------------------------------------------------------------
+    args.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    dest_tok = args.checkpoint_dir / "tokenizer.json"
+    if not dest_tok.exists():
+        shutil.copy2(str(tokenizer_path), str(dest_tok))
+
+    # ------------------------------------------------------------------
+    # Logger
+    # ------------------------------------------------------------------
+    log_fh = None
+    if args.log_file:
+        Path(args.log_file).parent.mkdir(parents=True, exist_ok=True)
+        log_fh = open(args.log_file, "a", encoding="utf-8", buffering=1)
+
+    def log(msg: str, level: str = "INFO") -> None:
+        ts   = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        line = f"[{ts}] [{level}] {msg}"
+        print(line, flush=True)
+        if log_fh:
+            log_fh.write(line + "\n")
+
+    # ------------------------------------------------------------------
+    # Training banner
+    # ------------------------------------------------------------------
+    eff_batch = args.batch_size * args.grad_accum
+    log("=" * 65)
+    log("ORPO Training — EVAFRILL-Mo")
+    log(f"  Pretrained ckpt : {args.pretrained_checkpoint}")
+    log(f"  Preference data : {args.preference_data} ({len(train_dataset):,} samples)")
+    log(f"  LoRA            : rank={args.lora_rank} alpha={args.lora_alpha} "
+        f"enabled={args.use_lora}")
+    log(f"  lambda_or={args.lambda_or}, lr={args.lr:.2e}, eff_batch={eff_batch}")
+    log(f"  max_steps={args.max_steps}, warmup={args.warmup_steps}, "
+        f"max_len={args.max_length}")
+    log(f"  device={device}")
+    log("=" * 65)
+
+    # ------------------------------------------------------------------
+    # Graceful shutdown handler
+    # ------------------------------------------------------------------
+    shutdown_requested = False
+
+    def shutdown_handler(signum, frame):
+        nonlocal shutdown_requested
+        shutdown_requested = True
+        log(f"Shutdown signal received (sig={signum}). Saving checkpoint ...", "WARN")
+
+    signal.signal(signal.SIGTERM, shutdown_handler)
+    signal.signal(signal.SIGINT,  shutdown_handler)
+    try:
+        signal.signal(signal.SIGHUP, shutdown_handler)
+    except AttributeError:
+        pass  # Windows does not have SIGHUP
+
+    # ------------------------------------------------------------------
+    # Data iterator (infinite, cycling through epochs)
+    # ------------------------------------------------------------------
+    import time
+
+    epoch = 0
+    loader_iter = iter(train_loader)
+
+    def next_batch() -> tuple[torch.Tensor, ...]:
+        nonlocal loader_iter, epoch
+        try:
+            return next(loader_iter)
+        except StopIteration:
+            epoch += 1
+            log(f"--- Epoch {epoch} begin ---")
+            loader_iter = iter(train_loader)
+            return next(loader_iter)
+
+    # ------------------------------------------------------------------
+    # Training loop
+    # ------------------------------------------------------------------
+    model.train()
+
+    # Running statistics (reset every log_interval steps)
+    running_total_loss = 0.0
+    running_sft_loss   = 0.0
+    running_or_loss    = 0.0
+    log_step_count     = 0
+    t0 = time.perf_counter()
+
+    # Keep track of the last loss value for the final checkpoint call
+    avg_loss = float("nan")
+
+    for step in range(start_step, args.max_steps):
+        optimizer.zero_grad(set_to_none=True)
+
+        accum_total = 0.0
+        accum_sft   = 0.0
+        accum_or    = 0.0
+
+        # ---- Gradient accumulation ----------------------------------------
+        for _micro in range(args.grad_accum):
+            batch = next_batch()
+            chosen_ids      = batch[0].to(device, dtype=torch.long, non_blocking=True)
+            chosen_labels   = batch[1].to(device, dtype=torch.long, non_blocking=True)
+            rejected_ids    = batch[2].to(device, dtype=torch.long, non_blocking=True)
+            rejected_labels = batch[3].to(device, dtype=torch.long, non_blocking=True)
+
+            loss, sft_l, or_l = compute_orpo_loss(
+                model,
+                chosen_ids, chosen_labels,
+                rejected_ids, rejected_labels,
+                lambda_or=args.lambda_or,
+            )
+
+            scaled_loss = loss / args.grad_accum
+            scaled_loss.backward()
+
+            accum_total += loss.item()
+            accum_sft   += sft_l
+            accum_or    += or_l
+
+        # ---- Gradient clipping --------------------------------------------
+        grad_norm = torch.nn.utils.clip_grad_norm_(
+            [p for p in model.parameters() if p.requires_grad],
+            max_norm=1.0,
+        ).item()
+
+        optimizer.step()
+        scheduler.step()
+
+        # ---- Accumulate stats ---------------------------------------------
+        avg_total = accum_total / args.grad_accum
+        avg_sft   = accum_sft   / args.grad_accum
+        avg_or    = accum_or    / args.grad_accum
+
+        running_total_loss += avg_total
+        running_sft_loss   += avg_sft
+        running_or_loss    += avg_or
+        log_step_count     += 1
+        avg_loss            = avg_total   # for use in checkpoint call
+
+        # ---- Graceful shutdown check --------------------------------------
+        if shutdown_requested:
+            log(f"Graceful shutdown at step {step + 1}", "WARN")
+            ckpt_path = save_checkpoint(
+                model, optimizer, scheduler,
+                step + 1, avg_loss, str(args.checkpoint_dir)
+            )
+            if args.use_lora:
+                save_lora(model, args.checkpoint_dir / f"lora-{step+1:07d}")
+            log(f"Checkpoint saved -> {ckpt_path}")
+            break
+
+        # ---- Logging ------------------------------------------------------
+        if (step + 1) % args.log_interval == 0:
+            t1      = time.perf_counter()
+            elapsed = t1 - t0
+
+            mean_total = running_total_loss / log_step_count
+            mean_sft   = running_sft_loss   / log_step_count
+            mean_or    = running_or_loss    / log_step_count
+            lr_now     = scheduler.get_last_lr()[0]
+            mem_gb     = torch.cuda.memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0
+            sps        = log_step_count / max(elapsed, 1e-6)  # steps per second
+
+            log(
+                f"step {step+1:>6d}/{args.max_steps} | "
+                f"loss {mean_total:.4f} "
+                f"(sft {mean_sft:.4f} or {mean_or:.4f}) | "
+                f"lr {lr_now:.2e} | "
+                f"gnorm {grad_norm:.3f} | "
+                f"mem {mem_gb:.1f}GB | "
+                f"{sps:.2f}step/s"
+            )
+
+            running_total_loss = 0.0
+            running_sft_loss   = 0.0
+            running_or_loss    = 0.0
+            log_step_count     = 0
+            t0 = t1
+
+        # ---- Periodic checkpoint ------------------------------------------
+        if (step + 1) % args.save_interval == 0:
+            ckpt_path = save_checkpoint(
+                model, optimizer, scheduler,
+                step + 1, avg_loss, str(args.checkpoint_dir)
+            )
+            if args.use_lora:
+                save_lora(model, args.checkpoint_dir / f"lora-{step+1:07d}")
+            log(f"Checkpoint saved -> {ckpt_path}")
+
+    # -----------------------------------------------------------------------
+    # Final checkpoint
+    # -----------------------------------------------------------------------
+    if not shutdown_requested:
+        final_path = save_checkpoint(
+            model, optimizer, scheduler,
+            args.max_steps, avg_loss, str(args.checkpoint_dir)
+        )
+        if args.use_lora:
+            save_lora(model, args.checkpoint_dir / "lora-final")
+        log(f"Final checkpoint -> {final_path}")
+
+    # -----------------------------------------------------------------------
+    # LoRA merge + save merged model
+    # -----------------------------------------------------------------------
+    if args.use_lora:
+        log("Merging LoRA weights into base model ...")
+        merge_lora(model)
+        merged_dir = args.checkpoint_dir / "checkpoint-merged"
+        model.save_pretrained(merged_dir)
+        # Also copy tokenizer into merged dir for easy inference
+        shutil.copy2(str(dest_tok), str(merged_dir / "tokenizer.json"))
+        log(f"Merged model saved -> {merged_dir}")
+
+    log("ORPO training complete.")
+
+    if log_fh:
+        log_fh.close()
+
+
+if __name__ == "__main__":
+    main()

scripts/sft.py

@@ -0,0 +1,854 @@
+"""
+train/sft.py — Supervised Fine-Tuning (SFT) entry point.
+
+Loads a pretrained checkpoint and fine-tunes it on instruction/conversation
+data using SFTDataset, which masks prompt tokens with ignore_index=-1 so only
+the assistant response tokens contribute to the loss.
+
+Launch single-GPU:
+    python train/sft.py \\
+        --base_checkpoint checkpoints/korean_1b_fp8_run1/checkpoint-0034000 \\
+        --sft_data data/sft/train.jsonl \\
+        --device cuda:0
+
+Launch multi-GPU (DDP via torchrun, 7 GPU):
+    torchrun --nproc_per_node=7 train/sft.py \\
+        --base_checkpoint checkpoints/3b_final/checkpoint-0319772 \\
+        --sft_data data/sft_combined/train_filtered.jsonl
+
+KEY DIFFERENCES from pretrain.py:
+  - Loads weights from a pretrained checkpoint via LLM.from_pretrained()
+  - Uses SFTDataset (JSONL instruction data) instead of PackedDataset
+  - Lower default learning rate (2e-5 vs 2e-4)
+  - Fewer default steps (3000 vs 100000)
+  - Copies tokenizer.json to checkpoint_dir for easy deployment
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import random
+import signal
+import shutil
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, DistributedSampler, RandomSampler
+
+# B200 Tensor Core 최대 활용: TF32 matmul + cuDNN
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+torch.set_float32_matmul_precision("high")  # TF32 precision for fp32 matmul
+
+# Allow imports from the project root regardless of working directory.
+_PROJECT_ROOT = Path(__file__).resolve().parent.parent
+if str(_PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(_PROJECT_ROOT))
+
+from model import LLM
+from train.trainer import TrainConfig, Trainer
+from train.utils import (
+    cleanup_ddp,
+    get_cosine_schedule_with_warmup,
+    is_main_process,
+    load_checkpoint,
+    setup_ddp,
+)
+
+# ---------------------------------------------------------------------------
+# Optional TransformerEngine import (FP8 support)
+# ---------------------------------------------------------------------------
+try:
+    import transformer_engine.pytorch as te  # type: ignore[import]
+    HAS_TE = True
+except ImportError:
+    te = None  # type: ignore[assignment]
+    HAS_TE = False
+
+
+# ---------------------------------------------------------------------------
+# Argument parsing
+# ---------------------------------------------------------------------------
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Supervised Fine-Tuning (SFT) of a pretrained decoder-only LLM.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+
+    # --- Required paths -----------------------------------------------------
+    parser.add_argument(
+        "--base_checkpoint",
+        type=Path,
+        required=True,
+        help=(
+            "Path to the pretrained checkpoint directory. "
+            "Must contain model.pt and config.yaml (produced by save_checkpoint)."
+        ),
+    )
+    parser.add_argument(
+        "--sft_data",
+        type=Path,
+        required=True,
+        help="Path to the JSONL SFT training data file.",
+    )
+
+    # --- Optional paths -----------------------------------------------------
+    parser.add_argument(
+        "--val_data",
+        type=Path,
+        default=None,
+        help="Optional path to JSONL SFT validation data file.",
+    )
+    parser.add_argument(
+        "--checkpoint_dir",
+        type=Path,
+        default=Path("checkpoints/korean_1b_sft"),
+        help="Root directory for saving SFT checkpoints.",
+    )
+    parser.add_argument(
+        "--resume",
+        type=Path,
+        default=None,
+        help="Path to an SFT checkpoint directory to resume fine-tuning from.",
+    )
+    parser.add_argument(
+        "--tokenizer",
+        type=Path,
+        default=None,
+        help=(
+            "Override path to tokenizer.json. "
+            "Defaults to <base_checkpoint>/tokenizer.json, "
+            "then falls back to tokenizer/korean_sp/tokenizer.json."
+        ),
+    )
+    parser.add_argument(
+        "--log_file",
+        type=Path,
+        default=None,
+        help=(
+            "Path to a text file for structured training logs (rank-0 only). "
+            "If omitted, logs go only to stdout."
+        ),
+    )
+
+    # --- Training hyper-parameters ------------------------------------------
+    parser.add_argument(
+        "--max_steps",
+        type=int,
+        default=3000,
+        help="Total number of optimiser steps.",
+    )
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=4,
+        help="Per-GPU micro-batch size.",
+    )
+    parser.add_argument(
+        "--lr",
+        type=float,
+        default=2e-5,
+        help=(
+            "Peak learning rate. "
+            "SFT uses a much lower lr than pretraining (2e-5 vs 2e-4) "
+            "to preserve pretrained representations."
+        ),
+    )
+    parser.add_argument(
+        "--weight_decay",
+        type=float,
+        default=0.01,
+        help="AdamW weight decay. Lower than pretrain (0.01 vs 0.1).",
+    )
+    parser.add_argument(
+        "--warmup_steps",
+        type=int,
+        default=100,
+        help="Number of linear LR warmup steps.",
+    )
+    parser.add_argument(
+        "--grad_accum",
+        type=int,
+        default=2,
+        help="Gradient accumulation steps.",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Base random seed (rank offset is added automatically in DDP).",
+    )
+    parser.add_argument(
+        "--use_fp8",
+        action="store_true",
+        default=False,
+        help=(
+            "Enable TransformerEngine FP8 training "
+            "(requires B200/H100, uses MXFP8BlockScaling)."
+        ),
+    )
+
+    # --- Single-GPU device override (ignored when using torchrun) -----------
+    parser.add_argument(
+        "--device",
+        type=str,
+        default=None,
+        help=(
+            "Explicit device string (e.g. 'cuda:0'). "
+            "Ignored when running under torchrun (DDP auto-assigns devices)."
+        ),
+    )
+
+    parser.add_argument(
+        "--config", type=Path, default=None,
+        help="YAML config file. Values under 'train:' section are used as CLI defaults.",
+    )
+    parser.add_argument("--save_interval", type=int, default=500, help="Checkpoint save interval (steps).")
+    parser.add_argument("--eval_interval", type=int, default=250, help="Validation eval interval (steps).")
+    parser.add_argument("--neftune_alpha", type=float, default=5.0, help="NEFTune noise magnitude (0 to disable).")
+    parser.add_argument("--no_fp8", action="store_true", default=False, help="Force disable FP8 even if pretrained config has use_fp8=True.")
+    parser.add_argument("--num_workers", type=int, default=4, help="Number of DataLoader worker processes.")
+    parser.add_argument("--max_val_batches", type=int, default=0, help="Max validation batches (0=unlimited).")
+
+    # First pass: just get --config
+    args, remaining = parser.parse_known_args()
+
+    # Load YAML config and apply values as defaults
+    if args.config is not None:
+        if not args.config.exists():
+            raise FileNotFoundError(f"Config file not found: {args.config}")
+        import yaml
+        with open(args.config, "r") as f:
+            yaml_cfg = yaml.safe_load(f)
+        train_section = yaml_cfg.get("train", {})
+        yaml_to_arg = {
+            "max_steps": "max_steps",
+            "batch_size": "batch_size",
+            "lr": "lr",
+            "weight_decay": "weight_decay",
+            "warmup_steps": "warmup_steps",
+            "grad_accum_steps": "grad_accum",
+            "save_interval": "save_interval",
+            "eval_interval": "eval_interval",
+            "neftune_alpha": "neftune_alpha",
+            "max_val_batches": "max_val_batches",
+        }
+        new_defaults = {}
+        for yaml_key, arg_name in yaml_to_arg.items():
+            if yaml_key in train_section:
+                new_defaults[arg_name] = train_section[yaml_key]
+        if new_defaults:
+            parser.set_defaults(**new_defaults)
+
+    return parser.parse_args()
+
+
+# ---------------------------------------------------------------------------
+# Seed helper
+# ---------------------------------------------------------------------------
+
+
+def set_seed(seed: int) -> None:
+    """Set deterministic seeds for Python, NumPy, and PyTorch."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+# ---------------------------------------------------------------------------
+# Optimizer parameter groups
+# (Copied from pretrain.py to avoid circular import; identical logic)
+# ---------------------------------------------------------------------------
+
+
+def build_optimizer_param_groups(
+    model: torch.nn.Module,
+    weight_decay: float,
+) -> list[dict]:
+    """
+    Split parameters into two groups:
+      - decay group   : weight tensors with ndim >= 2 (Linear, etc.)
+      - no-decay group: bias, LayerNorm/RMSNorm weights, and embedding weights
+
+    This follows standard practice (e.g. GPT-style training).
+    """
+    decay_params: list[torch.nn.Parameter] = []
+    no_decay_params: list[torch.nn.Parameter] = []
+
+    # Module types whose parameters should never be decayed.
+    no_decay_module_types = (
+        torch.nn.Embedding,
+        torch.nn.LayerNorm,
+    )
+    # Also skip any parameter whose name ends with '.bias'.
+    no_decay_name_suffixes = ("bias",)
+
+    # Collect module-level exclusions.
+    no_decay_module_params: set[int] = set()
+    for module in model.modules():
+        if isinstance(module, no_decay_module_types):
+            for param in module.parameters(recurse=False):
+                no_decay_module_params.add(id(param))
+
+    seen: set[int] = set()
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        if id(param) in seen:
+            continue
+        seen.add(id(param))
+
+        if (
+            id(param) in no_decay_module_params
+            or any(name.endswith(sfx) for sfx in no_decay_name_suffixes)
+            or param.ndim < 2
+        ):
+            no_decay_params.append(param)
+        else:
+            decay_params.append(param)
+
+    return [
+        {"params": decay_params, "weight_decay": weight_decay},
+        {"params": no_decay_params, "weight_decay": 0.0},
+    ]
+
+
+# ---------------------------------------------------------------------------
+# Tokenizer resolution helper
+# ---------------------------------------------------------------------------
+
+
+def _resolve_tokenizer_path(args: argparse.Namespace) -> Path:
+    """
+    Determine the tokenizer path in priority order:
+      1. Explicit --tokenizer argument
+      2. tokenizer.json inside the base_checkpoint directory
+      3. Project default: tokenizer/korean_sp/tokenizer.json
+    """
+    if args.tokenizer is not None:
+        p = Path(args.tokenizer)
+        if not p.exists():
+            raise FileNotFoundError(f"Tokenizer not found at --tokenizer path: {p}")
+        return p
+
+    ckpt_tok = args.base_checkpoint / "tokenizer.json"
+    if ckpt_tok.exists():
+        return ckpt_tok
+
+    default_tok = _PROJECT_ROOT / "tokenizer" / "korean_sp" / "tokenizer.json"
+    if default_tok.exists():
+        return default_tok
+
+    raise FileNotFoundError(
+        "Could not locate tokenizer.json. Tried:\n"
+        f"  1. {ckpt_tok}\n"
+        f"  2. {default_tok}\n"
+        "Use --tokenizer to specify an explicit path."
+    )
+
+
+# ---------------------------------------------------------------------------
+# Dynamic padding collate function
+# ---------------------------------------------------------------------------
+
+
+def dynamic_collate_fn(batch: list) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Collate function that pads each batch to its own maximum sequence length
+    instead of a fixed global max_seq_len.  This reduces wasted FLOPs on
+    short sequences and speeds up SFT which tends to have highly variable
+    response lengths.
+
+    Pads to the batch-local max, aligned to 64 tokens (for Flash Attention
+    efficiency), with a floor of 512 tokens so micro-batches are not too short.
+
+    Args:
+        batch: List of ``(input_ids, labels)`` tuples from SFTDataset.
+
+    Returns:
+        Tuple of ``(input_ids, labels, attention_mask)`` tensors shaped
+        ``[B, max_len]``.
+        ``input_ids``      is right-padded with 0 (pad token).
+        ``labels``         is right-padded with -1 (cross-entropy ignore_index).
+        ``attention_mask`` is 1 for real tokens, 0 for padding.
+    """
+    # 64-token alignment + minimum 512 floor
+    raw_max = max(item[0].size(0) for item in batch)
+    max_len = max(512, ((raw_max + 63) // 64) * 64)
+
+    input_ids_list, labels_list, mask_list = [], [], []
+    for ids, labs in batch:
+        pad_len = max_len - ids.size(0)
+        input_ids_list.append(F.pad(ids, (0, pad_len), value=0))
+        labels_list.append(F.pad(labs, (0, pad_len), value=-1))
+        mask_list.append(
+            F.pad(torch.ones(ids.size(0), dtype=torch.long), (0, pad_len), value=0)
+        )
+
+    return (
+        torch.stack(input_ids_list),
+        torch.stack(labels_list),
+        torch.stack(mask_list),
+    )
+
+
+# ---------------------------------------------------------------------------
+# NEFTune helper
+# ---------------------------------------------------------------------------
+
+
+def add_neftune_hook(model: torch.nn.Module, noise_alpha: float = 10.0):
+    """
+    Register a forward hook on the model's input embedding layer that adds
+    uniform noise scaled by noise_alpha during training (NEFTune).
+
+    Reference: "NEFTune: Noisy Embeddings Improve Instruction Finetuning"
+    (Jain et al., 2023). https://arxiv.org/abs/2310.05914
+
+    Args:
+        model:       Raw (non-DDP) model instance.
+        noise_alpha: Noise magnitude parameter (paper default: 10).
+
+    Returns:
+        The hook handle (call ``handle.remove()`` to deactivate), or None if
+        the embedding layer could not be located.
+    """
+    # Unwrap DDP if needed
+    raw = model.module if hasattr(model, "module") else model
+
+    # 1) Try the standard HuggingFace accessor first.
+    embedding: torch.nn.Embedding | None = None
+    if hasattr(raw, "get_input_embeddings"):
+        try:
+            emb = raw.get_input_embeddings()
+            if isinstance(emb, torch.nn.Embedding):
+                embedding = emb
+        except Exception:
+            pass
+
+    # 2) Fallback: walk common attribute paths found in open-source LLMs.
+    if embedding is None:
+        for attr_path in [
+            "embedding",
+            "embed_tokens",
+            "token_embedding",
+            "wte",
+            "word_embeddings",
+            "tok_embeddings",
+            "transformer.wte",
+            "model.embed_tokens",
+            "model.embedding",
+        ]:
+            obj = raw
+            for part in attr_path.split("."):
+                obj = getattr(obj, part, None)
+                if obj is None:
+                    break
+            if obj is not None and isinstance(obj, torch.nn.Embedding):
+                embedding = obj
+                break
+
+    if embedding is None:
+        print("[WARN] NEFTune: embedding layer을 찾지 못함, NEFTune 비활성화")
+        return None
+
+    print(
+        f"[INFO] NEFTune: {type(embedding).__name__} hook 등록 "
+        f"(shape={tuple(embedding.weight.shape)}, alpha={noise_alpha})"
+    )
+
+    def _hook(
+        module: torch.nn.Module,
+        inp: tuple,
+        out: torch.Tensor,
+    ) -> torch.Tensor:
+        if module.training:
+            # out shape: [B, seq_len, d_model]
+            mag = noise_alpha / ((out.size(1) * out.size(2)) ** 0.5)
+            out = out + torch.empty_like(out).uniform_(-mag, mag)
+        return out
+
+    return embedding.register_forward_hook(_hook)
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+
+def main() -> None:
+    args = parse_args()
+
+    # ---- Distributed setup -------------------------------------------------
+    is_ddp = "RANK" in os.environ
+    rank = 0
+    local_rank = 0
+    world_size = 1
+
+    if is_ddp:
+        rank, local_rank, world_size, device = setup_ddp()
+    else:
+        # Single-GPU: honour --device flag, else pick cuda:0 or cpu.
+        if args.device is not None:
+            device = torch.device(args.device)
+        elif torch.cuda.is_available():
+            device = torch.device("cuda:0")
+        else:
+            device = torch.device("cpu")
+
+    # Per-rank seed so data shuffling differs across replicas.
+    set_seed(args.seed + rank)
+
+    # ---- NUMA affinity for optimal GPU↔CPU memory locality ---------------
+    # B200 topology: GPU 0-3 → NUMA node 0 (cores 0-35)
+    #                GPU 4-6 → NUMA node 1 (cores 36-71)  [7 GPU 환경]
+    try:
+        if local_rank < 4:
+            os.sched_setaffinity(0, set(range(0, 36)))   # NUMA node 0
+        else:
+            os.sched_setaffinity(0, set(range(36, 72)))   # NUMA node 1
+        if is_main_process():
+            print(f"NUMA affinity: rank {rank} (GPU {local_rank}) → "
+                  f"{'NUMA0 cores 0-35' if local_rank < 4 else 'NUMA1 cores 36-71'}")
+    except (AttributeError, OSError) as e:
+        if is_main_process():
+            print(f"[WARN] NUMA affinity failed: {e}")
+
+    # ---- Validate base checkpoint ------------------------------------------
+    if not args.base_checkpoint.exists():
+        raise FileNotFoundError(
+            f"Base checkpoint directory not found: {args.base_checkpoint}"
+        )
+    for required_file in ("model.pt", "config.yaml"):
+        if not (args.base_checkpoint / required_file).exists():
+            raise FileNotFoundError(
+                f"Expected {required_file} inside base checkpoint: {args.base_checkpoint}"
+            )
+
+    # ---- Load pretrained model ---------------------------------------------
+    # LLM.from_pretrained() reads config.yaml + model.pt and returns the model on CPU.
+    # We move it to the target device immediately after loading.
+    #
+    # NOTE: fp8_model_init() is intentionally NOT used here (same as pretrain.py).
+    # MXFP8Tensor weights are incompatible with DDP's _broadcast_coalesced.
+    # Weights stay in float32; TransformerEngine quantizes on-the-fly inside fp8_autocast.
+    model = LLM.from_pretrained(args.base_checkpoint)
+
+    # FP8 override: --no_fp8 forces BF16 even if pretrained config had use_fp8=True.
+    # --use_fp8 enables FP8 if pretrained config had it disabled.
+    if args.no_fp8:
+        model.config.use_fp8 = False
+    elif args.use_fp8:
+        model.config.use_fp8 = True
+
+    # Move model to target device in bfloat16 (more memory-efficient than fp32
+    # for fine-tuning, and required when BF16 autocast + TE are active).
+    model = model.to(device=device, dtype=torch.bfloat16)
+
+    # ---- Gradient checkpointing ----------------------------------------
+    # Trades activation memory for recomputation during backward pass.
+    # Especially useful for large models / long sequences in SFT.
+    if hasattr(model, 'gradient_checkpointing_enable'):
+        model.gradient_checkpointing_enable()
+        if rank == 0:
+            print("[INFO] Gradient checkpointing enabled")
+
+    # FP8 alignment check: (batch_size × seq_len) must be divisible by 8.
+    if model.config.use_fp8:
+        seq_len = model.config.max_seq_len
+        if (args.batch_size * seq_len) % 8 != 0:
+            raise ValueError(
+                f"FP8: batch_size × max_seq_len = {args.batch_size} × {seq_len} "
+                f"= {args.batch_size * seq_len} must be divisible by 8."
+            )
+
+    if is_main_process():
+        total_params = sum(p.numel() for p in model.parameters())
+        print(f"Pretrained model loaded: {total_params:,} parameters")
+        print(f"LMConfig: {model.config}")
+
+    # ---- Wrap in DDP -------------------------------------------------------
+    if is_ddp:
+        from torch.nn.parallel import DistributedDataParallel as DDP
+
+        model = DDP(
+            model,
+            device_ids=[local_rank],
+            output_device=local_rank,
+            gradient_as_bucket_view=True,
+            bucket_cap_mb=800,
+            find_unused_parameters=False,
+        )
+
+    # ---- Tokenizer ---------------------------------------------------------
+    tokenizer_path = _resolve_tokenizer_path(args)
+    if is_main_process():
+        print(f"Loading tokenizer from: {tokenizer_path}")
+
+    # Use the fast tokenizers library (same as the rest of the project).
+    from tokenizers import Tokenizer  # type: ignore[import]
+    tokenizer = Tokenizer.from_file(str(tokenizer_path))
+
+    # ---- Dataset & DataLoader ----------------------------------------------
+    # Import SFTDataset (created separately alongside this file).
+    # SFTDataset returns (input_ids, targets) where prompt token positions in
+    # targets are filled with -1.  The Trainer._compute_loss already uses
+    # ignore_index=-1, so only response tokens contribute to the gradient.
+    from data.sft_dataset import SFTDataset  # type: ignore[import]
+
+    train_dataset = SFTDataset(
+        data_path=args.sft_data,
+        tokenizer=tokenizer,
+        max_seq_len=model.config.max_seq_len
+        if not isinstance(model, torch.nn.parallel.DistributedDataParallel)
+        else model.module.config.max_seq_len,
+    )
+
+    if is_ddp:
+        train_sampler: DistributedSampler | RandomSampler = DistributedSampler(
+            train_dataset,
+            num_replicas=world_size,
+            rank=rank,
+            shuffle=True,
+            seed=args.seed,
+        )
+        shuffle = False
+    else:
+        train_sampler = RandomSampler(train_dataset)
+        shuffle = False  # Sampler is provided; DataLoader must not also shuffle.
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=args.batch_size,
+        sampler=train_sampler,
+        # SFT datasets are typically small enough that 2–4 workers suffice.
+        # We use 4 to balance I/O with CPU parsing overhead from JSONL.
+        num_workers=args.num_workers,
+        pin_memory=True,
+        drop_last=True,
+        prefetch_factor=2,
+        persistent_workers=True,
+        collate_fn=dynamic_collate_fn,
+    )
+
+    # Optional validation loader.
+    # NOTE: The current Trainer implementation does not yet accept a val_loader
+    # argument; the eval_interval config field is reserved for future use.
+    # We construct the loader here so that once Trainer gains eval support,
+    # wiring it in requires only passing val_loader=val_loader below.
+    val_loader: DataLoader | None = None
+    if args.val_data is not None:
+        if not args.val_data.exists():
+            raise FileNotFoundError(f"Validation data not found: {args.val_data}")
+        val_dataset = SFTDataset(
+            data_path=args.val_data,
+            tokenizer=tokenizer,
+            max_seq_len=train_dataset.max_seq_len,
+        )
+        val_loader = DataLoader(
+            val_dataset,
+            batch_size=args.batch_size,
+            shuffle=False,
+            num_workers=2,
+            pin_memory=True,
+            drop_last=False,
+            collate_fn=dynamic_collate_fn,
+        )
+        if is_main_process():
+            print(f"Validation dataset: {len(val_dataset):,} samples")
+
+    # ---- Optimizer ---------------------------------------------------------
+    # Use the same two-group split (weight_decay / no weight_decay) as pretrain.
+    # Unwrap DDP to get the raw model's parameters.
+    raw_model = getattr(model, "module", model)
+    param_groups = build_optimizer_param_groups(raw_model, args.weight_decay)
+    optimizer = torch.optim.AdamW(
+        param_groups,
+        lr=args.lr,
+        betas=(0.9, 0.95),
+        eps=1e-8,
+        fused=torch.cuda.is_available(),  # Use fused kernel when on CUDA.
+    )
+
+    # ---- TrainConfig -------------------------------------------------------
+    # Set use_fp8 from the (possibly overridden) model config so Trainer builds
+    # the correct FP8 recipe and wraps forward passes in fp8_autocast.
+    use_fp8 = raw_model.config.use_fp8
+
+    train_config = TrainConfig(
+        max_steps=args.max_steps,
+        checkpoint_dir=str(args.checkpoint_dir),
+        grad_accum_steps=args.grad_accum,
+        use_fp8=use_fp8,
+        log_file=str(args.log_file) if args.log_file is not None else None,
+        save_interval=args.save_interval,
+        log_interval=10,
+        eval_interval=args.eval_interval,
+        max_val_batches=args.max_val_batches,
+    )
+
+    # ---- LR Scheduler ------------------------------------------------------
+    scheduler = get_cosine_schedule_with_warmup(
+        optimizer=optimizer,
+        warmup_steps=args.warmup_steps,
+        total_steps=train_config.max_steps,
+    )
+
+    # ---- Resume from SFT checkpoint ----------------------------------------
+    # When --resume is given we restore the SFT optimizer/scheduler state as
+    # well so learning rate, momentum buffers, etc. are correctly restored.
+    # NOTE: This resumes SFT training, NOT the pretrain checkpoint.
+    #       The pretrain weights were already loaded above via from_pretrained().
+    start_step = 0
+    if args.resume is not None:
+        if not args.resume.exists():
+            raise FileNotFoundError(f"Resume checkpoint not found: {args.resume}")
+        start_step, resume_loss = load_checkpoint(
+            path=args.resume,
+            model=model,
+            optimizer=optimizer,
+            scheduler=scheduler,
+        )
+        if is_main_process():
+            print(f"Resumed SFT from {args.resume} at step {start_step} (loss={resume_loss:.4f})")
+
+    if args.resume is not None and isinstance(train_sampler, DistributedSampler):
+        steps_per_epoch = len(train_loader)
+        approx_epoch = start_step // steps_per_epoch if steps_per_epoch > 0 else 0
+        train_sampler.set_epoch(approx_epoch)
+        if is_main_process():
+            print(f"[INFO] Resume: sampler epoch set to {approx_epoch}")
+
+    # ---- Checkpoint directory ----------------------------------------------
+    args.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+
+    # ---- Copy tokenizer to checkpoint dir for easy deployment later --------
+    # This mirrors the tokenizer into the SFT checkpoint root so that the
+    # final checkpoint directory is self-contained for convert_to_hf.py, etc.
+    if is_main_process():
+        dest_tok = args.checkpoint_dir / "tokenizer.json"
+        if not dest_tok.exists():
+            shutil.copy2(str(tokenizer_path), str(dest_tok))
+            print(f"Tokenizer copied to {dest_tok}")
+
+    # ---- Trainer -----------------------------------------------------------
+    trainer = Trainer(
+        model=model,
+        train_loader=train_loader,
+        optimizer=optimizer,
+        scheduler=scheduler,
+        config=train_config,
+        device=device,
+        rank=rank,
+        sampler=train_sampler if is_ddp else None,
+        val_loader=val_loader,
+    )
+
+    # ---- Signal handlers for graceful shutdown ----------------------------
+    import signal as _signal_mod
+
+    _trainer_ref = trainer
+
+    def _graceful_shutdown_handler(signum, frame):
+        sig_name = _signal_mod.Signals(signum).name
+        if is_main_process():
+            import datetime as _dt
+            ts = _dt.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            msg = (
+                f"[{ts}] [SIGNAL] Received {sig_name} (signum={signum}). "
+                f"Initiating graceful shutdown..."
+            )
+            print(f"\n{msg}")
+            if args.log_file is not None:
+                try:
+                    with open(args.log_file, "a", encoding="utf-8") as f:
+                        f.write(msg + "\n")
+                except Exception:
+                    pass
+        _trainer_ref.request_shutdown(sig_name)
+
+    for _sig in (_signal_mod.SIGHUP, _signal_mod.SIGTERM):
+        _signal_mod.signal(_sig, _graceful_shutdown_handler)
+
+    # ---- SFT banner --------------------------------------------------------
+    if is_main_process():
+        import datetime
+
+        inner_config = raw_model.config
+        eff_batch_seqs = args.batch_size * args.grad_accum * world_size
+        eff_tokens_per_step = eff_batch_seqs * inner_config.max_seq_len
+        train_samples = len(train_dataset)
+        precision_label = "FP8 (MXFP8BlockScaling)" if use_fp8 else "BF16"
+        nccl_debug = os.environ.get("NCCL_DEBUG", "not set")
+        omp_threads = os.environ.get("OMP_NUM_THREADS", "not set")
+
+        print(
+            f"\n{'='*70}\n"
+            f"  LLM Supervised Fine-Tuning — "
+            f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
+            f"{'='*70}\n"
+            f"  base ckpt : {args.base_checkpoint}\n"
+            f"  sft data  : {args.sft_data} ({train_samples:,} samples)\n"
+            f"  model     : {inner_config.num_params:,} params  |  "
+            f"d_model={inner_config.d_model}  n_layers={inner_config.n_layers}\n"
+            f"  precision : {precision_label}\n"
+            f"  GPUs      : {world_size}  |  batch/GPU={args.batch_size}  "
+            f"grad_accum={args.grad_accum}\n"
+            f"  eff_batch : {eff_batch_seqs} seqs  "
+            f"= {eff_tokens_per_step:,} tok/step\n"
+            f"  max_steps : {train_config.max_steps:,}\n"
+            f"  lr        : {args.lr:.2e}  "
+            f"warmup={args.warmup_steps}  weight_decay={args.weight_decay}\n"
+            f"  ckpt_dir  : {args.checkpoint_dir}\n"
+            f"  env       : OMP_NUM_THREADS={omp_threads}  NCCL_DEBUG={nccl_debug}\n"
+            f"{'='*70}\n"
+        )
+
+    # ---- NEFTune -----------------------------------------------------------
+    # Add uniform noise to embeddings during training to improve instruction
+    # following (Jain et al., 2023).  Hook is registered on the raw (non-DDP)
+    # model so it survives DDP's internal module wrapping.
+    neftune_alpha = getattr(args, 'neftune_alpha', 5.0)
+    neftune_handle = add_neftune_hook(raw_model, noise_alpha=neftune_alpha)
+    if rank == 0:
+        if neftune_handle is not None:
+            print(f"[INFO] NEFTune enabled (noise_alpha={neftune_alpha})")
+        else:
+            print("[WARN] NEFTune disabled - embedding layer not found")
+
+    # ---- Train -------------------------------------------------------------
+    try:
+        trainer.train(start_step=start_step)
+    except KeyboardInterrupt:
+        if is_main_process():
+            print("\n[INFO] SFT interrupted by user (KeyboardInterrupt).")
+    except Exception as e:
+        import traceback
+        if is_main_process():
+            tb = traceback.format_exc()
+            print(f"\n[ERROR] SFT failed at rank {rank}:\n{tb}")
+            if args.log_file is not None:
+                with open(args.log_file, "a", encoding="utf-8") as f:
+                    import datetime
+                    f.write(
+                        f"[{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] "
+                        f"[FATAL] {tb}\n"
+                    )
+        raise
+    finally:
+        # Remove NEFTune hook so the model is clean for inference/saving.
+        if neftune_handle is not None:
+            neftune_handle.remove()
+        if is_ddp:
+            cleanup_ddp()
+
+
+if __name__ == "__main__":
+    main()

sft-v2/config.json

@@ -0,0 +1,29 @@
+{
+  "vocab_size": 64000,
+  "d_model": 3072,
+  "n_layers": 26,
+  "n_heads": 24,
+  "n_kv_heads": 8,
+  "d_ffn": 9216,
+  "max_seq_len": 4096,
+  "rope_theta": 500000.0,
+  "dropout": 0.0,
+  "bias": false,
+  "use_flash_attn": true,
+  "use_fp8": false,
+  "use_hybrid": true,
+  "hybrid_pattern": "M M M M M M M M M M M M A M M M M M M M M M M M A M",
+  "mamba_d_state": 128,
+  "mamba_head_dim": 64,
+  "mamba_expand": 2,
+  "mamba_conv_kernel": 4,
+  "mamba_n_groups": 8,
+  "mamba_d_ffn": 4608,
+  "mamba_chunk_size": 256,
+  "model_type": "evafrill-mo",
+  "architectures": [
+    "EvafrillMoForCausalLM"
+  ],
+  "_variant": "sft-v2",
+  "_description": "SFT v2 (65K steps, val_loss 1.79, early stop)"
+}

sft-v2/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:144b4a8154e3c5fc13fa8d029d8016fe330fe8dc9083762a7ddbab62103d7073
+size 6301164272

slerp/config.json

@@ -0,0 +1,29 @@
+{
+  "vocab_size": 64000,
+  "d_model": 3072,
+  "n_layers": 26,
+  "n_heads": 24,
+  "n_kv_heads": 8,
+  "d_ffn": 9216,
+  "max_seq_len": 4096,
+  "rope_theta": 500000.0,
+  "dropout": 0.0,
+  "bias": false,
+  "use_flash_attn": true,
+  "use_fp8": false,
+  "use_hybrid": true,
+  "hybrid_pattern": "M M M M M M M M M M M M A M M M M M M M M M M M A M",
+  "mamba_d_state": 128,
+  "mamba_head_dim": 64,
+  "mamba_expand": 2,
+  "mamba_conv_kernel": 4,
+  "mamba_n_groups": 8,
+  "mamba_d_ffn": 4608,
+  "mamba_chunk_size": 256,
+  "model_type": "evafrill-mo",
+  "architectures": [
+    "EvafrillMoForCausalLM"
+  ],
+  "_variant": "slerp",
+  "_description": "SLERP merge alpha=0.5 (RECOMMENDED FINAL MODEL)"
+}

slerp/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b7fedbd0d0f8e33a1fb5e6c4e8e9393f729cc77b364d431e522857ce6a1c8d56
+size 6301164272