初始化项目，由ModelHub XC社区提供模型

Model: ayh015/myLightningOPD
Source: Original Platform

slime_plugins/mbridge/__init__.py

@@ -0,0 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from .glm4 import GLM4Bridge
+from .glm4moe import GLM4MoEBridge
+from .mimo import MimoBridge
+from .qwen3_next import Qwen3NextBridge
+
+__all__ = ["GLM4Bridge", "GLM4MoEBridge", "Qwen3NextBridge", "MimoBridge"]

slime_plugins/mbridge/glm4.py

@@ -0,0 +1,112 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
+
+from mbridge.core import LLMBridge, register_model
+
+
+@register_model("glm4")
+class GLM4Bridge(LLMBridge):
+    """
+    Bridge implementation for Qwen2 models.
+
+    This class extends LLMBridge to provide specific configurations and
+    optimizations for Qwen2 models, handling the conversion between
+    Hugging Face Qwen2 format and Megatron-Core.
+    """
+
+    _DIRECT_MAPPING = {
+        "embedding.word_embeddings.weight": "model.embed_tokens.weight",
+        "decoder.final_layernorm.weight": "model.norm.weight",
+        "output_layer.weight": "lm_head.weight",
+    }
+    _ATTENTION_MAPPING = {
+        "self_attention.linear_proj.weight": ["model.layers.{layer_number}.self_attn.o_proj.weight"],
+        "self_attention.linear_qkv.layer_norm_weight": ["model.layers.{layer_number}.input_layernorm.weight"],
+        "self_attention.q_layernorm.weight": ["model.layers.{layer_number}.self_attn.q_norm.weight"],
+        "self_attention.k_layernorm.weight": ["model.layers.{layer_number}.self_attn.k_norm.weight"],
+        "self_attention.linear_qkv.weight": [
+            "model.layers.{layer_number}.self_attn.q_proj.weight",
+            "model.layers.{layer_number}.self_attn.k_proj.weight",
+            "model.layers.{layer_number}.self_attn.v_proj.weight",
+        ],
+        "self_attention.linear_qkv.bias": [
+            "model.layers.{layer_number}.self_attn.q_proj.bias",
+            "model.layers.{layer_number}.self_attn.k_proj.bias",
+            "model.layers.{layer_number}.self_attn.v_proj.bias",
+        ],
+    }
+    _MLP_MAPPING = {
+        "mlp.linear_fc1.weight": [
+            "model.layers.{layer_number}.mlp.gate_up_proj.weight",
+        ],
+        "mlp.linear_fc1.layer_norm_weight": ["model.layers.{layer_number}.post_attention_layernorm.weight"],
+        "mlp.linear_fc2.weight": ["model.layers.{layer_number}.mlp.down_proj.weight"],
+    }
+
+    def _build_config(self):
+        """
+        Build the configuration for Qwen2 models.
+
+        Configures Qwen2-specific parameters such as QKV bias settings and
+        layer normalization options.
+
+        Returns:
+            TransformerConfig: Configuration object for Qwen2 models
+        """
+        return self._build_base_config(
+            # qwen2
+            add_qkv_bias=True,
+            qk_layernorm=False,
+            post_mlp_layernorm=True,
+            post_self_attn_layernorm=True,
+            rotary_interleaved=True,
+        )
+
+    def _get_transformer_layer_spec(self):
+        """
+        Gets the transformer layer specification.
+
+        Creates and returns a specification for the transformer layers based on
+        the current configuration.
+
+        Returns:
+            TransformerLayerSpec: Specification for transformer layers
+
+        Raises:
+            AssertionError: If normalization is not RMSNorm
+        """
+        transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(
+            post_self_attn_layernorm=True,
+            post_mlp_layernorm=True,
+        )
+        return transformer_layer_spec
+
+    def _weight_name_mapping_mcore_to_hf(self, mcore_weights_name: str) -> list[str]:
+        """
+        Map MCore weight names to Hugging Face weight names.
+
+        Args:
+            mcore_weights_name: MCore weight name
+
+        Returns:
+            list: Corresponding Hugging Face weight names
+        """
+        assert "_extra_state" not in mcore_weights_name, "extra_state should not be loaded"
+
+        if mcore_weights_name in self._DIRECT_MAPPING:
+            return [self._DIRECT_MAPPING[mcore_weights_name]]
+
+        if "post_self_attn_layernorm" in mcore_weights_name:
+            layer_number = mcore_weights_name.split(".")[2]
+            return [f"model.layers.{layer_number}.post_self_attn_layernorm.weight"]
+        elif "post_mlp_layernorm" in mcore_weights_name:
+            layer_number = mcore_weights_name.split(".")[2]
+            return [f"model.layers.{layer_number}.post_mlp_layernorm.weight"]
+        elif "self_attention" in mcore_weights_name:
+            return self._weight_name_mapping_attention(mcore_weights_name)
+        elif "mlp" in mcore_weights_name:
+            return self._weight_name_mapping_mlp(mcore_weights_name)
+        else:
+            raise NotImplementedError(f"Unsupported parameter name: {mcore_weights_name}")

slime_plugins/mbridge/glm4moe.py

@@ -0,0 +1,125 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+
+from mbridge.core import register_model
+from mbridge.models import Qwen2Bridge, Qwen2MoEBridge
+
+
+@register_model("glm4_moe")
+class GLM4MoEBridge(Qwen2MoEBridge):
+    """
+    Bridge implementation for Qwen2 models.
+
+    This class extends LLMBridge to provide specific configurations and
+    optimizations for Qwen2 models, handling the conversion between
+    Hugging Face Qwen2 format and Megatron-Core.
+    """
+
+    _MLP_MAPPING = {
+        **(Qwen2MoEBridge._MLP_MAPPING),
+        **(Qwen2Bridge._MLP_MAPPING),
+        "mlp.router.expert_bias": ["model.layers.{layer_number}.mlp.gate.e_score_correction_bias"],
+        "shared_experts.linear_fc1.weight": [
+            "model.layers.{layer_number}.mlp.shared_experts.gate_proj.weight",
+            "model.layers.{layer_number}.mlp.shared_experts.up_proj.weight",
+        ],
+        "shared_experts.linear_fc2.weight": ["model.layers.{layer_number}.mlp.shared_experts.down_proj.weight"],
+    }
+
+    _MTP_MAPPING = {
+        "enorm.weight": ["model.layers.{layer_number}.enorm.weight"],
+        "hnorm.weight": ["model.layers.{layer_number}.hnorm.weight"],
+        "eh_proj.weight": ["model.layers.{layer_number}.eh_proj.weight"],
+        "final_layernorm.weight": ["model.layers.{layer_number}.shared_head.norm.weight"],
+    }
+
+    def _weight_name_mapping_mtp(self, name: str, num_layers: int) -> str:
+        convert_names = []
+        for keyword, mapping_names in self._MTP_MAPPING.items():
+            if keyword in name:
+                convert_names.extend([x.format(layer_number=num_layers) for x in mapping_names])
+                break
+            elif "mlp" in name:
+                mtp_layer_index = int(re.findall(r"mtp\.layers\.(\d+)\.", name)[0])
+                name_ = re.sub(
+                    r"^mtp\.layers.\d+.transformer_layer", f"model.layers.{num_layers+mtp_layer_index}", name
+                )
+                convert_names = self._weight_name_mapping_mlp(name_)
+                break
+            elif "self_attention" in name:
+                mtp_layer_index = int(re.findall(r"mtp\.layers.(\d+)\.", name)[0])
+                name_ = re.sub(
+                    r"^mtp\.layers.\d+.transformer_layer", f"model.layers.{num_layers+mtp_layer_index}", name
+                )
+                convert_names = self._weight_name_mapping_attention(name_)
+                break
+
+        if len(convert_names) == 0:
+            raise NotImplementedError(f"Unsupported parameter name: {name}")
+        return convert_names
+
+    def _weight_name_mapping_mcore_to_hf(self, mcore_weights_name: str) -> list[str]:
+        """
+        Map MCore weight names to Hugging Face weight names.
+
+        Args:
+            mcore_weights_name: MCore weight name
+
+        Returns:
+            list: Corresponding Hugging Face weight names
+        """
+        assert "_extra_state" not in mcore_weights_name, "extra_state should not be loaded"
+        direct_name_mapping = {
+            "embedding.word_embeddings.weight": "model.embed_tokens.weight",
+            "decoder.final_layernorm.weight": "model.norm.weight",
+            "output_layer.weight": "lm_head.weight",
+        }
+        if mcore_weights_name in direct_name_mapping:
+            return [direct_name_mapping[mcore_weights_name]]
+
+        if "mtp" in mcore_weights_name:  # first check mtp
+            return self._weight_name_mapping_mtp(mcore_weights_name, self.config.num_layers)
+        elif "self_attention" in mcore_weights_name:
+            return self._weight_name_mapping_attention(mcore_weights_name)
+        elif "mlp" in mcore_weights_name:
+            return self._weight_name_mapping_mlp(mcore_weights_name)
+        else:
+            raise NotImplementedError(f"Unsupported parameter name: {mcore_weights_name}")
+
+    def _build_config(self):
+        """
+        Build the configuration for Qwen2 models.
+
+        Configures Qwen2-specific parameters such as QKV bias settings and
+        layer normalization options.
+
+        Returns:
+            TransformerConfig: Configuration object for Qwen2 models
+        """
+        return self._build_base_config(
+            use_cpu_initialization=False,
+            # MoE specific
+            moe_ffn_hidden_size=self.hf_config.moe_intermediate_size,
+            moe_router_bias_update_rate=0.001,
+            moe_router_topk=self.hf_config.num_experts_per_tok,
+            num_moe_experts=self.hf_config.n_routed_experts,
+            # moe_router_load_balancing_type="aux_loss",
+            moe_router_load_balancing_type="none",  # default None for RL
+            moe_grouped_gemm=True,
+            moe_router_score_function="sigmoid",
+            moe_router_enable_expert_bias=True,
+            moe_router_pre_softmax=True,
+            # Other optimizations
+            persist_layer_norm=True,
+            bias_activation_fusion=True,
+            bias_dropout_fusion=True,
+            # GLM specific
+            qk_layernorm=self.hf_config.use_qk_norm,
+            add_qkv_bias=True,
+            add_bias_linear=False,
+            # post_mlp_layernorm=True,
+            # post_self_attn_layernorm=True,
+            rotary_interleaved=True,
+        )

slime_plugins/mbridge/mimo.py

@@ -0,0 +1,123 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import torch
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_mtp_block_spec
+
+from mbridge.core import register_model
+from mbridge.models import Qwen2Bridge
+
+
+@register_model("mimo")
+class MimoBridge(Qwen2Bridge):
+    """
+    Bridge implementation for Mimo models.
+
+    This class extends Qwen2Bridge to provide specific configurations and
+    optimizations for Mimo models, handling the conversion between
+    Hugging Face Mimo format and Megatron-Core.
+
+    MiMo adds MTP (Multi-Token Prediction) layers on top of Qwen2 architecture.
+    """
+
+    def _build_config(self):
+        """Override to add MTP configuration."""
+        hf_config = self.hf_config
+
+        # Add MTP configuration if present
+        mtp_args = {}
+        if "num_nextn_predict_layers" in hf_config:
+            mtp_args["mtp_num_layers"] = hf_config.num_nextn_predict_layers
+
+        return self._build_base_config(
+            add_qkv_bias=True,
+            qk_layernorm=False,
+            **mtp_args,
+        )
+
+    def _get_gptmodel_args(self) -> dict:
+        """Override to add MTP block spec if needed."""
+        ret = super()._get_gptmodel_args()
+
+        # Add MTP block spec if MTP layers are present
+        if self.config.mtp_num_layers is not None:
+            transformer_layer_spec = self.config
+            mtp_block_spec = get_gpt_mtp_block_spec(self.config, transformer_layer_spec, use_transformer_engine=True)
+            ret["mtp_block_spec"] = mtp_block_spec
+
+        return ret
+
+    def _weight_name_mapping_mcore_to_hf(self, mcore_weights_name: str) -> list[str]:
+        """Override to handle MTP layer mappings."""
+        # Check if this is an MTP layer weight
+        if "mtp" in mcore_weights_name:
+            return self._convert_mtp_param(mcore_weights_name)
+
+        # Otherwise use parent class mapping
+        return super()._weight_name_mapping_mcore_to_hf(mcore_weights_name)
+
+    def _convert_mtp_param(self, name: str) -> list[str]:
+        """Convert MTP layer parameters from MCore to HF format."""
+        # For now, assume single MTP layer support
+        if "mtp.layers." not in name:
+            raise NotImplementedError(f"Invalid MTP parameter name: {name}")
+
+        # Get the MTP layer index
+        parts = name.split(".")
+        mtp_layer_idx = parts[2]  # mtp.layers.{idx}
+
+        # Direct mappings for MTP-specific components
+        direct_name_mapping = {
+            f"mtp.layers.{mtp_layer_idx}.enorm.weight": f"model.mtp_layers.{mtp_layer_idx}.token_layernorm.weight",
+            f"mtp.layers.{mtp_layer_idx}.hnorm.weight": f"model.mtp_layers.{mtp_layer_idx}.hidden_layernorm.weight",
+            f"mtp.layers.{mtp_layer_idx}.eh_proj.weight": f"model.mtp_layers.{mtp_layer_idx}.input_proj.weight",
+            f"mtp.layers.{mtp_layer_idx}.final_layernorm.weight": f"model.mtp_layers.{mtp_layer_idx}.final_layernorm.weight",
+        }
+
+        if name in direct_name_mapping:
+            return [direct_name_mapping[name]]
+
+        # Handle transformer components within MTP
+        # Check if this is a transformer_layer component
+        if "transformer_layer" in name:
+            # Create a proxy name to use with parent class methods
+            # Convert mtp.layers.{idx}.transformer_layer.* to decoder.layers.{idx}.*
+            proxy_name = name.replace(
+                f"mtp.layers.{mtp_layer_idx}.transformer_layer",
+                f"decoder.layers.{mtp_layer_idx}",
+            )
+
+            if "self_attention" in proxy_name or "input_layernorm.weight" in proxy_name:
+                convert_names = super()._weight_name_mapping_attention(proxy_name)
+            elif "mlp" in proxy_name:
+                convert_names = super()._weight_name_mapping_mlp(proxy_name)
+            else:
+                raise NotImplementedError(f"Unsupported transformer component in MTP: {name}")
+
+            # Replace the layer index in converted names to point to mtp_layers
+            convert_names = [
+                cn.replace(f"model.layers.{mtp_layer_idx}", f"model.mtp_layers.{mtp_layer_idx}")
+                for cn in convert_names
+            ]
+            return convert_names
+        else:
+            raise NotImplementedError(f"Unsupported MTP parameter name: {name}")
+        return convert_names
+
+    def _weight_to_mcore_format(self, mcore_weights_name: str, hf_weights: list[torch.Tensor]) -> torch.Tensor:
+        """Swap halves of eh_proj weights before handing off to Megatron-Core."""
+        weight = super()._weight_to_mcore_format(mcore_weights_name, hf_weights)
+        if mcore_weights_name.endswith("eh_proj.weight"):
+            first_half, second_half = weight.chunk(2, dim=1)
+            weight = torch.cat([second_half, first_half], dim=1)
+        return weight
+
+    def _weight_to_hf_format(
+        self, mcore_weights_name: str, mcore_weights: torch.Tensor
+    ) -> tuple[list[str], list[torch.Tensor]]:
+        """Swap halves back when exporting eh_proj weights to HuggingFace format."""
+        if mcore_weights_name.endswith("eh_proj.weight"):
+            first_half, second_half = mcore_weights.chunk(2, dim=1)
+            mcore_weights = torch.cat([second_half, first_half], dim=1)
+        return super()._weight_to_hf_format(mcore_weights_name, mcore_weights)

slime_plugins/mbridge/qwen3_next.py

@@ -0,0 +1,104 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+from mbridge.core import register_model
+from mbridge.models import Qwen2MoEBridge
+
+
+@register_model("qwen3_next")
+class Qwen3NextBridge(Qwen2MoEBridge):
+    _ATTENTION_MAPPING = (
+        Qwen2MoEBridge._ATTENTION_MAPPING
+        | {
+            f"self_attention.{weight_name}": ["model.layers.{layer_number}." + weight_name]
+            for weight_name in [
+                "input_layernorm.weight",
+                # linear attn
+                "linear_attn.A_log",
+                "linear_attn.conv1d.weight",
+                "linear_attn.dt_bias",
+                "linear_attn.in_proj_ba.weight",
+                "linear_attn.in_proj_qkvz.weight",
+                "linear_attn.norm.weight",
+                "linear_attn.out_proj.weight",
+                # gated attn
+                "self_attn.k_norm.weight",
+                "self_attn.k_proj.weight",
+                "self_attn.o_proj.weight",
+                "self_attn.q_norm.weight",
+                "self_attn.q_proj.weight",
+                "self_attn.v_proj.weight",
+            ]
+        }
+        | {
+            "self_attention.linear_qkv.layer_norm_weight": ["model.layers.{layer_number}.input_layernorm.weight"],
+            "self_attention.linear_qkv.weight": [
+                "model.layers.{layer_number}.self_attn.q_proj.weight",
+                "model.layers.{layer_number}.self_attn.k_proj.weight",
+                "model.layers.{layer_number}.self_attn.v_proj.weight",
+            ],
+        }
+    )
+
+    def _weight_to_mcore_format(
+        self, mcore_weights_name: str, hf_weights: list[torch.Tensor]
+    ) -> tuple[list[str], list[torch.Tensor]]:
+        if "self_attention.linear_qkv." in mcore_weights_name and "layer_norm" not in mcore_weights_name:
+            # merge qkv
+            assert len(hf_weights) == 3
+            num_key_value_heads = self.hf_config.num_key_value_heads
+            hidden_dim = self.hf_config.hidden_size
+            num_attention_heads = self.hf_config.num_attention_heads
+            num_querys_per_group = num_attention_heads // self.hf_config.num_key_value_heads
+            head_dim = getattr(self.hf_config, "head_dim", hidden_dim // num_attention_heads)
+            group_dim = head_dim * num_attention_heads // num_key_value_heads
+            q, k, v = hf_weights
+            # q k v might be tp split
+            real_num_key_value_heads = q.shape[0] // (2 * group_dim)
+            q = (
+                q.view(
+                    [
+                        real_num_key_value_heads,
+                        num_querys_per_group,
+                        2,
+                        head_dim,
+                        -1,
+                    ]
+                )
+                .transpose(1, 2)
+                .flatten(1, 3)
+            )
+            k = k.view([real_num_key_value_heads, head_dim, -1])
+            v = v.view([real_num_key_value_heads, head_dim, -1])
+            out_shape = [-1, hidden_dim] if ".bias" not in mcore_weights_name else [-1]
+
+            qgkv = torch.cat([q, k, v], dim=1).view(*out_shape).contiguous()
+            return qgkv
+
+        return super()._weight_to_mcore_format(mcore_weights_name, hf_weights)
+
+    def _build_config(self):
+        return self._build_base_config(
+            use_cpu_initialization=False,
+            # MoE specific
+            moe_ffn_hidden_size=self.hf_config.moe_intermediate_size,
+            moe_router_bias_update_rate=0.001,
+            moe_router_topk=self.hf_config.num_experts_per_tok,
+            num_moe_experts=self.hf_config.num_experts,
+            moe_aux_loss_coeff=self.hf_config.router_aux_loss_coef,
+            # moe_router_load_balancing_type="aux_loss",
+            moe_router_load_balancing_type="none",  # default None for RL
+            moe_grouped_gemm=True,
+            moe_router_score_function="softmax",
+            # Other optimizations
+            persist_layer_norm=True,
+            bias_activation_fusion=True,
+            bias_dropout_fusion=True,
+            # Qwen specific
+            moe_router_pre_softmax=False,
+            qk_layernorm=True,
+            # Qwen3 Next specific
+            attention_output_gate=True,
+            moe_shared_expert_gate=True,
+        )