From 02b61435b12dbeb0bd321226e3064a93b38c0797 Mon Sep 17 00:00:00 2001 From: Longfang Zhao Date: Tue, 30 Jun 2026 23:23:21 -0700 Subject: [PATCH] fix --- model/model_minimind.py | 20 ++++++++++++-------- trainer/train_full_sft.py | 4 ++++ trainer/train_pretrain.py | 37 +++++++++++++++++++++++-------------- trainer/train_qat.py | 4 ++++ trainer/trainer_utils.py | 3 +++ 5 files changed, 46 insertions(+), 22 deletions(-) diff --git a/model/model_minimind.py b/model/model_minimind.py index 419c4b2..0236639 100755 --- a/model/model_minimind.py +++ b/model/model_minimind.py @@ -206,20 +206,24 @@ class MiniMindModel(nn.Module): self.register_buffer("freqs_cos", freqs_cos, persistent=False) self.register_buffer("freqs_sin", freqs_sin, persistent=False) + 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 if hasattr(past_key_values, 'layers'): past_key_values = None 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. - # Hot path: skipped via cheap python bool (no GPU sync). The original `freqs_cos[0,0]==0` - # tensor compare forced a device-to-host sync on every forward. - if getattr(self, '_rope_ready', False) is False: - if self.freqs_cos.numel() == 0 or bool(self.freqs_cos[0, 0].eq(0).cpu()): - 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 + # RoPE buffers 已在 __init__ 里 precompute + register_buffer；forward 直接切片即可。 + # (原来这里有 lazy re-compute 兜底 transformers>=5.x meta-device init 掉 buffer 的情况， + # 但那段 `bool(freqs_cos[0,0].eq(0).cpu())` 是 torch.compile graph break——每步 forward + # 都要在 compiled ↔ eager 间切换 + guard 检查，短 step 下吃掉一半吞吐。如需 meta-init + # 兼容，在 model.to(device) 后手动调 warmup_rope() 即可，不要放 forward hot path。) position_embeddings = (self.freqs_cos[start_pos:start_pos + seq_length], self.freqs_sin[start_pos:start_pos + seq_length]) presents = [] for layer, past_key_value in zip(self.layers, past_key_values): diff --git a/trainer/train_full_sft.py b/trainer/train_full_sft.py index 0ee563e..e9348a9 100644 --- a/trainer/train_full_sft.py +++ b/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) diff --git a/trainer/train_pretrain.py b/trainer/train_pretrain.py index c2f64d7..9f2a033 100644 --- a/trainer/train_pretrain.py +++ b/trainer/train_pretrain.py @@ -30,6 +30,12 @@ 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 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) @@ -113,7 +119,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") @@ -127,17 +133,17 @@ if __name__ == "__main__": local_rank = init_distributed_mode() if dist.is_initialized(): args.device = f"cuda:{local_rank}" setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0)) - + # ========== 2. 配置目录、模型参数、检查ckp ========== os.makedirs(args.save_dir, exist_ok=True) lm_config = MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe)) ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume==1 else None - + # ========== 3. 设置混合精度 ========== device_type = "cuda" if "cuda" in args.device else "cpu" dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16 autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype) - + # ========== 4. 配wandb ========== wandb = None if args.use_wandb and is_main_process(): @@ -146,7 +152,7 @@ if __name__ == "__main__": resume = 'must' if wandb_id else None wandb_run_name = f"MiniMind-Pretrain-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}" wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume) - + # ========== 5. 定义模型、数据、优化器 ========== model, tokenizer = init_model(lm_config, args.from_weight, device=args.device) train_ds = PretrainDataset(args.data_path, tokenizer, max_length=args.max_seq_len) @@ -170,16 +176,19 @@ if __name__ == "__main__": scaler.load_state_dict(ckp_data['scaler']) start_epoch = ckp_data['epoch'] start_step = ckp_data.get('step', 0) - + # ========== 7. 编译和分布式包装 ========== if args.use_compile == 1: - # max-autotune 在 Triton 层为每个 fused kernel 试多套 tile/block 配置，挑最快的。 - # 编译时间多 ~30-60s，稳态可比 default 快 5-15%。 - model = torch.compile(model, mode="max-autotune") - Logger('torch.compile enabled (mode=max-autotune)') + # 新 CUDA 13 stack 上 max-autotune 会为每个 GEMM shape 试 21 种 triton kernel（分钟级 + # autotune 循环，且实测 cuBLAS aten mm 几乎每次都赢过 triton 变体，autotune 白花时间）。 + # 用 default 更实惠——inductor 融合仍在，只是不额外扫 triton kernel。 + # (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]) - + # ========== 8. 开始训练 ========== for epoch in range(start_epoch, args.epochs): train_sampler and train_sampler.set_epoch(epoch) @@ -188,13 +197,13 @@ if __name__ == "__main__": 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, persistent_workers=args.num_workers > 0) - if skip > 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) else: train_epoch(epoch, loader, len(loader), 0, wandb) - + # ========== 9. 清理分布进程 ========== if dist.is_initialized(): dist.barrier() - dist.destroy_process_group() \ No newline at end of file + dist.destroy_process_group() diff --git a/trainer/train_qat.py b/trainer/train_qat.py index 484138b..91e5592 100644 --- a/trainer/train_qat.py +++ b/trainer/train_qat.py @@ -31,6 +31,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 + CUDA graph 每步必须 mark step 边界，见 train_pretrain.py 同处注释。 + # 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 diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py index 9347d06..66608a1 100644 --- a/trainer/trainer_utils.py +++ b/trainer/trainer_utils.py @@ -45,6 +45,9 @@ def get_lr(current_step, total_steps, lr): # 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,