mcore-bridge 1.6.4__tar.gz → 1.8.0.dev0__tar.gz

{mcore_bridge-1.6.4 → mcore_bridge-1.8.0.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mcore_bridge
-Version: 1.6.4
+Version: 1.8.0.dev0
 Summary: MCore-Bridge: Making Megatron training as simple as Transformers
 Home-page: https://github.com/modelscope/mcore-bridge
 Author: ModelScope teams
@@ -19,12 +19,12 @@ Classifier: Programming Language :: Python :: 3.12
 Requires-Python: >=3.8.0
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: megatron-core<0.20,>=0.16
+Requires-Dist: megatron-core<0.21,>=0.18
 Requires-Dist: modelscope
 Requires-Dist: peft<0.21,>=0.11
 Requires-Dist: safetensors
 Requires-Dist: tqdm
-Requires-Dist: transformers<5.17.0,>=4.33
+Requires-Dist: transformers>=4.33
 
 # MCore-Bridge: Making Megatron training as simple as Transformers
 
@@ -48,7 +48,7 @@ Requires-Dist: transformers<5.17.0,>=4.33
 <p align="center">
 <img src="https://img.shields.io/badge/python-3.12-5be.svg">
 <img src="https://img.shields.io/badge/pytorch-%E2%89%A52.0-orange.svg">
-<a href="https://github.com/NVIDIA/Megatron-LM/"><img src="https://img.shields.io/badge/megatron--core-%E2%89%A50.16-76B900.svg"></a>
+<a href="https://github.com/NVIDIA/Megatron-LM/"><img src="https://img.shields.io/badge/megatron--core-%E2%89%A50.18-76B900.svg"></a>
 <!-- <a href="https://mcore-bridge.readthedocs.io/en/latest/"><img src="https://img.shields.io/badge/docs-latest-blue.svg"></a> -->
 <a href="https://pypi.org/project/mcore-bridge/"><img src="https://badge.fury.io/py/mcore-bridge.svg"></a>
 <a href="https://github.com/modelscope/mcore-bridge/blob/main/LICENSE"><img src="https://img.shields.io/github/license/modelscope/mcore-bridge"></a>
@@ -138,7 +138,7 @@ Recommended Runtime Environment:
 | torch        | >=2.0        | 2.8.0/2.11.0       |                    |
 | transformer-engine    | >=2.3       |  2.14.1    |                  |
 | apex |         | 0.1 | Optional |
-| megatron-core    |   >=0.16,<0.20    | 0.17.1      |                  |
+| megatron-core    |   >=0.18,<0.21    | 0.19.2      |                  |
 | flash-attn    |        | 2.8.3/3.0.0b1   |     Optional      |
 | transformers | >=4.33       | 4.57.6/5.8.1   |                    |
 | modelscope   | >=1.23       |             |                    |

{mcore_bridge-1.6.4 → mcore_bridge-1.8.0.dev0}/README.md

@@ -20,7 +20,7 @@
 <p align="center">
 <img src="https://img.shields.io/badge/python-3.12-5be.svg">
 <img src="https://img.shields.io/badge/pytorch-%E2%89%A52.0-orange.svg">
-<a href="https://github.com/NVIDIA/Megatron-LM/"><img src="https://img.shields.io/badge/megatron--core-%E2%89%A50.16-76B900.svg"></a>
+<a href="https://github.com/NVIDIA/Megatron-LM/"><img src="https://img.shields.io/badge/megatron--core-%E2%89%A50.18-76B900.svg"></a>
 <!-- <a href="https://mcore-bridge.readthedocs.io/en/latest/"><img src="https://img.shields.io/badge/docs-latest-blue.svg"></a> -->
 <a href="https://pypi.org/project/mcore-bridge/"><img src="https://badge.fury.io/py/mcore-bridge.svg"></a>
 <a href="https://github.com/modelscope/mcore-bridge/blob/main/LICENSE"><img src="https://img.shields.io/github/license/modelscope/mcore-bridge"></a>
@@ -110,7 +110,7 @@ Recommended Runtime Environment:
 | torch        | >=2.0        | 2.8.0/2.11.0       |                    |
 | transformer-engine    | >=2.3       |  2.14.1    |                  |
 | apex |         | 0.1 | Optional |
-| megatron-core    |   >=0.16,<0.20    | 0.17.1      |                  |
+| megatron-core    |   >=0.18,<0.21    | 0.19.2      |                  |
 | flash-attn    |        | 2.8.3/3.0.0b1   |     Optional      |
 | transformers | >=4.33       | 4.57.6/5.8.1   |                    |
 | modelscope   | >=1.23       |             |                    |

{mcore_bridge-1.6.4 → mcore_bridge-1.8.0.dev0}/setup.py

@@ -45,6 +45,9 @@ if __name__ == '__main__':
         url='https://github.com/modelscope/mcore-bridge',
         package_dir={'': 'src'},
         packages=find_packages('src'),
+        # The Megatron patch for GLM-5.3 ships inside the package so an installed mcore-bridge can
+        # apply it (see mcore_bridge/tools/apply_megatron_patch.py).
+        package_data={'mcore_bridge': ['patches/*.patch']},
         python_requires='>=3.8.0',
         classifiers=[
             'Development Status :: 4 - Beta',

{mcore_bridge-1.6.4 → mcore_bridge-1.8.0.dev0}/src/mcore_bridge/__init__.py

@@ -9,18 +9,18 @@ apply_patch()
 
 if TYPE_CHECKING:
     from .bridge import GPTBridge
-    from .config import ModelConfig, hf_to_mcore_config
-    from .model import get_mcore_model
+    from .config import MLAModelConfig, ModelConfig, hf_to_mcore_config
+    from .model import get_mcore_model, resolve_mtp_config, resolve_mtp_num_layers
     from .tuners import LoraParallelLinear
-    from .utils import get_logger, set_random_seed, split_cp_inputs, unwrap_model
+    from .utils import get_logger, set_random_seed, split_cp_inputs
     from .version import __release_datetime__, __version__
 else:
     _import_structure = {
         'bridge': ['GPTBridge'],
-        'config': ['ModelConfig', 'hf_to_mcore_config'],
-        'model': ['get_mcore_model'],
+        'config': ['MLAModelConfig', 'ModelConfig', 'hf_to_mcore_config'],
+        'model': ['get_mcore_model', 'resolve_mtp_config', 'resolve_mtp_num_layers'],
         'tuners': ['LoraParallelLinear'],
-        'utils': ['get_logger', 'set_random_seed', 'split_cp_inputs', 'unwrap_model'],
+        'utils': ['get_logger', 'set_random_seed', 'split_cp_inputs'],
         'version': ['__release_datetime__', '__version__'],
     }
 

{mcore_bridge-1.6.4 → mcore_bridge-1.8.0.dev0}/src/mcore_bridge/bridge/gpt_bridge.py

@@ -7,6 +7,7 @@ import torch.distributed as dist
 import torch.nn.functional as F
 from contextlib import contextmanager
 from megatron.core import mpu
+from megatron.core.utils import unwrap_model
 from peft import PeftModel
 from peft.utils import ModulesToSaveWrapper
 from tqdm import tqdm
@@ -17,7 +18,8 @@ from typing import Callable, List, Optional, Union
 from mcore_bridge.config import ModelConfig
 from mcore_bridge.tuners import LoraParallelLinear
 from mcore_bridge.utils import (MxFp4Dequantizer, PackedDequantizer, SafetensorLazyLoader, StreamingSafetensorSaver,
-                                deep_getattr, gc_collect, get_logger, is_master, unwrap_model)
+                                deep_getattr, gc_collect, get_logger, is_master)
+from mcore_bridge.utils.constants import EXPORT_CHUNK_BYTES
 
 logger = get_logger()
 
@@ -28,6 +30,9 @@ EP_PP_RANK = None
 
 class GPTBridge:
     fp8_block_size = 128
+    # Bound the per-collective GPU buffer when a full gathered tensor is
+    # streamed to CPU (checkpoint save / CPU-offloaded weight sync).
+    export_chunk_bytes = EXPORT_CHUNK_BYTES
     hf_layers_prefix = 'model.layers'
     hf_mtp_prefix = 'model.layers'
     hf_embed_key = 'model.embed_tokens.weight'
@@ -51,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
@@ -336,11 +346,57 @@ class GPTBridge:
                         hf_state_dict[k] = v.to(self._target_device)
             return self._add_prefix(hf_state_dict, hf_prefix)
 
+    def _stream_to_cpu(self) -> bool:
+        """Whether large gathers/broadcasts should be chunked straight into host
+        memory. Enabled when the exported tensor is requested on CPU (checkpoint
+        save / CPU-offloaded weight sync): materializing the full gathered tensor
+        on GPU can OOM while the training state is still resident. The decision
+        only depends on properties identical across the group, so the collective
+        sequence stays in sync on every rank. Every participating rank still
+        assembles the full CPU result (matching the pre-chunking semantics); only
+        the GPU-resident full-size buffer is avoided."""
+        return self._target_device == 'cpu'
+
+    def _chunk_rows(self, shape, elem_size: int) -> int:
+        """Rows (along dim0) per broadcast/gather chunk, ~export_chunk_bytes."""
+        inner = 1
+        for s in shape[1:]:
+            inner *= s
+        rows = max(1, self.export_chunk_bytes // max(1, inner * elem_size))
+        return max(1, min(rows, shape[0]))
+
+    def _chunked_all_gather_tp(self, tensor, tp_dim: int, tp_group, tp_size: int):
+        """All-gather `tensor` along tp_dim and assemble the result in host memory
+        chunk by chunk, so the full gathered tensor never exists on GPU at once.
+        The chunk count derives from the local shape, which is identical on every
+        rank of the group, so the collective sequence stays in sync."""
+        dim_size = tensor.shape[tp_dim]
+        chunk_bytes = max(1, self.export_chunk_bytes // tp_size)
+        inner_rows = max(1, chunk_bytes // max(1, tensor.numel() // dim_size * tensor.element_size()))
+        out_shape = list(tensor.shape)
+        out_shape[tp_dim] = dim_size * tp_size
+        output = torch.empty(out_shape, dtype=tensor.dtype, device='cpu')
+        for start in range(0, dim_size, inner_rows):
+            end = min(dim_size, start + inner_rows)
+            local = tensor.narrow(tp_dim, start, end - start).contiguous()
+            gathered = [torch.empty_like(local) for _ in range(tp_size)]
+            dist.all_gather(gathered, local, group=tp_group)
+            del local
+            for j in range(tp_size):
+                dst = tuple(
+                    slice(j * dim_size + start, j * dim_size + end) if ax == tp_dim else slice(None)
+                    for ax in range(tensor.ndim))
+                output[dst] = gathered[j].cpu()
+            del gathered
+        return output
+
     def _all_gather_tp(self, tensor, tp_dim, is_expert):
         tensor = None if tensor is None else tensor.to('cuda')
         tp_size = self.etp_size if is_expert else self.tp_size
         tp_group = self.etp_group if is_expert else self.tp_group
         if tensor is not None and tp_dim is not None and tp_size > 1:
+            if self._stream_to_cpu() and tensor.numel() * tensor.element_size() > self.export_chunk_bytes:
+                return self._chunked_all_gather_tp(tensor, tp_dim, tp_group, tp_size)
             if tp_dim == 0:
                 # save memory
                 tensor_shape = list(tensor.shape)
@@ -363,6 +419,33 @@ class GPTBridge:
             del output
         return tensor
 
+    def _chunked_broadcast_pp(self, tensor, shape, dtype, src_rank: int, pp_group):
+        """Chunked pp/ep-pp broadcast (the pp counterpart of _chunked_all_gather_tp):
+        stream the tensor chunk by chunk instead of materializing the full buffer
+        on GPU. On the holder rank `tensor` carries the data and is returned
+        as-is; receivers pass `shape`/`dtype` from the already-broadcast meta and
+        assemble the result in host memory (ranks that do not keep the export
+        still join every collective but skip the host assembly). Both sides
+        derive the same chunk count from the meta shape."""
+        rows = self._chunk_rows(shape, torch.tensor([], dtype=dtype).element_size())
+        if tensor is not None:
+            for start in range(0, shape[0], rows):
+                end = min(shape[0], start + rows)
+                send = tensor[start:end]
+                if not send.is_cuda or send.dtype != dtype or not send.is_contiguous():
+                    send = send.to(device='cuda', dtype=dtype).contiguous()
+                dist.broadcast(send, src=src_rank, group=pp_group)
+            return tensor
+        output = torch.empty(shape, dtype=dtype, device='cpu')
+        buf = None
+        for start in range(0, shape[0], rows):
+            end = min(shape[0], start + rows)
+            if buf is None or buf.shape[0] != end - start:
+                buf = torch.empty([end - start] + list(shape[1:]), device='cuda', dtype=dtype)
+            dist.broadcast(buf, src=src_rank, group=pp_group)
+            output[start:end] = buf.cpu()
+        return output
+
     def _broadcast_ep_pp(self, tensor, is_expert):
         pp_group = self.ep_pp_group if is_expert else self.pp_group
         pp_size = self.ep_pp_size if is_expert else self.pp_size
@@ -373,12 +456,20 @@ class GPTBridge:
             dist.all_reduce(src_rank, group=pp_group)
             src_rank = dist.get_global_rank(pp_group, src_rank.item())
             meta_data = torch.zeros(10, dtype=torch.int64, device='cuda')
-            dtype_mapping = [torch.float64, torch.float32, torch.float16, torch.bfloat16, torch.uint8, torch.int32]
+            # Append NCCL-supported buffer dtypes without renumbering existing metadata.
+            dtype_mapping = [
+                torch.float64, torch.float32, torch.float16, torch.bfloat16, torch.uint8, torch.int32, torch.int64,
+                torch.bool, torch.int8
+            ]
             dtype_mapping_r = {v: k for k, v in enumerate(dtype_mapping)}
             if tensor is None:
                 dist.broadcast(meta_data, src=src_rank, group=pp_group)
                 shape = meta_data[1:1 + meta_data[0]].tolist()
                 dtype = dtype_mapping[meta_data[-1].item()]
+                numel = math.prod(shape) if shape else 1
+                elem_size = torch.empty((), dtype=dtype).element_size()
+                if self._stream_to_cpu() and numel * elem_size > self.export_chunk_bytes and len(shape) > 0:
+                    return self._chunked_broadcast_pp(None, shape, dtype, src_rank, pp_group)
                 tensor = torch.empty(shape, device='cuda', dtype=dtype)
                 dist.broadcast(tensor, src=src_rank, group=pp_group)
             else:
@@ -386,6 +477,9 @@ class GPTBridge:
                 meta_data[1:1 + tensor.ndim] = torch.tensor(tensor.shape, dtype=torch.int64, device='cuda')
                 meta_data[-1] = dtype_mapping_r[tensor.dtype]
                 dist.broadcast(meta_data, src=src_rank, group=pp_group)
+                if self._stream_to_cpu() and tensor.numel() * tensor.element_size() > (
+                        self.export_chunk_bytes) and tensor.ndim > 0:
+                    return self._chunked_broadcast_pp(tensor, list(tensor.shape), tensor.dtype, src_rank, pp_group)
                 dist.broadcast(tensor, src=src_rank, group=pp_group)
         return tensor
 
@@ -405,7 +499,7 @@ class GPTBridge:
                 tensor = [tensor]
             if self._is_fp8_param(tensor[0]):
                 mg_scale_inv = [
-                    t._rowwise_scale_inv[..., :math.ceil(t._rowwise_data.shape[-1] / self.fp8_block_size)]
+                    t._rowwise_scale_inv[..., :math.ceil(t._rowwise_data.shape[-1] / self.fp8_block_size)].contiguous()
                     for t in tensor
                 ]
                 tensor = [t._rowwise_data for t in tensor]
@@ -428,14 +522,16 @@ class GPTBridge:
         if tensor.dtype == torch.uint8:
             mg_scale_inv = self._all_gather_tp(mg_scale_inv, tp_dim, is_expert)
             mg_scale_inv = self._broadcast_ep_pp(mg_scale_inv, is_expert)
+            if mg_scale_inv is not None and mg_scale_inv.device != tensor.device:
+                mg_scale_inv = mg_scale_inv.to(tensor.device)
             tensor = tensor.view(torch.float8_e4m3fn)
         assert tensor is not None, f'mg_key: {mg_key}'
         if offset:
             assert mg_scale_inv is None, f'mg_key: {mg_key}'
             tensor = tensor + offset
         is_embedding = mg_key in {'embedding.word_embeddings.weight', 'output_layer.weight'}
-        if is_embedding and self.config.padded_vocab_size < tensor.shape[0]:
-            tensor = tensor[:self.config.padded_vocab_size]
+        if is_embedding and self.config.actual_vocab_size < tensor.shape[0]:
+            tensor = tensor[:self.config.actual_vocab_size]
         if self._target_device is not None:
             tensor = tensor.to(device=self._target_device)
             if mg_scale_inv is not None:
@@ -1670,7 +1766,11 @@ class GPTBridge:
             self._set_state_dict(hyper_connection, 'bias', hf_state_dict, f'hc_{hf_key}_base', to_mcore)
             has_hyper_connection = hyper_connection is not None
             has_hyper_connection = self._reduce_tensor_pp_group(has_hyper_connection, to_mcore)
-            if has_hyper_connection:
+            # ``alpha_*`` are frozen base parameters written outside ``_set_state_dict``, so they
+            # need the peft guard the mapping_proj/bias calls above get for free -- otherwise a
+            # LoRA export writes base weights into ``adapter_model.safetensors`` and a LoRA load
+            # demands a key the adapter does not carry. Same shape as the Engram export guard.
+            if has_hyper_connection and not self._peft_format:
                 if to_mcore:
                     alpha = hf_state_dict[f'hc_{hf_key}_scale'].load()
                     for i, alpha_suffix in enumerate(['pre', 'post', 'res']):
@@ -1758,6 +1858,10 @@ class GPTBridge:
         self._set_state_dict(lm_model, 'decoder.final_layernorm.weight', hf_state_dict, self.hf_final_layernorm_key,
                              to_mcore)
 
+    def _convert_additional_layers(self, mg_model, hf_state_dict, hf_prefix, to_mcore, is_pp_last_stage):
+        """Extension point for model-specific auxiliary stacks outside standard MTP."""
+        return ()
+
     def _convert_hf_state_dict(self, hf_state_dict, to_mcore):
         res = {}
         for k, v in hf_state_dict.items():
@@ -1771,6 +1875,7 @@ class GPTBridge:
         return res
 
     def _convert(self, mg_models, hf_state_dict, hf_prefix: str, to_mcore: bool, tqdm_desc: str = 'Converting: '):
+        self._pending_export_iter = None
         if to_mcore:
             hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix)
             hf_state_dict = self._convert_hf_state_dict(hf_state_dict, to_mcore)
@@ -1820,7 +1925,11 @@ class GPTBridge:
                 yield
             else:
                 res = self._convert_hf_state_dict(res, to_mcore)
-                yield from list(self._add_prefix(res, hf_prefix).items())
+                # Drain any staged PLE table shards first: they are produced one
+                # at a time (see _iter_ple_table_export) so the 100GB-scale table
+                # never accumulates in host memory.
+                yield from self._drain_pending_export(hf_prefix)
+                yield from self._add_prefix(res, hf_prefix).items()
                 hf_state_dict = {}
 
         if (not to_mcore or is_pp_last_stage) and self.config.mtp_num_layers:
@@ -1828,8 +1937,14 @@ class GPTBridge:
             if to_mcore and self.pp_rank > 0:
                 self._set_state_dict(lm_model, 'embedding.word_embeddings.weight', hf_state_dict, self.hf_embed_key,
                                      to_mcore)
+            from mcore_bridge.model.mtp import get_physical_mtp_count
+            if not to_mcore:
+                phys_count = self._reduce_mtp_physical_count(lm_model)
+            else:
+                # Import: only the owning (last) stage reaches here.
+                phys_count = get_physical_mtp_count(lm_model)
             layer_idx = 0
-            while layer_idx < self.config.mtp_num_layers:
+            while layer_idx < phys_count:
                 res = self._convert_mtp_layer(lm_model, hf_state_dict, f'{self.hf_mtp_prefix}.', layer_idx, to_mcore)
                 layer_idx += 1
                 if to_mcore:
@@ -1838,6 +1953,7 @@ class GPTBridge:
                     res = self._convert_hf_state_dict(res, to_mcore)
                     yield from list(self._add_prefix(res, hf_prefix).items())
                     hf_state_dict = {}
+        yield from self._convert_additional_layers(mg_model, hf_state_dict, hf_prefix, to_mcore, is_pp_last_stage)
         if not to_mcore or is_pp_last_stage:
             hf_state_dict.update(self._convert_post_process(mg_model, hf_state_dict, '', to_mcore))
         if to_mcore:
@@ -1847,6 +1963,16 @@ class GPTBridge:
             yield from list(self._add_prefix(hf_state_dict, hf_prefix).items())
         prog_bar.close()
 
+    def _drain_pending_export(self, hf_prefix: str):
+        """Yield (and release) the PLE table shards staged by _set_layer_ple."""
+        it, self._pending_export_iter = self._pending_export_iter, None
+        if it is None:
+            return
+        for k, v in it:
+            if v is None:
+                continue
+            yield from self._add_prefix(self._convert_hf_state_dict({k: v}, False), hf_prefix).items()
+
     def _convert_mtp_extra(self, mtp_layer, hf_state_dict, to_mcore, origin_hf_state_dict):
         for key in ['enorm.weight', 'hnorm.weight', 'eh_proj.weight']:
             self._set_state_dict(mtp_layer, key, hf_state_dict, key, to_mcore)
@@ -1854,6 +1980,17 @@ class GPTBridge:
         self._set_state_dict(mtp_layer, 'final_layernorm.weight', hf_state_dict, self.hf_mtp_final_layernorm_key,
                              to_mcore)
 
+    def _reduce_mtp_physical_count(self, lm_model) -> int:
+        """PP export: obtain the physical MTP layer count from the owning stage
+        so every rank iterates the same number of times (collective safety)."""
+        from mcore_bridge.model.mtp import get_physical_mtp_count
+        local = get_physical_mtp_count(lm_model)
+        if self.pp_size <= 1:
+            return local
+        t = torch.tensor([local], dtype=torch.long, device='cuda')
+        dist.all_reduce(t, op=dist.ReduceOp.MAX, group=self.pp_group)
+        return int(t.item())
+
     def _convert_mtp_layer(self, lm_model, hf_state_dict, hf_prefix: str, layer_idx: int, to_mcore: bool):
         mtp_layer = lm_model.mtp.layers[layer_idx] if hasattr(lm_model, 'mtp') else None
         if self.hf_mtp_prefix == self.hf_layers_prefix:
@@ -1865,8 +2002,16 @@ class GPTBridge:
             origin_hf_state_dict = hf_state_dict
             hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix)
             if len(hf_state_dict) == 0:
-                logger.info(f'MTP Layer {mtp_layer.layer_number} safetensors weights not found, '
-                            'this part will be randomly initialized.')
+                if self._peft_format:
+                    return {}
+                expected = f'{hf_prefix}*'
+                if not self.config.mtp_allow_random_init:
+                    raise ValueError(f'MTP weights not found under {expected} (expected '
+                                     f'{self.config.mtp_num_layers} physical set(s) after the {self.config.num_layers} '
+                                     'decoder layers). Set mtp_shared_weights=True to reuse one set across depths, or '
+                                     'mtp_allow_random_init=True to train MTP from random init.')
+                logger.warning(f'MTP Layer {mtp_layer.layer_number} weights not found under {expected}; '
+                               'randomly initializing (mtp_allow_random_init=True).')
                 for param in mtp_layer.parameters():
                     if param.ndim == 2:
                         mtp_layer.config.init_method(param.data)
@@ -1875,7 +2020,7 @@ class GPTBridge:
             origin_hf_state_dict = {}
             hf_state_dict = {}
         self._convert_mtp_extra(mtp_layer, hf_state_dict, to_mcore, origin_hf_state_dict)
-        transformer_layer = None if mtp_layer is None else mtp_layer.transformer_layer
+        transformer_layer = None if mtp_layer is None else mtp_layer.mtp_model_layer
         self._convert_mtp_embeds(lm_model, hf_state_dict, to_mcore)
         hf_state_dict.update(self._set_layer_attn(transformer_layer, hf_state_dict, -1, to_mcore))
         hf_state_dict.update(self._set_layer_mlp(transformer_layer, hf_state_dict, -1, to_mcore, is_mtp=True))
@@ -1948,6 +2093,7 @@ class GPTBridge:
         tqdm_desc: str = 'Exporting: ',
         disable_tqdm: bool = True,
         _is_saving: bool = False,
+        skip_unsupported_export: bool = False,
     ):
         """Export Megatron model weights to safetensors (HuggingFace) format as a generator.
 
@@ -1965,6 +2111,11 @@ class GPTBridge:
             converter: Used to perform key-value conversion on the newly exported state_dict.
             tqdm_desc: Description text for the progress bar. Defaults to 'Exporting: '.
             disable_tqdm: Whether to disable the tqdm progress bar. Defaults to True.
+            skip_unsupported_export: When True, weights whose Megatron->HF export is not implemented
+                (e.g. DeepSeek-V4.1 Engram tables, which are frozen during on-policy RL and already
+                loaded in the rollout engine) are silently skipped instead of raising. Used by the RL
+                weight-sync path; the checkpoint-save path keeps the default (False) so a saved HF
+                checkpoint stays complete.
 
         Yields:
             Tuple[str, torch.Tensor]: Key-value pairs of parameter names and tensors.
@@ -1975,6 +2126,7 @@ class GPTBridge:
         self._adapter_name = adapter_name
         self._disable_tqdm = disable_tqdm
         self._is_saving = _is_saving
+        self._skip_unsupported_export = skip_unsupported_export
         self._peft_target_modules = set()
         self._peft_modules_to_save = set()
         self._fp8_skip_modules = set()
@@ -2045,6 +2197,10 @@ class GPTBridge:
         saver.finalize()
         dist.barrier()  # Ensure all weights are saved completely
 
+    def _normalize_missing_weight_key(self, key: str) -> str:
+        """Return the identity used to detect aliases while restoring source-only weights."""
+        return key
+
     def _save_missing_weights(self, saver, saved_keys, source_model_dir=None) -> None:
         """Copy tensors present in the source checkpoint but absent from the exported ones.
 
@@ -2061,7 +2217,9 @@ class GPTBridge:
             return
         with SafetensorLazyLoader(source_model_dir) as loader:
             state_dict = loader.get_state_dict()
-            missing_keys = sorted(set(state_dict.keys()) - saved_keys)
+            saved_identities = {self._normalize_missing_weight_key(key) for key in saved_keys}
+            missing_keys = sorted(
+                key for key in state_dict if self._normalize_missing_weight_key(key) not in saved_identities)
             if not missing_keys:
                 return
             logger.info(f'Restoring {len(missing_keys)} weights from the source checkpoint '

{mcore_bridge-1.6.4 → mcore_bridge-1.8.0.dev0}/src/mcore_bridge/config/__init__.py

@@ -1,3 +1,3 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
-from .model_config import ModelConfig
+from .model_config import MLAModelConfig, ModelConfig
 from .parser import hf_to_mcore_config