diff --git a/dataset/lm_dataset.py b/dataset/lm_dataset.py index bbc4197..f22360d 100644 --- a/dataset/lm_dataset.py +++ b/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 diff --git a/scripts/eval_perplexity.py b/scripts/eval_perplexity.py index 99061f8..c5cd022 100644 --- a/scripts/eval_perplexity.py +++ b/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') diff --git a/scripts/measure_sqnr.py b/scripts/measure_sqnr.py new file mode 100644 index 0000000..57a5bcd --- /dev/null +++ b/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() diff --git a/trainer/train_qat.py b/trainer/train_qat.py index 9d9c227..484138b 100644 --- a/trainer/train_qat.py +++ b/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)