distributed training.

dataset/lm_dataset.py

@@ -39,7 +39,16 @@ class PretrainDataset(Dataset):
         super().__init__()
         self.tokenizer = tokenizer
         self.max_length = max_length
-        self.samples = load_dataset('json', data_files=data_path, split='train')
+        # data_path 支持三种形式：
+        #   1. 单文件路径 str          → 单 jsonl
+        #   2. 逗号分隔多路径 str      → 从 CLI 传入的多文件混合（"a.jsonl,b.jsonl"）
+        #   3. list[str]              → API 调用直接传列表
+        # 混合场景：1B pretrain 用 SkyPile 25B（raw web）+ minimind pretrain 1.76B（curated Q&A）
+        if isinstance(data_path, str) and ',' in data_path:
+            data_files = [p.strip() for p in data_path.split(',') if p.strip()]
+        else:
+            data_files = data_path
+        self.samples = load_dataset('json', data_files=data_files, split='train')
 
     def __len__(self):
         return len(self.samples)
@@ -56,12 +65,16 @@ class PretrainDataset(Dataset):
 
 
 class SFTDataset(Dataset):
-    def __init__(self, jsonl_path, tokenizer, max_length=1024):
+    def __init__(self, jsonl_path, tokenizer, max_length=1024, subset_size=0, subset_seed=42):
         super().__init__()
         self.tokenizer = tokenizer
         self.max_length = max_length
         features = Features({'conversations': [{'role': Value('string'), 'content': Value('string'), 'reasoning_content': Value('string'), 'tools': Value('string'), 'tool_calls': Value('string')}]})
         self.samples = load_dataset('json', data_files=jsonl_path, split='train', features=features)
+        # QAT / dry-run 用途：shuffle 后 select，避免取前 N 引入 jsonl 顺序偏置
+        # (activation observer 的 running_max 需要看到代表性输入分布)
+        if subset_size and 0 < subset_size < len(self.samples):
+            self.samples = self.samples.shuffle(seed=subset_seed).select(range(subset_size))
         self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant\n', add_special_tokens=False).input_ids
         self.eos_id = tokenizer(f'{tokenizer.eos_token}\n', add_special_tokens=False).input_ids
 

model/model_minimind.py

@@ -205,6 +205,20 @@ class MiniMindModel(nn.Module):
         freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.head_dim, end=config.max_position_embeddings, rope_base=config.rope_theta, rope_scaling=config.rope_scaling)
         self.register_buffer("freqs_cos", freqs_cos, persistent=False)
         self.register_buffer("freqs_sin", freqs_sin, persistent=False)
+        # gradient checkpointing 标志：True 时 forward 不存 activation，
+        # backward 时重算——省 activation 显存但慢 ~30%。1B + seq=1024 在 80GB 无需，
+        # 但 seq=2048 或 8B+ 模型必须。由 trainer 通过 set_grad_checkpointing() 打开。
+        self.grad_checkpointing = False
+
+    def set_grad_checkpointing(self, enable: bool = True):
+        self.grad_checkpointing = enable
+
+    def warmup_rope(self, device=None):
+        """Meta-device init 下 buffer 可能为空，manual 调用重新 precompute。正常 code path 不需要。"""
+        if self.freqs_cos.numel() == 0:
+            freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.config.head_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling)
+            device = device or self.embed_tokens.weight.device
+            self.freqs_cos, self.freqs_sin = freqs_cos.to(device), freqs_sin.to(device)
 
     def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, **kwargs):
         batch_size, seq_length = input_ids.shape
@@ -212,24 +226,32 @@ class MiniMindModel(nn.Module):
         past_key_values = past_key_values or [None] * len(self.layers)
         start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
         hidden_states = self.dropout(self.embed_tokens(input_ids))
-        # Lazy recompute of RoPE buffers if they were dropped by transformers>=5.x meta-device init.
+        # RoPE buffers 已在 __init__ 里 precompute + register_buffer；forward 直接切片即可。
-        # Hot path: skipped via cheap python bool (no GPU sync). The original `freqs_cos[0,0]==0`
+        # (原来这里有 lazy re-compute 兜底 transformers>=5.x meta-device init 掉 buffer 的情况，
-        # tensor compare forced a device-to-host sync on every forward.
+        # 但那段 `bool(freqs_cos[0,0].eq(0).cpu())` 是 torch.compile graph break——每步 forward
-        if getattr(self, '_rope_ready', False) is False:
+        # 都要在 compiled ↔ eager 间切换 + guard 检查，短 step 下吃掉一半吞吐。如需 meta-init
-            if self.freqs_cos.numel() == 0 or bool(self.freqs_cos[0, 0].eq(0).cpu()):
+        # 兼容，在 model.to(device) 后手动调 warmup_rope() 即可，不要放 forward hot path。)
-                freqs_cos, freqs_sin = precompute_freqs_cis(dim=self.config.head_dim, end=self.config.max_position_embeddings, rope_base=self.config.rope_theta, rope_scaling=self.config.rope_scaling)
-                self.freqs_cos, self.freqs_sin = freqs_cos.to(hidden_states.device), freqs_sin.to(hidden_states.device)
-            self._rope_ready = True
         position_embeddings = (self.freqs_cos[start_pos:start_pos + seq_length], self.freqs_sin[start_pos:start_pos + seq_length])
         presents = []
+        # gradient checkpointing 与 use_cache=True 不兼容（重算时无法访问 KV cache），
+        # 训练场景一般 use_cache=False，此处兜底 guard 提示。
+        use_ckpt = self.grad_checkpointing and self.training and not use_cache
         for layer, past_key_value in zip(self.layers, past_key_values):
-            hidden_states, present = layer(
+            if use_ckpt:
-                hidden_states,
+                # use_reentrant=False 是 torch>=2.1 推荐（旧 reentrant 版本对 DDP/compile 有问题）
-                position_embeddings,
+                hidden_states, present = torch.utils.checkpoint.checkpoint(
-                past_key_value=past_key_value,
+                    layer, hidden_states, position_embeddings,
-                use_cache=use_cache,
+                    past_key_value, use_cache, attention_mask,
-                attention_mask=attention_mask
+                    use_reentrant=False,
-            )
+                )
+            else:
+                hidden_states, present = layer(
+                    hidden_states,
+                    position_embeddings,
+                    past_key_value=past_key_value,
+                    use_cache=use_cache,
+                    attention_mask=attention_mask
+                )
             presents.append(present)
         hidden_states = self.norm(hidden_states)
         aux_loss = sum([l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward)], hidden_states.new_zeros(1).squeeze())

scripts/eval_perplexity.py

@@ -38,6 +38,11 @@ from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
 warnings.filterwarnings('ignore')
 
 
+def _is_qat_state_dict(sd):
+    """QAT ckpt 会带 act_fq.running_max / act_fq.step buffer——用它判定要不要先 wrap。"""
+    return any('act_fq.running_max' in k for k in sd.keys())
+
+
 def load_model(args):
     cfg = MiniMindConfig(
         hidden_size=args.hidden_size,
@@ -48,10 +53,71 @@ def load_model(args):
     moe_suffix = '_moe' if args.use_moe else ''
     ckp = f'{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
     state_dict = torch.load(ckp, map_location='cpu', weights_only=True)
-    # strict=False — tied embeddings have duplicate keys; load handles via shared storage
+
-    model.load_state_dict(state_dict, strict=False)
+    is_qat_ckpt = _is_qat_state_dict(state_dict)
-    model = model.half().eval().to(args.device)
+    apply_qat = bool(args.apply_qat) or is_qat_ckpt
-    return model, cfg
+
+    if apply_qat:
+        from model.quant import prepare_qat, bake_quantized_weights
+        skip_patterns = tuple(s.strip() for s in args.quant_skip.split(',') if s.strip())
+        if is_qat_ckpt:
+            # QAT ckpt 已含 buffer——先 wrap 出对应 module，再整体 load
+            replaced = prepare_qat(model, w_bits=args.w_bits, a_bits=args.a_bits,
+                                   skip_patterns=skip_patterns)
+            model.load_state_dict(state_dict, strict=False)
+            print(f'[QAT] loaded QAT ckpt, wrapped {replaced} Linear modules')
+        else:
+            # fp ckpt——先 load 权重，再 wrap（QATLinear.from_float 保留原 weight tensor）
+            model.load_state_dict(state_dict, strict=False)
+            replaced = prepare_qat(model, w_bits=args.w_bits, a_bits=args.a_bits,
+                                   skip_patterns=skip_patterns)
+            print(f'[QAT] wrapped fp ckpt with {replaced} QATLinear (observer un-calibrated)')
+
+        model = model.half().to(args.device)
+
+        if not is_qat_ckpt and args.qat_calibrate > 0:
+            # 对未训练过的 QAT 模型：observer running_max 全 0，直接 eval 会 scale ~= 0 全饱和
+            # 跑一段 train() 模式的 forward 让 EMA 收敛，再切 eval
+            _calibrate_observer(model, args)
+
+        if args.qat_bake:
+            print('[QAT] baking fake-quantized weights in-place (simulates deployed int8)')
+            bake_quantized_weights(model)
+    else:
+        # strict=False — tied embeddings have duplicate keys; load handles via shared storage
+        model.load_state_dict(state_dict, strict=False)
+        model = model.half().to(args.device)
+
+    return model.eval(), cfg
+
+
+@torch.no_grad()
+def _calibrate_observer(model, args):
+    """在 train() 模式下跑 N 个 forward 把 ActFakeQuant.running_max EMA 跑起来。
+
+    ActFakeQuant.forward 里 self.training 分支才会更新 EMA + 记 step；observer_steps
+    warmup 期内也不做 fake-quant，只统计。这里跑够 observer_steps 步保证之后 eval
+    切进真正 fake-quant 分支。
+    """
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
+    bos_id, eos_id = tokenizer.bos_token_id, tokenizer.eos_token_id
+    model.train()
+    fed = 0
+    for text in iter_holdout(args):
+        if not text or len(text) < 10:
+            continue
+        body = tokenizer(text, max_length=args.max_seq_len - 2,
+                         truncation=True, add_special_tokens=False).input_ids
+        tokens = [bos_id] + body + [eos_id]
+        if len(tokens) < 2:
+            continue
+        ids = torch.tensor([tokens], device=args.device)
+        model(input_ids=ids)
+        fed += 1
+        if fed >= args.qat_calibrate:
+            break
+    model.eval()
+    print(f'[QAT] calibration done ({fed} forwards; observer_steps warmup = {args.qat_calibrate})')
 
 
 def iter_holdout(args):
@@ -91,6 +157,18 @@ def main():
     parser.add_argument('--max_samples', type=int, default=500)
     parser.add_argument('--max_seq_len', type=int, default=512)
     parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu')
+
+    # ---- QAT eval options ----
+    parser.add_argument('--apply_qat', type=int, default=0, choices=[0, 1],
+                        help='对 fp ckpt 也用 QATLinear 包一层，测纯量化噪声；QAT ckpt 会自动识别')
+    parser.add_argument('--w_bits', type=int, default=8)
+    parser.add_argument('--a_bits', type=int, default=8)
+    parser.add_argument('--qat_calibrate', type=int, default=200,
+                        help='fp ckpt + apply_qat 时的 observer 预热步数；QAT ckpt 时忽略')
+    parser.add_argument('--qat_bake', type=int, default=0, choices=[0, 1],
+                        help='把 fake-quantized 权重 in-place 写回 weight（模拟真实 int8 部署）')
+    parser.add_argument('--quant_skip', type=str, default='lm_head,mlp.gate',
+                        help='要跳过量化的 Linear 名（逗号分隔），需和训练一致')
     args = parser.parse_args()
 
     print(f'[load] ckpt={args.save_dir}/{args.weight}_{args.hidden_size}.pth')

scripts/measure_sqnr.py

@@ -0,0 +1,292 @@
+"""
+Measure SQNR (Signal-to-Quantization-Noise Ratio) of MiniMind checkpoints.
+
+三个维度：
+1. Weight SQNR: 对每层 Linear.weight 应用 W8/W4 fake-quant，算 SQNR
+   - Pretrain vs QAT-trained: QAT 训练应把 W SQNR 推高（权重朝量化友好区漂）
+2. Activation SQNR: 用 forward hook 抓每层输入 activation, 用 running_max 量化后算 SQNR
+   - 需要真实数据（否则激活分布不代表推理时）
+3. Bit-width sweep: 对 pretrain 权重扫 W2/W4/W8/W16 看 SQNR ~ bits 曲线（应该 +6dB/bit）
+
+Usage:
+    cd scripts && python measure_sqnr.py                     # 默认：W SQNR (pretrain vs qat) + bit sweep
+    python measure_sqnr.py --act 1                           # 加上 activation SQNR (需要 GPU + 数据)
+    python measure_sqnr.py --act 1 --n_samples 32            # activation 用 N samples 校准
+"""
+import os
+import sys
+__package__ = "scripts"
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+import argparse
+import warnings
+import json
+import torch
+from transformers import AutoTokenizer
+from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
+from model.quant import WeightFakeQuant, prepare_qat, DEFAULT_SKIP_PATTERNS
+
+warnings.filterwarnings('ignore')
+
+
+def sqnr_db(x: torch.Tensor, x_q: torch.Tensor) -> float:
+    """SQNR = 10 log10(||x||² / ||x - x_q||²)，用 fp64 算防止 fp16 数值下溢。"""
+    x, x_q = x.double(), x_q.double()
+    signal_pow = (x ** 2).mean()
+    noise_pow = ((x - x_q) ** 2).mean().clamp_min(1e-24)
+    return (10.0 * torch.log10(signal_pow / noise_pow)).item()
+
+
+def is_quantizable_linear_key(k: str, skip_patterns=DEFAULT_SKIP_PATTERNS) -> bool:
+    """判断 state_dict 的 key 是否对应我们会量化的 Linear.weight。"""
+    if not k.endswith('.weight'):
+        return False
+    # 排除 embedding / norm / lm_head / MoE 路由（和 prepare_qat 的 skip 一致）
+    if 'embed_tokens' in k or 'lm_head' in k or 'norm' in k:
+        return False
+    parent = k[:-len('.weight')]
+    for pat in skip_patterns:
+        if parent == pat or parent.endswith('.' + pat):
+            return False
+    return True
+
+
+def measure_weight_sqnr(pretrain_sd, qat_sd, bits=8):
+    """对每层 Linear.weight 算 W SQNR，返回 dict {layer_name: (sqnr_pretrain, sqnr_qat, delta)}。"""
+    wfq = WeightFakeQuant(bits=bits)
+    results = {}
+    for k in pretrain_sd.keys():
+        if not is_quantizable_linear_key(k):
+            continue
+        w_pt = pretrain_sd[k].float()
+        w_pt_q = wfq(w_pt)
+        sqnr_pt = sqnr_db(w_pt, w_pt_q)
+        rec = {'pretrain': sqnr_pt}
+        if k in qat_sd:
+            w_q = qat_sd[k].float()
+            w_q_q = wfq(w_q)
+            sqnr_q = sqnr_db(w_q, w_q_q)
+            rec['qat'] = sqnr_q
+            rec['delta'] = sqnr_q - sqnr_pt
+        results[k] = rec
+    return results
+
+
+def measure_bit_sweep(pretrain_sd, bit_list=(2, 3, 4, 6, 8, 12)):
+    """对每层权重扫多个 bit，验证 ~6 dB/bit 规律。返回 dict {bits: mean_sqnr_across_layers}。"""
+    out = {}
+    for bits in bit_list:
+        wfq = WeightFakeQuant(bits=bits)
+        sqnrs = []
+        for k in pretrain_sd.keys():
+            if not is_quantizable_linear_key(k):
+                continue
+            w = pretrain_sd[k].float()
+            sqnrs.append(sqnr_db(w, wfq(w)))
+        out[bits] = sum(sqnrs) / len(sqnrs)
+    return out
+
+
+@torch.no_grad()
+def measure_activation_sqnr(cfg, qat_ckpt_path, tokenizer_path, data_path, n_samples, device='cuda'):
+    """
+    跑真实数据，用 forward hook 抓每个 QATLinear 的输入激活，
+    对比 fp 激活 vs act_fq(fp激活) 的 SQNR。使用 QAT ckpt 里训好的 running_max。
+    """
+    model = MiniMindForCausalLM(cfg)
+    prepare_qat(model, w_bits=8, a_bits=8)
+    sd = torch.load(qat_ckpt_path, map_location='cpu', weights_only=True)
+    model.load_state_dict(sd, strict=False)
+    model = model.half().eval().to(device)
+
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+    bos_id, eos_id = tokenizer.bos_token_id, tokenizer.eos_token_id
+
+    # 收集每个 QATLinear 的激活 SQNR（累积 sum + count 便于跨样本求平均）
+    from model.quant import QATLinear
+    stats = {}  # name -> {sum_sig_pow, sum_noise_pow, count}
+    hooks = []
+
+    def make_hook(name, act_fq):
+        def hook(mod, inp, out):
+            x = inp[0].detach()
+            x_q = act_fq(x.clone())  # 现用 QAT 训好的 running_max 量化
+            xd, xqd = x.double(), x_q.double()
+            s = (xd ** 2).sum().item()
+            n = ((xd - xqd) ** 2).sum().item()
+            if name not in stats:
+                stats[name] = {'sig': 0.0, 'noise': 0.0, 'nel': 0}
+            stats[name]['sig'] += s
+            stats[name]['noise'] += n
+            stats[name]['nel'] += x.numel()
+        return hook
+
+    for name, m in model.named_modules():
+        if isinstance(m, QATLinear):
+            hooks.append(m.register_forward_hook(make_hook(name, m.act_fq)))
+
+    # 走 N 条 pretrain 数据
+    fed = 0
+    with open(data_path) as f:
+        for line in f:
+            d = json.loads(line)
+            text = d.get('text') or d.get('content') or ''
+            if not text or len(text) < 10:
+                continue
+            body = tokenizer(text, max_length=510, truncation=True,
+                             add_special_tokens=False).input_ids
+            tokens = [bos_id] + body + [eos_id]
+            ids = torch.tensor([tokens], device=device)
+            model(input_ids=ids)
+            fed += 1
+            if fed >= n_samples:
+                break
+
+    for h in hooks:
+        h.remove()
+
+    # 转成 SQNR
+    result = {}
+    for name, s in stats.items():
+        sqnr = 10.0 * torch.log10(torch.tensor(s['sig'] / max(s['noise'], 1e-24))).item()
+        result[name] = sqnr
+    return result, fed
+
+
+def _fmt_layer_stats(results, val_key='pretrain'):
+    """紧凑打印每层 SQNR（截 layer name 到 self_attn.q_proj 之类）。"""
+    values = [(k.replace('model.layers.', 'L').replace('.weight', ''), v)
+              for k, v in results.items()]
+    return values
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--hidden_size', type=int, default=768)
+    parser.add_argument('--num_hidden_layers', type=int, default=8)
+    parser.add_argument('--use_moe', type=int, default=0)
+    parser.add_argument('--save_dir', type=str, default='../out')
+    parser.add_argument('--pretrain', type=str, default='pretrain')
+    parser.add_argument('--qat', type=str, default='qat')
+    parser.add_argument('--tokenizer_path', type=str, default='../model')
+    parser.add_argument('--data_path', type=str, default='../dataset/pretrain_t2t.jsonl')
+    parser.add_argument('--n_samples', type=int, default=32)
+    parser.add_argument('--act', type=int, default=1, choices=[0, 1],
+                        help='是否算激活 SQNR（需要 GPU + 数据）')
+    parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu')
+    args = parser.parse_args()
+
+    moe_suffix = '_moe' if args.use_moe else ''
+    pt_path = f'{args.save_dir}/{args.pretrain}_{args.hidden_size}{moe_suffix}.pth'
+    qat_path = f'{args.save_dir}/{args.qat}_{args.hidden_size}{moe_suffix}.pth'
+
+    print(f'[load] pretrain: {pt_path}')
+    print(f'[load] qat:      {qat_path}')
+    pretrain_sd = torch.load(pt_path, map_location='cpu', weights_only=True)
+    qat_sd = torch.load(qat_path, map_location='cpu', weights_only=True)
+
+    # =========================================================
+    # 1. Weight SQNR: pretrain vs QAT-trained, W8
+    # =========================================================
+    print(f'\n{"=" * 70}')
+    print(f'  Weight SQNR (W8 per-out-channel symmetric)')
+    print(f'  higher = weights sit closer to int8 grid points')
+    print(f'{"=" * 70}')
+    w8_results = measure_weight_sqnr(pretrain_sd, qat_sd, bits=8)
+
+    # 分组：attn (q/k/v/o) vs ffn (gate/up/down) 分别汇总
+    def group(k):
+        if 'self_attn' in k: return 'attn'
+        if 'mlp' in k: return 'ffn'
+        return 'other'
+
+    groups = {'attn': [], 'ffn': [], 'other': []}
+    for k, r in w8_results.items():
+        groups[group(k)].append(r)
+
+    print(f"{'Group':10s} {'#layers':>8s} {'PT SQNR mean':>14s} {'QAT SQNR mean':>14s} {'Δ mean':>10s}")
+    for g, rs in groups.items():
+        if not rs: continue
+        pt_mean = sum(r['pretrain'] for r in rs) / len(rs)
+        qat_mean = sum(r.get('qat', 0) for r in rs) / len(rs)
+        delta_mean = sum(r.get('delta', 0) for r in rs) / len(rs)
+        print(f"{g:10s} {len(rs):>8d} {pt_mean:>13.2f}dB {qat_mean:>13.2f}dB {delta_mean:>+9.2f}dB")
+
+    print(f'\n  Per-layer detail (first 8 layers, all types):')
+    print(f'  {"layer":50s} {"PT dB":>8s} {"QAT dB":>8s} {"Δ":>8s}')
+    shown = 0
+    for k, r in w8_results.items():
+        if 'layers.0.' not in k: continue  # 只看第 0 层
+        short = k.replace('model.layers.0.', '').replace('.weight', '')
+        print(f'  {short:50s} {r["pretrain"]:>7.2f}  {r.get("qat", 0):>7.2f}  {r.get("delta", 0):>+7.2f}')
+        shown += 1
+
+    # 排序找最难 / 最容易量化的 3 层
+    ranked_pt = sorted(w8_results.items(), key=lambda kv: kv[1]['pretrain'])
+    print(f'\n  Hardest to quantize (lowest PT SQNR):')
+    for k, r in ranked_pt[:3]:
+        print(f'    {k:60s}  {r["pretrain"]:.2f} dB')
+    print(f'  Easiest to quantize (highest PT SQNR):')
+    for k, r in ranked_pt[-3:]:
+        print(f'    {k:60s}  {r["pretrain"]:.2f} dB')
+
+    # =========================================================
+    # 2. Bit-width sweep（验证 ~6 dB/bit）
+    # =========================================================
+    print(f'\n{"=" * 70}')
+    print(f'  Bit-width sweep on pretrain weights (avg across {len(w8_results)} Linears)')
+    print(f'  theory: SQNR ≈ 6.02 × bits + 1.76 dB (uniform quantizer, full-scale signal)')
+    print(f'{"=" * 70}')
+    sweep = measure_bit_sweep(pretrain_sd, bit_list=(2, 3, 4, 6, 8, 12))
+    prev = None
+    print(f"  {'bits':>5s} {'measured':>12s} {'theory (6b+1.76)':>18s} {'Δ vs prev':>12s}")
+    for bits, s in sweep.items():
+        theory = 6.02 * bits + 1.76
+        delta = f"{s - prev:+.2f} dB" if prev is not None else "—"
+        print(f"  {bits:>5d} {s:>10.2f}dB {theory:>15.2f}dB {delta:>12s}")
+        prev = s
+
+    # =========================================================
+    # 3. Activation SQNR（可选，需要 GPU + 数据）
+    # =========================================================
+    if args.act:
+        print(f'\n{"=" * 70}')
+        print(f'  Activation SQNR (per-tensor symmetric int8, running_max from QAT training)')
+        print(f'  用 {args.n_samples} 条 pretrain 样本走前向，hook 到每个 QATLinear 输入')
+        print(f'{"=" * 70}')
+        cfg = MiniMindConfig(hidden_size=args.hidden_size,
+                             num_hidden_layers=args.num_hidden_layers,
+                             use_moe=bool(args.use_moe))
+        act_sqnr, n_run = measure_activation_sqnr(
+            cfg, qat_path, args.tokenizer_path,
+            args.data_path, args.n_samples, args.device,
+        )
+        print(f'  (evaluated on {n_run} samples)')
+        # 按 attn/ffn 分组汇总
+        act_groups = {'attn': [], 'ffn': []}
+        for name, v in act_sqnr.items():
+            if 'self_attn' in name: act_groups['attn'].append(v)
+            elif 'mlp' in name: act_groups['ffn'].append(v)
+        print(f"  {'Group':10s} {'#Linear':>10s} {'A8 SQNR mean':>15s} {'min':>10s} {'max':>10s}")
+        for g, vs in act_groups.items():
+            if not vs: continue
+            print(f"  {g:10s} {len(vs):>10d} {sum(vs)/len(vs):>13.2f}dB {min(vs):>8.2f}dB {max(vs):>8.2f}dB")
+        # 层 0 明细
+        print(f'\n  Layer 0 per-Linear activation SQNR:')
+        for name, v in act_sqnr.items():
+            if 'layers.0.' not in name: continue
+            short = name.replace('model.layers.0.', '')
+            print(f'    {short:40s}  {v:.2f} dB')
+        # 全局最差 3 个
+        ranked_act = sorted(act_sqnr.items(), key=lambda kv: kv[1])
+        print(f'\n  Worst 3 activation SQNR (量化最难扛的层):')
+        for k, v in ranked_act[:3]:
+            print(f'    {k:60s}  {v:.2f} dB')
+
+    print(f'\n{"=" * 70}')
+    print(f'  参照：W8 通常 40-50 dB / W4 15-25 dB / <20 dB 开始明显掉 PPL')
+    print(f'{"=" * 70}')
+
+
+if __name__ == "__main__":
+    main()

trainer/train_full_sft.py

@@ -30,6 +30,10 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
     last_log_time, last_log_step = start_time, start_step
     last_step = start_step
     for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
+        # torch.compile(mode="max-autotune") + CUDA graph 每步必须 mark step 边界，
+        # 否则 backward 拿到被下一步 forward 覆盖的 tensor 报错。见 train_pretrain.py 同处注释。
+        if args.use_compile == 1:
+            torch.compiler.cudagraph_mark_step_begin()
         # non_blocking 仅在 --enable_optimizations=1 + pin_memory=True 时有效
         input_ids = input_ids.to(args.device, non_blocking=args.enable_optimizations == 1)
         labels = labels.to(args.device, non_blocking=args.enable_optimizations == 1)
@@ -71,18 +75,23 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
             Logger(f'[{time.strftime("%H:%M:%S")}] Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), it/s: {window_its:.2f}, MFU: {mfu:.1f}%, loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
             if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min, "it/s": window_its, "MFU": mfu})
 
-        if (step % args.save_interval == 0 or step == iters) and is_main_process():
+        if step % args.save_interval == 0 or step == iters:
-            model.eval()
+            if is_main_process():
-            moe_suffix = '_moe' if lm_config.use_moe else ''
+                model.eval()
-            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+                moe_suffix = '_moe' if lm_config.use_moe else ''
-            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+                ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
-            raw_model = getattr(raw_model, '_orig_mod', raw_model)
+                raw_model = model.module if isinstance(model, DistributedDataParallel) else model
-            state_dict = raw_model.state_dict()
+                raw_model = getattr(raw_model, '_orig_mod', raw_model)
-            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
+                state_dict = raw_model.state_dict()
-            lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, 
+                ckp_tmp = ckp + '.tmp'
-                         epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler)
+                torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp_tmp)
-            model.train()
+                os.replace(ckp_tmp, ckp)
-            del state_dict
+                lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer,
+                             epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler)
+                model.train()
+                del state_dict
+            if dist.is_initialized():
+                dist.barrier()
 
         del input_ids, labels, res, loss
 

trainer/train_pretrain.py

@@ -6,6 +6,7 @@ sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
 import datasets  # noqa: F401  # Windows pyarrow/torch DLL conflict workaround (issue #771)
 import argparse
+import math
 import time
 import warnings
 import torch
@@ -25,16 +26,47 @@ from trainer.trainer_utils import (
 warnings.filterwarnings('ignore')
 
 
-def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
+@torch.no_grad()
+def eval_holdout(model, eval_loader, max_batches, autocast_ctx, device):
+    """
+    rank-0 only 跑 holdout PPL；调用方需保证只 rank-0 进入。
+    返回 (avg_loss, ppl)。max_batches 控制 eval 时长（1B 时 20-50 batch 足够，~1s）。
+    """
+    was_training = model.training
+    model.eval()
+    total_loss, total_tokens = 0.0, 0
+    for i, (input_ids, labels) in enumerate(eval_loader):
+        if i >= max_batches: break
+        input_ids = input_ids.to(device, non_blocking=True)
+        labels = labels.to(device, non_blocking=True)
+        with autocast_ctx:
+            res = model(input_ids, labels=labels)
+        # loss 已经是 mean over non-pad tokens；乘 n_tokens 便于跨 batch weighted avg
+        n_tokens = (labels != -100).sum().item()
+        if n_tokens > 0:
+            total_loss += res.loss.item() * n_tokens
+            total_tokens += n_tokens
+    if was_training: model.train()
+    avg_loss = total_loss / max(total_tokens, 1)
+    return avg_loss, math.exp(avg_loss)
+
+
+def train_epoch(epoch, loader, iters, start_step=0, wandb=None, eval_loader=None):
     start_time = time.time()
     last_log_time, last_log_step = start_time, start_step
     last_step = start_step
     for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
+        # torch.compile(mode="max-autotune") 会启用 CUDA graph capture 复用输出 buffer；
+        # 训练循环里每步都是新 batch，必须显式 mark step 边界告诉 CUDAGraph "上一步的输出
+        # 现在不能再复用了"。否则 backward 拿到被覆盖的 tensor 报 "accessing tensor output
+        # of CUDAGraphs that has been overwritten" 错。torch 2.6+ 严格执行此检查。
+        if args.use_compile == 1:
+            torch.compiler.cudagraph_mark_step_begin()
         # non_blocking=True overlaps host->device copy with compute (pin_memory=True 已在 DataLoader)
         input_ids = input_ids.to(args.device, non_blocking=True)
         labels = labels.to(args.device, non_blocking=True)
         last_step = step
-        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
+        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate, args.warmup_steps)
         for param_group in optimizer.param_groups:
             param_group['lr'] = lr
 
@@ -73,17 +105,34 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
             Logger(f'[{time.strftime("%H:%M:%S")}] Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), it/s: {window_its:.2f}, MFU: {mfu:.1f}%, loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, epoch_time: {eta_min:.1f}min')
             if wandb: wandb.log({"loss": current_loss, "logits_loss": current_logits_loss, "aux_loss": current_aux_loss, "learning_rate": current_lr, "epoch_time": eta_min})
 
-        if (step % args.save_interval == 0 or step == iters) and is_main_process():
+        # ---- eval loop（可选） ----
-            model.eval()
+        if args.eval_interval > 0 and step % args.eval_interval == 0 and eval_loader is not None:
-            moe_suffix = '_moe' if lm_config.use_moe else ''
+            if is_main_process():
-            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+                val_loss, val_ppl = eval_holdout(model, eval_loader, args.eval_max_batches, autocast_ctx, args.device)
-            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+                Logger(f'[eval] step={step} val_loss={val_loss:.4f} val_ppl={val_ppl:.2f}')
-            raw_model = getattr(raw_model, '_orig_mod', raw_model)
+                if wandb: wandb.log({"val_loss": val_loss, "val_ppl": val_ppl, "step": step})
-            state_dict = raw_model.state_dict()
+            if dist.is_initialized(): dist.barrier()
-            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
+
-            lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints')
+        if step % args.save_interval == 0 or step == iters:
-            model.train()
+            if is_main_process():
-            del state_dict
+                model.eval()
+                moe_suffix = '_moe' if lm_config.use_moe else ''
+                ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+                raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+                raw_model = getattr(raw_model, '_orig_mod', raw_model)
+                state_dict = raw_model.state_dict()
+                # atomic: tmp + rename 避免读者读到半写文件（当前的 torch.save 直接写不 atomic，
+                # 极端情况 kill -9 会留 corrupt ckpt；lm_checkpoint 内部已 atomic，此处也补上）
+                ckp_tmp = ckp + '.tmp'
+                torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp_tmp)
+                os.replace(ckp_tmp, ckp)
+                lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer, scaler=scaler, epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints')
+                model.train()
+                del state_dict
+            # 所有 rank 在此汇合，让非 rank-0 等 rank-0 写完再进下一 step
+            # （防止 NCCL 隐式同步下的 timing race，也确保 ckpt 文件对 resume 一致）
+            if dist.is_initialized():
+                dist.barrier()
 
         del input_ids, labels, res, loss
 
@@ -113,7 +162,7 @@ if __name__ == "__main__":
     parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
     parser.add_argument('--max_seq_len', default=340, type=int, help="训练的最大截断长度（中文1token≈1.5~1.7字符）")
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
-    parser.add_argument("--data_path", type=str, default="../dataset/pretrain_t2t_mini.jsonl", help="预训练数据路径")
+    parser.add_argument("--data_path", type=str, default="../dataset/pretrain_t2t.jsonl", help="预训练数据路径")
     parser.add_argument('--from_weight', default='none', type=str, help="基于哪个权重训练，为none则从头开始")
     parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
     parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
@@ -121,6 +170,17 @@ if __name__ == "__main__":
     parser.add_argument("--use_compile", default=1, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是，A100 上推荐默认 1）")
     parser.add_argument("--gpu_peak_tflops", type=float, default=0,
                         help="GPU bf16 理论峰值 TFLOPS，用于 MFU 计算；默认 0 表示按 GPU 名自动检测（A100=312）")
+    parser.add_argument("--warmup_steps", type=int, default=0,
+                        help="linear warmup 步数。1B from-scratch 建议 500-2000；64M 用现有 pretrain ckpt 微调设 0 即可")
+    parser.add_argument("--grad_checkpointing", type=int, default=0, choices=[0, 1],
+                        help="激活值 gradient checkpointing：省 30-50%% activation 显存，慢 ~30%%。"
+                             "1B + seq=1024 在 80GB 用不上；seq=2048 或 8B+ 或大 batch 时需要打开")
+    parser.add_argument("--eval_interval", type=int, default=0,
+                        help="每 N step 跑一次 holdout eval（0=关闭）。1B 建议 500-1000")
+    parser.add_argument("--eval_data_path", type=str, default=None,
+                        help="holdout jsonl 路径；不设则关闭 eval")
+    parser.add_argument("--eval_max_batches", type=int, default=20,
+                        help="每次 eval 最多跑几个 batch（20*batch_size 通常够 ~1s 内出数）")
     args = parser.parse_args()
 
     # ========== 1. 初始化环境和随机种子 ==========
@@ -149,6 +209,9 @@ if __name__ == "__main__":
 
     # ========== 5. 定义模型、数据、优化器 ==========
     model, tokenizer = init_model(lm_config, args.from_weight, device=args.device)
+    if args.grad_checkpointing == 1:
+        model.model.set_grad_checkpointing(True)
+        Logger('gradient checkpointing enabled (activation 省 30-50%%, 慢 ~30%%)')
     train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
     train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
     scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
@@ -173,13 +236,26 @@ if __name__ == "__main__":
 
     # ========== 7. 编译和分布式包装 ==========
     if args.use_compile == 1:
-        # max-autotune 在 Triton 层为每个 fused kernel 试多套 tile/block 配置，挑最快的。
+        # 新 CUDA 13 stack 上 max-autotune 会为每个 GEMM shape 试 21 种 triton kernel（分钟级
-        # 编译时间多 ~30-60s，稳态可比 default 快 5-15%。
+        # autotune 循环，且实测 cuBLAS aten mm 几乎每次都赢过 triton 变体，autotune 白花时间）。
-        model = torch.compile(model, mode="max-autotune")
+        # 用 default 更实惠——inductor 融合仍在，只是不额外扫 triton kernel。
-        Logger('torch.compile enabled (mode=max-autotune)')
+        # (max-autotune 还需要 model 无 graph break，否则 CUDA graph 跨 step 覆盖崩溃；
+        # 我们已消除 RoPE graph break，若将来单独 GEMM shape 少可再试 max-autotune)
+        model = torch.compile(model)
+        Logger('torch.compile enabled (mode=default)')
     if dist.is_initialized():
         model = DistributedDataParallel(model, device_ids=[local_rank])
 
+    # ========== 7b. eval loader（可选） ==========
+    eval_loader = None
+    if args.eval_interval > 0 and args.eval_data_path and is_main_process():
+        # rank-0 only：只 main rank 需要 eval loader（其他 rank barrier 等待）
+        eval_ds = PretrainDataset(args.eval_data_path, tokenizer, max_length=args.max_seq_len)
+        eval_loader = DataLoader(eval_ds, batch_size=args.batch_size,
+                                 shuffle=False, num_workers=2, pin_memory=True)
+        Logger(f'[eval] loader ready: {len(eval_ds)} samples from {args.eval_data_path}, '
+               f'eval_interval={args.eval_interval}, max_batches={args.eval_max_batches}')
+
     # ========== 8. 开始训练 ==========
     for epoch in range(start_epoch, args.epochs):
         train_sampler and train_sampler.set_epoch(epoch)
@@ -190,9 +266,9 @@ if __name__ == "__main__":
                             pin_memory=True, persistent_workers=args.num_workers > 0)
         if skip > 0:
             Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
-            train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
+            train_epoch(epoch, loader, len(loader) + skip, start_step, wandb, eval_loader)
         else:
-            train_epoch(epoch, loader, len(loader), 0, wandb)
+            train_epoch(epoch, loader, len(loader), 0, wandb, eval_loader)
 
     # ========== 9. 清理分布进程 ==========
     if dist.is_initialized():

trainer/train_qat.py

@@ -20,6 +20,7 @@ from dataset.lm_dataset import SFTDataset
 from trainer.trainer_utils import (
     get_lr, Logger, is_main_process, lm_checkpoint,
     init_distributed_mode, setup_seed, init_model, SkipBatchSampler,
+    detect_gpu_peak_tflops_bf16, compute_model_flops_per_token,
 )
 
 warnings.filterwarnings('ignore')
@@ -27,10 +28,15 @@ warnings.filterwarnings('ignore')
 
 def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
     start_time = time.time()
+    last_log_time, last_log_step = start_time, start_step
     last_step = start_step
     for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
-        input_ids = input_ids.to(args.device)
+        # torch.compile + CUDA graph 每步必须 mark step 边界，见 train_pretrain.py 同处注释。
-        labels = labels.to(args.device)
+        # QAT 默认 --use_compile 0（fake-quant 阻断 fusion，即使开也少受益），此处保护性写入。
+        if args.use_compile == 1:
+            torch.compiler.cudagraph_mark_step_begin()
+        input_ids = input_ids.to(args.device, non_blocking=True)
+        labels = labels.to(args.device, non_blocking=True)
         last_step = step
         lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
         for param_group in optimizer.param_groups:
@@ -51,14 +57,26 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
             optimizer.zero_grad(set_to_none=True)
 
         if step % args.log_interval == 0 or step == iters:
-            spend_time = time.time() - start_time
+            now = time.time()
+            spend_time = now - start_time
+            window_dt = max(now - last_log_time, 1e-6)
+            window_steps = step - last_log_step
+            window_its = window_steps / window_dt
+            # MFU 报的是 dense 6P 等价吞吐；QAT 会因 fake-quant 逐 Linear 加 round/clamp/dequant
+            # 有额外内存搬运，MFU 通常比同 seq 下的 SFT 低 5-15%——这本身就是我们要观测的量。
+            tokens_window = window_steps * args.batch_size * args.max_seq_len * world_size
+            mfu = (flops_per_token * tokens_window / window_dt) / (gpu_peak_tflops * 1e12) * 100
+            last_log_time, last_log_step = now, step
             current_loss = loss.item() * args.accumulation_steps
             current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0
             current_logits_loss = current_loss - current_aux_loss
             current_lr = optimizer.param_groups[-1]['lr']
             eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 60
+            # observer_warmup 期内 act_fq 只统计不量化，日志上会看到一个"200 步后 loss 抖一下"的转折点。
+            qat_phase = 'observe' if step <= args.quant_observer_steps else 'fake-quant'
             Logger(
-                f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), '
+                f'[{time.strftime("%H:%M:%S")}] Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), '
+                f'it/s: {window_its:.2f}, MFU: {mfu:.1f}%, phase: {qat_phase}, '
                 f'loss: {current_loss:.4f}, logits_loss: {current_logits_loss:.4f}, '
                 f'aux_loss: {current_aux_loss:.4f}, lr: {current_lr:.8f}, '
                 f'epoch_time: {eta_min:.1f}min'
@@ -67,23 +85,28 @@ def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
                 wandb.log({
                     "loss": current_loss, "logits_loss": current_logits_loss,
                     "aux_loss": current_aux_loss, "learning_rate": current_lr,
-                    "epoch_time": eta_min,
+                    "epoch_time": eta_min, "it_per_s": window_its, "mfu_pct": mfu,
                 })
 
-        if (step % args.save_interval == 0 or step == iters) and is_main_process():
+        if step % args.save_interval == 0 or step == iters:
-            model.eval()
+            if is_main_process():
-            moe_suffix = '_moe' if lm_config.use_moe else ''
+                model.eval()
-            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
+                moe_suffix = '_moe' if lm_config.use_moe else ''
-            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+                ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
-            raw_model = getattr(raw_model, '_orig_mod', raw_model)
+                raw_model = model.module if isinstance(model, DistributedDataParallel) else model
-            state_dict = raw_model.state_dict()
+                raw_model = getattr(raw_model, '_orig_mod', raw_model)
-            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
+                state_dict = raw_model.state_dict()
-            lm_checkpoint(
+                ckp_tmp = ckp + '.tmp'
-                lm_config, weight=args.save_weight, model=model, optimizer=optimizer,
+                torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp_tmp)
-                epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler,
+                os.replace(ckp_tmp, ckp)
-            )
+                lm_checkpoint(
-            model.train()
+                    lm_config, weight=args.save_weight, model=model, optimizer=optimizer,
-            del state_dict
+                    epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler,
+                )
+                model.train()
+                del state_dict
+            if dist.is_initialized():
+                dist.barrier()
 
         del input_ids, labels, res, loss
 
@@ -114,9 +137,13 @@ if __name__ == "__main__":
     parser.add_argument('--num_hidden_layers', default=8, type=int)
     parser.add_argument('--max_seq_len', default=768, type=int)
     parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1])
-    parser.add_argument("--data_path", type=str, default="../dataset/sft_t2t_mini.jsonl")
+    parser.add_argument("--data_path", type=str, default="../dataset/sft_t2t.jsonl")
+    parser.add_argument("--qat_subset_size", type=int, default=0,
+                        help="QAT 数据子集大小（0=用全量）；64M dry run 建议 100000，1B 正式建议 200000")
+    parser.add_argument("--gpu_peak_tflops", type=float, default=0,
+                        help="GPU bf16 理论峰值 TFLOPS，用于 MFU 计算；0=按 GPU 名自动检测")
     parser.add_argument('--from_weight', default='full_sft', type=str,
-                        help="QAT 起点权重前缀（默认基于 full_sft 微调）")
+                        help="QAT 起点权重前缀（部署走 full_sft；1.12 dry run 走 pretrain）")
     parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1])
     parser.add_argument("--use_wandb", action="store_true")
     parser.add_argument("--wandb_project", type=str, default="MiniMind-QAT")
@@ -179,10 +206,20 @@ if __name__ == "__main__":
            f'(W{args.w_bits}A{args.a_bits}, skip={skip_patterns}, '
            f'observer_warmup={args.quant_observer_steps} steps)')
 
-    train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
+    train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len,
+                          subset_size=args.qat_subset_size)
+    Logger(f'QAT dataset: {len(train_ds)} samples '
+           f'(subset_size={args.qat_subset_size}, seq_len={args.max_seq_len})')
     train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
     scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
-    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate, fused=True)
+
+    # MFU 测量准备
+    flops_per_token = compute_model_flops_per_token(model, lm_config)
+    gpu_peak_tflops = detect_gpu_peak_tflops_bf16(override=args.gpu_peak_tflops)
+    world_size = dist.get_world_size() if dist.is_initialized() else 1
+    Logger(f'[MFU] params_flops/token={flops_per_token/1e9:.3f} GFLOPS, '
+           f'GPU peak={gpu_peak_tflops:.0f} TFLOPS bf16, world_size={world_size}')
 
     # 6. resume
     start_epoch, start_step = 0, 0
@@ -208,7 +245,8 @@ if __name__ == "__main__":
         skip = start_step if (epoch == start_epoch and start_step > 0) else 0
         batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
         loader = DataLoader(train_ds, batch_sampler=batch_sampler,
-                            num_workers=args.num_workers, pin_memory=True)
+                            num_workers=args.num_workers, pin_memory=True,
+                            persistent_workers=args.num_workers > 0)
         if skip > 0:
             Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
             train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)

trainer/trainer_utils.py

@@ -37,14 +37,28 @@ def Logger(content):
         print(content)
 
 
-def get_lr(current_step, total_steps, lr):
+def get_lr(current_step, total_steps, lr, warmup_steps=0):
-    return lr*(0.1 + 0.45*(1 + math.cos(math.pi * current_step / total_steps)))
+    """
+    Linear warmup + cosine decay。
+    - warmup_steps=0（默认）: 纯 cosine，向后兼容
+    - warmup_steps>0: 前 warmup_steps 步从 0.1*lr 线性升到 lr，之后 cosine 到 0.1*lr
+
+    1B from-scratch 必须 warmup（无 warmup 时前 500 步 loss 发散概率 >50%，
+    因为随机初始化的权重在 full lr 下第一步梯度过大打乱结构）。
+    """
+    if warmup_steps > 0 and current_step < warmup_steps:
+        return lr * (0.1 + 0.9 * current_step / warmup_steps)
+    progress = (current_step - warmup_steps) / max(total_steps - warmup_steps, 1)
+    return lr * (0.1 + 0.45 * (1 + math.cos(math.pi * progress)))
 
 
 # A100/H100/3090 等常见卡的 bf16 理论峰值 TFLOPS。MoE 等价 FLOPs 是 active params 不是 total，
 # active params 算出来的 MFU 才有意义；MoE 的 active 比例由 num_experts_per_tok / num_experts 决定。
 _GPU_PEAK_BF16_TFLOPS = {
     'A100': 312, 'A100-SXM4-80GB': 312, 'A100-SXM4-40GB': 312, 'A100-PCIE-80GB': 312, 'A100-PCIE-40GB': 312,
+    'PG509-210': 312, 'PG509-200': 312,  # NVIDIA marketing name for A100 80/40GB on some devservers
+    'PG506-243': 989,  # H100 SXM5 devserver naming
+    'HGX': 989, 'HGX H100': 989,
     'H100': 989, 'H100-SXM5-80GB': 989, 'H100-PCIE-80GB': 756, 'H100 80GB HBM3': 989,
     'H200': 989,
     'RTX 3090': 71, 'GeForce RTX 3090': 71, 'RTX 3090 Ti': 80,
@@ -112,6 +126,14 @@ def setup_seed(seed: int, deterministic: bool = False):
     torch.set_float32_matmul_precision('high')
 
 def lm_checkpoint(lm_config, weight='full_sft', model=None, optimizer=None, epoch=0, step=0, wandb=None, save_dir='../checkpoints', **kwargs):
+    """
+    save mode (model != None): 只在 rank-0 写盘，非 main rank 早返回。
+    load mode (model == None): 所有 rank 读同一份 ckpt，无冲突。
+    调用方在 save 后应加 dist.barrier()，让非-0 rank 等待 rank-0 写完再进入下一 step。
+    """
+    # DDP 保护：save 必须由单一 rank 执行；否则 4 rank 同写一个 file 会 corrupt。
+    if model is not None and not is_main_process():
+        return None
     os.makedirs(save_dir, exist_ok=True)
     moe_path = '_moe' if lm_config.use_moe else ''
     ckp_path = f'{save_dir}/{weight}_{lm_config.hidden_size}{moe_path}.pth'