diff --git a/model/model_minimind.py b/model/model_minimind.py index 1b34cee..3f23fdf 100755 --- a/model/model_minimind.py +++ b/model/model_minimind.py @@ -125,10 +125,6 @@ class Attention(nn.Module): if self.flash and (seq_len > 1) and (not self.is_causal or past_key_value is None) and (attention_mask is None or torch.all(attention_mask == 1)): output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=self.is_causal) else: - # 先转 float32 再加 mask:-1e9 超出 float16 的表示范围(min≈-65504), - # 在 float16 张量上累加会饱和成 -inf。若某个 query 在因果范围内可见的 - # key 全被 mask 掉(左 padding 时必然出现),softmax 整行就是 NaN, - # 并在下一层通过 K/V 投影扩散到同序列里有效的 token 上。 scores = (xq @ xk.transpose(-2, -1)).float() / math.sqrt(self.head_dim) if self.is_causal: scores[:, :, :, -seq_len:] += torch.full((seq_len, seq_len), float("-inf"), device=scores.device).triu(1) if attention_mask is not None: scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9 diff --git a/trainer/train_agent.py b/trainer/train_agent.py index 4a72a11..7ef6ed1 100644 --- a/trainer/train_agent.py +++ b/trainer/train_agent.py @@ -273,8 +273,6 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model rewards = calculate_rewards(prompts, completions, gt_batch, tools_batch, args.num_generations, reward_model, device=args.device, turn_outputs_batch=turn_outputs_batch, unfinished_batch=unfinished_batch) with autocast_ctx: - # 反向传播必须经过 DDP 包装后的模块:直接调用 .module 会跳过 - # DDP 的 prepare_for_backward,梯度不会 all-reduce,各卡静默发散。 res = model(input_ids, attention_mask=full_mask) aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device) logits = res.logits[:, :-1, :] diff --git a/trainer/train_grpo.py b/trainer/train_grpo.py index a7b01cf..21f5d9e 100755 --- a/trainer/train_grpo.py +++ b/trainer/train_grpo.py @@ -96,8 +96,6 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod rewards = calculate_rewards(prompts, completions, reward_model).to(args.device) # [B*num_gen] with autocast_ctx: - # 反向传播必须经过 DDP 包装后的模块:直接调用 .module 会跳过 - # DDP 的 prepare_for_backward,梯度不会 all-reduce,各卡静默发散。 res = model(outputs, attention_mask=full_mask) aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device) per_token_logps = F.log_softmax(res.logits[:, :-1, :], dim=-1).gather(2, outputs[:, 1:].unsqueeze(-1)).squeeze(-1).gather(1, logp_pos) diff --git a/trainer/train_ppo.py b/trainer/train_ppo.py index b66d4b6..4290b63 100644 --- a/trainer/train_ppo.py +++ b/trainer/train_ppo.py @@ -166,16 +166,12 @@ def ppo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, actor_sched for i in range(0, B, mb_size): inds = b_inds[i:i + mb_size] - # 反向传播必须经过 DDP 包装后的模块:直接调用 .module 会跳过 - # DDP 的 prepare_for_backward,梯度不会 all-reduce,各卡静默发散。 mb_values_seq = critic_model(input_ids=gen_out[inds], attention_mask=full_mask[inds]) mb_resp_values = mb_values_seq.gather(1, logp_pos[inds]) with autocast_ctx: res = actor_model(input_ids=gen_out[inds], attention_mask=full_mask[inds]) aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device) - # 在 autocast 内计算 log_softmax,避免直接对 fp16/bf16 logits - # 计算造成额外数值偏差。 mb_resp_logp = F.log_softmax(res.logits[:, :-1], dim=-1).gather(2, labels[inds].unsqueeze(-1)).squeeze(-1).gather(1, logp_pos[inds]) log_ratio = mb_resp_logp - old_resp_logp[inds]