[dev] Megatron Lite (1/4): shared primitive core

experimental/lite/megatron/lite/model/protocol_utils.py

@@ -0,0 +1,129 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Helpers shared by model protocol forward steps.
+
+The verl/runtime layers hand each protocol a raw, model-agnostic ``PackedBatch``
+(true per-sequence lengths, no padding, no ``PackedSeqParams``). Each model owns
+its pack/unpack pair: ``pack_thd_forward_kwargs`` pads + CP-splits the batch into
+model forward kwargs, and ``unpack_thd_forward_output`` reverses a model output
+back to jagged true-length form. THD models share the zigzag-CP pair below;
+models with a different CP layout (e.g. DeepSeek-V4 contiguous DSA) provide their
+own pair.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from megatron.lite.primitive.parallel import ParallelState
+from megatron.lite.primitive.parallel.thd import (
+    pack_nested_thd,
+    parallel_state_from_model,
+    prepare_packed_thd_kwargs_for_context_parallel,
+    thd_pack_meta,
+    unpack_thd_to_nested,
+)
+from megatron.lite.primitive.utils.packed_seq import PackedSeqParams
+from megatron.lite.runtime.contracts.data import PackedBatch
+from megatron.lite.runtime.contracts.loss import get_loss_context
+
+
+def _parallel_state(model) -> ParallelState:
+    return parallel_state_from_model(model) or ParallelState()
+
+
+def nested_from_packed(tensor: torch.Tensor | None, seq_lens: torch.Tensor):
+    """Split a 1-D packed (true, unpadded) tensor back into a jagged nested tensor."""
+    if tensor is None:
+        return None
+    if tensor.dim() == 2 and tensor.size(0) == 1:
+        tensor = tensor.squeeze(0)
+    if tensor.dim() != 1:
+        raise ValueError(f"PackedBatch tensor must be 1-D, got {tuple(tensor.shape)}.")
+    pieces = []
+    offset = 0
+    for length_t in seq_lens:
+        length = int(length_t.item())
+        pieces.append(tensor.narrow(0, offset, length))
+        offset += length
+    if offset != tensor.numel():
+        raise ValueError(f"PackedBatch sizes sum to {offset}, tensor has {tensor.numel()} tokens.")
+    return torch.nested.as_nested_tensor(pieces, layout=torch.jagged)
+
+
+def pack_thd_forward_kwargs(model, batch: PackedBatch) -> dict[str, Any]:
+    """Pad + zigzag-CP-split a raw THD batch into model forward kwargs.
+
+    Pads each sequence to the TE/zigzag alignment, then CP-splits tokens,
+    labels, masks and position ids through the shared THD primitive — the same
+    layout the model was validated against, now produced inside the protocol
+    rather than the connector.
+    """
+    ps = _parallel_state(model)
+    seq_lens = batch.seq_lens
+    packed = pack_nested_thd(
+        nested_from_packed(batch.input_ids, seq_lens),
+        tp_size=ps.tp_size,
+        cp_size=ps.cp_size,
+        cp_rank=ps.cp_rank,
+        cp_group=ps.cp_group if ps.cp_size > 1 else None,
+        split_cp=False,
+        labels=nested_from_packed(batch.labels, seq_lens),
+        roll_labels=batch.labels is not None,
+        loss_mask=nested_from_packed(batch.loss_mask, seq_lens),
+        roll_loss_mask=batch.loss_mask is not None,
+    )
+    max_seqlen = int(packed.padded_lengths.max().item()) if packed.padded_lengths.numel() else 0
+    # pack_nested_thd already returns [1, T] token rows; do not unsqueeze again.
+    kwargs: dict[str, Any] = {
+        "input_ids": packed.input_ids,
+        "labels": packed.labels,
+        "loss_mask": packed.loss_mask,
+        "position_ids": packed.position_ids,
+        "packed_seq_params": PackedSeqParams.from_cu_seqlens(
+            packed.cu_seqlens_padded, max_seqlen=max_seqlen
+        ),
+    }
+    prepare_packed_thd_kwargs_for_context_parallel(model, kwargs)
+    return kwargs
+
+
+def unpack_thd_forward_output(model, batch: PackedBatch, output: torch.Tensor) -> torch.Tensor:
+    """Reverse a zigzag-CP THD model output back to jagged true-length form."""
+    ps = _parallel_state(model)
+    meta = thd_pack_meta(
+        batch.seq_lens,
+        tp_size=ps.tp_size,
+        cp_size=ps.cp_size,
+        cp_group=ps.cp_group if ps.cp_size > 1 else None,
+    )
+    return unpack_thd_to_nested(output, meta, contiguous=False)
+
+
+def add_loss_context_kwargs(kwargs: dict[str, Any], *, include_return_log_probs: bool = False) -> None:
+    loss_context = get_loss_context()
+    if loss_context is None:
+        return
+    kwargs["temperature"] = loss_context.temperature
+    kwargs["calculate_entropy"] = loss_context.calculate_entropy
+    if include_return_log_probs:
+        kwargs["return_log_probs"] = loss_context.return_log_probs
+
+
+def add_cross_entropy_fusion(kwargs: dict[str, Any], model) -> None:
+    kwargs["use_fused_kernels"] = bool(getattr(model, "cross_entropy_fusion", False))
+
+
+def set_cross_entropy_fusion(chunks: list, enabled: bool) -> None:
+    for chunk in chunks:
+        chunk.cross_entropy_fusion = bool(enabled)
+
+
+__all__ = [
+    "add_cross_entropy_fusion",
+    "add_loss_context_kwargs",
+    "nested_from_packed",
+    "pack_thd_forward_kwargs",
+    "set_cross_entropy_fusion",
+    "unpack_thd_forward_output",
+]

experimental/lite/megatron/lite/model/qwen3_5/lite/checkpoint.py

@@ -35,9 +35,13 @@ def PLACEMENT_FN(param_name: str) -> list:
         return [Replicate(), Replicate(), Replicate(), Shard(0)]
     if "qkv" in param_name and "layer_norm" not in param_name:
         return [Replicate(), Replicate(), Replicate(), Shard(0)]
-    if ("proj" in param_name or "o_proj" in param_name) and (
-        "full_attn" in param_name or "linear_attn" in param_name
+    if (
+        ("proj" in param_name or "o_proj" in param_name)
+        and ("full_attn" in param_name or "linear_attn" in param_name)
+        and "layer_norm" not in param_name
     ):
+        # Row-parallel output proj weight: TP-shard on dim 1. Exclude layer_norm_weight (1-D, replicated
+        # under TP) which otherwise matches here ("in_proj" contains "proj") and gets an invalid Shard(1).
         return [Replicate(), Replicate(), Replicate(), Shard(1)]
     if "gate_up" in param_name and "shared" in param_name:
         return [Replicate(), Replicate(), Replicate(), Shard(0)]
@@ -760,10 +764,19 @@ def export_hf_weights(
     )
 
 
+def save_hf_weights(
+    model: nn.Module | list[nn.Module], path: str, config: Qwen35Config, ps: ParallelState
+) -> None:
+    from megatron.lite.primitive.ckpt.hf_weights import save_hf_weights as _save
+
+    _save(model, path, Qwen35WeightSpec(config), ps, vocab_size=config.vocab_size)
+
+
 __all__ = [
     "EXPERT_CLASSIFIER",
     "PLACEMENT_FN",
     "Qwen35WeightSpec",
     "export_hf_weights",
     "load_hf_weights",
+    "save_hf_weights",
 ]

experimental/lite/megatron/lite/model/qwen3_5/lite/model.py

@@ -2,9 +2,7 @@
 """Qwen3.5 lite native model.
 
 This implementation keeps the lightweight qwen3_moe/lite composition style
-and does not wrap Megatron-Core layer modules. It still reuses Megatron Lite
-parallel/TE primitives and small Megatron atomic RoPE helpers where those are
-already used by other native Megatron Lite modules.
+and uses Megatron Lite parallel, TE, RoPE, and MoE primitives directly.
 """
 
 from __future__ import annotations
@@ -16,15 +14,20 @@
 import torch.nn as nn
 import transformer_engine.pytorch as te
 
-from megatron.core.fusions.fused_bias_swiglu import bias_swiglu_impl
 from megatron.lite.model.qwen3_5.config import Qwen35Config
+from megatron.lite.primitive.kernels.swiglu import bias_swiglu_impl
 from megatron.lite.primitive.modules.dispatcher import TokenDispatcher
-from megatron.lite.primitive.modules.experts import Experts
+from megatron.lite.primitive.modules.experts import Experts, swiglu_with_probs
 from megatron.lite.primitive.modules.gated_delta_net import GatedDeltaNet
 from megatron.lite.primitive.modules.gqa import GQAttention as FullAttention
 from megatron.lite.primitive.modules.gqa import split_grouped_qkvg as _split_grouped_qkvg
 from megatron.lite.primitive.modules.mrope import MultimodalRotaryEmbedding as Qwen35MRoPE
-from megatron.lite.primitive.modules.mtp import MTPBlock, MTPDecoderLayer, MTPLossAutoScaler
+from megatron.lite.primitive.modules.mtp import (
+    MTPBlock,
+    MTPDecoderLayer,
+    MTPLossAutoScaler,
+    roll_mtp_tensor_left,
+)
 from megatron.lite.primitive.modules.router import TopKRouter
 from megatron.lite.primitive.ops.cross_entropy import vocab_parallel_cross_entropy
 from megatron.lite.primitive.ops.linear_cross_entropy import linear_cross_entropy
@@ -37,7 +40,6 @@
     VocabParallelOutput,
     build_pipeline_chunk_layout,
     gather_from_sequence_parallel,
-    roll_packed_thd_left,
     scatter_to_sequence_parallel,
 )
 from megatron.lite.primitive.utils import build_fp8_recipe
@@ -66,7 +68,7 @@ def _collect_sp_grad_params(model: nn.Module) -> list[nn.Parameter]:
 
 
 def _swiglu(x: torch.Tensor) -> torch.Tensor:
-    return bias_swiglu_impl(x, bias=None)
+    return swiglu_with_probs(x, probs=None)
 
 
 def _qwen_mrope_section(config: Qwen35Config) -> list[int]:
@@ -231,6 +233,7 @@ def __init__(
         moe_act_recompute: bool = False,
         use_thd: bool = False,
         deterministic: bool = False,
+        gdn_cp_mode: str = "fla_allgather",
     ):
         super().__init__()
         self.layer_idx = layer_idx
@@ -264,6 +267,7 @@ def __init__(
                 rms_norm_eps=config.rms_norm_eps,
                 ps=ps,
                 deterministic=deterministic,
+                cp_mode=gdn_cp_mode,
             )
         self.mlp_norm = te.RMSNorm(
             config.hidden_size, eps=config.rms_norm_eps, zero_centered_gamma=True
@@ -332,6 +336,7 @@ def __init__(
         mtp_enable_train: bool = False,
         mtp_detach_encoder: bool = False,
         mount_vision_model: bool = False,
+        gdn_cp_mode: str = "fla_allgather",
     ):
         super().__init__()
         del attention_backend_override
@@ -373,6 +378,7 @@ def __init__(
                     moe_act_recompute=moe_act_recompute,
                     use_thd=use_thd,
                     deterministic=getattr(train_config, "deterministic", False),
+                    gdn_cp_mode=gdn_cp_mode,
                 )
                 for idx in self.layer_indices
             ]
@@ -410,6 +416,7 @@ def make_mtp_layer(layer_idx: int) -> MTPDecoderLayer:
                         moe_act_recompute=moe_act_recompute,
                         use_thd=use_thd,
                         deterministic=getattr(train_config, "deterministic", False),
+                        gdn_cp_mode=gdn_cp_mode,
                     ),
                     detach_encoder=mtp_detach_encoder,
                 )
@@ -552,10 +559,10 @@ def _apply_mtp_loss(
         mtp_loss_mask = loss_mask.clone()
         mtp_loss_values = []
         for mtp_hidden in mtp_hidden_states:
-            mtp_labels, _ = roll_packed_thd_left(
+            mtp_labels, _ = roll_mtp_tensor_left(
                 mtp_labels, packed_seq_params=packed_seq_params, dims=-1
             )
-            mtp_loss_mask, num_tokens = roll_packed_thd_left(
+            mtp_loss_mask, num_tokens = roll_mtp_tensor_left(
                 mtp_loss_mask, packed_seq_params=packed_seq_params, dims=-1
             )
             labels_sb = mtp_labels.transpose(0, 1).contiguous()
@@ -658,7 +665,7 @@ def _hook_vision_params_avg_grad_across_tp(module: nn.Module) -> None:
         param.average_gradients_across_tp_domain = True  # type: ignore[assignment]
 
 
-def _build_native_vision_model(hf_path: str) -> nn.Module:
+def _build_native_vision_model(hf_path: str) -> nn.Module | None:
     if not hf_path:
         raise ValueError("mount_vision_model requires hf_path.")
     try:
@@ -669,12 +676,22 @@ def _build_native_vision_model(hf_path: str) -> nn.Module:
     hf_config = AutoConfig.from_pretrained(hf_path, trust_remote_code=True)
     vision_config = getattr(hf_config, "vision_config", None)
     if vision_config is None:
-        raise RuntimeError("HF config does not expose vision_config; cannot build vision_model.")
-    hf_vision_cls = _resolve_hf_vision_cls(hf_config, hf_path)
-    if hasattr(hf_vision_cls, "_from_config"):
-        vision = hf_vision_cls._from_config(vision_config)
+        # Text-only checkpoint (no vision tower): nothing to mount -> graceful skip.
+        return None
+    auto_map = getattr(hf_config, "auto_map", None) or {}
+    if auto_map:
+        # Remote-code checkpoint: resolve the vision class from its dynamic module.
+        hf_vision_cls = _resolve_hf_vision_cls(hf_config, hf_path)
+        if hasattr(hf_vision_cls, "_from_config"):
+            vision = hf_vision_cls._from_config(vision_config)
+        else:
+            vision = hf_vision_cls(vision_config)
     else:
-        vision = hf_vision_cls(vision_config)
+        # Native-transformers checkpoint (auto_map=None, e.g. Qwen3.5-35B-A3B): build the vision tower
+        # directly from its vision_config via the native AutoModel registry (no remote code).
+        from transformers import AutoModel
+
+        vision = AutoModel.from_config(vision_config)
     _hook_fp32_rotary_emb(vision)
     _hook_vision_params_avg_grad_across_tp(vision)
     return vision.to(torch.bfloat16)

experimental/lite/megatron/lite/primitive/bundle.py

@@ -24,6 +24,6 @@ class ModelBundle:
     parallel_state: ParallelState
     optimizer: Any | None = None
     finalize_grads: Callable[[], None] | None = None
-    forward_step: Callable[[nn.Module, dict], dict] | None = None
+    forward_step: Callable[..., dict] | None = None
     # extra metadata (expert_classifier, model_cfg, etc.)
     extras: dict[str, Any] = field(default_factory=dict)

experimental/lite/megatron/lite/primitive/ckpt/__init__.py

@@ -2,7 +2,6 @@
 """Checkpoint helpers."""
 
 from megatron.lite.primitive.ckpt.dcp import load_training_checkpoint, save_training_checkpoint
-from megatron.lite.primitive.ckpt.distckpt import attach_model_sharded_state_dict
 from megatron.lite.primitive.ckpt.hf_weights import HFWeights
 
 __all__ = [
@@ -11,3 +10,11 @@
     "load_training_checkpoint",
     "save_training_checkpoint",
 ]
+
+
+def __getattr__(name: str):
+    if name == "attach_model_sharded_state_dict":
+        from megatron.lite.primitive.ckpt.distckpt import attach_model_sharded_state_dict
+
+        return attach_model_sharded_state_dict
+    raise AttributeError(name)