update

dataset/lm_dataset.py

@@ -56,12 +56,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
 

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)
-    model = model.half().eval().to(args.device)
-    return model, cfg
+
+    is_qat_ckpt = _is_qat_state_dict(state_dict)
+    apply_qat = bool(args.apply_qat) or is_qat_ckpt
+
+    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_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,11 @@ 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)
-        labels = labels.to(args.device)
+        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 +53,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,7 +81,7 @@ 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():
@@ -114,9 +128,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 +197,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 +236,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)