From 14907897ef57259a7893aa1171665ce97a379e56 Mon Sep 17 00:00:00 2001 From: rj42 Date: Tue, 9 Dec 2025 13:50:31 +0300 Subject: [PATCH 1/2] Fix: don't enter branch if mtp_num_layers == 0 (#2581) Co-authored-by: Xin Yao --- megatron/core/models/gpt/gpt_model.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index e70221d2cfa..ef76281ac08 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -582,7 +582,8 @@ def _postprocess( if not self.post_process: return hidden_states - if self.config.mtp_num_layers is not None: + # Skip when mtp_num_layers is None or 0 + if self.config.mtp_num_layers: mtp_labels = labels.clone() hidden_states_list = torch.chunk(hidden_states, 1 + self.config.mtp_num_layers, dim=0) hidden_states = hidden_states_list[0] From 62a1cb4084e2f265bc707f39097b5fbb628b1ed9 Mon Sep 17 00:00:00 2001 From: Juntao Wang Date: Thu, 15 Jan 2026 12:06:28 +0800 Subject: [PATCH 2/2] Fix clip_qk for virtual pipeline size > 1 (#2776) Co-authored-by: Xin Yao --- megatron/core/optimizer/qk_clip.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/megatron/core/optimizer/qk_clip.py b/megatron/core/optimizer/qk_clip.py index 72127f94712..26b5787cd50 100644 --- a/megatron/core/optimizer/qk_clip.py +++ b/megatron/core/optimizer/qk_clip.py @@ -22,6 +22,11 @@ def clip_qk(model, log_max_only=False) -> float: for model_chunk in model: for transformer_layer in model_chunk.module.module.decoder.layers: if hasattr(transformer_layer.self_attention, 'clip_qk'): + if ( + transformer_layer.self_attention.core_attention.current_max_attn_logits + is None + ): + continue torch.distributed.all_reduce( transformer_layer.self_attention.core_attention.current_max_attn_logits, op=torch.distributed.ReduceOp.MAX,