compat megatron 0.20 release and latest dev branch

src/mcore_bridge/bridge/gpt_bridge.py

@@ -56,6 +56,11 @@ class GPTBridge:
     additional_dim1_keys = set()
     _support_hf_grouped_lora = True
 
+    def __deepcopy__(self, memo):
+        # Bridges hold process groups; per-layer config copies must share, not copy, them.
+        memo[id(self)] = self
+        return self
+
     def __init__(self, config: ModelConfig):
         self.config = config
         self._disable_tqdm = False

src/mcore_bridge/config/parser.py

@@ -329,6 +329,12 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]:
         # field); Megatron dev's hybrid transformer block, which expands the residual into mHC
         # streams, reads `hc_count` for the same thing.
         res['hc_count'] = res['num_residual_streams']
+        # dev splits mHC into two interfaces; the hybrid stack reads the new one.
+        from megatron.core.transformer.transformer_config import TransformerConfig
+        tc = TransformerConfig.__dataclass_fields__
+        if 'enable_mhc_connections' in tc:
+            res['enable_mhc_connections'] = True
+            res['mhc_num_residual_streams'] = res['num_residual_streams']
         pattern = ''.join(('K' if attn == 'linear_attention' else 'D') + ('E' if mlp == 'sparse' else '-')
                           for attn, mlp in zip(layer_types, mlp_layer_types))
         res['hybrid_layer_pattern'] = pattern

src/mcore_bridge/model/mm_gpts/glm5_next.py

@@ -13,7 +13,7 @@
 
 from ..constant import ModelType
 from ..mm_gpt_model import MultimodalGPTModel
-from ..mtp import get_physical_mtp_depth
+from ..mtp import get_mtp_capabilities, get_physical_mtp_depth
 from ..register import ModelLoader, ModelMeta, register_model
 from .utils import HuggingFaceVit
 
@@ -156,6 +156,9 @@ class Glm5NextMultiTokenPredictionLayer(McoreMultiTokenPredictionLayer):
     def __init__(self, config, *args, **kwargs):
         mtp_config = copy.copy(config)
         mtp_config.enable_hyper_connections = False
+        # Keep the fused `eh_proj` shape the checkpoint stores (no mHC in the MTP head).
+        if hasattr(mtp_config, 'enable_mhc_connections'):
+            mtp_config.enable_mhc_connections = False
         super().__init__(mtp_config, *args, **kwargs)
 
     def _get_embeddings(self,
@@ -166,33 +169,40 @@ def _get_embeddings(self,
                         packed_seq_params=None,
                         padding_mask=None,
                         sequence_roll_context=None,
-                        roll_depth=0):
+                        roll_depth=0,
+                        mtp_input_mask=None):
         # Decoder/MoE routing consumes the sequence-parallel shard, but MTP must roll the mask
         # beside the unsharded input_ids. Reconstruct it only around the roll, then restore the
         # shard expected by the nested HybridStack.
         sequence_parallel_mask = (
             padding_mask is not None and self.config.sequence_parallel and self.tp_group.size() > 1)
+        full_padding_len = None
         if sequence_parallel_mask:
             padding_mask = gather_from_sequence_parallel_region(
                 padding_mask.transpose(0, 1).contiguous(),
                 tensor_parallel_output_grad=False,
                 group=self.tp_group,
             ).transpose(0, 1).contiguous()
+            full_padding_len = padding_mask.shape[1]
 
-        input_ids, position_ids, padding_mask, decoder_input, hidden_states = super()._get_embeddings(
+        super_kwargs = dict(
             input_ids=input_ids,
             position_ids=position_ids,
             embedding=embedding,
             hidden_states=hidden_states,
             packed_seq_params=packed_seq_params,
             padding_mask=padding_mask,
             sequence_roll_context=sequence_roll_context,
-            roll_depth=roll_depth,
-        )
-
-        if sequence_parallel_mask:
-            padding_mask = padding_mask.chunk(self.tp_group.size(), dim=-1)[self.tp_group.rank()].contiguous()
-        return input_ids, position_ids, padding_mask, decoder_input, hidden_states
+            roll_depth=roll_depth)
+        # dev forwards `mtp_input_mask` unconditionally.
+        if get_mtp_capabilities().has_layer_input_mask:
+            super_kwargs['mtp_input_mask'] = mtp_input_mask
+        # dev returns six values (mask included), releases five; keep upstream's shape.
+        outputs = list(super()._get_embeddings(**super_kwargs))
+        # Upstream already scatters the mask under sequence parallel; chunk only if it did not.
+        if (sequence_parallel_mask and outputs[2] is not None and outputs[2].shape[1] == full_padding_len):
+            outputs[2] = outputs[2].chunk(self.tp_group.size(), dim=-1)[self.tp_group.rank()].contiguous()
+        return tuple(outputs)
 
 
 class Glm5NextHybridModel(HybridModel):
@@ -365,7 +375,9 @@ class Glm5NextVit(HuggingFaceVit):
 
     def prepare_model(self, hf_config: PretrainedConfig):
         from transformers.models.glm5_next import Glm5NextVisionModel
-        self.visual = Glm5NextVisionModel._from_config(hf_config.vision_config)
+
+        # This vision tower only implements sdpa; transformers would otherwise pick flash-attn.
+        self.visual = Glm5NextVisionModel._from_config(hf_config.vision_config, attn_implementation='sdpa')
 
     def prepare_language_model(self, hf_config: PretrainedConfig):
         self.visual = None

src/mcore_bridge/model/mtp.py

@@ -42,6 +42,15 @@ class MtpCapabilities:
 
     has_detach_heads: bool = False
 
+    # Knobs that differ across 0.20 / main / dev; gate on these, not on version numbers.
+    has_mtp_hybrid_override_pattern: bool = False
+    has_freeze_base_model_for_mtp: bool = False
+    has_mtp_loss_early_backward: bool = False
+    has_keep_mtp_in_bf16: bool = False
+
+    # `_get_embeddings` gained a trailing `mtp_input_mask`; 0.18/0.19 still return five values.
+    has_layer_input_mask: bool = False
+
     layer_forward_params: Tuple[str, ...] = ()
     layer_returns_padding_mask: bool = False
 
@@ -82,20 +91,30 @@ def get_mtp_capabilities() -> MtpCapabilities:
     tc = TransformerConfig.__dataclass_fields__
 
     layer_fwd = _params(McoreMultiTokenPredictionLayer.forward)
+    emb_fwd = _params(getattr(McoreMultiTokenPredictionLayer, '_get_embeddings', None))
     ckpt = _params(getattr(McoreMultiTokenPredictionLayer, '_checkpointed_forward', None))
     block_fwd = _params(McoreMultiTokenPredictionBlock.forward)
+    has_roll_context = hasattr(_mtp_mod, 'MTPSequenceRollContext')
 
     caps = MtpCapabilities(
         has_detach_heads='mtp_detach_heads' in tc,
+        has_mtp_hybrid_override_pattern='mtp_hybrid_override_pattern' in tc,
+        has_freeze_base_model_for_mtp='freeze_base_model_for_mtp' in tc,
+        has_mtp_loss_early_backward='mtp_loss_early_backward' in tc,
+        has_keep_mtp_in_bf16='keep_mtp_in_bf16' in tc,
+        has_layer_input_mask='mtp_input_mask' in emb_fwd,
         layer_forward_params=layer_fwd,
         layer_returns_padding_mask='padding_mask' in layer_fwd,
         checkpointed_forward_accepts_padding_mask='padding_mask' in ckpt,
         block_forward_params=block_fwd,
-        has_roll_context=hasattr(_mtp_mod, 'MTPSequenceRollContext'),
+        has_roll_context=has_roll_context,
         has_mhc_multistream='mhc_multistream' in block_fwd,
         fingerprint=(f'pm={"padding_mask" in layer_fwd} ckpt_pm={"padding_mask" in ckpt} '
-                     f'roll_ctx={hasattr(_mtp_mod, "MTPSequenceRollContext")} '
-                     f'mhc={"mhc_multistream" in block_fwd}'),
+                     f'roll_ctx={has_roll_context} mhc={"mhc_multistream" in block_fwd} '
+                     f'hybrid_pat={"mtp_hybrid_override_pattern" in tc} '
+                     f'early_bwd={"mtp_loss_early_backward" in tc} '
+                     f'keep_bf16={"keep_mtp_in_bf16" in tc} '
+                     f'layer_input_mask={"mtp_input_mask" in emb_fwd}'),
     )
     logger.debug(f'MTP capabilities: {caps.fingerprint}')
     return caps
@@ -653,31 +672,22 @@ def process_mtp_loss(
     packed_seq_params=None,
     **extra,
 ) -> torch.Tensor:
-    """Call upstream ``process_mtp_loss`` with the kwargs the installed version declares.
-
-    This wrapper cannot be dropped: 0.18 has no ``tp_group``/``input_ids`` parameters, while
-    0.19+ and dev need them, and ``gpt_model`` passes them unconditionally so that one call site
-    serves every supported version. *loss_mask* must already exclude ``-100`` tokens;
-    *output_layer* follows the upstream convention ``(hs, weight=..., runtime_gather_output=...)
-    -> (logits, bias)``.  ``scale_logits_fn`` is never forwarded so the MTP loss stays in sync
-    with the main loss (which does not apply MuP scaling).
-    """
     accepted_kwargs = set(_PROCESS_MTP_LOSS_PARAMS)
+    optional_kwargs = {'runtime_gather_output': runtime_gather_output, **extra}
     return _upstream_process_mtp_loss(
         hidden_states=hidden_states,
         labels=labels,
         loss_mask=loss_mask,
         output_layer=output_layer,
         output_weight=output_weight,
-        runtime_gather_output=runtime_gather_output,
         is_training=is_training,
         compute_language_model_loss=compute_language_model_loss,
         config=config,
         cp_group=cp_group,
         packed_seq_params=packed_seq_params,
         **{
             k: v
-            for k, v in extra.items() if k in accepted_kwargs
+            for k, v in optional_kwargs.items() if k in accepted_kwargs
         })