[dev] Megatron Lite (4/4) shared attention

experimental/lite/examples/bench/bench.py

@@ -156,18 +156,18 @@ def keep_experts_hook(model_cfg):
 
         def truncate_layers_hook(model_cfg):
             old_layers = getattr(model_cfg, "num_hidden_layers", None)
-            layer_types = getattr(model_cfg, "layer_types", None)
-            if old_layers is None or layer_types is None:
-                raise ValueError("truncate_layers requires num_hidden_layers and layer_types.")
+            if old_layers is None:
+                raise ValueError("truncate_layers requires num_hidden_layers.")
             if keep_layers <= 0 or keep_layers > old_layers:
                 raise ValueError(
                     f"truncate_layers must be in [1, {old_layers}], got {keep_layers}."
                 )
-            return replace(
-                model_cfg,
-                num_hidden_layers=keep_layers,
-                layer_types=list(layer_types[:keep_layers]),
-            )
+            updates = {"num_hidden_layers": keep_layers}
+            # layer_types is qwen-style metadata; deepseek_v4 / kimi_k2 configs don't carry it.
+            layer_types = getattr(model_cfg, "layer_types", None)
+            if layer_types is not None:
+                updates["layer_types"] = list(layer_types[:keep_layers])
+            return replace(model_cfg, **updates)
 
         hooks.append(truncate_layers_hook)
 
@@ -289,6 +289,7 @@ def build_runtime_config(cfg: BenchCliConfig) -> RuntimeConfig:
             model_name=cfg.model_name,
             parallel=parallel,
             optimizer=optimizer,
+            use_thd=cfg.use_thd,
             load_hf_weights=not cfg.skip_load_hf_weights,
             build_optimizer=not cfg.skip_optimizer_build,
             override_ddp_config=_json_mapping(cfg.override_ddp_json, name="override_ddp_json"),

experimental/lite/examples/bench/correctness.py

@@ -22,11 +22,12 @@
 if str(_REPO_ROOT) not in sys.path:
     sys.path.insert(1, str(_REPO_ROOT))
 
+from megatron.lite.primitive.deterministic import set_deterministic
+from megatron.lite.runtime import create_runtime
+
 from examples.bench.bench import BenchCliConfig, build_runtime_config, build_session_config
 from examples.bench.results import compare_correctness_artifacts, load_result_artifact
 from examples.bench.session import _make_data_iter
-from megatron.lite.primitive.deterministic import set_deterministic
-from megatron.lite.runtime import create_runtime
 
 
 def _distributed_rank() -> int:
@@ -299,48 +300,15 @@ def _weight_fingerprint(rt, handle) -> dict[str, Any]:
     return result
 
 
-def _batch_without_labels(batch: Any) -> dict[str, Any]:
-    if not isinstance(batch, dict):
-        return {
-            "input_ids": batch["input_ids"],
-            "position_ids": getattr(batch, "position_ids", None),
-            "packed_seq_params": getattr(batch, "packed_seq_params", None),
-        }
-    return {k: v for k, v in batch.items() if k != "labels"}
-
-
 def _forward_logits(rt, handle, batch: Any) -> torch.Tensor | None:
-    sample = _batch_without_labels(batch)
-    if "forward_step" in handle._extras:
-        try:
-            out = handle._model(**sample)
-        except (KeyError, TypeError):
-            out = handle._extras["forward_step"](handle._model, sample)
-        if isinstance(out, dict):
-            logits = out.get("logits")
-            if logits is not None:
-                return logits
-            return out.get("vocab_parallel_logits")
-        return out if isinstance(out, torch.Tensor) else None
-
-    model_list = handle._extras.get("model_list")
-    if model_list:
-        out = model_list[0](
-            input_ids=sample.get("input_ids"),
-            position_ids=sample.get("position_ids"),
-            attention_mask=sample.get("attention_mask"),
-            packed_seq_params=sample.get("packed_seq_params"),
-        )
-        if isinstance(out, tuple):
-            out = out[0]
-        if isinstance(out, dict):
-            logits = out.get("logits")
-            if logits is not None:
-                return logits
-            return out.get("vocab_parallel_logits")
-        return out if isinstance(out, torch.Tensor) else None
-
-    return None
+    result = rt.forward_backward(
+        handle,
+        iter([batch]),
+        loss_fn=None,
+        num_microbatches=1,
+        forward_only=True,
+    )
+    return result.model_output.vocab_parallel_logits
 
 
 def run_backend(

experimental/lite/examples/bench/session.py

@@ -10,7 +10,6 @@
 from typing import Any
 
 import torch
-
 from megatron.lite.runtime.backends import Runtime
 from megatron.lite.runtime.contracts.handle import ModelHandle
 
@@ -66,26 +65,33 @@ def _resolve_vocab_size(handle: ModelHandle) -> int:
     return 151936
 
 
-def _make_data_iter(handle: ModelHandle, cfg: PretrainSessionConfig):
-    data_seed = cfg.seed if cfg.same_data_across_dp else cfg.seed + handle.dp_rank
-    vocab_size = _resolve_vocab_size(handle)
-
-    if cfg.use_thd:
-        from megatron.lite.primitive.data import infinite_batches_thd
-
-        ps = handle._parallel_state
-        return infinite_batches_thd(
-            vocab_size,
-            cfg.seq_len,
-            cp_size=getattr(ps, "cp_size", 1),
-            cp_rank=getattr(ps, "cp_rank", 0),
-            device=cfg.device,
-            seed=data_seed,
+def _infinite_packed_batches(
+    vocab_size: int, seq_len: int, *, device: str, seed: int
+):
+    """Yield raw, model-agnostic :class:`PackedBatch` objects for the bench.
+
+    The bench is the single source of truth for one unpadded packed batch (1-D
+    ``input_ids``/``labels`` plus true per-sequence ``seq_lens``). Padding, CP
+    layout and THD metadata (``packed_seq_params``) are derived by whichever
+    runtime/model consumes the batch at the forward boundary, never baked into
+    bench data — that is what keeps the mlite-vs-bridge comparison fair.
+    """
+    from megatron.lite.runtime.contracts.data import PackedBatch
+
+    g = torch.Generator(device=device).manual_seed(seed)
+    seq_lens = torch.tensor([seq_len], dtype=torch.int64, device=device)
+    while True:
+        yield PackedBatch(
+            input_ids=torch.randint(0, vocab_size, (seq_len,), device=device, generator=g),
+            labels=torch.randint(0, vocab_size, (seq_len,), device=device, generator=g),
+            seq_lens=seq_lens.clone(),
         )
 
-    from megatron.lite.primitive.data import infinite_batches
 
-    return infinite_batches(vocab_size, cfg.seq_len, device=cfg.device, seed=data_seed)
+def _make_data_iter(handle: ModelHandle, cfg: PretrainSessionConfig):
+    data_seed = cfg.seed if cfg.same_data_across_dp else cfg.seed + handle.dp_rank
+    vocab_size = _resolve_vocab_size(handle)
+    return _infinite_packed_batches(vocab_size, cfg.seq_len, device=cfg.device, seed=data_seed)
 
 
 def _calc_tflops_per_gpu(

experimental/lite/examples/verl/scripts/run_qwen3moe_gsm8k_grpo.sh

@@ -85,9 +85,9 @@ MLITE_MODEL_NAME="${MLITE_MODEL_NAME:-auto}"
 MLITE_IMPL="${MLITE_IMPL:-lite}"
 ATTENTION_BACKEND="${ATTENTION_BACKEND:-flash}"
 # Optimizer backend:
-# - distopt (default): Megatron-Core DDP + distributed optimizer.
+# - dist_opt (default): Megatron-Core DDP + distributed optimizer.
 # - fsdp2: Megatron Lite FSDP2 wrapper + optimizer.
-MLITE_OPTIMIZER_BACKEND="${MLITE_OPTIMIZER_BACKEND:-distopt}"
+MLITE_OPTIMIZER_BACKEND="${MLITE_OPTIMIZER_BACKEND:-dist_opt}"
 
 ACTOR_LR="${ACTOR_LR:-1e-6}"
 POLICY_LOSS_MODE="${POLICY_LOSS_MODE:-vanilla}"
@@ -124,14 +124,14 @@ if [[ "${INFER_BACKEND}" != "vllm" && "${INFER_BACKEND}" != "sglang" && "${INFER
 fi
 
 case "${MLITE_OPTIMIZER_BACKEND}" in
-  distopt)
-    MLITE_IMPL_OPTIMIZER="mc"
+  dist_opt)
+    MLITE_IMPL_OPTIMIZER="dist_opt"
     ;;
   fsdp2)
     MLITE_IMPL_OPTIMIZER="fsdp2"
     ;;
   *)
-    echo "Unsupported MLITE_OPTIMIZER_BACKEND=${MLITE_OPTIMIZER_BACKEND}. Expected distopt or fsdp2." >&2
+    echo "Unsupported MLITE_OPTIMIZER_BACKEND=${MLITE_OPTIMIZER_BACKEND}. Expected dist_opt or fsdp2." >&2
     exit 1
     ;;
 esac
@@ -184,7 +184,6 @@ DATA=(
 MODEL=(
   "actor_rollout_ref.model.path=${MODEL_PATH}"
   "actor_rollout_ref.model.trust_remote_code=True"
-  "actor_rollout_ref.model.use_remove_padding=True"
   "actor_rollout_ref.model.use_fused_kernels=False"
 )
 

experimental/lite/examples/verl/scripts/run_qwen3moe_sft.sh

@@ -24,6 +24,10 @@ add_pythonpath "${VERL_ROOT:-}"
 add_pythonpath "${MEGATRON_ROOT:-}"
 
 export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}"
+if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then
+  unset ROCR_VISIBLE_DEVICES
+  unset HIP_VISIBLE_DEVICES
+fi
 
 : "${MODEL_PATH:?set MODEL_PATH to a Hugging Face checkpoint directory or model id}"
 : "${TRAIN_FILES:?set TRAIN_FILES to a messages parquet path or comma-separated parquet paths}"
@@ -46,13 +50,14 @@ SAVE_FREQ="${SAVE_FREQ:-${TOTAL_STEPS}}"
 TEST_FREQ="${TEST_FREQ:--1}"
 RESUME_MODE="${RESUME_MODE:-disable}"
 RESUME_FROM_PATH="${RESUME_FROM_PATH:-null}"
+CHECKPOINT_SAVE_CONTENTS="${CHECKPOINT_SAVE_CONTENTS:-[model,optimizer,extra]}"
+LOAD_HF_WEIGHTS="${LOAD_HF_WEIGHTS:-True}"
 TRAIN_BATCH_SIZE="${TRAIN_BATCH_SIZE:-64}"
 MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-1}"
 MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-8192}"
 MAX_LENGTH="${MAX_LENGTH:-${MAX_TOKENS_PER_GPU}}"
 PAD_MODE="${PAD_MODE:-no_padding}"
 USE_DYNAMIC_BSZ="${USE_DYNAMIC_BSZ:-True}"
-USE_REMOVE_PADDING="${USE_REMOVE_PADDING:-True}"
 IGNORE_INPUT_IDS_MISMATCH="${IGNORE_INPUT_IDS_MISMATCH:-True}"
 TRUST_REMOTE_CODE="${TRUST_REMOTE_CODE:-True}"
 MESSAGES_KEY="${MESSAGES_KEY:-messages}"
@@ -70,9 +75,9 @@ MLITE_MODEL_NAME="${MLITE_MODEL_NAME:-auto}"
 MLITE_IMPL="${MLITE_IMPL:-lite}"
 ATTENTION_BACKEND="${ATTENTION_BACKEND:-flash}"
 # Optimizer backend:
-# - distopt (default): Megatron-Core DDP + distributed optimizer.
+# - dist_opt (default): Megatron-Core DDP + distributed optimizer.
 # - fsdp2: Megatron Lite FSDP2 wrapper + optimizer.
-MLITE_OPTIMIZER_BACKEND="${MLITE_OPTIMIZER_BACKEND:-distopt}"
+MLITE_OPTIMIZER_BACKEND="${MLITE_OPTIMIZER_BACKEND:-dist_opt}"
 
 LR="${LR:-1e-5}"
 MIN_LR="${MIN_LR:-${LR}}"
@@ -102,14 +107,14 @@ if [[ "${PAD_MODE}" != "no_padding" ]]; then
 fi
 
 case "${MLITE_OPTIMIZER_BACKEND}" in
-  distopt)
-    MLITE_IMPL_OPTIMIZER="mc"
+  dist_opt)
+    MLITE_IMPL_OPTIMIZER="dist_opt"
     ;;
   fsdp2)
     MLITE_IMPL_OPTIMIZER="fsdp2"
     ;;
   *)
-    echo "Unsupported MLITE_OPTIMIZER_BACKEND=${MLITE_OPTIMIZER_BACKEND}. Expected distopt or fsdp2." >&2
+    echo "Unsupported MLITE_OPTIMIZER_BACKEND=${MLITE_OPTIMIZER_BACKEND}. Expected dist_opt or fsdp2." >&2
     exit 1
     ;;
 esac
@@ -149,7 +154,6 @@ COMMON_ARGS=(
   "model=hf_model"
   "model.path=${MODEL_PATH}"
   "model.trust_remote_code=${TRUST_REMOTE_CODE}"
-  "model.use_remove_padding=${USE_REMOVE_PADDING}"
   "optim=megatron"
   "optim.lr=${LR}"
   "optim.min_lr=${MIN_LR}"
@@ -172,7 +176,7 @@ COMMON_ARGS=(
   "trainer.resume_from_path=${RESUME_FROM_PATH}"
   "trainer.nnodes=${NNODES}"
   "trainer.n_gpus_per_node=${NPROC_PER_NODE}"
-  "checkpoint.save_contents=[model,optimizer,extra]"
+  "checkpoint.save_contents=${CHECKPOINT_SAVE_CONTENTS}"
 )
 
 if [[ -n "${VAL_FILES}" ]]; then
@@ -195,6 +199,7 @@ BACKEND_ARGS=(
   "engine.optimizer_offload=${OPTIMIZER_OFFLOAD}"
   "engine.grad_offload=${GRAD_OFFLOAD}"
   "engine.attention_backend_override=${ATTENTION_BACKEND}"
+  "engine.load_hf_weights=${LOAD_HF_WEIGHTS}"
   "engine.impl_cfg.use_thd=True"
   "+engine.impl_cfg.optimizer=${MLITE_IMPL_OPTIMIZER}"
 )
@@ -215,6 +220,7 @@ COMMAND=(
   --master_port="${MASTER_PORT}"
   --nproc_per_node="${NPROC_PER_NODE}"
   -m
+  verl_mlite.launch
   verl.trainer.sft_trainer
   "${COMMON_ARGS[@]}"
   "${BACKEND_ARGS[@]}"

experimental/lite/examples/verl/verl_mlite/compat.py

@@ -3,8 +3,11 @@
 
 from __future__ import annotations
 
+import importlib.util
+import sys
 from collections.abc import Iterable
 from functools import wraps
+from pathlib import Path
 from typing import Any
 
 
@@ -54,3 +57,40 @@ def patched(*args: Any, **kwargs: Any) -> Any:
 
 def apply_runtime_patches() -> None:
     _patch_transformers_rope_ignore_keys()
+
+
+def _load_verl_file(relative_path: str, module_name: str):
+    spec = importlib.util.find_spec("verl")
+    if spec is None or spec.submodule_search_locations is None:
+        raise ModuleNotFoundError("No module named 'verl'")
+
+    path = Path(next(iter(spec.submodule_search_locations))) / relative_path
+    file_spec = importlib.util.spec_from_file_location(module_name, path)
+    if file_spec is None or file_spec.loader is None:
+        raise ImportError(f"Unable to load VERL module from {path}")
+
+    module = importlib.util.module_from_spec(file_spec)
+    sys.modules[module_name] = module
+    file_spec.loader.exec_module(module)
+    return module
+
+
+def load_verl_engine_api():
+    # Prefer the canonical package import so the MLite engine registers into the
+    # SAME EngineRegistry that verl's trainers resolve against. Loading base.py as
+    # a standalone module (below) creates a *duplicate* registry, which silently
+    # drops the mlite backend ("Unknown backend: mlite"). The file-load path is
+    # only a fallback for environments where verl isn't importable as a package.
+    try:
+        from verl.workers.engine.base import BaseEngine, BaseEngineCtx, EngineRegistry
+        from verl.workers.engine.utils import postprocess_batch_func, prepare_micro_batches
+    except (ModuleNotFoundError, ImportError):
+        base = _load_verl_file("workers/engine/base.py", "_verl_mlite_verl_engine_base")
+        utils = _load_verl_file("workers/engine/utils.py", "_verl_mlite_verl_engine_utils")
+        BaseEngine = base.BaseEngine
+        BaseEngineCtx = base.BaseEngineCtx
+        EngineRegistry = base.EngineRegistry
+        postprocess_batch_func = utils.postprocess_batch_func
+        prepare_micro_batches = utils.prepare_micro_batches
+
+    return BaseEngine, BaseEngineCtx, EngineRegistry, postprocess_batch_func, prepare_micro_batches

experimental/lite/examples/verl/verl_mlite/config/engine/mlite.yaml

@@ -8,6 +8,7 @@ grad_offload: false
 forward_only: false
 dtype: bfloat16
 export_dtype: bfloat16
+load_hf_weights: true
 
 model_name: auto
 impl: lite

experimental/lite/examples/verl/verl_mlite/engine/config.py

@@ -28,7 +28,9 @@ class MegatronLiteEngineConfig(EngineConfig):
 
     attention_backend_override: str | None = "flash"
     router_aux_loss_coef: float | None = None
+    cross_entropy_fusion: bool | None = None
     export_dtype: str | None = "bfloat16"
+    load_hf_weights: bool = True
     impl_cfg: dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self) -> None: