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

Model: ayh015/myLightningOPD
Source: Original Platform

slime_plugins/models/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+

slime_plugins/models/glm4.py

@@ -0,0 +1,17 @@
+# 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
+
+
+def get_glm_spec(args, config, vp_stage):
+    transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(
+        num_experts=args.num_experts,
+        moe_grouped_gemm=args.moe_grouped_gemm,
+        qk_layernorm=args.qk_layernorm,
+        multi_latent_attention=args.multi_latent_attention,
+        moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm,
+        post_self_attn_layernorm=args.post_self_attn_layernorm,
+        post_mlp_layernorm=args.post_mlp_layernorm,
+    )
+    return transformer_layer_spec

slime_plugins/models/hf_attention.py

@@ -0,0 +1,118 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC, abstractmethod
+
+import torch
+import torch.distributed as dist
+from megatron.core import mpu, tensor_parallel
+from megatron.core.inference.contexts import BaseInferenceContext
+from megatron.core.packed_seq_params import PackedSeqParams
+from megatron.core.transformer.module import MegatronModule
+from transformers import AutoConfig
+
+
+class HuggingfaceAttention(MegatronModule, ABC):
+    """Attention layer abstract class.
+
+    This layer only contains common modules required for the "self attn" and
+    "cross attn" specializations.
+    """
+
+    def __init__(
+        self,
+        args,
+        config,
+        layer_number: int,
+        cp_comm_type: str = "p2p",
+        pg_collection=None,
+    ):
+        super().__init__(config=config)
+        self.args = args
+        self.config = config
+        # Note that megatron layer_number starts at 1
+        self.layer_number = layer_number
+        self.hf_layer_idx = layer_number - 1
+        self.hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True)
+        # hardcode to fa2 at the moment.
+        self.hf_config._attn_implementation = "flash_attention_2"
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        key_value_states: torch.Tensor | None = None,
+        inference_context: BaseInferenceContext | None = None,
+        rotary_pos_emb: torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None = None,
+        rotary_pos_cos: torch.Tensor | None = None,
+        rotary_pos_sin: torch.Tensor | None = None,
+        rotary_pos_cos_sin: torch.Tensor | None = None,
+        attention_bias: torch.Tensor | None = None,
+        packed_seq_params: PackedSeqParams | None = None,
+        sequence_len_offset: int | None = None,
+        *,
+        inference_params: BaseInferenceContext | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        assert packed_seq_params is not None
+        cu_seqlens = packed_seq_params.cu_seqlens_q
+
+        if self.args.sequence_parallel:
+            hidden_states = tensor_parallel.gather_from_sequence_parallel_region(
+                hidden_states, group=mpu.get_tensor_model_parallel_group()
+            )
+
+        if mpu.get_context_parallel_world_size() > 1:
+            cp_size = mpu.get_context_parallel_world_size()
+            hidden_states_list = dist.nn.all_gather(
+                hidden_states,
+                group=mpu.get_context_parallel_group(),
+            )
+
+            # TODO: preprocess this for each batch to prevent tolist in the training step
+            whole_hidden_states_list = []
+
+            local_cu_seqlens = cu_seqlens // cp_size
+            for i in range(len(cu_seqlens) - 1):
+                seqlen = cu_seqlens[i + 1] - cu_seqlens[i]
+                chunk_size = seqlen // 2 // cp_size
+                whole_hidden_states_list.extend(
+                    [
+                        hidden_states_list[cp_rank][local_cu_seqlens[i] : local_cu_seqlens[i] + chunk_size]
+                        for cp_rank in range(cp_size)
+                    ]
+                    + [
+                        hidden_states_list[cp_rank][local_cu_seqlens[i] + chunk_size : local_cu_seqlens[i + 1]]
+                        for cp_rank in range(cp_size)
+                    ][::-1],
+                )
+            hidden_states = torch.cat(whole_hidden_states_list, dim=0)
+
+        hidden_states = hidden_states.permute(1, 0, 2)  # [bsz, seq_len, hidden_dim]
+
+        output = self.hf_forward(hidden_states, packed_seq_params)
+        bias = None
+
+        output = output.permute(1, 0, 2)  # [seq_len, bsz, hidden_dim]
+
+        if mpu.get_context_parallel_world_size() > 1:
+            cp_rank = mpu.get_context_parallel_rank()
+            output_list = []
+            for i in range(len(cu_seqlens) - 1):
+                seqlen = cu_seqlens[i + 1] - cu_seqlens[i]
+                chunk_size = seqlen // 2 // cp_size
+                seq = output[cu_seqlens[i] : cu_seqlens[i + 1]]
+                chunks = torch.chunk(seq, 2 * cp_size, dim=0)
+                output_list.append(chunks[cp_rank])
+                output_list.append(chunks[2 * cp_size - 1 - cp_rank])
+            output = torch.cat(output_list, dim=0)
+
+        if self.args.sequence_parallel:
+            output = tensor_parallel.scatter_to_sequence_parallel_region(
+                output, group=mpu.get_tensor_model_parallel_group()
+            )
+
+        return output, bias
+
+    @abstractmethod
+    def hf_forward(self, hidden_states, packed_seq_params):
+        """Huggingface forward function"""

slime_plugins/models/qwen3_next.py

@@ -0,0 +1,229 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import copy
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_block import get_num_layers_to_build
+from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
+from transformers import AutoConfig
+from transformers.activations import ACT2FN
+
+try:
+    from fla.modules import FusedRMSNormGated, ShortConvolution
+    from fla.ops.gated_delta_rule import chunk_gated_delta_rule
+    from transformers.models.qwen3_next.modeling_qwen3_next import Qwen3NextAttention, Qwen3NextRMSNorm
+except ImportError:
+    pass
+
+from .hf_attention import HuggingfaceAttention
+
+
+# adapt from https://github.com/huggingface/transformers/blob/38a08b6e8ae35857109cedad75377997fecbf9d0/src/transformers/models/qwen3_next/modeling_qwen3_next.py#L564
+class Qwen3NextGatedDeltaNet(nn.Module):
+    """
+    Qwen3NextGatedDeltaNet with varlen support
+    """
+
+    def __init__(self, config, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.num_v_heads = config.linear_num_value_heads
+        self.num_k_heads = config.linear_num_key_heads
+        self.head_k_dim = config.linear_key_head_dim
+        self.head_v_dim = config.linear_value_head_dim
+        self.key_dim = self.head_k_dim * self.num_k_heads
+        self.value_dim = self.head_v_dim * self.num_v_heads
+
+        self.conv_kernel_size = config.linear_conv_kernel_dim
+        self.layer_idx = layer_idx
+        self.activation = config.hidden_act
+        self.act = ACT2FN[config.hidden_act]
+        self.layer_norm_epsilon = config.rms_norm_eps
+
+        # QKV
+        self.conv_dim = self.key_dim * 2 + self.value_dim
+        self.conv1d = ShortConvolution(
+            hidden_size=self.conv_dim,
+            bias=False,
+            kernel_size=self.conv_kernel_size,
+        )
+
+        # projection of the input hidden states
+        projection_size_qkvz = self.key_dim * 2 + self.value_dim * 2
+        projection_size_ba = self.num_v_heads * 2
+        self.in_proj_qkvz = nn.Linear(self.hidden_size, projection_size_qkvz, bias=False)
+        self.in_proj_ba = nn.Linear(self.hidden_size, projection_size_ba, bias=False)
+
+        # time step projection (discretization)
+        # instantiate once and copy inv_dt in init_weights of PretrainedModel
+        self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads))
+
+        A = torch.empty(self.num_v_heads).uniform_(0, 16)
+        self.A_log = nn.Parameter(torch.log(A))
+
+        self.norm = FusedRMSNormGated(
+            self.head_v_dim,
+            eps=self.layer_norm_epsilon,
+            activation=self.activation,
+            device=torch.cuda.current_device(),
+            dtype=config.dtype if config.dtype is not None else torch.get_current_dtype(),
+        )
+
+        self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
+
+    def fix_query_key_value_ordering(self, mixed_qkvz, mixed_ba):
+        """
+        Derives `query`, `key` and `value` tensors from `mixed_qkvz` and `mixed_ba`.
+        """
+
+        new_tensor_shape_qkvz = mixed_qkvz.size()[:-1] + (
+            self.num_k_heads,
+            2 * self.head_k_dim + 2 * self.head_v_dim * self.num_v_heads // self.num_k_heads,
+        )
+        new_tensor_shape_ba = mixed_ba.size()[:-1] + (self.num_k_heads, 2 * self.num_v_heads // self.num_k_heads)
+
+        mixed_qkvz = mixed_qkvz.view(*new_tensor_shape_qkvz)
+        mixed_ba = mixed_ba.view(*new_tensor_shape_ba)
+        split_arg_list_qkvz = [
+            self.head_k_dim,
+            self.head_k_dim,
+            (self.num_v_heads // self.num_k_heads * self.head_v_dim),
+            (self.num_v_heads // self.num_k_heads * self.head_v_dim),
+        ]
+        split_arg_list_ba = [self.num_v_heads // self.num_k_heads, self.num_v_heads // self.num_k_heads]
+        query, key, value, z = torch.split(mixed_qkvz, split_arg_list_qkvz, dim=3)
+        b, a = torch.split(mixed_ba, split_arg_list_ba, dim=3)
+        # [b, sq, ng, np/ng * hn] -> [b, sq, np, hn]
+        value = value.reshape(value.size(0), value.size(1), -1, self.head_v_dim)
+        z = z.reshape(z.size(0), z.size(1), -1, self.head_v_dim)
+        b = b.reshape(b.size(0), b.size(1), self.num_v_heads)
+        a = a.reshape(a.size(0), a.size(1), self.num_v_heads)
+        return query, key, value, z, b, a
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor = None,
+    ):
+        projected_states_qkvz = self.in_proj_qkvz(hidden_states)
+        projected_states_ba = self.in_proj_ba(hidden_states)
+        query, key, value, z, b, a = self.fix_query_key_value_ordering(projected_states_qkvz, projected_states_ba)
+        query, key, value = (x.reshape(x.shape[0], x.shape[1], -1) for x in (query, key, value))
+
+        mixed_qkv = torch.cat((query, key, value), dim=-1)
+
+        mixed_qkv, _ = self.conv1d(
+            x=mixed_qkv,
+            cu_seqlens=cu_seqlens,
+        )
+
+        query, key, value = torch.split(
+            mixed_qkv,
+            [
+                self.key_dim,
+                self.key_dim,
+                self.value_dim,
+            ],
+            dim=-1,
+        )
+        query = query.reshape(query.shape[0], query.shape[1], -1, self.head_k_dim)
+        key = key.reshape(key.shape[0], key.shape[1], -1, self.head_k_dim)
+        value = value.reshape(value.shape[0], value.shape[1], -1, self.head_v_dim)
+
+        beta = b.sigmoid()
+        # If the model is loaded in fp16, without the .float() here, A might be -inf
+        g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
+        if self.num_v_heads // self.num_k_heads > 1:
+            query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
+            key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2)
+
+        core_attn_out, last_recurrent_state = chunk_gated_delta_rule(
+            query,
+            key,
+            value,
+            g=g,
+            beta=beta,
+            initial_state=None,
+            output_final_state=False,
+            use_qk_l2norm_in_kernel=True,
+        )
+
+        z_shape_og = z.shape
+        # reshape input data into 2D tensor
+        core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
+        z = z.reshape(-1, z.shape[-1])
+        core_attn_out = self.norm(core_attn_out, z)
+        core_attn_out = core_attn_out.reshape(z_shape_og)
+        core_attn_out = core_attn_out.reshape(core_attn_out.shape[0], core_attn_out.shape[1], -1)
+
+        output = self.out_proj(core_attn_out)
+        return output
+
+
+class Attention(HuggingfaceAttention):
+    def __init__(
+        self,
+        args,
+        config,
+        layer_number: int,
+        cp_comm_type: str = "p2p",
+        pg_collection=None,
+    ):
+        super().__init__(
+            args,
+            config,
+            layer_number,
+            cp_comm_type,
+            pg_collection,
+        )
+        if Qwen3NextAttention is None:
+            raise ImportError("Please install transformers>=4.35.0 to use Qwen3NextAttention.")
+
+        self.linear_attn = Qwen3NextGatedDeltaNet(self.hf_config, self.hf_layer_idx)
+        self.input_layernorm = Qwen3NextRMSNorm(self.hf_config.hidden_size, eps=self.hf_config.rms_norm_eps)
+
+    def hf_forward(self, hidden_states, packed_seq_params):
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.linear_attn(
+            hidden_states=hidden_states,
+            cu_seqlens=packed_seq_params.cu_seqlens_q,
+        )
+        return hidden_states
+
+
+def get_qwen3_next_spec(args, config, vp_stage):
+    # always use the moe path
+    if not args.num_experts:
+        config.moe_layer_freq = [0] * config.num_layers
+
+    # Define the decoder block spec
+    kwargs = {
+        "use_transformer_engine": True,
+    }
+    if vp_stage is not None:
+        kwargs["vp_stage"] = vp_stage
+    transformer_layer_spec = get_gpt_decoder_block_spec(config, **kwargs)
+
+    assert config.pipeline_model_parallel_layout is None, "not support this at the moment"
+
+    # Slice the layer specs to only include the layers that are built in this pipeline stage.
+    # Note: MCore layer_number starts at 1
+    num_layers_to_build = get_num_layers_to_build(config, vp_stage=vp_stage)
+    offset = get_transformer_layer_offset(config, vp_stage=vp_stage)
+
+    hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True)
+
+    for layer_id in range(num_layers_to_build):
+        if hf_config.layer_types[layer_id + offset] == "linear_attention":
+            layer_specs = copy.deepcopy(transformer_layer_spec.layer_specs[layer_id])
+            layer_specs.submodules.self_attention = ModuleSpec(
+                module=Attention,
+                params={"args": args},
+            )
+            transformer_layer_spec.layer_specs[layer_id] = layer_specs
+    return transformer_layer_spec