From ab66c0c9ff74848e30af441589722dc268af1e81 Mon Sep 17 00:00:00 2001 From: Xialie Zhuang <62231346+WhichPaths@users.noreply.github.com> Date: Fri, 25 Sep 2026 10:46:32 +0800 Subject: [PATCH] fix MTP checkpoint saving with TP > 1 MultiTokenPredictionLayer.tp_group stays None on mcore 0.16, so sharded_state_dict computes tp rank 0 on every TP rank and marks the replicated enorm/hnorm weights as main replicas on all of them. Saving then fails with "Invalid access pattern". Fall back to the global TP group, which is what forward already resolves None to. --- src/mcore_bridge/model/modules/mtp_layer.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/mcore_bridge/model/modules/mtp_layer.py b/src/mcore_bridge/model/modules/mtp_layer.py index 1dd1eeb..7f280a8 100644 --- a/src/mcore_bridge/model/modules/mtp_layer.py +++ b/src/mcore_bridge/model/modules/mtp_layer.py @@ -42,6 +42,10 @@ def __init__(self, config: ModelConfig, submodules, *args, **kwargs): if replace_eh_proj: submodules.eh_proj = eh_proj self.tp_group = getattr(self, 'tp_group', None) + if self.tp_group is None: + # mcore 0.16 doesn't set tp_group on the MTP layer. If it stays None, sharded_state_dict uses + # tp rank 0 on every TP rank and the replicated enorm/hnorm end up with two main replicas. + self.tp_group = parallel_state.get_tensor_model_parallel_group(check_initialized=False) if not replace_eh_proj: return fp8_context = transformer_engine.pytorch.fp8_model_init(enabled=False)