fix
Longfang Zhao
parent
a6a19aa877
commit
02b61435b1
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
@ -173,10 +179,13 @@ if __name__ == "__main__":
|
|||
|
||||
# ========== 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])
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Reference in New Issue