diff --git a/experimental/lite/.pre-commit-config.yaml b/experimental/lite/.pre-commit-config.yaml new file mode 100644 index 00000000000..a1b8b2d897f --- /dev/null +++ b/experimental/lite/.pre-commit-config.yaml @@ -0,0 +1,16 @@ +repos: +- repo: local + hooks: + - id: isort + name: isort + entry: isort + language: python + additional_dependencies: ["isort==5.13.2"] + files: ^experimental/lite/.*\.py$ + - id: black + name: black + entry: black + language: python + additional_dependencies: ["black==24.4.2"] + files: ^experimental/lite/.*\.py$ + args: ["--skip-magic-trailing-comma", "--skip-string-normalization"] diff --git a/experimental/lite/README.md b/experimental/lite/README.md new file mode 100644 index 00000000000..e9b8f4d4082 --- /dev/null +++ b/experimental/lite/README.md @@ -0,0 +1,139 @@ +# Megatron Lite + +Megatron Lite is an experimental, agentic-native training runtime and native +model implementation layer for Megatron. It is designed for work that needs to +move quickly without giving up Megatron-Core performance: small composable +primitives, explicit model/runtime protocols, and validation recipes that make +changes easy to review and easy to reproduce. + +The source lives under `experimental/lite/megatron/lite`, and the public import +path is `megatron.lite`. + +Do not import `experimental.lite` from user code. Examples and public APIs should +refer to `megatron.lite`. + +## Scope + +This initial drop contains: + +- A lightweight runtime API in `megatron.lite.runtime`. +- Common training primitives in `megatron.lite.primitive`. +- Lite-only native model implementations for Qwen3 MoE and Qwen3.5 MoE. +- Hugging Face safetensors load/export helpers for the included models. +- Megatron-Core optimizer wrapping for the lite runtime. +- FSDP2 optimizer primitives for supported lite model protocols. +- Reference runtime backends for comparison runs: `mbridge` for the legacy + package and `bridge` for real Megatron-Bridge environments. +- A benchmark example that can dry-run or execute `mlite`, `mbridge`, and + `bridge` backends. + +This initial drop intentionally does not include: + +- Hybrid model implementations. +- Dense Qwen3 model support. The included Qwen3-family path is Qwen3 MoE only. + +## Why MLite + +- **Agentic-native development surface.** Runtime, model, and primitive code are + split into reviewable contracts so agents and humans can make targeted changes + without touching unrelated Megatron subsystems. +- **Native MLite models, not wrapper models.** `backend="mlite"` builds native + `megatron.lite` model code; reference backends are used only for comparison. +- **Megatron-Core distopt parity.** In deterministic correctness runs against the + `mbridge` reference backend on the Megatron-Core distributed optimizer path, + MLite matched loss and grad-norm exactly (`max_abs=0.0`) with no mismatches; + post-step weights and eval logits were checked by SHA256 fingerprints. +- **Speed-aligned with the Core path.** On an 8x H100 Qwen3.5 MoE benchmark using + `distopt`, MLite measured 309.433 ms/step and 105,896.935 tokens/s, compared + with 332.201 ms/step and 98,639.089 tokens/s for `mbridge`, and 334.936 + ms/step and 97,833.496 tokens/s for the real `bridge` path. + +The `mbridge` benchmark line is the validated Megatron-Core/distopt reference +used for this PR. The `bridge` line is a separate Megatron-Bridge environment +check and should not be confused with the Core/distopt parity claim. + +## Layout + +```text +experimental/lite/ + README.md + docs/ Design and usage notes + examples/ Optional integration and benchmark examples + skills/ Agent-agnostic maintenance skills + megatron/ + lite/ + runtime/ Runtime API, config, and backend registry + model/ Model registry and Qwen model implementations + primitive/ Parallel, checkpoint, optimizer, module, and op primitives +``` + +For local source-tree use: + +```bash +export PYTHONPATH=/path/to/Megatron-LM/experimental/lite:$PYTHONPATH +``` + +## Public API + +```python +from megatron.lite.runtime import MegatronLiteConfig, RuntimeConfig, create_runtime + +cfg = RuntimeConfig( + backend="mlite", + hf_path="/path/to/hf-model", + backend_cfg=MegatronLiteConfig(model_name="qwen3_moe", impl="lite"), +) +runtime = create_runtime(cfg) +handle = runtime.build_model() +``` + +`backend="mlite"` selects the Megatron Lite runtime backend. `impl="lite"` +selects the model implementation inside the registered model family. +`backend="mbridge"` selects the legacy `mbridge` reference backend used by the +validated benchmark example. `backend="bridge"` selects the Megatron-Bridge +runtime backend and requires an environment where `import megatron.bridge` works +when the model is built. + +Canonical model names currently registered by default: + +- `qwen3_moe`: Qwen3 MoE lite implementation. Use this name in new configs. +- `qwen3_5`: Qwen3.5 MoE lite implementation. + +Compatibility names: + +- `qwen3`: legacy alias for the Qwen3 MoE implementation only. It does not mean + dense Qwen3 support. HF `model_type` values `qwen3_moe` and `qwen2_moe` + currently resolve through this compatibility path. + +## Benchmark And Correctness Signoff + +The validated benchmark and correctness commands live in +[`examples/bench/README.md`](examples/bench/README.md). The signoff setup uses +Qwen3.5 MoE, `optimizer_backend=distopt`, deterministic mode for strict +correctness, and identical synthetic input streams for paired performance runs. + +Reproduce the strict MLite vs Megatron-Core/distopt comparison with: + +```bash +export MEGATRON_LITE_DETERMINISTIC=1 +export CUBLAS_WORKSPACE_CONFIG=:4096:8 + +HF_PATH=/models/Qwen3.5-35B-A3B \ +REFERENCE_BACKEND=mbridge \ +DRY_RUN=0 \ +bash experimental/lite/examples/bench/scripts/run_qwen35_correctness_pair.sh +``` + +## Docs + +- [Architecture](docs/architecture.md) +- [Runtime](docs/runtime.md) +- [Models](docs/models.md) +- [Porting Notes](docs/porting.md) +- [Skills](skills/README.md) +- [Bench Example](examples/bench/README.md) + +## Acknowledgements + +The Qwen3 MoE LoRA adapter support follows Mind-Lab's PEFT/Mint-compatible +adapter work. Thanks to Mind-Lab for the reference implementation and guidance. diff --git a/experimental/lite/docs/architecture.md b/experimental/lite/docs/architecture.md new file mode 100644 index 00000000000..a239b79997e --- /dev/null +++ b/experimental/lite/docs/architecture.md @@ -0,0 +1,46 @@ +# Architecture + +Megatron Lite source lives under `experimental/lite/megatron/lite` and is split +into three layers: + +- `runtime`: lifecycle and training-step orchestration. +- `model`: model registration plus model-specific build/load/export protocols. +- `primitive`: reusable lower-level pieces such as parallel state, tensor-parallel + layers, checkpoint conversion, MoE utilities, and optimizer wrapping. + +The runtime does not know Qwen implementation details. It imports a model +protocol from the model registry, builds the typed implementation config, then +delegates model construction to that protocol. + +## Import Boundary + +The source root for local use is `experimental/lite`; adding that directory to +`PYTHONPATH` exposes the package as `megatron.lite`. Internal imports also use +`megatron.lite` so user-facing code matches the final package path. + +## Runtime Boundary + +The runtime API owns: + +- Distributed initialization. +- Model protocol loading. +- Model checkpoint save/load dispatch. +- Forward/backward microbatch orchestration. +- Optimizer and learning-rate scheduler stepping. +- Optional model/optimizer offload hooks. + +The model protocol owns: + +- Architecture config creation. +- Model chunk construction. +- Model-specific recompute/offload wiring. +- Model-specific optimizer construction. +- HF checkpoint load/export mapping. + +## Current Deliberate Omissions + +This package currently includes only the lite model implementation path. It +intentionally excludes non-lite and hybrid model implementation packages. The +optional benchmark entrypoints under `examples/bench` are comparison tools and +are not imported by the package runtime. FSDP2 is included as an optimizer +primitive and can be selected by model protocols that support it. diff --git a/experimental/lite/docs/models.md b/experimental/lite/docs/models.md new file mode 100644 index 00000000000..98e17389a88 --- /dev/null +++ b/experimental/lite/docs/models.md @@ -0,0 +1,58 @@ +# Models + +Model code lives under `megatron.lite.model`. + +The model registry maps `(model_name, impl)` pairs to model protocol modules. +The runtime loads the protocol and expects the following required symbols: + +- `ImplConfig`: dataclass for model-specific knobs. +- `build_model_config(source, **overrides)`: returns a typed model config. +- `build_model(model_cfg, *, impl_cfg)`: returns a `ModelBundle`. + +Optional protocol symbols: + +- `load_hf_weights(chunk, hf_path, model_cfg, ps)` +- `export_hf_weights(chunks, model_cfg, ps, **kwargs)` +- `vocab_size(model_cfg)` + +## Included Models + +`qwen3_moe` is the canonical name for the Qwen3 MoE lite implementation: + +```text +megatron.lite.model.qwen3_moe.lite.protocol +``` + +`qwen3` is kept as a legacy compatibility alias for the same implementation. It +does not mean dense Qwen3 support. Hugging Face `model_type` values `qwen3_moe` +and `qwen2_moe` currently resolve through this compatibility path. + +`qwen3_5` maps to the Qwen3.5 MoE lite implementation: + +```text +megatron.lite.model.qwen3_5.lite.protocol +``` + +## Acknowledgements + +The Qwen3 MoE LoRA adapter support follows Mind-Lab's PEFT/Mint-compatible +adapter work. Thanks to Mind-Lab for the reference implementation and guidance. + +## Adding A Model + +Add a model package under `model/`, then register it in +`model/registry.py`: + +```python +register_model( + "my_model", + package="megatron.lite.model.my_model", + hf_model_types=["my_model"], + impls={ + "lite": "megatron.lite.model.my_model.lite.protocol", + }, +) +``` + +New models should keep heavyweight imports inside protocol functions when +possible so importing `megatron.lite` stays cheap. diff --git a/experimental/lite/docs/porting.md b/experimental/lite/docs/porting.md new file mode 100644 index 00000000000..2c762a9c3ba --- /dev/null +++ b/experimental/lite/docs/porting.md @@ -0,0 +1,43 @@ +# Porting Notes + +This tree is prepared as an experimental Megatron package. The package code +lives under `experimental/lite/megatron/lite`, so users can add +`experimental/lite` to `PYTHONPATH` and import `megatron.lite`. + +## Naming Rules + +- Use `Megatron Lite` for the component name in docs and comments. +- Use `megatron.lite` for public and internal imports. +- Do not introduce project-specific legacy branding. +- Use `mlite` for the runtime backend key. +- Use `lite` for model implementation names. + +## Included Surface + +Keep the PR focused on the lite model implementation path: + +- Runtime backend: `mlite`. +- Reference comparison backends: `mbridge` for the validated legacy + Megatron-Core/distopt path and `bridge` for real Megatron-Bridge environments. +- Models: Qwen3 MoE and Qwen3.5 MoE. Dense Qwen3 is not included. +- Model implementations: `lite` only. +- Optimizer primitives: Megatron-Core optimizer wrapping and FSDP2. +- Optional examples: benchmark and VERL launchers under `experimental/lite/examples`. + +Keep these out of the first PR unless the scope changes: + +- Hybrid model implementation packages. +- Megatron-Bridge model implementation packages. +- Undocumented experiment-specific entrypoints. + +## Package Integration + +No repository-level packaging changes are made in this experimental drop. The +current layout is importable from source with: + +```bash +export PYTHONPATH=/path/to/Megatron-LM/experimental/lite:$PYTHONPATH +``` + +A future integration step can decide whether to keep the experimental location +or move the tree into the final package location. diff --git a/experimental/lite/docs/runtime.md b/experimental/lite/docs/runtime.md new file mode 100644 index 00000000000..f74851f8fe2 --- /dev/null +++ b/experimental/lite/docs/runtime.md @@ -0,0 +1,83 @@ +# Runtime + +The public runtime entrypoint is `megatron.lite.runtime`. + +```python +from megatron.lite.runtime import MegatronLiteConfig, ParallelConfig, RuntimeConfig, create_runtime + +cfg = RuntimeConfig( + backend="mlite", + hf_path="/path/to/hf-model", + backend_cfg=MegatronLiteConfig( + model_name="qwen3_moe", + impl="lite", + parallel=ParallelConfig(tp=1, pp=1, cp=1, ep=1), + ), +) +runtime = create_runtime(cfg) +handle = runtime.build_model() +``` + +## API Tiers + +All runtime backends implement the pretraining tier: + +- `build_model` +- `save_checkpoint` +- `load_checkpoint` +- `train_mode` +- `eval_mode` +- `forward_backward` +- `zero_grad` +- `optimizer_step` +- `lr_scheduler_step` + +The lite runtime also implements `export_weights` and `to` when the underlying +model and optimizer support those operations. + +The `mbridge` runtime implements the same runtime contract through the legacy +`mbridge` package and Megatron-Core optimizer/checkpoint helpers. The benchmark +example currently uses this backend for validated reference runs. + +The `bridge` runtime is the real Megatron-Bridge path. It imports +`megatron.bridge` lazily from `build_model()`, so config construction and dry-run +examples can execute without Megatron-Bridge installed. + +## Config Types + +`RuntimeConfig` selects the backend and carries the Hugging Face model path. + +`MegatronLiteConfig` carries `mlite` backend settings: + +- `model_name`: `qwen3_moe` or `qwen3_5` for new configs. `qwen3` remains + accepted as a legacy alias for `qwen3_moe` only; dense Qwen3 is not included. +- `impl`: currently only `lite`. +- `parallel`: tensor, expert, pipeline, virtual pipeline, and context sizes. +- `optimizer`: Megatron-Core optimizer settings. +- `impl_cfg`: model-specific options consumed by each model protocol. + +`BridgeConfig` carries shared `mbridge` and `bridge` backend settings: + +- `model_name`: optional model identifier used for benchmark metadata. +- `parallel`: tensor, expert, pipeline, virtual pipeline, and context sizes. +- `optimizer`: Megatron-Core optimizer settings. +- `override_ddp_config`, `override_transformer_config`, and + `override_optimizer_config`: explicit reference-backend/Core override maps. +- `param_offload` and `optimizer_offload`: offload model/optimizer state between + train/eval contexts. + +## Backend Registry + +The built-in backend keys are `mlite`, `mbridge`, and `bridge`. Model +implementations for the native runtime remain selected through +`MegatronLiteConfig.impl`, which currently supports `impl="lite"`. + +Custom runtime backends can be registered with: + +```python +from megatron.lite.runtime import register_runtime + +register_runtime("my_backend", "my_package.my_runtime") +``` + +The target module must expose `create(hf_path, cfg)`. diff --git a/experimental/lite/examples/__init__.py b/experimental/lite/examples/__init__.py new file mode 100644 index 00000000000..742e27ee4bf --- /dev/null +++ b/experimental/lite/examples/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Megatron Lite examples package.""" diff --git a/experimental/lite/examples/bench/.gitignore b/experimental/lite/examples/bench/.gitignore new file mode 100644 index 00000000000..17aa483ab4e --- /dev/null +++ b/experimental/lite/examples/bench/.gitignore @@ -0,0 +1 @@ +outputs/ diff --git a/experimental/lite/examples/bench/README.md b/experimental/lite/examples/bench/README.md new file mode 100644 index 00000000000..a9a2ecf369e --- /dev/null +++ b/experimental/lite/examples/bench/README.md @@ -0,0 +1,205 @@ +# MLite Bench Example + +This example runs the same small pretrain-style benchmark through `backend=mlite` +and a reference backend. + +The validated reference backend in this PR is `backend=mbridge`, backed by the +legacy `mbridge` package. `backend=bridge` is reserved for the real +Megatron-Bridge package and requires an environment where `import +megatron.bridge` works. Dry-run mode does not import either reference package and +is safe for config validation. + +The `backend=mlite` path remains native MLite. For deterministic Qwen3.5 MoE +runs it mounts the Qwen3.5 vision module from the local Hugging Face model code +through `transformers`; it does not import `mbridge` or Megatron-Bridge to build +the MLite model. + +This benchmark separates two validation lines: + +- `mbridge`: the validated reference for MLite vs Megatron-Core distributed + optimizer (`distopt`) parity. +- `bridge`: a real Megatron-Bridge environment check when `import + megatron.bridge` works. + +Use the `mbridge` line for Core/distopt precision and speed claims. + +## Dry-Run + +```bash +export PYTHONPATH=/path/to/Megatron-LM/experimental/lite:$PYTHONPATH + +python experimental/lite/examples/bench/bench.py \ + --backend mlite \ + --hf-path /models/Qwen3.5-35B-A3B \ + --model-name qwen3_5 \ + --truncate-layers 2 \ + --disable-mtp \ + --dry-run + +python experimental/lite/examples/bench/bench.py \ + --backend mbridge \ + --hf-path /models/Qwen3.5-35B-A3B \ + --model-name qwen3_5 \ + --truncate-layers 2 \ + --disable-mtp \ + --dry-run +``` + +The dry-run output is JSON containing the resolved `RuntimeConfig` and session +settings. Use it to confirm the two backend runs differ only on the intended +axis. + +## Pair Script + +```bash +HF_PATH=/models/Qwen3.5-35B-A3B \ +REFERENCE_BACKEND=mbridge \ +DRY_RUN=1 \ +bash experimental/lite/examples/bench/scripts/run_qwen35_pair.sh +``` + +Set `DRY_RUN=0` to run the benchmark under `torchrun`. Results are written to +`experimental/lite/examples/bench/outputs/`. Set `REFERENCE_BACKEND=bridge` only +when Megatron-Bridge is installed and `import megatron.bridge` succeeds. + +## Validated Run + +The following paired run completed on 2026-06-07 with 8x NVIDIA H100 80GB GPUs: + +```bash +HF_PATH=/models/Qwen3.5-35B-A3B \ +OUTPUT_DIR=experimental/lite/examples/bench/outputs/qwen35_pair \ +REFERENCE_BACKEND=mbridge \ +DRY_RUN=0 \ +NPROC=8 \ +MASTER_PORT=31841 \ +MASTER_PORT_BRIDGE=31842 \ +STEPS=15 \ +WARMUP=5 \ +SEQ_LEN=1024 \ +NUM_MICROBATCHES=4 \ +TRUNCATE_LAYERS=8 \ +KEEP_EXPERTS=8 \ +SAME_DATA_ACROSS_DP=1 \ +bash experimental/lite/examples/bench/scripts/run_qwen35_pair.sh +``` + +Slurm job `12624917` completed with exit code `0:0`. The run used +`torch==2.10.0+cu129`; `transformer_engine`, `einops`, and `mbridge` were +available in the runtime environment. + +| Runtime | Impl | Optimizer backend | Measured steps | Avg step ms | Tokens/s | Tokens/s/GPU | Peak memory GB | TFLOPs/GPU | +| --- | --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | +| `mlite` | `lite` | `distopt` | 10 | 309.433 | 105896.935 | 13237.117 | 14.324 | 80.444 | +| `mbridge` | `bridge` | `distopt` | 10 | 332.201 | 98639.089 | 12329.886 | 17.987 | 74.931 | +| `bridge` | `bridge` | `distopt` | 10 | 334.936 | 97833.496 | 12229.187 | 16.403 | 74.319 | + +The two runs used the same synthetic input stream. Loss matched within +`atol=0.05, rtol=0.005` across 10 measured samples +(`max_abs_diff=0.000500`). This long benchmark is performance evidence; the +strict optimizer-metric evidence is the deterministic run below. + +## Deterministic mbridge Correctness + +Slurm job `12630675` completed a strict deterministic MLite vs `mbridge` run +with 1x GPU, `seed=42`, Qwen3.5 MoE, `seq_len=8`, `truncate_layers=1`, +`keep_experts=2`, `optimizer_backend=distopt`, and `steps=2`. + +The comparison passed with bitwise scalar parity and no mismatches: + +- `samples=2` +- `max_loss_abs=0.0` +- `max_grad_norm_abs=0.0` +- `mismatches=[]` +- step 0: `loss=13.027458190917969`, + `grad_norm=120.75512734973202`, post-step weight SHA256 + `1e3176a8cb18d68c5da9bfa5f31a507fa8d51a3a5ddc10fbcd821260a4c6c980` +- step 1: `loss=14.698704719543457`, + `grad_norm=96.15656991334498`, post-step weight SHA256 + `e6d034f7e05ee5ee6a42ceddda8970874c6742baf59cade38f53550abd7aec29` +- eval logits canonical bf16 SHA256 + `2f805802633927852c5ec87455b1afa3a68597e1e411b979f2678eaedfb1c710` + +## Real Run + +```bash +torchrun --nproc_per_node 1 experimental/lite/examples/bench/bench.py \ + --backend mlite \ + --hf-path /models/Qwen3.5-35B-A3B \ + --model-name qwen3_5 \ + --steps 5 \ + --warmup 1 \ + --seq-len 2048 \ + --num-microbatches 1 \ + --truncate-layers 2 \ + --disable-mtp \ + --output-json /tmp/qwen35_mlite_bench.json + +torchrun --nproc_per_node 1 experimental/lite/examples/bench/bench.py \ + --backend mbridge \ + --hf-path /models/Qwen3.5-35B-A3B \ + --model-name qwen3_5 \ + --steps 5 \ + --warmup 1 \ + --seq-len 2048 \ + --num-microbatches 1 \ + --truncate-layers 2 \ + --disable-mtp \ + --output-json /tmp/qwen35_mbridge_bench.json +``` + +Compare `loss`, `grad_norm`, `avg_step_ms`, `tok_per_s`, peak memory, and +`tflops_per_gpu` in the two JSON outputs. Benchmarks are performance evidence; +they are not a replacement for precision tests. + +## Deterministic Correctness + +Use `correctness.py` for strict deterministic parity. It emits exact scalar +fingerprints for `loss` and `grad_norm`, SHA256 fingerprints for logits and +post-step exported weights, and a strict comparison artifact. + +```bash +export MEGATRON_LITE_DETERMINISTIC=1 +export CUBLAS_WORKSPACE_CONFIG=:4096:8 + +torchrun --nproc_per_node 1 experimental/lite/examples/bench/correctness.py run \ + --backend mlite \ + --hf-path /models/Qwen3.5-35B-A3B \ + --model-name qwen3_5 \ + --steps 2 \ + --seq-len 128 \ + --num-microbatches 1 \ + --truncate-layers 2 \ + --disable-mtp \ + --same-data-across-dp \ + --output-json /tmp/qwen35_mlite_correctness.json + +torchrun --nproc_per_node 1 experimental/lite/examples/bench/correctness.py run \ + --backend mbridge \ + --hf-path /models/Qwen3.5-35B-A3B \ + --model-name qwen3_5 \ + --steps 2 \ + --seq-len 128 \ + --num-microbatches 1 \ + --truncate-layers 2 \ + --disable-mtp \ + --same-data-across-dp \ + --output-json /tmp/qwen35_mbridge_correctness.json + +python experimental/lite/examples/bench/correctness.py compare \ + /tmp/qwen35_mlite_correctness.json \ + /tmp/qwen35_mbridge_correctness.json \ + --fail-on-mismatch +``` + +For the PR signoff pair script: + +```bash +export MEGATRON_LITE_DETERMINISTIC=1 +export CUBLAS_WORKSPACE_CONFIG=:4096:8 + +HF_PATH=/models/Qwen3.5-35B-A3B \ +REFERENCE_BACKEND=mbridge \ +DRY_RUN=0 \ +bash experimental/lite/examples/bench/scripts/run_qwen35_correctness_pair.sh +``` diff --git a/experimental/lite/examples/bench/__init__.py b/experimental/lite/examples/bench/__init__.py new file mode 100644 index 00000000000..c910a4bfbf3 --- /dev/null +++ b/experimental/lite/examples/bench/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Benchmark example for Megatron Lite runtime backends.""" + +from .results import RunResult, StepTrace +from .session import PretrainSessionConfig, run_pretrain_session + +__all__ = ["PretrainSessionConfig", "RunResult", "StepTrace", "run_pretrain_session"] diff --git a/experimental/lite/examples/bench/bench.py b/experimental/lite/examples/bench/bench.py new file mode 100644 index 00000000000..be6192c9c20 --- /dev/null +++ b/experimental/lite/examples/bench/bench.py @@ -0,0 +1,443 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Benchmark MLite and reference runtime backends. + +Run from the Megatron-LM repo root after adding ``experimental/lite`` to +``PYTHONPATH``. ``--dry-run`` validates config construction without importing +reference backend packages or initializing distributed state. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from dataclasses import dataclass, fields, is_dataclass, replace +from pathlib import Path +from typing import Any + +_EXPERIMENTAL_LITE_ROOT = Path(__file__).resolve().parents[2] +_REPO_ROOT = Path(__file__).resolve().parents[4] +if str(_EXPERIMENTAL_LITE_ROOT) not in sys.path: + sys.path.insert(0, str(_EXPERIMENTAL_LITE_ROOT)) +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(1, str(_REPO_ROOT)) + +from examples.bench.results import StepTrace +from examples.bench.session import PretrainSessionConfig, run_pretrain_session +from megatron.lite.runtime import RuntimeConfig, create_runtime +from megatron.lite.runtime.backends.bridge.config import BridgeConfig +from megatron.lite.runtime.backends.mlite.config import MegatronLiteConfig +from megatron.lite.runtime.contracts.config import OptimizerConfig, ParallelConfig + + +@dataclass +class BenchCliConfig: + backend: str = "mlite" + hf_path: str = "" + model_name: str = "qwen3_moe" + impl: str = "lite" + tp: int = 1 + etp: int | None = None + ep: int = 1 + pp: int = 1 + vpp: int = 1 + cp: int = 1 + steps: int = 2 + warmup: int = 0 + num_microbatches: int = 1 + seq_len: int = 2048 + seed: int = 42 + device: str = "cuda" + use_thd: bool = False + same_data_across_dp: bool = False + no_optimizer: bool = False + skip_load_hf_weights: bool = False + skip_optimizer_build: bool = False + keep_experts: int | None = None + truncate_layers: int | None = None + disable_mtp: bool = False + optimizer_lr: float = 1e-4 + optimizer_weight_decay: float = 0.1 + optimizer_clip_grad: float = 1.0 + override_ddp_json: str = "{}" + override_transformer_json: str = "{}" + override_optimizer_json: str = "{}" + impl_cfg_json: str = "{}" + dry_run: bool = False + output_json: str | None = None + + +def _json_mapping(raw: str, *, name: str) -> dict[str, Any]: + try: + value = json.loads(raw) + except json.JSONDecodeError as exc: + raise ValueError(f"{name} must be a JSON object: {exc}") from exc + if not isinstance(value, dict): + raise ValueError(f"{name} must be a JSON object.") + return value + + +def _parallel_config(cfg: BenchCliConfig) -> ParallelConfig: + return ParallelConfig(tp=cfg.tp, etp=cfg.etp, ep=cfg.ep, pp=cfg.pp, vpp=cfg.vpp, cp=cfg.cp) + + +def _optimizer_config(cfg: BenchCliConfig) -> OptimizerConfig: + return OptimizerConfig( + lr=cfg.optimizer_lr, + weight_decay=cfg.optimizer_weight_decay, + clip_grad=cfg.optimizer_clip_grad, + adam_beta1=0.9, + adam_beta2=0.95, + adam_eps=1e-8, + use_precision_aware_optimizer=True, + decoupled_weight_decay=True, + ) + + +def _set_field(config_obj: Any, name: str, value: Any) -> Any: + if is_dataclass(config_obj): + return replace(config_obj, **{name: value}) + setattr(config_obj, name, value) + return config_obj + + +def _get_model_config_root(config_obj: Any) -> Any: + return getattr(config_obj, "text_config", config_obj) + + +def _get_bridge_config_root(bridge: Any) -> Any: + hf_pretrained = getattr(bridge, "hf_pretrained", None) + if hf_pretrained is not None: + return _get_model_config_root(getattr(hf_pretrained, "config", hf_pretrained)) + hf_config = getattr(bridge, "hf_config", None) + if hf_config is not None: + return _get_model_config_root(hf_config) + raise ValueError("Bridge object does not expose hf_pretrained or hf_config.") + + +def _refresh_bridge_config(bridge: Any) -> None: + if hasattr(bridge, "_build_config"): + bridge.config = bridge._build_config() + + +def _disable_mtp(config_obj: Any) -> Any: + root = _get_model_config_root(config_obj) + for attr in ("mtp_num_hidden_layers", "num_nextn_predict_layers"): + if hasattr(root, attr): + root = _set_field(root, attr, 0) + if hasattr(root, "mtp_layer_types"): + root = _set_field(root, "mtp_layer_types", []) + if root is not config_obj and hasattr(config_obj, "text_config"): + return _set_field(config_obj, "text_config", root) + return root + + +def _make_mlite_model_config_hook(cfg: BenchCliConfig): + hooks = [] + if cfg.keep_experts is not None: + keep_experts = cfg.keep_experts + + def keep_experts_hook(model_cfg): + old_num = getattr(model_cfg, "num_experts", None) + old_topk = getattr(model_cfg, "num_experts_per_tok", None) + if old_num is None or old_topk is None: + raise ValueError("keep_experts requires model config with MoE expert metadata.") + if keep_experts <= 0 or keep_experts > old_num: + raise ValueError(f"keep_experts must be in [1, {old_num}], got {keep_experts}.") + return replace( + model_cfg, num_experts=keep_experts, num_experts_per_tok=min(old_topk, keep_experts) + ) + + hooks.append(keep_experts_hook) + + if cfg.truncate_layers is not None: + keep_layers = cfg.truncate_layers + + 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 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]), + ) + + hooks.append(truncate_layers_hook) + + if cfg.disable_mtp: + hooks.append(_disable_mtp) + + if not hooks: + return None + + def composed(model_cfg): + for hook in hooks: + model_cfg = hook(model_cfg) + return model_cfg + + return composed + + +def _make_bridge_post_init_hook(cfg: BenchCliConfig): + hooks = [] + if cfg.keep_experts is not None: + keep_experts = cfg.keep_experts + + def keep_experts_hook(bridge) -> None: + hf_cfg = _get_bridge_config_root(bridge) + old_num = getattr(hf_cfg, "num_experts", None) + old_topk = getattr(hf_cfg, "num_experts_per_tok", None) + if old_num is None or old_topk is None: + raise ValueError("keep_experts requires HF config with MoE expert metadata.") + if keep_experts <= 0 or keep_experts > old_num: + raise ValueError(f"keep_experts must be in [1, {old_num}], got {keep_experts}.") + hf_cfg.num_experts = keep_experts + hf_cfg.num_experts_per_tok = min(old_topk, keep_experts) + _refresh_bridge_config(bridge) + + if hasattr(bridge, "_weight_to_mcore_format"): + original = bridge._weight_to_mcore_format + + def patched(name: str, hf_weights: list): + if "mlp.router.weight" in name and len(hf_weights) == 1: + hf_weights = [hf_weights[0][:keep_experts].contiguous()] + return original(name, hf_weights) + + bridge._weight_to_mcore_format = patched + + hooks.append(keep_experts_hook) + + if cfg.truncate_layers is not None: + keep_layers = cfg.truncate_layers + + def truncate_layers_hook(bridge) -> None: + hf_cfg = _get_bridge_config_root(bridge) + old_layers = getattr(hf_cfg, "num_hidden_layers", None) + if old_layers is None: + raise ValueError("truncate_layers requires HF config with 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}." + ) + hf_cfg.num_hidden_layers = keep_layers + if hasattr(hf_cfg, "layer_types"): + hf_cfg.layer_types = list(hf_cfg.layer_types[:keep_layers]) + _refresh_bridge_config(bridge) + + hooks.append(truncate_layers_hook) + + if cfg.disable_mtp: + + def disable_mtp_hook(bridge) -> None: + hf_cfg = _get_bridge_config_root(bridge) + for attr in ("mtp_num_hidden_layers", "num_nextn_predict_layers"): + if hasattr(hf_cfg, attr): + setattr(hf_cfg, attr, 0) + if hasattr(hf_cfg, "mtp_layer_types"): + hf_cfg.mtp_layer_types = [] + _refresh_bridge_config(bridge) + + hooks.append(disable_mtp_hook) + + if not hooks: + return None + + def composed(bridge) -> None: + for hook in hooks: + hook(bridge) + + return composed + + +def build_runtime_config(cfg: BenchCliConfig) -> RuntimeConfig: + parallel = _parallel_config(cfg) + optimizer = _optimizer_config(cfg) + optimizer_overrides = _json_mapping(cfg.override_optimizer_json, name="override_optimizer_json") + + if cfg.backend == "mlite": + for key, value in optimizer_overrides.items(): + setattr(optimizer, key, value) + impl_cfg = _json_mapping(cfg.impl_cfg_json, name="impl_cfg_json") + impl_cfg.setdefault("use_thd", cfg.use_thd) + if cfg.model_name == "qwen3_5": + from megatron.lite.primitive.deterministic import deterministic_requested + + if deterministic_requested(): + impl_cfg.setdefault("mount_vision_model", True) + backend_cfg = MegatronLiteConfig( + model_name=cfg.model_name, + impl=cfg.impl, + hf_path=cfg.hf_path, + parallel=parallel, + optimizer=optimizer, + load_hf_weights=not cfg.skip_load_hf_weights, + impl_cfg=impl_cfg, + model_config_hook=_make_mlite_model_config_hook(cfg), + ) + elif cfg.backend in {"bridge", "mbridge"}: + impl_cfg = _json_mapping(cfg.impl_cfg_json, name="impl_cfg_json") + if impl_cfg: + raise ValueError(f"{cfg.backend} backend does not accept impl_cfg_json.") + backend_cfg = BridgeConfig( + model_name=cfg.model_name, + parallel=parallel, + optimizer=optimizer, + 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"), + override_transformer_config=_json_mapping( + cfg.override_transformer_json, name="override_transformer_json" + ), + override_optimizer_config=optimizer_overrides, + bridge_post_init=_make_bridge_post_init_hook(cfg), + ) + else: + raise ValueError(f"backend must be 'mlite', 'bridge', or 'mbridge', got {cfg.backend!r}.") + + return RuntimeConfig(backend=cfg.backend, hf_path=cfg.hf_path, backend_cfg=backend_cfg) + + +def build_session_config(cfg: BenchCliConfig) -> PretrainSessionConfig: + return PretrainSessionConfig( + steps=cfg.steps, + warmup=cfg.warmup, + num_microbatches=cfg.num_microbatches, + seq_len=cfg.seq_len, + seed=cfg.seed, + device=cfg.device, + use_thd=cfg.use_thd, + same_data_across_dp=cfg.same_data_across_dp, + no_optimizer=cfg.no_optimizer, + ) + + +def _to_jsonable(value: Any) -> Any: + if is_dataclass(value): + return {field.name: _to_jsonable(getattr(value, field.name)) for field in fields(value)} + if isinstance(value, dict): + return {str(k): _to_jsonable(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [_to_jsonable(v) for v in value] + if callable(value): + return f"" + if isinstance(value, Path): + return str(value) + return value + + +def build_dry_run_plan(cfg: BenchCliConfig) -> dict[str, Any]: + return { + "dry_run": True, + "runtime": _to_jsonable(build_runtime_config(cfg)), + "session": _to_jsonable(build_session_config(cfg)), + "notes": [ + "Dry-run validates config construction only.", + "Run under torchrun for real benchmark execution.", + ], + } + + +def _distributed_rank() -> int: + for name in ("RANK", "SLURM_PROCID"): + raw = os.environ.get(name) + if raw is None: + continue + try: + return int(raw) + except ValueError: + continue + return 0 + + +def _step_reporter(trace: StepTrace) -> None: + parts = [ + "[MLITE_BENCH_STEP]", + f"step={trace.step}", + f"loss={trace.loss:.6f}", + f"grad_norm={trace.grad_norm:.6f}", + f"step_ms={trace.step_ms:.3f}", + f"peak_mem_gb={(trace.peak_mem_gb or 0.0):.3f}", + ] + if trace.tflops_per_gpu is not None: + parts.append(f"tflops_per_gpu={trace.tflops_per_gpu:.3f}") + print(" ".join(parts), flush=True) + + +def run(cfg: BenchCliConfig) -> dict[str, Any]: + if cfg.dry_run: + return build_dry_run_plan(cfg) + + rt_cfg = build_runtime_config(cfg) + rt = create_runtime(rt_cfg) + handle = rt.build_model() + result = run_pretrain_session( + rt, handle, build_session_config(cfg), step_reporter=_step_reporter + ) + return result.to_dict() + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--backend", choices=["mlite", "bridge", "mbridge"], default="mlite") + parser.add_argument("--hf-path", default="") + parser.add_argument("--model-name", default="qwen3_moe") + parser.add_argument("--impl", default="lite") + parser.add_argument("--tp", type=int, default=1) + parser.add_argument("--etp", type=int, default=None) + parser.add_argument("--ep", type=int, default=1) + parser.add_argument("--pp", type=int, default=1) + parser.add_argument("--vpp", type=int, default=1) + parser.add_argument("--cp", type=int, default=1) + parser.add_argument("--steps", type=int, default=2) + parser.add_argument("--warmup", type=int, default=0) + parser.add_argument("--num-microbatches", type=int, default=1) + parser.add_argument("--seq-len", type=int, default=2048) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--device", default="cuda") + parser.add_argument("--use-thd", action="store_true") + parser.add_argument("--same-data-across-dp", action="store_true") + parser.add_argument("--no-optimizer", action="store_true") + parser.add_argument("--skip-load-hf-weights", action="store_true") + parser.add_argument("--skip-optimizer-build", action="store_true") + parser.add_argument("--keep-experts", type=int, default=None) + parser.add_argument("--truncate-layers", type=int, default=None) + parser.add_argument("--disable-mtp", action="store_true") + parser.add_argument("--optimizer-lr", type=float, default=1e-4) + parser.add_argument("--optimizer-weight-decay", type=float, default=0.1) + parser.add_argument("--optimizer-clip-grad", type=float, default=1.0) + parser.add_argument("--override-ddp-json", default="{}") + parser.add_argument("--override-transformer-json", default="{}") + parser.add_argument("--override-optimizer-json", default="{}") + parser.add_argument("--impl-cfg-json", default="{}") + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--output-json", default=None) + return parser + + +def parse_args(argv: list[str] | None = None) -> BenchCliConfig: + ns = _parser().parse_args(argv) + return BenchCliConfig(**vars(ns)) + + +def main(argv: list[str] | None = None) -> dict[str, Any]: + cfg = parse_args(argv) + artifact = run(cfg) + if _distributed_rank() == 0: + text = json.dumps(artifact, indent=2, sort_keys=True) + print(text, flush=True) + if cfg.output_json is not None: + output_path = Path(cfg.output_json) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(text + "\n", encoding="utf-8") + return artifact + + +if __name__ == "__main__": + main() diff --git a/experimental/lite/examples/bench/correctness.py b/experimental/lite/examples/bench/correctness.py new file mode 100644 index 00000000000..fd534fdc05f --- /dev/null +++ b/experimental/lite/examples/bench/correctness.py @@ -0,0 +1,518 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Deterministic correctness runner for MLite and reference backends.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import struct +import sys +from contextlib import contextmanager +from pathlib import Path +from typing import Any + +import torch + +_EXPERIMENTAL_LITE_ROOT = Path(__file__).resolve().parents[2] +_REPO_ROOT = Path(__file__).resolve().parents[4] +if str(_EXPERIMENTAL_LITE_ROOT) not in sys.path: + sys.path.insert(0, str(_EXPERIMENTAL_LITE_ROOT)) +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(1, str(_REPO_ROOT)) + +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: + for name in ("RANK", "SLURM_PROCID"): + raw = os.environ.get(name) + if raw is None: + continue + try: + return int(raw) + except ValueError: + continue + return 0 + + +def _sync(device: str) -> None: + if device.startswith("cuda") and torch.cuda.is_available(): + torch.cuda.synchronize() + + +def _scalar(value: float | int | torch.Tensor | None) -> dict[str, Any]: + if isinstance(value, torch.Tensor): + if value.numel() != 1: + raise ValueError("scalar fingerprint requires a scalar tensor.") + value = float(value.detach().cpu().float().item()) + value_f = float(0.0 if value is None else value) + return { + "value": value_f, + "float_hex": value_f.hex(), + "sha256_f64_be": hashlib.sha256(struct.pack(">d", value_f)).hexdigest(), + } + + +def _hash_tensor(tensor: torch.Tensor | None) -> dict[str, Any] | None: + if tensor is None: + return None + t = tensor.detach().contiguous().cpu() + raw = t.view(torch.uint8).numpy().tobytes() + as_bf16 = t.to(torch.bfloat16).contiguous() if t.is_floating_point() else None + summary = t.float() if t.is_floating_point() else None + result = { + "shape": list(t.shape), + "dtype": str(t.dtype), + "sha256": hashlib.sha256(raw).hexdigest(), + } + if as_bf16 is not None: + result["sha256_as_bf16"] = hashlib.sha256( + as_bf16.view(torch.uint8).numpy().tobytes() + ).hexdigest() + if summary is not None: + flat = summary.reshape(-1) + result["summary"] = { + "min": float(flat.min().item()) if flat.numel() else 0.0, + "max": float(flat.max().item()) if flat.numel() else 0.0, + "mean": float(flat.mean().item()) if flat.numel() else 0.0, + "first8": [float(x) for x in flat[:8].tolist()], + } + return result + + +def _first_tensor(value: Any) -> torch.Tensor | None: + if isinstance(value, torch.Tensor): + return value + if isinstance(value, dict): + for item in value.values(): + tensor = _first_tensor(item) + if tensor is not None: + return tensor + return None + if isinstance(value, (list, tuple)): + for item in value: + tensor = _first_tensor(item) + if tensor is not None: + return tensor + return None + + +def _record_activation_probe( + records: list[dict[str, Any]], + name: str, + output: Any, + *, + record_grad: bool = False, + tensor_hooks: list[Any] | None = None, +) -> None: + tensor = _first_tensor(output) + record = {"name": name, "found": True, "tensor": _hash_tensor(tensor)} + if record_grad: + record["grad"] = None + record["grad_found"] = isinstance(tensor, torch.Tensor) and tensor.requires_grad + if isinstance(tensor, torch.Tensor) and tensor.requires_grad: + + def _grad_hook(grad, _record=record): + _record["grad"] = _hash_tensor(grad) + + hook = tensor.register_hook(_grad_hook) + if tensor_hooks is not None: + tensor_hooks.append(hook) + records.append(record) + + +def _resolve_probe_module(modules: dict[str, Any], name: str) -> tuple[str, Any] | None: + module = modules.get(name) + if module is not None: + return name, module + matches = [ + (candidate_name, candidate) + for candidate_name, candidate in modules.items() + if candidate_name.endswith(f".{name}") + ] + if len(matches) == 1: + return matches[0] + return None + + +@contextmanager +def _activation_probe_context(handle, probe_names: list[str], *, record_grad: bool = False): + records: list[dict[str, Any]] = [] + hooks = [] + tensor_hooks = [] + patched_methods = [] + modules = dict(handle._model.named_modules()) + for name in probe_names: + probe_name = name + record_input = probe_name.endswith(":input") + lookup_name = probe_name[:-6] if record_input else probe_name + if "::" in lookup_name: + module_name, method_name = lookup_name.split("::", 1) + resolved = _resolve_probe_module(modules, module_name) + if resolved is None or not hasattr(resolved[1], method_name): + records.append({"name": name, "found": False}) + continue + resolved_name, module = resolved + original = getattr(module, method_name) + + def _wrapped( + *args, + _original=original, + _probe_name=name, + _resolved_name=resolved_name, + _module_type=type(module).__module__ + "." + type(module).__qualname__, + _record_input=record_input, + _method_name=method_name, + **kwargs, + ): + if _record_input: + _record_activation_probe( + records, + _probe_name, + args, + record_grad=record_grad, + tensor_hooks=tensor_hooks, + ) + records[-1]["resolved_name"] = f"{_resolved_name}::{_method_name}:input" + records[-1]["module_type"] = _module_type + return _original(*args, **kwargs) + output = _original(*args, **kwargs) + _record_activation_probe( + records, _probe_name, output, record_grad=record_grad, tensor_hooks=tensor_hooks + ) + records[-1]["resolved_name"] = f"{_resolved_name}::{_method_name}" + records[-1]["module_type"] = _module_type + return output + + setattr(module, method_name, _wrapped) + patched_methods.append((module, method_name, original)) + continue + + resolved = _resolve_probe_module(modules, lookup_name) + if resolved is None: + records.append({"name": name, "found": False}) + continue + resolved_name, module = resolved + + if record_input: + + def _pre_hook(_module, args, probe_name=name, resolved_probe_name=resolved_name): + _record_activation_probe( + records, probe_name, args, record_grad=record_grad, tensor_hooks=tensor_hooks + ) + records[-1]["resolved_name"] = f"{resolved_probe_name}:input" + records[-1]["module_type"] = ( + type(_module).__module__ + "." + type(_module).__qualname__ + ) + + hooks.append(module.register_forward_pre_hook(_pre_hook)) + continue + + def _hook(_module, _args, output, probe_name=name, resolved_probe_name=resolved_name): + _record_activation_probe( + records, probe_name, output, record_grad=record_grad, tensor_hooks=tensor_hooks + ) + records[-1]["resolved_name"] = resolved_probe_name + records[-1]["module_type"] = type(_module).__module__ + "." + type(_module).__qualname__ + + hooks.append(module.register_forward_hook(_hook)) + try: + yield records + finally: + for hook in tensor_hooks: + hook.remove() + for hook in hooks: + hook.remove() + for module, method_name, original in patched_methods: + setattr(module, method_name, original) + + +def _update_hash_with_tensor(h: Any, name: str, tensor: torch.Tensor) -> None: + t = tensor.detach().contiguous().cpu() + h.update(name.encode("utf-8")) + h.update(b"\0") + h.update(str(t.dtype).encode("ascii")) + h.update(b"\0") + h.update(json.dumps(list(t.shape), separators=(",", ":")).encode("ascii")) + h.update(b"\0") + h.update(t.view(torch.uint8).numpy().tobytes()) + h.update(b"\0") + + +def _model_chunks(handle) -> list[Any]: + chunks = handle._extras.get("model_chunks") + if chunks is None: + chunks = handle._extras.get("model_list") + if chunks is None: + chunks = [handle._model] + return list(chunks) + + +def _grad_fingerprint(handle) -> dict[str, Any]: + h = hashlib.sha256() + count = 0 + details = [] + include_details = os.environ.get("MLITE_CORRECTNESS_GRAD_DETAILS") == "1" + for chunk_idx, chunk in enumerate(_model_chunks(handle)): + for name, param in sorted(chunk.named_parameters(), key=lambda item: item[0]): + grad = param.grad + if grad is None: + grad = getattr(param, "main_grad", None) + if grad is None: + continue + fingerprint_name = f"{chunk_idx}:{name}" + _update_hash_with_tensor(h, fingerprint_name, grad) + if include_details: + detail = _hash_tensor(grad) + assert detail is not None + detail["name"] = fingerprint_name + details.append(detail) + count += 1 + result = {"sha256": h.hexdigest(), "tensor_count": count} + if include_details: + result["details"] = details + return result + + +def _weight_fingerprint(rt, handle) -> dict[str, Any]: + h = hashlib.sha256() + count = 0 + details = [] + include_details = os.environ.get("MLITE_CORRECTNESS_WEIGHT_DETAILS") == "1" + for name, tensor in sorted(rt.export_weights(handle, cpu=True), key=lambda item: item[0]): + _update_hash_with_tensor(h, str(name), tensor) + if include_details: + detail = _hash_tensor(tensor) + assert detail is not None + detail["name"] = str(name) + details.append(detail) + count += 1 + result = {"sha256": h.hexdigest(), "tensor_count": count} + if include_details: + result["details"] = details + 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 + + +def run_backend( + cfg: BenchCliConfig, + *, + hash_weights: bool = True, + activation_probe_names: list[str] | None = None, +) -> dict[str, Any]: + os.environ["MEGATRON_LITE_DETERMINISTIC"] = "1" + set_deterministic(cfg.seed) + + rt_cfg = build_runtime_config(cfg) + rt = create_runtime(rt_cfg) + handle = rt.build_model() + session_cfg = build_session_config(cfg) + + eval_iter = _make_data_iter(handle, session_cfg) + eval_batch = next(eval_iter) + activation_probe_names = list(activation_probe_names or []) + with _activation_probe_context(handle, activation_probe_names) as activation_probes: + with rt.eval_mode(handle): + eval_logits = _hash_tensor(_forward_logits(rt, handle, eval_batch)) + + data_iter = _make_data_iter(handle, session_cfg) + steps: list[dict[str, Any]] = [] + with rt.train_mode(handle): + for step in range(session_cfg.steps): + with _activation_probe_context( + handle, activation_probe_names, record_grad=True + ) as train_activation_probes: + rt.zero_grad(handle) + _sync(session_cfg.device) + result = rt.forward_backward( + handle, data_iter, loss_fn=None, num_microbatches=session_cfg.num_microbatches + ) + _sync(session_cfg.device) + logits = _hash_tensor(result.model_output.vocab_parallel_logits) + grads = _grad_fingerprint(handle) + + if session_cfg.no_optimizer: + update_successful, grad_norm, num_zeros = True, 0.0, 0 + else: + update_successful, grad_norm, num_zeros = rt.optimizer_step(handle) + rt.lr_scheduler_step(handle) + _sync(session_cfg.device) + + steps.append( + { + "step": step, + "loss": _scalar(result.metrics.get("loss", 0.0)), + "logits": logits, + "grad_fingerprint": grads, + "grad_norm": _scalar(grad_norm), + "update_successful": bool(update_successful), + "num_zeros": None if num_zeros is None else int(num_zeros), + "post_step_weights": _weight_fingerprint(rt, handle) if hash_weights else None, + "train_activation_probes": train_activation_probes, + } + ) + + return { + "kind": "mlite_bench_correctness", + "backend": cfg.backend, + "model_name": cfg.model_name, + "seed": cfg.seed, + "seq_len": cfg.seq_len, + "num_microbatches": cfg.num_microbatches, + "steps": steps, + "eval_logits": eval_logits, + "activation_probes": activation_probes, + "metadata": { + "deterministic": True, + "hash_weights": hash_weights, + "same_data_across_dp": cfg.same_data_across_dp, + "use_thd": cfg.use_thd, + }, + } + + +def _add_run_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--backend", choices=["mlite", "bridge", "mbridge"], required=True) + parser.add_argument("--hf-path", required=True) + parser.add_argument("--model-name", default="qwen3_5") + parser.add_argument("--impl", default="lite") + parser.add_argument("--tp", type=int, default=1) + parser.add_argument("--etp", type=int, default=None) + parser.add_argument("--ep", type=int, default=1) + parser.add_argument("--pp", type=int, default=1) + parser.add_argument("--vpp", type=int, default=1) + parser.add_argument("--cp", type=int, default=1) + parser.add_argument("--steps", type=int, default=2) + parser.add_argument("--warmup", type=int, default=0) + parser.add_argument("--num-microbatches", type=int, default=1) + parser.add_argument("--seq-len", type=int, default=128) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--device", default="cuda") + parser.add_argument("--use-thd", action="store_true") + parser.add_argument("--same-data-across-dp", action="store_true") + parser.add_argument("--no-optimizer", action="store_true") + parser.add_argument("--skip-load-hf-weights", action="store_true") + parser.add_argument("--skip-optimizer-build", action="store_true") + parser.add_argument("--keep-experts", type=int, default=None) + parser.add_argument("--truncate-layers", type=int, default=None) + parser.add_argument("--disable-mtp", action="store_true") + parser.add_argument("--optimizer-lr", type=float, default=1e-4) + parser.add_argument("--optimizer-weight-decay", type=float, default=0.1) + parser.add_argument("--optimizer-clip-grad", type=float, default=1.0) + parser.add_argument("--override-ddp-json", default="{}") + parser.add_argument("--override-transformer-json", default="{}") + parser.add_argument("--override-optimizer-json", default="{}") + parser.add_argument("--impl-cfg-json", default="{}") + parser.add_argument("--output-json", required=True) + parser.add_argument("--skip-weight-hash", action="store_true") + parser.add_argument("--activation-probes-json", default="{}") + + +def _activation_probe_names(raw: str, backend: str) -> list[str]: + value = json.loads(raw) + if isinstance(value, list): + return [str(item) for item in value] + if isinstance(value, dict): + selected = value.get(backend, []) + if not isinstance(selected, list): + raise ValueError(f"activation probe list for {backend!r} must be a JSON list.") + return [str(item) for item in selected] + raise ValueError("activation_probes_json must be a JSON list or backend-to-list mapping.") + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="command", required=True) + run_p = sub.add_parser("run", help="run one backend and write a correctness artifact") + _add_run_args(run_p) + + cmp_p = sub.add_parser("compare", help="strictly compare two correctness artifacts") + cmp_p.add_argument("baseline") + cmp_p.add_argument("candidate") + cmp_p.add_argument("--output-json", default=None) + cmp_p.add_argument("--fail-on-mismatch", action="store_true") + return parser + + +def main(argv: list[str] | None = None) -> dict[str, Any]: + ns = _parser().parse_args(argv) + if ns.command == "compare": + result = compare_correctness_artifacts( + load_result_artifact(ns.baseline), load_result_artifact(ns.candidate) + ) + text = json.dumps(result, indent=2, sort_keys=True) + print(text, flush=True) + if ns.output_json: + Path(ns.output_json).write_text(text + "\n", encoding="utf-8") + if ns.fail_on_mismatch and not result["passed"]: + raise SystemExit(1) + return result + + cfg = BenchCliConfig( + **{k: v for k, v in vars(ns).items() if k in BenchCliConfig.__dataclass_fields__} + ) + artifact = run_backend( + cfg, + hash_weights=not ns.skip_weight_hash, + activation_probe_names=_activation_probe_names(ns.activation_probes_json, cfg.backend), + ) + if _distributed_rank() == 0: + text = json.dumps(artifact, indent=2, sort_keys=True) + print(text, flush=True) + output_path = Path(ns.output_json) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(text + "\n", encoding="utf-8") + return artifact + + +if __name__ == "__main__": + main() diff --git a/experimental/lite/examples/bench/results.py b/experimental/lite/examples/bench/results.py new file mode 100644 index 00000000000..6a015c10458 --- /dev/null +++ b/experimental/lite/examples/bench/results.py @@ -0,0 +1,239 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Benchmark result dataclasses.""" + +from __future__ import annotations + +import json +import math +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + + +@dataclass(slots=True) +class StepTrace: + step: int + loss: float + grad_norm: float + step_ms: float + peak_mem_gb: float | None = None + tflops_per_gpu: float | None = None + + def to_dict(self) -> dict[str, Any]: + result: dict[str, Any] = { + "step": self.step, + "loss": self.loss, + "grad_norm": self.grad_norm, + "step_ms": self.step_ms, + } + if self.peak_mem_gb is not None: + result["peak_mem_gb"] = self.peak_mem_gb + if self.tflops_per_gpu is not None: + result["tflops_per_gpu"] = self.tflops_per_gpu + return result + + +@dataclass(slots=True) +class RunResult: + backend: str + model_name: str + impl: str + optimizer_backend: str + tp: int + etp: int | None + ep: int + pp: int + vpp: int + cp: int + seq_len: int + num_microbatches: int + step_traces: list[StepTrace] = field(default_factory=list) + avg_step_ms: float = 0.0 + peak_mem_gb: float = 0.0 + tok_per_s: float = 0.0 + tok_per_s_per_gpu: float = 0.0 + tflops_per_gpu: float | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + def summary_dict(self) -> dict[str, Any]: + return { + "backend": self.backend, + "model_name": self.model_name, + "impl": self.impl, + "optimizer_backend": self.optimizer_backend, + "avg_step_ms": self.avg_step_ms, + "tok_per_s": self.tok_per_s, + "tok_per_s_per_gpu": self.tok_per_s_per_gpu, + "peak_mem_gb": self.peak_mem_gb, + "tflops_per_gpu": self.tflops_per_gpu, + "steps_measured": len(self.step_traces), + } + + def to_dict(self) -> dict[str, Any]: + return { + "summary": self.summary_dict(), + "result": { + "backend": self.backend, + "model_name": self.model_name, + "impl": self.impl, + "optimizer_backend": self.optimizer_backend, + "tp": self.tp, + "etp": self.etp, + "ep": self.ep, + "pp": self.pp, + "vpp": self.vpp, + "cp": self.cp, + "seq_len": self.seq_len, + "num_microbatches": self.num_microbatches, + "step_traces": [trace.to_dict() for trace in self.step_traces], + "avg_step_ms": self.avg_step_ms, + "peak_mem_gb": self.peak_mem_gb, + "tok_per_s": self.tok_per_s, + "tok_per_s_per_gpu": self.tok_per_s_per_gpu, + "tflops_per_gpu": self.tflops_per_gpu, + "metadata": dict(self.metadata), + }, + } + + +def load_result_artifact(path: str | Path) -> dict[str, Any]: + """Load a benchmark JSON artifact from ``bench.py --output-json``.""" + with Path(path).open(encoding="utf-8") as f: + value = json.load(f) + if not isinstance(value, dict): + raise ValueError(f"Benchmark artifact must be a JSON object: {path}") + return value + + +def result_summary(artifact: dict[str, Any]) -> dict[str, Any]: + """Return the summary block from a benchmark artifact.""" + summary = artifact.get("summary") + if isinstance(summary, dict): + return dict(summary) + result = artifact.get("result") + if not isinstance(result, dict): + raise ValueError("Benchmark artifact must contain `summary` or `result`.") + return { + "backend": result.get("backend"), + "model_name": result.get("model_name"), + "impl": result.get("impl"), + "optimizer_backend": result.get("optimizer_backend"), + "avg_step_ms": result.get("avg_step_ms"), + "tok_per_s": result.get("tok_per_s"), + "tok_per_s_per_gpu": result.get("tok_per_s_per_gpu"), + "peak_mem_gb": result.get("peak_mem_gb"), + "tflops_per_gpu": result.get("tflops_per_gpu"), + "steps_measured": len(result.get("step_traces", [])), + } + + +def compare_step_traces( + baseline: dict[str, Any], candidate: dict[str, Any], *, atol: float = 1e-4, rtol: float = 1e-4 +) -> dict[str, Any]: + """Compare loss and grad-norm traces from two benchmark artifacts.""" + base_steps = baseline.get("result", {}).get("step_traces", []) + cand_steps = candidate.get("result", {}).get("step_traces", []) + sample_count = min(len(base_steps), len(cand_steps)) + max_loss_abs = 0.0 + max_grad_norm_abs = 0.0 + for idx in range(sample_count): + base = base_steps[idx] + cand = cand_steps[idx] + max_loss_abs = max(max_loss_abs, abs(float(base["loss"]) - float(cand["loss"]))) + max_grad_norm_abs = max( + max_grad_norm_abs, abs(float(base["grad_norm"]) - float(cand["grad_norm"])) + ) + + lengths_match = sample_count == len(base_steps) == len(cand_steps) + loss_ref_max = max([abs(float(step["loss"])) for step in base_steps[:sample_count]] + [0.0]) + grad_norm_ref_max = max( + [abs(float(step["grad_norm"])) for step in base_steps[:sample_count]] + [0.0] + ) + loss_passed = lengths_match and max_loss_abs <= atol + rtol * loss_ref_max + grad_norm_passed = lengths_match and max_grad_norm_abs <= atol + rtol * grad_norm_ref_max + + return { + "samples": sample_count, + "atol": atol, + "rtol": rtol, + "passed": loss_passed and grad_norm_passed, + "loss_passed": loss_passed, + "grad_norm_passed": grad_norm_passed, + "max_loss_abs": max_loss_abs, + "max_grad_norm_abs": max_grad_norm_abs, + } + + +def compare_correctness_artifacts( + baseline: dict[str, Any], candidate: dict[str, Any] +) -> dict[str, Any]: + """Strict bitwise comparison for deterministic correctness artifacts.""" + base_steps = baseline.get("steps", []) + cand_steps = candidate.get("steps", []) + sample_count = min(len(base_steps), len(cand_steps)) + lengths_match = sample_count == len(base_steps) == len(cand_steps) + + max_loss_abs = 0.0 + max_grad_norm_abs = 0.0 + mismatches: list[dict[str, Any]] = [] + + def _tensor_fingerprint_matches(base: Any, cand: Any) -> bool: + if base == cand: + return True + if not isinstance(base, dict) or not isinstance(cand, dict): + return False + if base.get("shape") != cand.get("shape"): + return False + base_bf16 = base.get("sha256_as_bf16") + cand_bf16 = cand.get("sha256_as_bf16") + return bool(base_bf16 and base_bf16 == cand_bf16) + + base_eval = baseline.get("eval_logits") + cand_eval = candidate.get("eval_logits") + if base_eval is not None or cand_eval is not None: + if not _tensor_fingerprint_matches(base_eval, cand_eval): + mismatches.append({"field": "eval_logits"}) + + for idx in range(sample_count): + base = base_steps[idx] + cand = cand_steps[idx] + loss_abs = abs(float(base["loss"]["value"]) - float(cand["loss"]["value"])) + grad_abs = abs(float(base["grad_norm"]["value"]) - float(cand["grad_norm"]["value"])) + if math.isfinite(loss_abs): + max_loss_abs = max(max_loss_abs, loss_abs) + if math.isfinite(grad_abs): + max_grad_norm_abs = max(max_grad_norm_abs, grad_abs) + + for field in ("loss", "grad_norm", "post_step_weights", "update_successful", "num_zeros"): + if base.get(field) != cand.get(field): + mismatches.append({"step": idx, "field": field}) + if not _tensor_fingerprint_matches(base.get("logits"), cand.get("logits")): + mismatches.append({"step": idx, "field": "logits"}) + + if not lengths_match: + mismatches.append( + { + "field": "steps", + "baseline_count": len(base_steps), + "candidate_count": len(cand_steps), + } + ) + + return { + "samples": sample_count, + "passed": lengths_match and not mismatches, + "max_loss_abs": max_loss_abs, + "max_grad_norm_abs": max_grad_norm_abs, + "tensor_fingerprint_rule": "raw_sha256_or_bf16_canonical_sha256", + "mismatches": mismatches, + } + + +__all__ = [ + "RunResult", + "StepTrace", + "compare_correctness_artifacts", + "compare_step_traces", + "load_result_artifact", + "result_summary", +] diff --git a/experimental/lite/examples/bench/scripts/run_qwen35_correctness_pair.sh b/experimental/lite/examples/bench/scripts/run_qwen35_correctness_pair.sh new file mode 100755 index 00000000000..5fa5eff9c90 --- /dev/null +++ b/experimental/lite/examples/bench/scripts/run_qwen35_correctness_pair.sh @@ -0,0 +1,73 @@ +#!/usr/bin/env bash +set -euo pipefail + +HF_PATH=${HF_PATH:?set HF_PATH to a HuggingFace Qwen3.5 model directory} +REPO_ROOT=${REPO_ROOT:-$(pwd)} +PYTHON_BIN=${PYTHON_BIN:-python} +NPROC=${NPROC:-1} +OUTPUT_DIR=${OUTPUT_DIR:-"${REPO_ROOT}/experimental/lite/examples/bench/outputs/correctness"} +REFERENCE_BACKEND=${REFERENCE_BACKEND:-mbridge} + +export PYTHONPATH="${REPO_ROOT}/experimental/lite:${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" +export MEGATRON_LITE_DETERMINISTIC=1 +export CUBLAS_WORKSPACE_CONFIG=${CUBLAS_WORKSPACE_CONFIG:-:4096:8} + +mkdir -p "${OUTPUT_DIR}" + +COMMON_ARGS=( + --hf-path "${HF_PATH}" + --model-name qwen3_5 + --tp "${TP:-1}" + --etp "${ETP:-1}" + --ep "${EP:-1}" + --pp "${PP:-1}" + --cp "${CP:-1}" + --steps "${STEPS:-2}" + --num-microbatches "${NUM_MICROBATCHES:-1}" + --seq-len "${SEQ_LEN:-128}" + --seed "${SEED:-42}" + --truncate-layers "${TRUNCATE_LAYERS:-2}" + --disable-mtp + --same-data-across-dp +) + +if [[ -n "${KEEP_EXPERTS:-}" ]]; then + COMMON_ARGS+=(--keep-experts "${KEEP_EXPERTS}") +fi +if [[ "${SKIP_LOAD_HF_WEIGHTS:-0}" == "1" ]]; then + COMMON_ARGS+=(--skip-load-hf-weights) +fi +if [[ "${SKIP_WEIGHT_HASH:-0}" == "1" ]]; then + COMMON_ARGS+=(--skip-weight-hash) +fi +if [[ -n "${ACTIVATION_PROBES_JSON:-}" ]]; then + COMMON_ARGS+=(--activation-probes-json "${ACTIVATION_PROBES_JSON}") +fi + +MLITE_TORCHRUN_ARGS=(--nproc_per_node "${NPROC}") +BRIDGE_TORCHRUN_ARGS=(--nproc_per_node "${NPROC}") +if [[ -n "${MASTER_PORT:-}" ]]; then + MLITE_TORCHRUN_ARGS+=(--master_port "${MASTER_PORT}") +fi +if [[ -n "${MASTER_PORT_BRIDGE:-}" ]]; then + BRIDGE_TORCHRUN_ARGS+=(--master_port "${MASTER_PORT_BRIDGE}") +fi + +torchrun "${MLITE_TORCHRUN_ARGS[@]}" \ + "${REPO_ROOT}/experimental/lite/examples/bench/correctness.py" run \ + --backend mlite "${COMMON_ARGS[@]}" \ + --output-json "${OUTPUT_DIR}/qwen35_mlite_correctness.json" \ + 2>&1 | tee "${OUTPUT_DIR}/qwen35_mlite_correctness.log" + +torchrun "${BRIDGE_TORCHRUN_ARGS[@]}" \ + "${REPO_ROOT}/experimental/lite/examples/bench/correctness.py" run \ + --backend "${REFERENCE_BACKEND}" "${COMMON_ARGS[@]}" \ + --output-json "${OUTPUT_DIR}/qwen35_${REFERENCE_BACKEND}_correctness.json" \ + 2>&1 | tee "${OUTPUT_DIR}/qwen35_${REFERENCE_BACKEND}_correctness.log" + +"${PYTHON_BIN}" "${REPO_ROOT}/experimental/lite/examples/bench/correctness.py" compare \ + "${OUTPUT_DIR}/qwen35_mlite_correctness.json" \ + "${OUTPUT_DIR}/qwen35_${REFERENCE_BACKEND}_correctness.json" \ + --output-json "${OUTPUT_DIR}/qwen35_correctness_compare.json" \ + --fail-on-mismatch \ + 2>&1 | tee "${OUTPUT_DIR}/qwen35_correctness_compare.log" diff --git a/experimental/lite/examples/bench/scripts/run_qwen35_pair.sh b/experimental/lite/examples/bench/scripts/run_qwen35_pair.sh new file mode 100755 index 00000000000..2c2b7e8eb49 --- /dev/null +++ b/experimental/lite/examples/bench/scripts/run_qwen35_pair.sh @@ -0,0 +1,68 @@ +#!/usr/bin/env bash +set -euo pipefail + +HF_PATH=${HF_PATH:?set HF_PATH to a HuggingFace Qwen3.5 model directory} +REPO_ROOT=${REPO_ROOT:-$(pwd)} +PYTHON_BIN=${PYTHON_BIN:-python} +NPROC=${NPROC:-1} +DRY_RUN=${DRY_RUN:-1} +OUTPUT_DIR=${OUTPUT_DIR:-"${REPO_ROOT}/experimental/lite/examples/bench/outputs"} +REFERENCE_BACKEND=${REFERENCE_BACKEND:-mbridge} + +export PYTHONPATH="${REPO_ROOT}/experimental/lite:${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" + +COMMON_ARGS=( + --hf-path "${HF_PATH}" + --model-name qwen3_5 + --tp "${TP:-1}" + --etp "${ETP:-1}" + --ep "${EP:-1}" + --pp "${PP:-1}" + --cp "${CP:-1}" + --steps "${STEPS:-2}" + --warmup "${WARMUP:-0}" + --num-microbatches "${NUM_MICROBATCHES:-1}" + --seq-len "${SEQ_LEN:-2048}" + --truncate-layers "${TRUNCATE_LAYERS:-2}" + --disable-mtp +) + +if [[ -n "${KEEP_EXPERTS:-}" ]]; then + COMMON_ARGS+=(--keep-experts "${KEEP_EXPERTS}") +fi +if [[ "${SAME_DATA_ACROSS_DP:-0}" == "1" ]]; then + COMMON_ARGS+=(--same-data-across-dp) +fi +if [[ "${SKIP_LOAD_HF_WEIGHTS:-0}" == "1" ]]; then + COMMON_ARGS+=(--skip-load-hf-weights) +fi +if [[ "${SKIP_OPTIMIZER_BUILD:-0}" == "1" ]]; then + COMMON_ARGS+=(--skip-optimizer-build --no-optimizer) +fi + +if [[ "${DRY_RUN}" == "1" ]]; then + "${PYTHON_BIN}" "${REPO_ROOT}/experimental/lite/examples/bench/bench.py" \ + --backend mlite "${COMMON_ARGS[@]}" --dry-run + "${PYTHON_BIN}" "${REPO_ROOT}/experimental/lite/examples/bench/bench.py" \ + --backend "${REFERENCE_BACKEND}" "${COMMON_ARGS[@]}" --dry-run +else + mkdir -p "${OUTPUT_DIR}" + MLITE_TORCHRUN_ARGS=(--nproc_per_node "${NPROC}") + BRIDGE_TORCHRUN_ARGS=(--nproc_per_node "${NPROC}") + if [[ -n "${MASTER_PORT:-}" ]]; then + MLITE_TORCHRUN_ARGS+=(--master_port "${MASTER_PORT}") + fi + if [[ -n "${MASTER_PORT_BRIDGE:-}" ]]; then + BRIDGE_TORCHRUN_ARGS+=(--master_port "${MASTER_PORT_BRIDGE}") + fi + torchrun "${MLITE_TORCHRUN_ARGS[@]}" \ + "${REPO_ROOT}/experimental/lite/examples/bench/bench.py" \ + --backend mlite "${COMMON_ARGS[@]}" \ + --output-json "${OUTPUT_DIR}/qwen35_mlite.json" \ + 2>&1 | tee "${OUTPUT_DIR}/qwen35_mlite.log" + torchrun "${BRIDGE_TORCHRUN_ARGS[@]}" \ + "${REPO_ROOT}/experimental/lite/examples/bench/bench.py" \ + --backend "${REFERENCE_BACKEND}" "${COMMON_ARGS[@]}" \ + --output-json "${OUTPUT_DIR}/qwen35_${REFERENCE_BACKEND}.json" \ + 2>&1 | tee "${OUTPUT_DIR}/qwen35_${REFERENCE_BACKEND}.log" +fi diff --git a/experimental/lite/examples/bench/session.py b/experimental/lite/examples/bench/session.py new file mode 100644 index 00000000000..47e7a25789a --- /dev/null +++ b/experimental/lite/examples/bench/session.py @@ -0,0 +1,256 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Small pretrain benchmark session composed from runtime atoms.""" + +from __future__ import annotations + +import importlib +import time +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import torch + +from megatron.lite.runtime.backends import Runtime +from megatron.lite.runtime.contracts.handle import ModelHandle + +from .results import RunResult, StepTrace + + +@dataclass +class PretrainSessionConfig: + steps: int = 2 + warmup: int = 0 + num_microbatches: int = 1 + seq_len: int = 2048 + seed: int = 42 + device: str = "cuda" + use_thd: bool = False + same_data_across_dp: bool = False + no_optimizer: bool = False + + +def _is_cuda_device(device: str) -> bool: + return device.startswith("cuda") and torch.cuda.is_available() + + +def _sync(device: str) -> None: + if _is_cuda_device(device): + torch.cuda.synchronize() + + +def _reset_peak_memory(device: str) -> None: + if _is_cuda_device(device): + torch.cuda.reset_peak_memory_stats() + + +def _peak_memory_gb(device: str) -> float: + if _is_cuda_device(device): + return torch.cuda.max_memory_allocated() / 1e9 + return 0.0 + + +def _world_size() -> int: + if torch.distributed.is_available() and torch.distributed.is_initialized(): + return torch.distributed.get_world_size() + return 1 + + +def _resolve_vocab_size(handle: ModelHandle) -> int: + proto = handle._extras.get("protocol") + model_cfg = handle._extras.get("model_cfg") + if proto is not None and model_cfg is not None and hasattr(proto, "vocab_size"): + return int(proto.vocab_size(model_cfg)) + if model_cfg is not None and hasattr(model_cfg, "vocab_size"): + return int(model_cfg.vocab_size) + 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, + ) + + from megatron.lite.primitive.data import infinite_batches + + return infinite_batches(vocab_size, cfg.seq_len, device=cfg.device, seed=data_seed) + + +def _calc_tflops_per_gpu( + *, + num_floating_point_operations: int | None, + activated_params: int | None, + tokens_per_step: int, + step_s: float, + world_size: int, +) -> float | None: + if step_s <= 0: + return None + if num_floating_point_operations: + return num_floating_point_operations / (step_s * world_size * 1e12) + if activated_params: + return 6 * activated_params * tokens_per_step / (step_s * world_size * 1e12) + return None + + +def _resolve_model_stats(config: Any, proto: Any) -> Any | None: + model_name = getattr(config, "model_name", None) + if model_name and model_name != "auto": + stats_module = f"megatron.lite.model.{model_name}.stats" + try: + return importlib.import_module(stats_module) + except ModuleNotFoundError as exc: + if exc.name is not None and not stats_module.startswith(exc.name): + raise + return proto + + +def _resolve_step_flops( + handle: ModelHandle, cfg: PretrainSessionConfig +) -> tuple[int | None, int | None]: + config = handle.config + proto = handle._extras.get("protocol") + model_stats = _resolve_model_stats(config, proto) + model_cfg = handle._extras.get("model_cfg") + if model_stats is None or model_cfg is None: + return None, None + + step_flops = None + if hasattr(model_stats, "num_floating_point_operations"): + parallel_cfg = getattr(config, "parallel", None) + tp_size = getattr(parallel_cfg, "tp", 1) + step_flops = model_stats.num_floating_point_operations( + model_cfg, + seq_len=cfg.seq_len, + global_batch_size=cfg.num_microbatches * handle.dp_size, + tp_size=tp_size, + ) + + activated_params = None + if step_flops is None and hasattr(model_stats, "activated_params"): + activated_params = model_stats.activated_params(model_cfg) + + return step_flops, activated_params + + +def run_pretrain_session( + rt: Runtime, + handle: ModelHandle, + cfg: PretrainSessionConfig, + *, + data_iter: Any = None, + step_reporter: Callable[[StepTrace], None] | None = None, +) -> RunResult: + """Run a fixed-shape benchmark loop through the public runtime API.""" + if cfg.steps < 1: + raise ValueError("steps must be >= 1") + if cfg.warmup < 0 or cfg.warmup >= cfg.steps: + raise ValueError("warmup must satisfy 0 <= warmup < steps") + if cfg.num_microbatches < 1: + raise ValueError("num_microbatches must be >= 1") + + if data_iter is None: + data_iter = _make_data_iter(handle, cfg) + + world_size = _world_size() + tokens_per_step = cfg.num_microbatches * cfg.seq_len * world_size + step_flops, activated_params = _resolve_step_flops(handle, cfg) + + step_traces: list[StepTrace] = [] + timings: list[float] = [] + + _reset_peak_memory(cfg.device) + with rt.train_mode(handle): + for step in range(cfg.steps): + if step == cfg.warmup: + _reset_peak_memory(cfg.device) + + rt.zero_grad(handle) + _sync(cfg.device) + t0 = time.perf_counter() + result = rt.forward_backward( + handle, data_iter, loss_fn=None, num_microbatches=cfg.num_microbatches + ) + if cfg.no_optimizer: + grad_norm = 0.0 + else: + _, grad_norm, _ = rt.optimizer_step(handle) + rt.lr_scheduler_step(handle) + _sync(cfg.device) + + elapsed_ms = (time.perf_counter() - t0) * 1000 + tflops_per_gpu = _calc_tflops_per_gpu( + num_floating_point_operations=step_flops, + activated_params=activated_params, + tokens_per_step=tokens_per_step, + step_s=elapsed_ms / 1000, + world_size=world_size, + ) + trace = StepTrace( + step=step, + loss=float(result.metrics.get("loss", 0.0)), + grad_norm=float(grad_norm), + step_ms=elapsed_ms, + peak_mem_gb=_peak_memory_gb(cfg.device), + tflops_per_gpu=tflops_per_gpu, + ) + if step_reporter is not None: + step_reporter(trace) + if step >= cfg.warmup: + timings.append(elapsed_ms) + trace.step = step - cfg.warmup + step_traces.append(trace) + + avg_step_ms = sum(timings) / len(timings) if timings else 0.0 + avg_step_s = avg_step_ms / 1000 + tok_per_s = tokens_per_step / avg_step_s if avg_step_s > 0 else 0.0 + avg_tflops = _calc_tflops_per_gpu( + num_floating_point_operations=step_flops, + activated_params=activated_params, + tokens_per_step=tokens_per_step, + step_s=avg_step_s, + world_size=world_size, + ) + + config = handle.config + parallel = config.parallel + backend = "bridge" if type(config).__name__ == "BridgeConfig" else "mlite" + return RunResult( + backend=backend, + model_name=getattr(config, "model_name", "unknown"), + impl=getattr(config, "impl", "bridge"), + optimizer_backend=handle._extras.get( + "optimizer_backend", + getattr(handle._optimizer, "name", "none") if handle._optimizer is not None else "none", + ), + tp=parallel.tp, + etp=parallel.etp, + ep=parallel.ep, + pp=parallel.pp, + vpp=parallel.vpp, + cp=parallel.cp, + seq_len=cfg.seq_len, + num_microbatches=cfg.num_microbatches, + step_traces=step_traces, + avg_step_ms=avg_step_ms, + peak_mem_gb=_peak_memory_gb(cfg.device), + tok_per_s=tok_per_s, + tok_per_s_per_gpu=tok_per_s / world_size, + tflops_per_gpu=avg_tflops, + metadata={"warmup": cfg.warmup, "device": cfg.device, "use_thd": cfg.use_thd}, + ) + + +__all__ = ["PretrainSessionConfig", "run_pretrain_session"] diff --git a/experimental/lite/examples/verl/README.md b/experimental/lite/examples/verl/README.md new file mode 100644 index 00000000000..09a54454366 --- /dev/null +++ b/experimental/lite/examples/verl/README.md @@ -0,0 +1,192 @@ +# VERL Megatron Lite Example + +This directory contains a runnable VERL external engine integration for +Megatron Lite plus Qwen3.5-35B-A3B SFT and GRPO launch scripts. + +The Python package is `verl_mlite`. It registers VERL's language-model engine +backend as `mlite`, while Megatron Lite model implementations still use +`impl=lite`. + +## Layout + +- `verl_mlite/engine/mlite_engine.py`: VERL `BaseEngine` implementation backed + by `megatron.lite.runtime`. +- `verl_mlite/config/engine/mlite.yaml`: Hydra engine config for + `engine=mlite`. +- `scripts/run_qwen3moe_sft.sh`: Qwen MoE SFT launcher using + `verl.trainer.sft_trainer`. +- `scripts/run_qwen3moe_gsm8k_sft.sh`: GSM8K wrapper around the SFT launcher. +- `scripts/run_qwen3moe_gsm8k_grpo.sh`: GSM8K GRPO launcher with MLite actor + training and a standard VERL rollout backend. + +## Prerequisites + +Install or expose these packages before running: + +- VERL with the new engine worker path. + See [`REQUIRED_VERL.txt`](REQUIRED_VERL.txt) for the reference upstream + release tag. +- Megatron-LM from this repository, or another source tree via + `MEGATRON_ROOT=/path/to/Megatron-LM`. +- Megatron Lite from this repository. The script automatically adds + `experimental/lite` to `PYTHONPATH`. +- The examples directory is also added to `PYTHONPATH` and loads a local + compatibility hook for known VERL/vLLM/Transformers dependency gaps. + +Optional source-tree override: + +```bash +export VERL_ROOT=/path/to/verl +export MEGATRON_ROOT=/path/to/Megatron-LM +``` + +## SFT + +The SFT script expects VERL messages-format parquet input. + +```bash +export MODEL_PATH=/path/to/qwen3.5-35b-a3b-hf +export TRAIN_FILES=/path/to/train.parquet +export VAL_FILES=/path/to/val.parquet + +bash experimental/lite/examples/verl/scripts/run_qwen3moe_sft.sh +``` + +Useful knobs: + +- `TP_SIZE`, `PP_SIZE`, `VPP_SIZE`, `CP_SIZE`, `EP_SIZE`, `ETP_SIZE` +- `TOTAL_STEPS`, `TOTAL_EPOCHS`, `TRAIN_BATCH_SIZE`, `MICRO_BATCH_SIZE` +- `MAX_TOKENS_PER_GPU`, `MAX_LENGTH`, `MESSAGES_KEY` +- `PARAM_OFFLOAD`, `OPTIMIZER_OFFLOAD`, `GRAD_OFFLOAD` +- `MLITE_MODEL_NAME=auto`, `MLITE_IMPL=lite` +- `ATTENTION_BACKEND=flash` +- `DRY_RUN=1` to print the resolved `torchrun` command without launching + +FSDP2 supports two offload modes. `PARAM_OFFLOAD=True` and +`OPTIMIZER_OFFLOAD=True` move model parameters and optimizer state between CPU +and GPU when VERL switches execution contexts. `OPTIMIZER_OFFLOAD=True` also +sets `optim.override_optimizer_config.offload_fraction=1.0` by default, which +keeps FSDP2 optimizer update state on CPU during forward/backward to reduce GPU +memory pressure. + +Example dry run: + +```bash +MODEL_PATH=/path/to/qwen3.5-35b-a3b-hf \ +TRAIN_FILES=/path/to/train.parquet \ +DRY_RUN=1 \ +bash experimental/lite/examples/verl/scripts/run_qwen3moe_sft.sh +``` + +By default, logs, command snapshots, JSONL logger output, and checkpoints are +written under `experimental/lite/examples/verl/outputs/qwen3moe_sft`. Override +`OUTPUT_ROOT`, `LOG_FILE`, `JSONL_FILE`, `CMD_FILE`, or `CKPT_DIR` to redirect +artifacts. + +For local dry runs, prefer a temporary output directory if you do not want +command snapshots under the source tree: + +```bash +OUTPUT_ROOT="$(mktemp -d)" \ +MODEL_PATH=/path/to/qwen3.5-35b-a3b-hf \ +TRAIN_FILES=/path/to/train.parquet \ +DRY_RUN=1 \ +bash experimental/lite/examples/verl/scripts/run_qwen3moe_sft.sh +``` + +## GSM8K SFT + +Build messages-format GSM8K parquet files with VERL's SFT preprocessor: + +```bash +python3 /path/to/verl/examples/data_preprocess/gsm8k_multiturn_sft.py \ + --local_save_dir ~/data/gsm8k_sft +``` + +Run the MLite GSM8K SFT wrapper: + +```bash +MODEL_PATH=Qwen/Qwen3.5-35B-A3B \ +DRY_RUN=1 \ +bash experimental/lite/examples/verl/scripts/run_qwen3moe_gsm8k_sft.sh +``` + +The wrapper defaults to `Qwen/Qwen3.5-35B-A3B`, +`~/data/gsm8k_sft/train.parquet`, and +`~/data/gsm8k_sft/test.parquet`, then delegates to +`scripts/run_qwen3moe_sft.sh`. Override `DATASET_DIR`, `TRAIN_FILES`, or +`VAL_FILES` to use another location. + +By default, GSM8K SFT artifacts are written under +`experimental/lite/examples/verl/outputs/qwen35_gsm8k_sft`. + +## GSM8K GRPO + +Build RL-format GSM8K parquet files with VERL's GRPO/PPO preprocessor: + +```bash +python3 /path/to/verl/examples/data_preprocess/gsm8k.py \ + --local_save_dir ~/data/gsm8k +``` + +Run GRPO with the MLite actor and vLLM rollout: + +```bash +MODEL_PATH=Qwen/Qwen3.5-35B-A3B \ +DRY_RUN=1 \ +bash experimental/lite/examples/verl/scripts/run_qwen3moe_gsm8k_grpo.sh +``` + +Useful GRPO knobs: + +- `TRAIN_BATCH_SIZE`, `PPO_MINI_BATCH_SIZE`, + `ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU` +- `MAX_PROMPT_LENGTH`, `MAX_RESPONSE_LENGTH`, `PPO_MAX_TOKEN_LEN_PER_GPU` +- `ROLLOUT_N`, `ROLLOUT_TP`, `ROLLOUT_GPU_MEMORY_UTILIZATION` +- `ROLLOUT_MODE=async`, `ROLLOUT_MAX_MODEL_LEN`, `ROLLOUT_MAX_NUM_BATCHED_TOKENS` +- `ROLLOUT_LIMIT_IMAGES=0`, `ROLLOUT_LIMIT_VIDEOS=0` keep the vLLM rollout + backend in text-only mode for GSM8K by default. +- `ACTOR_TP`, `ACTOR_PP`, `ACTOR_VPP`, `ACTOR_CP`, `ACTOR_EP`, `ACTOR_ETP` +- `PARAM_OFFLOAD`, `OPTIMIZER_OFFLOAD`, `GRAD_OFFLOAD` +- `INFER_BACKEND=vllm` +- `USE_LEGACY_WORKER_IMPL=disable` to use VERL's new engine worker path +- `POLICY_LOSS_MODE=vanilla` and `LOSS_AGG_MODE=seq-mean-token-sum-norm` + select the pure GRPO baseline policy loss and aggregation mode. + +The GRPO launcher keeps the reference policy disabled by default +(`algorithm.use_kl_in_reward=False`, `actor_rollout_ref.actor.use_kl_loss=False`) +so the example exercises the current MLite actor path without expanding scope +to a separate reference model. It also disables VERL's legacy worker path by +default so `actor@actor_rollout_ref.actor=mlite_actor` is handled by the new +engine worker implementation. + +By default, GSM8K GRPO artifacts are written under +`experimental/lite/examples/verl/outputs/qwen35_gsm8k_grpo`. + +## Smoke / Dry-Run Checks + +Checked on this branch on 2026-06-07. These checks cover shell syntax, +Python import compilation, and resolved command construction only; they do not +cover end-to-end SFT or GRPO training. + +- Shell syntax: + - `bash -n experimental/lite/examples/verl/scripts/run_qwen3moe_sft.sh` + - `bash -n experimental/lite/examples/verl/scripts/run_qwen3moe_gsm8k_sft.sh` + - `bash -n experimental/lite/examples/verl/scripts/run_qwen3moe_gsm8k_grpo.sh` +- Python import compilation: + - `PYTHONPYCACHEPREFIX="$(mktemp -d)" python3 -m compileall -q experimental/lite/examples/verl/verl_mlite` +- GSM8K SFT dry run: + - `OUTPUT_ROOT="$(mktemp -d)" MODEL_PATH=Qwen/Qwen3.5-35B-A3B DRY_RUN=1 bash experimental/lite/examples/verl/scripts/run_qwen3moe_gsm8k_sft.sh` + - Dry-run output shows `torchrun -m verl.trainer.sft_trainer`, + `engine=mlite`, `model.path=Qwen/Qwen3.5-35B-A3B`, + `data.train_files=${HOME}/data/gsm8k_sft/train.parquet`, and + `data.val_files=${HOME}/data/gsm8k_sft/test.parquet`. +- GSM8K GRPO dry run: + - `OUTPUT_ROOT="$(mktemp -d)" MODEL_PATH=Qwen/Qwen3.5-35B-A3B DRY_RUN=1 bash experimental/lite/examples/verl/scripts/run_qwen3moe_gsm8k_grpo.sh` + - Dry-run output shows `python3 -m verl.trainer.main_ppo`, + `actor@actor_rollout_ref.actor=mlite_actor`, + `actor_rollout_ref.rollout.name=vllm`, + `actor_rollout_ref.actor.engine.impl=lite`, + `actor_rollout_ref.actor.engine.ep=8`, + `algorithm.adv_estimator=grpo`, `actor_rollout_ref.actor.policy_loss.loss_mode=vanilla`, + `critic.enable=False`, and `trainer.use_legacy_worker_impl=disable`. diff --git a/experimental/lite/examples/verl/REQUIRED_VERL.txt b/experimental/lite/examples/verl/REQUIRED_VERL.txt new file mode 100644 index 00000000000..0bec47f376d --- /dev/null +++ b/experimental/lite/examples/verl/REQUIRED_VERL.txt @@ -0,0 +1,3 @@ +# Reference VERL release for the Megatron Lite VERL example. +# Latest GitHub release checked on 2026-06-06: v0.8.0. +verl @ git+https://github.com/verl-project/verl.git@v0.8.0 diff --git a/experimental/lite/examples/verl/scripts/run_qwen3moe_gsm8k_grpo.sh b/experimental/lite/examples/verl/scripts/run_qwen3moe_gsm8k_grpo.sh new file mode 100755 index 00000000000..c20f1464db6 --- /dev/null +++ b/experimental/lite/examples/verl/scripts/run_qwen3moe_gsm8k_grpo.sh @@ -0,0 +1,319 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ "${VERBOSE:-0}" == "1" ]]; then + set -x +fi + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -L)" +EXAMPLE_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd -L)" +LITE_ROOT="$(cd "${EXAMPLE_ROOT}/../.." && pwd -L)" +REPO_ROOT="$(cd "${LITE_ROOT}/../.." && pwd -L)" + +add_pythonpath() { + local path="${1:-}" + if [[ -n "${path}" ]]; then + export PYTHONPATH="${path}:${PYTHONPATH:-}" + fi +} + +add_pythonpath "${EXAMPLE_ROOT}" +add_pythonpath "${LITE_ROOT}" +add_pythonpath "${REPO_ROOT}" +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 + +DATASET_DIR="${DATASET_DIR:-${HOME}/data/gsm8k}" +MODEL_PATH="${MODEL_PATH:-Qwen/Qwen3.5-35B-A3B}" +TRAIN_FILES="${TRAIN_FILES:-${DATASET_DIR}/train.parquet}" +VAL_FILES="${VAL_FILES:-${DATASET_DIR}/test.parquet}" + +OUTPUT_ROOT="${OUTPUT_ROOT:-${EXAMPLE_ROOT}/outputs/qwen35_gsm8k_grpo}" +PROJECT_NAME="${PROJECT_NAME:-verl-mlite-qwen35-gsm8k-grpo}" +INFER_BACKEND="${INFER_BACKEND:-vllm}" + +NNODES="${NNODES:-1}" +NGPUS_PER_NODE="${NGPUS_PER_NODE:-${NPROC_PER_NODE:-8}}" + +TRAIN_BATCH_SIZE="${TRAIN_BATCH_SIZE:-128}" +PPO_MINI_BATCH_SIZE="${PPO_MINI_BATCH_SIZE:-32}" +ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU="${ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU:-1}" +MAX_PROMPT_LENGTH="${MAX_PROMPT_LENGTH:-512}" +MAX_RESPONSE_LENGTH="${MAX_RESPONSE_LENGTH:-1024}" +PPO_MAX_TOKEN_LEN_PER_GPU="${PPO_MAX_TOKEN_LEN_PER_GPU:-8192}" + +ROLLOUT_N="${ROLLOUT_N:-5}" +ROLLOUT_MODE="${ROLLOUT_MODE:-async}" +ROLLOUT_TP="${ROLLOUT_TP:-2}" +ROLLOUT_GPU_MEMORY_UTILIZATION="${ROLLOUT_GPU_MEMORY_UTILIZATION:-0.6}" +ROLLOUT_LOG_PROB_MICRO_BATCH_SIZE_PER_GPU="${ROLLOUT_LOG_PROB_MICRO_BATCH_SIZE_PER_GPU:-1}" +ROLLOUT_LOG_PROB_MAX_TOKEN_LEN_PER_GPU="${ROLLOUT_LOG_PROB_MAX_TOKEN_LEN_PER_GPU:-${PPO_MAX_TOKEN_LEN_PER_GPU}}" +ROLLOUT_MAX_MODEL_LEN="${ROLLOUT_MAX_MODEL_LEN:-$((MAX_PROMPT_LENGTH + MAX_RESPONSE_LENGTH))}" +ROLLOUT_MAX_NUM_SEQS="${ROLLOUT_MAX_NUM_SEQS:-1024}" +ROLLOUT_DEFAULT_MAX_NUM_BATCHED_TOKENS="$((MAX_PROMPT_LENGTH + MAX_RESPONSE_LENGTH))" +if (( ROLLOUT_DEFAULT_MAX_NUM_BATCHED_TOKENS < ROLLOUT_MAX_NUM_SEQS )); then + ROLLOUT_DEFAULT_MAX_NUM_BATCHED_TOKENS="${ROLLOUT_MAX_NUM_SEQS}" +fi +ROLLOUT_MAX_NUM_BATCHED_TOKENS="${ROLLOUT_MAX_NUM_BATCHED_TOKENS:-${ROLLOUT_DEFAULT_MAX_NUM_BATCHED_TOKENS}}" +if (( ROLLOUT_MAX_NUM_BATCHED_TOKENS < ROLLOUT_MAX_NUM_SEQS )); then + ROLLOUT_MAX_NUM_BATCHED_TOKENS="${ROLLOUT_MAX_NUM_SEQS}" +fi +ROLLOUT_TEMPERATURE="${ROLLOUT_TEMPERATURE:-1.0}" +ROLLOUT_TOP_P="${ROLLOUT_TOP_P:-1.0}" +ROLLOUT_TOP_K="${ROLLOUT_TOP_K:--1}" +ROLLOUT_LIMIT_IMAGES="${ROLLOUT_LIMIT_IMAGES:-0}" +ROLLOUT_LIMIT_VIDEOS="${ROLLOUT_LIMIT_VIDEOS:-0}" +VAL_TEMPERATURE="${VAL_TEMPERATURE:-0.0}" +VAL_TOP_P="${VAL_TOP_P:-1.0}" +VAL_DO_SAMPLE="${VAL_DO_SAMPLE:-False}" +VAL_N="${VAL_N:-1}" + +ACTOR_TP="${ACTOR_TP:-2}" +ACTOR_PP="${ACTOR_PP:-1}" +ACTOR_VPP="${ACTOR_VPP:-null}" +ACTOR_CP="${ACTOR_CP:-1}" +ACTOR_EP="${ACTOR_EP:-8}" +ACTOR_ETP="${ACTOR_ETP:-1}" +DTYPE="${DTYPE:-bfloat16}" +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. +# - fsdp2: Megatron Lite FSDP2 wrapper + optimizer. +MLITE_OPTIMIZER_BACKEND="${MLITE_OPTIMIZER_BACKEND:-distopt}" + +ACTOR_LR="${ACTOR_LR:-1e-6}" +POLICY_LOSS_MODE="${POLICY_LOSS_MODE:-vanilla}" +LOSS_AGG_MODE="${LOSS_AGG_MODE:-seq-mean-token-sum-norm}" +WEIGHT_DECAY="${WEIGHT_DECAY:-0.1}" +BETAS="${BETAS:-[0.9,0.95]}" +CLIP_GRAD="${CLIP_GRAD:-1.0}" +LR_WARMUP_STEPS="${LR_WARMUP_STEPS:-0}" +LR_DECAY_STYLE="${LR_DECAY_STYLE:-constant}" +ENTROPY_COEFF="${ENTROPY_COEFF:-0}" +USE_DYNAMIC_BSZ="${USE_DYNAMIC_BSZ:-True}" +PARAM_OFFLOAD="${PARAM_OFFLOAD:-False}" +OPTIMIZER_OFFLOAD="${OPTIMIZER_OFFLOAD:-True}" +GRAD_OFFLOAD="${GRAD_OFFLOAD:-False}" +OPTIMIZER_STATE_OFFLOAD_FRACTION="${OPTIMIZER_STATE_OFFLOAD_FRACTION:-1.0}" +USE_PRECISION_AWARE_OPTIMIZER="${USE_PRECISION_AWARE_OPTIMIZER:-True}" +DECOUPLED_WEIGHT_DECAY="${DECOUPLED_WEIGHT_DECAY:-True}" + +TOTAL_EPOCHS="${TOTAL_EPOCHS:-15}" +TOTAL_TRAINING_STEPS="${TOTAL_TRAINING_STEPS:-null}" +SAVE_FREQ="${SAVE_FREQ:-20}" +TEST_FREQ="${TEST_FREQ:-5}" +RESUME_MODE="${RESUME_MODE:-auto}" +RESUME_FROM_PATH="${RESUME_FROM_PATH:-null}" +LOG_VAL_GENERATIONS="${LOG_VAL_GENERATIONS:-10}" +LOGGER="${LOGGER:-[console,file]}" +USE_LEGACY_WORKER_IMPL="${USE_LEGACY_WORKER_IMPL:-disable}" +DRY_RUN="${DRY_RUN:-0}" +EXTRA_ARGS=("$@") + +if [[ "${INFER_BACKEND}" != "vllm" && "${INFER_BACKEND}" != "sglang" && "${INFER_BACKEND}" != "trtllm" ]]; then + echo "Unsupported INFER_BACKEND=${INFER_BACKEND}. Expected vllm, sglang, or trtllm." >&2 + exit 1 +fi + +case "${MLITE_OPTIMIZER_BACKEND}" in + distopt) + MLITE_IMPL_OPTIMIZER="mc" + ;; + fsdp2) + MLITE_IMPL_OPTIMIZER="fsdp2" + ;; + *) + echo "Unsupported MLITE_OPTIMIZER_BACKEND=${MLITE_OPTIMIZER_BACKEND}. Expected distopt or fsdp2." >&2 + exit 1 + ;; +esac + +if [[ "${INFER_BACKEND}" == "vllm" ]]; then + export VLLM_USE_V1="${VLLM_USE_V1:-1}" + export VLLM_ALLREDUCE_USE_SYMM_MEM="${VLLM_ALLREDUCE_USE_SYMM_MEM:-0}" +fi + +MLITE_VPP_SIZE="${ACTOR_VPP}" +if [[ "${MLITE_VPP_SIZE}" == "null" ]]; then + MLITE_VPP_SIZE=1 +fi + +RUN_NAME="${RUN_NAME:-qwen35_gsm8k_grpo_mlite_${INFER_BACKEND}_tp${ACTOR_TP}_pp${ACTOR_PP}_cp${ACTOR_CP}_ep${ACTOR_EP}}" +CKPT_DIR="${CKPT_DIR:-${OUTPUT_ROOT}/checkpoints/${RUN_NAME}}" +LOG_FILE="${LOG_FILE:-${OUTPUT_ROOT}/${RUN_NAME}.log}" +JSONL_FILE="${JSONL_FILE:-${OUTPUT_ROOT}/${RUN_NAME}.jsonl}" +CMD_FILE="${CMD_FILE:-${OUTPUT_ROOT}/${RUN_NAME}.cmd.sh}" + +mkdir -p "${OUTPUT_ROOT}" "${CKPT_DIR}" "$(dirname "${LOG_FILE}")" "$(dirname "${JSONL_FILE}")" "$(dirname "${CMD_FILE}")" +export VERL_FILE_LOGGER_PATH="${JSONL_FILE}" + +CACHE_ROOT="${VERL_MLITE_CACHE_ROOT:-${TMPDIR:-/tmp}/verl_mlite}" +mkdir -p "${CACHE_ROOT}/pycache_${USER:-user}" "${CACHE_ROOT}/torchinductor_${USER:-user}" "${CACHE_ROOT}/triton_${USER:-user}" +export PYTHONPYCACHEPREFIX="${PYTHONPYCACHEPREFIX:-${CACHE_ROOT}/pycache_${USER:-user}}" +export TORCHINDUCTOR_CACHE_DIR="${TORCHINDUCTOR_CACHE_DIR:-${CACHE_ROOT}/torchinductor_${USER:-user}}" +export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-${CACHE_ROOT}/triton_${USER:-user}}" + +ALGORITHM=( + "algorithm.adv_estimator=grpo" + "algorithm.use_kl_in_reward=False" + "algorithm.kl_ctrl.kl_coef=0.0" + "algorithm.rollout_correction.bypass_mode=True" + "algorithm.norm_adv_by_std_in_grpo=False" +) + +DATA=( + "data.train_files=${TRAIN_FILES}" + "data.val_files=${VAL_FILES}" + "data.train_batch_size=${TRAIN_BATCH_SIZE}" + "data.prompt_key=prompt" + "data.return_raw_chat=True" + "data.max_prompt_length=${MAX_PROMPT_LENGTH}" + "data.max_response_length=${MAX_RESPONSE_LENGTH}" + "data.filter_overlong_prompts=True" + "data.truncation=error" +) + +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" +) + +ACTOR=( + "actor@actor_rollout_ref.actor=mlite_actor" + "actor_rollout_ref.actor.optim.lr=${ACTOR_LR}" + "actor_rollout_ref.actor.optim.weight_decay=${WEIGHT_DECAY}" + "actor_rollout_ref.actor.optim.betas=${BETAS}" + "actor_rollout_ref.actor.optim.clip_grad=${CLIP_GRAD}" + "actor_rollout_ref.actor.optim.lr_warmup_steps=${LR_WARMUP_STEPS}" + "actor_rollout_ref.actor.optim.lr_warmup_init=0" + "actor_rollout_ref.actor.optim.lr_decay_style=${LR_DECAY_STYLE}" + "actor_rollout_ref.actor.ppo_mini_batch_size=${PPO_MINI_BATCH_SIZE}" + "actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${ACTOR_PPO_MICRO_BATCH_SIZE_PER_GPU}" + "actor_rollout_ref.actor.use_dynamic_bsz=${USE_DYNAMIC_BSZ}" + "actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN_PER_GPU}" + "actor_rollout_ref.actor.use_kl_loss=False" + "actor_rollout_ref.actor.kl_loss_coef=0.0" + "actor_rollout_ref.actor.entropy_coeff=${ENTROPY_COEFF}" + "actor_rollout_ref.actor.policy_loss.loss_mode=${POLICY_LOSS_MODE}" + "actor_rollout_ref.actor.loss_agg_mode=${LOSS_AGG_MODE}" + "actor_rollout_ref.actor.engine.dtype=${DTYPE}" + "actor_rollout_ref.actor.engine.model_name=${MLITE_MODEL_NAME}" + "actor_rollout_ref.actor.engine.impl=${MLITE_IMPL}" + "actor_rollout_ref.actor.engine.tp=${ACTOR_TP}" + "actor_rollout_ref.actor.engine.pp=${ACTOR_PP}" + "actor_rollout_ref.actor.engine.vpp=${MLITE_VPP_SIZE}" + "actor_rollout_ref.actor.engine.cp=${ACTOR_CP}" + "actor_rollout_ref.actor.engine.ep=${ACTOR_EP}" + "actor_rollout_ref.actor.engine.etp=${ACTOR_ETP}" + "actor_rollout_ref.actor.engine.param_offload=${PARAM_OFFLOAD}" + "actor_rollout_ref.actor.engine.optimizer_offload=${OPTIMIZER_OFFLOAD}" + "actor_rollout_ref.actor.engine.grad_offload=${GRAD_OFFLOAD}" + "actor_rollout_ref.actor.engine.attention_backend_override=${ATTENTION_BACKEND}" + "actor_rollout_ref.actor.engine.impl_cfg.use_thd=True" + "+actor_rollout_ref.actor.engine.impl_cfg.optimizer=${MLITE_IMPL_OPTIMIZER}" +) + +if [[ "${OPTIMIZER_OFFLOAD}" == "True" || "${OPTIMIZER_OFFLOAD}" == "true" || "${OPTIMIZER_OFFLOAD}" == "1" ]]; then + ACTOR+=( + "+actor_rollout_ref.actor.optim.override_optimizer_config.offload_fraction=${OPTIMIZER_STATE_OFFLOAD_FRACTION}" + "+actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=${USE_PRECISION_AWARE_OPTIMIZER}" + "+actor_rollout_ref.actor.optim.override_optimizer_config.decoupled_weight_decay=${DECOUPLED_WEIGHT_DECAY}" + ) +fi + +ROLLOUT=( + "actor_rollout_ref.rollout.name=${INFER_BACKEND}" + "actor_rollout_ref.rollout.mode=${ROLLOUT_MODE}" + "actor_rollout_ref.rollout.tensor_model_parallel_size=${ROLLOUT_TP}" + "actor_rollout_ref.rollout.gpu_memory_utilization=${ROLLOUT_GPU_MEMORY_UTILIZATION}" + "actor_rollout_ref.rollout.n=${ROLLOUT_N}" + "actor_rollout_ref.rollout.calculate_log_probs=True" + "actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${USE_DYNAMIC_BSZ}" + "actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${ROLLOUT_LOG_PROB_MAX_TOKEN_LEN_PER_GPU}" + "actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${ROLLOUT_LOG_PROB_MICRO_BATCH_SIZE_PER_GPU}" + "actor_rollout_ref.rollout.prompt_length=${MAX_PROMPT_LENGTH}" + "actor_rollout_ref.rollout.response_length=${MAX_RESPONSE_LENGTH}" + "actor_rollout_ref.rollout.max_model_len=${ROLLOUT_MAX_MODEL_LEN}" + "actor_rollout_ref.rollout.max_num_seqs=${ROLLOUT_MAX_NUM_SEQS}" + "actor_rollout_ref.rollout.max_num_batched_tokens=${ROLLOUT_MAX_NUM_BATCHED_TOKENS}" + "actor_rollout_ref.rollout.temperature=${ROLLOUT_TEMPERATURE}" + "actor_rollout_ref.rollout.top_p=${ROLLOUT_TOP_P}" + "actor_rollout_ref.rollout.top_k=${ROLLOUT_TOP_K}" + "actor_rollout_ref.rollout.val_kwargs.temperature=${VAL_TEMPERATURE}" + "actor_rollout_ref.rollout.val_kwargs.top_p=${VAL_TOP_P}" + "actor_rollout_ref.rollout.val_kwargs.do_sample=${VAL_DO_SAMPLE}" + "actor_rollout_ref.rollout.val_kwargs.n=${VAL_N}" + "actor_rollout_ref.rollout.free_cache_engine=True" +) + +if [[ "${INFER_BACKEND}" == "vllm" ]]; then + ROLLOUT+=( + "+actor_rollout_ref.rollout.engine_kwargs.vllm.limit_mm_per_prompt.image=${ROLLOUT_LIMIT_IMAGES}" + "+actor_rollout_ref.rollout.engine_kwargs.vllm.limit_mm_per_prompt.video=${ROLLOUT_LIMIT_VIDEOS}" + ) +fi + +TRAINER=( + "critic.enable=False" + "trainer.balance_batch=True" + "trainer.logger=${LOGGER}" + "trainer.project_name=${PROJECT_NAME}" + "trainer.experiment_name=${RUN_NAME}" + "trainer.n_gpus_per_node=${NGPUS_PER_NODE}" + "trainer.nnodes=${NNODES}" + "trainer.save_freq=${SAVE_FREQ}" + "trainer.test_freq=${TEST_FREQ}" + "trainer.total_epochs=${TOTAL_EPOCHS}" + "trainer.total_training_steps=${TOTAL_TRAINING_STEPS}" + "trainer.resume_mode=${RESUME_MODE}" + "trainer.resume_from_path=${RESUME_FROM_PATH}" + "trainer.default_local_dir=${CKPT_DIR}" + "trainer.val_before_train=False" + "trainer.log_val_generations=${LOG_VAL_GENERATIONS}" + "trainer.use_legacy_worker_impl=${USE_LEGACY_WORKER_IMPL}" +) + +COMMAND=( + python3 + -m + verl.trainer.main_ppo + "hydra.searchpath=[pkg://verl_mlite.config]" + "${ALGORITHM[@]}" + "${DATA[@]}" + "${MODEL[@]}" + "${ACTOR[@]}" + "${ROLLOUT[@]}" + "${TRAINER[@]}" + "${EXTRA_ARGS[@]}" +) + +printf '%q ' "${COMMAND[@]}" > "${CMD_FILE}" +printf '\n' >> "${CMD_FILE}" + +if [[ "${DRY_RUN}" == "1" ]]; then + printf '%q ' "${COMMAND[@]}" + printf '\n' + exit 0 +fi + +echo "[mlite] output_root=${OUTPUT_ROOT}" +echo "[mlite] log=${LOG_FILE}" +echo "[mlite] jsonl=${JSONL_FILE}" +echo "[mlite] cmd=${CMD_FILE}" +echo "[mlite] optimizer_backend=${MLITE_OPTIMIZER_BACKEND} impl_cfg.optimizer=${MLITE_IMPL_OPTIMIZER}" + +set +e +"${COMMAND[@]}" 2>&1 | tee "${LOG_FILE}" +cmd_rc="${PIPESTATUS[0]}" +set -e +exit "${cmd_rc}" diff --git a/experimental/lite/examples/verl/scripts/run_qwen3moe_gsm8k_sft.sh b/experimental/lite/examples/verl/scripts/run_qwen3moe_gsm8k_sft.sh new file mode 100755 index 00000000000..9e59d530e01 --- /dev/null +++ b/experimental/lite/examples/verl/scripts/run_qwen3moe_gsm8k_sft.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ "${VERBOSE:-0}" == "1" ]]; then + set -x +fi + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -L)" + +DATASET_DIR="${DATASET_DIR:-${HOME}/data/gsm8k_sft}" +export MODEL_PATH="${MODEL_PATH:-Qwen/Qwen3.5-35B-A3B}" +export TRAIN_FILES="${TRAIN_FILES:-${DATASET_DIR}/train.parquet}" +export VAL_FILES="${VAL_FILES:-${DATASET_DIR}/test.parquet}" +export OUTPUT_ROOT="${OUTPUT_ROOT:-${SCRIPT_DIR}/../outputs/qwen35_gsm8k_sft}" +export PROJECT_NAME="${PROJECT_NAME:-verl-mlite-qwen35-gsm8k-sft}" +export RUN_NAME="${RUN_NAME:-qwen35_gsm8k_sft_mlite}" + +export TOTAL_STEPS="${TOTAL_STEPS:-100}" +export TOTAL_EPOCHS="${TOTAL_EPOCHS:-1}" +export TRAIN_BATCH_SIZE="${TRAIN_BATCH_SIZE:-64}" +export MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-1}" +export MAX_TOKENS_PER_GPU="${MAX_TOKENS_PER_GPU:-8192}" +export MAX_LENGTH="${MAX_LENGTH:-2048}" + +exec bash "${SCRIPT_DIR}/run_qwen3moe_sft.sh" "$@" diff --git a/experimental/lite/examples/verl/scripts/run_qwen3moe_sft.sh b/experimental/lite/examples/verl/scripts/run_qwen3moe_sft.sh new file mode 100755 index 00000000000..eb4c903b539 --- /dev/null +++ b/experimental/lite/examples/verl/scripts/run_qwen3moe_sft.sh @@ -0,0 +1,243 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ "${VERBOSE:-0}" == "1" ]]; then + set -x +fi + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -L)" +EXAMPLE_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd -L)" +LITE_ROOT="$(cd "${EXAMPLE_ROOT}/../.." && pwd -L)" +REPO_ROOT="$(cd "${LITE_ROOT}/../.." && pwd -L)" + +add_pythonpath() { + local path="${1:-}" + if [[ -n "${path}" ]]; then + export PYTHONPATH="${path}:${PYTHONPATH:-}" + fi +} + +add_pythonpath "${EXAMPLE_ROOT}" +add_pythonpath "${LITE_ROOT}" +add_pythonpath "${REPO_ROOT}" +add_pythonpath "${VERL_ROOT:-}" +add_pythonpath "${MEGATRON_ROOT:-}" + +export CUDA_DEVICE_MAX_CONNECTIONS="${CUDA_DEVICE_MAX_CONNECTIONS:-1}" + +: "${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}" + +BACKEND="${BACKEND:-mlite}" +VAL_FILES="${VAL_FILES:-}" +OUTPUT_ROOT="${OUTPUT_ROOT:-${EXAMPLE_ROOT}/outputs/qwen3moe_sft}" +PROJECT_NAME="${PROJECT_NAME:-verl-mlite-qwen3moe-sft}" + +NUM_GPUS="${NUM_GPUS:-${NPROC_PER_NODE:-8}}" +NPROC_PER_NODE="${NPROC_PER_NODE:-${NUM_GPUS}}" +NNODES="${NNODES:-1}" +NODE_RANK="${NODE_RANK:-0}" +MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}" +MASTER_PORT="${MASTER_PORT:-29500}" + +TOTAL_STEPS="${TOTAL_STEPS:-100}" +TOTAL_EPOCHS="${TOTAL_EPOCHS:-1}" +SAVE_FREQ="${SAVE_FREQ:-${TOTAL_STEPS}}" +TEST_FREQ="${TEST_FREQ:--1}" +RESUME_MODE="${RESUME_MODE:-disable}" +RESUME_FROM_PATH="${RESUME_FROM_PATH:-null}" +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}" +NUM_WORKERS="${NUM_WORKERS:-0}" +SEED="${SEED:-1}" + +TP_SIZE="${TP_SIZE:-2}" +PP_SIZE="${PP_SIZE:-1}" +VPP_SIZE="${VPP_SIZE:-null}" +CP_SIZE="${CP_SIZE:-1}" +EP_SIZE="${EP_SIZE:-8}" +ETP_SIZE="${ETP_SIZE:-1}" +DTYPE="${DTYPE:-bfloat16}" +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. +# - fsdp2: Megatron Lite FSDP2 wrapper + optimizer. +MLITE_OPTIMIZER_BACKEND="${MLITE_OPTIMIZER_BACKEND:-distopt}" + +LR="${LR:-1e-5}" +MIN_LR="${MIN_LR:-${LR}}" +WEIGHT_DECAY="${WEIGHT_DECAY:-0.1}" +BETAS="${BETAS:-[0.9,0.95]}" +CLIP_GRAD="${CLIP_GRAD:-1.0}" +LR_WARMUP_STEPS="${LR_WARMUP_STEPS:-0}" +LR_DECAY_STYLE="${LR_DECAY_STYLE:-constant}" + +PARAM_OFFLOAD="${PARAM_OFFLOAD:-False}" +OPTIMIZER_OFFLOAD="${OPTIMIZER_OFFLOAD:-True}" +GRAD_OFFLOAD="${GRAD_OFFLOAD:-False}" +OPTIMIZER_STATE_OFFLOAD_FRACTION="${OPTIMIZER_STATE_OFFLOAD_FRACTION:-1.0}" +USE_PRECISION_AWARE_OPTIMIZER="${USE_PRECISION_AWARE_OPTIMIZER:-True}" +DECOUPLED_WEIGHT_DECAY="${DECOUPLED_WEIGHT_DECAY:-True}" +DRY_RUN="${DRY_RUN:-0}" +EXTRA_ARGS=("$@") + +if [[ "${BACKEND}" != "mlite" ]]; then + echo "Unsupported BACKEND=${BACKEND}. This example is for BACKEND=mlite." >&2 + exit 1 +fi + +if [[ "${PAD_MODE}" != "no_padding" ]]; then + echo "Megatron Lite VERL example currently supports PAD_MODE=no_padding only." >&2 + exit 1 +fi + +case "${MLITE_OPTIMIZER_BACKEND}" in + distopt) + MLITE_IMPL_OPTIMIZER="mc" + ;; + fsdp2) + MLITE_IMPL_OPTIMIZER="fsdp2" + ;; + *) + echo "Unsupported MLITE_OPTIMIZER_BACKEND=${MLITE_OPTIMIZER_BACKEND}. Expected distopt or fsdp2." >&2 + exit 1 + ;; +esac + +MLITE_VPP_SIZE="${VPP_SIZE}" +if [[ "${MLITE_VPP_SIZE}" == "null" ]]; then + MLITE_VPP_SIZE=1 +fi + +RUN_NAME="${RUN_NAME:-qwen3moe_sft_mlite_tp${TP_SIZE}_pp${PP_SIZE}_cp${CP_SIZE}_ep${EP_SIZE}_etp${ETP_SIZE}}" +CKPT_DIR="${CKPT_DIR:-${OUTPUT_ROOT}/checkpoints/${RUN_NAME}}" +LOG_FILE="${LOG_FILE:-${OUTPUT_ROOT}/${RUN_NAME}.log}" +JSONL_FILE="${JSONL_FILE:-${OUTPUT_ROOT}/${RUN_NAME}.jsonl}" +CMD_FILE="${CMD_FILE:-${OUTPUT_ROOT}/${RUN_NAME}.cmd.sh}" + +mkdir -p "${OUTPUT_ROOT}" "${CKPT_DIR}" "$(dirname "${LOG_FILE}")" "$(dirname "${JSONL_FILE}")" "$(dirname "${CMD_FILE}")" +export VERL_FILE_LOGGER_PATH="${JSONL_FILE}" + +CACHE_ROOT="${VERL_MLITE_CACHE_ROOT:-${TMPDIR:-/tmp}/verl_mlite}" +mkdir -p "${CACHE_ROOT}/pycache_${USER:-user}" "${CACHE_ROOT}/torchinductor_${USER:-user}" "${CACHE_ROOT}/triton_${USER:-user}" +export PYTHONPYCACHEPREFIX="${PYTHONPYCACHEPREFIX:-${CACHE_ROOT}/pycache_${USER:-user}}" +export TORCHINDUCTOR_CACHE_DIR="${TORCHINDUCTOR_CACHE_DIR:-${CACHE_ROOT}/torchinductor_${USER:-user}}" +export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-${CACHE_ROOT}/triton_${USER:-user}}" + +COMMON_ARGS=( + "data.train_files=${TRAIN_FILES}" + "data.train_batch_size=${TRAIN_BATCH_SIZE}" + "data.micro_batch_size_per_gpu=${MICRO_BATCH_SIZE}" + "data.use_dynamic_bsz=${USE_DYNAMIC_BSZ}" + "data.max_token_len_per_gpu=${MAX_TOKENS_PER_GPU}" + "data.max_length=${MAX_LENGTH}" + "data.pad_mode=${PAD_MODE}" + "data.truncation=error" + "data.messages_key=${MESSAGES_KEY}" + "data.ignore_input_ids_mismatch=${IGNORE_INPUT_IDS_MISMATCH}" + "data.num_workers=${NUM_WORKERS}" + "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}" + "optim.weight_decay=${WEIGHT_DECAY}" + "optim.betas=${BETAS}" + "optim.clip_grad=${CLIP_GRAD}" + "optim.lr_warmup_steps=${LR_WARMUP_STEPS}" + "optim.lr_warmup_init=0" + "optim.lr_decay_style=${LR_DECAY_STYLE}" + "trainer.logger=[console,file]" + "trainer.project_name=${PROJECT_NAME}" + "trainer.experiment_name=${RUN_NAME}" + "trainer.default_local_dir=${CKPT_DIR}" + "trainer.total_epochs=${TOTAL_EPOCHS}" + "trainer.total_training_steps=${TOTAL_STEPS}" + "trainer.save_freq=${SAVE_FREQ}" + "trainer.test_freq=${TEST_FREQ}" + "trainer.seed=${SEED}" + "trainer.resume_mode=${RESUME_MODE}" + "trainer.resume_from_path=${RESUME_FROM_PATH}" + "trainer.nnodes=${NNODES}" + "trainer.n_gpus_per_node=${NPROC_PER_NODE}" + "checkpoint.save_contents=[model,optimizer,extra]" +) + +if [[ -n "${VAL_FILES}" ]]; then + COMMON_ARGS+=("data.val_files=${VAL_FILES}") +fi + +BACKEND_ARGS=( + "hydra.searchpath=[pkg://verl_mlite.config]" + "engine=mlite" + "engine.dtype=${DTYPE}" + "engine.model_name=${MLITE_MODEL_NAME}" + "engine.impl=${MLITE_IMPL}" + "engine.tp=${TP_SIZE}" + "engine.pp=${PP_SIZE}" + "engine.vpp=${MLITE_VPP_SIZE}" + "engine.cp=${CP_SIZE}" + "engine.ep=${EP_SIZE}" + "engine.etp=${ETP_SIZE}" + "engine.param_offload=${PARAM_OFFLOAD}" + "engine.optimizer_offload=${OPTIMIZER_OFFLOAD}" + "engine.grad_offload=${GRAD_OFFLOAD}" + "engine.attention_backend_override=${ATTENTION_BACKEND}" + "engine.impl_cfg.use_thd=True" + "+engine.impl_cfg.optimizer=${MLITE_IMPL_OPTIMIZER}" +) + +if [[ "${OPTIMIZER_OFFLOAD}" == "True" || "${OPTIMIZER_OFFLOAD}" == "true" || "${OPTIMIZER_OFFLOAD}" == "1" ]]; then + BACKEND_ARGS+=( + "+optim.override_optimizer_config.offload_fraction=${OPTIMIZER_STATE_OFFLOAD_FRACTION}" + "+optim.override_optimizer_config.use_precision_aware_optimizer=${USE_PRECISION_AWARE_OPTIMIZER}" + "+optim.override_optimizer_config.decoupled_weight_decay=${DECOUPLED_WEIGHT_DECAY}" + ) +fi + +COMMAND=( + torchrun + --nnodes="${NNODES}" + --node_rank="${NODE_RANK}" + --master_addr="${MASTER_ADDR}" + --master_port="${MASTER_PORT}" + --nproc_per_node="${NPROC_PER_NODE}" + -m + verl.trainer.sft_trainer + "${COMMON_ARGS[@]}" + "${BACKEND_ARGS[@]}" + "${EXTRA_ARGS[@]}" +) + +printf '%q ' "${COMMAND[@]}" > "${CMD_FILE}" +printf '\n' >> "${CMD_FILE}" + +if [[ "${DRY_RUN}" == "1" ]]; then + printf '%q ' "${COMMAND[@]}" + printf '\n' + exit 0 +fi + +echo "[${BACKEND}] output_root=${OUTPUT_ROOT}" +echo "[${BACKEND}] log=${LOG_FILE}" +echo "[${BACKEND}] jsonl=${JSONL_FILE}" +echo "[${BACKEND}] cmd=${CMD_FILE}" +echo "[${BACKEND}] optimizer_backend=${MLITE_OPTIMIZER_BACKEND} impl_cfg.optimizer=${MLITE_IMPL_OPTIMIZER}" + +set +e +"${COMMAND[@]}" 2>&1 | tee "${LOG_FILE}" +cmd_rc="${PIPESTATUS[0]}" +set -e +exit "${cmd_rc}" diff --git a/experimental/lite/examples/verl/sitecustomize.py b/experimental/lite/examples/verl/sitecustomize.py new file mode 100644 index 00000000000..dfe77bf5951 --- /dev/null +++ b/experimental/lite/examples/verl/sitecustomize.py @@ -0,0 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Process-wide compatibility hooks for the VERL MLite examples.""" + +from verl_mlite.compat import apply_runtime_patches + +apply_runtime_patches() diff --git a/experimental/lite/examples/verl/verl_mlite/__init__.py b/experimental/lite/examples/verl/verl_mlite/__init__.py new file mode 100644 index 00000000000..25500fa4d12 --- /dev/null +++ b/experimental/lite/examples/verl/verl_mlite/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""VERL integration package for Megatron Lite.""" diff --git a/experimental/lite/examples/verl/verl_mlite/compat.py b/experimental/lite/examples/verl/verl_mlite/compat.py new file mode 100644 index 00000000000..bb5ac89ba97 --- /dev/null +++ b/experimental/lite/examples/verl/verl_mlite/compat.py @@ -0,0 +1,56 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Small compatibility patches for dependency-version gaps in examples.""" + +from __future__ import annotations + +from collections.abc import Iterable +from functools import wraps +from typing import Any + + +def _patch_transformers_rope_ignore_keys() -> None: + try: + import transformers.modeling_rope_utils as rope_utils + except Exception: + return + + for cls in vars(rope_utils).values(): + if not isinstance(cls, type): + continue + if getattr(cls, "_verl_mlite_rope_ignore_keys_patch", False): + continue + descriptor = vars(cls).get("_check_received_keys") + if descriptor is None: + continue + + is_staticmethod = isinstance(descriptor, staticmethod) + is_classmethod = isinstance(descriptor, classmethod) + original = descriptor.__func__ if is_staticmethod or is_classmethod else descriptor + + def build_wrapper(check_received_keys: Any) -> Any: + @wraps(check_received_keys) + def patched(*args: Any, **kwargs: Any) -> Any: + ignore_keys = kwargs.get("ignore_keys") + if isinstance(ignore_keys, list): + kwargs["ignore_keys"] = set(ignore_keys) + elif ignore_keys is not None and not isinstance(ignore_keys, set): + if isinstance(ignore_keys, Iterable) and not isinstance( + ignore_keys, (str, bytes) + ): + kwargs["ignore_keys"] = set(ignore_keys) + return check_received_keys(*args, **kwargs) + + return patched + + patched = build_wrapper(original) + if is_staticmethod: + cls._check_received_keys = staticmethod(patched) + elif is_classmethod: + cls._check_received_keys = classmethod(patched) + else: + cls._check_received_keys = patched + cls._verl_mlite_rope_ignore_keys_patch = True + + +def apply_runtime_patches() -> None: + _patch_transformers_rope_ignore_keys() diff --git a/experimental/lite/examples/verl/verl_mlite/config/__init__.py b/experimental/lite/examples/verl/verl_mlite/config/__init__.py new file mode 100644 index 00000000000..15d962701b4 --- /dev/null +++ b/experimental/lite/examples/verl/verl_mlite/config/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Hydra config package for Verl MLite.""" diff --git a/experimental/lite/examples/verl/verl_mlite/config/actor/__init__.py b/experimental/lite/examples/verl/verl_mlite/config/actor/__init__.py new file mode 100644 index 00000000000..b7d8532bf20 --- /dev/null +++ b/experimental/lite/examples/verl/verl_mlite/config/actor/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Hydra actor config group for Verl MLite.""" diff --git a/experimental/lite/examples/verl/verl_mlite/config/actor/mlite_actor.yaml b/experimental/lite/examples/verl/verl_mlite/config/actor/mlite_actor.yaml new file mode 100644 index 00000000000..7efeab8431b --- /dev/null +++ b/experimental/lite/examples/verl/verl_mlite/config/actor/mlite_actor.yaml @@ -0,0 +1,10 @@ +# Megatron Lite actor config for VERL's new engine worker path. +defaults: + - /optim@optim: megatron + - /engine@engine: mlite + - actor + - _self_ + +_target_: verl.workers.config.ActorConfig + +strategy: mlite diff --git a/experimental/lite/examples/verl/verl_mlite/config/engine/__init__.py b/experimental/lite/examples/verl/verl_mlite/config/engine/__init__.py new file mode 100644 index 00000000000..fe007d793a8 --- /dev/null +++ b/experimental/lite/examples/verl/verl_mlite/config/engine/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Engine config group for Verl MLite.""" diff --git a/experimental/lite/examples/verl/verl_mlite/config/engine/mlite.yaml b/experimental/lite/examples/verl/verl_mlite/config/engine/mlite.yaml new file mode 100644 index 00000000000..7e3e4814239 --- /dev/null +++ b/experimental/lite/examples/verl/verl_mlite/config/engine/mlite.yaml @@ -0,0 +1,25 @@ +_target_: verl_mlite.engine.config.MegatronLiteEngineConfig + +strategy: mlite +custom_backend_module: verl_mlite.engine.mlite_engine +param_offload: false +optimizer_offload: false +grad_offload: false +forward_only: false +dtype: bfloat16 +export_dtype: bfloat16 + +model_name: auto +impl: lite + +tp: 1 +etp: null +ep: 1 +pp: 1 +vpp: 1 +cp: 1 + +attention_backend_override: flash +router_aux_loss_coef: null +impl_cfg: + use_thd: true diff --git a/experimental/lite/examples/verl/verl_mlite/engine/__init__.py b/experimental/lite/examples/verl/verl_mlite/engine/__init__.py new file mode 100644 index 00000000000..7662660591b --- /dev/null +++ b/experimental/lite/examples/verl/verl_mlite/engine/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Engine entrypoints for Verl MLite.""" + +from verl_mlite.engine.mlite_engine import MegatronLiteEngine + +__all__ = ["MegatronLiteEngine"] diff --git a/experimental/lite/examples/verl/verl_mlite/engine/config.py b/experimental/lite/examples/verl/verl_mlite/engine/config.py new file mode 100644 index 00000000000..fa2bb461ff4 --- /dev/null +++ b/experimental/lite/examples/verl/verl_mlite/engine/config.py @@ -0,0 +1,41 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Config objects for the Verl MLite Megatron Lite engine.""" + +from __future__ import annotations + +import importlib +from dataclasses import dataclass, field +from typing import Any + +from verl.workers.config.engine import EngineConfig + + +@dataclass +class MegatronLiteEngineConfig(EngineConfig): + """Minimal VERL-facing config for the external Megatron Lite engine.""" + + strategy: str = "mlite" + custom_backend_module: str | None = "verl_mlite.engine.mlite_engine" + model_name: str = "auto" + impl: str = "lite" + + tp: int = 1 + etp: int | None = None + ep: int = 1 + pp: int = 1 + vpp: int = 1 + cp: int = 1 + + attention_backend_override: str | None = "flash" + router_aux_loss_coef: float | None = None + export_dtype: str | None = "bfloat16" + impl_cfg: dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + super().__post_init__() + if self.strategy != "mlite": + raise ValueError( + f"MegatronLiteEngineConfig expects strategy='mlite', got {self.strategy!r}" + ) + if self.custom_backend_module: + importlib.import_module(self.custom_backend_module) diff --git a/experimental/lite/examples/verl/verl_mlite/engine/mlite_engine.py b/experimental/lite/examples/verl/verl_mlite/engine/mlite_engine.py new file mode 100644 index 00000000000..32e81f23fb3 --- /dev/null +++ b/experimental/lite/examples/verl/verl_mlite/engine/mlite_engine.py @@ -0,0 +1,778 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""External VERL engine backed by Megatron Lite runtime primitives.""" + +from __future__ import annotations + +import os +from typing import Any + +import torch +import torch.distributed as dist +from tensordict import TensorDict +from verl.trainer.config import CheckpointConfig +from verl.utils import tensordict_utils as tu +from verl.utils.dataset.dataset_utils import DatasetPadMode +from verl.utils.device import get_device_id, get_device_name +from verl.workers.config import HFModelConfig, OptimizerConfig +from verl.workers.engine.base import BaseEngine, BaseEngineCtx, EngineRegistry +from verl.workers.engine.utils import postprocess_batch_func, prepare_micro_batches + +from megatron.lite.model import resolve_model_type_from_hf +from megatron.lite.primitive.ckpt import load_training_checkpoint, save_training_checkpoint +from megatron.lite.primitive.parallel import pack_nested_thd, unpack_packed_thd_to_nested +from megatron.lite.primitive.protocols import default_expert_classifier, default_placement_fn +from megatron.lite.runtime import create_runtime +from megatron.lite.runtime.backends.mlite.config import MegatronLiteConfig +from megatron.lite.runtime.contracts.config import OptimizerConfig as MegatronLiteOptimizerConfig +from megatron.lite.runtime.contracts.config import ParallelConfig, RuntimeConfig + +from .config import MegatronLiteEngineConfig + +_LR_SCHEDULER_STATE = "lr_scheduler.pt" + + +def _isolate_compile_cache_per_rank() -> None: + """Avoid torchinductor/triton cache races between local torchrun ranks.""" + rank = os.environ.get("LOCAL_RANK") or os.environ.get("RANK") + if rank is None: + return + for var in ("TORCHINDUCTOR_CACHE_DIR", "TRITON_CACHE_DIR"): + base = os.environ.get(var) + if not base: + continue + base_var = f"VERL_MLITE_BASE_{var}" + root = os.environ.setdefault(base_var, base) + rank_dir = os.path.join(root, f"rank_{rank}") + os.makedirs(rank_dir, exist_ok=True) + os.environ[var] = rank_dir + + +def _build_lr_scheduler(optimizer, opt: MegatronLiteOptimizerConfig): + """Build a Megatron-style LR scheduler for Megatron Lite's optimizer.""" + total_steps = opt.total_training_steps + if total_steps <= 0: + return None + + from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler + + warmup_steps = opt.lr_warmup_steps if opt.lr_warmup_steps is not None else -1 + if warmup_steps <= 0 and opt.lr_warmup_steps_ratio > 0: + warmup_steps = int(opt.lr_warmup_steps_ratio * total_steps) + warmup_steps = max(warmup_steps, 0) + + decay_steps = opt.lr_decay_steps if opt.lr_decay_steps is not None else total_steps + min_lr = opt.min_lr if opt.min_lr is not None else 0.0 + for param_group in optimizer.param_groups: + if param_group.get("min_lr") is None: + param_group["min_lr"] = min_lr + + return OptimizerParamScheduler( + optimizer, + init_lr=opt.lr_warmup_init, + max_lr=opt.lr, + min_lr=min_lr, + lr_warmup_steps=warmup_steps, + lr_decay_steps=decay_steps, + lr_decay_style=opt.lr_decay_style, + start_wd=opt.weight_decay, + end_wd=opt.weight_decay, + wd_incr_steps=total_steps, + wd_incr_style=opt.weight_decay_incr_style, + use_checkpoint_opt_param_scheduler=opt.use_checkpoint_opt_param_scheduler, + override_opt_param_scheduler=not opt.use_checkpoint_opt_param_scheduler, + wsd_decay_steps=opt.lr_wsd_decay_steps, + lr_wsd_decay_style=opt.lr_wsd_decay_style, + ) + + +class _MegatronLiteModeCtx(BaseEngineCtx): + """Wrap Megatron Lite runtime contexts with VERL's offload behavior.""" + + def __init__(self, engine: MegatronLiteEngine, mode: str, **kwargs): + super().__init__(engine=engine, mode=mode, **kwargs) + self._runtime_ctx = None + + def __enter__(self): + super().__enter__() + assert self.engine.runtime is not None and self.engine.handle is not None + if self.mode == "train": + self._runtime_ctx = self.engine.runtime.train_mode(self.engine.handle) + else: + self._runtime_ctx = self.engine.runtime.eval_mode(self.engine.handle) + self._runtime_ctx.__enter__() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + assert self._runtime_ctx is not None + self._runtime_ctx.__exit__(exc_type, exc_val, exc_tb) + super().__exit__(exc_type, exc_val, exc_tb) + return False + + +@EngineRegistry.register(model_type="language_model", backend="mlite", device="cuda") +class MegatronLiteEngine(BaseEngine): + """VERL BaseEngine implementation that delegates model lifecycle to Megatron Lite.""" + + def __init__( + self, + model_config: HFModelConfig, + engine_config: MegatronLiteEngineConfig, + optimizer_config: OptimizerConfig, + checkpoint_config: CheckpointConfig, + ): + super().__init__() + _isolate_compile_cache_per_rank() + self.model_config = model_config + self.engine_config = engine_config + self.optimizer_config = optimizer_config + self.checkpoint_config = checkpoint_config + + self.mode = None + self.device_name = get_device_name() + self.runtime = None + self.handle = None + self.module = None + self._mlite_config = None + self._rank = dist.get_rank() if dist.is_initialized() else 0 + + @property + def is_param_offload_enabled(self) -> bool: + return self.engine_config.param_offload + + @property + def is_optimizer_offload_enabled(self) -> bool: + return self.engine_config.optimizer_offload + + def initialize(self): + if self.engine_config.full_determinism: + from verl.workers.engine.utils import enable_full_determinism + + enable_full_determinism(seed=self.engine_config.seed) + + self._mlite_config = self._build_mlite_config() + self.runtime = create_runtime( + RuntimeConfig( + backend="mlite", + hf_path=self.model_config.local_path, + backend_cfg=self._mlite_config, + ) + ) + self.handle = self.runtime.build_model() + self.module = self._extract_primary_module() + + if self.handle._optimizer is not None and self.handle._lr_scheduler is None: + self.handle._lr_scheduler = _build_lr_scheduler( + self.handle._optimizer, self._mlite_config.optimizer + ) + + self.to( + device="cpu", + model=self.is_param_offload_enabled, + optimizer=self.is_optimizer_offload_enabled, + grad=self.is_param_offload_enabled, + ) + + def train_mode(self, **kwargs): + self._require_initialized() + return _MegatronLiteModeCtx(self, mode="train", **kwargs) + + def eval_mode(self, **kwargs): + self._require_initialized() + return _MegatronLiteModeCtx(self, mode="eval", **kwargs) + + def optimizer_zero_grad(self): + self._require_initialized() + self.runtime.zero_grad(self.handle) + + def optimizer_step(self): + self._require_initialized() + _, grad_norm, _ = self.runtime.optimizer_step(self.handle) + return grad_norm + + def lr_scheduler_step(self): + self._require_initialized() + if self.handle._lr_scheduler is not None: + self.handle._lr_scheduler.step(1) + return self.handle._optimizer.param_groups[0]["lr"] + return 0.0 + + def forward_backward_batch( + self, data: TensorDict, loss_function, forward_only: bool = False + ) -> dict[str, Any]: + self._require_initialized() + pad_mode = tu.get_non_tensor_data( + data=data, key="pad_mode", default=DatasetPadMode.NO_PADDING + ) + if pad_mode != DatasetPadMode.NO_PADDING: + raise NotImplementedError( + "MegatronLiteEngine only supports pad_mode=no_padding for now." + ) + + tu.assign_non_tensor(data, sp_size=self.engine_config.cp) + + token_mask = data["loss_mask"] if "loss_mask" in data.keys() else data["response_mask"] + batch_num_tokens = token_mask.sum().to(get_device_id()) + torch.distributed.all_reduce( + batch_num_tokens, + op=torch.distributed.ReduceOp.SUM, + group=self.get_data_parallel_group(), + ) + tu.assign_non_tensor(data, batch_num_tokens=batch_num_tokens.item()) + tu.assign_non_tensor(data, dp_size=self.get_data_parallel_size()) + + micro_batches, indices = prepare_micro_batches( + data=data, dp_group=self.get_data_parallel_group(), same_micro_num_in_dp=True + ) + + if self._use_runtime_forward_backward(): + return self._forward_backward_batch_with_runtime( + data=data, + micro_batches=micro_batches, + indices=indices, + loss_function=loss_function, + forward_only=forward_only, + ) + + outputs = [] + num_micro_batches = len(micro_batches) + for micro_idx, micro_batch in enumerate(micro_batches): + tu.assign_non_tensor(micro_batch, micro_batch_idx=micro_idx) + micro_batch = micro_batch.to(get_device_id()) + model_inputs = self._make_model_inputs(micro_batch) + + pre_forward_hook = self.handle._extras.get("pre_forward_hook") + if pre_forward_hook is not None: + pre_forward_hook(torch.tensor(1.0 / num_micro_batches, device=get_device_id())) + + with torch.no_grad() if forward_only else torch.enable_grad(): + raw_output = self.module( + input_ids=model_inputs["input_ids"], + position_ids=model_inputs["position_ids"], + packed_seq_params=model_inputs["packed_seq_params"], + labels=model_inputs["labels"], + loss_mask=model_inputs.get("loss_mask"), + temperature=model_inputs["temperature"], + use_fused_kernels=model_inputs["use_fused_kernels"], + calculate_entropy=model_inputs["calculate_entropy"], + ) + + model_output = self._build_verl_model_output( + raw_output=raw_output, micro_batch=micro_batch, inputs=model_inputs + ) + + if loss_function is not None: + loss, metrics = loss_function( + model_output=model_output, + data=micro_batch, + dp_group=self.get_data_parallel_group(), + ) + else: + loss = torch.zeros((), device=get_device_id(), dtype=torch.float32) + metrics = {} + if raw_output.get("mtp_loss") is not None: + metrics = dict(metrics) + mtp_loss = self._reduce_mtp_metric(raw_output["mtp_loss"]) + metrics["mtp_losses/mtp_1_loss"] = ( + float(mtp_loss.item()) if mtp_loss.numel() == 1 else mtp_loss.cpu().tolist() + ) + + if not forward_only and loss_function is not None: + loss.backward() + + outputs.append( + {"model_output": model_output, "loss": loss.detach().item(), "metrics": metrics} + ) + + if not forward_only: + finalize_grads = self.handle._extras.get("finalize_grads") + if finalize_grads is not None: + finalize_grads() + + result = postprocess_batch_func(output_lst=outputs, indices=indices, data=data) + return result + + def get_per_tensor_param(self, **kwargs): + self._require_initialized() + if self.is_param_offload_enabled: + self.to("cuda", model=True, optimizer=False, grad=False) + export_kwargs = { + key: kwargs[key] + for key in ("limit", "include_mtp_only", "include_local_prefixes") + if key in kwargs + } + if self.engine_config.model_name == "qwen3_5": + export_kwargs["target"] = "vllm" + if self.engine_config.export_dtype: + export_kwargs["export_dtype"] = self.engine_config.export_dtype + return self.runtime.export_weights(self.handle, **export_kwargs), None + + def get_data_parallel_size(self): + if self.handle is None: + world_size = dist.get_world_size() if dist.is_initialized() else 1 + return world_size // ( + self.engine_config.tp * self.engine_config.cp * self.engine_config.pp + ) + return self.handle.dp_size + + def get_data_parallel_rank(self): + if self.handle is None: + rank = dist.get_rank() if dist.is_initialized() else 0 + dense_dp = self.get_data_parallel_size() + return (rank // (self.engine_config.tp * self.engine_config.cp)) % dense_dp + return self.handle.dp_rank + + def get_data_parallel_group(self): + if self.handle is None: + if ( + self.engine_config.tp == 1 + and self.engine_config.cp == 1 + and self.engine_config.pp == 1 + and dist.is_initialized() + ): + return dist.group.WORLD + return None + return self.handle.dp_group + + def to(self, device: str, model: bool = True, optimizer: bool = True, grad: bool = True): + self._require_initialized() + if model or not (optimizer or grad): + super().to(device=device, model=model, optimizer=optimizer, grad=grad) + self.runtime.to(self.handle, device, model=model, optimizer=optimizer, grad=grad) + + def save_checkpoint( + self, + local_path: str, + hdfs_path: str | None = None, + global_step: int = 0, + max_ckpt_to_keep: int | None = None, + **kwargs, + ) -> None: + del hdfs_path, max_ckpt_to_keep, kwargs + self._require_initialized() + + save_contents = self.checkpoint_config.get("save_contents", None) + save_model = save_contents is None or "model" in save_contents + save_optimizer = save_contents is None or "optimizer" in save_contents + if not save_model and not save_optimizer: + if self._rank == 0: + print( + f"Skipping Megatron Lite checkpoint save at step {global_step}: save_contents={save_contents}" + ) + if dist.is_initialized(): + dist.barrier() + return + + os.makedirs(local_path, exist_ok=True) + placement_fn, expert_classifier = self._checkpoint_hooks() + reload_params_for_save = self.is_param_offload_enabled + if reload_params_for_save: + self.to(device="cuda", model=True, optimizer=False, grad=False) + torch.cuda.synchronize() + try: + save_training_checkpoint( + self.module, + self.handle._optimizer, + global_step, + local_path, + self.handle._config.parallel, + self.handle._parallel_state, + get_placements=placement_fn, + is_expert=expert_classifier, + save_model=save_model, + save_optimizer=save_optimizer, + ) + if self.handle._lr_scheduler is not None and self._rank == 0: + torch.save( + self.handle._lr_scheduler.state_dict(), + os.path.join(local_path, _LR_SCHEDULER_STATE), + ) + if dist.is_initialized(): + dist.barrier() + finally: + if reload_params_for_save: + self.to(device="cpu", model=True, optimizer=False, grad=False) + + def load_checkpoint( + self, + local_path: str, + hdfs_path: str | None = None, + del_local_after_load: bool = True, + **kwargs, + ) -> None: + del hdfs_path, del_local_after_load, kwargs + self._require_initialized() + + placement_fn, expert_classifier = self._checkpoint_hooks() + reload_params_for_load = self.is_param_offload_enabled + if reload_params_for_load: + self.to(device="cuda", model=True, optimizer=False, grad=False) + torch.cuda.synchronize() + try: + load_training_checkpoint( + self.module, + self.handle._optimizer, + local_path, + self.handle._config.parallel, + self.handle._parallel_state, + get_placements=placement_fn, + is_expert=expert_classifier, + load_model=True, + load_optimizer=True, + ) + scheduler_path = os.path.join(local_path, _LR_SCHEDULER_STATE) + if self.handle._lr_scheduler is not None and os.path.exists(scheduler_path): + state = torch.load(scheduler_path, map_location="cpu", weights_only=False) + self.handle._lr_scheduler.load_state_dict(state) + if dist.is_initialized(): + dist.barrier() + finally: + if reload_params_for_load: + self.to(device="cpu", model=True, optimizer=False, grad=False) + + def is_mp_src_rank_with_outputs(self): + if self.handle is None: + rank = dist.get_rank() if dist.is_initialized() else 0 + dense_dp = self.get_data_parallel_size() + tp_rank = rank % self.engine_config.tp + cp_rank = (rank // self.engine_config.tp) % self.engine_config.cp + pp_rank = rank // (self.engine_config.tp * self.engine_config.cp * dense_dp) + return tp_rank == 0 and cp_rank == 0 and pp_rank == self.engine_config.pp - 1 + return self.runtime.is_mp_src_rank_with_outputs(self.handle) + + def _require_initialized(self) -> None: + if self.runtime is None or self.handle is None: + raise RuntimeError("MegatronLiteEngine is not initialized yet.") + + def _build_mlite_config(self) -> MegatronLiteConfig: + return MegatronLiteConfig( + model_name=self._resolve_model_name(), + impl=self.engine_config.impl, + hf_path=self.model_config.local_path, + parallel=ParallelConfig( + tp=self.engine_config.tp, + etp=self.engine_config.etp or 1, + ep=self.engine_config.ep, + pp=self.engine_config.pp, + vpp=self.engine_config.vpp, + cp=self.engine_config.cp, + ), + optimizer=self._build_mlite_optimizer_config(), + attention_backend_override=self.engine_config.attention_backend_override, + router_aux_loss_coef=self.engine_config.router_aux_loss_coef, + impl_cfg=self._build_impl_cfg(), + ) + + def _resolve_model_name(self) -> str: + if self.engine_config.model_name != "auto": + return self.engine_config.model_name + return resolve_model_type_from_hf(self.model_config.hf_config) + + def _build_impl_cfg(self) -> dict[str, Any]: + impl_cfg = dict(self.engine_config.impl_cfg) + if impl_cfg.get("use_thd", True) is not True: + raise ValueError( + "MegatronLiteEngine supports only THD/no-padding SFT; set engine.impl_cfg.use_thd=True." + ) + impl_cfg["use_thd"] = True + mtp_cfg = getattr(self.model_config, "mtp", None) + if mtp_cfg is not None: + mtp_enable = bool(getattr(mtp_cfg, "enable", False)) + mtp_enable_train = mtp_enable and bool(getattr(mtp_cfg, "enable_train", False)) + impl_cfg["mtp_enable"] = mtp_enable + impl_cfg["mtp_enable_train"] = mtp_enable_train + impl_cfg["mtp_detach_encoder"] = bool(getattr(mtp_cfg, "detach_encoder", False)) + impl_cfg["mtp_loss_scaling_factor"] = float( + getattr(mtp_cfg, "mtp_loss_scaling_factor", 0.1) + ) + if self.engine_config.full_determinism: + impl_cfg.setdefault("deterministic", True) + if self.engine_config.forward_only: + impl_cfg["optimizer"] = None + return impl_cfg + + def _build_mlite_optimizer_config(self) -> MegatronLiteOptimizerConfig: + optimizer_name = self._normalize_optimizer_name(self.optimizer_config) + betas = tuple(getattr(self.optimizer_config, "betas", (0.9, 0.999))) + override = getattr(self.optimizer_config, "override_optimizer_config", {}) or {} + offload_fraction = override.get( + "offload_fraction", override.get("optimizer_offload_fraction") + ) + if offload_fraction is None and override.get("optimizer_cpu_offload"): + offload_fraction = 1.0 + if offload_fraction is None and self.is_optimizer_offload_enabled: + offload_fraction = 1.0 + + min_lr = getattr(self.optimizer_config, "min_lr", None) + min_lr_ratio = getattr(self.optimizer_config, "min_lr_ratio", None) + if min_lr is None: + min_lr = 0.0 if min_lr_ratio is None else self.optimizer_config.lr * min_lr_ratio + + lr_decay_style = getattr(self.optimizer_config, "lr_decay_style", None) + if lr_decay_style is None: + lr_decay_style = getattr(self.optimizer_config, "lr_scheduler_type", "constant") + + return MegatronLiteOptimizerConfig( + optimizer=optimizer_name, + lr=self.optimizer_config.lr, + min_lr=min_lr, + clip_grad=self.optimizer_config.clip_grad, + weight_decay=self.optimizer_config.weight_decay, + lr_warmup_steps_ratio=self.optimizer_config.lr_warmup_steps_ratio, + total_training_steps=self.optimizer_config.total_training_steps, + lr_warmup_steps=self.optimizer_config.lr_warmup_steps, + lr_warmup_init=getattr(self.optimizer_config, "lr_warmup_init", 0.0), + lr_decay_steps=getattr(self.optimizer_config, "lr_decay_steps", None), + lr_decay_style=lr_decay_style, + weight_decay_incr_style=getattr( + self.optimizer_config, "weight_decay_incr_style", "constant" + ), + lr_wsd_decay_style=getattr(self.optimizer_config, "lr_wsd_decay_style", "exponential"), + lr_wsd_decay_steps=getattr(self.optimizer_config, "lr_wsd_decay_steps", None), + use_checkpoint_opt_param_scheduler=getattr( + self.optimizer_config, "use_checkpoint_opt_param_scheduler", False + ), + adam_beta1=betas[0], + adam_beta2=betas[1], + adam_eps=override.get("adam_eps", override.get("eps")), + offload_fraction=offload_fraction, + use_precision_aware_optimizer=override.get("use_precision_aware_optimizer"), + decoupled_weight_decay=override.get("decoupled_weight_decay"), + ) + + @staticmethod + def _normalize_optimizer_name(config: OptimizerConfig) -> str: + optimizer_name = getattr(config, "optimizer", "adam") + lower = str(optimizer_name).lower() + if "adam" in lower: + return "adam" + raise ValueError( + f"MegatronLiteEngine only supports Adam-style optimizers today, got {optimizer_name!r}" + ) + + def _extract_primary_module(self): + model = self.handle._model + if isinstance(model, list | tuple): + if not model: + raise RuntimeError("Megatron Lite runtime returned an empty model chunk list.") + if len(model) > 1: + return torch.nn.ModuleList(model) + return model[0] + return model + + def _use_runtime_forward_backward(self) -> bool: + ps = self.handle._parallel_state + return ps.pp_size > 1 + + def _forward_backward_batch_with_runtime( + self, + *, + data: TensorDict, + micro_batches: list[TensorDict], + indices, + loss_function, + forward_only: bool, + ) -> dict[str, Any]: + runtime_batches = [] + num_micro_batches = len(micro_batches) + batch_num_tokens = tu.get_non_tensor_data(data=data, key="batch_num_tokens", default=None) + if batch_num_tokens is None: + raise ValueError( + "MegatronLiteEngine PP/CP SFT requires batch_num_tokens for VERL-compatible loss scaling." + ) + if batch_num_tokens <= 0: + raise ValueError(f"batch_num_tokens must be positive, got {batch_num_tokens}.") + loss_scale = self.get_data_parallel_size() * num_micro_batches / float(batch_num_tokens) + for micro_idx, micro_batch in enumerate(micro_batches): + tu.assign_non_tensor(micro_batch, micro_batch_idx=micro_idx) + micro_batch = micro_batch.to(get_device_id()) + model_inputs = self._make_model_inputs(micro_batch) + runtime_batches.append( + { + "input_ids": model_inputs["input_ids"], + "position_ids": model_inputs["position_ids"], + "packed_seq_params": model_inputs["packed_seq_params"], + "labels": model_inputs["labels"], + "loss_mask": model_inputs.get("loss_mask"), + "loss_scale": loss_scale, + "temperature": model_inputs["temperature"], + "use_fused_kernels": model_inputs["use_fused_kernels"], + "calculate_entropy": model_inputs["calculate_entropy"], + "_verl_micro_batch": micro_batch, + "_verl_inputs": model_inputs, + } + ) + + runtime_loss_fn = None + if loss_function is not None or forward_only: + runtime_loss_fn = self._make_runtime_loss_fn(loss_function, forward_only=forward_only) + + result = self.runtime.forward_backward( + self.handle, + iter(runtime_batches), + loss_fn=runtime_loss_fn, + num_microbatches=num_micro_batches, + forward_only=forward_only, + ) + metrics = dict(result.metrics) + micro_outputs = metrics.pop("_micro_outputs", None) + if micro_outputs is not None and self.is_mp_src_rank_with_outputs(): + return postprocess_batch_func(output_lst=micro_outputs, indices=indices, data=data) + loss = float(metrics.get("loss", 0.0)) + return { + "model_output": {}, + "loss": [loss], + "metrics": {key: [value] for key, value in metrics.items()}, + } + + def _make_model_inputs(self, micro_batch: TensorDict) -> dict[str, torch.Tensor]: + input_ids = micro_batch["input_ids"] + if not getattr(input_ids, "is_nested", False): + raise NotImplementedError( + "MegatronLiteEngine supports only nested no-padding THD batches." + ) + + ps = self.handle._parallel_state + loss_mask = self._loss_mask_for_packing(micro_batch, input_ids) + packed_batch = pack_nested_thd( + input_ids, + 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, + labels=input_ids, + roll_labels=True, + loss_mask=loss_mask, + roll_loss_mask=True, + ) + use_fused_kernels = tu.get_non_tensor_data( + data=micro_batch, key="use_fused_kernels", default=self.engine_config.use_fused_kernels + ) + + return { + "input_ids": packed_batch.input_ids, + "labels": packed_batch.labels, + "loss_mask": packed_batch.loss_mask, + "position_ids": packed_batch.position_ids, + "packed_seq_params": packed_batch.packed_seq_params, + "packed_batch": packed_batch, + "temperature": self._scalar_temperature(micro_batch), + "use_fused_kernels": use_fused_kernels, + "calculate_entropy": tu.get_non_tensor_data( + data=micro_batch, key="calculate_entropy", default=False + ), + } + + @staticmethod + def _loss_mask_for_packing( + micro_batch: TensorDict, input_ids: torch.Tensor + ) -> torch.Tensor | None: + if "loss_mask" not in micro_batch.keys(): + return None + + loss_mask = micro_batch["loss_mask"] + if getattr(loss_mask, "is_nested", False): + return loss_mask + + rows = [] + for seq_ids, row_mask in zip(input_ids.unbind(0), loss_mask, strict=True): + seq_len = seq_ids.numel() + response_tokens = int(row_mask.sum().item()) + if response_tokens > seq_len: + raise ValueError( + f"response loss mask has {response_tokens} tokens but packed input sequence has {seq_len} tokens" + ) + full_mask = torch.zeros(seq_len, dtype=row_mask.dtype, device=row_mask.device) + if response_tokens: + full_mask[-response_tokens:] = row_mask[:response_tokens] + rows.append(full_mask) + return torch.nested.as_nested_tensor(rows, layout=torch.jagged) + + def _build_verl_model_output( + self, + *, + raw_output: dict[str, torch.Tensor], + micro_batch: TensorDict, + inputs: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + del micro_batch + log_probs = raw_output.get("log_probs") + if log_probs is None: + raise ValueError("Megatron Lite THD model output must contain token log_probs.") + nested_log_probs = unpack_packed_thd_to_nested(log_probs, inputs["packed_batch"]) + output = {"log_probs": nested_log_probs} + entropy = raw_output.get("entropy") + if entropy is not None: + output["entropy"] = unpack_packed_thd_to_nested(entropy, inputs["packed_batch"]) + return output + + def _make_runtime_loss_fn(self, loss_function, *, forward_only: bool): + def _loss_fn(raw_output: dict[str, torch.Tensor], runtime_batch: dict[str, Any]): + micro_batch = runtime_batch["_verl_micro_batch"] + inputs = runtime_batch["_verl_inputs"] + model_output = self._build_verl_model_output( + raw_output=raw_output, micro_batch=micro_batch, inputs=inputs + ) + raw_output["_verl_model_output"] = model_output + if loss_function is not None: + loss, metrics = loss_function( + model_output=model_output, + data=micro_batch, + dp_group=self.get_data_parallel_group(), + ) + else: + loss = torch.zeros((), device=get_device_id(), dtype=torch.float32) + metrics = {} + + if raw_output.get("mtp_loss") is not None: + metrics = dict(metrics) + mtp_loss = self._reduce_mtp_metric(raw_output["mtp_loss"]) + metrics["mtp_losses/mtp_1_loss"] = ( + float(mtp_loss.item()) if mtp_loss.numel() == 1 else mtp_loss.cpu().tolist() + ) + + raw_output["_verl_metrics"] = metrics + return loss, metrics + + return _loss_fn + + def _mtp_enable_train(self) -> bool: + mtp_cfg = getattr(self.model_config, "mtp", None) + return bool( + mtp_cfg is not None + and getattr(mtp_cfg, "enable", False) + and getattr(mtp_cfg, "enable_train", False) + ) + + def _reduce_mtp_metric(self, mtp_loss: torch.Tensor) -> torch.Tensor: + mtp_loss = mtp_loss.detach().float().clone() + dp_group = self.get_data_parallel_group() + if dist.is_initialized() and dp_group is not None: + dist.all_reduce(mtp_loss, op=dist.ReduceOp.AVG, group=dp_group) + return mtp_loss + + @staticmethod + def _scalar_temperature(micro_batch: TensorDict) -> float: + if "temperature" not in micro_batch.keys(): + return 1.0 + temperature = micro_batch["temperature"] + if not isinstance(temperature, torch.Tensor): + return float(temperature) + values = ( + temperature.values() + if getattr(temperature, "is_nested", False) + else temperature.reshape(-1) + ) + if values.numel() == 0: + return 1.0 + first = values[0].detach() + if not torch.all(values.detach() == first).item(): + raise NotImplementedError( + "MegatronLiteEngine currently supports scalar temperature only." + ) + return float(first.float().item()) + + def _checkpoint_hooks(self): + proto = self.handle._extras.get("protocol") + placement_fn = getattr(proto, "PLACEMENT_FN", default_placement_fn) + expert_classifier = getattr(proto, "EXPERT_CLASSIFIER", default_expert_classifier) + return placement_fn, expert_classifier diff --git a/experimental/lite/megatron/lite/__init__.py b/experimental/lite/megatron/lite/__init__.py new file mode 100644 index 00000000000..7478deab37b --- /dev/null +++ b/experimental/lite/megatron/lite/__init__.py @@ -0,0 +1,36 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Top-level Megatron Lite package exports.""" + +from __future__ import annotations + +import importlib +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from megatron.lite.runtime.backends.bridge.config import BridgeConfig + from megatron.lite.runtime.backends.mlite.config import DebugConfig, MegatronLiteConfig + from megatron.lite.runtime.contracts import OptimizerConfig, ParallelConfig, RuntimeConfig + +__all__ = [ + "BridgeConfig", + "DebugConfig", + "MegatronLiteConfig", + "OptimizerConfig", + "ParallelConfig", + "RuntimeConfig", +] + + +def __getattr__(name: str): + _lazy = { + "BridgeConfig": "megatron.lite.runtime.backends.bridge.config", + "DebugConfig": "megatron.lite.runtime.backends.mlite.config", + "MegatronLiteConfig": "megatron.lite.runtime.backends.mlite.config", + "OptimizerConfig": "megatron.lite.runtime.contracts", + "ParallelConfig": "megatron.lite.runtime.contracts", + "RuntimeConfig": "megatron.lite.runtime.contracts", + } + if name in _lazy: + mod = importlib.import_module(_lazy[name]) + return getattr(mod, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/experimental/lite/megatron/lite/model/__init__.py b/experimental/lite/megatron/lite/model/__init__.py new file mode 100644 index 00000000000..71a16072b71 --- /dev/null +++ b/experimental/lite/megatron/lite/model/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Public model-registry helpers for Megatron Lite.""" + +from megatron.lite.model.registry import ( + get_model_package, + get_train_runtime_module, + register_model, + resolve_model_type_from_hf, + resolve_runtime_model_name, +) + +__all__ = [ + "get_model_package", + "get_train_runtime_module", + "register_model", + "resolve_model_type_from_hf", + "resolve_runtime_model_name", +] diff --git a/experimental/lite/megatron/lite/model/qwen3_5/__init__.py b/experimental/lite/megatron/lite/model/qwen3_5/__init__.py new file mode 100644 index 00000000000..d009e1c6f13 --- /dev/null +++ b/experimental/lite/megatron/lite/model/qwen3_5/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Qwen3.5 model package.""" diff --git a/experimental/lite/megatron/lite/model/qwen3_5/config.py b/experimental/lite/megatron/lite/model/qwen3_5/config.py new file mode 100644 index 00000000000..246ca33dc54 --- /dev/null +++ b/experimental/lite/megatron/lite/model/qwen3_5/config.py @@ -0,0 +1,262 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Qwen3.5 model configuration.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +from megatron.lite.primitive.config import load_hf_config_dict + +_HF_FIELDS = frozenset( + { + "num_hidden_layers", + "hidden_size", + "num_attention_heads", + "num_key_value_heads", + "head_dim", + "vocab_size", + "rms_norm_eps", + "max_position_embeddings", + "router_aux_loss_coef", + "num_experts", + "num_experts_per_tok", + "moe_intermediate_size", + "shared_expert_intermediate_size", + "linear_num_key_heads", + "linear_key_head_dim", + "linear_num_value_heads", + "linear_value_head_dim", + "linear_conv_kernel_dim", + "layer_types", + "partial_rotary_factor", + "mrope_section", + "mtp_num_hidden_layers", + "mtp_use_dedicated_embeddings", + "num_nextn_predict_layers", + "mtp_loss_scaling_factor", + "mtp_use_repeated_layer", + "mtp_layer_types", + } +) + + +@dataclass +class Qwen35Config: + """Pure Qwen3.5-35B-A3B architecture parameters.""" + + num_hidden_layers: int = 40 + hidden_size: int = 2048 + num_attention_heads: int = 16 + num_key_value_heads: int = 2 + head_dim: int = 256 + vocab_size: int = 248320 + rms_norm_eps: float = 1e-6 + max_position_embeddings: int = 262144 + router_aux_loss_coef: float = 0.001 + num_experts: int = 256 + num_experts_per_tok: int = 8 + moe_intermediate_size: int = 512 + shared_expert_intermediate_size: int = 512 + linear_num_key_heads: int = 16 + linear_key_head_dim: int = 128 + linear_num_value_heads: int = 32 + linear_value_head_dim: int = 128 + linear_conv_kernel_dim: int = 4 + layer_types: list[str] = field( + default_factory=lambda: ( + ["linear_attention", "linear_attention", "linear_attention", "full_attention"] * 10 + ) + ) + partial_rotary_factor: float = 0.25 + rope_theta: float = 10_000_000.0 + mrope_section: list[int] | None = None + num_nextn_predict_layers: int = 0 + mtp_loss_scaling_factor: float = 0.1 + mtp_use_dedicated_embeddings: bool = False + mtp_use_repeated_layer: bool = False + mtp_layer_types: list[str] = field(default_factory=list) + + # ------------------------------------------------------------------ + # Derived properties + # ------------------------------------------------------------------ + + @property + def rotary_dim(self) -> int: + return int(self.head_dim * self.partial_rotary_factor) + + @property + def full_attn_qkv_size(self) -> int: + return (self.num_attention_heads + 2 * self.num_key_value_heads) * self.head_dim + + @property + def linear_conv_dim(self) -> int: + return self.linear_num_key_heads * self.linear_key_head_dim + + # ------------------------------------------------------------------ + # Validation + # ------------------------------------------------------------------ + + def __post_init__(self): + if ( + self.num_nextn_predict_layers > 0 + and not self.mtp_layer_types + and len(self.layer_types) == self.num_hidden_layers + self.num_nextn_predict_layers + ): + self.mtp_layer_types = self.layer_types[self.num_hidden_layers :] + self.layer_types = self.layer_types[: self.num_hidden_layers] + if self.num_nextn_predict_layers > 0 and not self.mtp_layer_types: + self.mtp_layer_types = ["full_attention"] * self.num_nextn_predict_layers + if len(self.layer_types) != self.num_hidden_layers: + raise ValueError( + f"len(layer_types)={len(self.layer_types)} != " + f"num_hidden_layers={self.num_hidden_layers}" + ) + if ( + self.num_nextn_predict_layers > 0 + and len(self.mtp_layer_types) != self.num_nextn_predict_layers + ): + raise ValueError( + f"len(mtp_layer_types)={len(self.mtp_layer_types)} != " + f"num_nextn_predict_layers={self.num_nextn_predict_layers}" + ) + self._validate() + + def _validate(self): + errors: list[str] = [] + + def _check(cond: bool, msg: str): + if not cond: + errors.append(msg) + + _check( + self.num_hidden_layers >= 1, + f"num_hidden_layers must be >= 1, got {self.num_hidden_layers}", + ) + _check( + self.num_nextn_predict_layers >= 0, + f"num_nextn_predict_layers must be >= 0, got {self.num_nextn_predict_layers}", + ) + _check( + not self.mtp_use_dedicated_embeddings, + "Qwen35Config only supports shared embeddings for MTP " + "(mtp_use_dedicated_embeddings=False)", + ) + _check(self.hidden_size > 0, f"hidden_size must be > 0, got {self.hidden_size}") + _check(self.head_dim > 0, f"head_dim must be > 0, got {self.head_dim}") + _check(self.vocab_size > 0, f"vocab_size must be > 0, got {self.vocab_size}") + _check( + self.num_attention_heads >= 1, + f"num_attention_heads must be >= 1, got {self.num_attention_heads}", + ) + _check( + self.num_attention_heads % self.num_key_value_heads == 0, + f"num_attention_heads({self.num_attention_heads}) must be divisible by " + f"num_key_value_heads({self.num_key_value_heads})", + ) + _check(self.num_experts >= 1, f"num_experts must be >= 1, got {self.num_experts}") + _check( + 1 <= self.num_experts_per_tok <= self.num_experts, + f"num_experts_per_tok({self.num_experts_per_tok}) must be in " + f"[1, num_experts({self.num_experts})]", + ) + _check( + self.moe_intermediate_size > 0, + f"moe_intermediate_size must be > 0, got {self.moe_intermediate_size}", + ) + _check( + self.shared_expert_intermediate_size > 0, + f"shared_expert_intermediate_size must be > 0, got {self.shared_expert_intermediate_size}", + ) + _check( + self.linear_num_key_heads >= 1, + f"linear_num_key_heads must be >= 1, got {self.linear_num_key_heads}", + ) + _check( + self.linear_key_head_dim > 0, + f"linear_key_head_dim must be > 0, got {self.linear_key_head_dim}", + ) + _check( + self.linear_num_value_heads >= 1, + f"linear_num_value_heads must be >= 1, got {self.linear_num_value_heads}", + ) + _check( + self.linear_value_head_dim > 0, + f"linear_value_head_dim must be > 0, got {self.linear_value_head_dim}", + ) + _check( + 0.0 < self.partial_rotary_factor <= 1.0, + f"partial_rotary_factor must be in (0, 1], got {self.partial_rotary_factor}", + ) + if self.mrope_section is not None: + _check( + len(self.mrope_section) == 3, + f"mrope_section must have three entries, got {self.mrope_section}", + ) + _check( + all(section >= 0 for section in self.mrope_section), + f"mrope_section entries must be non-negative, got {self.mrope_section}", + ) + _check( + 2 * sum(self.mrope_section) == self.rotary_dim, + f"sum(mrope_section)*2 must equal rotary_dim({self.rotary_dim}), " + f"got {self.mrope_section}", + ) + valid_types = {"linear_attention", "full_attention"} + for i, lt in enumerate(self.layer_types): + _check(lt in valid_types, f"layer_types[{i}] must be one of {valid_types}, got '{lt}'") + for i, lt in enumerate(self.mtp_layer_types): + _check( + lt in valid_types, f"mtp_layer_types[{i}] must be one of {valid_types}, got '{lt}'" + ) + + if errors: + raise ValueError( + f"Invalid Qwen35Config ({len(errors)} error" + f"{'s' if len(errors) > 1 else ''}):\n " + "\n ".join(errors) + ) + + # ------------------------------------------------------------------ + # Constructors + # ------------------------------------------------------------------ + + @classmethod + def from_hf(cls, path: str, **overrides) -> Qwen35Config: + hf = load_hf_config_dict(path) + return cls._from_hf_dict(hf, **overrides) + + @classmethod + def from_hf_config(cls, hf_config, **overrides) -> Qwen35Config: + return cls._from_hf_dict(hf_config.to_dict(), **overrides) + + @classmethod + def _from_hf_dict(cls, hf: dict, **overrides) -> Qwen35Config: + if "text_config" in hf and isinstance(hf["text_config"], dict): + hf = hf["text_config"] + kwargs = {k: v for k, v in hf.items() if k in _HF_FIELDS} + mtp_num_hidden_layers = kwargs.pop("mtp_num_hidden_layers", None) + if kwargs.get("num_nextn_predict_layers") is None and mtp_num_hidden_layers is not None: + kwargs["num_nextn_predict_layers"] = int(mtp_num_hidden_layers) + if "rope_parameters" in hf and isinstance(hf["rope_parameters"], dict): + rp = hf["rope_parameters"] + if "rope_theta" not in kwargs: + kwargs["rope_theta"] = float(rp.get("rope_theta", 10_000_000.0)) + if "partial_rotary_factor" not in kwargs and "partial_rotary_factor" in rp: + kwargs["partial_rotary_factor"] = float(rp["partial_rotary_factor"]) + if "mrope_section" not in kwargs and "mrope_section" in rp: + kwargs["mrope_section"] = list(rp["mrope_section"]) + if "head_dim" not in kwargs or kwargs.get("head_dim") is None: + kwargs["head_dim"] = kwargs.get("hidden_size", 2048) // kwargs.get( + "num_attention_heads", 16 + ) + if kwargs.get("num_nextn_predict_layers") is None: + kwargs["num_nextn_predict_layers"] = 0 + kwargs.update(overrides) + return cls(**kwargs) + + def layer_type_at(self, layer_idx: int) -> str: + if layer_idx < self.num_hidden_layers: + return self.layer_types[layer_idx] + mtp_idx = layer_idx - self.num_hidden_layers + if 0 <= mtp_idx < len(self.mtp_layer_types): + return self.mtp_layer_types[mtp_idx] + return "full_attention" diff --git a/experimental/lite/megatron/lite/model/qwen3_5/lite/__init__.py b/experimental/lite/megatron/lite/model/qwen3_5/lite/__init__.py new file mode 100644 index 00000000000..acffccfba76 --- /dev/null +++ b/experimental/lite/megatron/lite/model/qwen3_5/lite/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Qwen3.5 lite native sub-package.""" diff --git a/experimental/lite/megatron/lite/model/qwen3_5/lite/checkpoint.py b/experimental/lite/megatron/lite/model/qwen3_5/lite/checkpoint.py new file mode 100644 index 00000000000..3f55303ab68 --- /dev/null +++ b/experimental/lite/megatron/lite/model/qwen3_5/lite/checkpoint.py @@ -0,0 +1,769 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Qwen3.5 lite native checkpoint mapping. + +The loader reads HF safetensors directly into lite's native module names. +It intentionally does not require wrapper-specific state on the model. +""" + +from __future__ import annotations + +import re + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.distributed.tensor import Replicate, Shard + +from megatron.lite.model.qwen3_5.config import Qwen35Config +from megatron.lite.primitive.ckpt.hf_weights import SafeTensorReader, unwrap_model +from megatron.lite.primitive.parallel import ParallelState +from megatron.lite.primitive.utils import ensure_divisible, log_rank0 + + +def EXPERT_CLASSIFIER(name: str) -> bool: + return "experts" in name and "router" not in name and "shared" not in name + + +def PLACEMENT_FN(param_name: str) -> list: + if "experts" in param_name and "router" not in param_name and "shared" not in param_name: + if "fc1" in param_name: + return [Replicate(), Replicate(), Shard(0), Shard(0)] + if "fc2" in param_name: + return [Replicate(), Replicate(), Shard(0), Shard(1)] + return [Replicate(), Replicate(), Replicate(), Replicate()] + if "in_proj" in param_name and "layer_norm" not in param_name: + 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 + ): + return [Replicate(), Replicate(), Replicate(), Shard(1)] + if "gate_up" in param_name and "shared" in param_name: + return [Replicate(), Replicate(), Replicate(), Shard(0)] + if "down" in param_name and "shared" in param_name: + return [Replicate(), Replicate(), Replicate(), Shard(1)] + if "embed" in param_name or "head" in param_name: + return [Replicate(), Replicate(), Replicate(), Shard(0)] + if "conv1d" in param_name or "dt_bias" in param_name or "A_log" in param_name: + return [Replicate(), Replicate(), Replicate(), Shard(0)] + return [Replicate(), Replicate(), Replicate(), Replicate()] + + +def _tp(tensor: torch.Tensor, rank: int, size: int, dim: int = 0) -> torch.Tensor: + return tensor if size <= 1 else tensor.chunk(size, dim=dim)[rank].contiguous() + + +def _split_gate_up(tensor: torch.Tensor, rank: int, size: int) -> torch.Tensor: + if size <= 1: + return tensor + ffn = tensor.shape[0] // 2 + gate = tensor[:ffn].chunk(size, dim=0)[rank] + up = tensor[ffn:].chunk(size, dim=0)[rank] + return torch.cat([gate, up], dim=0).contiguous() + + +def _tp_linear_attn_in_proj( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + z: torch.Tensor, + b: torch.Tensor, + a: torch.Tensor, + *, + ps: ParallelState, +) -> torch.Tensor: + """Shard each GDN projection independently, then pack the local layout.""" + return torch.cat( + [ + _tp(q, ps.tp_rank, ps.tp_size), + _tp(k, ps.tp_rank, ps.tp_size), + _tp(v, ps.tp_rank, ps.tp_size), + _tp(z, ps.tp_rank, ps.tp_size), + _tp(b, ps.tp_rank, ps.tp_size), + _tp(a, ps.tp_rank, ps.tp_size), + ], + dim=0, + ).contiguous() + + +def _tp_linear_attn_conv1d( + tensor: torch.Tensor, *, cfg: Qwen35Config, ps: ParallelState +) -> torch.Tensor: + qk_dim = cfg.linear_num_key_heads * cfg.linear_key_head_dim + v_dim = cfg.linear_num_value_heads * cfg.linear_value_head_dim + q, k, v = tensor.split([qk_dim, qk_dim, v_dim], dim=0) + return torch.cat( + [ + _tp(q, ps.tp_rank, ps.tp_size), + _tp(k, ps.tp_rank, ps.tp_size), + _tp(v, ps.tp_rank, ps.tp_size), + ], + dim=0, + ).contiguous() + + +def _zero_centered_gamma_from_hf(tensor: torch.Tensor) -> torch.Tensor: + return tensor - 1 + + +def _get(reader: SafeTensorReader, name: str) -> torch.Tensor: + return reader.get_tensor(name) + + +def _has(reader: SafeTensorReader, name: str) -> bool: + if reader.index: + return name in reader.index + try: + reader.get_tensor(name) + except Exception: + return False + return True + + +def _load_vocab( + reader: SafeTensorReader, name: str, cfg: Qwen35Config, ps: ParallelState +) -> torch.Tensor: + from megatron.lite.primitive.parallel import pad_vocab_for_tp + + tensor = _get(reader, name) + padded = pad_vocab_for_tp(cfg.vocab_size, ps.tp_size) + if tensor.size(0) < padded: + pad = torch.zeros(padded - tensor.size(0), tensor.size(1), dtype=tensor.dtype) + tensor = torch.cat([tensor, pad], dim=0) + return _tp(tensor, ps.tp_rank, ps.tp_size) + + +def _load_full_attn( + out: dict[str, torch.Tensor], + *, + local_prefix: str, + hf_prefix: str, + input_ln: torch.Tensor, + cfg: Qwen35Config, + ps: ParallelState, + reader: SafeTensorReader, +) -> None: + out[f"{local_prefix}.full_attn.qkv.linear.layer_norm_weight"] = input_ln + q = _get(reader, f"{hf_prefix}.q_proj.weight") + k = _get(reader, f"{hf_prefix}.k_proj.weight") + v = _get(reader, f"{hf_prefix}.v_proj.weight") + out[f"{local_prefix}.full_attn.qkv.linear.weight"] = _tp( + _merge_full_attn_qkvg(q, k, v, cfg=cfg), ps.tp_rank, ps.tp_size + ) + out[f"{local_prefix}.full_attn.q_norm.weight"] = _get(reader, f"{hf_prefix}.q_norm.weight") + out[f"{local_prefix}.full_attn.k_norm.weight"] = _get(reader, f"{hf_prefix}.k_norm.weight") + out[f"{local_prefix}.full_attn.proj.linear.weight"] = _tp( + _get(reader, f"{hf_prefix}.o_proj.weight"), ps.tp_rank, ps.tp_size, dim=1 + ) + + +def _merge_full_attn_qkvg( + q_gate: torch.Tensor, key: torch.Tensor, value: torch.Tensor, *, cfg: Qwen35Config +) -> torch.Tensor: + kv_heads = cfg.num_key_value_heads + head_dim = cfg.head_dim + hidden = q_gate.shape[1] + q_gate = q_gate.reshape(cfg.num_attention_heads, 2 * head_dim, hidden) + query = q_gate.narrow(1, 0, head_dim).reshape(cfg.num_attention_heads * head_dim, hidden) + gate = q_gate.narrow(1, head_dim, head_dim).reshape(cfg.num_attention_heads * head_dim, hidden) + q_heads_per_group = ensure_divisible(cfg.num_attention_heads, cfg.num_key_value_heads) + q_group_width = q_heads_per_group * head_dim + query = query.reshape(kv_heads, q_group_width, hidden) + gate = gate.reshape(kv_heads, q_group_width, hidden) + key = key.reshape(kv_heads, head_dim, hidden) + value = value.reshape(kv_heads, head_dim, hidden) + return torch.cat([query, gate, key, value], dim=1).reshape(-1, hidden).contiguous() + + +def _unmerge_full_attn_qkvg( + tensor: torch.Tensor, *, cfg: Qwen35Config +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Invert Qwen35 lite's full-attention q/g/k/v packing.""" + q_heads_per_group = ensure_divisible(cfg.num_attention_heads, cfg.num_key_value_heads) + group_width = (2 * q_heads_per_group + 2) * cfg.head_dim + hidden = tensor.shape[-1] + packed = tensor.reshape(cfg.num_key_value_heads, group_width, hidden) + query, gate, key, value = packed.split( + [ + q_heads_per_group * cfg.head_dim, + q_heads_per_group * cfg.head_dim, + cfg.head_dim, + cfg.head_dim, + ], + dim=1, + ) + query = query.reshape(cfg.num_attention_heads, cfg.head_dim, hidden) + gate = gate.reshape(cfg.num_attention_heads, cfg.head_dim, hidden) + q_gate = torch.cat([query, gate], dim=1).reshape( + cfg.num_attention_heads * 2 * cfg.head_dim, hidden + ) + key = key.reshape(cfg.num_key_value_heads * cfg.head_dim, hidden) + value = value.reshape(cfg.num_key_value_heads * cfg.head_dim, hidden) + return q_gate.contiguous(), key.contiguous(), value.contiguous() + + +def _split_linear_attn_in_proj( + tensor: torch.Tensor, *, cfg: Qwen35Config +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + qk_dim = cfg.linear_num_key_heads * cfg.linear_key_head_dim + v_dim = cfg.linear_num_value_heads * cfg.linear_value_head_dim + return tensor.split( + [qk_dim, qk_dim, v_dim, v_dim, cfg.linear_num_value_heads, cfg.linear_num_value_heads], + dim=0, + ) + + +def _merge_linear_attn_in_proj_tp_shards( + shards: list[torch.Tensor], *, cfg: Qwen35Config +) -> torch.Tensor: + world_size = len(shards) + qk_dim = ensure_divisible(cfg.linear_num_key_heads * cfg.linear_key_head_dim, world_size) + v_dim = ensure_divisible(cfg.linear_num_value_heads * cfg.linear_value_head_dim, world_size) + value_heads = ensure_divisible(cfg.linear_num_value_heads, world_size) + + parts: list[list[torch.Tensor]] = [[] for _ in range(6)] + for shard in shards: + for bucket, part in zip( + parts, + shard.split([qk_dim, qk_dim, v_dim, v_dim, value_heads, value_heads], dim=0), + strict=True, + ): + bucket.append(part) + + return torch.cat([torch.cat(bucket, dim=0) for bucket in parts], dim=0).contiguous() + + +def _merge_linear_attn_conv1d_tp_shards( + shards: list[torch.Tensor], *, cfg: Qwen35Config +) -> torch.Tensor: + world_size = len(shards) + qk_dim = ensure_divisible(cfg.linear_num_key_heads * cfg.linear_key_head_dim, world_size) + v_dim = ensure_divisible(cfg.linear_num_value_heads * cfg.linear_value_head_dim, world_size) + + parts: list[list[torch.Tensor]] = [[] for _ in range(3)] + for shard in shards: + for bucket, part in zip(parts, shard.split([qk_dim, qk_dim, v_dim], dim=0), strict=True): + bucket.append(part) + + return torch.cat([torch.cat(bucket, dim=0) for bucket in parts], dim=0).contiguous() + + +def _merge_gate_up_tp_shards(shards: list[torch.Tensor]) -> torch.Tensor: + gates: list[torch.Tensor] = [] + ups: list[torch.Tensor] = [] + for shard in shards: + gate, up = shard.chunk(2, dim=0) + gates.append(gate) + ups.append(up) + return torch.cat([torch.cat(gates, dim=0), torch.cat(ups, dim=0)], dim=0).contiguous() + + +def _allgather_tp_shards(tensor: torch.Tensor, ps: ParallelState) -> list[torch.Tensor]: + shards = [torch.empty_like(tensor) for _ in range(ps.tp_size)] + dist.all_gather(shards, tensor.contiguous(), group=ps.tp_group) + return shards + + +class Qwen35WeightSpec: + """Export Qwen35 lite weights to HF checkpoint or vLLM runtime names.""" + + def __init__(self, config: Qwen35Config, target: str = "hf"): + if target not in {"hf", "vllm"}: + raise ValueError(f"Unsupported Qwen3.5 export target: {target!r}") + self.config = config + self.target = target + self._expert_export_buffers: dict[tuple[int, str], dict[int, torch.Tensor]] = {} + + @property + def num_experts(self) -> int: + return self.config.num_experts + + def weight_map(self) -> dict[str, list[str]]: + return {} + + def hf_to_native(self, native_name: str, hf_tensors: list[torch.Tensor]) -> torch.Tensor: + del native_name + return hf_tensors[0] + + def gather_dense( + self, native_name: str, tensor: torch.Tensor, ps: ParallelState + ) -> torch.Tensor | None: + if ps.tp_size <= 1: + return None + if native_name.endswith(".linear_attn.in_proj.linear.weight"): + return _merge_linear_attn_in_proj_tp_shards( + _allgather_tp_shards(tensor, ps), cfg=self.config + ) + if native_name.endswith(".linear_attn.conv1d.weight"): + return _merge_linear_attn_conv1d_tp_shards( + _allgather_tp_shards(tensor, ps), cfg=self.config + ) + if native_name.endswith(".moe.shared_expert.gate_up.linear.weight"): + return _merge_gate_up_tp_shards(_allgather_tp_shards(tensor, ps)) + return None + + def packed_expert_group_name(self, native_name: str) -> str | None: + if re.fullmatch(r"layers\.\d+\.moe\.experts\.fc[12]\.weight\d+", native_name) is None: + return None + return re.sub(r"\.weight\d+$", ".packed", native_name) + + def native_to_hf( + self, native_name: str, tensor: torch.Tensor + ) -> list[tuple[str, torch.Tensor]]: + if self.target == "vllm": + return self._native_to_vllm(native_name, tensor) + + if native_name == "embed.embedding.weight": + return [("model.language_model.embed_tokens.weight", tensor)] + if native_name == "norm.weight": + return [("model.language_model.norm.weight", tensor)] + if native_name == "head.col.linear.weight": + return [("lm_head.weight", tensor)] + if native_name == "mtp_embed.embedding.weight" or native_name.startswith("mtp."): + return [] + + match = re.match(r"layers\.(\d+)\.(.*)", native_name) + if match is None: + return [] + + layer_idx = int(match.group(1)) + suffix = match.group(2) + prefix = f"model.language_model.layers.{layer_idx}" + + if suffix == "full_attn.qkv.linear.layer_norm_weight": + return [(f"{prefix}.input_layernorm.weight", tensor)] + if suffix == "full_attn.qkv.linear.weight": + q_gate, key, value = _unmerge_full_attn_qkvg(tensor, cfg=self.config) + return [ + (f"{prefix}.self_attn.q_proj.weight", q_gate), + (f"{prefix}.self_attn.k_proj.weight", key), + (f"{prefix}.self_attn.v_proj.weight", value), + ] + if suffix == "full_attn.q_norm.weight": + return [(f"{prefix}.self_attn.q_norm.weight", tensor)] + if suffix == "full_attn.k_norm.weight": + return [(f"{prefix}.self_attn.k_norm.weight", tensor)] + if suffix == "full_attn.proj.linear.weight": + return [(f"{prefix}.self_attn.o_proj.weight", tensor)] + + if suffix == "linear_attn.in_proj.linear.layer_norm_weight": + return [(f"{prefix}.input_layernorm.weight", tensor)] + if suffix == "linear_attn.in_proj.linear.weight": + q, k, value, z, b, a = _split_linear_attn_in_proj(tensor, cfg=self.config) + return [ + ( + f"{prefix}.linear_attn.in_proj_qkv.weight", + torch.cat([q, k, value], dim=0).contiguous(), + ), + (f"{prefix}.linear_attn.in_proj_z.weight", z.contiguous()), + (f"{prefix}.linear_attn.in_proj_b.weight", b.contiguous()), + (f"{prefix}.linear_attn.in_proj_a.weight", a.contiguous()), + ] + if suffix == "linear_attn.conv1d.weight": + return [(f"{prefix}.linear_attn.conv1d.weight", tensor)] + if suffix == "linear_attn.dt_bias": + return [(f"{prefix}.linear_attn.dt_bias", tensor)] + if suffix == "linear_attn.A_log": + return [(f"{prefix}.linear_attn.A_log", tensor)] + if suffix == "linear_attn.norm.weight": + return [(f"{prefix}.linear_attn.norm.weight", tensor + 1)] + if suffix == "linear_attn.o_proj.linear.weight": + return [(f"{prefix}.linear_attn.out_proj.weight", tensor)] + + if suffix == "mlp_norm.weight": + return [(f"{prefix}.post_attention_layernorm.weight", tensor)] + if suffix == "moe.router.gate.weight": + return [(f"{prefix}.mlp.gate.weight", tensor)] + if suffix == "moe.shared_expert.gate_up.linear.weight": + gate, up = tensor.chunk(2, dim=0) + return [ + (f"{prefix}.mlp.shared_expert.gate_proj.weight", gate.contiguous()), + (f"{prefix}.mlp.shared_expert.up_proj.weight", up.contiguous()), + ] + if suffix == "moe.shared_expert.down.linear.weight": + return [(f"{prefix}.mlp.shared_expert.down_proj.weight", tensor)] + if suffix == "moe.shared_expert.shared_gate.weight": + return [(f"{prefix}.mlp.shared_expert_gate.weight", tensor)] + + expert_match = re.fullmatch(r"moe\.experts\.fc([12])\.weight(\d+)", suffix) + if expert_match is not None: + kind, expert_idx = expert_match.groups() + buffer_key = (layer_idx, "gate_up" if kind == "1" else "down") + buffer = self._expert_export_buffers.setdefault(buffer_key, {}) + buffer[int(expert_idx)] = tensor.contiguous() + if len(buffer) < self.config.num_experts: + return [] + packed = torch.stack( + [buffer[i] for i in range(self.config.num_experts)], dim=0 + ).contiguous() + del self._expert_export_buffers[buffer_key] + if kind == "1": + return [(f"{prefix}.mlp.experts.gate_up_proj", packed)] + return [(f"{prefix}.mlp.experts.down_proj", packed)] + + packed_expert_match = re.fullmatch(r"moe\.experts\.fc([12])\.packed", suffix) + if packed_expert_match is not None: + kind = packed_expert_match.group(1) + if kind == "1": + return [(f"{prefix}.mlp.experts.gate_up_proj", tensor.contiguous())] + return [(f"{prefix}.mlp.experts.down_proj", tensor.contiguous())] + + return [] + + def _native_to_vllm( + self, native_name: str, tensor: torch.Tensor + ) -> list[tuple[str, torch.Tensor]]: + if native_name == "embed.embedding.weight": + return [("language_model.model.embed_tokens.weight", tensor)] + if native_name == "norm.weight": + return [("language_model.model.norm.weight", tensor)] + if native_name == "head.col.linear.weight": + return [("language_model.lm_head.weight", tensor)] + if native_name == "mtp_embed.embedding.weight" or native_name.startswith("mtp."): + return [] + + match = re.match(r"layers\.(\d+)\.(.*)", native_name) + if match is None: + return [] + + layer_idx = int(match.group(1)) + suffix = match.group(2) + prefix = f"language_model.model.layers.{layer_idx}" + + if suffix == "full_attn.qkv.linear.layer_norm_weight": + return [(f"{prefix}.input_layernorm.weight", tensor)] + if suffix == "full_attn.qkv.linear.weight": + q_gate, key, value = _unmerge_full_attn_qkvg(tensor, cfg=self.config) + return [ + (f"{prefix}.self_attn.q_proj.weight", q_gate), + (f"{prefix}.self_attn.k_proj.weight", key), + (f"{prefix}.self_attn.v_proj.weight", value), + ] + if suffix == "full_attn.q_norm.weight": + return [(f"{prefix}.self_attn.q_norm.weight", tensor)] + if suffix == "full_attn.k_norm.weight": + return [(f"{prefix}.self_attn.k_norm.weight", tensor)] + if suffix == "full_attn.proj.linear.weight": + return [(f"{prefix}.self_attn.o_proj.weight", tensor)] + + if suffix == "linear_attn.in_proj.linear.layer_norm_weight": + return [(f"{prefix}.input_layernorm.weight", tensor)] + if suffix == "linear_attn.in_proj.linear.weight": + q, k, value, z, b, a = _split_linear_attn_in_proj(tensor, cfg=self.config) + return [ + ( + f"{prefix}.linear_attn.in_proj_qkv.weight", + torch.cat([q, k, value], dim=0).contiguous(), + ), + (f"{prefix}.linear_attn.in_proj_z.weight", z.contiguous()), + (f"{prefix}.linear_attn.in_proj_b.weight", b.contiguous()), + (f"{prefix}.linear_attn.in_proj_a.weight", a.contiguous()), + ] + if suffix == "linear_attn.conv1d.weight": + return [(f"{prefix}.linear_attn.conv1d.weight", tensor)] + if suffix == "linear_attn.dt_bias": + return [(f"{prefix}.linear_attn.dt_bias", tensor)] + if suffix == "linear_attn.A_log": + return [(f"{prefix}.linear_attn.A_log", tensor)] + if suffix == "linear_attn.norm.weight": + return [(f"{prefix}.linear_attn.norm.weight", tensor + 1)] + if suffix == "linear_attn.o_proj.linear.weight": + return [(f"{prefix}.linear_attn.out_proj.weight", tensor)] + + if suffix == "mlp_norm.weight": + return [(f"{prefix}.post_attention_layernorm.weight", tensor)] + if suffix == "moe.router.gate.weight": + return [(f"{prefix}.mlp.gate.weight", tensor)] + if suffix == "moe.shared_expert.gate_up.linear.weight": + gate, up = tensor.chunk(2, dim=0) + return [ + (f"{prefix}.mlp.shared_expert.gate_proj.weight", gate.contiguous()), + (f"{prefix}.mlp.shared_expert.up_proj.weight", up.contiguous()), + ] + if suffix == "moe.shared_expert.down.linear.weight": + return [(f"{prefix}.mlp.shared_expert.down_proj.weight", tensor)] + if suffix == "moe.shared_expert.shared_gate.weight": + return [(f"{prefix}.mlp.shared_expert_gate.weight", tensor)] + + expert_match = re.fullmatch(r"moe\.experts\.fc([12])\.weight(\d+)", suffix) + if expert_match is not None: + kind, expert_idx = expert_match.groups() + buffer_key = (layer_idx, "vllm_gate_up" if kind == "1" else "vllm_down") + buffer = self._expert_export_buffers.setdefault(buffer_key, {}) + buffer[int(expert_idx)] = tensor.contiguous() + if len(buffer) < self.config.num_experts: + return [] + packed = torch.stack( + [buffer[i] for i in range(self.config.num_experts)], dim=0 + ).contiguous() + del self._expert_export_buffers[buffer_key] + if kind == "1": + return [(f"{prefix}.mlp.experts.gate_up_proj", packed)] + return [(f"{prefix}.mlp.experts.down_proj", packed)] + + packed_expert_match = re.fullmatch(r"moe\.experts\.fc([12])\.packed", suffix) + if packed_expert_match is not None: + kind = packed_expert_match.group(1) + if kind == "1": + return [(f"{prefix}.mlp.experts.gate_up_proj", tensor.contiguous())] + return [(f"{prefix}.mlp.experts.down_proj", tensor.contiguous())] + + return [] + + def qkv_spec(self, native_name: str) -> tuple[int, int, int] | None: + del native_name + return None + + def tp_spec(self, native_name: str) -> tuple[int, int] | None: + if self.is_expert(native_name): + if ".fc1." in native_name: + return (0, 1) + if ".fc2." in native_name: + return (1, 1) + return None + if native_name in {"embed.embedding.weight", "head.col.linear.weight"}: + return (0, 0) + if native_name.endswith(".full_attn.qkv.linear.weight"): + return (0, 0) + if native_name.endswith(".full_attn.proj.linear.weight"): + return (1, 0) + if native_name.endswith(".linear_attn.in_proj.linear.weight"): + return (0, 0) + if native_name.endswith(".linear_attn.o_proj.linear.weight"): + return (1, 0) + if native_name.endswith(".moe.shared_expert.gate_up.linear.weight"): + return (0, 0) + if native_name.endswith(".moe.shared_expert.down.linear.weight"): + return (1, 0) + if any( + native_name.endswith(suffix) + for suffix in ( + ".linear_attn.conv1d.weight", + ".linear_attn.dt_bias", + ".linear_attn.A_log", + ) + ): + return (0, 0) + return None + + def is_expert(self, native_name: str) -> bool: + return ( + ".moe.experts." in native_name + and ".router." not in native_name + and ".shared" not in native_name + ) + + def expert_global_id(self, native_name: str) -> int | None: + match = re.search(r"\.weight(\d+)$", native_name) + return int(match.group(1)) if match is not None else None + + def expert_local_name(self, native_name: str, local_idx: int) -> str: + return re.sub(r"\.weight\d+$", f".weight{local_idx}", native_name) + + +def _load_linear_attn( + out: dict[str, torch.Tensor], + *, + local_prefix: str, + hf_prefix: str, + input_ln: torch.Tensor, + cfg: Qwen35Config, + ps: ParallelState, + reader: SafeTensorReader, +) -> None: + dk, dv = cfg.linear_key_head_dim, cfg.linear_value_head_dim + nk, nv = cfg.linear_num_key_heads, cfg.linear_num_value_heads + qk_dim, v_dim = nk * dk, nv * dv + qkv = _get(reader, f"{hf_prefix}.in_proj_qkv.weight") + q, k, v = qkv.split([qk_dim, qk_dim, v_dim], dim=0) + z = _get(reader, f"{hf_prefix}.in_proj_z.weight") + b = _get(reader, f"{hf_prefix}.in_proj_b.weight") + a = _get(reader, f"{hf_prefix}.in_proj_a.weight") + + out[f"{local_prefix}.linear_attn.in_proj.linear.weight"] = _tp_linear_attn_in_proj( + q, k, v, z, b, a, ps=ps + ) + out[f"{local_prefix}.linear_attn.in_proj.linear.layer_norm_weight"] = input_ln + out[f"{local_prefix}.linear_attn.conv1d.weight"] = _tp_linear_attn_conv1d( + _get(reader, f"{hf_prefix}.conv1d.weight"), cfg=cfg, ps=ps + ) + out[f"{local_prefix}.linear_attn.dt_bias"] = _tp( + _get(reader, f"{hf_prefix}.dt_bias"), ps.tp_rank, ps.tp_size + ) + out[f"{local_prefix}.linear_attn.A_log"] = _tp( + _get(reader, f"{hf_prefix}.A_log"), ps.tp_rank, ps.tp_size + ) + out[f"{local_prefix}.linear_attn.norm.weight"] = _zero_centered_gamma_from_hf( + _get(reader, f"{hf_prefix}.norm.weight") + ) + out[f"{local_prefix}.linear_attn.o_proj.linear.weight"] = _tp( + _get(reader, f"{hf_prefix}.out_proj.weight"), ps.tp_rank, ps.tp_size, dim=1 + ) + + +def _load_shared_expert( + out: dict[str, torch.Tensor], + *, + local_prefix: str, + hf_mlp_prefix: str, + ps: ParallelState, + reader: SafeTensorReader, +) -> None: + shared = f"{hf_mlp_prefix}.shared_expert" + gate_up = torch.cat( + [_get(reader, f"{shared}.gate_proj.weight"), _get(reader, f"{shared}.up_proj.weight")], + dim=0, + ) + out[f"{local_prefix}.moe.shared_expert.gate_up.linear.weight"] = _split_gate_up( + gate_up, ps.tp_rank, ps.tp_size + ) + out[f"{local_prefix}.moe.shared_expert.down.linear.weight"] = _tp( + _get(reader, f"{shared}.down_proj.weight"), ps.tp_rank, ps.tp_size, dim=1 + ) + out[f"{local_prefix}.moe.shared_expert.shared_gate.weight"] = _get( + reader, f"{hf_mlp_prefix}.shared_expert_gate.weight" + ) + + +def _load_experts( + out: dict[str, torch.Tensor], + *, + local_prefix: str, + hf_mlp_prefix: str, + cfg: Qwen35Config, + ps: ParallelState, + reader: SafeTensorReader, +) -> None: + num_local = ensure_divisible(cfg.num_experts, ps.ep_size) + local_start = ps.ep_rank * num_local + packed_gate_up = f"{hf_mlp_prefix}.experts.gate_up_proj" + packed_down = f"{hf_mlp_prefix}.experts.down_proj" + if _has(reader, packed_gate_up) and _has(reader, packed_down): + gate_up_all = _get(reader, packed_gate_up) + down_all = _get(reader, packed_down) + for local_idx in range(num_local): + global_idx = local_start + local_idx + fc1 = gate_up_all[global_idx] + fc2 = down_all[global_idx] + if ps.etp_size > 1: + fc1 = _split_gate_up(fc1, ps.etp_rank, ps.etp_size) + fc2 = _tp(fc2, ps.etp_rank, ps.etp_size, dim=1) + out[f"{local_prefix}.moe.experts.fc1.weight{local_idx}"] = fc1 + out[f"{local_prefix}.moe.experts.fc2.weight{local_idx}"] = fc2 + return + + for local_idx in range(num_local): + global_idx = local_start + local_idx + ep = f"{hf_mlp_prefix}.experts.{global_idx}" + fc1 = torch.cat( + [_get(reader, f"{ep}.gate_proj.weight"), _get(reader, f"{ep}.up_proj.weight")], dim=0 + ) + fc2 = _get(reader, f"{ep}.down_proj.weight") + if ps.etp_size > 1: + fc1 = _split_gate_up(fc1, ps.etp_rank, ps.etp_size) + fc2 = _tp(fc2, ps.etp_rank, ps.etp_size, dim=1) + out[f"{local_prefix}.moe.experts.fc1.weight{local_idx}"] = fc1 + out[f"{local_prefix}.moe.experts.fc2.weight{local_idx}"] = fc2 + + +def _copy_loaded_state(model: nn.Module, loaded: dict[str, torch.Tensor]) -> None: + state = model.state_dict() + resolved: dict[str, torch.Tensor] = {} + for name, tensor in loaded.items(): + actual = name if name in state else None + if actual is None: + for key in state: + if name in key: + actual = key + break + if actual is not None: + resolved[actual] = tensor + else: + log_rank0(f"WARNING: lite checkpoint tensor has no target param: {name}") + + fp32_names = ("A_log", "dt_bias") + for name, param in model.named_parameters(): + if name not in resolved: + log_rank0(f"WARNING: {name} not loaded from checkpoint") + continue + tensor = resolved[name].to(device=param.device) + if any(k in name for k in fp32_names): + param.data.copy_(tensor.float()) + else: + param.data.copy_(tensor.to(dtype=param.dtype)) + + +def load_hf_weights(model: nn.Module, path: str, config: Qwen35Config, ps: ParallelState) -> None: + base_model = unwrap_model(model) + reader = SafeTensorReader(path) + out: dict[str, torch.Tensor] = {} + + prefix = "model.language_model" + if getattr(base_model, "embed", None) is not None: + out["embed.embedding.weight"] = _load_vocab( + reader, f"{prefix}.embed_tokens.weight", config, ps + ) + if getattr(base_model, "norm", None) is not None: + out["norm.weight"] = _get(reader, f"{prefix}.norm.weight") + if getattr(base_model, "head", None) is not None: + out["head.col.linear.weight"] = _load_vocab(reader, "lm_head.weight", config, ps) + + for local_idx, global_idx in enumerate(base_model.layer_indices): + lp = f"layers.{local_idx}" + hp = f"{prefix}.layers.{global_idx}" + input_ln = _get(reader, f"{hp}.input_layernorm.weight") + if config.layer_type_at(global_idx) == "full_attention": + _load_full_attn( + out, + local_prefix=lp, + hf_prefix=f"{hp}.self_attn", + input_ln=input_ln, + cfg=config, + ps=ps, + reader=reader, + ) + else: + _load_linear_attn( + out, + local_prefix=lp, + hf_prefix=f"{hp}.linear_attn", + input_ln=input_ln, + cfg=config, + ps=ps, + reader=reader, + ) + out[f"{lp}.mlp_norm.weight"] = _get(reader, f"{hp}.post_attention_layernorm.weight") + out[f"{lp}.moe.router.gate.weight"] = _get(reader, f"{hp}.mlp.gate.weight")[ + : config.num_experts + ] + _load_shared_expert(out, local_prefix=lp, hf_mlp_prefix=f"{hp}.mlp", ps=ps, reader=reader) + _load_experts( + out, local_prefix=lp, hf_mlp_prefix=f"{hp}.mlp", cfg=config, ps=ps, reader=reader + ) + + _copy_loaded_state(base_model, out) + + +def export_hf_weights( + model: nn.Module | list[nn.Module], config: Qwen35Config, ps: ParallelState, **kwargs +): + from megatron.lite.primitive.ckpt.hf_weights import export_hf_weights as _export + + include_mtp_only = kwargs.pop("include_mtp_only", False) + kwargs.pop("include_local_prefixes", None) + target = kwargs.pop("target", "hf") + if include_mtp_only: + return + yield from _export( + model, Qwen35WeightSpec(config, target=target), ps, vocab_size=config.vocab_size, **kwargs + ) + + +__all__ = [ + "EXPERT_CLASSIFIER", + "PLACEMENT_FN", + "Qwen35WeightSpec", + "export_hf_weights", + "load_hf_weights", +] diff --git a/experimental/lite/megatron/lite/model/qwen3_5/lite/model.py b/experimental/lite/megatron/lite/model/qwen3_5/lite/model.py new file mode 100644 index 00000000000..9b8f36d6048 --- /dev/null +++ b/experimental/lite/megatron/lite/model/qwen3_5/lite/model.py @@ -0,0 +1,693 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""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. +""" + +from __future__ import annotations + +from contextlib import nullcontext + +import torch +import torch.distributed as dist +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.modules.dispatcher import TokenDispatcher +from megatron.lite.primitive.modules.experts import Experts +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.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 +from megatron.lite.primitive.ops.logprob import vocab_parallel_entropy +from megatron.lite.primitive.parallel import ( + ColumnParallelLinear, + ParallelState, + RowParallelLinear, + VocabParallelEmbedding, + 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 + +_SP_GRAD_SUFFIXES: tuple[str, ...] = ( + ".full_attn.qkv.linear.layer_norm_weight", + ".full_attn.q_norm.weight", + ".full_attn.k_norm.weight", + ".linear_attn.in_proj.linear.layer_norm_weight", + ".linear_attn.norm.weight", + ".mlp_norm.weight", + ".moe.router.gate.weight", + ".moe.shared_expert.shared_gate.weight", + ".enorm.weight", + ".hnorm.weight", + ".final_layernorm.weight", +) + + +def _collect_sp_grad_params(model: nn.Module) -> list[nn.Parameter]: + params = [] + for name, param in model.named_parameters(): + if any(name.endswith(s) for s in _SP_GRAD_SUFFIXES) or name == "norm.weight": + params.append(param) + return params + + +def _swiglu(x: torch.Tensor) -> torch.Tensor: + return bias_swiglu_impl(x, bias=None) + + +def _qwen_mrope_section(config: Qwen35Config) -> list[int]: + section = getattr(config, "mrope_section", None) + if section is not None: + return list(section) + rotary_half = max(int(config.rotary_dim // 2), 1) + base = rotary_half // 3 + return [base, base, rotary_half - 2 * base] + + +class SharedExpert(nn.Module): + _stream: torch.cuda.Stream | None = None + + class _CopyToTPRegion(torch.autograd.Function): + @staticmethod + def forward(ctx, x, group): + ctx.group = group + return x + + @staticmethod + def backward(ctx, grad): + group = ctx.group + if group is not None and dist.get_world_size(group) > 1: + dist.all_reduce(grad, group=group) + return grad, None + + class _PlainTELinear(nn.Module): + def __init__(self, in_features: int, out_features: int): + super().__init__() + self.linear = te.Linear( + in_features, + out_features, + bias=False, + params_dtype=torch.bfloat16, + parallel_mode=None, + sequence_parallel=False, + tp_group=None, + tp_size=1, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(x) + + def __init__( + self, config: Qwen35Config, ps: ParallelState, *, use_plain_te_linear: bool = False + ): + super().__init__() + ffn = config.shared_expert_intermediate_size + if use_plain_te_linear and ps.tp_size == 1: + self.gate_up = self._PlainTELinear(config.hidden_size, ffn * 2) + self.down = self._PlainTELinear(ffn, config.hidden_size) + else: + self.gate_up = ColumnParallelLinear(config.hidden_size, ffn * 2, ps, bias=False) + self.down = RowParallelLinear(ffn, config.hidden_size, ps, bias=False) + self.shared_gate = nn.Linear(config.hidden_size, 1, bias=False) + self.tp_group = ps.tp_group + self.use_mcore_overlap_graph = bool(use_plain_te_linear and ps.tp_size == 1) + + @staticmethod + def _get_stream() -> torch.cuda.Stream: + if SharedExpert._stream is None: + SharedExpert._stream = torch.cuda.Stream() + return SharedExpert._stream + + @staticmethod + def _set_grad_fn_sequence_sr(tensor: torch.Tensor) -> None: + grad_fn = getattr(tensor, "grad_fn", None) + if grad_fn is not None and hasattr(grad_fn, "_set_sequence_nr"): + grad_fn._set_sequence_nr(torch.iinfo(torch.int).max) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_val = self.shared_gate(x).sigmoid() + fc1_input = x + if self.use_mcore_overlap_graph: + fc1_input = self._CopyToTPRegion.apply(x, self.tp_group) + self._set_grad_fn_sequence_sr(fc1_input) + output = self.down(_swiglu(self.gate_up(fc1_input))) + if self.use_mcore_overlap_graph: + self._set_grad_fn_sequence_sr(output) + return output * gate_val + + +class MoELayer(nn.Module): + def __init__( + self, + config: Qwen35Config, + ps: ParallelState, + *, + use_deepep: bool, + router_bias_rate: float, + fp8: bool, + moe_act_recompute: bool, + router_dtype: torch.dtype | None = None, + preserve_3d_graph: bool = False, + shared_expert_plain_te: bool = False, + ): + super().__init__() + if fp8: + raise NotImplementedError("lite qwen35 MoE fp8 is not implemented yet.") + self.router = TopKRouter( + config, + ps, + router_bias_rate=router_bias_rate, + compute_aux_loss=True, + router_dtype=router_dtype, + ) + self.experts = Experts(config, ps, fp8=fp8, moe_act_recompute=moe_act_recompute) + self.dispatcher = TokenDispatcher( + config.num_experts, config.hidden_size, ps, use_deepep=use_deepep + ) + self.shared_expert = SharedExpert(config, ps, use_plain_te_linear=shared_expert_plain_te) + self.preserve_3d_graph = bool(preserve_3d_graph) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + input_shape = x.shape + x_2d = x.reshape(-1, x.size(-1)) + shared_input = x_2d.view(input_shape) if self.preserve_3d_graph else x_2d + router_input = x if self.preserve_3d_graph else x_2d + + shared_out = None + side_stream = None + if x_2d.is_cuda: + side_stream = SharedExpert._get_stream() + side_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side_stream): + shared_out = self.shared_expert(shared_input) + scores, indices = self.router(router_input) + dispatched, tpe, permuted_probs = self.dispatcher.dispatch(x_2d, scores, indices) + del scores, indices + self.dispatcher.wait_dispatch_event() + expert_out = self.experts( + dispatched, + tpe, + permuted_probs, + tokens_per_expert_list=getattr(self.dispatcher, "_local_tpe_list", None), + ) + routed_out = self.dispatcher.combine(expert_out) + + if shared_out is None: + shared_out = self.shared_expert(shared_input) + else: + assert side_stream is not None + torch.cuda.current_stream().wait_stream(side_stream) + output = routed_out.view(input_shape) + if shared_out.shape != input_shape: + shared_out = shared_out.view(input_shape) + output += shared_out + return output.to(x.dtype) + + +class Qwen35Layer(nn.Module): + def __init__( + self, + config: Qwen35Config, + ps: ParallelState, + layer_idx: int, + *, + use_deepep: bool = False, + router_bias_rate: float = 0.0, + fp8: bool = False, + moe_act_recompute: bool = False, + use_thd: bool = False, + deterministic: bool = False, + ): + super().__init__() + self.layer_idx = layer_idx + self._layer_type = config.layer_type_at(layer_idx) + self.full_attn: FullAttention | None = None + self.linear_attn: GatedDeltaNet | None = None + if self._layer_type == "full_attention": + self.full_attn = FullAttention( + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + head_dim=config.head_dim, + ps=ps, + rms_norm_eps=config.rms_norm_eps, + rope_theta=config.rope_theta, + rotary_percent=config.partial_rotary_factor, + use_thd=use_thd, + output_gate=True, + zero_centered_gamma=True, + qkv_layout="mcore", + mrope_section=_qwen_mrope_section(config), + ) + else: + self.linear_attn = GatedDeltaNet( + hidden_size=config.hidden_size, + linear_num_key_heads=config.linear_num_key_heads, + linear_key_head_dim=config.linear_key_head_dim, + linear_num_value_heads=config.linear_num_value_heads, + linear_value_head_dim=config.linear_value_head_dim, + linear_conv_kernel_dim=config.linear_conv_kernel_dim, + rms_norm_eps=config.rms_norm_eps, + ps=ps, + deterministic=deterministic, + ) + self.mlp_norm = te.RMSNorm( + config.hidden_size, eps=config.rms_norm_eps, zero_centered_gamma=True + ) + self.moe = MoELayer( + config, + ps, + use_deepep=use_deepep, + router_bias_rate=router_bias_rate, + fp8=fp8, + moe_act_recompute=moe_act_recompute, + router_dtype=torch.float32 if deterministic else None, + preserve_3d_graph=deterministic, + shared_expert_plain_te=deterministic, + ) + + def forward( + self, x: torch.Tensor, position_ids: torch.Tensor | None = None, packed_seq_params=None + ) -> torch.Tensor: + residual = x + if self.full_attn is not None: + h = self.full_attn(x, position_ids=position_ids, packed_seq_params=packed_seq_params) + else: + assert self.linear_attn is not None + h = self.linear_attn(x, position_ids=position_ids, packed_seq_params=packed_seq_params) + x = residual + h + residual = x + x = residual + self.moe(self.mlp_norm(x)) + return x + + +def _temperature_to_float(temperature: float | torch.Tensor) -> float: + if isinstance(temperature, torch.Tensor): + if temperature.numel() != 1: + raise ValueError("Qwen35Model fused/MTP SFT supports scalar temperature only.") + return float(temperature.detach().float().item()) + return float(temperature) + + +def _ensure_mrope_position_ids(position_ids: torch.Tensor | None) -> torch.Tensor | None: + if position_ids is None: + return None + if position_ids.dim() == 2: + return position_ids.unsqueeze(0).expand(3, -1, -1).contiguous() + if position_ids.dim() == 3: + if position_ids.shape[0] == 1: + return position_ids.expand(3, -1, -1).contiguous() + if position_ids.shape[0] == 3: + return position_ids + raise ValueError("Qwen3.5 MRoPE expects position_ids shape (B,S), (1,B,S), or (3,B,S).") + + +class Qwen35Model(nn.Module): + def __init__( + self, + config: Qwen35Config, + train_config, + ps: ParallelState, + *, + vpp_chunk_id: int | None = None, + router_bias_rate: float = 0.0, + use_thd: bool = False, + hf_path: str = "", + attention_backend_override: str | None = None, + mtp_enable: bool = False, + mtp_enable_train: bool = False, + mtp_detach_encoder: bool = False, + mount_vision_model: bool = False, + ): + super().__init__() + del attention_backend_override + self.config = config + self.train_config = train_config + self.ps = ps + self.mtp_enable_train = bool(mtp_enable and mtp_enable_train) + self.mtp_loss_scaling_factor = config.mtp_loss_scaling_factor + self._input_tensor: torch.Tensor | None = None + + layout = build_pipeline_chunk_layout( + config.num_hidden_layers, ps, train_config.vpp, vpp_chunk_id + ) + self.layer_indices = layout.layer_indices + has_embed = layout.has_embed + has_head = layout.has_head + self.pre_process = has_embed + self.post_process = has_head + self.share_embeddings_and_output_weights = False + self.vision_model: nn.Module | None = ( + _build_native_vision_model(hf_path) if has_embed and mount_vision_model else None + ) + + self.embed: VocabParallelEmbedding | None = None + if has_embed: + self.embed = VocabParallelEmbedding(config.vocab_size, config.hidden_size, ps) + + recompute_modules = getattr(train_config, "recompute_modules", []) + moe_act_recompute = "moe_act" in recompute_modules and "moe" not in recompute_modules + self.layers = nn.ModuleList( + [ + Qwen35Layer( + config, + ps, + idx, + use_deepep=train_config.use_deepep, + router_bias_rate=router_bias_rate, + fp8=train_config.fp8, + moe_act_recompute=moe_act_recompute, + use_thd=use_thd, + deterministic=getattr(train_config, "deterministic", False), + ) + for idx in self.layer_indices + ] + ) + + self.norm: nn.Module | None = None + self.head: VocabParallelOutput | None = None + if has_head: + self.norm = te.RMSNorm( + config.hidden_size, eps=config.rms_norm_eps, zero_centered_gamma=True + ) + self.head = VocabParallelOutput(config.vocab_size, config.hidden_size, ps) + + self.mtp_embed: VocabParallelEmbedding | None = None + self.mtp: MTPBlock | None = None + if mtp_enable and config.num_nextn_predict_layers > 0 and self.head is not None: + mtp_embedding = self.embed + if mtp_embedding is None: + mtp_embedding = VocabParallelEmbedding(config.vocab_size, config.hidden_size, ps) + self.mtp_embed = mtp_embedding + + def make_mtp_layer(layer_idx: int) -> MTPDecoderLayer: + return MTPDecoderLayer( + hidden_size=config.hidden_size, + rms_norm_eps=config.rms_norm_eps, + ps=ps, + embedding=mtp_embedding, + transformer_layer=Qwen35Layer( + config, + ps, + config.num_hidden_layers + layer_idx, + use_deepep=train_config.use_deepep, + router_bias_rate=router_bias_rate, + fp8=train_config.fp8, + moe_act_recompute=moe_act_recompute, + use_thd=use_thd, + deterministic=getattr(train_config, "deterministic", False), + ), + detach_encoder=mtp_detach_encoder, + ) + + self.mtp = MTPBlock( + num_layers=config.num_nextn_predict_layers, + repeated_layer=config.mtp_use_repeated_layer, + layer_factory=make_mtp_layer, + ) + + self.sp_params: list[nn.Parameter] = [] + if ps.tp_size > 1: + self.sp_params = _collect_sp_grad_params(self) + + def set_input_tensor(self, input_tensor): + if isinstance(input_tensor, list): + if len(input_tensor) > 1: + raise ValueError("Qwen35Model expects a single pipeline input tensor.") + input_tensor = input_tensor[0] if input_tensor else None + self._input_tensor = input_tensor + + def forward( + self, + input_ids: torch.Tensor | None = None, + hidden_states: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + packed_seq_params=None, + labels: torch.Tensor | None = None, + loss_mask: torch.Tensor | None = None, + temperature: float | torch.Tensor = 1.0, + use_fused_kernels: bool = False, + calculate_entropy: bool = False, + ) -> dict: + if self.embed is not None: + assert input_ids is not None + h = self.embed(input_ids) + else: + if hidden_states is None: + hidden_states = self._input_tensor + assert hidden_states is not None + h = hidden_states + + if packed_seq_params is not None: + assert position_ids is not None, "THD path requires caller-supplied MRoPE position_ids." + elif position_ids is None and input_ids is not None: + if self.ps.cp_size > 1: + raise ValueError("CP>1 requires caller-supplied FULL position_ids (3,B,S).") + batch, seq = input_ids.shape + pos = torch.arange(seq, device=input_ids.device, dtype=torch.long) + position_ids = pos.unsqueeze(0).unsqueeze(0).expand(3, batch, seq).contiguous() + position_ids = _ensure_mrope_position_ids(position_ids) + + fp8_ctx = ( + te.fp8_autocast(enabled=True, fp8_recipe=build_fp8_recipe(self.train_config)) + if self.train_config.fp8 + else nullcontext() + ) + with fp8_ctx: + if self.embed is not None: + h = scatter_to_sequence_parallel(h, self.ps) + for layer in self.layers: + h = layer(h, position_ids=position_ids, packed_seq_params=packed_seq_params) + + output = {"hidden_states": h} + if self.head is not None: + assert self.norm is not None + hidden_for_head = self.norm(h) + if labels is not None: + temperature_value = _temperature_to_float(temperature) + mtp_result = self._apply_mtp_loss( + hidden_for_head, + input_ids=input_ids, + position_ids=position_ids, + labels=labels, + loss_mask=loss_mask, + packed_seq_params=packed_seq_params, + temperature=temperature_value, + use_fused_kernels=use_fused_kernels, + ) + if mtp_result is not None: + hidden_for_head, mtp_loss = mtp_result + output["mtp_loss"] = mtp_loss + labels_sb = labels.transpose(0, 1).contiguous() + if use_fused_kernels: + hidden_full = gather_from_sequence_parallel(hidden_for_head, self.ps) + log_probs, entropy = linear_cross_entropy( + hidden_full, + self._head_weight_for_fused_ce(hidden_full), + labels_sb, + temperature_value, + self.ps.tp_group, + ) + output["loss"] = (-log_probs).mean() + output["log_probs"] = log_probs.transpose(0, 1).contiguous() + if calculate_entropy: + output["entropy"] = entropy.transpose(0, 1).contiguous() + else: + logits = self.head(hidden_for_head) + if temperature_value != 1.0: + logits = logits / temperature_value + loss = vocab_parallel_cross_entropy(logits, labels_sb, self.ps.tp_group) + output["loss"] = loss.mean() + output["log_probs"] = (-loss).transpose(0, 1).contiguous() + if calculate_entropy: + entropy = vocab_parallel_entropy(logits, self.ps.tp_group) + output["entropy"] = entropy.transpose(0, 1).contiguous() + else: + logits = self.head(hidden_for_head) + output["logits"] = self.head.gather(logits).transpose(0, 1).contiguous() + return output + + def _apply_mtp_loss( + self, + hidden_states: torch.Tensor, + *, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None, + labels: torch.Tensor, + loss_mask: torch.Tensor | None, + packed_seq_params, + temperature: float, + use_fused_kernels: bool, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + if self.mtp is None or not self.mtp_enable_train: + return None + if input_ids is None: + raise ValueError("MTP training requires input_ids.") + if loss_mask is None: + loss_mask = torch.ones_like(labels, dtype=torch.float32) + else: + loss_mask = loss_mask.to(dtype=torch.float32) + + mtp_hidden_states = self.mtp( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + packed_seq_params=packed_seq_params, + ) + mtp_labels = labels.clone() + mtp_loss_mask = loss_mask.clone() + mtp_loss_values = [] + for mtp_hidden in mtp_hidden_states: + mtp_labels, _ = roll_packed_thd_left( + mtp_labels, packed_seq_params=packed_seq_params, dims=-1 + ) + mtp_loss_mask, num_tokens = roll_packed_thd_left( + mtp_loss_mask, packed_seq_params=packed_seq_params, dims=-1 + ) + labels_sb = mtp_labels.transpose(0, 1).contiguous() + mask_sb = mtp_loss_mask.transpose(0, 1).contiguous() + if use_fused_kernels: + mtp_hidden_full = gather_from_sequence_parallel(mtp_hidden, self.ps) + log_probs, _entropy = linear_cross_entropy( + mtp_hidden_full, + self._head_weight_for_fused_ce(mtp_hidden_full), + labels_sb, + temperature, + self.ps.tp_group, + ) + token_loss = -log_probs + else: + assert self.head is not None + logits = self.head(mtp_hidden) + if temperature != 1.0: + logits = logits / temperature + token_loss = vocab_parallel_cross_entropy(logits, labels_sb, self.ps.tp_group) + token_loss = token_loss * mask_sb.to(dtype=token_loss.dtype) + num_tokens = num_tokens.to(dtype=token_loss.dtype).clamp_min(1.0) + mtp_loss_values.append(token_loss.sum() / num_tokens) + mtp_loss_scale = self.mtp_loss_scaling_factor / max(len(mtp_hidden_states), 1) + hidden_states = MTPLossAutoScaler.apply( + hidden_states, mtp_loss_scale * token_loss / num_tokens + ) + + if not mtp_loss_values: + return None + return ( + hidden_states, + torch.stack([loss.detach().float() for loss in mtp_loss_values]).mean(), + ) + + def _head_weight_for_fused_ce(self, hidden_states: torch.Tensor) -> torch.Tensor: + assert self.head is not None + weight = self.head.col.linear.weight + return ( + weight if weight.dtype == hidden_states.dtype else weight.to(dtype=hidden_states.dtype) + ) + + +def _iter_auto_model_class_refs(hf_config) -> list[str]: + auto_map = getattr(hf_config, "auto_map", None) or {} + preferred_keys = ( + "AutoModelForCausalLM", + "AutoModelForImageTextToText", + "AutoModelForVision2Seq", + "AutoModel", + ) + refs: list[str] = [] + for key in preferred_keys: + ref = auto_map.get(key) + if isinstance(ref, (list, tuple)): + ref = ref[0] if ref else None + if isinstance(ref, str) and ref not in refs: + refs.append(ref) + return refs + + +def _resolve_hf_vision_cls(hf_config, hf_path: str) -> type: + try: + from transformers.dynamic_module_utils import get_class_from_dynamic_module + except ImportError as exc: + raise ImportError( + "mount_vision_model=True requires transformers with dynamic module support." + ) from exc + + errors: list[str] = [] + for class_ref in _iter_auto_model_class_refs(hf_config): + try: + model_cls = get_class_from_dynamic_module(class_ref, hf_path, trust_remote_code=True) + except Exception as exc: + errors.append(f"{class_ref}: {exc}") + continue + vision_cls = getattr(model_cls, "HfVisionClass", None) + if vision_cls is not None: + return vision_cls + errors.append(f"{class_ref}: missing HfVisionClass") + detail = "; ".join(errors) if errors else "config auto_map has no supported AutoModel entry" + raise RuntimeError(f"Cannot resolve native HF vision class for Qwen3.5: {detail}.") + + +def _hook_fp32_rotary_emb(module: nn.Module) -> None: + for submodule in module.modules(): + if hasattr(submodule, "inv_freq") and submodule.inv_freq is not None: + submodule._inv_freq_fp32_original = submodule.inv_freq.detach().clone().float() + + def _hook(mod, args): + del args + if hasattr(mod, "_inv_freq_fp32_original"): + mod.inv_freq = mod._inv_freq_fp32_original.to(device=mod.inv_freq.device) + + submodule.register_forward_pre_hook(_hook) + + +def _hook_vision_params_avg_grad_across_tp(module: nn.Module) -> None: + for param in module.parameters(recurse=True): + param.average_gradients_across_tp_domain = True # type: ignore[assignment] + + +def _build_native_vision_model(hf_path: str) -> nn.Module: + if not hf_path: + raise ValueError("mount_vision_model requires hf_path.") + try: + from transformers import AutoConfig + except ImportError as exc: + raise ImportError("mount_vision_model=True requires transformers.") from exc + + 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) + else: + vision = hf_vision_cls(vision_config) + _hook_fp32_rotary_emb(vision) + _hook_vision_params_avg_grad_across_tp(vision) + return vision.to(torch.bfloat16) + + +__all__ = [ + "FullAttention", + "GatedDeltaNet", + "MoELayer", + "MTPLossAutoScaler", + "Qwen35MRoPE", + "Qwen35Layer", + "Qwen35Model", + "SharedExpert", + "_split_grouped_qkvg", +] diff --git a/experimental/lite/megatron/lite/model/qwen3_5/lite/protocol.py b/experimental/lite/megatron/lite/model/qwen3_5/lite/protocol.py new file mode 100644 index 00000000000..a3cbd558bbb --- /dev/null +++ b/experimental/lite/megatron/lite/model/qwen3_5/lite/protocol.py @@ -0,0 +1,267 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Qwen3.5 lite impl — native model protocol for Megatron Lite runtime.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import torch +import torch.nn as nn + +from megatron.lite.model.qwen3_5.config import Qwen35Config +from megatron.lite.model.qwen3_5.lite.checkpoint import EXPERT_CLASSIFIER, PLACEMENT_FN +from megatron.lite.model.qwen3_5.lite.checkpoint import export_hf_weights as _export_hf_weights_impl +from megatron.lite.model.qwen3_5.lite.checkpoint import load_hf_weights as _load_hf_weights_impl +from megatron.lite.primitive.bundle import ModelBundle +from megatron.lite.primitive.parallel import ParallelState, init_parallel +from megatron.lite.primitive.recompute import apply_recompute, parse_recompute_spec +from megatron.lite.runtime.contracts import OptimizerConfig, ParallelConfig + +__all__ = [ + "EXPERT_CLASSIFIER", + "ImplConfig", + "PLACEMENT_FN", + "build_model", + "build_model_config", + "export_hf_weights", + "load_hf_weights", + "vocab_size", +] + + +def is_expert_param(name: str) -> bool: + return "experts" in name and "router" not in name and "shared" not in name + + +@dataclass(frozen=True) +class ImplConfig: + parallel: ParallelConfig = field(default_factory=ParallelConfig) + optimizer: str | None = "mc_full" + recompute: list[str] = field(default_factory=list) + offload: list[str] = field(default_factory=list) + use_deepep: bool = False + use_thd: bool = False + hf_path: str = "" + attention_backend_override: str | None = None + router_aux_loss_coef: float | None = None + router_bias_rate: float = 0.0 + deterministic: bool = True + optimizer_config: OptimizerConfig | None = None + mtp_enable: bool = False + mtp_enable_train: bool = False + mtp_detach_encoder: bool = False + mtp_loss_scaling_factor: float = 0.1 + mtp_use_repeated_layer: bool | None = None + mount_vision_model: bool = False + + +def _full_attn_module(layer, name: str): + full_attn = getattr(layer, "full_attn", None) + return getattr(full_attn, name, None) if full_attn is not None else None + + +MODULE_MAP = { + "core_attn": lambda layer: _full_attn_module(layer, "core_attn"), + "experts": lambda layer: layer.moe.experts, + "moe": lambda layer: layer.moe, + "router": lambda layer: layer.moe.router, + "mlp_norm": lambda layer: layer.mlp_norm, + "attn_proj": lambda layer: _full_attn_module(layer, "proj"), + "linear_attn": lambda layer: layer.linear_attn, +} + + +def build_model_config(source: str | Path | dict, **overrides) -> Qwen35Config: + """Build Qwen3.5 architecture config from HF source.""" + if isinstance(source, dict): + cfg = Qwen35Config._from_hf_dict(source) + else: + cfg = Qwen35Config.from_hf(str(source)) + for k, v in overrides.items(): + if hasattr(cfg, k): + setattr(cfg, k, v) + return cfg + + +def _forward_step(model: nn.Module, batch: dict) -> dict: + kwargs: dict[str, Any] = {"input_ids": batch["input_ids"], "labels": batch["labels"]} + if "position_ids" in batch: + kwargs["position_ids"] = batch["position_ids"] + if "packed_seq_params" in batch: + kwargs["packed_seq_params"] = batch["packed_seq_params"] + for key in ("loss_mask", "temperature", "use_fused_kernels", "calculate_entropy"): + if key in batch: + kwargs[key] = batch[key] + if kwargs["input_ids"].dim() == 1: + kwargs["input_ids"] = kwargs["input_ids"].unsqueeze(0) + return model(**kwargs) + + +def _make_aux_loss_hook(): + from megatron.lite.primitive.modules.moe import MoEAuxLossAutoScaler + from megatron.lite.primitive.modules.mtp import MTPLossAutoScaler + + def hook(scale: torch.Tensor) -> None: + MoEAuxLossAutoScaler.set_loss_scale(scale) + MTPLossAutoScaler.set_loss_scale(scale) + + return hook + + +def _build_mc_optimizer(chunks, model_cfg: Qwen35Config, impl_cfg: ImplConfig, ps: ParallelState): + from megatron.lite.primitive.optimizers.megatron_wrap import build_mc_training_optimizer + + return build_mc_training_optimizer( + chunks, + model_cfg=model_cfg, + impl_cfg=impl_cfg, + ps=ps, + is_expert=is_expert_param, + model_name="qwen3_5", + deterministic=impl_cfg.deterministic, + ) + + +def build_model(model_cfg: Qwen35Config, *, impl_cfg: ImplConfig) -> ModelBundle: + from megatron.lite.model.qwen3_5.lite.model import Qwen35Model + + p = impl_cfg.parallel + + if impl_cfg.use_deepep and (p.etp is not None and p.etp > 1): + raise ValueError("use_deepep and etp>1 are mutually exclusive") + + if impl_cfg.router_aux_loss_coef is not None: + model_cfg.router_aux_loss_coef = impl_cfg.router_aux_loss_coef + mtp_enable = bool(impl_cfg.mtp_enable) + mtp_enable_train = mtp_enable and bool(impl_cfg.mtp_enable_train) + if mtp_enable: + if model_cfg.num_nextn_predict_layers <= 0: + raise ValueError("mtp_enable=True but HF config has no num_nextn_predict_layers.") + model_cfg.mtp_loss_scaling_factor = impl_cfg.mtp_loss_scaling_factor + if impl_cfg.mtp_use_repeated_layer is not None: + model_cfg.mtp_use_repeated_layer = impl_cfg.mtp_use_repeated_layer + else: + model_cfg.num_nextn_predict_layers = 0 + + ps = init_parallel(p) + recompute_spec = parse_recompute_spec(impl_cfg.recompute) + vpp = None if p.vpp == 1 else p.vpp + deterministic = impl_cfg.deterministic + if impl_cfg.use_thd and deterministic and "linear_attention" in model_cfg.layer_types: + deterministic = False + train_cfg = SimpleNamespace( + tp=ps.tp_size, + ep=ps.ep_size, + etp=ps.etp_size, + pp=ps.pp_size, + cp=ps.cp_size, + vpp=vpp, + use_deepep=impl_cfg.use_deepep, + fp8=False, + recompute_modules=recompute_spec, + deterministic=deterministic, + ) + model_kwargs: dict[str, Any] = dict( + router_bias_rate=impl_cfg.router_bias_rate, + use_thd=impl_cfg.use_thd, + hf_path=impl_cfg.hf_path, + attention_backend_override=impl_cfg.attention_backend_override, + mtp_enable=mtp_enable, + mtp_enable_train=mtp_enable_train, + mtp_detach_encoder=impl_cfg.mtp_detach_encoder, + mount_vision_model=impl_cfg.mount_vision_model, + ) + + if vpp is None: + chunks = [Qwen35Model(model_cfg, train_cfg, ps, **model_kwargs).to(torch.bfloat16).cuda()] + else: + chunks = [ + Qwen35Model(model_cfg, train_cfg, ps, vpp_chunk_id=i, **model_kwargs) + .to(torch.bfloat16) + .cuda() + for i in range(vpp) + ] + + if recompute_spec: + for chunk in chunks: + apply_recompute(chunk.layers, recompute_spec, MODULE_MAP) + + if impl_cfg.offload: + from megatron.lite.primitive.recompute import apply_offload + + for chunk in chunks: + apply_offload(chunk.layers, impl_cfg.offload, MODULE_MAP) + + optimizer = None + finalize_grads = None + post_model_load_hook = None + optimizer_backend = "none" + if impl_cfg.optimizer in {"mc", "mc_full"}: + optimizer, finalize_grads = _build_mc_optimizer(chunks, model_cfg, impl_cfg, ps) + from megatron.lite.primitive.ckpt import attach_model_sharded_state_dict + from megatron.lite.runtime.megatron_utils import register_training_hooks + + attach_model_sharded_state_dict( + chunks, ps, get_placements=PLACEMENT_FN, is_expert=is_expert_param + ) + register_training_hooks(chunks, optimizer) + optimizer_backend = "distopt" + elif impl_cfg.optimizer == "fsdp2": + optimizer_backend = "fsdp2" + + def _post_model_load_hook(): + from megatron.lite.model.qwen3_5.lite.model import Qwen35Layer + from megatron.lite.primitive.optimizers.fsdp2 import build_fsdp2_training_optimizer + + return { + "optimizer": build_fsdp2_training_optimizer( + chunks, + impl_cfg.optimizer_config, + ps, + unit_modules=(Qwen35Layer,), + expert_classifier=is_expert_param, + deterministic=deterministic, + vpp=impl_cfg.parallel.vpp, + leaf_module_names=(), + ) + } + + post_model_load_hook = _post_model_load_hook + elif impl_cfg.optimizer is not None: + raise ValueError(f"Unknown qwen3_5 lite optimizer: {impl_cfg.optimizer!r}.") + + return ModelBundle( + chunks=chunks, + parallel_state=ps, + optimizer=optimizer, + finalize_grads=finalize_grads, + forward_step=_forward_step, + extras={ + "model_cfg": model_cfg, + "optimizer_backend": optimizer_backend, + "post_model_load_hook": post_model_load_hook, + "pre_forward_hook": _make_aux_loss_hook(), + }, + ) + + +def load_hf_weights( + chunk: nn.Module, hf_path: str, model_cfg: Qwen35Config, ps: ParallelState +) -> None: + if not hf_path: + return + _load_hf_weights_impl(chunk, hf_path, model_cfg, ps) + + +def export_hf_weights( + chunks: list[nn.Module], model_cfg: Qwen35Config, ps: ParallelState, **kwargs +): + yield from _export_hf_weights_impl(chunks, model_cfg, ps, **kwargs) + + +def vocab_size(model_cfg) -> int | None: + cfg = getattr(model_cfg, "text_config", model_cfg) + return getattr(cfg, "vocab_size", None) diff --git a/experimental/lite/megatron/lite/model/qwen3_5/stats.py b/experimental/lite/megatron/lite/model/qwen3_5/stats.py new file mode 100644 index 00000000000..a22e3f64fcf --- /dev/null +++ b/experimental/lite/megatron/lite/model/qwen3_5/stats.py @@ -0,0 +1,197 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Model-level Qwen3.5 benchmark statistics.""" + +from __future__ import annotations + +from megatron.lite.model.qwen3_5.config import Qwen35Config + + +def num_floating_point_operations( + model_cfg: Qwen35Config, *, seq_len: int, global_batch_size: int, tp_size: int = 1 +) -> int | None: + """Megatron-aligned FLOPs estimate for one training step.""" + + cfg = getattr(model_cfg, "text_config", model_cfg) + + def _get(name: str): + if isinstance(cfg, dict): + return cfg.get(name) + return getattr(cfg, name, None) + + try: + if seq_len <= 0 or global_batch_size <= 0: + return None + + hidden_size = int(_get("hidden_size")) + num_hidden_layers = int(_get("num_hidden_layers")) + if num_hidden_layers <= 0: + return None + + layer_types_obj = _get("layer_types") + if layer_types_obj is not None: + layer_types = list(layer_types_obj) + if len(layer_types) < num_hidden_layers: + return None + layer_types = layer_types[:num_hidden_layers] + else: + full_attention_interval = int(_get("full_attention_interval")) + if full_attention_interval <= 0: + return None + layer_types = [ + "full_attention" if (i + 1) % full_attention_interval == 0 else "linear_attention" + for i in range(num_hidden_layers) + ] + + num_full_attention_layers = sum( + layer_type == "full_attention" for layer_type in layer_types + ) + num_linear_attention_layers = sum( + layer_type == "linear_attention" for layer_type in layer_types + ) + if num_full_attention_layers + num_linear_attention_layers != num_hidden_layers: + return None + + total_tokens = seq_len * global_batch_size + total_tokens_squared = seq_len * seq_len * global_batch_size + + num_attention_heads = int(_get("num_attention_heads")) + num_key_value_heads = int(_get("num_key_value_heads")) + head_dim = int(_get("head_dim")) + query_projection_size = head_dim * num_attention_heads + key_projection_size = head_dim * num_key_value_heads + value_projection_size = head_dim * num_key_value_heads + gate_projection_size = query_projection_size + standard_attention_flops = 6 * ( + total_tokens + * hidden_size + * ( + query_projection_size + + key_projection_size + + value_projection_size + + gate_projection_size + ) + + query_projection_size * total_tokens_squared + + total_tokens * query_projection_size * hidden_size + ) + + linear_num_key_heads = int(_get("linear_num_key_heads")) + linear_key_head_dim = int(_get("linear_key_head_dim")) + linear_num_value_heads = int(_get("linear_num_value_heads")) + linear_value_head_dim = int(_get("linear_value_head_dim")) + linear_conv_kernel_dim = int(_get("linear_conv_kernel_dim")) + qk_dim = linear_key_head_dim * linear_num_key_heads + v_dim = linear_value_head_dim * linear_num_value_heads + linear_attention_flops = ( + 6 + * total_tokens + * ( + hidden_size * (2 * qk_dim + 2 * v_dim + 2 * linear_num_value_heads) + + linear_conv_kernel_dim * (2 * qk_dim + v_dim) + + linear_num_value_heads * (linear_value_head_dim**2) * 4 + + hidden_size * v_dim + ) + ) + + moe_intermediate_size = int(_get("moe_intermediate_size")) + num_experts_per_tok = int(_get("num_experts_per_tok")) + shared_expert_intermediate_size = int(_get("shared_expert_intermediate_size")) + moe_flops = ( + 18 + * total_tokens + * hidden_size + * (moe_intermediate_size * num_experts_per_tok + shared_expert_intermediate_size) + * num_hidden_layers + ) + + from megatron.lite.primitive.parallel.linear import pad_vocab_for_tp + + padded_vocab_size = pad_vocab_for_tp(int(_get("vocab_size")), max(tp_size, 1)) + logits_flops = 6 * total_tokens * hidden_size * padded_vocab_size + + return int( + standard_attention_flops * num_full_attention_layers + + linear_attention_flops * num_linear_attention_layers + + moe_flops + + logits_flops + ) + except (TypeError, ValueError): + return None + + +def activated_params(model_cfg: Qwen35Config) -> int | None: + """Approximate active params/token for benchmark TFLOPS reporting.""" + + cfg = getattr(model_cfg, "text_config", model_cfg) + + def _get(name: str): + if isinstance(cfg, dict): + return cfg.get(name) + return getattr(cfg, name, None) + + try: + hidden_size = int(_get("hidden_size")) + num_hidden_layers = int(_get("num_hidden_layers")) + layer_types_obj = _get("layer_types") + if layer_types_obj is None: + return None + layer_types = list(layer_types_obj) + if not layer_types: + return None + if len(layer_types) < num_hidden_layers: + return None + layer_types = layer_types[:num_hidden_layers] + + num_attention_heads = int(_get("num_attention_heads")) + num_key_value_heads = int(_get("num_key_value_heads")) + head_dim = int(_get("head_dim")) + full_qkv_dim = (num_attention_heads + 2 * num_key_value_heads) * head_dim + full_attention = hidden_size * full_qkv_dim + (num_attention_heads * head_dim) * hidden_size + + linear_num_key_heads = int(_get("linear_num_key_heads")) + linear_key_head_dim = int(_get("linear_key_head_dim")) + linear_num_value_heads = int(_get("linear_num_value_heads")) + linear_value_head_dim = int(_get("linear_value_head_dim")) + linear_conv_kernel_dim = int(_get("linear_conv_kernel_dim")) + qk_dim = linear_num_key_heads * linear_key_head_dim + v_dim = linear_num_value_heads * linear_value_head_dim + gdn_in_proj_dim = 2 * qk_dim + 2 * v_dim + 2 * linear_num_value_heads + gdn_conv_dim = 2 * qk_dim + v_dim + linear_attention = ( + hidden_size * gdn_in_proj_dim + + gdn_conv_dim * linear_conv_kernel_dim + + 2 * linear_num_value_heads + + linear_value_head_dim + + v_dim * hidden_size + ) + + num_experts = int(_get("num_experts")) + num_experts_per_tok = int(_get("num_experts_per_tok")) + moe_intermediate_size = int(_get("moe_intermediate_size")) + shared_expert_intermediate_size = int(_get("shared_expert_intermediate_size")) + router = hidden_size * num_experts + routed_expert = ( + hidden_size * (2 * moe_intermediate_size) + moe_intermediate_size * hidden_size + ) + shared_expert = ( + hidden_size * (2 * shared_expert_intermediate_size) + + shared_expert_intermediate_size * hidden_size + + hidden_size + ) + moe = router + num_experts_per_tok * routed_expert + shared_expert + + total = 0 + for layer_type in layer_types: + if layer_type == "full_attention": + total += full_attention + elif layer_type == "linear_attention": + total += linear_attention + else: + return None + total += moe + + return int(total) + except (TypeError, ValueError): + return None + + +__all__ = ["activated_params", "num_floating_point_operations"] diff --git a/experimental/lite/megatron/lite/model/qwen3_moe/__init__.py b/experimental/lite/megatron/lite/model/qwen3_moe/__init__.py new file mode 100644 index 00000000000..1d6ee75b62b --- /dev/null +++ b/experimental/lite/megatron/lite/model/qwen3_moe/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""qwen3_moe model package for Megatron Lite.""" diff --git a/experimental/lite/megatron/lite/model/qwen3_moe/common.py b/experimental/lite/megatron/lite/model/qwen3_moe/common.py new file mode 100644 index 00000000000..ad291bb5dbb --- /dev/null +++ b/experimental/lite/megatron/lite/model/qwen3_moe/common.py @@ -0,0 +1,11 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Shared Qwen3MoE model helpers.""" + +from __future__ import annotations + + +def is_expert_param(name: str) -> bool: + return "experts" in name and "router" not in name + + +__all__ = ["is_expert_param"] diff --git a/experimental/lite/megatron/lite/model/qwen3_moe/config.py b/experimental/lite/megatron/lite/model/qwen3_moe/config.py new file mode 100644 index 00000000000..7395c2497c5 --- /dev/null +++ b/experimental/lite/megatron/lite/model/qwen3_moe/config.py @@ -0,0 +1,117 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Qwen3MoE model configuration — pure architecture parameters. + +Like HuggingFace's model config: only describes the model architecture. +Impl-specific knobs live in protocol.py. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from dataclasses import fields as dc_fields +from typing import Any + +from megatron.lite.primitive.config import load_hf_config_dict + + +@dataclass +class Qwen3MoEConfig: + """Pure Qwen3MoE architecture parameters (Qwen3-30B-A3B defaults).""" + + num_hidden_layers: int = 48 + hidden_size: int = 2048 + num_attention_heads: int = 32 + num_key_value_heads: int = 4 + head_dim: int = 128 + vocab_size: int = 151936 + num_experts: int = 128 + num_experts_per_tok: int = 8 + moe_intermediate_size: int = 768 + rope_theta: float = 1_000_000.0 + rms_norm_eps: float = 1e-6 + max_position_embeddings: int = 32768 + router_aux_loss_coef: float = 0.001 + num_nextn_predict_layers: int = 0 + mtp_loss_scaling_factor: float = 0.1 + mtp_use_repeated_layer: bool = False + layer_types: list[str] = field(default_factory=lambda: ["full_attention"] * 48) + + def __post_init__(self): + self._validate() + + def _validate(self): + errors: list[str] = [] + + def _check(cond: bool, msg: str): + if not cond: + errors.append(msg) + + _check( + self.hidden_size % self.num_attention_heads == 0, + f"hidden_size({self.hidden_size}) % num_attention_heads({self.num_attention_heads}) != 0", + ) + _check( + self.num_attention_heads % self.num_key_value_heads == 0, + f"num_attention_heads({self.num_attention_heads}) % num_key_value_heads({self.num_key_value_heads}) != 0", + ) + _check(self.head_dim > 0, f"head_dim must be > 0, got {self.head_dim}") + _check(self.num_experts >= 1, f"num_experts must be >= 1, got {self.num_experts}") + _check( + 1 <= self.num_experts_per_tok <= self.num_experts, + f"num_experts_per_tok({self.num_experts_per_tok}) not in [1, {self.num_experts}]", + ) + _check(self.moe_intermediate_size > 0, "moe_intermediate_size must be > 0") + _check(self.vocab_size > 0, "vocab_size must be > 0") + _check(self.num_hidden_layers >= 1, "num_hidden_layers must be >= 1") + _check(self.num_nextn_predict_layers >= 0, "num_nextn_predict_layers must be >= 0") + _check( + len(self.layer_types) == self.num_hidden_layers, + f"len(layer_types)={len(self.layer_types)} != " + f"num_hidden_layers={self.num_hidden_layers}", + ) + valid_types = {"full_attention"} + for i, lt in enumerate(self.layer_types): + _check(lt in valid_types, f"layer_types[{i}] must be one of {valid_types}, got '{lt}'") + + if errors: + raise ValueError( + f"Invalid Qwen3MoEConfig ({len(errors)} errors):\n " + "\n ".join(errors) + ) + + @property + def qkv_size(self) -> int: + return (self.num_attention_heads + 2 * self.num_key_value_heads) * self.head_dim + + def to_dict(self) -> dict[str, Any]: + return {f.name: getattr(self, f.name) for f in dc_fields(self)} + + @classmethod + def from_hf(cls, path_or_name: str, **overrides) -> Qwen3MoEConfig: + hf_dict = load_hf_config_dict(path_or_name) + return cls._from_hf_dict(hf_dict, **overrides) + + @classmethod + def from_hf_config(cls, hf_config, **overrides) -> Qwen3MoEConfig: + hf_dict = hf_config.to_dict() if hasattr(hf_config, "to_dict") else vars(hf_config) + return cls._from_hf_dict(hf_dict, **overrides) + + @classmethod + def _from_hf_dict(cls, hf: dict[str, Any], **overrides) -> Qwen3MoEConfig: + valid_fields = {f.name for f in dc_fields(cls)} + kwargs = {k: v for k, v in hf.items() if k in valid_fields} + + if "rope_theta" not in kwargs: + if "rope_parameters" in hf and isinstance(hf["rope_parameters"], dict): + kwargs["rope_theta"] = float(hf["rope_parameters"].get("rope_theta", 1_000_000.0)) + + if "head_dim" not in kwargs or kwargs["head_dim"] is None: + hs = kwargs.get("hidden_size", 2048) + nh = kwargs.get("num_attention_heads", 32) + kwargs["head_dim"] = hs // nh + if kwargs.get("num_nextn_predict_layers") is None: + kwargs["num_nextn_predict_layers"] = 0 + + if "layer_types" not in kwargs: + kwargs["layer_types"] = ["full_attention"] * kwargs.get("num_hidden_layers", 48) + kwargs.update(overrides) + return cls(**kwargs) diff --git a/experimental/lite/megatron/lite/model/qwen3_moe/lite/__init__.py b/experimental/lite/megatron/lite/model/qwen3_moe/lite/__init__.py new file mode 100644 index 00000000000..2764e0f5880 --- /dev/null +++ b/experimental/lite/megatron/lite/model/qwen3_moe/lite/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Native Qwen3MoE implementation.""" diff --git a/experimental/lite/megatron/lite/model/qwen3_moe/lite/checkpoint.py b/experimental/lite/megatron/lite/model/qwen3_moe/lite/checkpoint.py new file mode 100644 index 00000000000..9f72e67a05b --- /dev/null +++ b/experimental/lite/megatron/lite/model/qwen3_moe/lite/checkpoint.py @@ -0,0 +1,294 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Qwen3MoE WeightSpec — name mapping + format conversion. + +Orchestration (TP scatter, EP sharding, PP remap) lives in +primitive/ckpt/weight_loader.py. This file only defines what's +model-specific: the weight map and tensor conversions. +""" + +from __future__ import annotations + +import torch +from torch.distributed.tensor import Replicate, Shard + +from megatron.lite.model.qwen3_moe.config import Qwen3MoEConfig +from megatron.lite.primitive.ckpt.dcp import ( # noqa: F401 — re-export + canonicalize_fc1_for_dcp, + canonicalize_qkv_for_dcp, + decanon_fc1_after_dcp, + decanon_qkv_after_dcp, +) +from megatron.lite.primitive.ckpt.hf_weights import extract_layer_idx, parse_expert_idx + + +def _pack_mcore_qkv( + q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, config: Qwen3MoEConfig +) -> torch.Tensor: + q_per_group = config.num_attention_heads // config.num_key_value_heads + q = q.view(config.num_key_value_heads, q_per_group * config.head_dim, -1) + k = k.view(config.num_key_value_heads, config.head_dim, -1) + v = v.view(config.num_key_value_heads, config.head_dim, -1) + return torch.cat([q, k, v], dim=1).reshape(-1, q.shape[-1]).contiguous() + + +def _unpack_mcore_qkv( + tensor: torch.Tensor, config: Qwen3MoEConfig +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q_per_group = config.num_attention_heads // config.num_key_value_heads + group_width = (q_per_group + 2) * config.head_dim + packed = tensor.view(config.num_key_value_heads, group_width, -1) + q_end = q_per_group * config.head_dim + k_end = q_end + config.head_dim + q = packed[:, :q_end].reshape(config.num_attention_heads * config.head_dim, -1) + k = packed[:, q_end:k_end].reshape(config.num_key_value_heads * config.head_dim, -1) + v = packed[:, k_end:].reshape(config.num_key_value_heads * config.head_dim, -1) + return q, k, v + + +class Qwen3MoEWeightSpec: + """WeightSpec for Qwen3MoE native impl.""" + + def __init__(self, config: Qwen3MoEConfig): + self.config = config + + @property + def num_experts(self) -> int: + return self.config.num_experts + + def weight_map(self) -> dict[str, list[str]]: + c = self.config + wm: dict[str, list[str]] = { + "embed.embedding.weight": ["model.embed_tokens.weight"], + "mtp_embed.embedding.weight": ["model.embed_tokens.weight"], + "norm.weight": ["model.norm.weight"], + "head.col.linear.weight": ["lm_head.weight"], + } + for li in range(c.num_hidden_layers): + ap = f"model.layers.{li}.self_attn" + mp = f"model.layers.{li}.mlp" + lp = f"layers.{li}" + wm.update( + { + f"{lp}.attn.qkv.linear.layer_norm_weight": [ + f"model.layers.{li}.input_layernorm.weight" + ], + f"{lp}.attn.qkv.linear.weight": [ + f"{ap}.q_proj.weight", + f"{ap}.k_proj.weight", + f"{ap}.v_proj.weight", + ], + f"{lp}.attn.q_norm.weight": [f"{ap}.q_norm.weight"], + f"{lp}.attn.k_norm.weight": [f"{ap}.k_norm.weight"], + f"{lp}.attn.proj.linear.weight": [f"{ap}.o_proj.weight"], + f"{lp}.mlp_norm.weight": [f"model.layers.{li}.post_attention_layernorm.weight"], + f"{lp}.moe.router.gate.weight": [f"{mp}.gate.weight"], + } + ) + for e in range(c.num_experts): + wm[f"{lp}.moe.experts._fc1_weight_{e}"] = [ + f"{mp}.experts.{e}.gate_proj.weight", + f"{mp}.experts.{e}.up_proj.weight", + ] + wm[f"{lp}.moe.experts._fc2_weight_{e}"] = [f"{mp}.experts.{e}.down_proj.weight"] + for mi in range(c.num_nextn_predict_layers): + hf_li = c.num_hidden_layers + mi + hp = f"model.layers.{hf_li}" + ap = f"{hp}.self_attn" + mp = f"{hp}.mlp" + lp = f"mtp.layers.{mi}" + tlp = f"{lp}.transformer_layer" + wm.update( + { + f"{lp}.enorm.weight": [f"{hp}.enorm.weight"], + f"{lp}.hnorm.weight": [f"{hp}.hnorm.weight"], + f"{lp}.eh_proj.linear.weight": [f"{hp}.eh_proj.weight"], + f"{lp}.final_layernorm.weight": [f"{hp}.shared_head.norm.weight"], + f"{tlp}.attn.qkv.linear.layer_norm_weight": [f"{hp}.input_layernorm.weight"], + f"{tlp}.attn.qkv.linear.weight": [ + f"{ap}.q_proj.weight", + f"{ap}.k_proj.weight", + f"{ap}.v_proj.weight", + ], + f"{tlp}.attn.q_norm.weight": [f"{ap}.q_norm.weight"], + f"{tlp}.attn.k_norm.weight": [f"{ap}.k_norm.weight"], + f"{tlp}.attn.proj.linear.weight": [f"{ap}.o_proj.weight"], + f"{tlp}.mlp_norm.weight": [f"{hp}.post_attention_layernorm.weight"], + f"{tlp}.moe.router.gate.weight": [f"{mp}.gate.weight"], + } + ) + for e in range(c.num_experts): + wm[f"{tlp}.moe.experts._fc1_weight_{e}"] = [ + f"{mp}.experts.{e}.gate_proj.weight", + f"{mp}.experts.{e}.up_proj.weight", + ] + wm[f"{tlp}.moe.experts._fc2_weight_{e}"] = [f"{mp}.experts.{e}.down_proj.weight"] + return wm + + def hf_to_native(self, native_name: str, hf_tensors: list[torch.Tensor]) -> torch.Tensor: + if len(hf_tensors) == 3 and "qkv" in native_name: + # Match MCore SelfAttention's local qkv packing: + # [q heads for kv-group 0, k0, v0, q heads for kv-group 1, k1, v1, ...]. + return _pack_mcore_qkv(*hf_tensors, self.config) + if len(hf_tensors) == 2: + # gate + up → concat + return torch.cat(hf_tensors, dim=0) + t = hf_tensors[0] + if "router.gate.weight" in native_name: + return t[: self.config.num_experts] + return t + + def native_to_hf( + self, native_name: str, tensor: torch.Tensor + ) -> list[tuple[str, torch.Tensor]]: + c = self.config + if native_name == "mtp_embed.embedding.weight": + return [] + if native_name.startswith("mtp.layers."): + parts = native_name.split(".") + mtp_idx = int(parts[2]) + hf_li = c.num_hidden_layers + mtp_idx + hp = f"model.layers.{hf_li}" + if native_name.endswith(".enorm.weight"): + return [(f"{hp}.enorm.weight", tensor)] + if native_name.endswith(".hnorm.weight"): + return [(f"{hp}.hnorm.weight", tensor)] + if native_name.endswith(".eh_proj.linear.weight"): + return [(f"{hp}.eh_proj.weight", tensor)] + if native_name.endswith(".final_layernorm.weight"): + return [(f"{hp}.shared_head.norm.weight", tensor)] + proxy = native_name.replace( + f"mtp.layers.{mtp_idx}.transformer_layer", f"layers.{hf_li}" + ) + return self.native_to_hf(proxy, tensor) + if "embed" in native_name and "embedding" in native_name: + return [("model.embed_tokens.weight", tensor)] + if ( + native_name.endswith("norm.weight") + and "layers" not in native_name + and "attn" not in native_name + and "mlp" not in native_name + ): + return [("model.norm.weight", tensor)] + if "head" in native_name: + return [("lm_head.weight", tensor)] + if "layer_norm_weight" in native_name and "qkv" in native_name: + li = extract_layer_idx(native_name) + return [(f"model.layers.{li}.input_layernorm.weight", tensor)] + if "mlp_norm" in native_name: + li = extract_layer_idx(native_name) + return [(f"model.layers.{li}.post_attention_layernorm.weight", tensor)] + if "qkv" in native_name and "layer_norm" not in native_name: + li = extract_layer_idx(native_name) + ap = f"model.layers.{li}.self_attn" + q, k, v = _unpack_mcore_qkv(tensor, c) + return [ + (f"{ap}.q_proj.weight", q), + (f"{ap}.k_proj.weight", k), + (f"{ap}.v_proj.weight", v), + ] + if "q_norm" in native_name: + li = extract_layer_idx(native_name) + return [(f"model.layers.{li}.self_attn.q_norm.weight", tensor)] + if "k_norm" in native_name: + li = extract_layer_idx(native_name) + return [(f"model.layers.{li}.self_attn.k_norm.weight", tensor)] + if "proj.linear" in native_name: + li = extract_layer_idx(native_name) + return [(f"model.layers.{li}.self_attn.o_proj.weight", tensor)] + if "router.gate" in native_name: + li = extract_layer_idx(native_name) + return [(f"model.layers.{li}.mlp.gate.weight", tensor)] + if "experts" in native_name and "fc1" in native_name: + li = extract_layer_idx(native_name) + ei = parse_expert_idx(native_name) + mp = f"model.layers.{li}.mlp" + gate, up = tensor.chunk(2, dim=0) + return [ + (f"{mp}.experts.{ei}.gate_proj.weight", gate), + (f"{mp}.experts.{ei}.up_proj.weight", up), + ] + if "experts" in native_name and "fc2" in native_name: + li = extract_layer_idx(native_name) + ei = parse_expert_idx(native_name) + return [(f"model.layers.{li}.mlp.experts.{ei}.down_proj.weight", tensor)] + return [(native_name, tensor)] + + def qkv_spec(self, native_name: str) -> tuple[int, int, int] | None: + return None + + def tp_spec(self, native_name: str) -> tuple[int, int] | None: + if self.is_expert(native_name): + if "fc1" in native_name: + return (0, 1) # ETP dim 0 + if "fc2" in native_name: + return (1, 1) # ETP dim 1 + return None + if "eh_proj" in native_name: + return (0, 0) + if "qkv" in native_name and "layer_norm" not in native_name: + return (0, 0) + if "proj" in native_name and "attn" in native_name: + return (1, 0) + if "embed" in native_name or "head" in native_name: + return (0, 0) + return None + + def is_expert(self, native_name: str) -> bool: + return "experts" in native_name and "router" not in native_name + + def expert_global_id(self, native_name: str) -> int | None: + if "_fc1_weight_" in native_name or "_fc2_weight_" in native_name: + return int(native_name.split("_")[-1]) + return None + + def expert_local_name(self, native_name: str, local_idx: int) -> str: + prefix = native_name.rsplit("._fc", 1)[0] + fc_tag = "fc1" if "_fc1_weight_" in native_name else "fc2" + return f"{prefix}.{fc_tag}.weight{local_idx}" + + +# --------------------------------------------------------------------------- +# Convenience: standalone functions wrapping WeightSpec + generic loader +# --------------------------------------------------------------------------- + + +def load_hf_weights(model, path: str, config: Qwen3MoEConfig, ps) -> None: + from megatron.lite.primitive.ckpt.hf_weights import load_hf_weights as _load + + _load(model, path, Qwen3MoEWeightSpec(config), ps, vocab_size=config.vocab_size) + + +def export_hf_weights(model, config: Qwen3MoEConfig, ps, **kwargs): + from megatron.lite.primitive.ckpt.hf_weights import export_hf_weights as _export + + yield from _export( + model, Qwen3MoEWeightSpec(config), ps, vocab_size=config.vocab_size, **kwargs + ) + + +def save_hf_weights(model, path: str, config: Qwen3MoEConfig, ps) -> None: + from megatron.lite.primitive.ckpt.hf_weights import save_hf_weights as _save + + _save(model, path, Qwen3MoEWeightSpec(config), ps, vocab_size=config.vocab_size) + + +def EXPERT_CLASSIFIER(name: str) -> bool: + return "experts" in name and "router" not in name + + +def PLACEMENT_FN(param_name: str) -> list: + if "experts" in param_name and "router" not in param_name: + if "fc1" in param_name: + return [Replicate(), Replicate(), Shard(0), Shard(0)] + if "fc2" in param_name: + return [Replicate(), Replicate(), Shard(0), Shard(1)] + return [Replicate(), Replicate(), Replicate(), Replicate()] + if "eh_proj" in param_name: + 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 and "attn" in param_name: + return [Replicate(), Replicate(), Replicate(), Shard(1)] + if "embed" in param_name or "head" in param_name: + return [Replicate(), Replicate(), Replicate(), Shard(0)] + return [Replicate(), Replicate(), Replicate(), Replicate()] diff --git a/experimental/lite/megatron/lite/model/qwen3_moe/lite/lora_adapter.py b/experimental/lite/megatron/lite/model/qwen3_moe/lite/lora_adapter.py new file mode 100644 index 00000000000..ce839108aa8 --- /dev/null +++ b/experimental/lite/megatron/lite/model/qwen3_moe/lite/lora_adapter.py @@ -0,0 +1,803 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""PEFT adapter import/export for Qwen3-MoE lite native LoRA.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import torch +import torch.distributed as dist +import torch.nn as nn + +from megatron.lite.model.qwen3_moe.config import Qwen3MoEConfig +from megatron.lite.primitive.modules.lora import LoraConfig, normalize_lora_config +from megatron.lite.primitive.parallel import ParallelState + +_PEFT_PREFIX = "base_model.model.model" + + +def _rank() -> int: + return dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 + + +def _world_size(group=None) -> int: + if not dist.is_available() or not dist.is_initialized(): + return 1 + return dist.get_world_size(group) + + +def _unwrap_model(module: nn.Module) -> nn.Module: + current = module + seen: set[int] = set() + while hasattr(current, "module") and id(current) not in seen: + seen.add(id(current)) + inner = getattr(current, "module") + if isinstance(inner, list): + if len(inner) != 1: + break + inner = inner[0] + if inner is current or not isinstance(inner, nn.Module): + break + current = inner + return current + + +def _iter_qwen_chunks(chunks: list[nn.Module] | tuple[nn.Module, ...]) -> list[nn.Module]: + return [_unwrap_model(chunk) for chunk in chunks] + + +def _all_gather_cat(tensor: torch.Tensor, group, dim: int) -> torch.Tensor: + if _world_size(group) == 1: + return tensor.contiguous() + gathered = [torch.empty_like(tensor) for _ in range(_world_size(group))] + dist.all_gather(gathered, tensor.contiguous(), group=group) + return torch.cat(gathered, dim=dim).contiguous() + + +def _select_tp_replicated(tensor: torch.Tensor, ps: ParallelState) -> torch.Tensor: + if _world_size(ps.tp_group) == 1: + return tensor.contiguous() + gathered = [torch.empty_like(tensor) for _ in range(_world_size(ps.tp_group))] + dist.all_gather(gathered, tensor.contiguous(), group=ps.tp_group) + return gathered[0].contiguous() + + +def _is_rank_partitioned_lora_a(lora: Any, ps: ParallelState) -> bool: + if getattr(lora, "rank_partitioned_a", False): + return True + rank = int(getattr(lora, "rank", lora.lora_b.shape[1])) + return ps.tp_size > 1 and lora.lora_a.shape[0] != rank + + +def _gather_lora_rank_partition(tensor: torch.Tensor, ps: ParallelState) -> torch.Tensor: + return _all_gather_cat(tensor, ps.tp_group, dim=0) + + +def _slice_lora_rank_partition(tensor: torch.Tensor, ps: ParallelState) -> torch.Tensor: + if ps.tp_size == 1: + return tensor.contiguous() + if tensor.shape[0] % ps.tp_size != 0: + raise ValueError(f"Cannot shard LoRA rank dim {tensor.shape[0]} over TP={ps.tp_size}.") + local_rank = tensor.shape[0] // ps.tp_size + start = ps.tp_rank * local_rank + return tensor[start : start + local_rank].contiguous() + + +def _is_output_partitioned_lora_b(lora: Any, ps: ParallelState) -> bool: + if getattr(lora, "output_partitioned_b", False): + return True + return ps.tp_size > 1 and lora.lora_b.shape[0] * ps.tp_size == getattr(lora, "out_features", -1) + + +def _expert_lora_is_shared(lora: Any) -> bool: + return bool(getattr(lora, "shared_across_experts", False)) or lora.lora_a.dim() == 2 + + +def _expand_shared_expert_lora( + lora: Any, num_local_experts: int +) -> tuple[torch.Tensor, torch.Tensor]: + if _expert_lora_is_shared(lora): + return ( + lora.lora_a.detach().unsqueeze(0).expand(num_local_experts, -1, -1).contiguous(), + lora.lora_b.detach().unsqueeze(0).expand(num_local_experts, -1, -1).contiguous(), + ) + return lora.lora_a.detach(), lora.lora_b.detach() + + +def _split_local_mcore_qkv_b( + qkv_b: torch.Tensor, *, num_heads_local: int, num_kv_heads_local: int, head_dim: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q_per_group = num_heads_local // num_kv_heads_local + group_width = (q_per_group + 2) * head_dim + packed = qkv_b.view(num_kv_heads_local, group_width, -1) + q_end = q_per_group * head_dim + k_end = q_end + head_dim + q = packed[:, :q_end].reshape(num_heads_local * head_dim, -1) + k = packed[:, q_end:k_end].reshape(num_kv_heads_local * head_dim, -1) + v = packed[:, k_end:].reshape(num_kv_heads_local * head_dim, -1) + return q.contiguous(), k.contiguous(), v.contiguous() + + +def _pack_local_mcore_qkv_b( + q_b: torch.Tensor, + k_b: torch.Tensor, + v_b: torch.Tensor, + *, + num_heads_local: int, + num_kv_heads_local: int, + head_dim: int, +) -> torch.Tensor: + q_per_group = num_heads_local // num_kv_heads_local + q = q_b.view(num_kv_heads_local, q_per_group * head_dim, -1) + k = k_b.view(num_kv_heads_local, head_dim, -1) + v = v_b.view(num_kv_heads_local, head_dim, -1) + return torch.cat([q, k, v], dim=1).reshape(-1, q_b.shape[-1]).contiguous() + + +def _layer_prefix(layer_idx: int) -> str: + return f"{_PEFT_PREFIX}.layers.{layer_idx}" + + +def _attn_key(layer_idx: int, module: str, suffix: str) -> str: + return f"{_layer_prefix(layer_idx)}.self_attn.{module}.{suffix}.weight" + + +def _expert_key(layer_idx: int, expert_idx: int, module: str, suffix: str) -> str: + return f"{_layer_prefix(layer_idx)}.mlp.experts.{expert_idx}.{module}.{suffix}.weight" + + +def _target_modules_from_lora_config(lora_config: LoraConfig) -> list[str]: + targets = lora_config.targets() + out: list[str] = [] + if "linear_qkv" in targets: + out += ["q_proj", "k_proj", "v_proj"] + if "linear_proj" in targets: + out.append("o_proj") + if "linear_fc1" in targets: + out += ["gate_proj", "up_proj"] + if "linear_fc2" in targets: + out.append("down_proj") + return out + + +def _state_target_modules(state: dict[str, torch.Tensor]) -> set[str]: + out: set[str] = set() + for key in state: + if ".q_proj." in key: + out.add("q_proj") + elif ".k_proj." in key: + out.add("k_proj") + elif ".v_proj." in key: + out.add("v_proj") + elif ".o_proj." in key: + out.add("o_proj") + elif ".gate_proj." in key: + out.add("gate_proj") + elif ".up_proj." in key: + out.add("up_proj") + elif ".down_proj." in key: + out.add("down_proj") + return out + + +def _effective_lora_alpha(lora_config: LoraConfig) -> int: + return lora_config.rank if lora_config.alpha is None else int(lora_config.alpha) + + +def _peft_target_set(value: Any) -> set[str] | None: + if value is None: + return None + if isinstance(value, str): + return {value} + return {str(item) for item in value} + + +def _infer_state_rank(state: dict[str, torch.Tensor]) -> int | None: + ranks = { + int(tensor.shape[0]) + for key, tensor in state.items() + if key.endswith(".lora_A.weight") and tensor.ndim >= 2 + } + if not ranks: + return None + if len(ranks) != 1: + raise ValueError(f"Adapter contains inconsistent LoRA ranks: {sorted(ranks)}.") + return next(iter(ranks)) + + +def _iter_native_lora_modules(chunks: list[nn.Module] | tuple[nn.Module, ...]): + for chunk in _iter_qwen_chunks(list(chunks)): + for layer in chunk.layers: + attn = layer.attn + if attn.qkv_lora is not None: + yield attn.qkv_lora + if attn.proj_lora is not None: + yield attn.proj_lora + experts = layer.moe.experts + if experts.fc1_lora is not None: + yield experts.fc1_lora + if experts.fc2_lora is not None: + yield experts.fc2_lora + + +def _infer_native_alpha(chunks: list[nn.Module] | tuple[nn.Module, ...]) -> int | None: + alphas: set[int] = set() + for module in _iter_native_lora_modules(chunks): + rank = getattr(module, "rank", None) + scale = getattr(module, "scale", None) + if rank is None or scale is None: + continue + alphas.add(int(round(float(scale) * int(rank)))) + if not alphas: + return None + if len(alphas) != 1: + raise ValueError( + f"Native LoRA modules have inconsistent effective alpha values: {sorted(alphas)}." + ) + return next(iter(alphas)) + + +def _validate_adapter_config( + chunks: list[nn.Module] | tuple[nn.Module, ...], + state: dict[str, torch.Tensor], + adapter_config: dict[str, Any], + *, + lora_config: LoraConfig | dict[str, Any] | None = None, +) -> None: + peft_type = adapter_config.get("peft_type") + if peft_type is not None and str(peft_type).upper() != "LORA": + raise ValueError(f"Expected PEFT adapter_config peft_type='LORA', got {peft_type!r}.") + + state_rank = _infer_state_rank(state) + config_rank = adapter_config.get("r") + if config_rank is not None and state_rank is not None and int(config_rank) != state_rank: + raise ValueError( + f"Adapter config rank r={config_rank} does not match tensor rank {state_rank}." + ) + + config_targets = _peft_target_set(adapter_config.get("target_modules")) + state_targets = _state_target_modules(state) + if config_targets is not None and config_targets != state_targets: + raise ValueError( + "Adapter config target_modules do not match adapter tensors: " + f"config={sorted(config_targets)}, tensors={sorted(state_targets)}." + ) + + config_alpha = adapter_config.get("lora_alpha") + native_alpha = _infer_native_alpha(chunks) + if config_alpha is not None and native_alpha is not None and int(config_alpha) != native_alpha: + raise ValueError( + f"Adapter config lora_alpha={config_alpha} does not match native model alpha={native_alpha}." + ) + + if lora_config is not None: + expected = normalize_lora_config(lora_config) + expected_targets = set(_target_modules_from_lora_config(expected)) + if config_rank is not None and int(config_rank) != expected.rank: + raise ValueError( + f"Adapter config rank r={config_rank} does not match expected rank {expected.rank}." + ) + if config_alpha is not None and int(config_alpha) != _effective_lora_alpha(expected): + raise ValueError( + "Adapter config lora_alpha=" + f"{config_alpha} does not match expected alpha {_effective_lora_alpha(expected)}." + ) + if config_targets is not None and config_targets != expected_targets: + raise ValueError( + "Adapter config target_modules do not match expected LoRA config: " + f"config={sorted(config_targets)}, expected={sorted(expected_targets)}." + ) + + +def _validate_attention_tp(model_cfg: Qwen3MoEConfig, ps: ParallelState) -> None: + if ps.tp_size <= 0: + raise ValueError(f"TP size must be positive, got {ps.tp_size}.") + if model_cfg.num_attention_heads % ps.tp_size != 0: + raise ValueError( + "LoRA adapter import/export requires num_attention_heads " + f"({model_cfg.num_attention_heads}) to be divisible by TP={ps.tp_size}." + ) + if model_cfg.num_key_value_heads % ps.tp_size != 0: + raise ValueError( + "LoRA adapter import/export requires num_key_value_heads " + f"({model_cfg.num_key_value_heads}) to be divisible by TP={ps.tp_size}." + ) + q_heads_local = model_cfg.num_attention_heads // ps.tp_size + kv_heads_local = model_cfg.num_key_value_heads // ps.tp_size + if kv_heads_local <= 0: + raise ValueError( + "LoRA adapter import/export requires at least one local KV head; " + f"got num_key_value_heads={model_cfg.num_key_value_heads}, TP={ps.tp_size}." + ) + if q_heads_local % kv_heads_local != 0: + raise ValueError( + "LoRA adapter import/export requires local query heads to be divisible " + f"by local KV heads, got q={q_heads_local}, kv={kv_heads_local}." + ) + + +def export_lora_adapter_state( + chunks: list[nn.Module] | tuple[nn.Module, ...], + model_cfg: Qwen3MoEConfig, + ps: ParallelState, + *, + cpu: bool = True, +) -> dict[str, torch.Tensor]: + """Export native LoRA tensors to a full PEFT-style adapter state dict. + + All distributed ranks must call this function. Rank 0 returns the full + adapter state; other ranks return an empty dict. + """ + + if ps.pp_size != 1: + raise NotImplementedError("LoRA adapter export currently supports pp=1.") + if ps.etp_size != 1: + raise NotImplementedError("LoRA adapter export currently supports etp=1.") + + _validate_attention_tp(model_cfg, ps) + state: dict[str, torch.Tensor] = {} + q_heads_local = model_cfg.num_attention_heads // ps.tp_size + kv_heads_local = model_cfg.num_key_value_heads // ps.tp_size + + for chunk in _iter_qwen_chunks(list(chunks)): + for layer in chunk.layers: + layer_idx = int(layer.layer_idx) + attn = layer.attn + + if attn.qkv_lora is not None: + if _is_rank_partitioned_lora_a(attn.qkv_lora, ps): + qkv_a = _gather_lora_rank_partition(attn.qkv_lora.lora_a.detach(), ps) + else: + qkv_a = _select_tp_replicated(attn.qkv_lora.lora_a.detach(), ps) + q_b_local, k_b_local, v_b_local = _split_local_mcore_qkv_b( + attn.qkv_lora.lora_b.detach(), + num_heads_local=q_heads_local, + num_kv_heads_local=kv_heads_local, + head_dim=model_cfg.head_dim, + ) + q_b = _all_gather_cat(q_b_local, ps.tp_group, dim=0) + k_b = _all_gather_cat(k_b_local, ps.tp_group, dim=0) + v_b = _all_gather_cat(v_b_local, ps.tp_group, dim=0) + if _rank() == 0: + state[_attn_key(layer_idx, "q_proj", "lora_A")] = qkv_a + state[_attn_key(layer_idx, "q_proj", "lora_B")] = q_b + state[_attn_key(layer_idx, "k_proj", "lora_A")] = qkv_a.clone() + state[_attn_key(layer_idx, "k_proj", "lora_B")] = k_b + state[_attn_key(layer_idx, "v_proj", "lora_A")] = qkv_a.clone() + state[_attn_key(layer_idx, "v_proj", "lora_B")] = v_b + + if attn.proj_lora is not None: + proj_a = _all_gather_cat(attn.proj_lora.lora_a.detach(), ps.tp_group, dim=1) + if _is_output_partitioned_lora_b(attn.proj_lora, ps): + proj_b = _all_gather_cat(attn.proj_lora.lora_b.detach(), ps.tp_group, dim=0) + else: + proj_b = _select_tp_replicated(attn.proj_lora.lora_b.detach(), ps) + if _rank() == 0: + state[_attn_key(layer_idx, "o_proj", "lora_A")] = proj_a + state[_attn_key(layer_idx, "o_proj", "lora_B")] = proj_b + + experts = layer.moe.experts + if experts.fc1_lora is not None: + fc1_a_local, fc1_b_local = _expand_shared_expert_lora( + experts.fc1_lora, experts.num_local_experts + ) + fc1_a = _all_gather_cat(fc1_a_local, ps.ep_group, dim=0) + fc1_b = _all_gather_cat(fc1_b_local, ps.ep_group, dim=0) + gate_b, up_b = fc1_b.chunk(2, dim=1) + if _rank() == 0: + for expert_idx in range(model_cfg.num_experts): + state[_expert_key(layer_idx, expert_idx, "gate_proj", "lora_A")] = fc1_a[ + expert_idx + ] + state[_expert_key(layer_idx, expert_idx, "gate_proj", "lora_B")] = gate_b[ + expert_idx + ] + state[_expert_key(layer_idx, expert_idx, "up_proj", "lora_A")] = fc1_a[ + expert_idx + ].clone() + state[_expert_key(layer_idx, expert_idx, "up_proj", "lora_B")] = up_b[ + expert_idx + ] + + if experts.fc2_lora is not None: + fc2_a_local, fc2_b_local = _expand_shared_expert_lora( + experts.fc2_lora, experts.num_local_experts + ) + fc2_a = _all_gather_cat(fc2_a_local, ps.ep_group, dim=0) + fc2_b = _all_gather_cat(fc2_b_local, ps.ep_group, dim=0) + if _rank() == 0: + for expert_idx in range(model_cfg.num_experts): + state[_expert_key(layer_idx, expert_idx, "down_proj", "lora_A")] = fc2_a[ + expert_idx + ] + state[_expert_key(layer_idx, expert_idx, "down_proj", "lora_B")] = fc2_b[ + expert_idx + ] + + if _rank() != 0: + return {} + if cpu: + return {name: tensor.detach().cpu().contiguous() for name, tensor in state.items()} + return {name: tensor.detach().contiguous() for name, tensor in state.items()} + + +def save_lora_adapter( + chunks: list[nn.Module] | tuple[nn.Module, ...], + model_cfg: Qwen3MoEConfig, + ps: ParallelState, + output_dir: str | Path, + *, + base_model_name_or_path: str = "", + lora_config: LoraConfig | dict[str, Any] | None = None, + metadata: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Save a PEFT/Mint-compatible LoRA adapter directory.""" + + from safetensors.torch import save_file + + if lora_config is None: + raise ValueError("save_lora_adapter requires the LoRA config used to build the model.") + lora = normalize_lora_config(lora_config) + if not lora.enabled: + raise ValueError("save_lora_adapter requires an enabled LoRA config.") + state = export_lora_adapter_state(chunks, model_cfg, ps, cpu=True) + output = Path(output_dir) + if _rank() == 0: + output.mkdir(parents=True, exist_ok=True) + save_file(state, str(output / "adapter_model.safetensors")) + config = { + "peft_type": "LORA", + "task_type": "CAUSAL_LM", + "base_model_name_or_path": base_model_name_or_path, + "inference_mode": False, + "r": lora.rank, + "lora_alpha": _effective_lora_alpha(lora), + "lora_dropout": lora.dropout, + "target_modules": _target_modules_from_lora_config(lora), + "bias": "none", + "fan_in_fan_out": False, + "init_lora_weights": True, + "modules_to_save": None, + } + (output / "adapter_config.json").write_text(json.dumps(config, indent=2) + "\n") + meta = { + "format": "megatron.lite_qwen3_moe_lora_peft_v1", + "expert_lora_representation": ( + "shared_local_expert_group" + if any( + _expert_lora_is_shared(getattr(layer.moe.experts, attr)) + for chunk in _iter_qwen_chunks(list(chunks)) + for layer in chunk.layers + for attr in ("fc1_lora", "fc2_lora") + if getattr(layer.moe.experts, attr) is not None + ) + else "per_expert" + ), + "num_tensors": len(state), + "num_parameters": int(sum(t.numel() for t in state.values())), + "parallel": {"tp": ps.tp_size, "ep": ps.ep_size, "etp": ps.etp_size, "pp": ps.pp_size}, + "model": { + "num_hidden_layers": model_cfg.num_hidden_layers, + "hidden_size": model_cfg.hidden_size, + "num_attention_heads": model_cfg.num_attention_heads, + "num_key_value_heads": model_cfg.num_key_value_heads, + "head_dim": model_cfg.head_dim, + "num_experts": model_cfg.num_experts, + "moe_intermediate_size": model_cfg.moe_intermediate_size, + }, + "metadata": metadata or {}, + } + (output / "megatron.lite_adapter_meta.json").write_text(json.dumps(meta, indent=2) + "\n") + result = { + "path": str(output), + "adapter_model": str(output / "adapter_model.safetensors"), + "adapter_config": str(output / "adapter_config.json"), + **meta, + } + else: + result = {} + if dist.is_available() and dist.is_initialized(): + dist.barrier() + return result + + +def _require_tensor(state: dict[str, torch.Tensor], key: str) -> torch.Tensor: + try: + return state[key] + except KeyError as exc: + raise KeyError(f"Missing adapter tensor {key!r}") from exc + + +def _slice_tp_output(tensor: torch.Tensor, local_width: int, ps: ParallelState) -> torch.Tensor: + start = ps.tp_rank * local_width + return tensor[start : start + local_width].contiguous() + + +def _slice_tp_input(tensor: torch.Tensor, local_width: int, ps: ParallelState) -> torch.Tensor: + start = ps.tp_rank * local_width + return tensor[:, start : start + local_width].contiguous() + + +def load_lora_adapter_state( + chunks: list[nn.Module] | tuple[nn.Module, ...], + state: dict[str, torch.Tensor], + model_cfg: Qwen3MoEConfig, + ps: ParallelState, + *, + strict: bool = True, +) -> dict[str, Any]: + """Load a PEFT adapter state into this rank's native LoRA shards.""" + + if ps.pp_size != 1: + raise NotImplementedError("LoRA adapter import currently supports pp=1.") + if ps.etp_size != 1: + raise NotImplementedError("LoRA adapter import currently supports etp=1.") + + _validate_attention_tp(model_cfg, ps) + q_heads_local = model_cfg.num_attention_heads // ps.tp_size + kv_heads_local = model_cfg.num_key_value_heads // ps.tp_size + q_width_local = q_heads_local * model_cfg.head_dim + kv_width_local = kv_heads_local * model_cfg.head_dim + attn_in_width_local = q_width_local + + loaded = 0 + for chunk in _iter_qwen_chunks(list(chunks)): + for layer in chunk.layers: + layer_idx = int(layer.layer_idx) + attn = layer.attn + if attn.qkv_lora is not None: + q_a = _require_tensor(state, _attn_key(layer_idx, "q_proj", "lora_A")).to( + device=attn.qkv_lora.lora_a.device, dtype=attn.qkv_lora.lora_a.dtype + ) + k_a = _require_tensor(state, _attn_key(layer_idx, "k_proj", "lora_A")).to(q_a) + v_a = _require_tensor(state, _attn_key(layer_idx, "v_proj", "lora_A")).to(q_a) + if strict and (not torch.equal(q_a, k_a) or not torch.equal(q_a, v_a)): + raise ValueError( + "Megatron Lite fused qkv_lora requires q/k/v lora_A tensors to match." + ) + q_a_local = ( + _slice_lora_rank_partition(q_a, ps) + if _is_rank_partitioned_lora_a(attn.qkv_lora, ps) + else q_a.contiguous() + ) + q_b = _slice_tp_output( + _require_tensor(state, _attn_key(layer_idx, "q_proj", "lora_B")).to( + device=attn.qkv_lora.lora_b.device, dtype=attn.qkv_lora.lora_b.dtype + ), + q_width_local, + ps, + ) + k_b = _slice_tp_output( + _require_tensor(state, _attn_key(layer_idx, "k_proj", "lora_B")).to( + device=attn.qkv_lora.lora_b.device, dtype=attn.qkv_lora.lora_b.dtype + ), + kv_width_local, + ps, + ) + v_b = _slice_tp_output( + _require_tensor(state, _attn_key(layer_idx, "v_proj", "lora_B")).to( + device=attn.qkv_lora.lora_b.device, dtype=attn.qkv_lora.lora_b.dtype + ), + kv_width_local, + ps, + ) + attn.qkv_lora.lora_a.data.copy_(q_a_local) + attn.qkv_lora.lora_b.data.copy_( + _pack_local_mcore_qkv_b( + q_b, + k_b, + v_b, + num_heads_local=q_heads_local, + num_kv_heads_local=kv_heads_local, + head_dim=model_cfg.head_dim, + ) + ) + loaded += 2 + + if attn.proj_lora is not None: + proj_a = _slice_tp_input( + _require_tensor(state, _attn_key(layer_idx, "o_proj", "lora_A")).to( + device=attn.proj_lora.lora_a.device, dtype=attn.proj_lora.lora_a.dtype + ), + attn_in_width_local, + ps, + ) + proj_b = _require_tensor(state, _attn_key(layer_idx, "o_proj", "lora_B")).to( + device=attn.proj_lora.lora_b.device, dtype=attn.proj_lora.lora_b.dtype + ) + if _is_output_partitioned_lora_b(attn.proj_lora, ps): + proj_b = _slice_tp_output(proj_b, attn.proj_lora.lora_b.shape[0], ps) + attn.proj_lora.lora_a.data.copy_(proj_a) + attn.proj_lora.lora_b.data.copy_(proj_b) + loaded += 2 + + experts = layer.moe.experts + expert_start = ps.ep_rank * experts.num_local_experts + expert_stop = expert_start + experts.num_local_experts + if experts.fc1_lora is not None: + if _expert_lora_is_shared(experts.fc1_lora): + local_gate_a = [] + local_gate_b = [] + local_up_b = [] + for expert_idx in range(expert_start, expert_stop): + gate_a = _require_tensor( + state, _expert_key(layer_idx, expert_idx, "gate_proj", "lora_A") + ).to( + device=experts.fc1_lora.lora_a.device, + dtype=experts.fc1_lora.lora_a.dtype, + ) + up_a = _require_tensor( + state, _expert_key(layer_idx, expert_idx, "up_proj", "lora_A") + ).to(gate_a) + if strict and not torch.equal(gate_a, up_a): + raise ValueError( + "Megatron Lite fused fc1_lora requires gate/up lora_A tensors to match." + ) + local_gate_a.append(gate_a) + local_gate_b.append( + _require_tensor( + state, _expert_key(layer_idx, expert_idx, "gate_proj", "lora_B") + ).to( + device=experts.fc1_lora.lora_b.device, + dtype=experts.fc1_lora.lora_b.dtype, + ) + ) + local_up_b.append( + _require_tensor( + state, _expert_key(layer_idx, expert_idx, "up_proj", "lora_B") + ).to( + device=experts.fc1_lora.lora_b.device, + dtype=experts.fc1_lora.lora_b.dtype, + ) + ) + if strict: + if any( + not torch.equal(local_gate_a[0], value) for value in local_gate_a[1:] + ): + raise ValueError( + "Megatron Lite shared expert fc1_lora can only import PEFT adapters " + "whose local expert gate lora_A tensors are identical." + ) + if any( + not torch.equal(local_gate_b[0], value) for value in local_gate_b[1:] + ): + raise ValueError( + "Megatron Lite shared expert fc1_lora can only import PEFT adapters " + "whose local expert gate lora_B tensors are identical." + ) + if any(not torch.equal(local_up_b[0], value) for value in local_up_b[1:]): + raise ValueError( + "Megatron Lite shared expert fc1_lora can only import PEFT adapters " + "whose local expert up lora_B tensors are identical." + ) + experts.fc1_lora.lora_a.data.copy_(local_gate_a[0]) + experts.fc1_lora.lora_b.data.copy_( + torch.cat([local_gate_b[0], local_up_b[0]], dim=0) + ) + loaded += 2 + else: + for local_idx, expert_idx in enumerate(range(expert_start, expert_stop)): + gate_a = _require_tensor( + state, _expert_key(layer_idx, expert_idx, "gate_proj", "lora_A") + ).to( + device=experts.fc1_lora.lora_a.device, + dtype=experts.fc1_lora.lora_a.dtype, + ) + up_a = _require_tensor( + state, _expert_key(layer_idx, expert_idx, "up_proj", "lora_A") + ).to(gate_a) + if strict and not torch.equal(gate_a, up_a): + raise ValueError( + "Megatron Lite fused fc1_lora requires gate/up lora_A tensors to match." + ) + gate_b = _require_tensor( + state, _expert_key(layer_idx, expert_idx, "gate_proj", "lora_B") + ).to( + device=experts.fc1_lora.lora_b.device, + dtype=experts.fc1_lora.lora_b.dtype, + ) + up_b = _require_tensor( + state, _expert_key(layer_idx, expert_idx, "up_proj", "lora_B") + ).to( + device=experts.fc1_lora.lora_b.device, + dtype=experts.fc1_lora.lora_b.dtype, + ) + experts.fc1_lora.lora_a.data[local_idx].copy_(gate_a) + experts.fc1_lora.lora_b.data[local_idx].copy_( + torch.cat([gate_b, up_b], dim=0) + ) + loaded += 2 + + if experts.fc2_lora is not None: + if _expert_lora_is_shared(experts.fc2_lora): + local_a = [] + local_b = [] + for expert_idx in range(expert_start, expert_stop): + local_a.append( + _require_tensor( + state, _expert_key(layer_idx, expert_idx, "down_proj", "lora_A") + ).to( + device=experts.fc2_lora.lora_a.device, + dtype=experts.fc2_lora.lora_a.dtype, + ) + ) + local_b.append( + _require_tensor( + state, _expert_key(layer_idx, expert_idx, "down_proj", "lora_B") + ).to( + device=experts.fc2_lora.lora_b.device, + dtype=experts.fc2_lora.lora_b.dtype, + ) + ) + if strict: + if any(not torch.equal(local_a[0], value) for value in local_a[1:]): + raise ValueError( + "Megatron Lite shared expert fc2_lora can only import PEFT adapters " + "whose local expert down lora_A tensors are identical." + ) + if any(not torch.equal(local_b[0], value) for value in local_b[1:]): + raise ValueError( + "Megatron Lite shared expert fc2_lora can only import PEFT adapters " + "whose local expert down lora_B tensors are identical." + ) + experts.fc2_lora.lora_a.data.copy_(local_a[0]) + experts.fc2_lora.lora_b.data.copy_(local_b[0]) + loaded += 2 + else: + for local_idx, expert_idx in enumerate(range(expert_start, expert_stop)): + experts.fc2_lora.lora_a.data[local_idx].copy_( + _require_tensor( + state, _expert_key(layer_idx, expert_idx, "down_proj", "lora_A") + ).to( + device=experts.fc2_lora.lora_a.device, + dtype=experts.fc2_lora.lora_a.dtype, + ) + ) + experts.fc2_lora.lora_b.data[local_idx].copy_( + _require_tensor( + state, _expert_key(layer_idx, expert_idx, "down_proj", "lora_B") + ).to( + device=experts.fc2_lora.lora_b.device, + dtype=experts.fc2_lora.lora_b.dtype, + ) + ) + loaded += 2 + + return {"loaded_tensors": loaded} + + +def load_lora_adapter( + chunks: list[nn.Module] | tuple[nn.Module, ...], + adapter_dir: str | Path, + model_cfg: Qwen3MoEConfig, + ps: ParallelState, + *, + strict: bool = True, + lora_config: LoraConfig | dict[str, Any] | None = None, +) -> dict[str, Any]: + from safetensors.torch import load_file + + path = Path(adapter_dir) / "adapter_model.safetensors" + state = load_file(str(path), device="cpu") + config_path = Path(adapter_dir) / "adapter_config.json" + if config_path.exists(): + adapter_config = json.loads(config_path.read_text()) + _validate_adapter_config(chunks, state, adapter_config, lora_config=lora_config) + elif lora_config is not None: + expected = normalize_lora_config(lora_config) + state_rank = _infer_state_rank(state) + if state_rank is not None and state_rank != expected.rank: + raise ValueError( + f"Adapter tensor rank {state_rank} does not match expected rank {expected.rank}." + ) + return load_lora_adapter_state(chunks, state, model_cfg, ps, strict=strict) + + +__all__ = [ + "export_lora_adapter_state", + "load_lora_adapter", + "load_lora_adapter_state", + "save_lora_adapter", +] diff --git a/experimental/lite/megatron/lite/model/qwen3_moe/lite/model.py b/experimental/lite/megatron/lite/model/qwen3_moe/lite/model.py new file mode 100644 index 00000000000..cc5ee916a14 --- /dev/null +++ b/experimental/lite/megatron/lite/model/qwen3_moe/lite/model.py @@ -0,0 +1,621 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Native Qwen3MoE: TransformerLayer + Qwen3MoEModel. + +Attention and MoE come from primitive/modules; this file only +defines the model-specific composition (Layer stacking, PP layout, +loss computation). +""" + +from __future__ import annotations + +from contextlib import nullcontext + +import torch +import torch.nn as nn +import transformer_engine.pytorch as te + +from megatron.lite.model.qwen3_moe.config import Qwen3MoEConfig +from megatron.lite.primitive.modules.dispatcher import TokenDispatcher +from megatron.lite.primitive.modules.experts import Experts +from megatron.lite.primitive.modules.gqa import GQAttention +from megatron.lite.primitive.modules.lora import LoraConfig +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 +from megatron.lite.primitive.ops.logprob import vocab_parallel_entropy +from megatron.lite.primitive.parallel import ( + ParallelState, + VanillaColumnParallelLinear, + VocabParallelEmbedding, + 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 + +# --------------------------------------------------------------------------- +# MoE Layer (thin assembly over megatron.lite.primitive.modules) +# --------------------------------------------------------------------------- + + +class MoELayer(nn.Module): + def __init__( + self, + config: Qwen3MoEConfig, + ps: ParallelState, + *, + use_deepep: bool = True, + router_bias_rate: float = 0.0, + fp8: bool = False, + moe_act_recompute: bool = False, + lora_config: LoraConfig | dict | None = None, + ): + super().__init__() + # Match Qwen3-MoE's `load_balancing_type="none"` setting: no aux loss. + self.router = TopKRouter( + config, ps, router_bias_rate=router_bias_rate, compute_aux_loss=False + ) + self.experts = Experts( + config, ps, fp8=fp8, moe_act_recompute=moe_act_recompute, lora_config=lora_config + ) + self.dispatcher = TokenDispatcher( + config.num_experts, config.hidden_size, ps, use_deepep=use_deepep + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + input_shape = x.shape + if x.dim() == 3: + x_2d = x.view(-1, x.size(-1)) + else: + x_2d = x + + scores, indices = self.router(x_2d) + dispatched, tpe, permuted_probs = self.dispatcher.dispatch(x_2d, scores, indices) + del scores, indices + self.dispatcher.wait_dispatch_event() + expert_out = self.experts( + dispatched, + tpe, + permuted_probs, + tokens_per_expert_list=getattr(self.dispatcher, "_local_tpe_list", None), + ) + del dispatched, tpe, permuted_probs + combined = self.dispatcher.combine(expert_out) + del expert_out + + return combined.view(input_shape).to(x.dtype) + + +# --------------------------------------------------------------------------- +# Transformer Layer + Model +# --------------------------------------------------------------------------- + +_SP_GRAD_SUFFIXES: tuple[str, ...] = ( + ".attn.qkv.linear.layer_norm_weight", + ".mlp_norm.weight", + ".q_norm.weight", + ".k_norm.weight", + ".moe.router.gate.weight", + ".enorm.weight", + ".hnorm.weight", + ".final_layernorm.weight", +) + + +def _collect_sp_grad_params(model: nn.Module) -> list[nn.Parameter]: + """Collect non-TP-sharded params needing coalesced all_reduce after backward.""" + params = [] + for name, p in model.named_parameters(): + if any(name.endswith(s) for s in _SP_GRAD_SUFFIXES) or name == "norm.weight": + params.append(p) + return params + + +class TransformerLayer(nn.Module): + def __init__( + self, + config: Qwen3MoEConfig, + ps: ParallelState, + layer_idx: int, + *, + use_deepep: bool = True, + router_bias_rate: float = 0.0, + fp8: bool = False, + moe_act_recompute: bool = False, + use_thd: bool = False, + lora_config: LoraConfig | dict | None = None, + ): + super().__init__() + self.layer_idx = layer_idx + + # Declaration order follows MC's TransformerLayer (self_attention → + # pre_mlp_layernorm → mlp). `named_parameters()` iterates in + # declaration order, and MC's `DistributedDataParallel` lays out + # gradient buckets by that order; mismatching it changes fp32 master + # shard layouts and breaks bitwise alignment from step 1 onwards. + self.attn = GQAttention( + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + head_dim=config.head_dim, + ps=ps, + rms_norm_eps=config.rms_norm_eps, + rope_theta=config.rope_theta, + use_thd=use_thd, + qkv_layout="mcore", + lora_config=lora_config, + ) + self.mlp_norm = te.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.moe = MoELayer( + config, + ps, + use_deepep=use_deepep, + router_bias_rate=router_bias_rate, + fp8=fp8, + moe_act_recompute=moe_act_recompute, + lora_config=lora_config, + ) + + def forward( + self, x: torch.Tensor, position_ids: torch.Tensor | None = None, packed_seq_params=None + ) -> torch.Tensor: + residual = x + h = self.attn(x, position_ids=position_ids, packed_seq_params=packed_seq_params) + x = residual + h + + residual = x + h = self.mlp_norm(x) + moe_out = self.moe(h) + x = residual + moe_out + + return x + + +class MTPLossAutoScaler(torch.autograd.Function): + """Attach MTP loss gradients to the main LM hidden state.""" + + main_loss_backward_scale: float = 1.0 + + @staticmethod + def forward(ctx, output: torch.Tensor, mtp_loss: torch.Tensor): + ctx.save_for_backward(mtp_loss) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + (mtp_loss,) = ctx.saved_tensors + scaled_mtp_grad = torch.ones_like(mtp_loss) * MTPLossAutoScaler.main_loss_backward_scale + return grad_output, scaled_mtp_grad + + @staticmethod + def set_loss_scale(scale: torch.Tensor | float) -> None: + if isinstance(scale, torch.Tensor): + scale = float(scale.detach().float().item()) + MTPLossAutoScaler.main_loss_backward_scale = float(scale) + + +class MultiTokenPredictionLayer(nn.Module): + """MCore-style MTP layer for the THD SFT lite path.""" + + def __init__( + self, + config: Qwen3MoEConfig, + ps: ParallelState, + layer_idx: int, + *, + embedding: VocabParallelEmbedding, + use_deepep: bool, + router_bias_rate: float, + fp8: bool, + moe_act_recompute: bool, + use_thd: bool, + detach_encoder: bool, + lora_config: LoraConfig | dict | None, + ): + super().__init__() + self.ps = ps + self.embedding = embedding + self.detach_encoder = detach_encoder + self.enorm = te.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = te.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.eh_proj = VanillaColumnParallelLinear( + config.hidden_size * 2, config.hidden_size, ps, sp=ps.tp_size > 1, gather_output=True + ) + self.transformer_layer = TransformerLayer( + config, + ps, + config.num_hidden_layers + layer_idx, + use_deepep=use_deepep, + router_bias_rate=router_bias_rate, + fp8=fp8, + moe_act_recompute=moe_act_recompute, + use_thd=use_thd, + lora_config=lora_config, + ) + self.final_layernorm = te.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + *, + input_ids: torch.Tensor, + position_ids: torch.Tensor | None, + hidden_states: torch.Tensor, + rotary_position_ids: torch.Tensor | None = None, + packed_seq_params=None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + attention_position_ids = ( + rotary_position_ids if rotary_position_ids is not None else position_ids + ) + input_ids, _ = roll_packed_thd_left(input_ids, packed_seq_params=packed_seq_params, dims=-1) + if position_ids is not None: + position_ids, _ = roll_packed_thd_left( + position_ids, packed_seq_params=packed_seq_params, dims=-1 + ) + decoder_input = self.embedding(input_ids) + decoder_input = scatter_to_sequence_parallel(decoder_input, self.ps) + + if self.detach_encoder: + decoder_input = decoder_input.detach() + hidden_states = hidden_states.detach() + + decoder_input = self.enorm(decoder_input) + hidden_states = self.hnorm(hidden_states) + hidden_states = torch.cat((decoder_input, hidden_states), dim=-1) + hidden_states = self.eh_proj(hidden_states) + hidden_states = scatter_to_sequence_parallel(hidden_states, self.ps) + hidden_states = self.transformer_layer( + hidden_states, position_ids=attention_position_ids, packed_seq_params=packed_seq_params + ) + hidden_states = self.final_layernorm(hidden_states) + return hidden_states, input_ids, position_ids + + +class MultiTokenPredictionBlock(nn.Module): + def __init__( + self, + config: Qwen3MoEConfig, + ps: ParallelState, + *, + embedding: VocabParallelEmbedding, + use_deepep: bool, + router_bias_rate: float, + fp8: bool, + moe_act_recompute: bool, + use_thd: bool, + detach_encoder: bool, + repeated_layer: bool, + lora_config: LoraConfig | dict | None, + ): + super().__init__() + self.num_layers = config.num_nextn_predict_layers + self.repeated_layer = repeated_layer + layers_to_build = 1 if repeated_layer else self.num_layers + self.layers = nn.ModuleList( + [ + MultiTokenPredictionLayer( + config, + ps, + idx, + embedding=embedding, + use_deepep=use_deepep, + router_bias_rate=router_bias_rate, + fp8=fp8, + moe_act_recompute=moe_act_recompute, + use_thd=use_thd, + detach_encoder=detach_encoder, + lora_config=lora_config, + ) + for idx in range(layers_to_build) + ] + ) + + def forward( + self, + *, + input_ids: torch.Tensor, + position_ids: torch.Tensor | None, + hidden_states: torch.Tensor, + packed_seq_params=None, + ) -> list[torch.Tensor]: + outputs: list[torch.Tensor] = [] + rotary_position_ids = position_ids + for depth in range(self.num_layers): + layer = self.layers[0] if self.repeated_layer else self.layers[depth] + hidden_states, input_ids, position_ids = layer( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + rotary_position_ids=rotary_position_ids, + packed_seq_params=packed_seq_params, + ) + outputs.append(hidden_states) + return outputs + + +def _temperature_to_float(temperature: float | torch.Tensor) -> float: + if isinstance(temperature, torch.Tensor): + if temperature.numel() != 1: + raise ValueError( + "Megatron Lite fused/MTP SFT currently supports scalar temperature only." + ) + return float(temperature.detach().float().item()) + return float(temperature) + + +class Qwen3MoEModel(nn.Module): + def __init__( + self, + config: Qwen3MoEConfig, + ps: ParallelState, + vpp: int | None = None, + vpp_chunk_id: int | None = None, + *, + use_deepep: bool = False, + fp8: bool = False, + recompute_modules: list[str] | None = None, + router_bias_rate: float = 0.0, + use_thd: bool = False, + mtp_enable: bool = False, + mtp_enable_train: bool = False, + mtp_detach_encoder: bool = False, + lora_config: LoraConfig | dict | None = None, + ): + super().__init__() + self.config = config + self.ps = ps + self.fp8 = fp8 + self.mtp_enable_train = bool(mtp_enable and mtp_enable_train) + self.mtp_loss_scaling_factor = config.mtp_loss_scaling_factor + self._input_tensor: torch.Tensor | None = None + layout = build_pipeline_chunk_layout(config.num_hidden_layers, ps, vpp, vpp_chunk_id) + self.layer_indices = layout.layer_indices + has_embed = layout.has_embed + has_head = layout.has_head + self.pre_process = has_embed + self.post_process = has_head + self.share_embeddings_and_output_weights = False + + self.embed: VocabParallelEmbedding | None = None + if has_embed: + self.embed = VocabParallelEmbedding(config.vocab_size, config.hidden_size, ps) + + _recompute = recompute_modules or [] + moe_act_recompute = "moe_act" in _recompute and "moe" not in _recompute + self.layers = nn.ModuleList( + [ + TransformerLayer( + config, + ps, + idx, + use_deepep=use_deepep, + router_bias_rate=router_bias_rate, + fp8=fp8, + moe_act_recompute=moe_act_recompute, + use_thd=use_thd, + lora_config=lora_config, + ) + for idx in self.layer_indices + ] + ) + + self.norm: nn.Module | None = None + self.head: VocabParallelOutput | None = None + if has_head: + self.norm = te.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.head = VocabParallelOutput(config.vocab_size, config.hidden_size, ps) + + self.mtp_embed: VocabParallelEmbedding | None = None + self.mtp: MultiTokenPredictionBlock | None = None + if mtp_enable and config.num_nextn_predict_layers > 0 and self.head is not None: + mtp_embedding = self.embed + if mtp_embedding is None: + mtp_embedding = VocabParallelEmbedding(config.vocab_size, config.hidden_size, ps) + self.mtp_embed = mtp_embedding + self.mtp = MultiTokenPredictionBlock( + config, + ps, + embedding=mtp_embedding, + use_deepep=use_deepep, + router_bias_rate=router_bias_rate, + fp8=fp8, + moe_act_recompute=moe_act_recompute, + use_thd=use_thd, + detach_encoder=mtp_detach_encoder, + repeated_layer=config.mtp_use_repeated_layer, + lora_config=lora_config, + ) + + self.sp_params: list[nn.Parameter] = [] + if ps.tp_size > 1: + self.sp_params = _collect_sp_grad_params(self) + + def set_input_tensor(self, input_tensor): + if isinstance(input_tensor, list): + if len(input_tensor) > 1: + raise ValueError("Qwen3MoEModel expects a single pipeline input tensor.") + input_tensor = input_tensor[0] if input_tensor else None + self._input_tensor = input_tensor + + def forward( + self, + input_ids: torch.Tensor | None = None, + hidden_states: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + packed_seq_params=None, + labels: torch.Tensor | None = None, + loss_mask: torch.Tensor | None = None, + temperature: float | torch.Tensor = 1.0, + use_fused_kernels: bool = False, + calculate_entropy: bool = False, + return_log_probs: bool = True, + ) -> dict: + if self.embed is not None: + assert input_ids is not None + h = self.embed(input_ids) + else: + if hidden_states is None: + hidden_states = self._input_tensor + assert hidden_states is not None + h = hidden_states + + fp8_ctx = ( + te.fp8_autocast(enabled=True, fp8_recipe=build_fp8_recipe()) + if self.fp8 + else nullcontext() + ) + + with fp8_ctx: + if self.embed is not None: + h = scatter_to_sequence_parallel(h, self.ps) + for layer in self.layers: + h = layer(h, position_ids=position_ids, packed_seq_params=packed_seq_params) + # Head path is SP-aware: norm runs on SP-sharded [S/tp, B, H] and + # head's internal all-gather happens inside VocabParallelOutput. + # Mirrors MC GPTModel's final_layernorm → output_layer(sp=True). + + output = {"hidden_states": h} + + if self.head is not None: + hidden_for_head = self.norm(h) + + if labels is not None: + temperature_value = _temperature_to_float(temperature) + mtp_result = self._apply_mtp_loss( + hidden_for_head, + input_ids=input_ids, + position_ids=position_ids, + labels=labels, + loss_mask=loss_mask, + packed_seq_params=packed_seq_params, + temperature=temperature_value, + use_fused_kernels=use_fused_kernels, + ) + if mtp_result is not None: + hidden_for_head, mtp_loss = mtp_result + output["mtp_loss"] = mtp_loss + labels_sb = labels.transpose(0, 1).contiguous() + if use_fused_kernels: + hidden_full = gather_from_sequence_parallel(hidden_for_head, self.ps) + log_probs, entropy = linear_cross_entropy( + hidden_full, + self._head_weight_for_fused_ce(hidden_full), + labels_sb, + temperature_value, + self.ps.tp_group, + ) + token_loss = -log_probs + output["loss"] = token_loss.mean() + if return_log_probs: + output["log_probs"] = log_probs.transpose(0, 1).contiguous() + if calculate_entropy: + output["entropy"] = entropy.transpose(0, 1).contiguous() + else: + logits = self.head(hidden_for_head) + if temperature_value != 1.0: + logits = logits / temperature_value + token_loss = vocab_parallel_cross_entropy(logits, labels_sb, self.ps.tp_group) + output["loss"] = token_loss.mean() + if return_log_probs: + output["log_probs"] = (-token_loss).transpose(0, 1).contiguous() + if calculate_entropy: + entropy = vocab_parallel_entropy(logits, self.ps.tp_group) + output["entropy"] = entropy.transpose(0, 1).contiguous() + + if labels is None: + logits = self.head(hidden_for_head) + output["logits"] = self.head.gather(logits) + + return output + + def _apply_mtp_loss( + self, + hidden_states: torch.Tensor, + *, + input_ids: torch.Tensor | None, + position_ids: torch.Tensor | None, + labels: torch.Tensor, + loss_mask: torch.Tensor | None, + packed_seq_params, + temperature: float, + use_fused_kernels: bool, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + if self.mtp is None: + return None + if not self.mtp_enable_train: + return None + if input_ids is None: + raise ValueError("MTP training requires input_ids.") + if loss_mask is None: + loss_mask = torch.ones_like(labels, dtype=torch.float32) + else: + loss_mask = loss_mask.to(dtype=torch.float32) + + mtp_hidden_states = self.mtp( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + packed_seq_params=packed_seq_params, + ) + + mtp_labels = labels.clone() + mtp_loss_mask = loss_mask.clone() + mtp_loss_values = [] + for mtp_hidden in mtp_hidden_states: + mtp_labels, _ = roll_packed_thd_left( + mtp_labels, packed_seq_params=packed_seq_params, dims=-1 + ) + mtp_loss_mask, num_tokens = roll_packed_thd_left( + mtp_loss_mask, packed_seq_params=packed_seq_params, dims=-1 + ) + labels_sb = mtp_labels.transpose(0, 1).contiguous() + mask_sb = mtp_loss_mask.transpose(0, 1).contiguous() + + if use_fused_kernels: + mtp_hidden_full = gather_from_sequence_parallel(mtp_hidden, self.ps) + log_probs, _entropy = linear_cross_entropy( + mtp_hidden_full, + self._head_weight_for_fused_ce(mtp_hidden_full), + labels_sb, + temperature, + self.ps.tp_group, + ) + token_loss = -log_probs + else: + logits = self.head(mtp_hidden) + if temperature != 1.0: + logits = logits / temperature + token_loss = vocab_parallel_cross_entropy(logits, labels_sb, self.ps.tp_group) + token_loss = token_loss * mask_sb.to(dtype=token_loss.dtype) + num_tokens = num_tokens.to(dtype=token_loss.dtype).clamp_min(1.0) + mtp_loss_values.append(token_loss.sum() / num_tokens) + + mtp_loss_scale = self.mtp_loss_scaling_factor / max(len(mtp_hidden_states), 1) + hidden_states = MTPLossAutoScaler.apply( + hidden_states, mtp_loss_scale * token_loss / num_tokens + ) + + if not mtp_loss_values: + return None + return ( + hidden_states, + torch.stack([loss.detach().float() for loss in mtp_loss_values]).mean(), + ) + + def _head_weight_for_fused_ce(self, hidden_states: torch.Tensor) -> torch.Tensor: + assert self.head is not None + weight = self.head.col.linear.weight + if weight.dtype == hidden_states.dtype: + return weight + return weight.to(dtype=hidden_states.dtype) + + +__all__ = [ + "MoELayer", + "MTPLossAutoScaler", + "MultiTokenPredictionBlock", + "MultiTokenPredictionLayer", + "Qwen3MoEModel", + "TransformerLayer", +] diff --git a/experimental/lite/megatron/lite/model/qwen3_moe/lite/protocol.py b/experimental/lite/megatron/lite/model/qwen3_moe/lite/protocol.py new file mode 100644 index 00000000000..e8de45d63d6 --- /dev/null +++ b/experimental/lite/megatron/lite/model/qwen3_moe/lite/protocol.py @@ -0,0 +1,322 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Qwen3MoE lite impl — model protocol for Megatron Lite runtime. + +This file is the reference implementation of the Megatron Lite model protocol. +New model authors: copy this file and adapt. + +Protocol convention (what runtime calls): + Required: + ImplConfig — @dataclass, per-impl knobs + build_model_config(source, **overrides) → ModelConfig + build_model(model_cfg, *, impl_cfg) → ModelBundle + Optional (in ModelBundle.extras or module-level): + load_hf_weights(chunk, hf_path, model_cfg, ps) — HF weight loading + export_hf_weights(chunks, model_cfg, ps) — HF weight export + vocab_size(model_cfg) -> int — benchmark metadata + Escape hatch: + create_runtime(hf_path, cfg) -> Runtime — fully override runtime +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn + +from megatron.lite.model.qwen3_moe.common import is_expert_param +from megatron.lite.model.qwen3_moe.config import Qwen3MoEConfig +from megatron.lite.model.qwen3_moe.lite.checkpoint import EXPERT_CLASSIFIER, PLACEMENT_FN +from megatron.lite.model.qwen3_moe.lite.checkpoint import load_hf_weights as _load_hf_weights_impl +from megatron.lite.model.qwen3_moe.lite.model import MTPLossAutoScaler, Qwen3MoEModel +from megatron.lite.primitive.bundle import ModelBundle +from megatron.lite.primitive.modules.lora import ( + LoraConfig, + freeze_non_lora_params, + normalize_lora_config, + trainable_param_stats, +) +from megatron.lite.primitive.parallel import ParallelState, init_parallel +from megatron.lite.primitive.recompute import apply_recompute, parse_recompute_spec +from megatron.lite.runtime.contracts import OptimizerConfig, ParallelConfig + +__all__ = [ + "EXPERT_CLASSIFIER", + "ImplConfig", + "PLACEMENT_FN", + "build_model", + "build_model_config", + "export_hf_weights", + "load_hf_weights", + "vocab_size", +] + +# --------------------------------------------------------------------------- +# ImplConfig +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class ImplConfig: + """Lite impl knobs. Constructed by runtime from user config.""" + + parallel: ParallelConfig = field(default_factory=ParallelConfig) + optimizer: str | None = "mc" # None = no optimizer (inference) + recompute: list[str] = field(default_factory=list) + offload: list[str] = field(default_factory=list) + use_deepep: bool = False + use_thd: bool = False + router_aux_loss_coef: float | None = None + router_bias_rate: float = 0.0 + # User-level OptimizerConfig threaded through the runtime. + optimizer_config: OptimizerConfig | None = None + mtp_enable: bool = False + mtp_enable_train: bool = False + mtp_detach_encoder: bool = False + mtp_loss_scaling_factor: float = 0.1 + mtp_use_repeated_layer: bool | None = None + deterministic: bool = True + lora: LoraConfig | dict | None = None + + +# --------------------------------------------------------------------------- +# Module map for recompute/offload +# --------------------------------------------------------------------------- + +MODULE_MAP = { + "core_attn": lambda layer: layer.attn.core_attn, + "experts": lambda layer: layer.moe.experts, + "moe": lambda layer: layer.moe, + "router": lambda layer: layer.moe.router, + "mlp_norm": lambda layer: layer.mlp_norm, + "attn_proj": lambda layer: layer.attn.proj, +} + + +# --------------------------------------------------------------------------- +# Required: build_model_config +# --------------------------------------------------------------------------- + + +def build_model_config(source: str | Path | dict, **overrides) -> Qwen3MoEConfig: + """Build Qwen3MoE architecture config from HF source.""" + if isinstance(source, dict): + cfg = Qwen3MoEConfig._from_hf_dict(source) + else: + cfg = Qwen3MoEConfig.from_hf(str(source)) + for k, v in overrides.items(): + if hasattr(cfg, k): + setattr(cfg, k, v) + return cfg + + +# --------------------------------------------------------------------------- +# Required: build_model +# --------------------------------------------------------------------------- + + +def _forward_step(model: nn.Module, batch: dict) -> dict: + kwargs = {"input_ids": batch["input_ids"], "labels": batch["labels"]} + if "packed_seq_params" in batch: + kwargs["packed_seq_params"] = batch["packed_seq_params"] + if "position_ids" in batch: + kwargs["position_ids"] = batch["position_ids"] + for key in ( + "loss_mask", + "temperature", + "use_fused_kernels", + "calculate_entropy", + "return_log_probs", + ): + if key in batch: + kwargs[key] = batch[key] + if kwargs["input_ids"].dim() == 1: + kwargs["input_ids"] = kwargs["input_ids"].unsqueeze(0) + return model(**kwargs) + + +def build_model(model_cfg: Qwen3MoEConfig, *, impl_cfg: ImplConfig) -> ModelBundle: + """Build lite Qwen3MoE: model, parallel state, optimizer — everything. + + Model owns all construction. Runtime just consumes the ModelBundle. + """ + p = impl_cfg.parallel + lora_config = normalize_lora_config(impl_cfg.lora) + + # ── validation ── + if impl_cfg.use_deepep and (p.etp is not None and p.etp > 1): + raise ValueError("use_deepep and etp>1 are mutually exclusive") + + # ── override model config from impl_cfg ── + if impl_cfg.router_aux_loss_coef is not None: + model_cfg.router_aux_loss_coef = impl_cfg.router_aux_loss_coef + mtp_enable = bool(impl_cfg.mtp_enable) + mtp_enable_train = mtp_enable and bool(impl_cfg.mtp_enable_train) + if mtp_enable: + if model_cfg.num_nextn_predict_layers <= 0: + raise ValueError("mtp_enable=True but HF config has no num_nextn_predict_layers.") + model_cfg.mtp_loss_scaling_factor = impl_cfg.mtp_loss_scaling_factor + if impl_cfg.mtp_use_repeated_layer is not None: + model_cfg.mtp_use_repeated_layer = impl_cfg.mtp_use_repeated_layer + else: + model_cfg.num_nextn_predict_layers = 0 + + # ── parallel state (model creates its own) ── + ps = init_parallel(p) + deterministic = impl_cfg.deterministic + + # ── build chunks ── + recompute_spec = parse_recompute_spec(impl_cfg.recompute) + model_kwargs: dict[str, Any] = dict( + use_deepep=impl_cfg.use_deepep, + fp8=False, + recompute_modules=recompute_spec, + router_bias_rate=impl_cfg.router_bias_rate, + use_thd=impl_cfg.use_thd, + mtp_enable=mtp_enable, + mtp_enable_train=mtp_enable_train, + mtp_detach_encoder=impl_cfg.mtp_detach_encoder, + lora_config=lora_config, + ) + + vpp = None if p.vpp == 1 else p.vpp + if vpp is None: + chunks = [Qwen3MoEModel(model_cfg, ps, **model_kwargs).to(torch.bfloat16).cuda()] + else: + chunks = [] + for i in range(vpp): + chunks.append( + Qwen3MoEModel(model_cfg, ps, vpp=vpp, vpp_chunk_id=i, **model_kwargs) + .to(torch.bfloat16) + .cuda() + ) + + # ── recompute ── + if recompute_spec: + for chunk in chunks: + apply_recompute(chunk.layers, recompute_spec, MODULE_MAP) + + # ── offload ── + if impl_cfg.offload: + from megatron.lite.primitive.recompute import apply_offload + + for chunk in chunks: + apply_offload(chunk.layers, impl_cfg.offload, MODULE_MAP) + + lora_stats = None + if lora_config.enabled: + lora_stats = {"chunks": []} + for chunk in chunks: + freeze_stats = freeze_non_lora_params(chunk) + trainable_stats = trainable_param_stats(chunk) + lora_stats["chunks"].append({**freeze_stats, **trainable_stats}) + + # ── optimizer (model chooses which primitive) ── + optimizer = None + finalize_grads = None + post_model_load_hook = None + if impl_cfg.optimizer == "mc": + from megatron.lite.primitive.optimizers.megatron_wrap import build_mc_training_optimizer + + optimizer, finalize_grads = build_mc_training_optimizer( + chunks, + model_cfg=model_cfg, + impl_cfg=impl_cfg, + ps=ps, + model_name="qwen3_moe", + is_expert=is_expert_param, + deterministic=deterministic, + ) + from megatron.lite.primitive.ckpt import attach_model_sharded_state_dict + + attach_model_sharded_state_dict( + chunks, ps, get_placements=PLACEMENT_FN, is_expert=is_expert_param + ) + optimizer_backend = "distopt" + elif impl_cfg.optimizer == "fsdp2": + optimizer_backend = "fsdp2" + + def _post_model_load_hook(): + from megatron.lite.model.qwen3_moe.lite.model import TransformerLayer + from megatron.lite.primitive.optimizers.fsdp2 import build_fsdp2_training_optimizer + + return { + "optimizer": build_fsdp2_training_optimizer( + chunks, + impl_cfg.optimizer_config, + ps, + unit_modules=(TransformerLayer,), + expert_classifier=is_expert_param, + deterministic=deterministic, + vpp=impl_cfg.parallel.vpp, + # Non-layer params stay under the root FSDP2 unit. The fused + # CE path reads head.col.linear.weight directly, and the + # embedding path is also driven from model.forward(). + leaf_module_names=(), + ) + } + + post_model_load_hook = _post_model_load_hook + elif impl_cfg.optimizer is None: + optimizer_backend = "none" + else: + raise ValueError(f"Unknown qwen3_moe lite optimizer: {impl_cfg.optimizer!r}.") + + from megatron.lite.primitive.modules.moe import MoEAuxLossAutoScaler + + def _pre_forward_hook(loss_scale): + MoEAuxLossAutoScaler.set_loss_scale(loss_scale) + MTPLossAutoScaler.set_loss_scale(loss_scale) + + return ModelBundle( + chunks=chunks, + parallel_state=ps, + optimizer=optimizer, + finalize_grads=finalize_grads, + forward_step=_forward_step, + extras={ + "model_cfg": model_cfg, + # Lite's router uses megatron.lite's MoEAuxLossAutoScaler; hand the + # classmethod directly as the per-microbatch hook. + "pre_forward_hook": _pre_forward_hook, + "optimizer_backend": optimizer_backend, + "post_model_load_hook": post_model_load_hook, + "lora_config": lora_config, + "lora_stats": lora_stats, + }, + ) + + +# --------------------------------------------------------------------------- +# Optional: load_hf_weights +# --------------------------------------------------------------------------- + + +def load_hf_weights( + chunk: nn.Module, hf_path: str, model_cfg: Qwen3MoEConfig, ps: ParallelState +) -> None: + """Load HF pretrained weights into model chunk.""" + if not hf_path: + return + _load_hf_weights_impl(chunk, hf_path, model_cfg, ps) + + +def export_hf_weights( + chunks: list[nn.Module], model_cfg: Qwen3MoEConfig, ps: ParallelState, **kwargs +): + """Export HF weights from model chunks.""" + from megatron.lite.model.qwen3_moe.lite.checkpoint import export_hf_weights as _export + + for chunk in chunks: + yield from _export(chunk, model_cfg, ps, **kwargs) + + +# --------------------------------------------------------------------------- +# Tooling metadata (benchmark / debug) +# --------------------------------------------------------------------------- + + +def vocab_size(model_cfg: Qwen3MoEConfig) -> int | None: + return getattr(model_cfg, "vocab_size", None) diff --git a/experimental/lite/megatron/lite/model/qwen3_moe/stats.py b/experimental/lite/megatron/lite/model/qwen3_moe/stats.py new file mode 100644 index 00000000000..77c3fd6215d --- /dev/null +++ b/experimental/lite/megatron/lite/model/qwen3_moe/stats.py @@ -0,0 +1,26 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Model-level Qwen3MoE benchmark statistics.""" + +from __future__ import annotations + +from megatron.lite.model.qwen3_moe.config import Qwen3MoEConfig + + +def activated_params(model_cfg: Qwen3MoEConfig) -> int | None: + try: + h = model_cfg.hidden_size + n_layers = model_cfg.num_hidden_layers + n_q = model_cfg.num_attention_heads + n_kv = model_cfg.num_key_value_heads + d = model_cfg.head_dim + attn = h * (n_q + n_kv + n_kv) * d + (n_q * d) * h + router = h * model_cfg.num_experts + inter = model_cfg.moe_intermediate_size + expert = h * (inter * 2) + inter * h + experts_active = model_cfg.num_experts_per_tok * expert + return int((attn + router + experts_active) * n_layers) + except (AttributeError, TypeError, ValueError): + return None + + +__all__ = ["activated_params"] diff --git a/experimental/lite/megatron/lite/model/registry.py b/experimental/lite/megatron/lite/model/registry.py new file mode 100644 index 00000000000..2de6f05357a --- /dev/null +++ b/experimental/lite/megatron/lite/model/registry.py @@ -0,0 +1,165 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Model and protocol registry.""" + +from __future__ import annotations + +import importlib +import json +from pathlib import Path + +# --------------------------------------------------------------------------- +# Registry tables (populated by register_model) +# --------------------------------------------------------------------------- + +# model_name → model package module path +MODEL_PACKAGES: dict[str, str] = {} + +# HF model_type string → Megatron Lite model_name +_HF_MODEL_TYPE_MAP: dict[str, str] = {} + +# (model_name, impl) → runtime_model_name +_IMPL_TO_RUNTIME_MODEL: dict[tuple[str, str], str] = {} + +# runtime_model_name → protocol module path +TRAIN_RUNTIME_MODULES: dict[str, str] = {} + +# --------------------------------------------------------------------------- +# Registration API +# --------------------------------------------------------------------------- + + +def register_model( + model_name: str, + *, + package: str, + hf_model_types: list[str] | None = None, + impls: dict[str, str] | None = None, +) -> None: + """Register a model and all its implementations in one call. + + Args: + model_name: Megatron Lite model name (e.g. ``"qwen3"``). + package: Model package module path (e.g. ``"megatron.lite.model.qwen3_moe"``). + hf_model_types: HF ``model_type`` strings that map to this model. + impls: ``{impl_name: protocol_module_path}``. + The first impl is also registered as the bare ``model_name`` runtime key. + + Example:: + + register_model( + "qwen3", + package="megatron.lite.model.qwen3_moe", + hf_model_types=["qwen3_moe", "qwen2_moe"], + impls={ + "lite": "megatron.lite.model.qwen3_moe.lite.protocol", + }, + ) + """ + MODEL_PACKAGES[model_name] = package + + if hf_model_types: + for hf_type in hf_model_types: + _HF_MODEL_TYPE_MAP[hf_type] = model_name + + if impls: + for i, (impl_name, proto_module) in enumerate(impls.items()): + runtime_key = model_name if i == 0 else f"{model_name}_{impl_name}" + _IMPL_TO_RUNTIME_MODEL[(model_name, impl_name)] = runtime_key + TRAIN_RUNTIME_MODULES[runtime_key] = proto_module + + +# --------------------------------------------------------------------------- +# Built-in models +# --------------------------------------------------------------------------- + +_QWEN3_MOE_LITE = "megatron.lite.model.qwen3_moe.lite.protocol" + +register_model( + "qwen3", + package="megatron.lite.model.qwen3_moe", + hf_model_types=["qwen3_moe", "qwen2_moe"], + impls={"lite": _QWEN3_MOE_LITE}, +) + +register_model( + "qwen3_moe", package="megatron.lite.model.qwen3_moe", impls={"lite": _QWEN3_MOE_LITE} +) + +register_model( + "qwen3_5", + package="megatron.lite.model.qwen3_5", + hf_model_types=["qwen3_5_moe"], + impls={"lite": "megatron.lite.model.qwen3_5.lite.protocol"}, +) + + +# --------------------------------------------------------------------------- +# Lookup functions +# --------------------------------------------------------------------------- + + +def get_model_package(model_name: str): + if model_name not in MODEL_PACKAGES: + raise ValueError(f"Unknown model: {model_name!r}. Available: {list(MODEL_PACKAGES)}") + return importlib.import_module(MODEL_PACKAGES[model_name]) + + +def get_train_runtime_module(model_name: str): + if model_name in TRAIN_RUNTIME_MODULES: + return importlib.import_module(TRAIN_RUNTIME_MODULES[model_name]) + raise ValueError(f"No protocol module for: {model_name!r}") + + +def resolve_runtime_model_name(model_name: str, impl: str) -> str: + key = (model_name, impl) + if key not in _IMPL_TO_RUNTIME_MODEL: + raise ValueError( + f"No runtime for ({model_name!r}, {impl!r}). " f"Known: {list(_IMPL_TO_RUNTIME_MODEL)}" + ) + return _IMPL_TO_RUNTIME_MODEL[key] + + +def resolve_model_type_from_hf(source: str | Path | dict) -> str: + """Resolve Megatron Lite model_name from an HF source. + + Args: + source: One of: + - Directory path (str/Path) containing ``config.json`` + - Path to a ``config.json`` file directly + - A dict / HF config object with a ``model_type`` key + """ + if isinstance(source, dict): + hf_config = source + elif hasattr(source, "model_type"): + # HF PretrainedConfig object + hf_config = {"model_type": source.model_type} + else: + p = Path(source) + if p.is_file(): + config_path = p + elif p.is_dir(): + config_path = p / "config.json" + else: + raise FileNotFoundError(f"Not a file or directory: {source}") + if not config_path.exists(): + raise FileNotFoundError(f"No config.json found at {source}") + with open(config_path) as f: + hf_config = json.load(f) + + hf_model_type = hf_config.get("model_type", "") + native_name = _HF_MODEL_TYPE_MAP.get(hf_model_type) + if native_name is not None: + return native_name + raise ValueError( + f"Cannot resolve model_type={hf_model_type!r}. " + f"Known: {list(_HF_MODEL_TYPE_MAP)}. Set model_name explicitly." + ) + + +__all__ = [ + "get_model_package", + "get_train_runtime_module", + "register_model", + "resolve_model_type_from_hf", + "resolve_runtime_model_name", +] diff --git a/experimental/lite/megatron/lite/primitive/__init__.py b/experimental/lite/megatron/lite/primitive/__init__.py new file mode 100644 index 00000000000..17c992fc88d --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Public primitive-layer entrypoints.""" diff --git a/experimental/lite/megatron/lite/primitive/bundle.py b/experimental/lite/megatron/lite/primitive/bundle.py new file mode 100644 index 00000000000..e74603d060e --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/bundle.py @@ -0,0 +1,29 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""ModelBundle — return type of protocol.build_model().""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +import torch.nn as nn + +from megatron.lite.primitive.parallel.state import ParallelState + + +@dataclass +class ModelBundle: + """Everything runtime needs to run a training loop. + + Returned by protocol.build_model(). Model owns the construction + of all fields — runtime just consumes them. + """ + + chunks: list[nn.Module] + parallel_state: ParallelState + optimizer: Any | None = None + finalize_grads: Callable[[], None] | None = None + forward_step: Callable[[nn.Module, dict], dict] | None = None + # extra metadata (expert_classifier, model_cfg, etc.) + extras: dict[str, Any] = field(default_factory=dict) diff --git a/experimental/lite/megatron/lite/primitive/ckpt/__init__.py b/experimental/lite/megatron/lite/primitive/ckpt/__init__.py new file mode 100644 index 00000000000..14ff5b4d884 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/ckpt/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""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__ = [ + "HFWeights", + "attach_model_sharded_state_dict", + "load_training_checkpoint", + "save_training_checkpoint", +] diff --git a/experimental/lite/megatron/lite/primitive/ckpt/dcp.py b/experimental/lite/megatron/lite/primitive/ckpt/dcp.py new file mode 100644 index 00000000000..cb5fb1ebc2c --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/ckpt/dcp.py @@ -0,0 +1,623 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +""" +DCP (Distributed Checkpoint) framework for training checkpoints. + +Model-agnostic: takes a placement function to describe how each parameter is sharded. +HF weight loading/saving is model-specific and lives in models//checkpoint.py. +""" + +from __future__ import annotations + +import os +import random +from collections.abc import Iterable +from pathlib import Path +from typing import Any + +import numpy as np +import torch # pyright: ignore[reportMissingImports] +import torch.distributed as dist # pyright: ignore[reportMissingImports] +import torch.distributed.checkpoint as dcp # pyright: ignore[reportMissingImports] +import torch.nn as nn # pyright: ignore[reportMissingImports] +from torch.distributed.device_mesh import DeviceMesh # pyright: ignore[reportMissingImports] +from torch.distributed.tensor import DTensor # pyright: ignore[reportMissingImports] + +from megatron.lite.primitive.parallel import ParallelState +from megatron.lite.primitive.protocols import ( + ExpertClassifierFn, + PlacementFn, + default_expert_classifier, + default_placement_fn, +) + + +def save_training_checkpoint( + model: nn.Module | Iterable[nn.Module], + optimizer, + step: int | str, + path: str | None = None, + config=None, + ps: ParallelState | None = None, + get_placements: PlacementFn = default_placement_fn, + is_expert: ExpertClassifierFn = default_expert_classifier, + *, + use_dcp: bool | None = True, + save_rng: bool = True, + save_model: bool = True, + save_optimizer: bool = True, +) -> None: + """Save training checkpoint using DTensor + DCP for automatic resharding.""" + if path is None and isinstance(step, str): + path = step + step = 0 + if path is None: + raise ValueError("checkpoint path is required") + step = int(step) + if use_dcp is None: + use_dcp = True + if not use_dcp: + _save_local_training_checkpoint(model, optimizer, step, path, save_rng=save_rng) + return + if _supports_distopt_distckpt(model, optimizer): + ckpt_path = os.path.join(path, f"step_{step}") + os.makedirs(ckpt_path, exist_ok=True) + _save_distopt_checkpoint( + model, optimizer, step, ckpt_path, save_model=save_model, save_optimizer=save_optimizer + ) + if save_rng: + _save_rng_sidecar(ckpt_path) + log_rank0(f"Saved distopt checkpoint at step {step} to {ckpt_path}") + return + if config is None or ps is None: + raise ValueError("DCP checkpointing requires config and ParallelState.") + if not isinstance(model, nn.Module): + raise TypeError("DCP checkpointing currently expects a single nn.Module.") + dense_mesh, expert_mesh = _build_meshes(config) + state_dict: dict = {"step": step} + + if save_model: + for name, param in model.named_parameters(): + placements = get_placements(name) + mesh = expert_mesh if is_expert(name) else dense_mesh + state_dict[f"model.{name}"] = _dcp_tensor_from_param(param, mesh, placements) + + ckpt_path = os.path.join(path, f"step_{step}") + os.makedirs(ckpt_path, exist_ok=True) + dcp.save(state_dict, checkpoint_id=ckpt_path) + if save_optimizer: + _save_optimizer_checkpoint(optimizer, ckpt_path) + if save_rng: + _save_rng_sidecar(ckpt_path) + log_rank0(f"Saved training checkpoint at step {step} to {ckpt_path}") + + +def load_training_checkpoint( + model: nn.Module | Iterable[nn.Module], + optimizer, + path: str, + config=None, + ps: ParallelState | None = None, + get_placements: PlacementFn = default_placement_fn, + is_expert: ExpertClassifierFn = default_expert_classifier, + *, + use_dcp: bool | None = True, + load_rng: bool = True, + load_parameter_state_update_legacy_format: bool = False, + load_model: bool = True, + load_optimizer: bool = True, +) -> int: + """Load training checkpoint with automatic resharding across different parallel configs.""" + if use_dcp is None: + use_dcp = True + if not use_dcp: + return _load_local_training_checkpoint( + model, + optimizer, + path, + load_rng=load_rng, + load_parameter_state_update_legacy_format=load_parameter_state_update_legacy_format, + ) + ckpt_path = _resolve_step_checkpoint_path(path) + if _supports_distopt_distckpt(model, optimizer): + step = _load_distopt_checkpoint( + model, optimizer, ckpt_path, load_model=load_model, load_optimizer=load_optimizer + ) + if load_rng: + _load_rng_sidecar(ckpt_path) + log_rank0(f"Loaded distopt checkpoint from {path} at step {step}") + return step + if config is None or ps is None: + raise ValueError("DCP checkpointing requires config and ParallelState.") + if not isinstance(model, nn.Module): + raise TypeError("DCP checkpointing currently expects a single nn.Module.") + dense_mesh, expert_mesh = _build_meshes(config) + + state_dict: dict = {"step": 0} + + if load_model: + for name, param in model.named_parameters(): + placements = get_placements(name) + mesh = expert_mesh if is_expert(name) else dense_mesh + state_dict[f"model.{name}"] = _empty_dcp_tensor_like_param(param, mesh, placements) + + dcp.load(state_dict, checkpoint_id=ckpt_path) + + if load_model: + for name, param in model.named_parameters(): + key = f"model.{name}" + if key in state_dict: + t = state_dict[key] + with torch.no_grad(): + _copy_tensor_(param, t) + + if load_optimizer: + _load_optimizer_checkpoint(optimizer, ckpt_path) + + step = state_dict.get("step", 0) + if load_rng: + _load_rng_sidecar(ckpt_path) + log_rank0(f"Loaded training checkpoint from {path} at step {step}") + return step + + +def _resolve_step_checkpoint_path(path: str) -> str: + if os.path.basename(path).startswith("step_"): + return path + + step_dirs = sorted( + [d for d in os.listdir(path) if d.startswith("step_")], key=lambda d: int(d.split("_")[1]) + ) + if step_dirs: + return os.path.join(path, step_dirs[-1]) + return path + + +def _supports_distopt_distckpt(model: nn.Module | Iterable[nn.Module], optimizer) -> bool: + from megatron.lite.primitive.ckpt.distckpt import supports_distopt_distckpt + + return supports_distopt_distckpt(model, optimizer) + + +def _save_distopt_checkpoint( + model: nn.Module | Iterable[nn.Module], + optimizer, + step: int, + path: str, + *, + save_model: bool, + save_optimizer: bool, +) -> None: + from megatron.lite.primitive.ckpt.distckpt import save_distopt_checkpoint + + save_distopt_checkpoint( + model, optimizer, step, path, save_model=save_model, save_optimizer=save_optimizer + ) + + +def _load_distopt_checkpoint( + model: nn.Module | Iterable[nn.Module], + optimizer, + path: str, + *, + load_model: bool, + load_optimizer: bool, +) -> int: + from megatron.lite.primitive.ckpt.distckpt import load_distopt_checkpoint + + return load_distopt_checkpoint( + model, optimizer, path, load_model=load_model, load_optimizer=load_optimizer + ) + + +def _optimizer_checkpoint_path(path: str) -> str: + rank = dist.get_rank() if dist.is_initialized() else 0 + return os.path.join(path, f"optimizer_rank_{rank}.pt") + + +def _save_optimizer_checkpoint(optimizer, path: str) -> None: + if optimizer is None: + log_rank0("Skipping optimizer checkpoint save because optimizer is None") + return + state_dict_fn = getattr(optimizer, "state_dict", None) + if not callable(state_dict_fn): + raise TypeError(f"Optimizer {type(optimizer).__name__} does not provide state_dict().") + torch.save(state_dict_fn(), _optimizer_checkpoint_path(path)) + + +def _load_optimizer_checkpoint(optimizer, path: str) -> None: + if optimizer is None: + log_rank0("Skipping optimizer checkpoint load because optimizer is None") + return + ckpt_path = _optimizer_checkpoint_path(path) + if not os.path.exists(ckpt_path): + log_rank0(f"No optimizer checkpoint found at {ckpt_path}; loading model state only") + return + load_state_dict_fn = getattr(optimizer, "load_state_dict", None) + if not callable(load_state_dict_fn): + raise TypeError(f"Optimizer {type(optimizer).__name__} does not provide load_state_dict().") + state = torch.load(ckpt_path, map_location="cpu", weights_only=False) + load_state_dict_fn(state) + + +def _model_chunks(model: nn.Module | Iterable[nn.Module]) -> list[nn.Module]: + if isinstance(model, nn.Module): + return [model] + chunks = list(model) + if not all(isinstance(chunk, nn.Module) for chunk in chunks): + raise TypeError("checkpoint model chunks must be nn.Module instances.") + return chunks + + +def _to_local_tensor(tensor: Any) -> torch.Tensor: + local_tensor = getattr(tensor, "_local_tensor", None) + if isinstance(local_tensor, torch.Tensor): + return local_tensor + to_local = getattr(tensor, "to_local", None) + if callable(to_local): + return to_local() + return tensor + + +def _is_dtensor_like(tensor: Any) -> bool: + return ( + callable(getattr(tensor, "to_local", None)) + and hasattr(tensor, "device_mesh") + and hasattr(tensor, "placements") + ) + + +def _dcp_tensor_from_param(param: torch.Tensor, mesh: DeviceMesh, placements: list) -> DTensor: + if _is_dtensor_like(param): + return _dtensor_from_dtensor_like_param(param, _to_local_tensor(param).detach()) + return DTensor.from_local(_to_local_tensor(param).detach(), mesh, placements) + + +def _empty_dcp_tensor_like_param( + param: torch.Tensor, mesh: DeviceMesh, placements: list +) -> DTensor: + if _is_dtensor_like(param): + return _dtensor_from_dtensor_like_param(param, torch.empty_like(_to_local_tensor(param))) + return DTensor.from_local(torch.empty_like(_to_local_tensor(param)), mesh, placements) + + +def _dtensor_from_dtensor_like_param(param: torch.Tensor, local_tensor: torch.Tensor) -> DTensor: + return DTensor.from_local( + local_tensor, + param.device_mesh, + param.placements, + shape=tuple(param.shape), + stride=tuple(param.stride()), + ) + + +def _copy_tensor_(target: torch.Tensor, src: torch.Tensor) -> None: + local_target = _to_local_tensor(target) + local_src = _to_local_tensor(src).to(device=local_target.device, dtype=local_target.dtype) + if isinstance(local_target, torch.Tensor) and local_target is not target: + local_target.copy_(local_src) + else: + target.copy_(local_src) + + +def _chunk_tensor_state(module: nn.Module) -> dict[str, torch.Tensor]: + state: dict[str, torch.Tensor] = {} + for name, param in module.named_parameters(): + state[f"param.{name}"] = _to_local_tensor(param.detach()).cpu().clone() + for name, buffer in module.named_buffers(): + state[f"buffer.{name}"] = _to_local_tensor(buffer.detach()).cpu().clone() + return state + + +def _load_chunk_tensor_state(module: nn.Module, state: dict[str, torch.Tensor]) -> None: + params = dict(module.named_parameters()) + buffers = dict(module.named_buffers()) + missing: list[str] = [] + for key, src in state.items(): + kind, name = key.split(".", 1) + if kind == "param" and name in params: + with torch.no_grad(): + _copy_tensor_(params[name], src) + elif kind == "buffer" and name in buffers: + with torch.no_grad(): + _copy_tensor_(buffers[name], src) + else: + missing.append(key) + if missing: + raise RuntimeError(f"checkpoint contains unknown tensor keys: {missing}") + + +def _local_checkpoint_file(path: str | os.PathLike[str]) -> Path: + ckpt_path = Path(path) + if ckpt_path.is_dir() or ckpt_path.suffix == "": + if _is_distributed_checkpoint_ranked(): + return ckpt_path / f"training_state_{_rank_suffix()}.pt" + return ckpt_path / "training_state.pt" + return ckpt_path + + +def _local_optimizer_parameter_state_file(ckpt_file: Path) -> Path: + return ckpt_file.with_name(f"{ckpt_file.stem}.optimizer_parameter_state{ckpt_file.suffix}") + + +def _rank_suffix() -> str: + if dist.is_available() and dist.is_initialized(): + return f"rank_{dist.get_rank():05d}" + return "rank_00000" + + +def _is_distributed_checkpoint_ranked() -> bool: + return dist.is_available() and dist.is_initialized() + + +def _rng_sidecar_file(path: str | os.PathLike[str]) -> Path: + return Path(path) / f"rng_state_{_rank_suffix()}.pt" + + +def _cpu_clone(tensor: torch.Tensor | None) -> torch.Tensor | None: + if tensor is None: + return None + return tensor.detach().cpu().clone() + + +def _get_cuda_rng_state() -> torch.Tensor | None: + if not torch.cuda.is_initialized(): + return None + return _cpu_clone(torch.cuda.get_rng_state()) + + +def _get_cuda_rng_tracker_states() -> dict[str, torch.Tensor]: + if not torch.cuda.is_initialized(): + return {} + + from megatron.core import tensor_parallel + + states = tensor_parallel.get_cuda_rng_tracker().get_states() + return {name: _cpu_clone(state) for name, state in states.items() if state is not None} + + +def _get_rng_state() -> dict[str, Any]: + return { + "random_rng_state": random.getstate(), + "np_rng_state": np.random.get_state(), + "torch_rng_state": _cpu_clone(torch.get_rng_state()), + "cuda_rng_state": _get_cuda_rng_state(), + "rng_tracker_states": _get_cuda_rng_tracker_states(), + } + + +def _restore_cuda_rng_tracker_states(states: dict[str, torch.Tensor]) -> None: + if not states or not torch.cuda.is_initialized(): + return + try: + from megatron.core import tensor_parallel + + tracker = tensor_parallel.get_cuda_rng_tracker() + graph_safe = tensor_parallel.is_graph_safe_cuda_rng_tracker(tracker) + restored = { + name: tensor_parallel.convert_cuda_rng_state(state, to_graphable=graph_safe) + for name, state in states.items() + } + tracker.set_states(restored) + except Exception as exc: + raise RuntimeError("Failed to restore Megatron tensor-parallel RNG tracker state.") from exc + + +def _restore_rng_state(state: dict[str, Any] | None) -> None: + if not state: + return + random.setstate(state["random_rng_state"]) + np.random.set_state(state["np_rng_state"]) + torch.set_rng_state(state["torch_rng_state"]) + cuda_rng_state = state.get("cuda_rng_state") + if cuda_rng_state is not None and torch.cuda.is_initialized(): + torch.cuda.set_rng_state(cuda_rng_state) + _restore_cuda_rng_tracker_states(state.get("rng_tracker_states", {})) + + +def _save_rng_sidecar(path: str | os.PathLike[str]) -> None: + rng_file = _rng_sidecar_file(path) + rng_file.parent.mkdir(parents=True, exist_ok=True) + torch.save(_get_rng_state(), rng_file) + + +def _load_rng_sidecar(path: str | os.PathLike[str]) -> None: + rng_file = _rng_sidecar_file(path) + if not rng_file.exists(): + log_rank0(f"RNG sidecar not found at {rng_file}; skipping RNG restore.") + return + _restore_rng_state(torch.load(rng_file, map_location="cpu", weights_only=False)) + + +def _save_local_training_checkpoint( + model: nn.Module | Iterable[nn.Module], + optimizer, + step: int, + path: str, + *, + save_rng: bool = True, +) -> None: + chunks = _model_chunks(model) + ckpt_file = _local_checkpoint_file(path) + ckpt_file.parent.mkdir(parents=True, exist_ok=True) + save_parameter_state = getattr(optimizer, "save_parameter_state", None) + optimizer_parameter_state_file = ( + _local_optimizer_parameter_state_file(ckpt_file) if callable(save_parameter_state) else None + ) + state = { + "format": "megatron_lite.local_training.v1", + "step": int(step), + "model": [_chunk_tensor_state(chunk) for chunk in chunks], + "optimizer": optimizer.state_dict() if optimizer is not None else None, + "optimizer_parameter_state": ( + optimizer_parameter_state_file.name + if optimizer_parameter_state_file is not None + else None + ), + "rng_state": _get_rng_state() if save_rng else None, + } + torch.save(state, ckpt_file) + if optimizer_parameter_state_file is not None: + save_parameter_state(str(optimizer_parameter_state_file)) + log_rank0(f"Saved local training checkpoint at step {step} to {ckpt_file}") + + +def _load_local_training_checkpoint( + model: nn.Module | Iterable[nn.Module], + optimizer, + path: str, + *, + load_rng: bool = True, + load_parameter_state_update_legacy_format: bool = False, +) -> int: + ckpt_file = _local_checkpoint_file(path) + state = torch.load(ckpt_file, map_location="cpu", weights_only=False) + if state.get("format") != "megatron_lite.local_training.v1": + raise RuntimeError(f"Unsupported local checkpoint format in {ckpt_file}") + chunks = _model_chunks(model) + chunk_states = state.get("model") + if not isinstance(chunk_states, list) or len(chunk_states) != len(chunks): + raise RuntimeError("Checkpoint model chunk count does not match target model.") + for chunk, chunk_state in zip(chunks, chunk_states, strict=True): + _load_chunk_tensor_state(chunk, chunk_state) + if optimizer is not None and state.get("optimizer") is not None: + optimizer.load_state_dict(state["optimizer"]) + parameter_state_name = state.get("optimizer_parameter_state") + load_parameter_state = getattr(optimizer, "load_parameter_state", None) + if parameter_state_name is not None and callable(load_parameter_state): + load_parameter_state( + str(ckpt_file.with_name(parameter_state_name)), + update_legacy_format=load_parameter_state_update_legacy_format, + ) + else: + reload_model_params = getattr(optimizer, "reload_model_params", None) + if callable(reload_model_params): + reload_model_params() + if load_rng: + _restore_rng_state(state.get("rng_state")) + step = int(state.get("step", 0)) + log_rank0(f"Loaded local training checkpoint from {ckpt_file} at step {step}") + return step + + +def _build_meshes(config): + """Build separate meshes for dense and expert parameters. + + Dense mesh [PP, DP, CP, TP] — matches init_parallel dense decomposition. + Expert mesh [PP, EDP, EP, ETP] — matches init_parallel expert decomposition. + + Both meshes use C-order layout so the innermost (rightmost) dimension + corresponds to the fastest-changing rank index, consistent with + init_parallel's rank = (...) * inner_size + inner_rank formula. + """ + ws = dist.get_world_size() + tp = int(config.tp or 1) + ep = int(config.ep or 1) + etp = max(int(config.etp or 1), 1) + cp = max(int(config.cp or 1), 1) + pp = max(int(config.pp or 1), 1) + + dense_dp = ws // (tp * cp * pp) + expert_dp = ws // (etp * ep * pp) + + ranks = torch.arange(ws) + dense_mesh = DeviceMesh("cuda", ranks.reshape(pp, dense_dp, cp, tp)) + expert_mesh = DeviceMesh("cuda", ranks.reshape(pp, expert_dp, ep, etp)) + return dense_mesh, expert_mesh + + +def log_rank0(msg: str) -> None: + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: + print(f"[megatron.lite] {msg}", flush=True) + + +# ====================================================================== +# QKV / FC1 canonicalize for DCP (interleaved-TP ↔ canonical layout) +# ====================================================================== + + +def _ag(data, size, group, dim=0): + from megatron.lite.primitive.ckpt.hf_weights import allgather_concat + + return allgather_concat(data, size, group, dim) + + +def canonicalize_qkv_for_dcp(model, num_attention_heads, num_key_value_heads, head_dim, ps): + """Rearrange fused QKV from interleaved-TP to canonical (Q|K|V) for DCP save.""" + if ps.tp_size <= 1: + return + from megatron.lite.primitive.utils import ensure_divisible + + nq = ensure_divisible(num_attention_heads, ps.tp_size) * head_dim + nkv = ensure_divisible(num_key_value_heads, ps.tp_size) * head_dim + for name, param in model.named_parameters(): + if "qkv" not in name or "layer_norm" in name: + continue + full = _ag(param.data, ps.tp_size, ps.tp_group) + cs = param.data.shape[0] + q, k, v = [], [], [] + for r in range(ps.tp_size): + s = full[r * cs : (r + 1) * cs] + q.append(s[:nq]) + k.append(s[nq : nq + nkv]) + v.append(s[nq + nkv :]) + canon = torch.cat([torch.cat(q), torch.cat(k), torch.cat(v)], dim=0) + param.data.copy_(canon.chunk(ps.tp_size, dim=0)[ps.tp_rank]) + + +def decanon_qkv_after_dcp(model, num_attention_heads, num_key_value_heads, head_dim, ps): + """Reverse of canonicalize_qkv_for_dcp.""" + if ps.tp_size <= 1: + return + qs = num_attention_heads * head_dim + kvs = num_key_value_heads * head_dim + for name, param in model.named_parameters(): + if "qkv" not in name or "layer_norm" in name: + continue + full = _ag(param.data, ps.tp_size, ps.tp_group) + ql = full[:qs].chunk(ps.tp_size)[ps.tp_rank] + kl = full[qs : qs + kvs].chunk(ps.tp_size)[ps.tp_rank] + vl = full[qs + kvs :].chunk(ps.tp_size)[ps.tp_rank] + param.data.copy_(torch.cat([ql, kl, vl], dim=0)) + + +def canonicalize_fc1_for_dcp(model, ps): + """Rearrange fused gate-up FC1 from interleaved-ETP to canonical for DCP save.""" + if ps.etp_size <= 1: + return + for name, param in model.named_parameters(): + if "experts" not in name or "fc1" not in name: + continue + full = _ag(param.data, ps.etp_size, ps.etp_group) + cs = param.data.shape[0] + ffn = cs // 2 + g, u = [], [] + for r in range(ps.etp_size): + s = full[r * cs : (r + 1) * cs] + g.append(s[:ffn]) + u.append(s[ffn:]) + canon = torch.cat([torch.cat(g), torch.cat(u)], dim=0) + param.data.copy_(canon.chunk(ps.etp_size, dim=0)[ps.etp_rank]) + + +def decanon_fc1_after_dcp(model, ps): + """Reverse of canonicalize_fc1_for_dcp.""" + if ps.etp_size <= 1: + return + for name, param in model.named_parameters(): + if "experts" not in name or "fc1" not in name: + continue + full = _ag(param.data, ps.etp_size, ps.etp_group) + ffn = full.shape[0] // 2 + gl = full[:ffn].chunk(ps.etp_size)[ps.etp_rank] + ul = full[ffn:].chunk(ps.etp_size)[ps.etp_rank] + param.data.copy_(torch.cat([gl, ul], dim=0)) + + +__all__ = [ + "canonicalize_fc1_for_dcp", + "canonicalize_qkv_for_dcp", + "decanon_fc1_after_dcp", + "decanon_qkv_after_dcp", + "load_training_checkpoint", + "save_training_checkpoint", +] diff --git a/experimental/lite/megatron/lite/primitive/ckpt/distckpt.py b/experimental/lite/megatron/lite/primitive/ckpt/distckpt.py new file mode 100644 index 00000000000..e50769dba68 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/ckpt/distckpt.py @@ -0,0 +1,542 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Megatron Core distributed checkpoint bridge for MLite distopt.""" + +from __future__ import annotations + +import os +from collections.abc import Callable, Iterable, MutableMapping +from dataclasses import replace +from types import MethodType +from typing import Any + +import torch +import torch.nn as nn + +from megatron.core import dist_checkpointing +from megatron.core.dist_checkpointing.mapping import ShardedTensor +from megatron.lite.primitive.parallel import ParallelState +from megatron.lite.primitive.protocols import ( + ExpertClassifierFn, + PlacementFn, + default_expert_classifier, + default_placement_fn, +) + +_DISTOPT_METADATA = { + "distrib_optim_sharding_type": "fully_reshardable", + "distrib_optim_fully_reshardable_mem_efficient": False, + "chained_optim_avoid_prefix": True, +} + + +def attach_model_sharded_state_dict( + model_chunks: Iterable[nn.Module], + ps: ParallelState, + *, + get_placements: PlacementFn = default_placement_fn, + is_expert: ExpertClassifierFn = default_expert_classifier, +) -> None: + """Attach an MLite-local mcore sharded_state_dict method to distopt chunks.""" + + for chunk in model_chunks: + chunk.sharded_state_dict = MethodType( # type: ignore[method-assign] + _build_bound_sharded_state_dict(ps, get_placements, is_expert), chunk + ) + chunk._mlite_distopt_sharded_state_dict = True # type: ignore[attr-defined] + chunk._mlite_distopt_parallel_state = ps # type: ignore[attr-defined] + + +def supports_distopt_distckpt(model: nn.Module | Iterable[nn.Module], optimizer: Any) -> bool: + """Return whether this model/optimizer pair can use mcore dist_checkpointing.""" + + if optimizer is not None and not callable(getattr(optimizer, "sharded_state_dict", None)): + return False + return all( + bool(getattr(chunk, "_mlite_distopt_sharded_state_dict", False)) + and callable(getattr(chunk, "sharded_state_dict", None)) + for chunk in _model_chunks(model) + ) + + +def save_distopt_checkpoint( + model: nn.Module | Iterable[nn.Module], + optimizer: Any, + step: int, + checkpoint_dir: str, + *, + save_model: bool = True, + save_optimizer: bool = True, +) -> None: + """Save model and DistributedOptimizer state through mcore dist_checkpointing.""" + + os.makedirs(checkpoint_dir, exist_ok=True) + model_sd = _model_sharded_state_dict(model) if save_model or save_optimizer else {} + state_dict: dict[str, Any] = {"step": int(step)} + if save_model: + state_dict.update(model_sd) + if save_optimizer and optimizer is not None: + _synchronize_native_optimizer_steps(optimizer) + patches = _patch_empty_native_optimizer_state_dicts(optimizer, fallback_step=step) + try: + state_dict["optimizer"] = optimizer.sharded_state_dict( + _single_or_all_model_state(model_sd), metadata=_DISTOPT_METADATA + ) + finally: + _restore_state_dict_patches(patches) + dist_checkpointing.save( + state_dict, + checkpoint_dir, + validate_access_integrity=False, + content_metadata=_DISTOPT_METADATA, + ) + + +def load_distopt_checkpoint( + model: nn.Module | Iterable[nn.Module], + optimizer: Any, + checkpoint_dir: str, + *, + load_model: bool = True, + load_optimizer: bool = True, +) -> int: + """Load a mcore dist_checkpointing checkpoint into model and DistributedOptimizer.""" + + model_sd = _model_sharded_state_dict(model) if load_model or load_optimizer else {} + load_sd: dict[str, Any] = {"step": 0} + if load_model: + load_sd.update(model_sd) + if load_optimizer and optimizer is not None: + patches = _patch_empty_native_optimizer_state_dicts(optimizer, fallback_step=0) + try: + load_sd["optimizer"] = optimizer.sharded_state_dict( + _single_or_all_model_state(model_sd), is_loading=True, metadata=_DISTOPT_METADATA + ) + finally: + _restore_state_dict_patches(patches) + state_dict = dist_checkpointing.load(load_sd, checkpoint_dir, validate_access_integrity=False) + if load_model: + _load_model_state_dict(model, state_dict) + if load_optimizer and optimizer is not None and "optimizer" in state_dict: + load_patches = _patch_native_optimizer_step_load(optimizer) + try: + optimizer.load_state_dict(state_dict["optimizer"]) + finally: + _restore_set_state_patches(load_patches) + _synchronize_native_optimizer_steps(optimizer) + elif load_model and optimizer is not None: + reload_model_params = getattr(optimizer, "reload_model_params", None) + if callable(reload_model_params): + reload_model_params() + return int(state_dict.get("step", 0)) + + +def _synchronize_native_optimizer_steps(optimizer: Any) -> None: + """Align torch optimizer per-parameter steps before mcore fallback checkpointing.""" + + seen: set[int] = set() + + def visit(obj: Any) -> None: + obj_id = id(obj) + if obj_id in seen: + return + seen.add(obj_id) + + for child in _iter_optimizer_children(obj): + visit(child) + + state = getattr(obj, "state", None) + if isinstance(state, MutableMapping): + _synchronize_step_mapping(state) + + visit(optimizer) + + +def _patch_empty_native_optimizer_state_dicts( + optimizer: Any, *, fallback_step: int +) -> list[tuple[Any, Any]]: + patches: list[tuple[Any, Any]] = [] + for distopt in _iter_distributed_optimizers(optimizer): + inner = getattr(distopt, "optimizer", None) + state = getattr(inner, "state", None) + if not isinstance(state, MutableMapping) or state: + continue + original_state_dict = distopt.state_dict + + def patched_state_dict( + original_state_dict=original_state_dict, distopt=distopt, fallback_step=fallback_step + ): + try: + return original_state_dict() + except AssertionError: + return _empty_native_optimizer_state_dict(distopt, fallback_step) + + distopt.state_dict = patched_state_dict # type: ignore[method-assign] + patches.append((distopt, original_state_dict)) + return patches + + +def _restore_state_dict_patches(patches: list[tuple[Any, Any]]) -> None: + for distopt, original_state_dict in patches: + distopt.state_dict = original_state_dict # type: ignore[method-assign] + + +def _patch_native_optimizer_step_load(optimizer: Any) -> list[tuple[Any, Any]]: + patches: list[tuple[Any, Any]] = [] + for distopt in _iter_distributed_optimizers(optimizer): + original_set_state = distopt._set_main_param_and_optimizer_states + + def patched_set_state( + model_param, tensors, distopt=distopt, original_set_state=original_set_state + ): + removed_step = _pop_optimizer_step_for_model_param(distopt, model_param, tensors) + try: + return original_set_state(model_param, tensors) + finally: + if removed_step is not None: + state, step = removed_step + state["step"] = step + + distopt._set_main_param_and_optimizer_states = patched_set_state # type: ignore[method-assign] + patches.append((distopt, original_set_state)) + return patches + + +def _restore_set_state_patches(patches: list[tuple[Any, Any]]) -> None: + for distopt, original_set_state in patches: + distopt._set_main_param_and_optimizer_states = original_set_state # type: ignore[method-assign] + + +def _pop_optimizer_step_for_model_param( + distopt: Any, model_param, tensors: dict[str, Any] +) -> tuple[MutableMapping, Any] | None: + if "step" in tensors: + return None + try: + group_index, group_order = distopt.model_param_group_index_map[model_param] + main_param = distopt.optimizer.param_groups[group_index]["params"][group_order] + state = distopt.optimizer.state[main_param] + except (KeyError, IndexError, TypeError): + return None + if not isinstance(state, MutableMapping) or "step" not in state: + return None + return state, state.pop("step") + + +def _iter_distributed_optimizers(optimizer: Any) -> Iterable[Any]: + seen: set[int] = set() + + def visit(obj: Any): + obj_id = id(obj) + if obj_id in seen: + return + seen.add(obj_id) + + inner = _safe_inner_optimizer(obj) + if ( + callable(getattr(obj, "sharded_state_dict", None)) + and hasattr(obj, "gbuf_ranges") + and hasattr(obj, "buffers") + and inner is not None + ): + yield obj + + for child in _iter_optimizer_children(obj, known_inner=inner): + yield from visit(child) + + yield from visit(optimizer) + + +def _iter_optimizer_children(obj: Any, *, known_inner: Any | None = None) -> Iterable[Any]: + chained = getattr(obj, "chained_optimizers", None) + if isinstance(chained, Iterable): + yield from chained + + sub_optimizers = getattr(obj, "sub_optimizers", None) + if isinstance(sub_optimizers, Iterable): + yield from sub_optimizers + + inner = _safe_inner_optimizer(obj) if known_inner is None else known_inner + if inner is not None and inner is not obj: + yield inner + + +def _safe_inner_optimizer(obj: Any) -> Any | None: + if isinstance(getattr(obj, "chained_optimizers", None), Iterable): + # Megatron-Core ChainedOptimizer exposes `.optimizer` only for the + # single-optimizer compatibility case; multi-optimizer PP/EP chains + # assert on access. The children above are the real traversal targets. + return None + return getattr(obj, "optimizer", None) + + +def _empty_native_optimizer_state_dict(distopt: Any, fallback_step: int) -> dict[str, Any]: + inner_state_dict = distopt.optimizer.state_dict() + optimizer_state = { + key: ([group.copy() for group in value] if key == "param_groups" else value) + for key, value in inner_state_dict.items() + if key != "state" + } + for param_group in optimizer_state["param_groups"]: + param_group.pop("params", None) + param_group["step"] = int(fallback_step) + state_dict: dict[str, Any] = {"optimizer": optimizer_state} + grad_scaler = getattr(distopt, "grad_scaler", None) + if grad_scaler: + state_dict["grad_scaler"] = grad_scaler.state_dict() + return state_dict + + +def _synchronize_step_mapping(state: MutableMapping) -> None: + steps: list[Any] = [] + for param_state in state.values(): + if isinstance(param_state, MutableMapping) and "step" in param_state: + steps.append(param_state["step"]) + if not steps: + return + target = max(_step_as_int(step) for step in steps) + for param_state in state.values(): + if isinstance(param_state, MutableMapping) and "step" in param_state: + param_state["step"] = _step_like(param_state["step"], target) + + +def _step_as_int(step: Any) -> int: + if isinstance(step, torch.Tensor): + return int(step.detach().cpu().item()) + return int(step) + + +def _step_like(reference: Any, value: int) -> Any: + if isinstance(reference, torch.Tensor): + return torch.full_like(reference, value) + return value + + +def _build_bound_sharded_state_dict( + ps: ParallelState, get_placements: PlacementFn, is_expert: ExpertClassifierFn +) -> Callable: + def sharded_state_dict( + self, + prefix: str = "", + sharded_offsets: tuple[tuple[int, int, int], ...] = (), + metadata: dict | None = None, + ) -> dict[str, ShardedTensor]: + del metadata + return _module_sharded_state_dict( + _wrapped_module(self), + ps, + get_placements=get_placements, + is_expert=is_expert, + prefix=prefix, + sharded_offsets=sharded_offsets, + ) + + return sharded_state_dict + + +def _module_sharded_state_dict( + module: nn.Module, + ps: ParallelState, + *, + get_placements: PlacementFn, + is_expert: ExpertClassifierFn, + prefix: str = "", + sharded_offsets: tuple[tuple[int, int, int], ...] = (), +) -> dict[str, ShardedTensor]: + state: dict[str, ShardedTensor] = {} + for name, param in module.named_parameters(): + state[f"{prefix}{name}"] = _make_sharded_tensor( + f"{prefix}{name}", + param, + ps, + placements=get_placements(name), + expert=is_expert(name), + sharded_offsets=sharded_offsets, + ) + for name, buffer in module.named_buffers(): + state[f"{prefix}{name}"] = _make_sharded_tensor( + f"{prefix}{name}", + buffer, + ps, + placements=get_placements(name), + expert=is_expert(name), + sharded_offsets=sharded_offsets, + ) + return state + + +def _make_sharded_tensor( + key: str, + tensor: torch.Tensor, + ps: ParallelState, + *, + placements: list, + expert: bool, + sharded_offsets: tuple[tuple[int, int, int], ...] = (), +) -> ShardedTensor: + rank_offsets, replica_id = _rank_offsets_and_replica_id(placements, ps, expert=expert) + return ShardedTensor.from_rank_offsets( + key, tensor, *sharded_offsets, *rank_offsets, replica_id=replica_id + ) + + +def _rank_offsets_and_replica_id( + placements: list, ps: ParallelState, *, expert: bool +) -> tuple[tuple[tuple[int, int, int], ...], tuple[int, ...]]: + ranks, sizes = _mesh_ranks_and_sizes(ps, expert=expert) + axis_fragments: dict[int, tuple[int, int]] = {} + for placement, rank, size in zip(placements, ranks, sizes, strict=True): + if _is_shard_placement(placement): + dim = _shard_dim(placement) + if dim is None: + raise ValueError(f"Unsupported Shard placement without dim: {placement!r}.") + prev_rank, prev_size = axis_fragments.get(dim, (0, 1)) + axis_fragments[dim] = (prev_rank * size + rank, prev_size * size) + rank_offsets = tuple((dim, rank, size) for dim, (rank, size) in axis_fragments.items()) + return rank_offsets, _replica_id(placements, ps, expert=expert) + + +def _replica_id(placements: list, ps: ParallelState, *, expert: bool) -> tuple[int, int, int]: + # PP stages own different parameters. They are not replicas of one + # another, so PP rank must not make a shard non-main. + if expert: + return ( + 0, + _replica_axis_rank(placements, 2, ps.ep_rank), + _replica_axis_rank(placements, 1, ps.expert_dp_rank), + ) + dp_cp_rank = ( + 0 + if _placement_is_sharded(placements, 1) or _placement_is_sharded(placements, 2) + else ps.dp_cp_rank + ) + return (0, _replica_axis_rank(placements, 3, ps.tp_rank), int(dp_cp_rank)) + + +def _replica_axis_rank(placements: list, axis: int, rank: int) -> int: + return 0 if _placement_is_sharded(placements, axis) else int(rank) + + +def _placement_is_sharded(placements: list, axis: int) -> bool: + return axis < len(placements) and _is_shard_placement(placements[axis]) + + +def _mesh_ranks_and_sizes(ps: ParallelState, *, expert: bool) -> tuple[list[int], list[int]]: + if expert: + return ( + [ps.pp_rank, ps.expert_dp_rank, ps.ep_rank, ps.etp_rank], + [ps.pp_size, ps.expert_dp_size, ps.ep_size, ps.etp_size], + ) + return ( + [ps.pp_rank, ps.dp_rank, ps.cp_rank, ps.tp_rank], + [ps.pp_size, ps.dp_size, ps.cp_size, ps.tp_size], + ) + + +def _is_shard_placement(placement: Any) -> bool: + return type(placement).__name__ == "Shard" + + +def _shard_dim(placement: Any) -> int | None: + dim = getattr(placement, "dim", None) + if dim is None: + dim = getattr(placement, "_dim", None) + return None if dim is None else int(dim) + + +def _model_sharded_state_dict(model: nn.Module | Iterable[nn.Module]) -> dict[str, Any]: + chunks = _model_chunks(model) + ps = _chunk_parallel_state(chunks[0]) if chunks else None + return { + _model_chunk_key(ps, idx, len(chunks)): _chunk_sharded_state_dict( + chunk, _model_chunk_sharded_key_prefix(ps, idx, len(chunks)) + ) + for idx, chunk in enumerate(chunks) + } + + +def _model_chunk_key(ps: ParallelState | None, idx: int, num_chunks: int) -> str: + if ps is not None and ps.pp_size > 1: + key = f"model_pp{ps.pp_rank}" + if num_chunks > 1: + key = f"{key}_vpp{idx}" + return key + if num_chunks == 1: + return "model" + return f"model{idx}" + + +def _model_chunk_sharded_key_prefix(ps: ParallelState | None, idx: int, num_chunks: int) -> str: + if ps is None and num_chunks == 1: + return "" + if ps is not None and ps.pp_size <= 1 and num_chunks == 1: + return "" + return f"{_model_chunk_key(ps, idx, num_chunks)}." + + +def _chunk_sharded_state_dict(chunk: nn.Module, sharded_key_prefix: str) -> dict[str, Any]: + chunk_sd = chunk.sharded_state_dict() # type: ignore[attr-defined] + if not sharded_key_prefix: + return chunk_sd + return { + key: ( + replace(value, key=f"{sharded_key_prefix}{value.key}") + if isinstance(value, ShardedTensor) + else value + ) + for key, value in chunk_sd.items() + } + + +def _single_or_all_model_state(model_sd: dict[str, Any]) -> dict[str, Any]: + if "model" in model_sd: + return model_sd["model"] + return model_sd + + +def _load_model_state_dict( + model: nn.Module | Iterable[nn.Module], state_dict: dict[str, Any] +) -> None: + chunks = _model_chunks(model) + if len(chunks) == 1 and "model" in state_dict: + _wrapped_module(chunks[0]).load_state_dict(state_dict["model"], strict=False) + return + ps = _chunk_parallel_state(chunks[0]) if chunks else None + if len(chunks) == 1 and ps is not None and ps.pp_size > 1: + key = _model_chunk_key(ps, 0, 1) + if key in state_dict: + _wrapped_module(chunks[0]).load_state_dict(state_dict[key], strict=False) + return + for idx, chunk in enumerate(chunks): + key = _model_chunk_key(ps, idx, len(chunks)) + if key in state_dict: + _wrapped_module(chunk).load_state_dict(state_dict[key], strict=False) + + +def _chunk_parallel_state(chunk: nn.Module) -> ParallelState | None: + ps = getattr(chunk, "_mlite_distopt_parallel_state", None) + if ps is not None: + return ps + wrapped = _wrapped_module(chunk) + return getattr(wrapped, "_mlite_distopt_parallel_state", None) + + +def _model_chunks(model: nn.Module | Iterable[nn.Module]) -> list[nn.Module]: + if isinstance(model, nn.Module): + return list(model) if isinstance(model, nn.ModuleList) else [model] + chunks = list(model) + if not all(isinstance(chunk, nn.Module) for chunk in chunks): + raise TypeError("distckpt model chunks must be nn.Module instances.") + return chunks + + +def _wrapped_module(model: nn.Module) -> nn.Module: + module = getattr(model, "module", None) + return module if isinstance(module, nn.Module) else model + + +__all__ = [ + "attach_model_sharded_state_dict", + "load_distopt_checkpoint", + "save_distopt_checkpoint", + "supports_distopt_distckpt", +] diff --git a/experimental/lite/megatron/lite/primitive/ckpt/hf_weights.py b/experimental/lite/megatron/lite/primitive/ckpt/hf_weights.py new file mode 100644 index 00000000000..aa94fd0d5ce --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/ckpt/hf_weights.py @@ -0,0 +1,593 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""HF ↔ Megatron Lite weight conversion. + +Everything needed for HuggingFace safetensors ↔ Megatron Lite model conversion: +- HFWeights protocol (model implements this) +- SafeTensorReader / save_safetensors (file I/O) +- Tensor utilities (split_dim, allgather_concat, remap_layer_index, ...) +- Generic load_hf_weights / export_hf_weights / save_hf_weights (orchestration) +""" + +from __future__ import annotations + +import json +import os +import re +from collections.abc import Generator +from pathlib import Path +from typing import Protocol, runtime_checkable + +import torch +import torch.distributed as dist +import torch.nn as nn +from safetensors import safe_open +from safetensors.torch import save_file as _safe_save + + +@runtime_checkable +class HFWeights(Protocol): + """Protocol for HF ↔ Megatron Lite weight conversion. + + Model-specific implementations only do tensor math, never distributed comm. + """ + + def weight_map(self) -> dict[str, list[str]]: + """Megatron Lite param name → [HF param names]. Multiple = concat (QKV, gate+up).""" + ... + + def hf_to_native(self, native_name: str, hf_tensors: list[torch.Tensor]) -> torch.Tensor: + """Convert HF tensors → single Megatron Lite tensor (e.g. merge QKV).""" + ... + + def native_to_hf( + self, native_name: str, tensor: torch.Tensor + ) -> list[tuple[str, torch.Tensor]]: + """Convert Megatron Lite tensor → [(hf_name, hf_tensor)] (e.g. split QKV back).""" + ... + + def tp_spec(self, native_name: str) -> tuple[int, int] | None: + """TP sharding: ``(split_dim, 0=TP|1=ETP)``, or None if replicated.""" + ... + + def qkv_spec(self, native_name: str) -> tuple[int, int, int] | None: + """If native_name is a fused QKV weight, return (num_q_heads, num_kv_heads, head_dim). + + Needed for correct GQA TP sharding — Q/K/V must be split independently. + Return None for non-QKV parameters. + """ + return None + + @property + def num_experts(self) -> int: + """Total number of experts (needed for EP gather index math).""" + ... + + def is_expert(self, native_name: str) -> bool: + """Whether this param belongs to an expert (for EP sharding).""" + ... + + def expert_global_id(self, native_name: str) -> int | None: + """Global expert ID from synthetic name. None if not expert.""" + ... + + def expert_local_name(self, native_name: str, local_idx: int) -> str: + """Synthetic expert name → actual model param name.""" + ... + + +# ====================================================================== +# SafeTensors I/O +# ====================================================================== + + +class SafeTensorReader: + """Read individual tensors from an HF safetensors directory.""" + + def __init__(self, path: str): + self.path = Path(path) + self.index = self._load_index() + + def _load_index(self) -> dict[str, str]: + idx_file = self.path / "model.safetensors.index.json" + if idx_file.exists(): + with open(idx_file) as f: + return json.load(f)["weight_map"] + return {} + + def get_tensor(self, name: str) -> torch.Tensor: + if self.index: + filepath = self.path / self.index[name] + else: + filepath = self.path / "model.safetensors" + with safe_open(str(filepath), framework="pt", device="cpu") as f: + return f.get_tensor(name) + + +def unwrap_model(model: nn.Module) -> nn.Module: + """Strip nested wrapper modules like DDP -> model.""" + base_model = model + seen: set[int] = set() + while hasattr(base_model, "module"): + ident = id(base_model) + if ident in seen: + break + seen.add(ident) + next_model = base_model.module + if not isinstance(next_model, nn.Module) or next_model is base_model: + break + base_model = next_model + return base_model + + +def save_safetensors( + tensors: dict[str, torch.Tensor], path: str, filename: str = "model.safetensors" +) -> None: + os.makedirs(path, exist_ok=True) + _safe_save(tensors, os.path.join(path, filename)) + + +def _resolve_export_dtype(export_dtype: str | torch.dtype | None) -> torch.dtype | None: + if export_dtype is None: + return None + if isinstance(export_dtype, torch.dtype): + return export_dtype + normalized = str(export_dtype).lower() + aliases = { + "bf16": torch.bfloat16, + "bfloat16": torch.bfloat16, + "fp16": torch.float16, + "float16": torch.float16, + "half": torch.float16, + "fp32": torch.float32, + "float32": torch.float32, + "float": torch.float32, + } + if normalized not in aliases: + raise ValueError(f"Unsupported export_dtype={export_dtype!r}") + return aliases[normalized] + + +def _cast_export_tensor(tensor: torch.Tensor, export_dtype: torch.dtype | None) -> torch.Tensor: + if export_dtype is None or not tensor.is_floating_point(): + return tensor + return tensor.to(dtype=export_dtype) + + +# ====================================================================== +# Tensor utilities +# ====================================================================== + + +def split_dim(tensor: torch.Tensor, rank: int, world: int, dim: int = 0) -> torch.Tensor: + if world <= 1: + return tensor + return tensor.chunk(world, dim=dim)[rank].contiguous() + + +def split_qkv( + tensor: torch.Tensor, rank: int, world: int, num_q_heads: int, num_kv_heads: int, head_dim: int +) -> torch.Tensor: + """TP-shard a fused [Q, K, V] weight, splitting Q/K/V heads independently. + + Naive ``split_dim`` would slice across the Q/K/V boundary incorrectly + when num_q_heads != num_kv_heads (GQA). + """ + if world <= 1: + return tensor + q_size = num_q_heads * head_dim + kv_size = num_kv_heads * head_dim + q = tensor[:q_size] + k = tensor[q_size : q_size + kv_size] + v = tensor[q_size + kv_size :] + q_shard = q.chunk(world, dim=0)[rank] + k_shard = k.chunk(world, dim=0)[rank] + v_shard = v.chunk(world, dim=0)[rank] + return torch.cat([q_shard, k_shard, v_shard], dim=0).contiguous() + + +def split_gate_up(tensor: torch.Tensor, rank: int, world: int) -> torch.Tensor: + if world <= 1: + return tensor + ffn = tensor.shape[0] // 2 + gate = tensor[:ffn].chunk(world, dim=0)[rank] + up = tensor[ffn:].chunk(world, dim=0)[rank] + return torch.cat([gate, up], dim=0).contiguous() + + +def allgather_concat( + tensor: torch.Tensor, world_size: int, group: dist.ProcessGroup | None, dim: int +) -> torch.Tensor: + gathered = [torch.empty_like(tensor) for _ in range(world_size)] + dist.all_gather(gathered, tensor.contiguous(), group=group) + return torch.cat(gathered, dim=dim) + + +def remap_layer_index(name: str, global_to_local: dict[int, int]) -> str | None: + if not global_to_local: + return name + m = re.match(r"(layers\.)(\d+)(\..*)", name) + if not m: + return name + gidx = int(m.group(2)) + if gidx not in global_to_local: + return None + return f"{m.group(1)}{global_to_local[gidx]}{m.group(3)}" + + +def extract_layer_idx(name: str) -> int: + m = re.search(r"layers\.(\d+)\.", name) + return int(m.group(1)) if m else 0 + + +def parse_expert_idx(name: str) -> int: + m = re.search(r"weight(\d+)$", name) + return int(m.group(1)) if m else 0 + + +def set_expert_idx(name: str, idx: int) -> str: + return re.sub(r"weight\d+$", f"weight{idx}", name) + + +def to_global_layer_name(name: str, layer_map: dict[int, int]) -> str: + if not layer_map: + return name + + def _replace(m: re.Match) -> str: + return f"layers.{layer_map.get(int(m.group(1)), int(m.group(1)))}." + + return re.sub(r"layers\.(\d+)\.", _replace, name) + + +def gather_gate_up(tensor: torch.Tensor, world_size: int, group: dist.ProcessGroup) -> torch.Tensor: + ffn_local = tensor.shape[0] // 2 + gate_full = allgather_concat(tensor[:ffn_local], world_size, group, dim=0) + up_full = allgather_concat(tensor[ffn_local:], world_size, group, dim=0) + return torch.cat([gate_full, up_full], dim=0) + + +# ====================================================================== +# Generic load / export / save using HFWeights +# ====================================================================== + + +def load_hf_weights( + model: nn.Module, hf_path: str, spec: HFWeights, ps, *, vocab_size: int | None = None +) -> None: + """Load HF safetensors into a Megatron Lite model using HFWeights. + + Handles PP layer filtering, TP split, EP shard assignment. + ``ps`` is a ParallelState (lazy import to avoid GPU dep at module level). + """ + from megatron.lite.primitive.parallel import pad_vocab_for_tp + from megatron.lite.primitive.utils import log_rank0 + + base_model = unwrap_model(model) + reader = SafeTensorReader(hf_path) + wmap = spec.weight_map() + + global_to_local: dict[int, int] = ( + {gi: li for li, gi in enumerate(base_model.layer_indices)} + if hasattr(base_model, "layer_indices") + else {} + ) + + state = base_model.state_dict() + loaded: dict[str, torch.Tensor] = {} + num_experts_total = getattr(spec, "num_experts", None) + expert_shard = None + if num_experts_total is None: + expert_ids = [spec.expert_global_id(name) for name in wmap] + expert_ids = [expert_id for expert_id in expert_ids if expert_id is not None] + if expert_ids: + num_experts_total = max(expert_ids) + 1 + if num_experts_total: + from megatron.lite.primitive.utils import ensure_divisible + + experts_per_rank = ensure_divisible(num_experts_total, ps.ep_size) + local_start = ps.ep_rank * experts_per_rank + expert_shard = (experts_per_rank, local_start) + + for native_name, hf_names in wmap.items(): + mapped = remap_layer_index(native_name, global_to_local) + if mapped is None: + continue + + expert_gid = spec.expert_global_id(mapped) + if expert_gid is not None: + _load_expert_weight( + mapped, hf_names, reader, spec, ps, loaded, expert_gid, expert_shard + ) + continue + + hf_tensors = [reader.get_tensor(n) for n in hf_names] + tensor = spec.hf_to_native(mapped, hf_tensors) + + tp_info = spec.tp_spec(mapped) + if tp_info is not None: + split_d, tp_or_etp = tp_info + if tp_or_etp == 0: + if vocab_size is not None and ("embed" in mapped or "head" in mapped): + padded = pad_vocab_for_tp(vocab_size, ps.tp_size) + if tensor.size(0) < padded: + pad = torch.zeros( + padded - tensor.size(0), *tensor.shape[1:], dtype=tensor.dtype + ) + tensor = torch.cat([tensor, pad], dim=0) + qkv = spec.qkv_spec(mapped) if hasattr(spec, "qkv_spec") else None + if qkv is not None: + tensor = split_qkv(tensor, ps.tp_rank, ps.tp_size, *qkv) + else: + tensor = split_dim(tensor, ps.tp_rank, ps.tp_size, dim=split_d) + else: + tensor = split_dim(tensor, ps.etp_rank, ps.etp_size, dim=split_d) + + actual = _resolve_param_name(mapped, state) + if actual: + loaded[actual] = tensor.to(dtype=torch.bfloat16) + + for name, param in base_model.named_parameters(): + if name in loaded: + param.data.copy_(loaded[name]) + elif "lora" in name.lower() or "adapter" in name.lower(): + continue + else: + log_rank0(f"WARNING: {name} not loaded from checkpoint") + + +def _load_expert_weight(native_name, hf_names, reader, spec, ps, loaded, expert_gid, expert_shard): + if expert_shard is None: + raise RuntimeError("Expert weight encountered but expert shard metadata is unavailable.") + experts_per_rank, local_start = expert_shard + if expert_gid < local_start or expert_gid >= local_start + experts_per_rank: + return + + hf_tensors = [reader.get_tensor(n) for n in hf_names] + tensor = spec.hf_to_native(native_name, hf_tensors) + + if ps.etp_size > 1: + tp_info = spec.tp_spec(native_name) + if tp_info is not None: + split_d, _ = tp_info + if "fc1" in native_name: + tensor = split_gate_up(tensor, ps.etp_rank, ps.etp_size) + else: + tensor = split_dim(tensor, ps.etp_rank, ps.etp_size, dim=split_d) + + loaded[spec.expert_local_name(native_name, expert_gid - local_start)] = tensor.to( + dtype=torch.bfloat16 + ) + + +def _resolve_param_name(name: str, state_dict: dict) -> str | None: + if name in state_dict: + return name + for key in state_dict: + if name in key: + return key + return None + + +def export_hf_weights( + model: nn.Module | list[nn.Module], + spec: HFWeights, + ps, + *, + vocab_size: int | None = None, + limit: int | None = None, + rank0_only: bool = False, + export_dtype: str | torch.dtype | None = None, +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Export model weights as HF-format (name, tensor) pairs. + + Gathers across TP/ETP/EP/PP so the output is the full unsharded HF state on + every participating rank. RL weight sync needs every colocated rollout rank + to receive weights; save paths can pass ``rank0_only=True`` to avoid + materializing duplicate writers. + """ + if isinstance(model, nn.ModuleList): + chunks: list[nn.Module] = list(model) + elif isinstance(model, list): + chunks = model + else: + chunks = [model] + + rank = dist.get_rank() if dist.is_initialized() else 0 + resolved_export_dtype = _resolve_export_dtype(export_dtype) + + if ps.pp_size <= 1: + exported_params = 0 + expert_groups: dict[str, list[tuple[int, str, torch.Tensor]]] = {} + for chunk in chunks: + base_chunk = unwrap_model(chunk) + layer_map = ( + {i: base_chunk.layer_indices[i] for i in range(len(base_chunk.layer_indices))} + if hasattr(base_chunk, "layer_indices") + else {} + ) + for name, param in base_chunk.named_parameters(): + gname = to_global_layer_name(name, layer_map) + tensor = param.data.detach() + + gathered_one: dict[str, torch.Tensor] = {} + if spec.is_expert(gname): + if limit is None: + expert_groups.setdefault(_expert_group_key(gname), []).append( + (parse_expert_idx(gname), gname, tensor) + ) + exported_params += 1 + continue + _gather_expert(gname, tensor, spec, ps, gathered_one) + else: + gathered_one[gname] = _gather_dense(gname, tensor, spec, ps) + + exported_params += 1 + if not rank0_only or rank == 0: + for native_name, gathered_tensor in gathered_one.items(): + if vocab_size is not None and ( + "embed" in native_name or "head" in native_name + ): + gathered_tensor = gathered_tensor[:vocab_size] + for hf_name, hf_tensor in spec.native_to_hf(native_name, gathered_tensor): + yield hf_name, _cast_export_tensor(hf_tensor, resolved_export_dtype) + + if limit is not None and exported_params >= limit: + return + + for group_key in sorted(expert_groups): + gathered_group: dict[str, torch.Tensor] = {} + _gather_expert_group(expert_groups[group_key], spec, ps, gathered_group) + if not rank0_only or rank == 0: + for native_name in sorted(gathered_group, key=parse_expert_idx): + gathered_tensor = gathered_group[native_name] + for hf_name, hf_tensor in spec.native_to_hf(native_name, gathered_tensor): + yield hf_name, _cast_export_tensor(hf_tensor, resolved_export_dtype) + return + + gathered: dict[str, torch.Tensor] = {} + for chunk in chunks: + base_chunk = unwrap_model(chunk) + # Map local layer indices to global for PP + layer_map = ( + {i: base_chunk.layer_indices[i] for i in range(len(base_chunk.layer_indices))} + if hasattr(base_chunk, "layer_indices") + else {} + ) + for name, param in base_chunk.named_parameters(): + gname = to_global_layer_name(name, layer_map) + t = param.data.detach() + + if spec.is_expert(gname): + _gather_expert(gname, t, spec, ps, gathered) + else: + gathered[gname] = _gather_dense(gname, t, spec, ps) + + # PP gather + if ps.pp_size > 1: + all_states: list[dict | None] = [None] * ps.pp_size + dist.all_gather_object(all_states, gathered, group=ps.pp_group) + gathered = {} + for s in all_states: + if s is not None: + gathered.update(s) + + rank = dist.get_rank() if dist.is_initialized() else 0 + if rank0_only and rank != 0: + return + + # Vocab trim + if vocab_size is not None: + for key in list(gathered.keys()): + if "embed" in key or "head" in key: + gathered[key] = gathered[key][:vocab_size] + + # Convert Megatron Lite names → HF names via spec + for native_name, tensor in gathered.items(): + for hf_name, hf_tensor in spec.native_to_hf(native_name, tensor): + yield hf_name, _cast_export_tensor(hf_tensor, resolved_export_dtype) + + +def _gather_dense(name: str, tensor: torch.Tensor, spec: HFWeights, ps) -> torch.Tensor: + """Gather a dense (non-expert) param across TP.""" + custom_gather = getattr(spec, "gather_dense", None) + if callable(custom_gather): + gathered = custom_gather(name, tensor, ps) + if gathered is not None: + return gathered.cpu() + + tp_info = spec.tp_spec(name) + if tp_info is not None and ps.tp_size > 1: + split_d, tp_or_etp = tp_info + if tp_or_etp == 0: + tensor = allgather_concat(tensor, ps.tp_size, ps.tp_group, dim=split_d) + return tensor.cpu() + + +def _gather_expert( + name: str, tensor: torch.Tensor, spec: HFWeights, ps, out: dict[str, torch.Tensor] +) -> None: + """Gather an expert param across ETP + EP.""" + tensor = _gather_expert_etp(name, tensor, spec, ps) + + # EP gather: global_id = ep_rank * n_local + local_id. + local_idx = parse_expert_idx(name) + if ps.ep_size > 1 and ps.ep_group is not None: + n_local = spec.num_experts // ps.ep_size + ep_gathered = [torch.empty_like(tensor) for _ in range(ps.ep_size)] + dist.all_gather(ep_gathered, tensor.contiguous(), group=ps.ep_group) + for ep_rank, ep_tensor in enumerate(ep_gathered): + global_idx = ep_rank * n_local + local_idx + out[set_expert_idx(name, global_idx)] = ep_tensor.cpu() + else: + out[name] = tensor.cpu() + + +def _gather_expert_etp(name: str, tensor: torch.Tensor, spec: HFWeights, ps) -> torch.Tensor: + # ETP gather + if ps.etp_size > 1 and ps.etp_group is not None: + tp_info = spec.tp_spec(name) + if tp_info is not None: + split_d, _ = tp_info + if "fc1" in name: + return gather_gate_up(tensor, ps.etp_size, ps.etp_group) + return allgather_concat(tensor, ps.etp_size, ps.etp_group, dim=split_d) + return tensor + + +def _expert_group_key(name: str) -> str: + return re.sub(r"weight\d+$", "weight", name) + + +def _gather_expert_group( + entries: list[tuple[int, str, torch.Tensor]], spec: HFWeights, ps, out: dict[str, torch.Tensor] +) -> None: + """Gather local experts in one EP collective per layer/kind.""" + prepared = [ + (local_idx, name, _gather_expert_etp(name, tensor, spec, ps)) + for local_idx, name, tensor in sorted(entries) + ] + packed_group_name = getattr(spec, "packed_expert_group_name", None) + if callable(packed_group_name): + packed_name = packed_group_name(prepared[0][1]) + if packed_name is not None: + if ps.ep_size <= 1 or ps.ep_group is None: + out[packed_name] = torch.stack( + [tensor.contiguous() for _, _, tensor in prepared], dim=0 + ).cpu() + return + + stacked = torch.stack([tensor.contiguous() for _, _, tensor in prepared], dim=0) + ep_gathered = [torch.empty_like(stacked) for _ in range(ps.ep_size)] + dist.all_gather(ep_gathered, stacked, group=ps.ep_group) + out[packed_name] = torch.cat(ep_gathered, dim=0).cpu() + return + + if ps.ep_size <= 1 or ps.ep_group is None: + for _, name, tensor in prepared: + out[name] = tensor.cpu() + return + + n_local = spec.num_experts // ps.ep_size + stacked = torch.stack([tensor.contiguous() for _, _, tensor in prepared], dim=0) + ep_gathered = [torch.empty_like(stacked) for _ in range(ps.ep_size)] + dist.all_gather(ep_gathered, stacked, group=ps.ep_group) + for ep_rank, ep_tensor in enumerate(ep_gathered): + for slot, (local_idx, name, _) in enumerate(prepared): + global_idx = ep_rank * n_local + local_idx + out[set_expert_idx(name, global_idx)] = ep_tensor[slot].cpu() + + +def save_hf_weights( + model: nn.Module | list[nn.Module], + hf_path: str, + spec: HFWeights, + ps, + *, + vocab_size: int | None = None, +) -> None: + """Export + write to safetensors.""" + rank = dist.get_rank() if dist.is_initialized() else 0 + out = dict(export_hf_weights(model, spec, ps, vocab_size=vocab_size, rank0_only=True)) + if rank == 0 and out: + save_safetensors(out, hf_path) + if dist.is_initialized(): + dist.barrier() diff --git a/experimental/lite/megatron/lite/primitive/config.py b/experimental/lite/megatron/lite/primitive/config.py new file mode 100644 index 00000000000..c13bbf4fb86 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/config.py @@ -0,0 +1,54 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Configuration utilities. + +Model-specific configs live under `megatron.lite/models/*/config.py` +as standalone dataclasses. +""" + +from __future__ import annotations + +import json +from dataclasses import fields as dc_fields +from pathlib import Path +from typing import Any + + +def to_dataclass(oc_cfg, cls): + """Convert an OmegaConf DictConfig back to a typed dataclass instance.""" + from omegaconf import OmegaConf + + init_names = {f.name for f in dc_fields(cls) if f.init} + d = OmegaConf.to_container(oc_cfg, resolve=True) + return cls(**{k: v for k, v in d.items() if k in init_names}) + + +def load_hf_config_dict(path_or_name: str) -> dict[str, Any]: + """Load HF config dict from local path or Hub.""" + p = Path(path_or_name) + + if p.is_file() and p.name == "config.json": + with open(p) as f: + return json.load(f) + + if p.is_dir(): + cfg_file = p / "config.json" + if cfg_file.exists(): + with open(cfg_file) as f: + return json.load(f) + raise FileNotFoundError(f"No config.json in {p}") + + try: + from transformers import AutoConfig + + hf_config = AutoConfig.from_pretrained(path_or_name, trust_remote_code=True) + return hf_config.to_dict() + except ImportError as err: + raise ImportError( + f"'{path_or_name}' is not a local path. " + "Install transformers to load from HuggingFace Hub: pip install transformers" + ) from err + except Exception as e: + raise ValueError(f"Failed to load config from '{path_or_name}': {e}") from e + + +__all__ = ["load_hf_config_dict", "to_dataclass"] diff --git a/experimental/lite/megatron/lite/primitive/data.py b/experimental/lite/megatron/lite/primitive/data.py new file mode 100644 index 00000000000..776f7ba9cc8 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/data.py @@ -0,0 +1,198 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Batch generation utilities for Megatron Lite runtime.""" + +from __future__ import annotations + +import os + +import torch # pyright: ignore[reportMissingImports] + +_TRUE_ENV_VALUES = {"1", "true", "yes", "on"} +_FALSE_ENV_VALUES = {"0", "false", "no", "off"} + + +def _read_bool_env(name: str) -> bool | None: + value = os.environ.get(name) + if value is None or value.strip() == "": + return None + normalized = value.strip().lower() + if normalized in _TRUE_ENV_VALUES: + return True + if normalized in _FALSE_ENV_VALUES: + return False + raise ValueError(f"{name} must be a boolean value, got {value!r}") + + +def _resolve_thd_padding(seq_len: int, cp_size: int) -> tuple[int, int, bool]: + """Return (padded_seq_len, alignment_multiple, pad_enabled) for THD input.""" + if seq_len <= 0: + raise ValueError(f"seq_len must be positive, got {seq_len}") + if cp_size < 1: + raise ValueError(f"cp_size must be >= 1, got {cp_size}") + + pad_to_alignment = _read_bool_env("MEGATRON_LITE_THD_PAD_TO_ALIGNMENT") + if pad_to_alignment is None: + pad_to_alignment = cp_size > 1 + + pad_multiple_env = os.environ.get("MEGATRON_LITE_THD_PAD_MULTIPLE", "auto").strip().lower() + if pad_multiple_env in ("", "auto", "cp", "min_cp"): + align_size = cp_size * 2 if cp_size > 1 else 1 + else: + try: + align_size = int(pad_multiple_env) + except ValueError as exc: + raise ValueError( + "MEGATRON_LITE_THD_PAD_MULTIPLE must be 'auto' or a positive integer, " + f"got {pad_multiple_env!r}" + ) from exc + if align_size <= 0: + raise ValueError( + "MEGATRON_LITE_THD_PAD_MULTIPLE must be 'auto' or a positive integer, " + f"got {pad_multiple_env!r}" + ) + + if not pad_to_alignment: + return seq_len, align_size, False + + seq_len_padded = ((seq_len + align_size - 1) // align_size) * align_size + return seq_len_padded, align_size, True + + +def fixed_batches( + vocab_size: int, + seq_len: int, + num_steps: int, + batch_size: int = 1, + device: str = "cuda", + seed: int = 42, +) -> list[tuple[torch.Tensor, torch.Tensor]]: + """Deterministic (input_ids, labels) pairs, identical across all ranks.""" + g = torch.Generator().manual_seed(seed) + batches = [] + for _ in range(num_steps): + ids = torch.randint(0, vocab_size, (batch_size, seq_len), generator=g).to(device) + labels = torch.randint(0, vocab_size, (batch_size, seq_len), generator=g).to(device) + batches.append((ids, labels)) + return batches + + +def infinite_batches( + vocab_size: int, seq_len: int, batch_size: int = 1, device: str = "cuda", seed: int = 42 +): + """Infinite deterministic batch generator (for throughput benchmarks).""" + g = torch.Generator(device=device).manual_seed(seed) + while True: + yield { + "input_ids": torch.randint( + 0, vocab_size, (batch_size, seq_len), device=device, generator=g + ), + "labels": torch.randint( + 0, vocab_size, (batch_size, seq_len), device=device, generator=g + ), + } + + +def infinite_batches_thd( + vocab_size: int, + seq_len: int, + *, + cp_size: int = 1, + cp_rank: int = 0, + device: str = "cuda", + seed: int = 42, +): + """Infinite deterministic batch generator in THD (packed variable-length) format. + + Produces plain dicts consumed directly by native lite model forwards + (input_ids 2-D batch=1, mrope position_ids (3,1,T), pre-built PackedSeqParams). + + cp_size/cp_rank should be supplied explicitly by the caller (for example + from handle._parallel_state). The defaults keep older single-CP call sites + working, but CP runs must pass the real CP rank and world size. + + When context parallelism is active (cp_size > 1), each rank's packed tokens are + split via zigzag striping. position_ids stay FULL because lite's RoPE + (is_thd_format=False) auto-slices emb internally, same as BSH path. + """ + from megatron.core.packed_seq_params import ( # pyright: ignore[reportMissingImports] # noqa: I001 + PackedSeqParams, + ) + + if cp_size < 1: + raise ValueError(f"cp_size must be >= 1, got {cp_size}") + if cp_rank < 0 or cp_rank >= cp_size: + raise ValueError(f"cp_rank must be in [0, {cp_size}), got {cp_rank}") + + g = torch.Generator(device=device).manual_seed(seed) + seq_len_padded, _align_size, _pad_to_alignment = _resolve_thd_padding(seq_len, cp_size) + + # position_ids always stay full length; RoPE auto-slices for CP. + position_ids = ( + torch.arange(seq_len_padded, device=device) + .view(1, 1, -1) + .expand(3, 1, seq_len_padded) + .contiguous() + ) + cu_seqlens = torch.tensor([0, seq_len], dtype=torch.int32, device=device) + cu_seqlens_padded = torch.tensor([0, seq_len_padded], dtype=torch.int32, device=device) + + if cp_size > 1: + from megatron.lite.primitive.parallel.cp import zigzag_split_for_cp # noqa: I001 + + # cu_seqlens stays full; MC GDN internally divides by cp_size. + packed_seq_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=seq_len_padded, + max_seqlen_kv=seq_len_padded, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + ) + while True: + ids_full = torch.randint(0, vocab_size, (seq_len,), device=device, generator=g) + lbl_full = torch.randint(0, vocab_size, (seq_len,), device=device, generator=g) + if seq_len_padded != seq_len: + ids_padded = torch.zeros(seq_len_padded, dtype=ids_full.dtype, device=device) + lbl_padded = torch.zeros(seq_len_padded, dtype=lbl_full.dtype, device=device) + ids_padded[:seq_len] = ids_full + lbl_padded[:seq_len] = lbl_full + else: + ids_padded = ids_full + lbl_padded = lbl_full + ids_cp = zigzag_split_for_cp(ids_padded, cp_rank, cp_size, seq_dim=0) + lbl_cp = zigzag_split_for_cp(lbl_padded, cp_rank, cp_size, seq_dim=0) + yield { + "input_ids": ids_cp.unsqueeze(0), # (1, T/cp) + "labels": lbl_cp.unsqueeze(0), # (1, T/cp) + "position_ids": position_ids, # FULL (3, 1, T) + "packed_seq_params": packed_seq_params, + } + else: + packed_seq_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=seq_len_padded, + max_seqlen_kv=seq_len_padded, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + ) + while True: + input_ids = torch.zeros((1, seq_len_padded), dtype=torch.long, device=device) + labels = torch.zeros((1, seq_len_padded), dtype=torch.long, device=device) + input_ids[:, :seq_len] = torch.randint( + 0, vocab_size, (1, seq_len), device=device, generator=g + ) + labels[:, :seq_len] = torch.randint( + 0, vocab_size, (1, seq_len), device=device, generator=g + ) + yield { + "input_ids": input_ids, + "labels": labels, + "position_ids": position_ids, + "packed_seq_params": packed_seq_params, + } + + +__all__ = ["fixed_batches", "infinite_batches", "infinite_batches_thd"] diff --git a/experimental/lite/megatron/lite/primitive/deterministic.py b/experimental/lite/megatron/lite/primitive/deterministic.py new file mode 100644 index 00000000000..f1a736bf690 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/deterministic.py @@ -0,0 +1,35 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Deterministic computation utilities for bitwise reproducible experiments.""" + +from __future__ import annotations + +import os + +import torch + +_is_deterministic: bool = False +_TRUE_VALUES = {"1", "true", "yes", "on"} +_DETERMINISTIC_ENV = "MEGATRON_LITE_DETERMINISTIC" + + +def deterministic_requested() -> bool: + """Return whether bitwise-validation deterministic mode was requested.""" + return os.environ.get(_DETERMINISTIC_ENV, "").strip().lower() in _TRUE_VALUES + + +def set_deterministic(seed: int = 42) -> None: + """Enable all deterministic flags. Must call before first CUDA op in process.""" + global _is_deterministic + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0" + torch.use_deterministic_algorithms(True, warn_only=False) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + _is_deterministic = True + + +def is_deterministic() -> bool: + return _is_deterministic diff --git a/experimental/lite/megatron/lite/primitive/modules/__init__.py b/experimental/lite/megatron/lite/primitive/modules/__init__.py new file mode 100644 index 00000000000..ff35613ced6 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/modules/__init__.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Shared model modules owned by Megatron Lite.""" + +from __future__ import annotations + +_EXPORTS = { + "Experts": ("megatron.lite.primitive.modules.experts", "Experts"), + "GatedDeltaNet": ("megatron.lite.primitive.modules.gated_delta_net", "GatedDeltaNet"), + "GQAttention": ("megatron.lite.primitive.modules.gqa", "GQAttention"), + "MTPBlock": ("megatron.lite.primitive.modules.mtp", "MTPBlock"), + "MTPDecoderLayer": ("megatron.lite.primitive.modules.mtp", "MTPDecoderLayer"), + "MTPLossAutoScaler": ("megatron.lite.primitive.modules.mtp", "MTPLossAutoScaler"), + "MoEAuxLossAutoScaler": ("megatron.lite.primitive.modules.moe", "MoEAuxLossAutoScaler"), + "MultimodalRotaryEmbedding": ( + "megatron.lite.primitive.modules.mrope", + "MultimodalRotaryEmbedding", + ), + "SigmoidTopKRouter": ("megatron.lite.primitive.modules.router", "SigmoidTopKRouter"), + "TokenDispatcher": ("megatron.lite.primitive.modules.dispatcher", "TokenDispatcher"), + "TopKRouter": ("megatron.lite.primitive.modules.router", "TopKRouter"), + "_AllToAll": ("megatron.lite.primitive.modules.moe", "_AllToAll"), + "split_grouped_qkvg": ("megatron.lite.primitive.modules.gqa_utils", "split_grouped_qkvg"), +} + + +def __getattr__(name: str): + if name not in _EXPORTS: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + import importlib + + module_name, attr_name = _EXPORTS[name] + module = importlib.import_module(module_name) + value = getattr(module, attr_name) + globals()[name] = value + return value + + +__all__ = [ + "Experts", + "GatedDeltaNet", + "GQAttention", + "MTPBlock", + "MTPDecoderLayer", + "MTPLossAutoScaler", + "MoEAuxLossAutoScaler", + "MultimodalRotaryEmbedding", + "SigmoidTopKRouter", + "split_grouped_qkvg", + "TokenDispatcher", + "TopKRouter", + "_AllToAll", +] diff --git a/experimental/lite/megatron/lite/primitive/modules/dispatcher.py b/experimental/lite/megatron/lite/primitive/modules/dispatcher.py new file mode 100644 index 00000000000..e18c6ec5e26 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/modules/dispatcher.py @@ -0,0 +1,577 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Token dispatcher: AllToAll and DeepEP dispatch/combine.""" + +from __future__ import annotations + +import os + +import torch # pyright: ignore[reportMissingImports] +import torch.distributed as dist # pyright: ignore[reportMissingImports] + +from megatron.core.transformer.moe.moe_utils import ( # pyright: ignore[reportMissingImports] + permute, + unpermute, +) +from megatron.lite.primitive.modules.moe import _AllToAll +from megatron.lite.primitive.parallel import ParallelState +from megatron.lite.primitive.utils import ensure_divisible + +try: + import deep_ep # pyright: ignore[reportMissingImports] + from deep_ep.utils import EventHandle, EventOverlap # pyright: ignore[reportMissingImports] +except ImportError: + deep_ep = None # type: ignore + EventHandle = None # type: ignore + EventOverlap = None # type: ignore + + +def _hidden_bytes(hidden_size: int) -> int: + return hidden_size * 2 + + +def _build_deepep_buffer(group: dist.ProcessGroup, hidden_size: int): + if deep_ep is None: + raise RuntimeError("DeepEP buffer requested but deep_ep is not installed.") + + group_size = dist.get_world_size(group=group) + hidden_bytes = _hidden_bytes(hidden_size) + num_nvl_bytes = 0 + num_rdma_bytes = 0 + + for config in ( + deep_ep.Buffer.get_dispatch_config(group_size), + deep_ep.Buffer.get_combine_config(group_size), + ): + num_nvl_bytes = max( + config.get_nvl_buffer_size_hint(hidden_bytes, group_size), num_nvl_bytes + ) + num_rdma_bytes = max( + config.get_rdma_buffer_size_hint(hidden_bytes, group_size), num_rdma_bytes + ) + + return deep_ep.Buffer(group=group, num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=num_rdma_bytes) + + +def _use_moe_permute_fusion() -> bool: + return os.environ.get("MEGATRON_LITE_MOE_PERMUTE_FUSION", "0") == "1" + + +def _tensor_hidden_bytes(x: torch.Tensor) -> int: + return x.size(1) * max(x.element_size(), 2) + + +class _DeepEPDispatch(torch.autograd.Function): + @staticmethod + def forward( + ctx, + buffer, + hidden_states: torch.Tensor, + topk_indices: torch.Tensor, + topk_scores: torch.Tensor, + num_experts: int, + async_finish: bool, + allocate_on_comm_stream: bool, + ): + previous_event = ( + EventOverlap(EventHandle()) + if async_finish and EventHandle is not None and EventOverlap is not None + else None + ) + ( + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + is_token_in_rank, + event, + ) = buffer.get_dispatch_layout( + topk_indices, + num_experts=num_experts, + previous_event=previous_event, + async_finish=async_finish, + allocate_on_comm_stream=allocate_on_comm_stream, + ) + (recv_hidden, recv_indices, recv_probs, recv_per_expert, handle, after_event) = ( + buffer.dispatch( + hidden_states.contiguous(), + topk_idx=topk_indices, + topk_weights=topk_scores.float(), + num_tokens_per_rank=num_tokens_per_rank, + num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, + is_token_in_rank=is_token_in_rank, + num_tokens_per_expert=num_tokens_per_expert, + previous_event=event, + async_finish=async_finish, + allocate_on_comm_stream=allocate_on_comm_stream, + ) + ) + if async_finish: + after_event.current_stream_wait() + + ctx.buffer = buffer + ctx.handle = handle + ctx.async_finish = async_finish + ctx.allocate_on_comm_stream = allocate_on_comm_stream + recv_per_expert_tensor = torch.tensor( + recv_per_expert, dtype=torch.int64, device=recv_hidden.device + ) + return recv_hidden, recv_indices, recv_probs, recv_per_expert_tensor, handle + + @staticmethod + def backward( + ctx, grad_recv_hidden, grad_recv_indices, grad_recv_probs, grad_recv_per_expert, grad_handle + ): + del grad_recv_indices, grad_recv_per_expert, grad_handle + previous_event = ( + EventOverlap(EventHandle()) + if ctx.async_finish and EventHandle is not None and EventOverlap is not None + else None + ) + grad_scores = None if grad_recv_probs is None else grad_recv_probs.float() + grad_hidden, grad_topk_scores, after_event = ctx.buffer.combine( + grad_recv_hidden.contiguous(), + ctx.handle, + topk_weights=grad_scores, + previous_event=previous_event, + async_finish=ctx.async_finish, + allocate_on_comm_stream=ctx.allocate_on_comm_stream, + ) + if ctx.async_finish: + after_event.current_stream_wait() + return None, grad_hidden, None, grad_topk_scores, None, None, None + + +class _DeepEPCombine(torch.autograd.Function): + @staticmethod + def forward( + ctx, + buffer, + rank_grouped: torch.Tensor, + handle, + async_finish: bool, + allocate_on_comm_stream: bool, + ): + previous_event = ( + EventOverlap(EventHandle()) + if async_finish and EventHandle is not None and EventOverlap is not None + else None + ) + combined, _, after_event = buffer.combine( + rank_grouped, + handle, + previous_event=previous_event, + async_finish=async_finish, + allocate_on_comm_stream=allocate_on_comm_stream, + ) + if async_finish: + after_event.current_stream_wait() + ctx.buffer = buffer + ctx.handle = handle + ctx.async_finish = async_finish + ctx.allocate_on_comm_stream = allocate_on_comm_stream + return combined + + @staticmethod + def backward(ctx, grad_output): + previous_event = ( + EventOverlap(EventHandle()) + if ctx.async_finish and EventHandle is not None and EventOverlap is not None + else None + ) + grad_rank_grouped, _, _, _, _, after_event = ctx.buffer.dispatch( + grad_output.contiguous(), + handle=ctx.handle, + previous_event=previous_event, + async_finish=ctx.async_finish, + allocate_on_comm_stream=ctx.allocate_on_comm_stream, + ) + if ctx.async_finish: + after_event.current_stream_wait() + return None, grad_rank_grouped, None, None, None + + +class TokenDispatcher: + + def __init__( + self, num_experts: int, hidden_size: int, ps: ParallelState, *, use_deepep: bool = True + ): + self.ps = ps + self.num_experts = num_experts + self.ep_size = ps.ep_size + self.num_local_experts = ensure_divisible(num_experts, ps.ep_size) + + self.use_deepep = use_deepep and deep_ep is not None and ps.ep_size > 1 + if self.use_deepep: + assert ps.tp_ep_group is not None + self.buffer = _build_deepep_buffer(ps.tp_ep_group, hidden_size) + + self._row_id_map: torch.Tensor | None = None + self._restore_shape: tuple | None = None + self._input_splits: list[int] | None = None + self._output_splits: list[int] | None = None + self._handle = None + self._deepep_event = None + + if self.ep_size > 1 and self.num_local_experts > 1: + chunk_idxs = torch.arange(self.ep_size * self.num_local_experts) + self._sort_by_experts = ( + chunk_idxs.reshape(self.ep_size, self.num_local_experts).T.ravel().tolist() + ) + self._restore_by_ranks = ( + chunk_idxs.reshape(self.num_local_experts, self.ep_size).T.ravel().tolist() + ) + + def dispatch( + self, hidden_states: torch.Tensor, topk_scores: torch.Tensor, topk_indices: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + if self.ep_size <= 1: + return self._dispatch_local(hidden_states, topk_scores, topk_indices) + if self.use_deepep: + return self._dispatch_deepep(hidden_states, topk_scores, topk_indices) + dispatched, tpe, sorted_scores = self._dispatch_alltoall( + hidden_states, topk_scores, topk_indices + ) + return dispatched, tpe, sorted_scores + + def combine(self, expert_output: torch.Tensor) -> torch.Tensor: + if self.ep_size <= 1: + return self._combine_local(expert_output) + if self.use_deepep: + return self._combine_deepep(expert_output) + return self._combine_alltoall(expert_output) + + def submit_deepep_combine( + self, expert_output: torch.Tensor, *, allocate_on_comm_stream: bool = False + ): + if not self.use_deepep: + raise RuntimeError("submit_deepep_combine requires DeepEP combine.") + rank_grouped = unpermute( + expert_output, + self._row_id_map, + restore_shape=self._restore_shape, + fused=_use_moe_permute_fusion(), + ) + previous_event = ( + EventOverlap(EventHandle()) + if EventHandle is not None and EventOverlap is not None + else None + ) + combined = self.buffer.combine( + rank_grouped, + self._handle, + previous_event=previous_event, + async_finish=True, + allocate_on_comm_stream=allocate_on_comm_stream, + ) + event = None + if isinstance(combined, tuple): + if len(combined) >= 3: + event = combined[2] + combined = combined[0] + return {"combined": combined, "event": event} + + def finish_deepep_combine(self, state): + if not self.use_deepep: + raise RuntimeError("finish_deepep_combine requires DeepEP combine.") + event = state.get("event") + if event is not None: + event.current_stream_wait() + self._row_id_map = None + self._restore_shape = None + self._handle = None + self._local_tpe_list = None + return state["combined"] + + def _dispatch_local(self, hidden_states, topk_scores, topk_indices): + t, h = hidden_states.shape + e = self.num_experts + + routing_map = torch.zeros(t, e, dtype=torch.bool, device=hidden_states.device) + routing_map.scatter_(1, topk_indices, True) + num_out = int(routing_map.sum().item()) + + probs_2d = torch.zeros(t, e, dtype=topk_scores.dtype, device=hidden_states.device) + probs_2d.scatter_(1, topk_indices, topk_scores) + + permuted, permuted_probs, sorted_indices = permute( + hidden_states, + routing_map, + probs=probs_2d, + num_out_tokens=num_out, + fused=_use_moe_permute_fusion(), + )[:3] + + self._row_id_map = sorted_indices + self._restore_shape = hidden_states.shape + + tokens_per_expert = routing_map.sum(dim=0).to(torch.int64) + return permuted, tokens_per_expert, permuted_probs + + def _combine_local(self, expert_output): + result = unpermute( + expert_output, + self._row_id_map, + restore_shape=self._restore_shape, + fused=_use_moe_permute_fusion(), + ) + self._row_id_map = None + self._restore_shape = None + return result + + def _dispatch_alltoall(self, hidden_states, topk_scores, topk_indices): + t, h = hidden_states.shape + e = self.num_experts + + routing_map = torch.zeros(t, e, dtype=torch.bool, device=hidden_states.device) + routing_map.scatter_(1, topk_indices, True) + num_out = t * topk_indices.size(1) + + probs_2d = torch.zeros(t, e, dtype=topk_scores.dtype, device=hidden_states.device) + probs_2d.scatter_(1, topk_indices, topk_scores) + + permuted, permuted_probs, sorted_indices = permute( + hidden_states, + routing_map, + probs=probs_2d, + num_out_tokens=num_out, + fused=_use_moe_permute_fusion(), + )[:3] + self._row_id_map = sorted_indices + self._restore_shape = hidden_states.shape + + tokens_per_expert = routing_map.sum(dim=0).to(torch.int64) + tpe_by_rank = tokens_per_expert.view(self.ep_size, self.num_local_experts).sum(dim=1) + self._input_splits = tpe_by_rank.tolist() + + global_tpe_flat = tokens_per_expert.new_empty(self.ep_size * e) + dist.all_gather_into_tensor(global_tpe_flat, tokens_per_expert, group=self.ps.ep_group) + global_tpe_2d = global_tpe_flat.view(self.ep_size, e) + ep_rank = dist.get_rank(group=self.ps.ep_group) + my_start = ep_rank * self.num_local_experts + recv_tpe_2d = global_tpe_2d[:, my_start : my_start + self.num_local_experts].contiguous() + self._output_splits = recv_tpe_2d.sum(dim=1).tolist() + + recv_flat = _AllToAll.apply( + permuted, self._input_splits, self._output_splits, self.ps.ep_group + ) + recv_scores = _AllToAll.apply( + permuted_probs.unsqueeze(-1), self._input_splits, self._output_splits, self.ps.ep_group + ) + + if self.num_local_experts > 1: + chunk_sizes = recv_tpe_2d.ravel().tolist() + chunks = torch.split(recv_flat, chunk_sizes, dim=0) + score_chunks = torch.split(recv_scores, chunk_sizes, dim=0) + sort_idxs = self._sort_by_experts + restore_idxs = self._restore_by_ranks + dispatched = torch.cat([chunks[i] for i in sort_idxs], dim=0) + permuted_probs_out = torch.cat([score_chunks[i] for i in sort_idxs], dim=0) + self._combine_chunk_sizes = [chunk_sizes[i] for i in sort_idxs] + self._combine_restore_idxs = restore_idxs + else: + dispatched = recv_flat + permuted_probs_out = recv_scores + self._combine_chunk_sizes = None + self._combine_restore_idxs = None + + recv_tpe = recv_tpe_2d.sum(dim=0) + return dispatched, recv_tpe, permuted_probs_out.squeeze(-1) + + def _combine_alltoall(self, expert_output): + if self._combine_chunk_sizes is not None: + chunks = torch.split(expert_output, self._combine_chunk_sizes, dim=0) + restore_idxs = ( + self._combine_restore_idxs + if self._combine_restore_idxs is not None + else self._restore_by_ranks + ) + rank_grouped = torch.cat([chunks[i] for i in restore_idxs], dim=0) + else: + rank_grouped = expert_output + + combined = _AllToAll.apply( + rank_grouped, self._output_splits, self._input_splits, self.ps.ep_group + ) + result = unpermute( + combined, + self._row_id_map, + restore_shape=self._restore_shape, + fused=_use_moe_permute_fusion(), + ) + self._row_id_map = None + self._restore_shape = None + self._input_splits = None + self._output_splits = None + self._combine_chunk_sizes = None + self._combine_restore_idxs = None + self._local_tpe_list = None + return result + + def submit_deepep_dispatch( + self, hidden_states, topk_scores, topk_indices, *, allocate_on_comm_stream: bool = False + ): + if not self.use_deepep: + raise RuntimeError("submit_deepep_dispatch requires DeepEP dispatch.") + previous_event = ( + EventOverlap(EventHandle()) + if EventHandle is not None and EventOverlap is not None + else None + ) + ( + num_tokens_per_rank, + num_tokens_per_rdma_rank, + num_tokens_per_expert, + is_token_in_rank, + event, + ) = self.buffer.get_dispatch_layout( + topk_indices, + num_experts=self.num_experts, + previous_event=previous_event, + async_finish=True, + allocate_on_comm_stream=allocate_on_comm_stream, + ) + + topk_scores = topk_scores.float() + recv_hidden, recv_indices, recv_probs, recv_per_expert, handle, event = ( + self.buffer.dispatch( + hidden_states, + topk_idx=topk_indices, + topk_weights=topk_scores, + num_tokens_per_rank=num_tokens_per_rank, + num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, + is_token_in_rank=is_token_in_rank, + num_tokens_per_expert=num_tokens_per_expert, + previous_event=event, + async_finish=True, + allocate_on_comm_stream=allocate_on_comm_stream, + ) + ) + return { + "recv_hidden": recv_hidden, + "recv_indices": recv_indices, + "recv_probs": recv_probs, + "recv_per_expert": recv_per_expert, + "handle": handle, + "event": event, + } + + def finish_deepep_dispatch(self, state): + if not self.use_deepep: + raise RuntimeError("finish_deepep_dispatch requires DeepEP dispatch.") + self._handle = state["handle"] + self._deepep_event = state["event"] + self.wait_dispatch_event() + return self._finish_deepep_dispatch( + state["recv_hidden"], + state["recv_indices"], + state["recv_probs"], + state["recv_per_expert"], + ) + + def _finish_deepep_dispatch( + self, + recv_hidden: torch.Tensor, + recv_indices: torch.Tensor, + recv_probs: torch.Tensor, + recv_per_expert, + ): + if isinstance(recv_per_expert, torch.Tensor): + recv_per_expert = [int(x) for x in recv_per_expert.detach().cpu().tolist()] + local_tpe = torch.tensor( + recv_per_expert[: self.num_local_experts], dtype=torch.int64, device=recv_hidden.device + ) + self._local_tpe_list = [int(x) for x in recv_per_expert[: self.num_local_experts]] + rows = recv_hidden.size(0) + recv_indices = recv_indices.to(torch.long) + routing_map = torch.zeros( + rows, self.num_local_experts, dtype=torch.bool, device=recv_hidden.device + ) + probs_2d = torch.zeros( + rows, self.num_local_experts, dtype=recv_probs.dtype, device=recv_hidden.device + ) + valid = recv_indices >= 0 + row_ids = torch.arange(rows, device=recv_hidden.device).unsqueeze(1) + row_ids = row_ids.expand_as(recv_indices)[valid] + expert_ids = recv_indices[valid] + routing_map[row_ids, expert_ids] = True + probs_2d[row_ids, expert_ids] = recv_probs[valid] + num_out = sum(int(x) for x in recv_per_expert) + dispatched, permuted_probs, sorted_indices = permute( + recv_hidden, + routing_map, + probs=probs_2d, + num_out_tokens=num_out, + fused=_use_moe_permute_fusion(), + )[:3] + self._row_id_map = sorted_indices + self._restore_shape = recv_hidden.shape + if os.environ.get("MEGATRON_LITE_DEEPEP_DEBUG_METADATA") == "1": + ep_rank = dist.get_rank(group=self.ps.ep_group) + print( + "[DEEPEP_METADATA] " + f"ep_rank={ep_rank} recv_rows={int(recv_hidden.shape[0])} " + f"expert_rows={int(dispatched.shape[0])} " + f"recv_indices_shape={tuple(recv_indices.shape)} " + f"recv_per_expert_len={len(recv_per_expert)} " + f"recv_per_expert_sum={sum(int(x) for x in recv_per_expert)} " + f"recv_per_expert_head={recv_per_expert[: self.num_local_experts]} " + f"local_tpe_sum={int(local_tpe.sum().item())}", + flush=True, + ) + if os.environ.get("MEGATRON_LITE_DEEPEP_SKIP_DISPATCH_METADATA_CHECK") != "1" and int( + local_tpe.sum().item() + ) != int(dispatched.shape[0]): + ep_rank = dist.get_rank(group=self.ps.ep_group) + raise RuntimeError( + "DeepEP dispatch metadata mismatch: " + f"ep_rank={ep_rank} dispatched_tokens={int(dispatched.shape[0])} " + f"local_tpe={local_tpe.tolist()} recv_per_expert_len={len(recv_per_expert)}" + ) + return dispatched, local_tpe, permuted_probs + + def _dispatch_deepep(self, hidden_states, topk_scores, topk_indices): + if torch.is_grad_enabled(): + recv_hidden, recv_indices, recv_probs, recv_per_expert, handle = _DeepEPDispatch.apply( + self.buffer, + hidden_states, + topk_indices, + topk_scores.float(), + self.num_experts, + False, + False, + ) + self._handle = handle + self._deepep_event = None + return self._finish_deepep_dispatch( + recv_hidden, recv_indices, recv_probs, recv_per_expert + ) + state = self.submit_deepep_dispatch( + hidden_states, topk_scores, topk_indices, allocate_on_comm_stream=False + ) + return self.finish_deepep_dispatch(state) + + def wait_dispatch_event(self): + if self._deepep_event is not None: + self._deepep_event.current_stream_wait() + self._deepep_event = None + + def _combine_deepep(self, expert_output): + rank_grouped = unpermute( + expert_output, + self._row_id_map, + restore_shape=self._restore_shape, + fused=_use_moe_permute_fusion(), + ) + if torch.is_grad_enabled(): + combined = _DeepEPCombine.apply(self.buffer, rank_grouped, self._handle, False, False) + else: + combined = self.buffer.combine(rank_grouped, self._handle) + if isinstance(combined, tuple): + combined = combined[0] + self._row_id_map = None + self._restore_shape = None + self._handle = None + self._local_tpe_list = None + return combined + + +__all__ = ["TokenDispatcher"] diff --git a/experimental/lite/megatron/lite/primitive/modules/experts.py b/experimental/lite/megatron/lite/primitive/modules/experts.py new file mode 100644 index 00000000000..2ab182b9905 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/modules/experts.py @@ -0,0 +1,279 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""MoE expert compute: SwiGLU fusions, _AllReduceETP, and Experts.""" + +from __future__ import annotations + +import os +from contextlib import contextmanager +from typing import Any + +import torch # pyright: ignore[reportMissingImports] +import torch.distributed as dist # pyright: ignore[reportMissingImports] +import torch.nn as nn # pyright: ignore[reportMissingImports] +import torch.nn.functional as F # pyright: ignore[reportMissingImports] +import transformer_engine.pytorch as te # pyright: ignore[reportMissingImports] + +from megatron.lite.primitive.modules.lora import ( + LoraConfig, + SharedGroupedLinearLoRA, + normalize_lora_config, +) +from megatron.lite.primitive.parallel import ParallelState +from megatron.lite.primitive.recompute import CheckpointWithoutOutput +from megatron.lite.primitive.utils import ensure_divisible + +__all__ = ["Experts", "_AllReduceETP"] + + +@contextmanager +def _expert_nvtx_range(name: str): + if os.environ.get("MEGATRON_LITE_EP_EXPERT_NVTX") != "1" or not torch.cuda.is_available(): + yield + return + torch.cuda.nvtx.range_push(name) + try: + yield + finally: + torch.cuda.nvtx.range_pop() + + +@torch.compile +def _swiglu(y): + y_1, y_2 = torch.chunk(y, 2, -1) + return F.silu(y_1) * y_2 + + +@torch.compile +def _weighted_swiglu(y, weights): + dtype = y.dtype + res = _swiglu(y) * weights + return res.to(dtype) + + +@torch.compile +def _swiglu_back(g, y): + y_1, y_2 = torch.chunk(y, 2, -1) + return torch.cat( + (g * torch.sigmoid(y_1) * (1 + y_1 * (1 - torch.sigmoid(y_1))) * y_2, g * F.silu(y_1)), -1 + ) + + +@torch.compile +def _weighted_swiglu_back(g, y, weights): + input_dtype = y.dtype + w_dtype = weights.dtype + input_grad = _swiglu_back(g * weights, y) + weights_grad = _swiglu(y) * g.to(w_dtype) + weights_grad = torch.sum(weights_grad, dim=-1, keepdim=True) + return input_grad.to(input_dtype), weights_grad.to(w_dtype) + + +class _WeightedSwiGLUFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, input, weights, fp8_input_store): + input_for_backward = input.to(torch.float8_e4m3fn) if fp8_input_store else input + ctx.save_for_backward(input_for_backward, weights) + ctx.ori_input_dtype = input.dtype + ctx.fp8_input_store = fp8_input_store + return _weighted_swiglu(input, weights) + + @staticmethod + def backward(ctx, grad_output): + input, weights = ctx.saved_tensors + input = input.to(ctx.ori_input_dtype) if ctx.fp8_input_store else input + tmp, wgrad = _weighted_swiglu_back(grad_output, input, weights) + return tmp, wgrad, None + + +def weighted_bias_swiglu_impl(input, bias, weights, fp8_input_store=False): + """Token-wise-weighted bias swiglu fusion (copied from MC).""" + ori_shape = input.shape + assert len(ori_shape) in [2, 3] + input = input.view(-1, ori_shape[-1]) + if bias is not None: + raise NotImplementedError("Bias is not supported for weighted swiglu fusion") + output = _WeightedSwiGLUFunction.apply(input, weights, fp8_input_store) + return output if len(ori_shape) == 2 else output.view(ori_shape[0], ori_shape[1], -1) + + +def swiglu_with_probs(y: torch.Tensor, probs: torch.Tensor | None) -> torch.Tensor: + """SwiGLU with optional expert probability scaling.""" + if probs is not None: + return weighted_bias_swiglu_impl(y, bias=None, weights=probs) + y1, y2 = torch.chunk(y, 2, -1) + return F.silu(y1) * y2 + + +class _AllReduceETP(torch.autograd.Function): + """AllReduce with proper autograd: grad(AllReduce) = AllReduce.""" + + @staticmethod + def forward(ctx, x, group): + ctx.group = group + dist.all_reduce(x, group=group) + return x + + @staticmethod + def backward(ctx, grad): + return grad, None + + +class Experts(nn.Module): + + def __init__( + self, + config: Any, + ps: ParallelState, + *, + fp8: bool = False, + moe_act_recompute: bool = False, + lora_config: LoraConfig | dict | None = None, + ): + super().__init__() + self.num_local_experts = ensure_divisible(config.num_experts, ps.ep_size) + self.fp8 = fp8 + self.moe_act_recompute = moe_act_recompute + self.etp_group = ps.etp_group if ps.etp_size > 1 else None + + self.fc1 = te.GroupedLinear( + self.num_local_experts, + config.hidden_size, + config.moe_intermediate_size * 2 // ps.etp_size, + bias=False, + params_dtype=torch.bfloat16, + ) + self.fc2 = te.GroupedLinear( + self.num_local_experts, + config.moe_intermediate_size // ps.etp_size, + config.hidden_size, + bias=False, + params_dtype=torch.bfloat16, + ) + lora = normalize_lora_config(lora_config) + self.fc1_lora: SharedGroupedLinearLoRA | None = None + self.fc2_lora: SharedGroupedLinearLoRA | None = None + if lora.enabled and lora.targets_module("linear_fc1"): + self.fc1_lora = SharedGroupedLinearLoRA( + self.num_local_experts, + config.hidden_size, + config.moe_intermediate_size * 2 // ps.etp_size, + lora.rank, + alpha=lora.alpha, + dropout=lora.dropout, + ) + if lora.enabled and lora.targets_module("linear_fc2"): + self.fc2_lora = SharedGroupedLinearLoRA( + self.num_local_experts, + config.moe_intermediate_size // ps.etp_size, + config.hidden_size, + lora.rank, + alpha=lora.alpha, + dropout=lora.dropout, + ) + if ps.tp_size > 1 and ps.ep_size == 1 and ps.etp_size == 1: + tp_group = ps.tp_group + for module in (self.fc1, self.fc2, self.fc1_lora, self.fc2_lora): + if module is None: + continue + for param in module.parameters(): + + def _ar(grad, g=tp_group): + dist.all_reduce(grad, op=dist.ReduceOp.SUM, group=g) + return grad + + param.register_hook(_ar) + + def forward( + self, + x: torch.Tensor, + tokens_per_expert: torch.Tensor, + permuted_probs: torch.Tensor | None = None, + tokens_per_expert_list: list[int] | None = None, + ) -> torch.Tensor: + m_splits = ( + tokens_per_expert.tolist() + if tokens_per_expert_list is None + else list(tokens_per_expert_list) + ) + pad_mask = None + if self.fp8: + x, permuted_probs, m_splits, pad_mask = self._fp8_pad(x, permuted_probs, m_splits) + + etp_real_len = x.shape[0] + if self.etp_group is not None: + max_len = torch.tensor([etp_real_len], device=x.device, dtype=torch.int64) + dist.all_reduce(max_len, op=dist.ReduceOp.MAX, group=self.etp_group) + max_len = int(max_len.item()) + if etp_real_len < max_len: + x = torch.cat( + [ + x, + torch.zeros( + max_len - etp_real_len, x.shape[1], dtype=x.dtype, device=x.device + ), + ], + dim=0, + ) + if permuted_probs is not None: + permuted_probs = torch.cat( + [ + permuted_probs, + torch.zeros( + max_len - etp_real_len, dtype=permuted_probs.dtype, device=x.device + ), + ], + dim=0, + ) + m_splits = list(m_splits) + m_splits[-1] += max_len - etp_real_len + + probs = permuted_probs.unsqueeze(-1) if permuted_probs is not None else None + with _expert_nvtx_range("ep_experts.forward"): + if self.moe_act_recompute and probs is not None: + act_ckpt = CheckpointWithoutOutput(preserve_rng_state=True) + fc1_out = self.fc1(x, m_splits) + if self.fc1_lora is not None: + fc1_out = fc1_out + self.fc1_lora(x, m_splits) + h = act_ckpt.checkpoint(swiglu_with_probs, fc1_out, probs) + out = self.fc2(h, m_splits) + if self.fc2_lora is not None: + out = out + self.fc2_lora(h, m_splits) + act_ckpt.discard_output_and_register_recompute(out) + else: + fc1_out = self.fc1(x, m_splits) + if self.fc1_lora is not None: + fc1_out = fc1_out + self.fc1_lora(x, m_splits) + h = swiglu_with_probs(fc1_out, probs) + out = self.fc2(h, m_splits) + if self.fc2_lora is not None: + out = out + self.fc2_lora(h, m_splits) + + if self.etp_group is not None: + out = _AllReduceETP.apply(out, self.etp_group) + out = out[:etp_real_len] + + if pad_mask is not None: + out = out[pad_mask] + return out + + @staticmethod + def _fp8_pad(x, permuted_probs, m_splits): + padded = [(s + 15) // 16 * 16 for s in m_splits] + if padded == m_splits: + return x, permuted_probs, m_splits, None + device, dtype = x.device, x.dtype + total_padded = sum(padded) + x_pad = torch.zeros(total_padded, x.size(1), device=device, dtype=dtype) + mask = torch.zeros(total_padded, dtype=torch.bool, device=device) + probs_pad = None + if permuted_probs is not None: + probs_pad = torch.zeros(total_padded, device=device, dtype=permuted_probs.dtype) + src_off, dst_off = 0, 0 + for real, pad in zip(m_splits, padded, strict=True): + x_pad[dst_off : dst_off + real] = x[src_off : src_off + real] + mask[dst_off : dst_off + real] = True + if probs_pad is not None: + probs_pad[dst_off : dst_off + real] = permuted_probs[src_off : src_off + real] + src_off += real + dst_off += pad + return x_pad, probs_pad, padded, mask diff --git a/experimental/lite/megatron/lite/primitive/modules/gated_delta_net.py b/experimental/lite/megatron/lite/primitive/modules/gated_delta_net.py new file mode 100644 index 00000000000..87b315f8f0d --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/modules/gated_delta_net.py @@ -0,0 +1,304 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Qwen-style Gated DeltaNet primitive.""" + +from __future__ import annotations + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +import transformer_engine.pytorch as te + +from megatron.core.jit import jit_fuser +from megatron.lite.primitive.ops.gated_delta_rule import l2norm, torch_chunk_gated_delta_rule +from megatron.lite.primitive.parallel import ColumnParallelLinear, ParallelState, RowParallelLinear +from megatron.lite.primitive.parallel.cp import ( + zigzag_reconstruct_from_cp_parts, + zigzag_slice_for_cp, +) +from megatron.lite.primitive.parallel.thd import ( + reconstruct_packed_from_cp_parts, + split_packed_to_cp_local, +) +from megatron.lite.primitive.utils import ensure_divisible + +try: + from fla.modules.convolution import ( + causal_conv1d as _fla_causal_conv1d, # pyright: ignore[reportMissingImports] + ) + from fla.ops.gated_delta_rule import ( + chunk_gated_delta_rule as _fla_chunk_gated_delta_rule, # pyright: ignore[reportMissingImports] + ) + + _HAS_FLA = True +except ImportError: + _HAS_FLA = False + + +class GatedDeltaNet(nn.Module): + """Native Gated DeltaNet with dense/packed CP reconstruction.""" + + def __init__( + self, + *, + hidden_size: int, + linear_num_key_heads: int, + linear_key_head_dim: int, + linear_num_value_heads: int, + linear_value_head_dim: int, + linear_conv_kernel_dim: int, + rms_norm_eps: float, + ps: ParallelState, + deterministic: bool = False, + ): + super().__init__() + self.ps = ps + self.deterministic = bool(deterministic) + self.num_k_heads = linear_num_key_heads + self.num_v_heads = linear_num_value_heads + self.dk = linear_key_head_dim + self.dv = linear_value_head_dim + self.v_heads_per_k_head = ensure_divisible(self.num_v_heads, self.num_k_heads) + self.num_k_heads_local = ensure_divisible(self.num_k_heads, ps.tp_size) + self.num_v_heads_local = ensure_divisible(self.num_v_heads, ps.tp_size) + self.qk_dim = self.num_k_heads * self.dk + self.v_dim = self.num_v_heads * self.dv + self.qk_dim_local = self.num_k_heads_local * self.dk + self.v_dim_local = self.num_v_heads_local * self.dv + self.in_proj_dim = self.qk_dim * 2 + self.v_dim * 2 + self.num_v_heads * 2 + + self.in_proj = ColumnParallelLinear( + hidden_size, + self.in_proj_dim, + ps, + bias=False, + normalization="RMSNorm", + eps=rms_norm_eps, + zero_centered_gamma=True, + ) + conv_dim_local = self.qk_dim_local * 2 + self.v_dim_local + self.conv1d = nn.Conv1d( + in_channels=conv_dim_local, + out_channels=conv_dim_local, + kernel_size=linear_conv_kernel_dim, + groups=conv_dim_local, + bias=False, + padding=linear_conv_kernel_dim - 1, + ) + self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads_local, dtype=torch.float32)) + self.A_log = nn.Parameter(torch.zeros(self.num_v_heads_local, dtype=torch.float32)) + self.norm = te.RMSNorm(self.dv, eps=rms_norm_eps, zero_centered_gamma=True) + self.o_proj = RowParallelLinear(self.v_dim, hidden_size, ps, bias=False) + + def forward( + self, x: torch.Tensor, position_ids: torch.Tensor | None = None, packed_seq_params=None + ) -> torch.Tensor: + del position_ids + qkvzba = self.in_proj(x).transpose(0, 1).contiguous() + cp_restore = None + if self.ps.cp_size > 1: + qkvzba, cp_restore = self._gather_cp_qkvzba(qkvzba, packed_seq_params) + batch, seq_len = qkvzba.shape[:2] + query, key, value, gate, beta, alpha = self._split_proj(qkvzba) + qkv = torch.cat( + [ + query.reshape(batch, seq_len, -1), + key.reshape(batch, seq_len, -1), + value.reshape(batch, seq_len, -1), + ], + dim=-1, + ) + + cu_seqlens = None + if packed_seq_params is not None: + cu_seqlens = ( + packed_seq_params.cu_seqlens_q_padded + if getattr(packed_seq_params, "cu_seqlens_q_padded", None) is not None + else packed_seq_params.cu_seqlens_q + ) + if not _HAS_FLA: + raise NotImplementedError("GatedDeltaNet packed THD requires FLA kernels.") + + qkv = self._causal_conv1d(qkv, seq_len, cu_seqlens=cu_seqlens) + query, key, value, gate, beta, alpha = self._prepare_qkv( + qkv, gate, beta, alpha, batch, seq_len + ) + g, beta = self._compute_g_and_beta(self.A_log, self.dt_bias, alpha, beta) + out, _ = self._gated_delta_rule( + query, + key, + value, + g, + beta, + initial_state=None, + output_final_state=False, + cu_seqlens=cu_seqlens, + ) + + if cp_restore is not None: + out = self._slice_cp_output(out, cp_restore) + gate = self._slice_cp_output(gate, cp_restore) + batch, seq_len = out.shape[:2] + out = self._apply_gated_norm(out, gate) + out = out.reshape(batch, seq_len, self.v_dim_local).transpose(0, 1).contiguous() + return self.o_proj(out) + + def _all_gather_cp(self, tensor: torch.Tensor) -> list[torch.Tensor]: + if self.ps.cp_group is None: + raise RuntimeError("CP>1 requires ParallelState.cp_group.") + try: + from torch.distributed.nn.functional import all_gather + + return list(all_gather(tensor, group=self.ps.cp_group)) + except Exception: + parts = [torch.empty_like(tensor) for _ in range(self.ps.cp_size)] + dist.all_gather(parts, tensor, group=self.ps.cp_group) + return parts + + def _gather_cp_qkvzba(self, qkvzba: torch.Tensor, packed_seq_params): + parts = self._all_gather_cp(qkvzba) + if packed_seq_params is not None: + cu_seqlens = self._packed_cu_seqlens(packed_seq_params) + full = reconstruct_packed_from_cp_parts( + parts, cu_seqlens_padded=cu_seqlens, cp_size=self.ps.cp_size, dim=1 + ) + return full, ("packed", cu_seqlens) + full = zigzag_reconstruct_from_cp_parts(parts, seq_dim=1) + return full, ("dense",) + + def _slice_cp_output(self, out: torch.Tensor, cp_restore) -> torch.Tensor: + kind = cp_restore[0] + if kind == "packed": + return split_packed_to_cp_local( + out, + cu_seqlens_padded=cp_restore[1], + cp_size=self.ps.cp_size, + cp_rank=self.ps.cp_rank, + dim=1, + ) + if kind == "dense": + return zigzag_slice_for_cp(out, self.ps.cp_rank, self.ps.cp_size, seq_dim=1) + raise RuntimeError(f"Unknown CP restore kind: {kind!r}") + + def _causal_conv1d( + self, qkv: torch.Tensor, seq_len: int, *, cu_seqlens: torch.Tensor | None + ) -> torch.Tensor: + if cu_seqlens is None: + qkv_t = qkv.transpose(1, 2).contiguous() + return F.silu(self.conv1d(qkv_t)[:, :, :seq_len].transpose(1, 2)) + if _HAS_FLA: + qkv, _ = _fla_causal_conv1d( + x=qkv, + weight=self.conv1d.weight.squeeze(1), + bias=None, + activation="silu", + cu_seqlens=cu_seqlens, + ) + return qkv + raise NotImplementedError("GatedDeltaNet packed THD requires FLA causal conv.") + + def _gated_delta_rule( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + initial_state: torch.Tensor | None, + output_final_state: bool, + cu_seqlens: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if _HAS_FLA and not self.deterministic: + return _fla_chunk_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=False, + cu_seqlens=cu_seqlens, + ) + return torch_chunk_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=output_final_state, + ) + + @staticmethod + def _packed_cu_seqlens(packed_seq_params) -> torch.Tensor: + cu_seqlens = ( + packed_seq_params.cu_seqlens_q_padded + if getattr(packed_seq_params, "cu_seqlens_q_padded", None) is not None + else packed_seq_params.cu_seqlens_q + ) + if cu_seqlens is None: + raise ValueError("packed_seq_params must carry cu_seqlens_q for CP GatedDeltaNet.") + return cu_seqlens + + def _split_proj(self, qkvzba: torch.Tensor): + q, k, v, z, b, a = qkvzba.split( + [ + self.qk_dim_local, + self.qk_dim_local, + self.v_dim_local, + self.v_dim_local, + self.num_v_heads_local, + self.num_v_heads_local, + ], + dim=-1, + ) + batch, seq_len = qkvzba.shape[:2] + return ( + q.reshape(batch, seq_len, self.num_k_heads_local, self.dk), + k.reshape(batch, seq_len, self.num_k_heads_local, self.dk), + v.reshape(batch, seq_len, self.num_v_heads_local, self.dv), + z.reshape(batch, seq_len, self.num_v_heads_local, self.dv), + b.reshape(batch, seq_len, self.num_v_heads_local), + a.reshape(batch, seq_len, self.num_v_heads_local), + ) + + def _prepare_qkv(self, qkv: torch.Tensor, gate, beta, alpha, batch: int, seq_len: int): + query_key, value = qkv.split([2 * self.qk_dim_local, self.v_dim_local], dim=-1) + query_key = query_key.reshape(batch, seq_len, 2 * self.num_k_heads_local, self.dk) + value = value.reshape(batch, seq_len, self.num_v_heads_local, self.dv) + query, key = query_key.split(self.num_k_heads_local, dim=2) + query = self._l2norm(query.contiguous()) + key = self._l2norm(key.contiguous()) + if self.v_heads_per_k_head > 1: + query = query.repeat_interleave(self.v_heads_per_k_head, dim=2) + key = key.repeat_interleave(self.v_heads_per_k_head, dim=2) + return ( + query.contiguous(), + key.contiguous(), + value.contiguous(), + gate.contiguous(), + beta.contiguous(), + alpha.contiguous(), + ) + + def _l2norm(self, x: torch.Tensor) -> torch.Tensor: + return l2norm(x) + + @staticmethod + def _compute_g_and_beta(A_log, dt_bias, alpha, beta): + g = -A_log.exp() * F.softplus(alpha.float() + dt_bias) + return g, beta.sigmoid() + + @jit_fuser + def _apply_gated_norm(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + x_dtype = x.dtype + x = x.reshape(-1, x.shape[-1]) + y = self.norm(x) + gate = gate.reshape(-1, gate.shape[-1]) + y = y * F.silu(gate.float()) + return y.to(x_dtype) + + +__all__ = ["GatedDeltaNet"] diff --git a/experimental/lite/megatron/lite/primitive/modules/gqa.py b/experimental/lite/megatron/lite/primitive/modules/gqa.py new file mode 100644 index 00000000000..3618379591e --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/modules/gqa.py @@ -0,0 +1,357 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Grouped Query Attention + Rotary Embedding (MC-atoms). + +Model-agnostic: takes explicit params instead of model-specific config. +Supports sequence parallel, context parallel, and THD (packed sequences). + +RoPE internals call Megatron-Core's atomic +`megatron.core.models.common.embeddings.rotary_pos_embedding.RotaryEmbedding` +and `rope_utils._apply_rotary_pos_emb_bshd / _apply_rotary_pos_emb_thd` — this +matches Megatron-Core's unfused rotate-half path because +``config.apply_rope_fusion`` defaults to ``False``. See +`docs/gqa_mc_atoms_plan.md` section 4 (Option A). +""" + +from __future__ import annotations + +import inspect + +import torch +import torch.nn as nn +import transformer_engine.pytorch as te + +from megatron.core.models.common.embeddings.rope_utils import ( # pyright: ignore[reportMissingImports] + _apply_rotary_pos_emb_bshd, + _apply_rotary_pos_emb_thd, +) +from megatron.core.models.common.embeddings.rotary_pos_embedding import ( + RotaryEmbedding as MCoreRotaryEmbedding, # pyright: ignore[reportMissingImports] +) +from megatron.lite.primitive.modules.gqa_utils import split_grouped_qkvg +from megatron.lite.primitive.modules.lora import LinearLoRA, LoraConfig, normalize_lora_config +from megatron.lite.primitive.modules.mrope import MultimodalRotaryEmbedding +from megatron.lite.primitive.parallel import ColumnParallelLinear, ParallelState, RowParallelLinear +from megatron.lite.primitive.utils import ensure_divisible + +# Whitelist of MC PackedSeqParams fields accepted by TE DotProductAttention.forward(). +# MC-only fields (local_cp_size, cp_group, total_tokens, seq_idx) are excluded. +# Mirror MC TEDotProductAttention.kept_packed_seq_params pattern +# (Megatron-LM/megatron/core/extensions/transformer_engine.py:1501-1593). +_KEPT_PSP_FIELDS = ( + "qkv_format", + "cu_seqlens_q", + "cu_seqlens_kv", + "cu_seqlens_q_padded", + "cu_seqlens_kv_padded", + "max_seqlen_q", + "max_seqlen_kv", +) + + +def _callable_accepts_kwarg(fn, kwarg: str) -> bool: + try: + parameters = inspect.signature(fn).parameters.values() + except (TypeError, ValueError): + return False + return any( + param.kind is inspect.Parameter.VAR_KEYWORD or param.name == kwarg for param in parameters + ) + + +class GQAttention(nn.Module): + """Grouped Query Attention with TE DotProductAttention. + + Model-agnostic: all architecture params passed explicitly. + """ + + _cp_stream: torch.cuda.Stream | None = None + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + ps: ParallelState, + *, + rms_norm_eps: float = 1e-6, + rope_theta: float = 1_000_000.0, + rotary_percent: float = 1.0, + use_thd: bool = False, + output_gate: bool = False, + use_fp32_rope: bool = False, + zero_centered_gamma: bool = False, + qkv_layout: str = "flat", + lora_config: LoraConfig | dict | None = None, + mrope_section: list[int] | None = None, + ): + super().__init__() + self.num_heads_local = ensure_divisible(num_attention_heads, ps.tp_size) + self.num_kv_heads_local = ensure_divisible(num_key_value_heads, ps.tp_size) + self.head_dim = head_dim + self.ps = ps + self._output_gate = output_gate + self._use_fp32_rope = use_fp32_rope + self._qkv_eps = rms_norm_eps + self._qkv_zero_centered_gamma = zero_centered_gamma + if qkv_layout not in {"flat", "mcore"}: + raise ValueError(f"Unsupported qkv_layout={qkv_layout!r}") + self._qkv_layout = qkv_layout + self._mrope_section = list(mrope_section) if mrope_section is not None else None + + # Declaration order follows MC's `SelfAttention` submodule order + # (linear_proj → linear_qkv → q_layernorm → k_layernorm). `named_ + # parameters()` iterates in registration order, and MC's + # `DistributedDataParallel` partitions gradient buckets by that order. + # Mismatched order would put bucket boundaries in different places, + # producing different per-rank fp32 master shard layouts and + # non-bitwise step-1 divergence. + self.proj = RowParallelLinear(num_attention_heads * head_dim, hidden_size, ps, bias=False) + q_cols = num_attention_heads * (2 if output_gate else 1) + qkv_size = (q_cols + 2 * num_key_value_heads) * head_dim + self.qkv = ColumnParallelLinear( + hidden_size, + qkv_size, + ps, + bias=False, + normalization="RMSNorm", + eps=rms_norm_eps, + zero_centered_gamma=zero_centered_gamma, + ) + self.q_norm = te.RMSNorm( + head_dim, eps=rms_norm_eps, zero_centered_gamma=zero_centered_gamma + ) + self.k_norm = te.RMSNorm( + head_dim, eps=rms_norm_eps, zero_centered_gamma=zero_centered_gamma + ) + + lora = normalize_lora_config(lora_config) + self.qkv_lora: LinearLoRA | None = None + self.proj_lora: LinearLoRA | None = None + if lora.enabled and lora.targets_module("linear_qkv"): + self.qkv_lora = LinearLoRA( + hidden_size, + self.qkv.local_out, + lora.rank, + alpha=lora.alpha, + dropout=lora.dropout, + sequence_parallel_input=self.qkv.use_sp, + tp_group=ps.tp_group, + rank_partition_size=ps.tp_size, + rank_partitioned_a=ps.tp_size > 1, + a_tensor_model_parallel=ps.tp_size > 1, + b_tensor_model_parallel=ps.tp_size > 1, + ) + if lora.enabled and lora.targets_module("linear_proj"): + self.proj_lora = LinearLoRA( + self.proj.local_in, + hidden_size, + lora.rank, + alpha=lora.alpha, + dropout=lora.dropout, + tp_group=ps.tp_group, + tp_rank=ps.tp_rank, + sequence_parallel_scatter_output=self.proj.use_sp, + input_parallel_reduce=ps.tp_size > 1, + output_partition_size=ps.tp_size, + output_partitioned_b=ps.tp_size > 1, + a_tensor_model_parallel=ps.tp_size > 1, + b_tensor_model_parallel=ps.tp_size > 1, + ) + + if self._mrope_section is None: + # MC's RotaryEmbedding is atomic (flat kwargs, no TransformerConfig). + # cp_group is read from self.cp_group inside forward() when not passed + # — no manual CP shard needed on our side. + self.rotary = MCoreRotaryEmbedding( + kv_channels=head_dim, + rotary_percent=rotary_percent, + rotary_interleaved=False, + rotary_base=int(rope_theta), + use_cpu_initialization=False, + cp_group=ps.cp_group if ps.cp_size > 1 else None, + ) + else: + self.rotary = MultimodalRotaryEmbedding( + kv_channels=head_dim, + rotary_percent=rotary_percent, + rotary_base=rope_theta, + cp_group=ps.cp_group if ps.cp_size > 1 else None, + ) + self._rotary_accepts_packed_seq = self._mrope_section is None and _callable_accepts_kwarg( + self.rotary.forward, "packed_seq" + ) + + cp_kwargs = {} + if ps.cp_size > 1: + if GQAttention._cp_stream is None: + GQAttention._cp_stream = torch.cuda.Stream() + cp_kwargs = dict( + cp_group=ps.cp_group, + cp_global_ranks=ps.cp_global_ranks, + cp_stream=GQAttention._cp_stream, + ) + + self.core_attn = te.DotProductAttention( + num_attention_heads=self.num_heads_local, + kv_channels=head_dim, + num_gqa_groups=self.num_kv_heads_local, + attention_dropout=0.0, + attn_mask_type="causal", + qkv_format="thd" if use_thd else "sbhd", + **cp_kwargs, + ) + + def forward( + self, x: torch.Tensor, position_ids: torch.Tensor | None = None, packed_seq_params=None + ) -> torch.Tensor: + qkv = self.qkv(x) + if self.qkv_lora is not None: + qkv = qkv + self.qkv_lora(self._qkv_lora_input(x)) + q, gate, k, v = self._split_qkv(qkv) + + is_thd = packed_seq_params is not None + if is_thd: + q, k, v = q.squeeze(1), k.squeeze(1), v.squeeze(1) + + q = self.q_norm(q) + k = self.k_norm(k) + + # RoPE — unfused bshd/thd to match MC's default apply_rope_fusion=False. + # MC's RotaryEmbedding.forward takes only `max_seq_len` + optional + # `offset`; position_ids is NOT consumed here (MC handles position via + # offset for inference / mRoPE via a separate class). + if self._use_fp32_rope: + orig_dtype = q.dtype + q, k = q.float(), k.float() + if self._mrope_section is not None: + if position_ids is None: + raise ValueError("MRoPE attention requires position_ids.") + # For MRoPE the packed THD path applies RoPE directly through the + # bshd helper, so the rotary module must slice freqs for this CP + # rank before q/k are rotated. + freqs = self.rotary(position_ids, self._mrope_section, packed_seq=False) + if is_thd: + q = _apply_rotary_pos_emb_bshd(q[:, None], freqs).squeeze(1) + k = _apply_rotary_pos_emb_bshd(k[:, None], freqs).squeeze(1) + else: + q = _apply_rotary_pos_emb_bshd(q, freqs) + k = _apply_rotary_pos_emb_bshd(k, freqs) + elif is_thd: + max_q = getattr(packed_seq_params, "max_seqlen_q", None) + max_kv = getattr(packed_seq_params, "max_seqlen_kv", None) + if max_q is None or max_kv is None: + seq_len_for_rope = int(packed_seq_params.cu_seqlens_q[-1]) + else: + seq_len_for_rope = int(max(max_q, max_kv)) + # Match MC RotaryEmbedding.get_rotary_seq_len for packed THD: the + # rotary length is the max per-sequence padded length, not total + # packed tokens. Using total tokens makes rope_utils switch to + # offset mapping, so later packed sequences do not restart at pos 0. + # + # MC contract (gpt_model.py:380-381): THD path passes packed_seq=True so the + # rotary skips its internal cp-slice; _apply_rotary_pos_emb_thd does the + # cp-zigzag slice itself via _get_thd_freqs_on_this_cp_rank. Older MC + # runtimes do not expose this kwarg; callers without context + # parallelism can use the legacy call shape. + if self._rotary_accepts_packed_seq: + freqs = self.rotary(seq_len_for_rope, packed_seq=True) + else: + freqs = self.rotary(seq_len_for_rope) + q = _apply_rotary_pos_emb_thd( + q, + packed_seq_params.cu_seqlens_q, + freqs, + rotary_interleaved=False, + mscale=1.0, + cp_group=self.ps.cp_group, + ) + k = _apply_rotary_pos_emb_thd( + k, + packed_seq_params.cu_seqlens_kv, + freqs, + rotary_interleaved=False, + mscale=1.0, + cp_group=self.ps.cp_group, + ) + else: + # q is CP-zigzag pre-sliced; rotary needs FULL seq len, + # its internal get_pos_emb_on_this_cp_rank re-slices to local len. + local_seq_len = q.size(0) + seq_len_for_rope = local_seq_len * self.ps.cp_size + freqs = self.rotary(seq_len_for_rope) + q = _apply_rotary_pos_emb_bshd(q, freqs, rotary_interleaved=False, mscale=1.0) + k = _apply_rotary_pos_emb_bshd(k, freqs, rotary_interleaved=False, mscale=1.0) + if self._use_fp32_rope: + q, k = q.to(orig_dtype), k.to(orig_dtype) + + if is_thd: + psp_kwargs = { + k: getattr(packed_seq_params, k) + for k in _KEPT_PSP_FIELDS + if getattr(packed_seq_params, k, None) is not None + } + attn_out = self.core_attn( + q, + k, + v, + core_attention_bias_type="no_bias", + attn_mask_type="padding_causal", + **psp_kwargs, + ) + attn_out = attn_out.reshape(attn_out.size(0), 1, -1) + else: + attn_out = self.core_attn(q, k, v, core_attention_bias_type="no_bias") + if attn_out.dim() > x.dim(): + shape = attn_out.shape + attn_out = attn_out.reshape(*shape[:-2], self.num_heads_local * self.head_dim) + + if gate is not None: + gate_fp32 = gate.reshape(attn_out.shape).float().sigmoid() + attn_out = (attn_out.float() * gate_fp32).to(attn_out.dtype) + output = self.proj(attn_out) + if self.proj_lora is not None: + output = output + self.proj_lora(attn_out) + return output + + def _qkv_lora_input(self, x: torch.Tensor) -> torch.Tensor: + linear = self.qkv.linear + if not hasattr(linear, "layer_norm_weight"): + return x + weight = linear.layer_norm_weight + if self._qkv_zero_centered_gamma: + weight = weight + 1 + variance = x.float().pow(2).mean(dim=-1, keepdim=True) + x_norm = x.float() * torch.rsqrt(variance + self._qkv_eps) + return (x_norm * weight.float()).to(x.dtype) + + def _split_qkv(self, qkv: torch.Tensor): + nq, nkv, hd = self.num_heads_local, self.num_kv_heads_local, self.head_dim + lead = qkv.shape[:-1] + if self._qkv_layout == "mcore": + q_per_group = ensure_divisible(nq, nkv) + if self._output_gate: + return split_grouped_qkvg(qkv, num_heads=nq, num_kv_heads=nkv, head_dim=hd) + + qkv = qkv.view(*lead, nkv, (q_per_group + 2) * hd) + q = qkv[..., : q_per_group * hd].reshape(*lead, nq, hd) + k = qkv[..., q_per_group * hd : (q_per_group + 1) * hd] + v = qkv[..., (q_per_group + 1) * hd : (q_per_group + 2) * hd] + return q, None, k, v + + if self._output_gate: + q_block = qkv[..., : nq * 2 * hd].reshape(*lead, nq, 2 * hd) + kv_block = qkv[..., nq * 2 * hd :].reshape(*lead, 2 * nkv, hd) + q = q_block[..., :hd] + gate = q_block[..., hd:] + k = kv_block[..., :nkv, :] + v = kv_block[..., nkv:, :] + return q, gate, k, v + qkv = qkv.view(*lead, nq + 2 * nkv, hd) + # Match MCore SelfAttention's split path: keep q/k/v as views instead + # of inserting copy nodes, since those copy nodes alter backward + # accumulation at the qkv boundary under TP/CP. + q = qkv[..., :nq, :] + k = qkv[..., nq : nq + nkv, :] + v = qkv[..., nq + nkv :, :] + return q, None, k, v diff --git a/experimental/lite/megatron/lite/primitive/modules/gqa_utils.py b/experimental/lite/megatron/lite/primitive/modules/gqa_utils.py new file mode 100644 index 00000000000..2ef329b2c20 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/modules/gqa_utils.py @@ -0,0 +1,29 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Pure grouped-query attention helpers.""" + +from __future__ import annotations + +import torch + +from megatron.lite.primitive.utils import ensure_divisible + + +def split_grouped_qkvg( + qkv: torch.Tensor, *, num_heads: int, num_kv_heads: int, head_dim: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + lead = qkv.shape[:-1] + q_heads_per_group = ensure_divisible(num_heads, num_kv_heads) + group_width = (2 * q_heads_per_group + 2) * head_dim + grouped = qkv.reshape(*lead, num_kv_heads, group_width) + query, gate, key, value = grouped.split( + [q_heads_per_group * head_dim, q_heads_per_group * head_dim, head_dim, head_dim], dim=-1 + ) + return ( + query.reshape(*lead, num_heads, head_dim), + gate.reshape(*lead, num_heads, head_dim), + key.reshape(*lead, num_kv_heads, head_dim), + value.reshape(*lead, num_kv_heads, head_dim), + ) + + +__all__ = ["split_grouped_qkvg"] diff --git a/experimental/lite/megatron/lite/primitive/modules/lora.py b/experimental/lite/megatron/lite/primitive/modules/lora.py new file mode 100644 index 00000000000..943e5066eb5 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/modules/lora.py @@ -0,0 +1,556 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""LoRA helpers for Megatron Lite native model implementations. + +This module is intentionally narrow: it supports the Qwen3-MoE lite path's +Megatron-style sharded linear surfaces, not arbitrary PEFT injection. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F + +_DEFAULT_TARGET_MODULES = ("linear_qkv", "linear_proj", "linear_fc1", "linear_fc2") +_TARGET_ALIASES = { + "qkv": "linear_qkv", + "proj": "linear_proj", + "fc1": "linear_fc1", + "fc2": "linear_fc2", +} + + +@dataclass(frozen=True) +class LoraConfig: + rank: int = 0 + alpha: int | None = None + dropout: float = 0.0 + target_modules: tuple[str, ...] = field(default_factory=lambda: _DEFAULT_TARGET_MODULES) + + @property + def enabled(self) -> bool: + return self.rank > 0 + + @property + def scale(self) -> float: + return float(self.rank if self.alpha is None else self.alpha) / float(self.rank) + + def targets(self) -> set[str]: + out = set() + for target in self.target_modules: + out.add(_TARGET_ALIASES.get(target, target)) + return out + + def targets_module(self, name: str) -> bool: + canonical = _TARGET_ALIASES.get(name, name) + return canonical in self.targets() + + +def normalize_lora_config(config: LoraConfig | dict[str, Any] | None) -> LoraConfig: + if config is None: + return LoraConfig() + if isinstance(config, LoraConfig): + return config + if not isinstance(config, dict): + raise TypeError(f"LoRA config must be LoraConfig, dict, or None, got {type(config)!r}.") + values = dict(config) + enabled = values.pop("enabled", None) + if enabled is False: + values["rank"] = 0 + if "targets" in values and "target_modules" not in values: + values["target_modules"] = values.pop("targets") + else: + values.pop("targets", None) + if "target_modules" in values and not isinstance(values["target_modules"], tuple): + values["target_modules"] = tuple(values["target_modules"]) + return LoraConfig(**values) + + +def freeze_non_lora_params(model: nn.Module) -> dict[str, int]: + """Freeze base parameters and leave adapter parameters trainable.""" + + lora_tensors = 0 + lora_numel = 0 + frozen_tensors = 0 + frozen_numel = 0 + for name, param in model.named_parameters(): + if "lora" in name.lower() or "adapter" in name.lower(): + param.requires_grad_(True) + lora_tensors += 1 + lora_numel += param.numel() + else: + param.requires_grad_(False) + frozen_tensors += 1 + frozen_numel += param.numel() + return { + "lora_tensors": lora_tensors, + "lora_numel": lora_numel, + "frozen_tensors": frozen_tensors, + "frozen_numel": frozen_numel, + } + + +def trainable_param_stats(model: nn.Module) -> dict[str, int]: + tensors = 0 + numel = 0 + for param in model.parameters(): + if param.requires_grad: + tensors += 1 + numel += param.numel() + return {"trainable_tensors": tensors, "trainable_numel": numel} + + +def _gather_sequence_parallel(x: torch.Tensor, group) -> torch.Tensor: + if group is None or dist.get_world_size(group) == 1: + return x + return _AllGatherSequence.apply(x, group) + + +def _reduce_scatter_sequence_parallel(x: torch.Tensor, group) -> torch.Tensor: + if group is None or dist.get_world_size(group) == 1: + return x + return _ReduceScatterSequence.apply(x, group) + + +def _scatter_sequence_parallel(x: torch.Tensor, group, group_rank: int) -> torch.Tensor: + if group is None or dist.get_world_size(group) == 1: + return x + return _ScatterSequence.apply(x, group, group_rank) + + +def _all_reduce_sum(x: torch.Tensor, group) -> torch.Tensor: + if group is None or dist.get_world_size(group) == 1: + return x + return _AllReduceSum.apply(x, group) + + +class _AllGatherSequence(torch.autograd.Function): + @staticmethod + def forward(ctx, x: torch.Tensor, group) -> torch.Tensor: + world_size = dist.get_world_size(group) + ctx.group = group + ctx.local_seq = x.shape[0] + out = torch.empty((x.shape[0] * world_size, *x.shape[1:]), dtype=x.dtype, device=x.device) + dist.all_gather_into_tensor(out, x.contiguous(), group=group) + return out + + @staticmethod + def backward(ctx, grad: torch.Tensor): + out = torch.empty((ctx.local_seq, *grad.shape[1:]), dtype=grad.dtype, device=grad.device) + dist.reduce_scatter_tensor(out, grad.contiguous(), group=ctx.group) + return out, None + + +class _ReduceScatterSequence(torch.autograd.Function): + @staticmethod + def forward(ctx, x: torch.Tensor, group) -> torch.Tensor: + world_size = dist.get_world_size(group) + if x.shape[0] % world_size != 0: + raise ValueError( + f"Cannot reduce-scatter sequence dim {x.shape[0]} over TP={world_size}." + ) + ctx.group = group + ctx.world_size = world_size + out = torch.empty((x.shape[0] // world_size, *x.shape[1:]), dtype=x.dtype, device=x.device) + dist.reduce_scatter_tensor(out, x.contiguous(), group=group) + return out + + @staticmethod + def backward(ctx, grad: torch.Tensor): + out = torch.empty( + (grad.shape[0] * ctx.world_size, *grad.shape[1:]), dtype=grad.dtype, device=grad.device + ) + dist.all_gather_into_tensor(out, grad.contiguous(), group=ctx.group) + return out, None + + +class _ScatterSequence(torch.autograd.Function): + @staticmethod + def forward(ctx, x: torch.Tensor, group, group_rank: int) -> torch.Tensor: + world_size = dist.get_world_size(group) + if x.shape[0] % world_size != 0: + raise ValueError(f"Cannot scatter sequence dim {x.shape[0]} over TP={world_size}.") + ctx.group = group + ctx.world_size = world_size + local_seq = x.shape[0] // world_size + start = int(group_rank) * local_seq + return x[start : start + local_seq].contiguous() + + @staticmethod + def backward(ctx, grad: torch.Tensor): + out = torch.empty( + (grad.shape[0] * ctx.world_size, *grad.shape[1:]), dtype=grad.dtype, device=grad.device + ) + dist.all_gather_into_tensor(out, grad.contiguous(), group=ctx.group) + return out, None, None + + +class _AllReduceSum(torch.autograd.Function): + @staticmethod + def forward(ctx, x: torch.Tensor, group) -> torch.Tensor: + ctx.group = group + out = x.contiguous() + dist.all_reduce(out, op=dist.ReduceOp.SUM, group=group) + return out + + @staticmethod + def backward(ctx, grad: torch.Tensor): + out = grad.contiguous() + dist.all_reduce(out, op=dist.ReduceOp.SUM, group=ctx.group) + return out, None + + +def _all_gather_last_dim(x: torch.Tensor, group, *, reduce_backward: bool = False) -> torch.Tensor: + if group is None or dist.get_world_size(group) == 1: + return x + return _AllGatherLastDim.apply(x, group, reduce_backward) + + +class _AllGatherLastDim(torch.autograd.Function): + """All-gather last dim with Megatron tensor-parallel split backward.""" + + @staticmethod + def forward(ctx, x: torch.Tensor, group, reduce_backward: bool) -> torch.Tensor: + world_size = dist.get_world_size(group) + ctx.group = group + ctx.local_width = x.shape[-1] + ctx.group_rank = dist.get_rank(group) + ctx.reduce_backward = bool(reduce_backward) + flat = x.movedim(-1, 0).contiguous().view(ctx.local_width, -1) + gathered = torch.empty( + (ctx.local_width * world_size, flat.shape[1]), dtype=x.dtype, device=x.device + ) + dist.all_gather_into_tensor(gathered, flat, group=group) + return ( + gathered.view(ctx.local_width * world_size, *x.shape[:-1]).movedim(0, -1).contiguous() + ) + + @staticmethod + def backward(ctx, grad: torch.Tensor): + flat = grad.movedim(-1, 0).contiguous().view(grad.shape[-1], -1) + start = ctx.group_rank * ctx.local_width + out = flat.narrow(0, start, ctx.local_width).contiguous() + if ctx.reduce_backward: + dist.all_reduce(out, op=dist.ReduceOp.SUM, group=ctx.group) + return out.view(ctx.local_width, *grad.shape[:-1]).movedim(0, -1).contiguous(), None, None + + +class _SequenceParallelRankPartitionedLoRA(torch.autograd.Function): + """QKV LoRA path that recomputes gathered activations in backward. + + The ordinary composition of all-gather + matmul saves the full + sequence-parallel gathered input for every layer. For QKV LoRA that input + is much larger than the low-rank hidden activation. This function saves + only the local input plus LoRA weights, then repeats the small gather/matmul + sequence during backward. + """ + + @staticmethod + def forward( + ctx, x: torch.Tensor, lora_a: torch.Tensor, lora_b: torch.Tensor, scale: float, group + ): + world_size = dist.get_world_size(group) if group is not None else 1 + if world_size > 1: + gathered = _all_gather_sequence_forward(x, group, world_size) + else: + gathered = x + hidden_local = gathered.matmul(lora_a.t()) + hidden = _all_gather_last_dim_forward(hidden_local, group, world_size) + out = hidden.matmul(lora_b.t()) * scale + ctx.save_for_backward(x, lora_a, lora_b) + ctx.group = group + ctx.world_size = world_size + ctx.local_seq = x.shape[0] + ctx.local_rank_width = hidden_local.shape[-1] + ctx.scale = float(scale) + return out + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): + x, lora_a, lora_b = ctx.saved_tensors + world_size = ctx.world_size + group = ctx.group + if world_size > 1: + gathered = _all_gather_sequence_forward(x, group, world_size) + else: + gathered = x + hidden_local = gathered.matmul(lora_a.t()) + hidden = _all_gather_last_dim_forward(hidden_local, group, world_size) + + grad_out_scaled = grad_out * ctx.scale + grad_b = ( + grad_out_scaled.reshape(-1, grad_out_scaled.shape[-1]) + .t() + .matmul(hidden.reshape(-1, hidden.shape[-1])) + ) + grad_hidden = grad_out_scaled.matmul(lora_b) + if world_size > 1: + grad_hidden_local = _split_last_dim( + grad_hidden, dist.get_rank(group), ctx.local_rank_width + ) + dist.all_reduce(grad_hidden_local, op=dist.ReduceOp.SUM, group=group) + else: + grad_hidden_local = grad_hidden + grad_a = ( + grad_hidden_local.reshape(-1, grad_hidden_local.shape[-1]) + .t() + .matmul(gathered.reshape(-1, gathered.shape[-1])) + ) + grad_gathered = grad_hidden_local.matmul(lora_a) + if world_size > 1: + grad_x = _reduce_scatter_sequence_forward(grad_gathered, group, ctx.local_seq) + else: + grad_x = grad_gathered + return grad_x, grad_a, grad_b, None, None + + +def _all_gather_sequence_forward(x: torch.Tensor, group, world_size: int) -> torch.Tensor: + out = torch.empty((x.shape[0] * world_size, *x.shape[1:]), dtype=x.dtype, device=x.device) + dist.all_gather_into_tensor(out, x.contiguous(), group=group) + return out + + +def _reduce_scatter_sequence_forward(x: torch.Tensor, group, local_seq: int) -> torch.Tensor: + out = torch.empty((local_seq, *x.shape[1:]), dtype=x.dtype, device=x.device) + dist.reduce_scatter_tensor(out, x.contiguous(), group=group) + return out + + +def _all_gather_last_dim_forward(x: torch.Tensor, group, world_size: int) -> torch.Tensor: + if world_size == 1: + return x + local_width = x.shape[-1] + flat = x.movedim(-1, 0).contiguous().view(local_width, -1) + gathered = torch.empty( + (local_width * world_size, flat.shape[1]), dtype=x.dtype, device=x.device + ) + dist.all_gather_into_tensor(gathered, flat, group=group) + return gathered.view(local_width * world_size, *x.shape[:-1]).movedim(0, -1).contiguous() + + +def _split_last_dim(x: torch.Tensor, group_rank: int, local_width: int) -> torch.Tensor: + start = int(group_rank) * local_width + return x.narrow(-1, start, local_width).contiguous() + + +class LinearLoRA(nn.Module): + """Low-rank delta for a sharded linear layer. + + `a` is replicated unless the caller feeds a row-parallel local input. `b` + has the local output shard for column-parallel surfaces, and the replicated + full output for row-parallel surfaces. + """ + + def __init__( + self, + in_features: int, + out_features: int, + rank: int, + *, + alpha: int | None = None, + dropout: float = 0.0, + sequence_parallel_input: bool = False, + row_parallel_output: bool = False, + sequence_parallel_scatter_output: bool = False, + tp_group=None, + tp_rank: int = 0, + rank_partition_size: int | None = None, + rank_partitioned_a: bool = False, + input_parallel_reduce: bool = False, + output_partition_size: int | None = None, + output_partitioned_b: bool = False, + a_tensor_model_parallel: bool = False, + b_tensor_model_parallel: bool = False, + ): + super().__init__() + if rank <= 0: + raise ValueError("LoRA rank must be positive for LinearLoRA.") + self.rank = int(rank) + self.rank_partitioned_a = bool(rank_partitioned_a) + if self.rank_partitioned_a: + partition_size = ( + int(rank_partition_size) + if rank_partition_size is not None + else (dist.get_world_size(tp_group) if tp_group is not None else 1) + ) + if partition_size <= 0: + raise ValueError("LoRA rank partition size must be positive.") + if self.rank % partition_size != 0: + raise ValueError( + f"LoRA rank {self.rank} must be divisible by rank partition size {partition_size}." + ) + self.rank_partition_size = partition_size + self.local_rank = self.rank // partition_size + else: + self.rank_partition_size = 1 + self.local_rank = self.rank + self.scale = float(rank if alpha is None else alpha) / float(rank) + self.dropout_p = float(dropout) + self.sequence_parallel_input = bool(sequence_parallel_input) + self.row_parallel_output = bool(row_parallel_output) + self.sequence_parallel_scatter_output = bool(sequence_parallel_scatter_output) + if self.row_parallel_output and self.sequence_parallel_scatter_output: + raise ValueError( + "Use either row_parallel_output or sequence_parallel_scatter_output, not both." + ) + self.tp_group = tp_group + self.tp_rank = int(tp_rank) + self.input_parallel_reduce = bool(input_parallel_reduce) + self.output_partitioned_b = bool(output_partitioned_b) + if self.output_partitioned_b: + partition_size = ( + int(output_partition_size) + if output_partition_size is not None + else (dist.get_world_size(tp_group) if tp_group is not None else 1) + ) + if partition_size <= 0: + raise ValueError("LoRA output partition size must be positive.") + if out_features % partition_size != 0: + raise ValueError( + f"LoRA output features {out_features} must be divisible by {partition_size}." + ) + self.output_partition_size = partition_size + self.local_out_features = out_features // partition_size + else: + self.output_partition_size = 1 + self.local_out_features = out_features + self.lora_a = nn.Parameter(torch.empty(self.local_rank, in_features)) + self.lora_b = nn.Parameter(torch.empty(self.local_out_features, rank)) + self.lora_a.tensor_model_parallel = bool(a_tensor_model_parallel) + self.lora_b.tensor_model_parallel = bool(b_tensor_model_parallel) + nn.init.kaiming_uniform_(self.lora_a, a=5**0.5) + nn.init.zeros_(self.lora_b) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.sequence_parallel_input and self.rank_partitioned_a and not self.training: + # Keep eval/inference on the simple path; the memory optimization + # matters only when autograd needs to retain forward activations. + pass + elif ( + self.sequence_parallel_input + and self.rank_partitioned_a + and not self.input_parallel_reduce + and not self.output_partitioned_b + and not self.row_parallel_output + and not self.sequence_parallel_scatter_output + and self.dropout_p == 0.0 + ): + return _SequenceParallelRankPartitionedLoRA.apply( + x, self.lora_a, self.lora_b, self.scale, self.tp_group + ) + if self.sequence_parallel_input: + x = _gather_sequence_parallel(x, self.tp_group) + dropped = F.dropout(x, p=self.dropout_p, training=self.training) if self.dropout_p else x + hidden = dropped.matmul(self.lora_a.t()) + if self.rank_partitioned_a: + hidden = _all_gather_last_dim(hidden, self.tp_group, reduce_backward=True) + if self.input_parallel_reduce: + hidden = _all_reduce_sum(hidden, self.tp_group) + out = hidden.matmul(self.lora_b.t()) * self.scale + if self.output_partitioned_b: + out = _all_gather_last_dim(out, self.tp_group) + if self.row_parallel_output: + out = _reduce_scatter_sequence_parallel(out, self.tp_group) + if self.sequence_parallel_scatter_output: + out = _scatter_sequence_parallel(out, self.tp_group, self.tp_rank) + return out + + +class GroupedLinearLoRA(nn.Module): + """Per-local-expert LoRA delta for `te.GroupedLinear` expert surfaces.""" + + def __init__( + self, + num_local_experts: int, + in_features: int, + out_features: int, + rank: int, + *, + alpha: int | None = None, + dropout: float = 0.0, + ): + super().__init__() + if rank <= 0: + raise ValueError("LoRA rank must be positive for GroupedLinearLoRA.") + self.num_local_experts = int(num_local_experts) + self.rank = int(rank) + self.scale = float(rank if alpha is None else alpha) / float(rank) + self.dropout_p = float(dropout) + self.lora_a = nn.Parameter(torch.empty(num_local_experts, rank, in_features)) + self.lora_b = nn.Parameter(torch.empty(num_local_experts, out_features, rank)) + nn.init.kaiming_uniform_(self.lora_a, a=5**0.5) + nn.init.zeros_(self.lora_b) + + def forward(self, x: torch.Tensor, splits: list[int]) -> torch.Tensor: + if len(splits) != self.num_local_experts: + raise ValueError( + f"GroupedLinearLoRA expected {self.num_local_experts} splits, got {len(splits)}." + ) + outputs = [] + offset = 0 + for expert_idx, size in enumerate(splits): + x_i = x[offset : offset + size] + if size == 0: + outputs.append(x_i.new_empty((0, self.lora_b.shape[1]))) + else: + dropped = ( + F.dropout(x_i, p=self.dropout_p, training=self.training) + if self.dropout_p + else x_i + ) + h_i = dropped.matmul(self.lora_a[expert_idx].t()) + outputs.append(h_i.matmul(self.lora_b[expert_idx].t()) * self.scale) + offset += size + return torch.cat(outputs, dim=0) if outputs else x.new_empty((0, self.lora_b.shape[1])) + + +class SharedGroupedLinearLoRA(nn.Module): + """LoRA delta shared by all local experts in a GroupedLinear.""" + + def __init__( + self, + num_local_experts: int, + in_features: int, + out_features: int, + rank: int, + *, + alpha: int | None = None, + dropout: float = 0.0, + ): + super().__init__() + if rank <= 0: + raise ValueError("LoRA rank must be positive for SharedGroupedLinearLoRA.") + self.num_local_experts = int(num_local_experts) + self.rank = int(rank) + self.scale = float(rank if alpha is None else alpha) / float(rank) + self.dropout_p = float(dropout) + self.shared_across_experts = True + self.lora_a = nn.Parameter(torch.empty(rank, in_features)) + self.lora_b = nn.Parameter(torch.empty(out_features, rank)) + self.lora_a.tensor_model_parallel = False + self.lora_b.tensor_model_parallel = False + nn.init.kaiming_uniform_(self.lora_a, a=5**0.5) + nn.init.zeros_(self.lora_b) + + def forward(self, x: torch.Tensor, splits: list[int]) -> torch.Tensor: + if len(splits) != self.num_local_experts: + raise ValueError( + f"SharedGroupedLinearLoRA expected {self.num_local_experts} splits, got {len(splits)}." + ) + dropped = F.dropout(x, p=self.dropout_p, training=self.training) if self.dropout_p else x + return dropped.matmul(self.lora_a.t()).matmul(self.lora_b.t()) * self.scale + + +__all__ = [ + "GroupedLinearLoRA", + "LinearLoRA", + "LoraConfig", + "SharedGroupedLinearLoRA", + "freeze_non_lora_params", + "normalize_lora_config", + "trainable_param_stats", +] diff --git a/experimental/lite/megatron/lite/primitive/modules/moe.py b/experimental/lite/megatron/lite/primitive/modules/moe.py new file mode 100644 index 00000000000..c7d294195c5 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/modules/moe.py @@ -0,0 +1,79 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Shared MoE utilities: _AllToAll and MoEAuxLossAutoScaler. + +Extracted from models/*/moe.py (Level 0 Option C — pure extraction, no behavior change). +All three models (qwen3_moe, qwen3_5, deepseek_v3) had identical _AllToAll +implementations and functionally identical MoEAuxLossAutoScaler implementations. +The qwen3_moe version is used as the canonical form (adds docstring and named +intermediate variable for clarity). + +Note: this is megatron.lite's own MoEAuxLossAutoScaler, kept deliberately separate +from MC's `megatron.core.transformer.moe.moe_utils.MoEAuxLossAutoScaler`. +`runtime/backends/mlite/runtime.py` calls `set_loss_scale` on this class to +apply the 1/num_microbatches aux-loss gradient scale. Megatron-Core MoE modules +use MC's class directly when they are imported by MC internally. +""" + +from __future__ import annotations + +import torch # pyright: ignore[reportMissingImports] +import torch.distributed as dist # pyright: ignore[reportMissingImports] + +__all__ = ["MoEAuxLossAutoScaler", "_AllToAll"] + + +class MoEAuxLossAutoScaler(torch.autograd.Function): + """Piggyback aux_loss onto autograd so main_loss.backward() triggers it.""" + + main_loss_backward_scale: torch.Tensor | None = None + + @staticmethod + def forward(ctx, output: torch.Tensor, aux_loss: torch.Tensor) -> torch.Tensor: + ctx.save_for_backward(aux_loss) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + (aux_loss,) = ctx.saved_tensors + scale = ( + MoEAuxLossAutoScaler.main_loss_backward_scale.to(aux_loss.device) + if MoEAuxLossAutoScaler.main_loss_backward_scale is not None + else torch.ones(1, device=aux_loss.device) + ) + scaled_aux_loss_grad = torch.ones_like(aux_loss) * scale + return grad_output, scaled_aux_loss_grad + + @staticmethod + def set_loss_scale(scale: torch.Tensor) -> None: + MoEAuxLossAutoScaler.main_loss_backward_scale = scale + + +class _AllToAll(torch.autograd.Function): + @staticmethod + def forward(ctx, input_tensor, input_splits, output_splits, group): + ctx.input_splits = input_splits + ctx.output_splits = output_splits + ctx.group = group + input_tensor = input_tensor.contiguous() + output = input_tensor.new_empty([sum(output_splits)] + list(input_tensor.shape[1:])) + dist.all_to_all_single( + output, + input_tensor, + output_split_sizes=output_splits, + input_split_sizes=input_splits, + group=group, + ) + return output + + @staticmethod + def backward(ctx, grad_output): + grad_output = grad_output.contiguous() + grad_input = grad_output.new_empty([sum(ctx.input_splits)] + list(grad_output.shape[1:])) + dist.all_to_all_single( + grad_input, + grad_output, + output_split_sizes=ctx.input_splits, + input_split_sizes=ctx.output_splits, + group=ctx.group, + ) + return grad_input, None, None, None diff --git a/experimental/lite/megatron/lite/primitive/modules/mrope.py b/experimental/lite/megatron/lite/primitive/modules/mrope.py new file mode 100644 index 00000000000..71718d87adc --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/modules/mrope.py @@ -0,0 +1,55 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Multimodal rotary embedding primitive.""" + +from __future__ import annotations + +import torch +import torch.distributed as dist +import torch.nn as nn + +from megatron.core.models.common.embeddings.rope_utils import ( # pyright: ignore[reportMissingImports] + get_pos_emb_on_this_cp_rank, +) + +__all__ = ["MultimodalRotaryEmbedding"] + + +class MultimodalRotaryEmbedding(nn.Module): + """Qwen-style multimodal RoPE embedding with optional CP slicing.""" + + def __init__( + self, + *, + kv_channels: int, + rotary_percent: float, + rotary_base: float, + cp_group: dist.ProcessGroup | None, + ): + super().__init__() + dim = int(kv_channels * rotary_percent) + inv_freq = 1.0 / (rotary_base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.cp_group = cp_group + + @staticmethod + def _apply_interleaved_mrope(freqs: torch.Tensor, mrope_section: list[int]) -> torch.Tensor: + freqs_t = freqs[0].clone() + for dim, offset in enumerate((1, 2), start=1): + length = mrope_section[dim] * 3 + freqs_t[..., offset:length:3] = freqs[dim, ..., offset:length:3] + return freqs_t + + def forward( + self, position_ids: torch.Tensor, mrope_section: list[int], *, packed_seq: bool = False + ) -> torch.Tensor: + seq = position_ids.to(device=self.inv_freq.device) + inv_freq = self.inv_freq.float() + inv = inv_freq[None, None, :, None].expand(3, seq.shape[1], -1, 1) + seq_expanded = seq[:, :, None, :].float() + freqs = (inv @ seq_expanded).transpose(2, 3) + freqs = self._apply_interleaved_mrope(freqs, mrope_section) + emb = torch.cat((freqs, freqs), dim=-1) + emb = emb[..., None, :].transpose(0, 1).contiguous() + if not packed_seq and self.cp_group is not None and dist.get_world_size(self.cp_group) > 1: + emb = get_pos_emb_on_this_cp_rank(emb, 0, self.cp_group) + return emb diff --git a/experimental/lite/megatron/lite/primitive/modules/mtp.py b/experimental/lite/megatron/lite/primitive/modules/mtp.py new file mode 100644 index 00000000000..299e143618c --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/modules/mtp.py @@ -0,0 +1,133 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Multi-token prediction primitives.""" + +from __future__ import annotations + +from collections.abc import Callable + +import torch +import torch.nn as nn +import transformer_engine.pytorch as te + +from megatron.lite.primitive.parallel import ( + ParallelState, + VanillaColumnParallelLinear, + VocabParallelEmbedding, + roll_packed_thd_left, + scatter_to_sequence_parallel, +) + +__all__ = ["MTPBlock", "MTPDecoderLayer", "MTPLossAutoScaler"] + + +class MTPLossAutoScaler(torch.autograd.Function): + main_loss_backward_scale: float = 1.0 + + @staticmethod + def forward(ctx, output: torch.Tensor, mtp_loss: torch.Tensor): + ctx.save_for_backward(mtp_loss) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + (mtp_loss,) = ctx.saved_tensors + scaled_mtp_grad = torch.ones_like(mtp_loss) * MTPLossAutoScaler.main_loss_backward_scale + return grad_output, scaled_mtp_grad + + @staticmethod + def set_loss_scale(scale: torch.Tensor | float) -> None: + if isinstance(scale, torch.Tensor): + scale = float(scale.detach().float().item()) + MTPLossAutoScaler.main_loss_backward_scale = float(scale) + + +class MTPDecoderLayer(nn.Module): + def __init__( + self, + *, + hidden_size: int, + rms_norm_eps: float, + ps: ParallelState, + embedding: VocabParallelEmbedding, + transformer_layer: nn.Module, + detach_encoder: bool, + ): + super().__init__() + self.ps = ps + self.embedding = embedding + self.detach_encoder = detach_encoder + self.enorm = te.RMSNorm(hidden_size, eps=rms_norm_eps, zero_centered_gamma=True) + self.hnorm = te.RMSNorm(hidden_size, eps=rms_norm_eps, zero_centered_gamma=True) + self.eh_proj = VanillaColumnParallelLinear( + hidden_size * 2, hidden_size, ps, sp=ps.tp_size > 1, gather_output=True + ) + self.transformer_layer = transformer_layer + self.final_layernorm = te.RMSNorm(hidden_size, eps=rms_norm_eps, zero_centered_gamma=True) + + def forward( + self, + *, + input_ids: torch.Tensor, + position_ids: torch.Tensor | None, + hidden_states: torch.Tensor, + rotary_position_ids: torch.Tensor | None = None, + packed_seq_params=None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + attention_position_ids = ( + rotary_position_ids if rotary_position_ids is not None else position_ids + ) + input_ids, _ = roll_packed_thd_left(input_ids, packed_seq_params=packed_seq_params, dims=-1) + if position_ids is not None: + position_ids, _ = roll_packed_thd_left( + position_ids, packed_seq_params=packed_seq_params, dims=-1 + ) + decoder_input = scatter_to_sequence_parallel(self.embedding(input_ids), self.ps) + if self.detach_encoder: + decoder_input = decoder_input.detach() + hidden_states = hidden_states.detach() + decoder_input = self.enorm(decoder_input) + hidden_states = self.hnorm(hidden_states) + hidden_states = torch.cat((decoder_input, hidden_states), dim=-1) + hidden_states = scatter_to_sequence_parallel(self.eh_proj(hidden_states), self.ps) + hidden_states = self.transformer_layer( + hidden_states, position_ids=attention_position_ids, packed_seq_params=packed_seq_params + ) + hidden_states = self.final_layernorm(hidden_states) + return hidden_states, input_ids, position_ids + + +class MTPBlock(nn.Module): + def __init__( + self, + *, + num_layers: int, + repeated_layer: bool, + layer_factory: Callable[[int], MTPDecoderLayer], + ): + super().__init__() + self.num_layers = num_layers + self.repeated_layer = repeated_layer + layers_to_build = 1 if repeated_layer else num_layers + self.layers = nn.ModuleList([layer_factory(idx) for idx in range(layers_to_build)]) + + def forward( + self, + *, + input_ids: torch.Tensor, + position_ids: torch.Tensor | None, + hidden_states: torch.Tensor, + packed_seq_params=None, + ) -> list[torch.Tensor]: + outputs: list[torch.Tensor] = [] + rotary_position_ids = position_ids + for depth in range(self.num_layers): + layer = self.layers[0] if self.repeated_layer else self.layers[depth] + hidden_states, input_ids, position_ids = layer( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + rotary_position_ids=rotary_position_ids, + packed_seq_params=packed_seq_params, + ) + outputs.append(hidden_states) + return outputs diff --git a/experimental/lite/megatron/lite/primitive/modules/router.py b/experimental/lite/megatron/lite/primitive/modules/router.py new file mode 100644 index 00000000000..821407a2a7b --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/modules/router.py @@ -0,0 +1,208 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""MoE router implementations: TopKRouter (softmax) and SigmoidTopKRouter. + +Internals call the atomic free functions in Megatron-Core's +`megatron.core.transformer.moe.moe_utils` (plan `docs/moe_mc_wrap_plan.md` +D3/D4). The outer classes keep the flat-kwargs + `ParallelState` constructor +style of megatron.lite primitives — no `TransformerConfig`, no mpu globals. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch # pyright: ignore[reportMissingImports] +import torch.distributed as dist # pyright: ignore[reportMissingImports] +import torch.nn as nn # pyright: ignore[reportMissingImports] + +from megatron.core.transformer.moe.moe_utils import ( # pyright: ignore[reportMissingImports] + compute_routing_scores_for_aux_loss, + router_gating_linear, + switch_load_balancing_loss_func, + topk_routing_with_score_function, +) +from megatron.lite.primitive.modules.moe import MoEAuxLossAutoScaler + +if TYPE_CHECKING: + from megatron.lite.primitive.parallel import ParallelState + + +def _ordered_topk_from_routing_map( + probs_dense: torch.Tensor, routing_map: torch.Tensor, topk: int +) -> tuple[torch.Tensor, torch.Tensor]: + expert_ids = torch.arange( + probs_dense.size(-1), device=probs_dense.device, dtype=torch.long + ).expand_as(routing_map) + masked_ids = torch.where( + routing_map, expert_ids, torch.full_like(expert_ids, probs_dense.size(-1)) + ) + topk_indices = torch.sort(masked_ids, dim=-1).values[:, :topk] + topk_scores = torch.gather(probs_dense, dim=-1, index=topk_indices) + return topk_scores, topk_indices + + +class TopKRouter(nn.Module): + """TopK gating with optional high-precision router logits/probabilities.""" + + def __init__( + self, + config, + ps: ParallelState, + *, + router_bias_rate: float = 0.0, + compute_aux_loss: bool = True, + use_pre_softmax: bool = False, + moe_router_fusion: bool = False, + router_dtype: torch.dtype | None = None, + ): + super().__init__() + if router_bias_rate > 0: + raise NotImplementedError( + "expert-bias EMA is not implemented in the primitive router; " + "use load_balancing_type='none' or extend ParallelState." + ) + self.topk = config.num_experts_per_tok + self.num_experts = config.num_experts + self.aux_loss_coeff = config.router_aux_loss_coef + self.router_bias_rate = router_bias_rate + self.compute_aux_loss = compute_aux_loss + self.use_pre_softmax = use_pre_softmax + self.moe_router_fusion = moe_router_fusion + self.router_dtype = router_dtype + + self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) + self.register_buffer( + "expert_bias", torch.zeros(config.num_experts, dtype=torch.float32), persistent=False + ) + + self._aux_loss_group = ps.tp_group if ps.tp_size > 1 else None + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + router_dtype = self.router_dtype or x.dtype + logits = router_gating_linear(x, self.gate.weight, None, router_dtype) + logits = logits.view(-1, self.num_experts) + num_tokens = logits.size(0) + if self.moe_router_fusion: + probs_dense, _ = topk_routing_with_score_function( + logits, + self.topk, + use_pre_softmax=self.use_pre_softmax, + score_function="softmax", + fused=True, + ) + topk_scores, topk_indices = torch.topk(probs_dense, k=self.topk, dim=-1) + else: + probs_dense, routing_map = topk_routing_with_score_function( + logits, + self.topk, + use_pre_softmax=self.use_pre_softmax, + score_function="softmax", + fused=False, + ) + topk_scores, topk_indices = _ordered_topk_from_routing_map( + probs_dense, routing_map, self.topk + ) + if self.router_dtype is None: + topk_scores = topk_scores.to(x.dtype) + + if self.compute_aux_loss and self.training and torch.is_grad_enabled(): + routing_map, aux_scores = compute_routing_scores_for_aux_loss( + logits, self.topk, score_function="softmax", fused=self.moe_router_fusion + ) + tokens_per_expert = routing_map.sum(dim=0).to(torch.int64) + total_num_tokens = num_tokens + if self._aux_loss_group is not None: + dist.all_reduce(tokens_per_expert, group=self._aux_loss_group) + total_num_tokens = num_tokens * dist.get_world_size(group=self._aux_loss_group) + aux_loss = switch_load_balancing_loss_func( + aux_scores, + tokens_per_expert, + total_num_tokens, + self.topk, + self.num_experts, + self.aux_loss_coeff, + fused=False, + ) + topk_scores = MoEAuxLossAutoScaler.apply(topk_scores, aux_loss) + + return topk_scores, topk_indices + + +class SigmoidTopKRouter(nn.Module): + """Sigmoid-based TopK router for DeepSeek V3.""" + + def __init__( + self, + config, + ps: ParallelState, + *, + router_bias_rate: float = 0.0, + compute_aux_loss: bool = True, + use_pre_softmax: bool = False, + moe_router_fusion: bool = False, + ): + super().__init__() + if router_bias_rate > 0: + raise NotImplementedError( + "expert-bias EMA is not implemented in the primitive router; " + "use load_balancing_type='none' or extend ParallelState." + ) + self.topk = config.num_experts_per_tok + self.num_experts = config.n_routed_experts + self.aux_loss_coeff = config.aux_loss_alpha + self.scaling_factor = config.routed_scaling_factor + self.router_bias_rate = router_bias_rate + self.compute_aux_loss = compute_aux_loss + self.use_pre_softmax = use_pre_softmax + self.moe_router_fusion = moe_router_fusion + + self.gate = nn.Linear(config.hidden_size, config.n_routed_experts, bias=False) + self.register_buffer( + "expert_bias", + torch.zeros(config.n_routed_experts, dtype=torch.float32), + persistent=False, + ) + + self._aux_loss_group = ps.tp_group if ps.tp_size > 1 else None + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + logits = self.gate(x) + logits = logits.view(-1, self.num_experts) + num_tokens = logits.size(0) + probs_dense, routing_map = topk_routing_with_score_function( + logits, + self.topk, + score_function="sigmoid", + expert_bias=self.expert_bias.to(logits.dtype), + scaling_factor=(self.scaling_factor or None), + fused=self.moe_router_fusion, + ) + topk_scores, topk_indices = _ordered_topk_from_routing_map( + probs_dense, routing_map, self.topk + ) + topk_scores = topk_scores.to(logits.dtype) + + if self.compute_aux_loss and self.training and torch.is_grad_enabled(): + _, aux_scores = compute_routing_scores_for_aux_loss( + logits, self.topk, score_function="sigmoid", fused=self.moe_router_fusion + ) + tokens_per_expert = routing_map.sum(dim=0).to(torch.int64) + total_num_tokens = num_tokens + if self._aux_loss_group is not None: + dist.all_reduce(tokens_per_expert, group=self._aux_loss_group) + total_num_tokens = num_tokens * dist.get_world_size(group=self._aux_loss_group) + aux_loss = switch_load_balancing_loss_func( + aux_scores, + tokens_per_expert, + total_num_tokens, + self.topk, + self.num_experts, + self.aux_loss_coeff, + fused=False, + ) + topk_scores = MoEAuxLossAutoScaler.apply(topk_scores, aux_loss) + + return topk_scores, topk_indices + + +__all__ = ["SigmoidTopKRouter", "TopKRouter"] diff --git a/experimental/lite/megatron/lite/primitive/ops/__init__.py b/experimental/lite/megatron/lite/primitive/ops/__init__.py new file mode 100644 index 00000000000..473fa29a7c0 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/ops/__init__.py @@ -0,0 +1,27 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Runtime math primitives for Megatron Lite.""" + +from megatron.lite.primitive.ops.cross_entropy import vocab_parallel_cross_entropy +from megatron.lite.primitive.ops.gated_delta_rule import l2norm, torch_chunk_gated_delta_rule +from megatron.lite.primitive.ops.logprob import ( + vocab_parallel_entropy, + vocab_parallel_log_probs_from_logits, +) +from megatron.lite.primitive.ops.sp_ops import ( + AllGatherDim0, + AllGatherDim0ForNonSPConsumer, + ReduceScatterDim0, + ScatterToSP, +) + +__all__ = [ + "AllGatherDim0", + "AllGatherDim0ForNonSPConsumer", + "ReduceScatterDim0", + "ScatterToSP", + "l2norm", + "torch_chunk_gated_delta_rule", + "vocab_parallel_cross_entropy", + "vocab_parallel_entropy", + "vocab_parallel_log_probs_from_logits", +] diff --git a/experimental/lite/megatron/lite/primitive/ops/cross_entropy.py b/experimental/lite/megatron/lite/primitive/ops/cross_entropy.py new file mode 100644 index 00000000000..9bdc40d3867 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/ops/cross_entropy.py @@ -0,0 +1,113 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Vocab-parallel cross entropy loss (copied from Megatron-Core). + +Computes cross entropy when logits are split across TP ranks. +With TP=1 the all-reduce calls are no-ops and this degenerates to a +standard cross-entropy with in-place ops and a memory-efficient custom backward. +""" + +from __future__ import annotations + +import torch # pyright: ignore[reportMissingImports] +import torch.distributed as dist # pyright: ignore[reportMissingImports] + + +def _vocab_range(partition_vocab_size: int, rank: int, world_size: int): + start = rank * partition_vocab_size + return start, start + partition_vocab_size + + +class _VocabParallelCrossEntropy(torch.autograd.Function): + + @staticmethod + def forward(ctx, vocab_parallel_logits, target, tp_group): + # Cast to float32 and compute max for numerical stability. + vocab_parallel_logits = vocab_parallel_logits.float() + logits_max = torch.max(vocab_parallel_logits, dim=-1)[0] + + if tp_group is not None and dist.get_world_size(tp_group) > 1: + dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=tp_group) + + # In-place subtract max. + vocab_parallel_logits -= logits_max.unsqueeze(dim=-1) + + # Partition info. + partition_vocab_size = vocab_parallel_logits.size(-1) + if tp_group is not None and dist.get_world_size(tp_group) > 1: + rank = dist.get_rank(tp_group) + world_size = dist.get_world_size(tp_group) + else: + rank = 0 + world_size = 1 + vocab_start_index, vocab_end_index = _vocab_range(partition_vocab_size, rank, world_size) + + # Mask targets outside this partition's vocab range. + target_mask = (target < vocab_start_index) | (target >= vocab_end_index) + masked_target = target.clone() - vocab_start_index + masked_target[target_mask] = 0 + + # Get predicted logits = logits[target]. + logits_2d = vocab_parallel_logits.view(-1, partition_vocab_size) + masked_target_1d = masked_target.view(-1) + arange_1d = torch.arange(logits_2d.size(0), device=logits_2d.device) + predicted_logits_1d = logits_2d[arange_1d, masked_target_1d] + predicted_logits_1d = predicted_logits_1d.clone().contiguous() + predicted_logits = predicted_logits_1d.view_as(target) + predicted_logits[target_mask] = 0.0 + + # Sum of exp(logits). + exp_logits = vocab_parallel_logits + torch.exp(vocab_parallel_logits, out=exp_logits) + sum_exp_logits = exp_logits.sum(dim=-1) + + # All-reduce predicted_logits and sum_exp_logits across TP. + if tp_group is not None and dist.get_world_size(tp_group) > 1: + dist.all_reduce(predicted_logits, op=dist.ReduceOp.SUM, group=tp_group) + dist.all_reduce(sum_exp_logits, op=dist.ReduceOp.SUM, group=tp_group) + + # Loss = log(sum(exp(logits))) - predicted_logit. + loss = torch.log(sum_exp_logits) - predicted_logits + + # Normalize exp_logits to get softmax (reused in backward). + exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1)) + + # Save for backward. + ctx.save_for_backward(exp_logits, target_mask, masked_target_1d) + + return loss + + @staticmethod + def backward(ctx, grad_output): + softmax, target_mask, masked_target_1d = ctx.saved_tensors + + # grad_input = softmax (copy is implicit since softmax is saved). + grad_input = softmax + partition_vocab_size = softmax.size(-1) + grad_2d = grad_input.view(-1, partition_vocab_size) + + arange_1d = torch.arange(grad_2d.size(0), device=grad_2d.device) + softmax_update = 1.0 - target_mask.view(-1).float() + + grad_2d[arange_1d, masked_target_1d] -= softmax_update + + # Scale by upstream gradient. + grad_input.mul_(grad_output.unsqueeze(dim=-1)) + + return grad_input, None, None + + +def vocab_parallel_cross_entropy(vocab_parallel_logits, target, tp_group=None): + """Cross entropy loss for vocab-parallel logits. + + Args: + vocab_parallel_logits: [S, B, V/tp] logits split across TP ranks. + target: [S, B] integer target token ids. + tp_group: TP process group (None or single-rank group → no communication). + + Returns: + Per-token loss tensor of shape [S, B]. + """ + return _VocabParallelCrossEntropy.apply(vocab_parallel_logits, target, tp_group) + + +__all__ = ["vocab_parallel_cross_entropy"] diff --git a/experimental/lite/megatron/lite/primitive/ops/gated_delta_rule.py b/experimental/lite/megatron/lite/primitive/ops/gated_delta_rule.py new file mode 100644 index 00000000000..62c21e441f0 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/ops/gated_delta_rule.py @@ -0,0 +1,108 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Gated Delta Rule math helpers.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +try: + from fla.modules.l2norm import l2norm as _fla_l2norm # pyright: ignore[reportMissingImports] + + _HAS_FLA_L2NORM = True +except ImportError: + _HAS_FLA_L2NORM = False + +__all__ = ["l2norm", "torch_chunk_gated_delta_rule"] + + +def l2norm(x: torch.Tensor, dim: int = -1, eps: float = 1e-6) -> torch.Tensor: + if _HAS_FLA_L2NORM and dim == -1 and eps == 1e-6: + return _fla_l2norm(x) + return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps) + + +def torch_chunk_gated_delta_rule( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + chunk_size: int = 64, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Pure PyTorch Gated Delta Rule fallback for correctness smoke tests.""" + initial_dtype = query.dtype + if use_qk_l2norm_in_kernel: + query = l2norm(query) + key = l2norm(key) + query, key, value, beta, g = [ + x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g) + ] + + batch_size, num_heads, sequence_length, key_dim = key.shape + value_dim = value.shape[-1] + pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size + query = F.pad(query, (0, 0, 0, pad_size)) + key = F.pad(key, (0, 0, 0, pad_size)) + value = F.pad(value, (0, 0, 0, pad_size)) + beta = F.pad(beta, (0, pad_size)) + g = F.pad(g, (0, pad_size)) + total_sequence_length = sequence_length + pad_size + scale = 1 / (query.shape[-1] ** 0.5) + query = query * scale + + v_beta = value * beta.unsqueeze(-1) + k_beta = key * beta.unsqueeze(-1) + query, key, value, k_beta, v_beta = [ + x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) + for x in (query, key, value, k_beta, v_beta) + ] + g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size) + + mask = torch.triu( + torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0 + ) + g = g.cumsum(dim=-1) + decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() + attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0) + for i in range(1, chunk_size): + row = attn[..., i, :i].clone() + sub = attn[..., :i, :i].clone() + attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) + attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device) + + value = attn @ v_beta + k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) + last_recurrent_state = ( + torch.zeros(batch_size, num_heads, key_dim, value_dim).to(value) + if initial_state is None + else initial_state.to(value) + ) + core_attn_out = torch.zeros_like(value) + mask = torch.triu( + torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1 + ) + + for i in range(0, total_sequence_length // chunk_size): + q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] + attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0) + v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state + v_new = v_i - v_prime + attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state + core_attn_out[:, :, i] = attn_inter + attn @ v_new + last_recurrent_state = ( + last_recurrent_state * g[:, :, i, -1, None, None].exp() + + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new + ) + + if not output_final_state: + last_recurrent_state = None + core_attn_out = core_attn_out.reshape( + core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1] + ) + core_attn_out = core_attn_out[:, :, :sequence_length] + core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype) + return core_attn_out, last_recurrent_state diff --git a/experimental/lite/megatron/lite/primitive/ops/linear_cross_entropy.py b/experimental/lite/megatron/lite/primitive/ops/linear_cross_entropy.py new file mode 100644 index 00000000000..453702dc05e --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/ops/linear_cross_entropy.py @@ -0,0 +1,70 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Linear + vocab-parallel cross entropy helpers. + +The fast path delegates to VERL's Triton-backed fused kernel when available. +The fallback keeps the same contract and is used by unit/smoke tests that do +not have the fused extension on ``PYTHONPATH``. +""" + +from __future__ import annotations + +import torch +import torch.distributed as dist + +from megatron.lite.primitive.ops.cross_entropy import vocab_parallel_cross_entropy + + +def _all_reduce_if_needed(tensor: torch.Tensor, group, op=dist.ReduceOp.SUM) -> torch.Tensor: + if group is not None and dist.is_initialized() and dist.get_world_size(group) > 1: + dist.all_reduce(tensor, op=op, group=group) + return tensor + + +def _vocab_parallel_entropy(logits: torch.Tensor, tp_group=None) -> torch.Tensor: + logits = logits.float() + logits_max = logits.max(dim=-1).values + _all_reduce_if_needed(logits_max, tp_group, op=dist.ReduceOp.MAX) + + shifted = logits - logits_max.unsqueeze(-1) + exp_logits = torch.exp(shifted) + sum_exp = exp_logits.sum(dim=-1) + _all_reduce_if_needed(sum_exp, tp_group) + + weighted_logits = (exp_logits * logits).sum(dim=-1) + _all_reduce_if_needed(weighted_logits, tp_group) + expected_logits = weighted_logits / sum_exp + return torch.log(sum_exp) + logits_max - expected_logits + + +def _reshape_like_labels(values: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + if values.shape != labels.shape and values.numel() == labels.numel(): + return values.reshape(labels.shape) + return values + + +def linear_cross_entropy( + hidden: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + temperature: float = 1.0, + tp_group=None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return token log-probs and entropy without changing VERL's fused API.""" + try: + from verl.utils.kernel.linear_cross_entropy import linear_cross_entropy as _verl_lce + except Exception: + _verl_lce = None + + if _verl_lce is not None and hidden.is_cuda: + log_probs, entropy = _verl_lce(hidden, weight, labels, float(temperature), "none", tp_group) + return _reshape_like_labels(log_probs, labels), _reshape_like_labels(entropy, labels) + + logits = torch.matmul(hidden, weight.t()) + if temperature != 1.0: + logits = logits / float(temperature) + loss = vocab_parallel_cross_entropy(logits.clone(), labels, tp_group) + entropy = _vocab_parallel_entropy(logits, tp_group) + return _reshape_like_labels(-loss, labels), _reshape_like_labels(entropy, labels) + + +__all__ = ["linear_cross_entropy"] diff --git a/experimental/lite/megatron/lite/primitive/ops/logprob.py b/experimental/lite/megatron/lite/primitive/ops/logprob.py new file mode 100644 index 00000000000..c2fc295ef8b --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/ops/logprob.py @@ -0,0 +1,73 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Probability-first primitives for Megatron Lite phase 1.""" + +from __future__ import annotations + +import torch +import torch.distributed as dist + + +def vocab_parallel_log_probs_from_logits( + logits: torch.Tensor, labels: torch.Tensor | None = None +) -> torch.Tensor: + """Compute log-probabilities from already materialized logits. + + Phase 1 assumes logits are already full-vocab tensors on the local rank. + If `labels` are provided, this returns token-selected log-probabilities with + the same shape as `labels` (modulo an internal transpose when logits are + `[S, B, V]` and labels are `[B, S]`). + """ + + log_probs = torch.log_softmax(logits.float(), dim=-1) + if labels is None: + return log_probs + + aligned_labels, transposed = _align_labels_to_logits(logits, labels) + gathered = log_probs.gather(dim=-1, index=aligned_labels.unsqueeze(-1)).squeeze(-1) + return gathered.transpose(0, 1).contiguous() if transposed else gathered + + +def _all_reduce_if_needed(tensor: torch.Tensor, group=None, op=dist.ReduceOp.SUM) -> torch.Tensor: + if group is not None and dist.is_initialized() and dist.get_world_size(group) > 1: + dist.all_reduce(tensor, op=op, group=group) + return tensor + + +def vocab_parallel_entropy(logits: torch.Tensor, tp_group=None) -> torch.Tensor: + """Compute per-token entropy from logits.""" + + logits = logits.float() + logits_max = logits.max(dim=-1).values + _all_reduce_if_needed(logits_max, tp_group, op=dist.ReduceOp.MAX) + + shifted = logits - logits_max.unsqueeze(-1) + exp_logits = torch.exp(shifted) + sum_exp = exp_logits.sum(dim=-1) + _all_reduce_if_needed(sum_exp, tp_group) + + weighted_logits = (exp_logits * logits).sum(dim=-1) + _all_reduce_if_needed(weighted_logits, tp_group) + expected_logits = weighted_logits / sum_exp + return torch.log(sum_exp) + logits_max - expected_logits + + +def _align_labels_to_logits( + logits: torch.Tensor, labels: torch.Tensor +) -> tuple[torch.Tensor, bool]: + if logits.ndim != labels.ndim + 1: + raise ValueError( + f"logits rank must be labels rank + 1, got logits={logits.shape}, labels={labels.shape}." + ) + + if logits.shape[:-1] == labels.shape: + return labels, False + + if ( + logits.ndim == 3 + and labels.ndim == 2 + and logits.shape[0] == labels.shape[1] + and logits.shape[1] == labels.shape[0] + ): + return labels.transpose(0, 1).contiguous(), True + + raise ValueError(f"Could not align labels {labels.shape} with logits {logits.shape}.") diff --git a/experimental/lite/megatron/lite/primitive/ops/sp_ops.py b/experimental/lite/megatron/lite/primitive/ops/sp_ops.py new file mode 100644 index 00000000000..0d52cd0daaf --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/ops/sp_ops.py @@ -0,0 +1,90 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Sequence-parallel autograd primitives for non-TE layers (embedding, scatter/gather). + +AllGather/ReduceScatter operate on the sequence dimension (dim 0) with layout +[S, B, H]. NCCL operates on dim 0 directly — no transpose needed. +TE-based layers use TE's native sequence_parallel=True instead. +""" + +from __future__ import annotations + +import torch # pyright: ignore[reportMissingImports] +import torch.distributed as dist # pyright: ignore[reportMissingImports] + + +def _ag_dim0(x: torch.Tensor, tp_size: int, group: dist.ProcessGroup) -> torch.Tensor: + """AllGather on dim 0: [S/tp, B, H] → [S, B, H].""" + out = torch.empty(tp_size * x.shape[0], x.shape[1], x.shape[2], dtype=x.dtype, device=x.device) + dist.all_gather_into_tensor(out, x.contiguous(), group=group) + return out + + +def _rs_dim0(x: torch.Tensor, local_seq: int, group: dist.ProcessGroup) -> torch.Tensor: + """ReduceScatter on dim 0: [S, B, H] → [S/tp, B, H].""" + out = torch.empty(local_seq, x.shape[1], x.shape[2], dtype=x.dtype, device=x.device) + dist.reduce_scatter_tensor(out, x.contiguous(), group=group) + return out + + +class AllGatherDim0(torch.autograd.Function): + """AllGather [S/tp, B, H] → [S, B, H]. Backward: ReduceScatter.""" + + @staticmethod + def forward(ctx, x, tp_size, tp_rank, group): + ctx.group = group + ctx.local_seq = x.shape[0] + return _ag_dim0(x, tp_size, group) + + @staticmethod + def backward(ctx, grad): + out = _rs_dim0(grad, ctx.local_seq, ctx.group) + return out, None, None, None + + +class ReduceScatterDim0(torch.autograd.Function): + """ReduceScatter [S, B, H] → [S/tp, B, H]. Backward: AllGather.""" + + @staticmethod + def forward(ctx, x, tp_size, tp_rank, group): + ctx.tp_size = tp_size + ctx.group = group + local_seq = x.shape[0] // tp_size + return _rs_dim0(x, local_seq, ctx.group) + + @staticmethod + def backward(ctx, grad): + return _ag_dim0(grad, ctx.tp_size, ctx.group), None, None, None + + +class AllGatherDim0ForNonSPConsumer(torch.autograd.Function): + """AllGather [S/tp, B, H] → [S, B, H]. Backward: Scatter (no reduce).""" + + @staticmethod + def forward(ctx, x, tp_size, tp_rank, group): + ctx.tp_rank = tp_rank + ctx.local_seq = x.shape[0] + return _ag_dim0(x, tp_size, group) + + @staticmethod + def backward(ctx, grad): + start = ctx.tp_rank * ctx.local_seq + return grad[start : start + ctx.local_seq].contiguous(), None, None, None + + +class ScatterToSP(torch.autograd.Function): + """Scatter [S, B, H] → [S/tp, B, H] (no comm). Backward: AllGather.""" + + @staticmethod + def forward(ctx, x, tp_size, tp_rank, group): + ctx.tp_size = tp_size + ctx.group = group + local_seq = x.shape[0] // tp_size + start = tp_rank * local_seq + return x[start : start + local_seq, :, :].contiguous() + + @staticmethod + def backward(ctx, grad): + return _ag_dim0(grad, ctx.tp_size, ctx.group), None, None, None + + +__all__ = ["AllGatherDim0", "AllGatherDim0ForNonSPConsumer", "ReduceScatterDim0", "ScatterToSP"] diff --git a/experimental/lite/megatron/lite/primitive/optimizers/__init__.py b/experimental/lite/megatron/lite/primitive/optimizers/__init__.py new file mode 100644 index 00000000000..5bbfdd7cd00 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/optimizers/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Optimizer backend registry.""" + +from __future__ import annotations + +import importlib + +BACKENDS = { + "mc": "megatron.lite.primitive.optimizers.megatron_wrap", + "fsdp2": "megatron.lite.primitive.optimizers.fsdp2", +} + + +def get_optimizer_backend(name: str): + if name not in BACKENDS: + raise ValueError(f"Unknown Megatron Lite optimizer backend: {name!r}.") + return importlib.import_module(BACKENDS[name]).BACKEND + + +__all__ = ["get_optimizer_backend"] diff --git a/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/__init__.py b/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/__init__.py new file mode 100644 index 00000000000..b6dd3404d14 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/__init__.py @@ -0,0 +1,54 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""FSDP2 primitive surface.""" + +from __future__ import annotations + +from megatron.lite.primitive.optimizers.fsdp2.grad_clip import ( + all_reduce_scalar_, + clip_grads_with_sharded_norm_, + resolve_torch_dtype, + sharded_grad_abs_max, + sharded_grad_norm, + sharded_grad_sq_sum, +) +from megatron.lite.primitive.optimizers.fsdp2.optimizer import ( + BACKEND, + FSDP2Optimizer, + FSDP2OptimizerBackend, + build_fsdp2_adamw, + build_fsdp2_training_optimizer, +) +from megatron.lite.primitive.optimizers.fsdp2.wrap import ( + FSDP2Config, + build_fsdp2_device_mesh, + build_fsdp2_process_group_mesh, + build_fsdp2_shard_placement_fn, + fsdp2_available, + promote_fsdp2_trainable_params_to_fp32, + set_fsdp2_requires_gradient_sync, + wrap_fsdp2, + wrap_fsdp2_module, +) + +__all__ = [ + "BACKEND", + "FSDP2Config", + "FSDP2Optimizer", + "FSDP2OptimizerBackend", + "all_reduce_scalar_", + "build_fsdp2_adamw", + "build_fsdp2_training_optimizer", + "build_fsdp2_device_mesh", + "build_fsdp2_process_group_mesh", + "build_fsdp2_shard_placement_fn", + "clip_grads_with_sharded_norm_", + "fsdp2_available", + "promote_fsdp2_trainable_params_to_fp32", + "resolve_torch_dtype", + "set_fsdp2_requires_gradient_sync", + "sharded_grad_abs_max", + "sharded_grad_norm", + "sharded_grad_sq_sum", + "wrap_fsdp2", + "wrap_fsdp2_module", +] diff --git a/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/adamw.py b/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/adamw.py new file mode 100644 index 00000000000..84eb029b530 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/adamw.py @@ -0,0 +1,529 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""AdamW helpers for the FSDP2 optimizer primitive.""" + +from __future__ import annotations + +import inspect +from collections.abc import Callable, Iterable +from typing import Any + +import torch +import torch.distributed as dist +import torch.nn as nn + + +def local_grad_sq_sum( + params: Iterable[nn.Parameter], + *, + dtype: torch.dtype, + default_device: torch.device | None = None, +) -> torch.Tensor: + total: torch.Tensor | None = None + for param in params: + grad = param.grad + if grad is None: + continue + grad = to_local_tensor(grad) + if total is None: + total = torch.zeros((), device=grad.device, dtype=dtype) + total += grad.detach().to(dtype).pow(2).sum() + if total is None: + return torch.zeros((), device=default_device or torch.device("cpu"), dtype=dtype) + return total + + +def to_local_tensor(tensor): + local_tensor = getattr(tensor, "_local_tensor", None) + if isinstance(local_tensor, torch.Tensor): + return local_tensor + to_local = getattr(tensor, "to_local", None) + if callable(to_local): + return to_local() + return tensor + + +def fsdp2_model_param_dtype(param: nn.Parameter) -> torch.dtype | None: + dtype = getattr(param, "_fsdp2_model_param_dtype", None) + return dtype if isinstance(dtype, torch.dtype) else None + + +def has_dtensor_grad_or_param(param: nn.Parameter) -> bool: + grad = param.grad + return is_dtensor_like(param) or (grad is not None and is_dtensor_like(grad)) + + +def is_dtensor_like(tensor: Any) -> bool: + return ( + callable(getattr(tensor, "to_local", None)) + and hasattr(tensor, "device_mesh") + and hasattr(tensor, "placements") + ) + + +def copy_local_tensor_to_param_(param: nn.Parameter, local_tensor: torch.Tensor) -> None: + if not is_dtensor_like(param): + param.detach().copy_(local_tensor.to(device=param.device, dtype=param.dtype)) + return + + local_param = to_local_tensor(param) + local_value = local_tensor.to(device=local_param.device, dtype=local_param.dtype) + from torch.distributed.tensor import DTensor + + param.detach().copy_(DTensor.from_local(local_value, param.device_mesh, param.placements)) + + +def all_reduce_grad_(grad: torch.Tensor, *, group: dist.ProcessGroup) -> None: + local_grad = to_local_tensor(grad) + dist.all_reduce(local_grad, op=dist.ReduceOp.SUM, group=group) + if local_grad is grad: + return + from torch.distributed.tensor import DTensor + + grad.copy_(DTensor.from_local(local_grad, grad.device_mesh, grad.placements)) + + +class ChainedOptimizer: + def __init__(self, optimizers: Iterable[torch.optim.Optimizer]): + self.optimizers = list(optimizers) + + @property + def param_groups(self) -> list[dict[str, Any]]: + groups: list[dict[str, Any]] = [] + for optimizer in self.optimizers: + groups.extend(optimizer.param_groups) + return groups + + def zero_grad(self, *args, **kwargs) -> None: + for optimizer in self.optimizers: + optimizer.zero_grad(*args, **kwargs) + + def step(self) -> None: + for optimizer in self.optimizers: + optimizer.step() + + def state_dict(self) -> dict[str, Any]: + return { + "type": "chained_torch_optimizer", + "optimizers": [optimizer.state_dict() for optimizer in self.optimizers], + } + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + optimizer_states = state_dict.get("optimizers") + if not isinstance(optimizer_states, list) or len(optimizer_states) != len(self.optimizers): + raise ValueError("Invalid chained torch optimizer state_dict.") + for optimizer, optimizer_state in zip(self.optimizers, optimizer_states, strict=True): + optimizer.load_state_dict(optimizer_state) + + +class FP32AdamW: + """AdamW with FP32 master params for BF16/DTensor model weights.""" + + def __init__( + self, + params: Iterable[nn.Parameter] | Iterable[dict[str, Any]], + *, + lr: float, + weight_decay: float, + betas: tuple[float, float], + eps: float, + cpu_update: bool = False, + model_param_dtypes: dict[int, torch.dtype] | None = None, + ): + self.param_groups = normalize_param_groups(params, default_weight_decay=weight_decay) + self.params: list[nn.Parameter] = [] + self.lr = lr + self.weight_decay = weight_decay + self.betas = betas + self.eps = eps + self.cpu_update = bool(cpu_update) + self.step_count = 0 + self.state: dict[nn.Parameter, dict[str, torch.Tensor]] = {} + self._master_for_param: dict[nn.Parameter, torch.Tensor] = {} + self._model_param_dtypes_by_id = dict(model_param_dtypes or {}) + self._model_dtype_for_param: dict[nn.Parameter, torch.dtype] = {} + + for group in self.param_groups: + group.setdefault("lr", lr) + group.setdefault("wd_mult", 1.0) + group_weight_decay = float(group.get("weight_decay", weight_decay)) + group["weight_decay"] = group_weight_decay + for param in group["params"]: + self.params.append(param) + model_dtype = self._model_param_dtypes_by_id.get(id(param)) + if model_dtype is not None: + self._model_dtype_for_param[param] = model_dtype + master = self._init_master_param(param) + self.state[param] = { + "master_param": master, + "exp_avg": torch.zeros_like(master, dtype=torch.float32), + "exp_avg_sq": torch.zeros_like(master, dtype=torch.float32), + "step": 0, + } + self._master_for_param[param] = master + + def _init_master_param(self, param: nn.Parameter) -> torch.Tensor: + if self.cpu_update: + local_param = to_local_tensor(param.detach()) + return local_param.detach().to(device="cpu", dtype=torch.float32).clone() + if self._model_param_dtype(param) is not None: + return param.detach().to(dtype=torch.float32).clone() + return ( + param.detach() + if param.dtype is torch.float32 + else param.detach().to(dtype=torch.float32).clone() + ) + + def _model_param_dtype(self, param: nn.Parameter) -> torch.dtype | None: + return self._model_dtype_for_param.get(param) or fsdp2_model_param_dtype(param) + + def zero_grad(self, *args, **kwargs) -> None: + set_to_none = kwargs.get("set_to_none", False) + if args: + set_to_none = bool(args[0]) + for param in self.params: + if set_to_none: + param.grad = None + elif param.grad is not None: + param.grad.detach_() + param.grad.zero_() + + def step(self) -> None: + self._step_param_groups() + + def _step_param_groups(self) -> None: + self.step_count += 1 + beta1, beta2 = self.betas + + for group in self.param_groups: + group_lr = float(group.get("lr", self.lr)) + group_weight_decay = float(group.get("weight_decay", self.weight_decay)) + for param in group["params"]: + grad = param.grad + if grad is None: + continue + state = self.state[param] + state["step"] = int(state["step"]) + 1 + param_step = int(state["step"]) + bias_correction1 = 1.0 - beta1**param_step + bias_correction2_sqrt = (1.0 - beta2**param_step) ** 0.5 + group_step_size = group_lr / bias_correction1 + master = state["master_param"] + exp_avg = state["exp_avg"] + exp_avg_sq = state["exp_avg_sq"] + if group_weight_decay != 0.0: + master.mul_(1.0 - group_lr * group_weight_decay) + grad = self._prepare_grad(grad, master) + exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2) + denom = exp_avg_sq.sqrt().div_(bias_correction2_sqrt).add_(self.eps) + master.addcdiv_(exp_avg.to(dtype=torch.float32), denom, value=-group_step_size) + self._copy_master_to_param(param, master) + + def _prepare_grad(self, grad: torch.Tensor, master: torch.Tensor) -> torch.Tensor: + if self.cpu_update: + grad = to_local_tensor(grad) + return grad.detach().to(device=master.device, dtype=torch.float32) + return grad.detach().to(dtype=torch.float32) + + def _copy_master_to_param(self, param: nn.Parameter, master: torch.Tensor) -> None: + model_dtype = self._model_param_dtype(param) + if model_dtype is not None: + master = master.to(dtype=model_dtype).to(dtype=param.dtype) + if not self.cpu_update: + param.detach().copy_(master.to(dtype=param.dtype)) + return + copy_local_tensor_to_param_(param, master) + + def state_dict(self) -> dict[str, Any]: + return { + "type": "fp32_adamw", + "step_count": self.step_count, + "master_params": [self.state[param]["master_param"] for param in self.params], + "exp_avgs": [self.state[param]["exp_avg"] for param in self.params], + "exp_avg_sqs": [self.state[param]["exp_avg_sq"] for param in self.params], + "steps": [int(self.state[param]["step"]) for param in self.params], + "weight_decays": [ + float(group.get("weight_decay", self.weight_decay)) + for group in self.param_groups + for _param in group["params"] + ], + } + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + if state_dict.get("type") != "fp32_adamw": + raise ValueError("Invalid FP32 AdamW state_dict.") + self.step_count = int(state_dict.get("step_count", 0)) + for target_name, key in ( + ("master_params", "master_param"), + ("exp_avgs", "exp_avg"), + ("exp_avg_sqs", "exp_avg_sq"), + ): + loaded = state_dict.get(target_name) + if not isinstance(loaded, list) or len(loaded) != len(self.params): + raise ValueError(f"Invalid FP32 AdamW {target_name} state.") + for param, src in zip(self.params, loaded, strict=True): + self.state[param][key].copy_(src) + loaded_steps = state_dict.get("steps") + if loaded_steps is not None: + if not isinstance(loaded_steps, list) or len(loaded_steps) != len(self.params): + raise ValueError("Invalid FP32 AdamW steps state.") + for param, step in zip(self.params, loaded_steps, strict=True): + self.state[param]["step"] = int(step) + else: + for param in self.params: + self.state[param]["step"] = self.step_count + loaded_weight_decays = state_dict.get("weight_decays") + if loaded_weight_decays is not None: + if not isinstance(loaded_weight_decays, list) or len(loaded_weight_decays) != len( + self.params + ): + raise ValueError("Invalid FP32 AdamW weight_decay state.") + idx = 0 + for group in self.param_groups: + if not group["params"]: + continue + group["weight_decay"] = float(loaded_weight_decays[idx]) + idx += len(group["params"]) + + +def build_adamw_optimizer( + params: Iterable[nn.Parameter] | Iterable[dict[str, Any]], + *, + all_params: Iterable[nn.Parameter], + lr: float, + weight_decay: float, + betas: tuple[float, float], + eps: float, + foreach: bool | str, + use_fp32_master: bool, + cpu_update: bool, + model_param_dtypes: dict[int, torch.dtype] | None, + opt, +) -> Any: + param_groups = normalize_param_groups(params, default_weight_decay=weight_decay) + fused_adam = maybe_build_te_fused_adam_optimizer( + param_groups, + all_params=all_params, + lr=lr, + weight_decay=weight_decay, + betas=betas, + eps=eps, + opt=opt, + use_fp32_master=use_fp32_master, + ) + if fused_adam is not None: + return fused_adam + if use_fp32_master: + return FP32AdamW( + param_groups, + lr=lr, + weight_decay=weight_decay, + betas=betas, + eps=eps, + cpu_update=cpu_update, + model_param_dtypes=model_param_dtypes, + ) + if foreach not in {True, False, "auto"}: + raise ValueError(f"adamw_foreach must be True, False, or 'auto', got {foreach!r}.") + if foreach is False: + return torch.optim.AdamW( + param_groups, lr=lr, weight_decay=weight_decay, betas=betas, eps=eps, foreach=False + ) + + dtensor_param_groups, tensor_param_groups = split_dtensor_and_tensor_param_groups( + param_groups, default_weight_decay=weight_decay + ) + split_param_groups = [group for group in (dtensor_param_groups, tensor_param_groups) if group] + if foreach == "auto" and not dtensor_param_groups: + return torch.optim.AdamW( + param_groups, lr=lr, weight_decay=weight_decay, betas=betas, eps=eps, foreach=False + ) + if len(split_param_groups) <= 1: + return torch.optim.AdamW( + split_param_groups[0] if split_param_groups else param_groups, + lr=lr, + weight_decay=weight_decay, + betas=betas, + eps=eps, + foreach=True, + ) + return ChainedOptimizer( + torch.optim.AdamW( + group, lr=lr, weight_decay=weight_decay, betas=betas, eps=eps, foreach=True + ) + for group in split_param_groups + ) + + +def maybe_build_te_fused_adam_optimizer( + param_groups: list[dict[str, Any]], + *, + all_params: Iterable[nn.Parameter], + lr: float, + weight_decay: float, + betas: tuple[float, float], + eps: float, + opt, + use_fp32_master: bool, +) -> Any | None: + if not get_bool_opt(opt, "fsdp2_use_te_fused_adam", default=False): + return None + try: + from transformer_engine.pytorch.optimizers.fused_adam import FusedAdam + except ImportError: + return None + + all_param_list = list(all_params) + master_weights = get_bool_opt( + opt, "master_weights", default=use_fp32_master and should_use_master_weights(all_param_list) + ) + kwargs = dict( + lr=lr, + weight_decay=weight_decay, + betas=betas, + eps=eps, + adam_w_mode=True, + master_weights=master_weights, + master_weight_dtype=get_dtype_opt(opt, "master_weight_dtype", default=torch.float32), + store_param_remainders=get_bool_opt(opt, "store_param_remainders", default=master_weights), + exp_avg_dtype=get_dtype_opt(opt, "exp_avg_dtype", default=torch.float32), + exp_avg_sq_dtype=get_dtype_opt(opt, "exp_avg_sq_dtype", default=torch.float32), + ) + return FusedAdam(param_groups, **filter_supported_kwargs(FusedAdam.__init__, kwargs)) + + +def filter_supported_kwargs(fn: Callable[..., Any], kwargs: dict[str, Any]) -> dict[str, Any]: + try: + params = inspect.signature(fn).parameters + except (TypeError, ValueError): + return kwargs + if any(param.kind is inspect.Parameter.VAR_KEYWORD for param in params.values()): + return kwargs + return {key: value for key, value in kwargs.items() if key in params} + + +def should_use_master_weights(params: Iterable[nn.Parameter]) -> bool: + return any(param.is_floating_point() and param.dtype is not torch.float32 for param in params) + + +def get_bool_opt(opt, attr: str, *, default: bool) -> bool: + value = get_opt_value(opt, attr) + if value is None: + return bool(default) + if isinstance(value, bool): + return value + if isinstance(value, (int, float)): + return bool(value) + return str(value).strip().lower() in {"1", "true", "yes", "on"} + + +def get_dtype_opt(opt, attr: str, *, default: torch.dtype) -> torch.dtype: + value = get_opt_value(opt, attr) + if value is None: + return default + if isinstance(value, torch.dtype): + return value + name = str(value).removeprefix("torch.") + resolved = getattr(torch, name, None) + if not isinstance(resolved, torch.dtype): + raise ValueError(f"Unsupported dtype for FSDP2 TE FusedAdam: {value!r}.") + return resolved + + +def get_opt_value(opt, attr: str): + if opt is None: + return None + if isinstance(opt, dict): + value = opt.get(attr) + override = opt.get("override_optimizer_config") + else: + value = getattr(opt, attr, None) + override = getattr(opt, "override_optimizer_config", None) + if value is not None: + return value + if isinstance(override, dict): + return override.get(attr) + return None + + +def normalize_param_groups( + params: Iterable[nn.Parameter] | Iterable[dict[str, Any]], *, default_weight_decay: float +) -> list[dict[str, Any]]: + items = list(params) + if not items: + return [] + if all(isinstance(item, dict) for item in items): + groups: list[dict[str, Any]] = [] + for item in items: + group = dict(item) + group_params = list(group.get("params", ())) + if not group_params: + continue + group["params"] = group_params + group.setdefault("weight_decay", default_weight_decay) + groups.append(group) + return groups + return [{"params": items, "weight_decay": default_weight_decay}] + + +def split_dtensor_and_tensor_param_groups( + param_groups: Iterable[dict[str, Any]], *, default_weight_decay: float +) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + dtensor_groups: list[dict[str, Any]] = [] + tensor_groups: list[dict[str, Any]] = [] + for group in param_groups: + dtensor_params, tensor_params = split_dtensor_and_tensor_params(group["params"]) + metadata = {key: value for key, value in group.items() if key != "params"} + metadata.setdefault("weight_decay", default_weight_decay) + if dtensor_params: + dtensor_groups.append({**metadata, "params": dtensor_params}) + if tensor_params: + tensor_groups.append({**metadata, "params": tensor_params}) + return dtensor_groups, tensor_groups + + +def split_dtensor_and_tensor_params( + params: Iterable[nn.Parameter], +) -> tuple[list[nn.Parameter], list[nn.Parameter]]: + dtensor_params: list[nn.Parameter] = [] + tensor_params: list[nn.Parameter] = [] + for param in params: + if is_dtensor_like(param): + dtensor_params.append(param) + else: + tensor_params.append(param) + return dtensor_params, tensor_params + + +def iter_torch_optimizers(optimizer: Any) -> Iterable[torch.optim.Optimizer]: + if isinstance(optimizer, ChainedOptimizer): + yield from optimizer.optimizers + else: + yield optimizer + + +def dtensor_from_local( + local_tensor: torch.Tensor, device_mesh: Any, placements: Any +) -> torch.Tensor: + from torch.distributed.tensor import DTensor + + return DTensor.from_local(local_tensor, device_mesh, placements) + + +__all__ = [ + "all_reduce_grad_", + "build_adamw_optimizer", + "copy_local_tensor_to_param_", + "dtensor_from_local", + "filter_supported_kwargs", + "fsdp2_model_param_dtype", + "get_bool_opt", + "get_dtype_opt", + "get_opt_value", + "has_dtensor_grad_or_param", + "is_dtensor_like", + "iter_torch_optimizers", + "local_grad_sq_sum", + "normalize_param_groups", + "to_local_tensor", +] diff --git a/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/grad_clip.py b/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/grad_clip.py new file mode 100644 index 00000000000..f14bdd7fd83 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/grad_clip.py @@ -0,0 +1,335 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Sharded gradient norm and clipping primitive for FSDP2 optimizers.""" + +from __future__ import annotations + +import math +from collections import defaultdict +from collections.abc import Callable, Iterable +from numbers import Number +from typing import Any + +import torch +import torch.distributed as dist +import torch.nn as nn + +try: # pragma: no cover - import availability is PyTorch-version dependent. + from torch.distributed.tensor import DTensor, Partial, Replicate +except ImportError: # pragma: no cover + DTensor = Partial = Replicate = None # type: ignore[assignment] + + +def sharded_grad_sq_sum( + params: Iterable[nn.Parameter], + *, + accum_dtype: str | torch.dtype = torch.float32, + default_device: torch.device | None = None, + chunk_size_numel: int = 0, + scalar_all_reduce: ( + Callable[[torch.Tensor, dist.ProcessGroup, dist.ReduceOp], None] | None + ) = None, +) -> torch.Tensor: + """Return global L2 grad squared-sum for Tensor/DTensor parameters. + + The primitive reduces one scalar per DTensor sharding group over mesh + dimensions whose placement is not replicated. Pipeline/expert reductions + are intentionally left to the runtime adapter because they are model-layout + policy, not a DTensor property. + """ + + dtype = resolve_torch_dtype(accum_dtype) + groups = _group_grads(params) + total: torch.Tensor | None = None + for group in groups.values(): + local_sq = _group_local_sq_sum(group, dtype=dtype, chunk_size_numel=chunk_size_numel) + meta = group[0][2] + if meta is not None and not _has_partial_placement(meta) and dist.is_initialized(): + _reduce_dtensor_scalar_( + local_sq, meta, op=dist.ReduceOp.SUM, scalar_all_reduce=scalar_all_reduce + ) + total = local_sq if total is None else total.to(local_sq.device) + local_sq + + if total is None: + return torch.zeros((), device=default_device or torch.device("cpu"), dtype=dtype) + return total + + +def sharded_grad_norm( + params: Iterable[nn.Parameter], + *, + norm_type: float = 2.0, + pp_group: dist.ProcessGroup | None = None, + accum_dtype: str | torch.dtype = torch.float32, + default_device: torch.device | None = None, +) -> torch.Tensor: + """Return global grad norm for Tensor/DTensor parameters. + + Only L2 and infinity norms are implemented because they cover Megatron Lite's + optimizer path today and avoid ambiguous cross-placement semantics for + arbitrary p-norms. + """ + + if math.isinf(float(norm_type)): + total = sharded_grad_abs_max( + params, pp_group=pp_group, accum_dtype=accum_dtype, default_device=default_device + ) + return total + if float(norm_type) != 2.0: + raise ValueError(f"sharded_grad_norm supports norm_type=2.0 or inf, got {norm_type!r}.") + sq_sum = sharded_grad_sq_sum(params, accum_dtype=accum_dtype, default_device=default_device) + if pp_group is not None and dist.is_initialized() and dist.get_world_size(pp_group) > 1: + all_reduce_scalar_(sq_sum, op=dist.ReduceOp.SUM, group=pp_group) + return sq_sum.sqrt() + + +def sharded_grad_abs_max( + params: Iterable[nn.Parameter], + *, + pp_group: dist.ProcessGroup | None = None, + accum_dtype: str | torch.dtype = torch.float32, + default_device: torch.device | None = None, +) -> torch.Tensor: + """Return global infinity grad norm for Tensor/DTensor parameters.""" + + dtype = resolve_torch_dtype(accum_dtype) + groups = _group_grads(params) + total: torch.Tensor | None = None + for group in groups.values(): + local_max = _group_local_abs_max(group, dtype=dtype) + meta = group[0][2] + if meta is not None and not _has_partial_placement(meta) and dist.is_initialized(): + _reduce_dtensor_scalar_(local_max, meta, op=dist.ReduceOp.MAX) + total = local_max if total is None else torch.maximum(total.to(local_max.device), local_max) + + if total is None: + total = torch.zeros((), device=default_device or torch.device("cpu"), dtype=dtype) + if pp_group is not None and dist.is_initialized() and dist.get_world_size(pp_group) > 1: + all_reduce_scalar_(total, op=dist.ReduceOp.MAX, group=pp_group) + return total + + +def all_reduce_scalar_( + value: torch.Tensor, + *, + op: dist.ReduceOp, + group: dist.ProcessGroup, +) -> None: + """All-reduce a scalar on a device compatible with the process group backend.""" + + reduced = _scalar_for_process_group(value, group) + dist.all_reduce(reduced, op=op, group=group) + if reduced is not value: + value.copy_(reduced.to(device=value.device, dtype=value.dtype)) + + +@torch.no_grad() +def clip_grads_with_sharded_norm_( + params: Iterable[nn.Parameter], max_norm: float, total_norm: torch.Tensor | float +) -> None: + """Scale gradients in-place using a precomputed global norm.""" + + max_norm = float(max_norm) + if max_norm <= 0: + return + if isinstance(total_norm, torch.Tensor): + if not bool(torch.isfinite(total_norm).item()): + return + clip_coef = (max_norm / (total_norm + 1.0e-6)).clamp(max=1.0) + if float(clip_coef.item()) >= 1.0: + return + else: + norm_value = float(total_norm) + if not math.isfinite(norm_value): + return + clip_coef = max_norm / (norm_value + 1.0e-6) + if clip_coef >= 1.0: + return + for param in params: + if param.grad is not None: + _scale_grad_(param.grad, clip_coef) + + +def resolve_torch_dtype(dtype: str | torch.dtype) -> torch.dtype: + if isinstance(dtype, torch.dtype): + resolved = dtype + else: + name = dtype.removeprefix("torch.") + resolved = getattr(torch, name, None) + if not isinstance(resolved, torch.dtype): + raise ValueError(f"Unsupported torch dtype for grad norm accumulation: {dtype!r}") + if not torch.empty((), dtype=resolved).is_floating_point(): + raise ValueError(f"Grad norm accumulation dtype must be floating point: {dtype!r}") + return resolved + + +def _group_grads( + params: Iterable[nn.Parameter], +) -> dict[tuple[Any, ...], list[tuple[nn.Parameter, torch.Tensor, Any | None]]]: + groups: dict[tuple[Any, ...], list[tuple[nn.Parameter, torch.Tensor, Any | None]]] = ( + defaultdict(list) + ) + for param in params: + grad = param.grad + if grad is None: + continue + meta = _dtensor_meta(param, grad) + if meta is None: + key = ("tensor", grad.device) + else: + key = ( + "dtensor", + id(meta.device_mesh), + tuple((type(placement).__name__, repr(placement)) for placement in meta.placements), + ) + groups[key].append((param, grad, meta)) + return groups + + +def _group_local_sq_sum( + group: list[tuple[nn.Parameter, torch.Tensor, Any | None]], + *, + dtype: torch.dtype, + chunk_size_numel: int = 0, +) -> torch.Tensor: + device = _local_grad(group[0][1], group[0][2]).device + total = torch.zeros((), device=device, dtype=dtype) + for _param, grad, meta in group: + local_grad = _local_grad(grad, meta) + total += _tensor_sq_sum(local_grad.detach(), dtype=dtype, chunk_size_numel=chunk_size_numel) + return total + + +def _tensor_sq_sum( + tensor: torch.Tensor, *, dtype: torch.dtype, chunk_size_numel: int = 0 +) -> torch.Tensor: + if chunk_size_numel <= 0 or tensor.numel() <= chunk_size_numel: + return tensor.to(dtype).pow(2).sum() + try: + flat = tensor.view(-1) + except RuntimeError: + flat = tensor.reshape(-1) + total = torch.zeros((), device=tensor.device, dtype=dtype) + for start in range(0, flat.numel(), chunk_size_numel): + chunk = flat.narrow(0, start, min(chunk_size_numel, flat.numel() - start)) + total += chunk.to(dtype).pow(2).sum() + return total + + +def _group_local_abs_max( + group: list[tuple[nn.Parameter, torch.Tensor, Any | None]], *, dtype: torch.dtype +) -> torch.Tensor: + device = _local_grad(group[0][1], group[0][2]).device + total = torch.zeros((), device=device, dtype=dtype) + for _param, grad, meta in group: + local_grad = _local_grad(grad, meta) + if local_grad.numel() > 0: + total = torch.maximum(total, local_grad.detach().to(dtype).abs().max()) + return total + + +def _local_grad(grad: torch.Tensor, meta: Any | None) -> torch.Tensor: + if meta is not None and _has_partial_placement(meta): + full_tensor = getattr(grad, "full_tensor", None) + if callable(full_tensor): + return full_tensor() + to_local = getattr(grad, "to_local", None) + if callable(to_local): + return to_local() + return grad + + +def _scale_grad_(grad: torch.Tensor, scale: float | torch.Tensor) -> None: + to_local = getattr(grad, "to_local", None) + if not callable(to_local): + grad.mul_(_scale_for_tensor(scale, grad)) + return + local_grad = to_local() + local_grad.mul_(_scale_for_tensor(scale, local_grad)) + if DTensor is not None and isinstance(grad, DTensor): + grad.copy_(DTensor.from_local(local_grad, grad.device_mesh, grad.placements)) + + +def _scale_for_tensor(scale: float | torch.Tensor, tensor: torch.Tensor) -> float | torch.Tensor: + if isinstance(scale, Number): + return float(scale) + return scale.to(device=tensor.device, dtype=tensor.dtype) + + +def _dtensor_meta(param: nn.Parameter, grad: torch.Tensor) -> Any | None: + if _is_dtensor_like(grad): + return grad + if _is_dtensor_like(param): + return param + return None + + +def _is_dtensor_like(tensor: Any) -> bool: + if DTensor is not None and isinstance(tensor, DTensor): + return True + return ( + callable(getattr(tensor, "to_local", None)) + and hasattr(tensor, "device_mesh") + and hasattr(tensor, "placements") + ) + + +def _has_partial_placement(dtensor: Any) -> bool: + return any(_placement_name(placement) == "Partial" for placement in dtensor.placements) + + +def _is_replicate_placement(placement: Any) -> bool: + return _placement_name(placement) == "Replicate" + + +def _placement_name(placement: Any) -> str: + return type(placement).__name__ + + +def _reduce_dtensor_scalar_( + value: torch.Tensor, + dtensor: Any, + *, + op: dist.ReduceOp, + scalar_all_reduce: ( + Callable[[torch.Tensor, dist.ProcessGroup, dist.ReduceOp], None] | None + ) = None, +) -> None: + for mesh_dim, placement in enumerate(dtensor.placements): + if _is_replicate_placement(placement): + continue + group = dtensor.device_mesh.get_group(mesh_dim) + if dist.get_world_size(group) > 1: + if scalar_all_reduce is None: + all_reduce_scalar_(value, op=op, group=group) + else: + scalar_all_reduce(value, group, op) + + +def _scalar_for_process_group( + value: torch.Tensor, group: dist.ProcessGroup +) -> torch.Tensor: + backend = _process_group_backend(group) + if "nccl" in backend and value.device.type != "cuda": + if not torch.cuda.is_available(): + raise RuntimeError("NCCL scalar all_reduce requires a CUDA tensor.") + return value.to(device=torch.device("cuda", torch.cuda.current_device())) + if "gloo" in backend and value.device.type != "cpu": + return value.to(device=torch.device("cpu")) + return value + + +def _process_group_backend(group: dist.ProcessGroup) -> str: + try: + return str(dist.get_backend(group)).lower() + except (RuntimeError, ValueError, TypeError): + return "" + + +__all__ = [ + "all_reduce_scalar_", + "clip_grads_with_sharded_norm_", + "resolve_torch_dtype", + "sharded_grad_abs_max", + "sharded_grad_norm", + "sharded_grad_sq_sum", +] diff --git a/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/optimizer.py b/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/optimizer.py new file mode 100644 index 00000000000..63f2935c372 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/optimizer.py @@ -0,0 +1,680 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""FSDP2 optimizer adapter for Megatron Lite runtime contracts.""" + +from __future__ import annotations + +import math +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any + +import torch +import torch.distributed as dist +import torch.nn as nn + +from megatron.lite.primitive.optimizers.fsdp2.adamw import ( + all_reduce_grad_, + build_adamw_optimizer, + fsdp2_model_param_dtype, + get_bool_opt, + has_dtensor_grad_or_param, + local_grad_sq_sum, +) +from megatron.lite.primitive.optimizers.fsdp2.grad_clip import ( + all_reduce_scalar_, + clip_grads_with_sharded_norm_, + resolve_torch_dtype, + sharded_grad_sq_sum, +) +from megatron.lite.primitive.optimizers.fsdp2.state import ( + OffloadedStateEntry, + move_offloaded_optimizer_state_to_device, + move_optimizer_state_to_cpu, +) +from megatron.lite.primitive.optimizers.fsdp2.wrap import ( + FSDP2Config, + build_fsdp2_process_group_mesh, + build_fsdp2_shard_placement_fn, + promote_fsdp2_trainable_params_to_fp32, + wrap_fsdp2, + wrap_fsdp2_module, +) +from megatron.lite.primitive.parallel.state import ParallelState + +_DEFAULT_RESHARD_AFTER_FORWARD: bool | int | None = True +_DEFAULT_WRAP_ROOT = True +_DEFAULT_LEAF_MODULE_NAMES = ("embed", "head") +_DEFAULT_FORWARD_PREFETCH_DEPTH = 1 +_DEFAULT_BACKWARD_PREFETCH_DEPTH = 0 +_DEFAULT_PARAM_DTYPE: str | None = "bfloat16" +_DEFAULT_REDUCE_DTYPE: str | None = "float32" +_DEFAULT_USE_FP32_SHARDS = True +_DEFAULT_USE_FP32_MASTER = True +_DEFAULT_ADAMW_FOREACH: bool | str = "auto" + + +class FSDP2Optimizer: + """Adapt an optimizer to Megatron Lite's FSDP2 optimizer contract.""" + + name = "fsdp2" + + def __init__( + self, + optimizer: Any, + params: Iterable[nn.Parameter], + ps: ParallelState | None = None, + *, + clip_grad: float = 1.0, + replicated_grad_params: Iterable[nn.Parameter] | None = None, + replicated_grad_sync_group: dist.ProcessGroup | None = None, + replicated_grad_sync_divisor: float | None = None, + replicated_grad_norm_group: dist.ProcessGroup | None = None, + expert_sharded_grad_params: Iterable[nn.Parameter] | None = None, + expert_sharded_grad_scale: float | None = None, + expert_sharded_grad_norm_group: dist.ProcessGroup | None = None, + tp_replicated_grad_params: Iterable[nn.Parameter] | None = None, + tp_replicated_grad_sync_group: dist.ProcessGroup | None = None, + grad_norm_accum_dtype: str | torch.dtype = torch.float32, + param_names: dict[int, str] | None = None, + ): + self.optimizer = optimizer + self.params = list(params) + self.param_names = dict(param_names or {}) + self.ps = ps + self.clip_grad = float(clip_grad) + self.grad_norm_accum_dtype = resolve_torch_dtype(grad_norm_accum_dtype) + self.replicated_grad_params = list(replicated_grad_params or ()) + self._replicated_grad_param_ids = {id(param) for param in self.replicated_grad_params} + self.replicated_grad_sync_group = replicated_grad_sync_group + self.replicated_grad_sync_divisor = replicated_grad_sync_divisor + self.replicated_grad_norm_group = replicated_grad_norm_group + self.expert_sharded_grad_params = list(expert_sharded_grad_params or ()) + self._expert_sharded_grad_param_ids = { + id(param) for param in self.expert_sharded_grad_params + } + self.expert_sharded_grad_scale = ( + 1.0 if expert_sharded_grad_scale is None else float(expert_sharded_grad_scale) + ) + self.expert_sharded_grad_norm_group = expert_sharded_grad_norm_group + self.tp_replicated_grad_params = list(tp_replicated_grad_params or ()) + self._tp_replicated_grad_param_ids = {id(param) for param in self.tp_replicated_grad_params} + self.tp_replicated_grad_sync_group = tp_replicated_grad_sync_group + self._cpu_offloaded_state: dict[tuple[int, str], OffloadedStateEntry] = {} + self.grad_sync_enabled = False + + @property + def param_groups(self): + return self.optimizer.param_groups + + def zero_grad(self) -> None: + self.grad_sync_enabled = False + self.optimizer.zero_grad(set_to_none=True) + + def step(self) -> tuple[bool, float, int]: + self.sync_tp_replicated_grads() + self.sync_replicated_grads() + self.scale_expert_sharded_grads() + grad_norm = self.clip_grad_norm() + if not math.isfinite(grad_norm): + self.grad_sync_enabled = False + return False, float(grad_norm), 0 + self.optimizer.step() + self.grad_sync_enabled = False + return True, float(grad_norm), 0 + + def sync_replicated_grads(self) -> None: + group = self.replicated_grad_sync_group + if not self.replicated_grad_params: + return + group_size = 1 + if group is not None and dist.is_initialized(): + group_size = dist.get_world_size(group) + if group is None and self.replicated_grad_sync_divisor is None: + return + divisor = self.replicated_grad_sync_divisor + if divisor is None: + divisor = float(group_size) + for param in self.replicated_grad_params: + grad = param.grad + if grad is None: + continue + if group is not None and dist.is_initialized() and group_size > 1: + dist.all_reduce(grad, op=dist.ReduceOp.SUM, group=group) + if divisor != 1.0: + grad.div_(divisor) + + def sync_tp_replicated_grads(self) -> None: + group = self.tp_replicated_grad_sync_group + if not self.tp_replicated_grad_params: + return + if group is None or not dist.is_initialized() or dist.get_world_size(group) <= 1: + return + for param in self.tp_replicated_grad_params: + grad = param.grad + if grad is not None: + all_reduce_grad_(grad, group=group) + + def scale_expert_sharded_grads(self) -> None: + if not self.expert_sharded_grad_params or self.expert_sharded_grad_scale == 1.0: + return + for param in self.expert_sharded_grad_params: + grad = param.grad + if grad is not None: + grad.mul_(self.expert_sharded_grad_scale) + + def clip_grad_norm(self) -> float: + excluded_sharded_param_ids = ( + self._replicated_grad_param_ids + | self._tp_replicated_grad_param_ids + | self._expert_sharded_grad_param_ids + ) + sharded_params = [ + param for param in self.params if id(param) not in excluded_sharded_param_ids + ] + dtensor_sharded_params = [ + param for param in sharded_params if has_dtensor_grad_or_param(param) + ] + plain_sharded_params = [ + param for param in sharded_params if not has_dtensor_grad_or_param(param) + ] + total_sq = sharded_grad_sq_sum( + dtensor_sharded_params, accum_dtype=self.grad_norm_accum_dtype + ) + plain_sharded_sq = local_grad_sq_sum( + plain_sharded_params, + dtype=resolve_torch_dtype(self.grad_norm_accum_dtype), + default_device=total_sq.device, + ) + plain_sharded_group = None + if self.ps is not None: + plain_sharded_group = self.ps.dp_cp_group or self.ps.dp_group + if ( + plain_sharded_group is not None + and dist.is_initialized() + and dist.get_world_size(plain_sharded_group) > 1 + ): + all_reduce_scalar_(plain_sharded_sq, op=dist.ReduceOp.SUM, group=plain_sharded_group) + total_sq = total_sq.to(plain_sharded_sq.device) + plain_sharded_sq + if ( + self.ps is not None + and self.ps.tp_group is not None + and dist.is_initialized() + and dist.get_world_size(self.ps.tp_group) > 1 + ): + all_reduce_scalar_(total_sq, op=dist.ReduceOp.SUM, group=self.ps.tp_group) + + tp_replicated_sq = sharded_grad_sq_sum( + self.tp_replicated_grad_params, + accum_dtype=self.grad_norm_accum_dtype, + default_device=total_sq.device, + ) + replicated_sq = local_grad_sq_sum( + self.replicated_grad_params, + dtype=self.grad_norm_accum_dtype, + default_device=total_sq.device, + ) + if ( + self.replicated_grad_norm_group is not None + and dist.is_initialized() + and dist.get_world_size(self.replicated_grad_norm_group) > 1 + ): + all_reduce_scalar_( + replicated_sq, op=dist.ReduceOp.SUM, group=self.replicated_grad_norm_group + ) + + expert_sharded_sq = sharded_grad_sq_sum( + self.expert_sharded_grad_params, + accum_dtype=self.grad_norm_accum_dtype, + default_device=total_sq.device, + ) + if ( + self.expert_sharded_grad_norm_group is not None + and dist.is_initialized() + and dist.get_world_size(self.expert_sharded_grad_norm_group) > 1 + ): + all_reduce_scalar_( + expert_sharded_sq, + op=dist.ReduceOp.SUM, + group=self.expert_sharded_grad_norm_group, + ) + + total_sq = ( + total_sq + + tp_replicated_sq.to(total_sq.device) + + replicated_sq.to(total_sq.device) + + expert_sharded_sq.to(total_sq.device) + ) + if ( + self.ps is not None + and self.ps.pp_group is not None + and dist.is_initialized() + and dist.get_world_size(self.ps.pp_group) > 1 + ): + all_reduce_scalar_(total_sq, op=dist.ReduceOp.SUM, group=self.ps.pp_group) + + grad_norm = total_sq.sqrt() + if torch.isfinite(grad_norm): + clip_grads_with_sharded_norm_(self.params, self.clip_grad, grad_norm) + return float(grad_norm.float().item()) + + def state_dict(self) -> dict[str, Any]: + return self.optimizer.state_dict() + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + self.optimizer.load_state_dict(state_dict) + + def offload_state_to_cpu(self) -> None: + move_optimizer_state_to_cpu( + self.optimizer, self._cpu_offloaded_state, include_dtensor_state=True + ) + + def load_state_to_device(self) -> None: + move_offloaded_optimizer_state_to_device(self.optimizer, self._cpu_offloaded_state) + + +def build_fsdp2_adamw( + model_chunks: list[nn.Module], + opt, + ps: ParallelState, + *, + replicated_grad_params: Iterable[nn.Parameter] | None = None, + replicated_grad_sync_group: dist.ProcessGroup | None = None, + replicated_grad_sync_divisor: float | None = None, + replicated_grad_norm_group: dist.ProcessGroup | None = None, + expert_sharded_grad_params: Iterable[nn.Parameter] | None = None, + expert_sharded_grad_scale: float | None = None, + expert_sharded_grad_norm_group: dist.ProcessGroup | None = None, + tp_replicated_grad_params: Iterable[nn.Parameter] | None = None, + tp_replicated_grad_sync_group: dist.ProcessGroup | None = None, + grad_norm_accum_dtype: str | torch.dtype = torch.float32, + adamw_foreach: bool | str = "auto", + use_fp32_master: bool = False, + model_param_dtypes: dict[tuple[int, str], torch.dtype] | None = None, +) -> FSDP2Optimizer: + """Build AdamW from Megatron Lite's shared OptimizerConfig-like object.""" + + optimizer_name = getattr(opt, "optimizer", "adam") + if optimizer_name not in {"adam", "adamw"}: + raise ValueError(f"fsdp2 supports adam/adamw, got {optimizer_name!r}.") + + params, param_groups, param_names, param_model_dtypes = _build_adamw_param_groups( + model_chunks, + weight_decay=float(getattr(opt, "weight_decay", 0.01)), + apply_wd_to_qk_layernorm=bool(getattr(opt, "apply_wd_to_qk_layernorm", False)), + model_param_dtypes=model_param_dtypes, + ) + beta1 = getattr(opt, "adam_beta1", None) + beta2 = getattr(opt, "adam_beta2", None) + eps = getattr(opt, "adam_eps", None) + offload_fraction = getattr(opt, "offload_fraction", None) or 0.0 + optimizer = build_adamw_optimizer( + param_groups, + all_params=params, + lr=float(getattr(opt, "lr", 1.0e-4)), + weight_decay=float(getattr(opt, "weight_decay", 0.01)), + betas=(0.9 if beta1 is None else beta1, 0.999 if beta2 is None else beta2), + eps=1.0e-8 if eps is None else eps, + foreach=adamw_foreach, + use_fp32_master=use_fp32_master, + cpu_update=use_fp32_master and float(offload_fraction) > 0.0, + model_param_dtypes=param_model_dtypes, + opt=opt, + ) + return FSDP2Optimizer( + optimizer, + params, + ps, + clip_grad=float(getattr(opt, "clip_grad", 1.0)), + replicated_grad_params=replicated_grad_params, + replicated_grad_sync_group=replicated_grad_sync_group, + replicated_grad_sync_divisor=replicated_grad_sync_divisor, + replicated_grad_norm_group=replicated_grad_norm_group, + expert_sharded_grad_params=expert_sharded_grad_params, + expert_sharded_grad_scale=expert_sharded_grad_scale, + expert_sharded_grad_norm_group=expert_sharded_grad_norm_group, + tp_replicated_grad_params=tp_replicated_grad_params, + tp_replicated_grad_sync_group=tp_replicated_grad_sync_group, + grad_norm_accum_dtype=grad_norm_accum_dtype, + param_names=param_names, + ) + + +def build_fsdp2_training_optimizer( + model_chunks: list[nn.Module], + opt, + ps: ParallelState, + *, + unit_modules: tuple[type[nn.Module] | str, ...], + expert_classifier: Callable[[str], bool] | None = None, + expert_module_leaf_name: str = "experts", + deterministic: bool | None = None, + vpp: int | None = 1, + leaf_module_names: Iterable[str] = _DEFAULT_LEAF_MODULE_NAMES, + reshard_after_forward: bool | int | None = _DEFAULT_RESHARD_AFTER_FORWARD, + wrap_root: bool = _DEFAULT_WRAP_ROOT, + forward_prefetch_depth: int = _DEFAULT_FORWARD_PREFETCH_DEPTH, + backward_prefetch_depth: int = _DEFAULT_BACKWARD_PREFETCH_DEPTH, + param_dtype: str | torch.dtype | None = _DEFAULT_PARAM_DTYPE, + reduce_dtype: str | torch.dtype | None = _DEFAULT_REDUCE_DTYPE, + use_fp32_shards: bool | None = None, + use_fp32_master: bool | None = None, + adamw_foreach: bool | str = _DEFAULT_ADAMW_FOREACH, +) -> FSDP2Optimizer: + """Wrap model chunks with FSDP2 and build the matching AdamW adapter.""" + + if (vpp or 1) > 1 and ps.pp_size <= 1: + raise ValueError("optimizer='fsdp2' requires pp>1 when vpp>1.") + + expert_params = _collect_expert_params(model_chunks, ps, expert_classifier) + expert_modules = _collect_expert_modules( + model_chunks, ps, expert_classifier, expert_module_leaf_name=expert_module_leaf_name + ) + if expert_params and not expert_modules: + raise RuntimeError("FSDP2 expert parameters were found but no expert module was found.") + + if opt is None: + opt = SimpleNamespace( + optimizer="adam", + lr=1e-4, + weight_decay=0.01, + clip_grad=1.0, + adam_beta1=None, + adam_beta2=None, + adam_eps=None, + ) + if deterministic is None: + from megatron.lite.primitive.deterministic import deterministic_requested + + deterministic = deterministic_requested() + effective_use_fp32_shards = ( + bool(use_fp32_shards) + if use_fp32_shards is not None + else get_bool_opt(opt, "fsdp2_use_fp32_shards", default=_DEFAULT_USE_FP32_SHARDS) + ) + effective_use_fp32_master = ( + bool(use_fp32_master) + if use_fp32_master is not None + else get_bool_opt(opt, "fsdp2_use_fp32_master", default=_DEFAULT_USE_FP32_MASTER) + ) + + unit_reshard_after_forward = _fsdp2_unit_reshard_after_forward( + ps, reshard_after_forward=reshard_after_forward + ) + fsdp2_config = FSDP2Config( + unit_modules=unit_modules, + leaf_module_names=tuple(leaf_module_names), + reshard_after_forward=unit_reshard_after_forward, + last_unit_reshard_after_forward=unit_reshard_after_forward, + root_reshard_after_forward=False, + wrap_root=wrap_root, + forward_prefetch_depth=_fsdp2_prefetch_depth(ps, default_depth=forward_prefetch_depth), + backward_prefetch_depth=_fsdp2_prefetch_depth(ps, default_depth=backward_prefetch_depth), + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + ) + if effective_use_fp32_shards: + model_param_dtypes = _collect_model_param_dtypes(model_chunks) + for chunk in model_chunks: + promote_fsdp2_trainable_params_to_fp32(chunk) + else: + model_param_dtypes = {} + + tp_replicated_grad_param_names = _collect_tp_replicated_grad_param_names(model_chunks) + + dense_shard_placement_fn = build_fsdp2_shard_placement_fn(ps.dp_cp_size) + expert_shard_placement_fn = build_fsdp2_shard_placement_fn(ps.expert_dp_size) + + if expert_modules: + if ps.ep_dp_group is None: + raise RuntimeError("FSDP2 expert sharding requires ParallelState.ep_dp_group.") + expert_mesh = build_fsdp2_process_group_mesh( + ps.ep_dp_group, mesh_dim_name="expert_dp", device_type=fsdp2_config.device_type + ) + for module in expert_modules: + wrap_fsdp2_module( + module, + ps, + fsdp2_config, + mesh=expert_mesh, + shard_placement_fn=expert_shard_placement_fn, + reshard_after_forward=unit_reshard_after_forward, + ) + + ignored_expert_params = _collect_module_params(expert_modules) + for chunk in model_chunks: + wrap_fsdp2( + chunk, + ps, + fsdp2_config, + ignored_params=ignored_expert_params or None, + shard_placement_fn=dense_shard_placement_fn, + ) + if model_param_dtypes: + _restore_model_param_dtypes(model_chunks, model_param_dtypes) + + tp_replicated_grad_params = _collect_tp_replicated_grad_params( + model_chunks, param_names=tp_replicated_grad_param_names + ) + return build_fsdp2_adamw( + model_chunks, + opt, + ps, + expert_sharded_grad_params=list(ignored_expert_params), + expert_sharded_grad_scale=( + float(ps.expert_dp_size) / float(ps.dp_cp_size) if ignored_expert_params else None + ), + expert_sharded_grad_norm_group=ps.ep_group if ignored_expert_params else None, + tp_replicated_grad_params=tp_replicated_grad_params, + tp_replicated_grad_sync_group=ps.tp_group if tp_replicated_grad_params else None, + grad_norm_accum_dtype="float32", + adamw_foreach=False if deterministic else adamw_foreach, + use_fp32_master=effective_use_fp32_master, + model_param_dtypes=model_param_dtypes, + ) + + +def _collect_expert_params( + chunks: Iterable[nn.Module], ps: ParallelState, expert_classifier: Callable[[str], bool] | None +) -> set[nn.Parameter]: + if ps.ep_size <= 1 or expert_classifier is None: + return set() + return { + param + for chunk in chunks + for name, param in chunk.named_parameters() + if expert_classifier(name) + } + + +def _collect_expert_modules( + chunks: Iterable[nn.Module], + ps: ParallelState, + expert_classifier: Callable[[str], bool] | None, + *, + expert_module_leaf_name: str, +) -> list[nn.Module]: + if ps.ep_size <= 1 or expert_classifier is None: + return [] + modules: list[nn.Module] = [] + seen: set[int] = set() + for chunk in chunks: + for module_name, module in chunk.named_modules(): + if module_name.rsplit(".", 1)[-1] != expert_module_leaf_name: + continue + prefix = f"{module_name}." + if not any( + expert_classifier(prefix + name) + for name, _param in module.named_parameters(recurse=True) + ): + continue + module_id = id(module) + if module_id not in seen: + modules.append(module) + seen.add(module_id) + return modules + + +def _collect_module_params(modules: Iterable[nn.Module]) -> set[nn.Parameter]: + return {param for module in modules for param in module.parameters()} + + +def _collect_model_param_dtypes(chunks: Iterable[nn.Module]) -> dict[tuple[int, str], torch.dtype]: + return { + (chunk_idx, name): param.dtype + for chunk_idx, chunk in enumerate(chunks) + for name, param in chunk.named_parameters() + if param.requires_grad and param.is_floating_point() and param.dtype != torch.float32 + } + + +def _restore_model_param_dtypes( + chunks: Iterable[nn.Module], model_param_dtypes: dict[tuple[int, str], torch.dtype] +) -> None: + for chunk_idx, chunk in enumerate(chunks): + for name, param in chunk.named_parameters(): + model_dtype = model_param_dtypes.get((chunk_idx, name)) + if model_dtype is not None: + param._fsdp2_model_param_dtype = model_dtype + + +def _collect_tp_replicated_grad_param_names(chunks: Iterable[nn.Module]) -> list[tuple[int, str]]: + names: list[tuple[int, str]] = [] + for chunk_idx, chunk in enumerate(chunks): + sp_param_ids = {id(param) for param in getattr(chunk, "sp_params", ())} + if not sp_param_ids: + continue + for name, param in chunk.named_parameters(): + if id(param) in sp_param_ids: + names.append((chunk_idx, name)) + return names + + +def _collect_tp_replicated_grad_params( + chunks: Iterable[nn.Module], *, param_names: Iterable[tuple[int, str]] | None = None +) -> list[nn.Parameter]: + chunk_list = list(chunks) + params: list[nn.Parameter] = [] + seen: set[int] = set() + if param_names is not None: + named_by_chunk = [dict(chunk.named_parameters()) for chunk in chunk_list] + for chunk_idx, name in param_names: + if chunk_idx >= len(named_by_chunk): + continue + param = named_by_chunk[chunk_idx].get(name) + if param is None or not param.requires_grad or id(param) in seen: + continue + params.append(param) + seen.add(id(param)) + return params + for chunk in chunk_list: + for param in getattr(chunk, "sp_params", ()): + if not param.requires_grad or id(param) in seen: + continue + params.append(param) + seen.add(id(param)) + return params + + +def _fsdp2_unit_reshard_after_forward( + ps: ParallelState, *, reshard_after_forward: bool | int | None +) -> bool | int | None: + if ps.pp_size > 1: + return False + return reshard_after_forward + + +def _fsdp2_prefetch_depth(ps: ParallelState, *, default_depth: int) -> int: + if ps.pp_size > 1: + return 0 + return default_depth + + +def _build_adamw_param_groups( + model_chunks: Iterable[nn.Module], + *, + weight_decay: float, + apply_wd_to_qk_layernorm: bool, + model_param_dtypes: dict[tuple[int, str], torch.dtype] | None = None, +) -> tuple[list[nn.Parameter], list[dict[str, Any]], dict[int, str], dict[int, torch.dtype]]: + params: list[nn.Parameter] = [] + decay_params: list[nn.Parameter] = [] + no_decay_params: list[nn.Parameter] = [] + param_names: dict[int, str] = {} + param_model_dtypes: dict[int, torch.dtype] = {} + seen_param_ids: set[int] = set() + + for chunk_idx, chunk in enumerate(model_chunks): + for name, param in chunk.named_parameters(): + if not param.requires_grad or id(param) in seen_param_ids: + continue + seen_param_ids.add(id(param)) + params.append(param) + param_names[id(param)] = f"chunk{chunk_idx}.{name}" + model_dtype = None + if model_param_dtypes is not None: + model_dtype = model_param_dtypes.get((chunk_idx, name)) + if model_dtype is None: + model_dtype = fsdp2_model_param_dtype(param) + if model_dtype is not None: + param_model_dtypes[id(param)] = model_dtype + if _matches_megatron_no_weight_decay(name, param, apply_wd_to_qk_layernorm): + no_decay_params.append(param) + else: + decay_params.append(param) + + # Match Megatron's default get_megatron_optimizer(config_overrides=None): + # 1D params and bias skip AdamW decay unless the Q/K layernorm override is set. + param_groups: list[dict[str, Any]] = [] + if decay_params: + param_groups.append({"params": decay_params, "weight_decay": weight_decay, "wd_mult": 1.0}) + if no_decay_params: + param_groups.append({"params": no_decay_params, "weight_decay": 0.0, "wd_mult": 0.0}) + return params, param_groups, param_names, param_model_dtypes + + +def _matches_megatron_no_weight_decay( + name: str, param: nn.Parameter, apply_wd_to_qk_layernorm: bool +) -> bool: + if len(param.shape) != 1 and not name.endswith(".bias"): + return False + if not apply_wd_to_qk_layernorm: + return True + return not ( + "q_layernorm." in name or "k_layernorm." in name or "q_norm." in name or "k_norm." in name + ) + + +@dataclass(frozen=True, slots=True) +class FSDP2OptimizerBackend: + name: str = "fsdp2" + runtime_backend: str = "fsdp2" + + def zero_grad(self, optimizer: FSDP2Optimizer) -> None: + optimizer.zero_grad() + + def finish_grad_sync(self, optimizer: FSDP2Optimizer) -> None: + return None + + def clip_grad_norm(self, optimizer: FSDP2Optimizer): + return optimizer.clip_grad_norm() + + def step(self, optimizer: FSDP2Optimizer): + return optimizer.step() + + def state_dict(self, optimizer: FSDP2Optimizer) -> dict: + return optimizer.state_dict() + + def load_state_dict(self, optimizer: FSDP2Optimizer, state_dict: dict) -> None: + optimizer.load_state_dict(state_dict) + + +BACKEND = FSDP2OptimizerBackend() + +__all__ = [ + "BACKEND", + "FSDP2OptimizerBackend", + "FSDP2Optimizer", + "build_fsdp2_adamw", + "build_fsdp2_training_optimizer", +] diff --git a/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/state.py b/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/state.py new file mode 100644 index 00000000000..3dd0f06df4c --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/state.py @@ -0,0 +1,96 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Optimizer state movement helpers for the FSDP2 primitive.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch + +from megatron.lite.primitive.optimizers.fsdp2.adamw import ( + dtensor_from_local, + is_dtensor_like, + iter_torch_optimizers, +) + + +@dataclass +class OffloadedStateEntry: + device: torch.device + is_dtensor: bool = False + device_mesh: Any | None = None + placements: Any | None = None + + +def move_optimizer_state_to_cpu( + optimizer: Any, + offloaded: dict[tuple[int, str], OffloadedStateEntry], + *, + include_dtensor_state: bool, +) -> None: + for child in iter_torch_optimizers(optimizer): + state = getattr(child, "state", None) + if not isinstance(state, dict): + continue + for param, param_state in state.items(): + if not isinstance(param_state, dict): + continue + for key, value in list(param_state.items()): + if is_dtensor_like(value): + if not include_dtensor_state: + continue + local_value = value.to_local() + if not isinstance(local_value, torch.Tensor) or not local_value.is_cuda: + continue + offloaded[(id(param), key)] = OffloadedStateEntry( + device=local_value.device, + is_dtensor=True, + device_mesh=value.device_mesh, + placements=value.placements, + ) + param_state[key] = local_value.detach().to("cpu") + continue + + if not isinstance(value, torch.Tensor) or not value.is_cuda: + continue + offloaded[(id(param), key)] = OffloadedStateEntry(device=value.device) + param_state[key] = value.detach().to("cpu") + + +def move_offloaded_optimizer_state_to_device( + optimizer: Any, offloaded: dict[tuple[int, str], OffloadedStateEntry] +) -> None: + if not offloaded: + return + remaining = dict(offloaded) + for child in iter_torch_optimizers(optimizer): + state = getattr(child, "state", None) + if not isinstance(state, dict): + continue + for param, param_state in state.items(): + if not isinstance(param_state, dict): + continue + for key, entry in list(remaining.items()): + param_id, state_key = key + if param_id != id(param) or state_key not in param_state: + continue + value = param_state[state_key] + if isinstance(value, torch.Tensor) and not is_dtensor_like(value): + device_value = value.to(entry.device, non_blocking=True) + if entry.is_dtensor: + param_state[state_key] = dtensor_from_local( + device_value, entry.device_mesh, entry.placements + ) + else: + param_state[state_key] = device_value + remaining.pop(key, None) + for key in set(offloaded) - set(remaining): + offloaded.pop(key, None) + + +__all__ = [ + "OffloadedStateEntry", + "move_offloaded_optimizer_state_to_device", + "move_optimizer_state_to_cpu", +] diff --git a/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/wrap.py b/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/wrap.py new file mode 100644 index 00000000000..149920d8a4a --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/optimizers/fsdp2/wrap.py @@ -0,0 +1,498 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""PyTorch FSDP2 wrapping primitive. + +This module is intentionally independent from Megatron Lite model and runtime +packages. Model protocols can call it after building modules and before +building the optimizer. +""" + +from __future__ import annotations + +import importlib +import warnings +from collections.abc import Callable, Iterable +from dataclasses import dataclass, field +from typing import Any + +import torch +import torch.distributed as dist +import torch.nn as nn + +from megatron.lite.primitive.parallel.state import ParallelState + +UnitModule = type[nn.Module] | str + +_WARNED_TP_NOT_RECOMMENDED = False + + +@dataclass(frozen=True) +class FSDP2Config: + """Configuration for ``wrap_fsdp2``.""" + + unit_modules: tuple[UnitModule, ...] = field(default_factory=tuple) + leaf_module_names: tuple[str, ...] = field(default_factory=tuple) + reshard_after_forward: bool | int | None = None + last_unit_reshard_after_forward: bool | int | None = False + root_reshard_after_forward: bool | int | None = False + wrap_root: bool = True + preserve_param_attrs: bool = True + forward_prefetch_depth: int = 1 + backward_prefetch_depth: int = 0 + mesh_dim_name: str = "dp_cp" + device_type: str = "cuda" + param_dtype: str | torch.dtype | None = None + reduce_dtype: str | torch.dtype | None = None + output_dtype: str | torch.dtype | None = None + cast_forward_inputs: bool | None = True + + def __post_init__(self) -> None: + object.__setattr__(self, "unit_modules", tuple(self.unit_modules)) + object.__setattr__(self, "leaf_module_names", tuple(self.leaf_module_names)) + if not self.wrap_root and not self.unit_modules and not self.leaf_module_names: + raise ValueError( + "FSDP2Config requires wrap_root=True, at least one unit module, " + "or at least one leaf module name." + ) + for name in self.leaf_module_names: + if not isinstance(name, str) or not name: + raise ValueError("leaf_module_names entries must be non-empty strings.") + if not self.mesh_dim_name: + raise ValueError("mesh_dim_name must be non-empty.") + if not self.device_type: + raise ValueError("device_type must be non-empty.") + if self.forward_prefetch_depth < 0: + raise ValueError("forward_prefetch_depth must be >= 0.") + if self.backward_prefetch_depth < 0: + raise ValueError("backward_prefetch_depth must be >= 0.") + + +def fsdp2_available() -> bool: + """Return whether the installed PyTorch exposes FSDP2 ``fully_shard``.""" + + try: + from torch.distributed import DeviceMesh # noqa: F401 + from torch.distributed.fsdp import fully_shard # noqa: F401 + except ImportError: + return False + return True + + +def build_fsdp2_device_mesh(ps: ParallelState, config: FSDP2Config | None = None) -> Any: + """Build the default one-dimensional FSDP2 DeviceMesh from ``ParallelState``.""" + + cfg = config or FSDP2Config() + if not dist.is_initialized(): + raise RuntimeError("FSDP2 requires torch.distributed to be initialized.") + + group = ps.dp_cp_group or ps.dp_group + if group is None: + raise RuntimeError("FSDP2 requires ParallelState.dp_cp_group or dp_group.") + + from torch.distributed import DeviceMesh + + return DeviceMesh.from_group( + group, device_type=cfg.device_type, mesh_dim_names=(cfg.mesh_dim_name,) + ) + + +def build_fsdp2_process_group_mesh( + group: dist.ProcessGroup, *, mesh_dim_name: str, device_type: str = "cuda" +) -> Any: + """Build a one-dimensional FSDP2 DeviceMesh from an explicit process group.""" + + if not dist.is_initialized(): + raise RuntimeError("FSDP2 requires torch.distributed to be initialized.") + if group is None: + raise RuntimeError("FSDP2 requires a non-null process group.") + if not mesh_dim_name: + raise ValueError("mesh_dim_name must be non-empty.") + if not device_type: + raise ValueError("device_type must be non-empty.") + + from torch.distributed import DeviceMesh + + return DeviceMesh.from_group(group, device_type=device_type, mesh_dim_names=(mesh_dim_name,)) + + +def build_fsdp2_shard_placement_fn(fsdp_size: int) -> Callable[[nn.Parameter], Any]: + """Build AutoModel-style FSDP2 shard placement. + + Choose the first tensor dimension divisible by the FSDP group size to avoid + padded DTensor shards when possible. + """ + + if fsdp_size <= 0: + raise ValueError(f"fsdp_size must be positive, got {fsdp_size}.") + + def shard_placement_fn(param: nn.Parameter) -> Any: + from torch.distributed.tensor import Shard + + for dim, size in enumerate(param.shape): + if int(size) % fsdp_size == 0: + return Shard(dim) + return Shard(0) + + return shard_placement_fn + + +def wrap_fsdp2( + model: nn.Module, + ps: ParallelState, + config: FSDP2Config | None = None, + *, + mesh: Any | None = None, + ignored_params: set[nn.Parameter] | None = None, + mp_policy: Any | None = None, + offload_policy: Any | None = None, + shard_placement_fn: Callable[[nn.Parameter], Any] | None = None, +) -> nn.Module: + """Apply PyTorch FSDP2 ``fully_shard`` to selected modules and the root. + + The model is mutated in place and returned for call-site convenience. + ``wrap_fsdp2`` must run before constructing a regular PyTorch optimizer. + """ + + cfg = config or FSDP2Config() + _warn_tp_not_recommended(ps) + fully_shard = _load_fully_shard() + fsdp_mesh = mesh if mesh is not None else build_fsdp2_device_mesh(ps, cfg) + unit_types = _resolve_unit_module_types(cfg.unit_modules) + + saved_attrs = _save_param_attrs(model) if cfg.preserve_param_attrs else {} + common_kwargs = _fully_shard_kwargs( + mesh=fsdp_mesh, + reshard_after_forward=None, + ignored_params=ignored_params, + mp_policy=mp_policy or _mixed_precision_policy_from_config(cfg), + offload_policy=offload_policy, + shard_placement_fn=shard_placement_fn, + ) + + wrapped_units: list[nn.Module] = [] + unit_modules = list(_iter_fsdp2_unit_modules(model, unit_types, cfg.leaf_module_names)) + for idx, sub_module in enumerate(unit_modules): + kwargs = dict(common_kwargs) + _set_optional_reshard_after_forward( + kwargs, _unit_reshard_after_forward(cfg, idx, len(unit_modules)) + ) + fully_shard(sub_module, **kwargs) + wrapped_units.append(sub_module) + + if cfg.wrap_root: + kwargs = dict(common_kwargs) + _set_optional_reshard_after_forward(kwargs, cfg.root_reshard_after_forward) + fully_shard(model, **kwargs) + + _apply_fsdp2_prefetch( + model, + wrapped_units, + forward_depth=cfg.forward_prefetch_depth, + backward_depth=cfg.backward_prefetch_depth, + ) + if saved_attrs: + _restore_param_attrs(model, saved_attrs) + return model + + +def wrap_fsdp2_module( + module: nn.Module, + ps: ParallelState, + config: FSDP2Config | None = None, + *, + mesh: Any | None = None, + ignored_params: set[nn.Parameter] | None = None, + mp_policy: Any | None = None, + offload_policy: Any | None = None, + shard_placement_fn: Callable[[nn.Parameter], Any] | None = None, + reshard_after_forward: bool | int | None = None, +) -> nn.Module: + """Apply FSDP2 ``fully_shard`` to exactly one module. + + This is used for nested modules that need a different sharding mesh from + their parent, e.g. EP-local MoE experts sharded over expert-DP while the + transformer block is sharded over dense DP/CP. + """ + + cfg = config or FSDP2Config() + _warn_tp_not_recommended(ps) + fully_shard = _load_fully_shard() + fsdp_mesh = mesh if mesh is not None else build_fsdp2_device_mesh(ps, cfg) + + saved_attrs = _save_param_attrs(module) if cfg.preserve_param_attrs else {} + kwargs = _fully_shard_kwargs( + mesh=fsdp_mesh, + reshard_after_forward=( + cfg.root_reshard_after_forward + if reshard_after_forward is None + else reshard_after_forward + ), + ignored_params=ignored_params, + mp_policy=mp_policy or _mixed_precision_policy_from_config(cfg), + offload_policy=offload_policy, + shard_placement_fn=shard_placement_fn, + ) + fully_shard(module, **kwargs) + + if saved_attrs: + _restore_param_attrs(module, saved_attrs) + return module + + +def _warn_tp_not_recommended(ps: ParallelState) -> None: + global _WARNED_TP_NOT_RECOMMENDED + if _WARNED_TP_NOT_RECOMMENDED or ps.tp_size <= 1: + return + if dist.is_available() and dist.is_initialized() and dist.get_rank() != 0: + return + _WARNED_TP_NOT_RECOMMENDED = True + warnings.warn( + f"FSDP2 with tp={ps.tp_size} is supported, but TP is not recommended " + "for FSDP2 V1; prefer tp=1 etp=1 for precision and speed signoff.", + RuntimeWarning, + stacklevel=3, + ) + + +def promote_fsdp2_trainable_params_to_fp32( + model: nn.Module, *, ignored_params: set[nn.Parameter] | None = None +) -> int: + """Promote FSDP2-owned trainable floating parameters to FP32 shards. + + Model protocols still use ``FSDP2Config.param_dtype`` to run compute in + BF16. Keeping the sharded parameters in FP32 makes the torch optimizer path + closer to MCore dist-opt's main-param semantics and avoids BF16 grad-norm + and update drift before FSDP2 wrapping. + """ + + ignored_param_ids = {id(param) for param in ignored_params or ()} + promoted = 0 + with torch.no_grad(): + for param in model.parameters(): + if id(param) in ignored_param_ids: + continue + if not param.requires_grad or not param.is_floating_point(): + continue + if param.dtype == torch.float32: + continue + param._fsdp2_model_param_dtype = param.dtype + param.data = param.data.to(torch.float32) + if param.grad is not None: + param.grad = param.grad.to(torch.float32) + promoted += 1 + return promoted + + +def set_fsdp2_requires_gradient_sync( + module: nn.Module, requires_gradient_sync: bool, *, recurse: bool = True +) -> int: + """Set FSDP2 gradient sync state and return the number of touched roots.""" + + setter = getattr(module, "set_requires_gradient_sync", None) + if callable(setter): + setter(requires_gradient_sync, recurse=recurse) + return 1 + if not recurse: + return 0 + + touched = 0 + for child in module.modules(): + if child is module: + continue + setter = getattr(child, "set_requires_gradient_sync", None) + if callable(setter): + setter(requires_gradient_sync, recurse=False) + touched += 1 + return touched + + +def _load_fully_shard(): + try: + from torch.distributed.fsdp import fully_shard + except ImportError as exc: + raise RuntimeError( + "PyTorch FSDP2 is unavailable; install a PyTorch build with " + "torch.distributed.fsdp.fully_shard." + ) from exc + return fully_shard + + +def _resolve_unit_module_types(unit_modules: Iterable[UnitModule]) -> tuple[type[nn.Module], ...]: + resolved: list[type[nn.Module]] = [] + for item in unit_modules: + if isinstance(item, str): + item = _import_module_type(item) + if not isinstance(item, type) or not issubclass(item, nn.Module): + raise TypeError( + "FSDP2 unit_modules entries must be nn.Module subclasses or import paths." + ) + resolved.append(item) + return tuple(resolved) + + +def _import_module_type(path: str) -> type[nn.Module]: + module_name, sep, attr_name = path.rpartition(".") + if not sep or not module_name or not attr_name: + raise ValueError(f"Invalid FSDP2 unit module path: {path!r}") + module = importlib.import_module(module_name) + obj = getattr(module, attr_name) + if not isinstance(obj, type) or not issubclass(obj, nn.Module): + raise TypeError(f"FSDP2 unit module path does not resolve to nn.Module: {path!r}") + return obj + + +def _iter_fsdp2_unit_modules( + root: nn.Module, unit_types: tuple[type[nn.Module], ...], leaf_module_names: tuple[str, ...] +) -> Iterable[nn.Module]: + leaf_names = set(leaf_module_names) + if not unit_types and not leaf_names: + return () + ordered_units: list[nn.Module] = [] + seen: set[int] = set() + + def visit(module: nn.Module, module_name: str) -> None: + leaf_name = module_name.rsplit(".", 1)[-1] if module_name else "" + selected = module is not root and ( + isinstance(module, unit_types) or leaf_name in leaf_names + ) + if selected and _module_has_trainable_param(module): + module_id = id(module) + if module_id not in seen: + ordered_units.append(module) + seen.add(module_id) + return + for child_name, child in _iter_ordered_named_children(module): + full_name = child_name if not module_name else f"{module_name}.{child_name}" + visit(child, full_name) + + visit(root, "") + return tuple(ordered_units) + + +def _iter_ordered_named_children(module: nn.Module) -> Iterable[tuple[str, nn.Module]]: + if isinstance(module, nn.ModuleDict): + return module.items() + if isinstance(module, nn.ModuleList): + return ((str(idx), module[idx]) for idx in range(len(module))) + return module.named_children() + + +def _module_has_trainable_param(module: nn.Module) -> bool: + return any(param.requires_grad for param in module.parameters()) + + +def _is_fsdp2_module(module: nn.Module) -> bool: + return hasattr(module, "set_modules_to_forward_prefetch") and hasattr( + module, "set_modules_to_backward_prefetch" + ) + + +def _apply_fsdp2_prefetch( + root: nn.Module, wrapped_units: list[nn.Module], *, forward_depth: int, backward_depth: int +) -> None: + fsdp_units = [module for module in wrapped_units if _is_fsdp2_module(module)] + fsdp_root = root if _is_fsdp2_module(root) else None + + if fsdp_units and forward_depth > 0: + if fsdp_root is not None: + fsdp_root.set_modules_to_forward_prefetch(fsdp_units[:forward_depth]) + for idx, current in enumerate(fsdp_units[:-1]): + targets = fsdp_units[idx + 1 : idx + 1 + forward_depth] + if targets: + current.set_modules_to_forward_prefetch(targets) + + if len(fsdp_units) > 1 and backward_depth > 0: + for idx in range(1, len(fsdp_units)): + start = max(0, idx - backward_depth) + targets = list(reversed(fsdp_units[start:idx])) + if targets: + fsdp_units[idx].set_modules_to_backward_prefetch(targets) + + +def _unit_reshard_after_forward(cfg: FSDP2Config, idx: int, total_units: int) -> bool | int | None: + if total_units > 0 and idx == total_units - 1: + return cfg.last_unit_reshard_after_forward + return cfg.reshard_after_forward + + +def _set_optional_reshard_after_forward( + kwargs: dict[str, Any], reshard_after_forward: bool | int | None +) -> None: + if reshard_after_forward is not None: + kwargs["reshard_after_forward"] = reshard_after_forward + + +def _fully_shard_kwargs( + *, + mesh: Any, + reshard_after_forward: bool | int | None, + ignored_params: set[nn.Parameter] | None, + mp_policy: Any | None, + offload_policy: Any | None, + shard_placement_fn: Callable[[nn.Parameter], Any] | None, +) -> dict[str, Any]: + kwargs: dict[str, Any] = {"mesh": mesh} + if reshard_after_forward is not None: + kwargs["reshard_after_forward"] = reshard_after_forward + if ignored_params is not None: + kwargs["ignored_params"] = ignored_params + if mp_policy is not None: + kwargs["mp_policy"] = mp_policy + if offload_policy is not None: + kwargs["offload_policy"] = offload_policy + if shard_placement_fn is not None: + kwargs["shard_placement_fn"] = shard_placement_fn + return kwargs + + +def _mixed_precision_policy_from_config(cfg: FSDP2Config) -> Any | None: + if cfg.param_dtype is None and cfg.reduce_dtype is None and cfg.output_dtype is None: + return None + try: + from torch.distributed.fsdp import MixedPrecisionPolicy + except ImportError as exc: + raise RuntimeError("FSDP2 mixed precision policy is unavailable.") from exc + kwargs = dict( + param_dtype=_resolve_torch_dtype(cfg.param_dtype), + reduce_dtype=_resolve_torch_dtype(cfg.reduce_dtype), + output_dtype=_resolve_torch_dtype(cfg.output_dtype), + ) + if cfg.cast_forward_inputs is not None: + kwargs["cast_forward_inputs"] = cfg.cast_forward_inputs + try: + return MixedPrecisionPolicy(**kwargs) + except TypeError: + kwargs.pop("cast_forward_inputs", None) + return MixedPrecisionPolicy(**kwargs) + + +def _resolve_torch_dtype(dtype: str | torch.dtype | None) -> torch.dtype | None: + if dtype is None or isinstance(dtype, torch.dtype): + return dtype + name = dtype.removeprefix("torch.") + resolved = getattr(torch, name, None) + if not isinstance(resolved, torch.dtype): + raise ValueError(f"Unsupported torch dtype for FSDP2 mixed precision: {dtype!r}") + return resolved + + +def _save_param_attrs(module: nn.Module) -> dict[str, dict[str, Any]]: + return {name: dict(vars(param)) for name, param in module.named_parameters()} + + +def _restore_param_attrs(module: nn.Module, saved_attrs: dict[str, dict[str, Any]]) -> None: + for name, param in module.named_parameters(): + for attr_name, attr_value in saved_attrs.get(name, {}).items(): + setattr(param, attr_name, attr_value) + + +__all__ = [ + "FSDP2Config", + "build_fsdp2_device_mesh", + "build_fsdp2_process_group_mesh", + "build_fsdp2_shard_placement_fn", + "fsdp2_available", + "promote_fsdp2_trainable_params_to_fp32", + "set_fsdp2_requires_gradient_sync", + "wrap_fsdp2", + "wrap_fsdp2_module", +] diff --git a/experimental/lite/megatron/lite/primitive/optimizers/megatron_wrap.py b/experimental/lite/megatron/lite/primitive/optimizers/megatron_wrap.py new file mode 100644 index 00000000000..ec182609fe1 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/optimizers/megatron_wrap.py @@ -0,0 +1,444 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Megatron-Core optimizer wrap backend for Megatron Lite.""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any + +import torch # pyright: ignore[reportMissingImports] +import torch.nn as nn # pyright: ignore[reportMissingImports] + +from megatron.lite.primitive.protocols import ExpertClassifierFn, default_expert_classifier + + +def validate_mc_config(engine_cfg) -> None: + """Validate dist_opt constraints owned by this optimizer primitive.""" + p = engine_cfg.parallel + if p.vpp > 1 and p.pp == 1: + raise ValueError("dist_opt requires pp>1 when vpp>1.") + + +# Legacy alias — kept for compat shim path +validate_mc_session = validate_mc_config + + +def _effective_etp(parallel) -> int: + return int(parallel.etp if parallel.etp is not None else 1) + + +def _ensure_mc_mpu_parallel_state(engine_cfg) -> None: + """Initialize Megatron-Core mpu globals when MC fallback groups are used.""" + + from megatron.core import parallel_state as mpu # pyright: ignore[reportMissingImports] + + p = engine_cfg.parallel + expected = (int(p.tp), int(p.ep), _effective_etp(p), int(p.pp), int(p.cp)) + if mpu.is_initialized(): + current = ( + int(mpu.get_tensor_model_parallel_world_size()), + int(mpu.get_expert_model_parallel_world_size()), + int(mpu.get_expert_tensor_parallel_world_size() or 1), + int(mpu.get_pipeline_model_parallel_world_size()), + int(mpu.get_context_parallel_world_size()), + ) + if current != expected: + raise RuntimeError( + "dist_opt found an incompatible existing Megatron-Core parallel state: " + f"current={current}, expected={expected}." + ) + return + + mpu.initialize_model_parallel( + tensor_model_parallel_size=p.tp, + pipeline_model_parallel_size=p.pp, + virtual_pipeline_model_parallel_size=None if int(p.vpp or 1) <= 1 else p.vpp, + context_parallel_size=p.cp, + expert_model_parallel_size=p.ep, + expert_tensor_parallel_size=_effective_etp(p), + create_gloo_process_groups=bool(getattr(engine_cfg, "deterministic", False)), + ) + + +def build_mc_optimizer_config(opt, *, override_optimizer_config: dict[str, Any] | None = None): + """Build MC OptimizerConfig from user's OptimizerConfig (duck-typed). + + Single source of truth for Megatron Lite's Megatron-Core optimizer stack. + + Works on either `runtime.contracts.config.OptimizerConfig` (real dataclass) + or a `SimpleNamespace` with the same field names (legacy lite path). + """ + from megatron.core.optimizer.optimizer_config import ( + OptimizerConfig as MCOptimizerConfig, # pyright: ignore[reportMissingImports] + ) + + offload = getattr(opt, "offload_fraction", None) or 0.0 + args: dict[str, Any] = { + "optimizer": opt.optimizer, + "lr": opt.lr, + "min_lr": getattr(opt, "min_lr", 0.0), + "weight_decay": opt.weight_decay, + "clip_grad": opt.clip_grad, + "use_distributed_optimizer": True, + "bf16": True, + "params_dtype": torch.bfloat16, + } + if offload > 0: + args["optimizer_offload_fraction"] = offload + args["overlap_cpu_optimizer_d2h_h2d"] = True + args["optimizer_cpu_offload"] = True + if getattr(opt, "adam_beta1", None) is not None: + args["adam_beta1"] = opt.adam_beta1 + if getattr(opt, "adam_beta2", None) is not None: + args["adam_beta2"] = opt.adam_beta2 + if getattr(opt, "adam_eps", None) is not None: + args["adam_eps"] = opt.adam_eps + if getattr(opt, "use_precision_aware_optimizer", None) is not None: + args["use_precision_aware_optimizer"] = opt.use_precision_aware_optimizer + if getattr(opt, "decoupled_weight_decay", None) is not None: + args["decoupled_weight_decay"] = opt.decoupled_weight_decay + if override_optimizer_config: + args.update(override_optimizer_config) + return MCOptimizerConfig(**args) + + +def build_mc_stack( + model_chunks: list[nn.Module], + *, + model_cfg, + engine_cfg, + ps, + is_expert: ExpertClassifierFn | None = None, + proto=None, + skip_ddp_wrap: bool = False, +): + """Wrap ML model chunks with MC DDP and build the matching MC optimizer. + + Args: + skip_ddp_wrap: when True, ``model_chunks`` are assumed to already be + MC ``DistributedDataParallel``-wrapped; we skip our own wrapping + and feed them directly to the optimizer. The bucket layout + influences optimizer master-grad sharding, so callers that prewrap + chunks own the DDP config compatibility. + """ + from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig + from megatron.core.distributed.finalize_model_grads import finalize_model_grads + from megatron.core.optimizer import get_megatron_optimizer + from megatron.core.transformer.enums import ModelType + + validate_mc_config(engine_cfg) + + p = engine_cfg.parallel + opt = engine_cfg.optimizer + + mc_transformer_cfg = _build_transformer_config(model_cfg, engine_cfg) + mc_transformer_cfg.finalize_model_grads_func = finalize_model_grads + if is_expert is not None: + is_expert_param = is_expert + elif proto is not None and hasattr(proto, "EXPERT_CLASSIFIER"): + is_expert_param = proto.EXPERT_CLASSIFIER + else: + is_expert_param = default_expert_classifier + use_mpu_groups = bool(getattr(engine_cfg, "deterministic", False)) + if use_mpu_groups: + _ensure_mc_mpu_parallel_state(engine_cfg) + pg_collection = None if use_mpu_groups else _build_pg_collection(ps, engine_cfg) + + if skip_ddp_wrap: + # Caller already wrapped and marked every param. Our helper setting + # `param.allreduce` on dense params could clash with MC code paths that + # distinguish `hasattr(param,'allreduce')` from `getattr(..., True)`. + wrapped_chunks = list(model_chunks) + else: + ddp_config = DistributedDataParallelConfig( + use_distributed_optimizer=True, overlap_grad_reduce=False, grad_reduce_in_fp32=True + ) + wrapped_chunks = [] + for chunk_idx, chunk in enumerate(model_chunks): + chunk.model_type = ModelType.encoder_or_decoder + _mark_mc_parallel_attrs(chunk, is_expert_param, tp_size=p.tp) + ddp_kwargs = {} + if pg_collection is not None: + ddp_kwargs["pg_collection"] = pg_collection + wrapped_chunks.append( + DistributedDataParallel( + mc_transformer_cfg, + ddp_config, + chunk, + disable_bucketing=(chunk_idx > 0), + **ddp_kwargs, + ) + ) + + # Single-source-of-truth OptimizerConfig construction for native lite + # model protocols. + opt_config = build_mc_optimizer_config(opt) + + # This branch falls back to MC mpu globals for the optimizer's process + # groups. Long term, this primitive should always pass its own + # `pg_collection`. + if skip_ddp_wrap or use_mpu_groups: + optimizer = get_megatron_optimizer(config=opt_config, model_chunks=wrapped_chunks) + optimizer._mc_pg_collection = None # pyright: ignore[reportAttributeAccessIssue] + else: + optimizer = get_megatron_optimizer( + config=opt_config, + model_chunks=wrapped_chunks, + use_gloo_process_groups=False, + pg_collection=pg_collection, + ) + optimizer._mc_pg_collection = pg_collection # pyright: ignore[reportAttributeAccessIssue] + return wrapped_chunks, optimizer + + +def build_mc_training_optimizer( + model_chunks: list[nn.Module], + *, + model_cfg, + impl_cfg, + ps, + model_name: str, + is_expert: ExpertClassifierFn | None = None, + skip_ddp_wrap: bool = False, + deterministic: bool | None = None, +): + """Build the MC DDP+optimizer stack from a Megatron Lite model ImplConfig.""" + + opt = impl_cfg.optimizer_config + if opt is None: + opt = SimpleNamespace( + optimizer="adam", + lr=1e-4, + weight_decay=0.01, + clip_grad=1.0, + offload_fraction=None, + adam_beta1=None, + adam_beta2=None, + adam_eps=None, + ) + if deterministic is None: + from megatron.lite.primitive.deterministic import deterministic_requested + + deterministic = deterministic_requested() + + engine_cfg = SimpleNamespace( + model_name=model_name, + parallel=impl_cfg.parallel, + optimizer=opt, + deterministic=bool(deterministic), + ) + model_chunks[:], optimizer = build_mc_stack( + model_chunks, + model_cfg=model_cfg, + engine_cfg=engine_cfg, + ps=ps, + is_expert=is_expert, + skip_ddp_wrap=skip_ddp_wrap, + ) + + def finalize_grads() -> None: + finalize_mc_grads(model_chunks, optimizer) + + return optimizer, finalize_grads + + +def finalize_mc_grads(model_chunks: list[nn.Module], optimizer) -> None: + """Run MC gradient finalization to match the optimizer's expected contract.""" + from megatron.core.distributed.finalize_model_grads import finalize_model_grads + + finalize_model_grads(model_chunks, pg_collection=optimizer._mc_pg_collection) + + +def _build_transformer_config(model_cfg, engine_cfg): + from megatron.core.transformer.transformer_config import TransformerConfig + + p = engine_cfg.parallel + kwargs = dict( + num_layers=max(getattr(model_cfg, "num_hidden_layers", 1), 1), + hidden_size=max(getattr(model_cfg, "hidden_size", 1), 1), + num_attention_heads=max(getattr(model_cfg, "num_attention_heads", 1), 1), + num_query_groups=getattr(model_cfg, "num_key_value_heads", None), + num_moe_experts=getattr(model_cfg, "num_experts", None), + moe_ffn_hidden_size=getattr(model_cfg, "moe_intermediate_size", None), + tensor_model_parallel_size=p.tp, + pipeline_model_parallel_size=p.pp, + context_parallel_size=p.cp, + expert_model_parallel_size=p.ep, + expert_tensor_parallel_size=p.etp if p.etp is not None else 1, + sequence_parallel=p.tp > 1, + bf16=True, + params_dtype=torch.bfloat16, + ) + if hasattr(model_cfg, "add_bias_linear"): + kwargs["add_bias_linear"] = bool(model_cfg.add_bias_linear) + elif kwargs["num_moe_experts"] is not None and kwargs["expert_tensor_parallel_size"] > 1: + kwargs["add_bias_linear"] = False + if p.pp > 1: + kwargs["pipeline_dtype"] = torch.bfloat16 + return TransformerConfig(**kwargs) + + +def _mark_mc_parallel_attrs( + model: nn.Module, is_expert_param: ExpertClassifierFn, *, tp_size: int +) -> None: + """Mark per-param MC metadata (allreduce / tensor_model_parallel / sequence_parallel). + + IMPORTANT: respect attrs that are already set. Prewrapped MC models may + mark these correctly per-param (e.g. `moe.router.weight` is 2D but + TP-replicated, and must NOT have `tensor_model_parallel=True`). Blind + override would cause MC grad-norm to over-count replicated params. + """ + sp_param_ids = {id(param) for param in getattr(model, "sp_params", [])} + for name, param in model.named_parameters(): + # MC uses `allreduce=False` to route expert params into expert-DP buffers. + if not hasattr(param, "allreduce"): + param.allreduce = not is_expert_param(name) + if tp_size > 1 and id(param) not in sp_param_ids and param.ndim > 1: + # vision params are replicated across TP (AVG all-reduce, not TP-split). + # tensor_model_parallel=True would cause MC to wrong-account their grad-norm. + if getattr(param, "average_gradients_across_tp_domain", False): + continue + # Skip params already marked sequence_parallel=True: they are TP-replicated + # with SP-sharded input (e.g. shared_experts.gate_weight, RMSNorm weights). + # Stacking tensor_model_parallel=True on top would cause double all-reduce. + if getattr(param, "sequence_parallel", False): + continue + # MC excludes TP replicas from grad-norm accounting via this metadata. + if not hasattr(param, "tensor_model_parallel"): + param.tensor_model_parallel = True + + for param in getattr(model, "sp_params", []): + if not hasattr(param, "sequence_parallel"): + param.sequence_parallel = True + param.allreduce = True + param.tensor_model_parallel = False + + +def _build_pg_collection(ps, engine_cfg): + import torch.distributed as dist # pyright: ignore[reportMissingImports] + + from megatron.core.process_groups_config import ProcessGroupCollection + + if ps.pp_group is None: + raise ValueError("dist_opt requires a local pp_group.") + + def _dense_rank(tp_i: int, cp_i: int, dp_i: int, pp_i: int) -> int: + return ((pp_i * ps.dp_size + dp_i) * ps.cp_size + cp_i) * ps.tp_size + tp_i + + def _expert_rank(etp_i: int, ep_i: int, edp_i: int, pp_i: int) -> int: + return ((pp_i * ps.expert_dp_size + edp_i) * ps.ep_size + ep_i) * ps.etp_size + etp_i + + rank = dist.get_rank() + world = dist.get_world_size() + + singleton_group = None + for singleton_rank in range(world): + group = dist.new_group([singleton_rank]) + if rank == singleton_rank: + singleton_group = group + if singleton_group is None: + raise RuntimeError( + "Failed to construct singleton process group for optional MC reductions." + ) + + if engine_cfg.parallel.pp == 1: + mp_group = ps.tp_group + tp_ep_pp_group = ps.tp_ep_group + else: + mp_group = None + for dp_idx in range(ps.dp_size): + for cp_idx in range(ps.cp_size): + ranks = [ + _dense_rank(tp_idx, cp_idx, dp_idx, pp_idx) + for pp_idx in range(ps.pp_size) + for tp_idx in range(ps.tp_size) + ] + group = dist.new_group(ranks) + if rank in ranks: + mp_group = group + + tp_ep_pp_group = None + for expert_dp_idx in range(ps.expert_dp_size): + ranks = [ + _expert_rank(etp_idx, ep_idx, expert_dp_idx, pp_idx) + for pp_idx in range(ps.pp_size) + for ep_idx in range(ps.ep_size) + for etp_idx in range(ps.etp_size) + ] + group = dist.new_group(ranks) + if rank in ranks: + tp_ep_pp_group = group + + if mp_group is None or tp_ep_pp_group is None: + raise RuntimeError("Failed to construct mc pipeline-aware process groups.") + + return ProcessGroupCollection( + tp=ps.tp_group, + cp=ps.cp_group, + pp=ps.pp_group, + ep=ps.ep_group, + mp=mp_group, + dp=ps.dp_group, + dp_cp=ps.dp_cp_group, + expt_dp=ps.ep_dp_group, + expt_tp=ps.etp_group, + tp_ep=ps.tp_ep_group, + tp_ep_pp=tp_ep_pp_group, + # For MC distributed optimizer, grad stats are reduced over the full optimizer instance. + # With a single dist-opt instance in this benchmark proof, that is the global world group. + intra_dist_opt=dist.group.WORLD, + # ML models do not expose MC's embedding/position-embedding sharing surface. + # Use singleton groups so MC's optional embedding reductions become no-ops + # without falling back to the global MCore embedding group. + embd=singleton_group, + pos_embd=singleton_group, + ) + + +# --------------------------------------------------------------------------- +# Backend adapter (consumed by runtime/session.py) +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True, slots=True) +class MCBackend: + name: str = "mc" + runtime_backend: str = "mc" + + def zero_grad(self, optimizer: Any) -> None: + optimizer.zero_grad() + + def finish_grad_sync(self, optimizer: Any) -> None: + if hasattr(optimizer, "finish_grad_sync"): + optimizer.finish_grad_sync() + + def clip_grad_norm(self, optimizer: Any): + if hasattr(optimizer, "clip_grad_norm"): + return optimizer.clip_grad_norm() + return None + + def step(self, optimizer: Any): + return optimizer.step() + + def state_dict(self, optimizer: Any) -> dict: + return optimizer.state_dict() + + def load_state_dict(self, optimizer: Any, state_dict: dict) -> None: + optimizer.load_state_dict(state_dict) + + def finalize_grads(self, finalize_fn, model_chunks: list[Any], optimizer: Any) -> None: + finalize_fn(model_chunks, optimizer) + + +BACKEND = MCBackend() + +__all__ = [ + "BACKEND", + "MCBackend", + "build_mc_stack", + "build_mc_training_optimizer", + "finalize_mc_grads", + "validate_mc_config", + "validate_mc_session", +] diff --git a/experimental/lite/megatron/lite/primitive/parallel/__init__.py b/experimental/lite/megatron/lite/primitive/parallel/__init__.py new file mode 100644 index 00000000000..c6fc5701e8a --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/parallel/__init__.py @@ -0,0 +1,76 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Megatron Lite-owned parallel runtime exports.""" + +from __future__ import annotations + +from megatron.lite.primitive.parallel.cp import ( + split_packed_for_cp, + zigzag_position_ids_for_cp, + zigzag_reconstruct_from_cp_parts, + zigzag_slice_for_cp, + zigzag_split_for_cp, +) +from megatron.lite.primitive.parallel.pipeline import forward_backward_pipelining +from megatron.lite.primitive.parallel.pp import PipelineChunkLayout, build_pipeline_chunk_layout +from megatron.lite.primitive.parallel.sp import ( + gather_for_non_sp_head, + gather_from_sequence_parallel, + scatter_to_sequence_parallel, +) +from megatron.lite.primitive.parallel.state import ParallelState, init_parallel +from megatron.lite.primitive.parallel.thd import ( + PackedSeqParams, + PackedTHDBatch, + pack_nested_thd, + reconstruct_packed_from_cp_parts, + roll_packed_thd_left, + split_packed_to_cp_local, + unpack_packed_thd_to_nested, +) + +_LAZY_LINEAR_EXPORTS = { + "ColumnParallelLinear", + "RowParallelLinear", + "VanillaColumnParallelLinear", + "VocabParallelEmbedding", + "VocabParallelOutput", + "pad_vocab_for_tp", +} + + +def __getattr__(name: str): + if name in _LAZY_LINEAR_EXPORTS: + from megatron.lite.primitive.parallel import linear as _linear + + return getattr(_linear, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +__all__ = [ + "ColumnParallelLinear", + "PackedSeqParams", + "PackedTHDBatch", + "PipelineChunkLayout", + "ParallelState", + "RowParallelLinear", + "VanillaColumnParallelLinear", + "VocabParallelEmbedding", + "VocabParallelOutput", + "build_pipeline_chunk_layout", + "forward_backward_pipelining", + "gather_for_non_sp_head", + "gather_from_sequence_parallel", + "init_parallel", + "pad_vocab_for_tp", + "pack_nested_thd", + "reconstruct_packed_from_cp_parts", + "roll_packed_thd_left", + "scatter_to_sequence_parallel", + "split_packed_to_cp_local", + "split_packed_for_cp", + "unpack_packed_thd_to_nested", + "zigzag_position_ids_for_cp", + "zigzag_reconstruct_from_cp_parts", + "zigzag_slice_for_cp", + "zigzag_split_for_cp", +] diff --git a/experimental/lite/megatron/lite/primitive/parallel/cp.py b/experimental/lite/megatron/lite/primitive/parallel/cp.py new file mode 100644 index 00000000000..2b1f4c21490 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/parallel/cp.py @@ -0,0 +1,154 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Context parallel zigzag sequence splitting helpers.""" + +from __future__ import annotations + +import torch + + +def zigzag_split_for_cp( + tensor: torch.Tensor, cp_rank: int, cp_size: int, seq_dim: int = 1 +) -> torch.Tensor: + """Split tensor along sequence dim using zigzag (striped) pattern for CP. + + Splits into 2*cp_size chunks; GPU i gets chunk[i] + chunk[2*cp_size-1-i]. + This balances causal-mask workload across CP ranks. + + Example (CP=2, seq=8): chunks [0,1,2,3] -> + GPU0: chunk[0]+chunk[3] = tokens [0,1,6,7] + GPU1: chunk[1]+chunk[2] = tokens [2,3,4,5] + """ + if cp_size <= 1: + return tensor + seq_len = tensor.shape[seq_dim] + assert ( + seq_len % (2 * cp_size) == 0 + ), f"seq_len={seq_len} must be divisible by 2*cp_size={2 * cp_size}" + shape = list(tensor.shape) + shape[seq_dim : seq_dim + 1] = [2 * cp_size, seq_len // (2 * cp_size)] + tensor = tensor.view(*shape) + idx = torch.tensor([cp_rank, 2 * cp_size - cp_rank - 1], dtype=torch.long, device=tensor.device) + tensor = tensor.index_select(seq_dim, idx) + shape[seq_dim : seq_dim + 2] = [seq_len // cp_size] + return tensor.reshape(*shape) + + +def zigzag_reconstruct_from_cp_parts( + parts: list[torch.Tensor] | tuple[torch.Tensor, ...], seq_dim: int = 1 +) -> torch.Tensor: + """Reconstruct a full sequence from per-rank zigzag CP shards.""" + cp_size = len(parts) + if cp_size <= 1: + return parts[0] + local_len = parts[0].shape[seq_dim] + assert ( + local_len % 2 == 0 + ), f"local seq_len={local_len} must be divisible by 2 for zigzag CP reconstruction" + for idx, part in enumerate(parts): + assert ( + part.shape == parts[0].shape + ), f"CP part {idx} shape {tuple(part.shape)} != {tuple(parts[0].shape)}" + + chunk = local_len // 2 + full_len = local_len * cp_size + out_shape = list(parts[0].shape) + out_shape[seq_dim] = full_len + full = torch.zeros(out_shape, dtype=parts[0].dtype, device=parts[0].device) + for rank, part in enumerate(parts): + first = part.narrow(seq_dim, 0, chunk) + second = part.narrow(seq_dim, chunk, chunk) + full.narrow(seq_dim, rank * chunk, chunk).copy_(first) + full.narrow(seq_dim, full_len - (rank + 1) * chunk, chunk).copy_(second) + return full + + +def zigzag_slice_for_cp( + tensor: torch.Tensor, cp_rank: int, cp_size: int, seq_dim: int = 1 +) -> torch.Tensor: + """Return one rank's zigzag CP shard from a full sequence tensor.""" + if cp_size <= 1: + return tensor + seq_len = tensor.shape[seq_dim] + assert ( + seq_len % (2 * cp_size) == 0 + ), f"seq_len={seq_len} must be divisible by 2*cp_size={2 * cp_size}" + chunk = seq_len // (2 * cp_size) + first = tensor.narrow(seq_dim, cp_rank * chunk, chunk) + second_start = seq_len - (cp_rank + 1) * chunk + second = tensor.narrow(seq_dim, second_start, chunk) + return torch.cat((first, second), dim=seq_dim).contiguous() + + +def zigzag_position_ids_for_cp( + seq_len: int, cp_rank: int, cp_size: int, device: torch.device +) -> torch.Tensor: + """Return global position IDs for this CP rank under zigzag splitting. + + Returns shape [1, seq_len // cp_size] matching batch dim convention. + """ + if cp_size <= 1: + return torch.arange(seq_len, device=device).unsqueeze(0) + chunk = seq_len // (2 * cp_size) + first = torch.arange(cp_rank * chunk, (cp_rank + 1) * chunk, device=device) + second_start = (2 * cp_size - cp_rank - 1) * chunk + second = torch.arange(second_start, second_start + chunk, device=device) + return torch.cat([first, second]).unsqueeze(0) + + +def split_packed_for_cp( + input_ids: torch.Tensor, + position_ids: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + cp_rank: int, + cp_size: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Zigzag-split a packed sequence batch for context parallelism. + + Each sample defined by *cu_seqlens* is split with the same zigzag + (striped) pattern as :func:`zigzag_split_for_cp`: tokens are divided + into ``2 * cp_size`` chunks, and this rank keeps + ``chunk[cp_rank] + chunk[2*cp_size - 1 - cp_rank]``. + + Returns: + ``(input_ids, position_ids, cu_seqlens, max_seqlen)`` for this CP rank. + """ + if cp_size <= 1: + return input_ids, position_ids, cu_seqlens, max_seqlen + + num_seqs = cu_seqlens.size(0) - 1 + ids_parts: list[torch.Tensor] = [] + pos_parts: list[torch.Tensor] = [] + new_lengths: list[int] = [] + + for i in range(num_seqs): + start = int(cu_seqlens[i].item()) + end = int(cu_seqlens[i + 1].item()) + seq_len = end - start + assert ( + seq_len % (2 * cp_size) == 0 + ), f"Sample {i} length {seq_len} not divisible by 2*cp_size={2 * cp_size}" + chunk = seq_len // (2 * cp_size) + c1 = start + cp_rank * chunk + c2 = start + (2 * cp_size - cp_rank - 1) * chunk + ids_parts.append(input_ids[c1 : c1 + chunk]) + ids_parts.append(input_ids[c2 : c2 + chunk]) + pos_parts.append(position_ids[c1 : c1 + chunk]) + pos_parts.append(position_ids[c2 : c2 + chunk]) + new_lengths.append(2 * chunk) + + new_ids = torch.cat(ids_parts) + new_pos = torch.cat(pos_parts) + lens = torch.tensor(new_lengths, dtype=torch.int32, device=cu_seqlens.device) + new_cu = torch.zeros(num_seqs + 1, dtype=torch.int32, device=cu_seqlens.device) + torch.cumsum(lens, dim=0, out=new_cu[1:]) + return new_ids, new_pos, new_cu, max(new_lengths) + + +__all__ = [ + "split_packed_for_cp", + "zigzag_reconstruct_from_cp_parts", + "zigzag_position_ids_for_cp", + "zigzag_slice_for_cp", + "zigzag_split_for_cp", +] diff --git a/experimental/lite/megatron/lite/primitive/parallel/linear.py b/experimental/lite/megatron/lite/primitive/parallel/linear.py new file mode 100644 index 00000000000..7ab907b8afd --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/parallel/linear.py @@ -0,0 +1,414 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""TP-parallel linear layers and vocab-parallel embedding/output.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch # pyright: ignore[reportMissingImports] +import torch.distributed as dist # pyright: ignore[reportMissingImports] +import torch.nn as nn # pyright: ignore[reportMissingImports] +import transformer_engine.pytorch as te # pyright: ignore[reportMissingImports] + +from megatron.lite.primitive.utils import ensure_divisible + +if TYPE_CHECKING: + from megatron.lite.primitive.parallel.state import ParallelState + + +# --------------------------------------------------------------------------- +# Vanilla column-parallel linear (torch.matmul kernel, NOT TE). +# +# Matches Megatron-Core `tensor_parallel.ColumnParallelLinear` bit-for-bit +# in bf16: same `torch.matmul(input, weight.t())` forward, same +# all-reduce-on-backward-grad_input pattern. Use this for heads like the +# vocab LM projection where MC's GPT model uses vanilla torch matmul +# (hardcoded in `LinearCrossEntropyModule(tensor_parallel.ColumnParallelLinear)`) +# — TE's `te.Linear` uses a different cuBLAS algo selection that introduces +# ~3e-4 loss-level drift under bf16 vs torch.matmul. +# +# For QKV / MoE experts we still prefer the TE path (fused LN+linear, FP8 +# readiness) — this vanilla path is a drop-in substitute only when kernel +# parity with the reference backend is required. +# --------------------------------------------------------------------------- + + +class _VanillaColParallelMatmul(torch.autograd.Function): + """forward: output = matmul(input, weight.t()) — replicated input, sharded output. + backward: grad_input = grad_output @ weight (+ all-reduce across TP); + grad_weight = grad_output.T @ input. + Mirrors MC `LinearWithGradAccumulationAndAsyncCommunication`. + """ + + @staticmethod + def forward(ctx, input_, weight, tp_group): + ctx.save_for_backward(input_, weight) + ctx.tp_group = tp_group + return torch.matmul(input_, weight.t()) + + @staticmethod + def backward(ctx, grad_output): + input_, weight = ctx.saved_tensors + grad_input = grad_output.matmul(weight) + if ctx.tp_group is not None and dist.get_world_size(ctx.tp_group) > 1: + dist.all_reduce(grad_input, group=ctx.tp_group) + # grad_weight = grad_output^T @ input (sum over leading dims). + gi = grad_output.reshape(-1, grad_output.shape[-1]) + xi = input_.reshape(-1, input_.shape[-1]) + grad_weight = gi.t().matmul(xi) + return grad_input, grad_weight, None + + +class _VanillaColParallelMatmulSP(torch.autograd.Function): + """SP-aware column-parallel matmul matching MC's + `ColumnParallelLinear(sequence_parallel=True)` kernel bit-for-bit. + + forward: + - all-gather input along dim-0 from [S/tp, B, H] → [S, B, H] + - matmul(gathered, weight.t()) → [S, B, V/tp] + backward: + - grad_input_full = grad_output @ weight → [S, B, H] + - reduce-scatter dim-0 → [S/tp, B, H] + - grad_weight = grad_output^T @ gathered_input + """ + + @staticmethod + def forward(ctx, input_, weight, tp_group): + ws = dist.get_world_size(tp_group) if tp_group is not None else 1 + if ws > 1: + s_local = input_.shape[0] + total_shape = (s_local * ws, *input_.shape[1:]) + total_input = torch.empty(total_shape, dtype=input_.dtype, device=input_.device) + dist.all_gather_into_tensor(total_input, input_.contiguous(), group=tp_group) + else: + total_input = input_ + ctx.save_for_backward(total_input, weight) + ctx.tp_group = tp_group + ctx.tp_size = ws + ctx.local_s = input_.shape[0] + return torch.matmul(total_input, weight.t()) + + @staticmethod + def backward(ctx, grad_output): + total_input, weight = ctx.saved_tensors + grad_input_full = grad_output.matmul(weight) + if ctx.tp_size > 1: + out_shape = (ctx.local_s, *grad_input_full.shape[1:]) + grad_input = torch.empty( + out_shape, dtype=grad_input_full.dtype, device=grad_input_full.device + ) + dist.reduce_scatter_tensor(grad_input, grad_input_full.contiguous(), group=ctx.tp_group) + else: + grad_input = grad_input_full + gi = grad_output.reshape(-1, grad_output.shape[-1]) + xi = total_input.reshape(-1, total_input.shape[-1]) + grad_weight = gi.t().matmul(xi) + return grad_input, grad_weight, None + + +class _VanillaColLinear(nn.Module): + """Drop-in for `te.Linear(parallel_mode='column')` using torch.matmul. + + Shape: `self.weight` is (out_features_per_tp, in_features). Exposes + `.weight` at the same attribute path as TE Linear for checkpoint-loader + compatibility. + + When `sp=True`, input is assumed SP-sharded on dim-0; forward gathers + before matmul and backward reduce-scatters grad_input — matching MC's + `ColumnParallelLinear(sequence_parallel=True)` bit-for-bit. When + `sp=False`, input is assumed replicated and grad_input is all-reduced. + """ + + def __init__(self, in_features: int, out_features: int, ps: ParallelState, *, sp: bool = False): + super().__init__() + self.tp_group = ps.tp_group + self.tp_size = ps.tp_size + self.sp = sp + local_out = ensure_divisible(out_features, ps.tp_size) + self.weight = nn.Parameter(torch.empty(local_out, in_features, dtype=torch.bfloat16)) + nn.init.xavier_uniform_(self.weight) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.sp: + return _VanillaColParallelMatmulSP.apply(x, self.weight, self.tp_group) + return _VanillaColParallelMatmul.apply(x, self.weight, self.tp_group) + + +class VanillaColumnParallelLinear(nn.Module): + """Public vanilla column-parallel linear matching MCore torch matmul.""" + + def __init__( + self, + in_features: int, + out_features: int, + ps: ParallelState, + *, + sp: bool = False, + gather_output: bool = False, + ): + super().__init__() + self.linear = _VanillaColLinear(in_features, out_features, ps, sp=sp) + self.tp_size = ps.tp_size + self.tp_group = ps.tp_group + self.gather_output = gather_output + + @property + def weight(self) -> torch.nn.Parameter: + return self.linear.weight + + def forward(self, x: torch.Tensor) -> torch.Tensor: + out = self.linear(x) + if self.gather_output and self.tp_size > 1: + out = _AllGatherLastDim.apply(out, self.tp_size, self.tp_group) + return out + + +class ColumnParallelLinear(nn.Module): + """TE-based column-parallel linear. Splits output dim across TP.""" + + def __init__( + self, + in_features: int, + out_features: int, + ps: ParallelState, + bias: bool = False, + gather_output: bool = False, + normalization: str | None = None, + eps: float = 1e-6, + zero_centered_gamma: bool = False, + sequence_parallel: bool | None = None, + ): + super().__init__() + self.tp_size = ps.tp_size + self.tp_rank = ps.tp_rank + self.tp_group = ps.tp_group + self.local_out = ensure_divisible(out_features, ps.tp_size) + self.use_sp = ( + ps.tp_size > 1 and not gather_output if sequence_parallel is None else sequence_parallel + ) + if normalization is not None: + self.linear = te.LayerNormLinear( + in_features, + out_features, + bias=bias, + normalization=normalization, + eps=eps, + zero_centered_gamma=zero_centered_gamma, + params_dtype=torch.bfloat16, + parallel_mode="column", + sequence_parallel=self.use_sp, + tp_group=ps.tp_group, + tp_size=ps.tp_size, + ) + else: + self.linear = te.Linear( + in_features, + out_features, + bias=bias, + params_dtype=torch.bfloat16, + parallel_mode="column", + sequence_parallel=self.use_sp, + tp_group=ps.tp_group, + tp_size=ps.tp_size, + ) + self.gather_output = gather_output + + def forward(self, x: torch.Tensor) -> torch.Tensor: + out = self.linear(x) + if self.gather_output and self.tp_size > 1: + out = _AllGatherLastDim.apply(out, self.tp_size, self.tp_group) + return out + + +class RowParallelLinear(nn.Module): + """TE-based row-parallel linear. Splits input dim across TP.""" + + def __init__( + self, + in_features: int, + out_features: int, + ps: ParallelState, + bias: bool = False, + input_is_parallel: bool = True, + ): + super().__init__() + self.tp_size = ps.tp_size + self.tp_rank = ps.tp_rank + self.tp_group = ps.tp_group + self.use_sp = ps.tp_size > 1 + self.local_in = ensure_divisible(in_features, ps.tp_size) + self.linear = te.Linear( + in_features, + out_features, + bias=bias, + params_dtype=torch.bfloat16, + parallel_mode="row", + sequence_parallel=self.use_sp, + tp_group=ps.tp_group, + tp_size=ps.tp_size, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(x) + + +def pad_vocab_for_tp(vocab_size: int, tp_size: int) -> int: + """Round vocab up to be divisible by `lcm(128, tp_size)`. + + Matches MC's `_vocab_size_with_padding(..., make_vocab_size_divisible_by=128)`: + pad to 128-multiple for GEMM alignment, and also require tp-divisibility. + For typical `tp_size ∈ {1,2,4,...,128}`, `lcm = 128` so a vocab already + divisible by 128 (e.g. Qwen3-MoE's 151936) stays unchanged — which is + what MC's `output_layer` sees. Using `128 * tp_size` instead would + over-pad (e.g. 151936 -> 152064 + at tp=2), introducing 128 extra logits into the vocab-parallel cross- + entropy log-sum-exp and driving a ~3e-4 loss drift. + """ + import math + + divisor = math.lcm(128, tp_size) + return ((vocab_size + divisor - 1) // divisor) * divisor + + +class VocabParallelEmbedding(nn.Module): + """Embedding table split across TP on the vocab dimension.""" + + def __init__( + self, vocab_size: int, hidden_size: int, ps: ParallelState, *, deterministic: bool = False + ): + super().__init__() + self.tp_size = ps.tp_size + self.tp_rank = ps.tp_rank + self.deterministic = bool(deterministic) + padded_vocab = pad_vocab_for_tp(vocab_size, ps.tp_size) + self.local_vocab = ensure_divisible(padded_vocab, ps.tp_size) + self.vocab_start = self.tp_rank * self.local_vocab + self.vocab_end = self.vocab_start + self.local_vocab + self.embedding = nn.Embedding(self.local_vocab, hidden_size) + self.tp_group = ps.tp_group + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + # input_ids: [B, S] → out: [S, B, H] + mask = (input_ids >= self.vocab_start) & (input_ids < self.vocab_end) + local_ids = (input_ids - self.vocab_start).clamp(min=0, max=self.local_vocab - 1) + if self.deterministic: + out = self.embedding.weight[local_ids] + else: + out = self.embedding(local_ids) + out = out * mask.unsqueeze(-1) + if self.tp_size > 1: + out = _ReduceFromTP.apply(out, self.tp_group) + return out.transpose(0, 1).contiguous() + + +class _ColForLMHead(nn.Module): + """Thin wrapper exposing `.linear` (with `.weight`) for checkpoint-loader + compat, switchable between TE and vanilla torch.matmul kernel. + + `sp=True` threads to the underlying linear: vanilla uses the SP-aware + matmul (gather-in / reduce-scatter-on-bwd); TE uses `sequence_parallel=True` + on `te.Linear`. Both match MC's `output_layer(sequence_parallel=True)` + semantics so the upstream final_layernorm can run on SP-sharded input. + """ + + def __init__( + self, + in_features: int, + out_features: int, + ps: ParallelState, + *, + backend: str = "vanilla", + sp: bool = False, + ): + super().__init__() + if backend == "vanilla": + # Matches MC's LinearCrossEntropyModule → tensor_parallel.ColumnParallelLinear + # kernel bit-for-bit in bf16. Preferred for LM head parity with the reference backend. + self.linear = _VanillaColLinear(in_features, out_features, ps, sp=sp) + elif backend == "te": + # TE-backed path — cuBLAS algo may differ from MC's torch.matmul + # under bf16, producing ~3e-4 loss-level drift. Keep for future + # FP8 / fused-norm paths. + _wrapper = ColumnParallelLinear( + in_features, out_features, ps, bias=False, gather_output=not sp + ) + self.linear = _wrapper.linear + else: + raise ValueError(f"Unknown LM-head backend: {backend!r}") + + +class VocabParallelOutput(nn.Module): + """Output projection split across TP on the vocab dimension (column parallel). + + Default backend is `"vanilla"` (torch.matmul) to match the reference backend's + `tensor_parallel.ColumnParallelLinear` bit-for-bit. Pass `backend="te"` + to use TE's `te.Linear` kernel (e.g. for FP8 inference paths). + """ + + def __init__( + self, vocab_size: int, hidden_size: int, ps: ParallelState, *, backend: str = "vanilla" + ): + super().__init__() + padded_vocab = pad_vocab_for_tp(vocab_size, ps.tp_size) + # SP-aware head: when tp>1 we run on SP-sharded input (matches MC + # GPTModel where final_layernorm runs on SP-sharded hiddens and + # output_layer gathers internally + reduce-scatters on backward). + self.col = _ColForLMHead(hidden_size, padded_vocab, ps, backend=backend, sp=ps.tp_size > 1) + self.padded_vocab = padded_vocab + self.local_vocab = padded_vocab // ps.tp_size + self.vocab_size = vocab_size + self.ps = ps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.col.linear(x) + + def gather(self, logits: torch.Tensor) -> torch.Tensor: + """All-gather TP-sharded logits and trim to actual vocab_size.""" + if self.ps.tp_size > 1: + chunks = [torch.empty_like(logits) for _ in range(self.ps.tp_size)] + dist.all_gather(chunks, logits, group=self.ps.tp_group) + logits = torch.cat(chunks, dim=-1) + return logits[..., : self.vocab_size] + + +class _AllGatherLastDim(torch.autograd.Function): + """all-gather along last dim; backward = take local shard.""" + + @staticmethod + def forward(ctx, x, tp_size, group): + ctx.tp_size = tp_size + ctx.group = group + ctx.local_dim = x.shape[-1] + ctx.rank = dist.get_rank(group) + chunks = [torch.empty_like(x) for _ in range(tp_size)] + dist.all_gather(chunks, x.contiguous(), group=group) + return torch.cat(chunks, dim=-1) + + @staticmethod + def backward(ctx, grad_output): + start = ctx.rank * ctx.local_dim + return grad_output[..., start : start + ctx.local_dim].contiguous(), None, None + + +class _ReduceFromTP(torch.autograd.Function): + """all-reduce forward; identity backward.""" + + @staticmethod + def forward(ctx, x, group): + out = x.clone() + dist.all_reduce(out, group=group) + return out + + @staticmethod + def backward(ctx, grad_output): + return grad_output, None + + +__all__ = [ + "ColumnParallelLinear", + "RowParallelLinear", + "VanillaColumnParallelLinear", + "VocabParallelEmbedding", + "VocabParallelOutput", + "pad_vocab_for_tp", +] diff --git a/experimental/lite/megatron/lite/primitive/parallel/pipeline.py b/experimental/lite/megatron/lite/primitive/parallel/pipeline.py new file mode 100644 index 00000000000..d52271862a8 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/parallel/pipeline.py @@ -0,0 +1,794 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Pipeline parallel: 1F1B schedule with correct backward handling.""" + +from __future__ import annotations + +import os +from collections.abc import Callable +from typing import TYPE_CHECKING + +import torch # pyright: ignore[reportMissingImports] +import torch.distributed as dist # pyright: ignore[reportMissingImports] + +from megatron.lite.primitive.utils import ensure_divisible + +if TYPE_CHECKING: + from megatron.lite.primitive.parallel.state import ParallelState + + +def forward_backward_pipelining( + forward_step_fn: Callable, + model_chunks: list, + data_iter, + config, + ps: ParallelState, + tensor_shape: tuple[int, ...] | None = None, + grad_sync_fn: Callable[[], None] | None = None, + pre_forward_hook: Callable[[torch.Tensor], None] | None = None, + loss_fn: Callable | None = None, + forward_only: bool = False, +) -> list[dict]: + """ + Run forward and backward passes with pipeline parallelism. + + Args: + forward_step_fn: Callable(model, batch) -> output_dict with "loss" and "hidden_states" + model_chunks: local model chunks. ``len>1`` enables interleaved VPP. + data_iter: iterator yielding micro-batches + config: training config + ps: parallel state + tensor_shape: shape of hidden states passed between stages [B, S, H] + grad_sync_fn: called before the last micro-batch's backward to enable + overlapped gradient ReduceScatter in DistributedOptimizer. + + Returns: + List of output dicts from forward passes. + """ + if ps.pp_size <= 1: + if forward_only: + return _forward_only_no_pipeline( + forward_step_fn, + model_chunks[0], + data_iter, + config, + ps, + pre_forward_hook=pre_forward_hook, + loss_fn=loss_fn, + ) + return _no_pipeline( + forward_step_fn, + model_chunks[0], + data_iter, + config, + ps, + grad_sync_fn=grad_sync_fn, + pre_forward_hook=pre_forward_hook, + loss_fn=loss_fn, + ) + + if tensor_shape is None: + raise ValueError("tensor_shape is required when PP > 1") + + num_microbatches = _num_microbatches_from_config(config, ps) + + if forward_only: + return _forward_only_pipeline_schedule( + forward_step_fn, + model_chunks, + data_iter, + num_microbatches, + ps, + tensor_shape, + pre_forward_hook=pre_forward_hook, + loss_fn=loss_fn, + ) + + if len(model_chunks) > 1: + return _interleaved_1f1b_schedule( + forward_step_fn, + model_chunks, + data_iter, + num_microbatches, + config, + ps, + tensor_shape, + grad_sync_fn=grad_sync_fn, + pre_forward_hook=pre_forward_hook, + loss_fn=loss_fn, + ) + + return _1f1b_schedule( + forward_step_fn, + model_chunks[0], + data_iter, + num_microbatches, + config, + ps, + tensor_shape, + grad_sync_fn=grad_sync_fn, + pre_forward_hook=pre_forward_hook, + loss_fn=loss_fn, + ) + + +# ══════════════════════════════════════════════════════════════════════ +# No pipeline (PP=1) +# ══════════════════════════════════════════════════════════════════════ +def _num_microbatches_from_config(config, ps: ParallelState) -> int: + explicit = getattr(config, "num_microbatches", None) + if explicit is not None: + return int(explicit) + return ensure_divisible(config.gbs, config.mbs * ps.dp_size) + + +def _set_aux_loss_scale(pre_forward_hook, num_microbatches: int) -> None: + if pre_forward_hook is not None: + scale = torch.tensor(1.0 / num_microbatches, device="cuda") + pre_forward_hook(scale) + + +def _batch_get(batch, key: str): + if isinstance(batch, dict): + return batch.get(key) + return getattr(batch, key, None) + + +def _batch_input_ids(batch): + input_ids = _batch_get(batch, "input_ids") + if input_ids is not None and input_ids.dim() == 1: + input_ids = input_ids.unsqueeze(0) + return input_ids + + +def _apply_external_loss( + out: dict, batch, loss_fn +) -> tuple[torch.Tensor, dict] | tuple[None, None]: + if loss_fn is None: + return None, None + loss, metrics = loss_fn(out, batch) + out["loss"] = loss + out["_loss_fn_metrics"] = metrics + return loss, metrics + + +def _compact_pipeline_output(out: dict | None) -> dict: + if not out: + return {} + compact: dict = {} + if "_verl_model_output" in out: + compact["model_output"] = out["_verl_model_output"] + if "loss" in out and out["loss"] is not None: + loss = out["loss"] + compact["loss"] = loss.detach().item() if isinstance(loss, torch.Tensor) else float(loss) + if "_loss_fn_metrics" in out: + compact["metrics"] = out["_loss_fn_metrics"] + elif "_verl_metrics" in out: + compact["metrics"] = out["_verl_metrics"] + return compact + + +def _no_pipeline( + forward_step_fn, + model, + data_iter, + config, + ps, + *, + grad_sync_fn=None, + pre_forward_hook=None, + loss_fn=None, +): + num_microbatches = _num_microbatches_from_config(config, ps) + outputs = [] + for i in range(num_microbatches): + if grad_sync_fn and i == num_microbatches - 1: + grad_sync_fn() + batch = next(data_iter) + _set_aux_loss_scale(pre_forward_hook, num_microbatches) + output = forward_step_fn(model, batch) + if loss_fn is not None: + loss, _metrics = _apply_external_loss(output, batch, loss_fn) + assert loss is not None + else: + loss = output["loss"] + loss = loss / num_microbatches + loss.backward() + outputs.append(_compact_pipeline_output(output)) + return outputs + + +def _forward_only_no_pipeline( + forward_step_fn, model, data_iter, config, ps, *, pre_forward_hook=None, loss_fn=None +): + num_microbatches = _num_microbatches_from_config(config, ps) + del ps + outputs = [] + for _ in range(num_microbatches): + batch = next(data_iter) + _set_aux_loss_scale(pre_forward_hook, num_microbatches) + output = forward_step_fn(model, batch) + _apply_external_loss(output, batch, loss_fn) + outputs.append(_compact_pipeline_output(output)) + return outputs + + +# ══════════════════════════════════════════════════════════════════════ +# 1F1B Schedule +# ══════════════════════════════════════════════════════════════════════ +def _1f1b_schedule( + forward_step_fn, + model, + data_iter, + num_microbatches: int, + config, + ps: ParallelState, + tensor_shape: tuple[int, ...], + *, + grad_sync_fn=None, + pre_forward_hook=None, + loss_fn=None, +): + """ + 1-Forward-1-Backward pipeline schedule using batch_isend_irecv. + + Communication is always done via combined send+recv to avoid deadlocks. + """ + num_warmup = min(ps.pp_size - ps.pp_rank - 1, num_microbatches) + num_steady = num_microbatches - num_warmup + + batches = [next(data_iter) for _ in range(num_microbatches)] + mb_idx = 0 + + input_tensors: list[torch.Tensor | None] = [] + output_hiddens: list[torch.Tensor | None] = [] + losses: list[torch.Tensor | None] = [] + outputs: list[dict] = [] + + # Fix 3: Pre-allocate recv buffers to avoid torch.empty per P2P call. + _fwd_recv_buf = ( + torch.empty(tensor_shape, dtype=_PIPELINE_TENSOR_DTYPE, device="cuda") + if not ps.pp_is_first + else None + ) + _bwd_recv_buf = ( + torch.empty(tensor_shape, dtype=_PIPELINE_TENSOR_DTYPE, device="cuda") + if not ps.pp_is_last + else None + ) + + def _run_forward(input_tensor, batch): + _set_aux_loss_scale(pre_forward_hook, num_microbatches) + position_ids = _batch_get(batch, "position_ids") + packed_seq_params = _batch_get(batch, "packed_seq_params") + if ps.pp_is_first: + return forward_step_fn(model, batch) + if ps.pp_is_last: + out = model( + input_ids=_batch_input_ids(batch), + hidden_states=input_tensor, + position_ids=position_ids, + packed_seq_params=packed_seq_params, + labels=_batch_get(batch, "labels"), + loss_mask=_batch_get(batch, "loss_mask"), + temperature=_batch_get(batch, "temperature") or 1.0, + use_fused_kernels=bool(_batch_get(batch, "use_fused_kernels") or False), + calculate_entropy=bool(_batch_get(batch, "calculate_entropy") or False), + ) + _apply_external_loss(out, batch, loss_fn) + return out + return model( + hidden_states=input_tensor, + position_ids=position_ids, + packed_seq_params=packed_seq_params, + ) + + def _run_backward(inp_t, hid_t, loss_t, grad_t): + if ps.pp_is_last: + if loss_t is not None: + loss_t.backward() + else: + if hid_t is not None and hid_t.requires_grad: + torch.autograd.backward(hid_t, grad_t) + return inp_t.grad if inp_t is not None else None + + def _p2p(send_fwd=None, send_bwd=None, recv_fwd=False, recv_bwd=False): + return _send_recv_pipeline( + send_fwd, + send_bwd, + recv_fwd, + recv_bwd, + ps, + tensor_shape, + fwd_recv_buf=_fwd_recv_buf, + bwd_recv_buf=_bwd_recv_buf, + ) + + # ── Warmup: pure forward passes ── + fwd_input: torch.Tensor | None = None + for k in range(num_warmup): + if not ps.pp_is_first and k == 0: + fwd_input, _ = _p2p(recv_fwd=True) + + batch = batches[mb_idx] + mb_idx += 1 + current_input = fwd_input + out = _run_forward(fwd_input, batch) + hidden = out.get("hidden_states") + loss_s = out["loss"] / num_microbatches if "loss" in out and ps.pp_is_last else None + + need_recv_next = not ps.pp_is_first and k < num_warmup - 1 + if not ps.pp_is_last: + fwd_input, _ = _p2p(send_fwd=hidden, recv_fwd=need_recv_next) + elif need_recv_next: + fwd_input, _ = _p2p(recv_fwd=True) + + input_tensors.append(current_input if not ps.pp_is_first else None) + output_hiddens.append(hidden) + losses.append(loss_s) + # Fix 4: only keep loss from output, drop logits/hidden references + outputs.append(_compact_pipeline_output(out)) + + # ── Steady: interleaved forward + backward ── + for k in range(num_steady): + if grad_sync_fn and num_warmup == 0 and k == num_steady - 1: + grad_sync_fn() + + if not ps.pp_is_first and k == 0 and num_warmup == 0: + fwd_input, _ = _p2p(recv_fwd=True) + + batch = batches[mb_idx] + mb_idx += 1 + out = _run_forward(fwd_input, batch) + hidden = out.get("hidden_states") + loss_s = out["loss"] / num_microbatches if "loss" in out and ps.pp_is_last else None + + input_tensors.append(fwd_input if not ps.pp_is_first else None) + output_hiddens.append(hidden) + losses.append(loss_s) + outputs.append(_compact_pipeline_output(out)) + + send_fwd = hidden if not ps.pp_is_last else None + need_bwd = not ps.pp_is_last + _, bwd_grad = _p2p(send_fwd=send_fwd, recv_bwd=need_bwd) + + old_inp = input_tensors.pop(0) + old_hid = output_hiddens.pop(0) + old_loss = losses.pop(0) + in_grad = _run_backward(old_inp, old_hid, old_loss, bwd_grad) + + send_bwd = in_grad if not ps.pp_is_first else None + need_fwd = not ps.pp_is_first and k < num_steady - 1 + fwd_input, _ = _p2p(send_bwd=send_bwd, recv_fwd=need_fwd) + + # ── Cooldown: drain remaining backwards ── + for k in range(num_warmup): + if grad_sync_fn and k == num_warmup - 1: + grad_sync_fn() + + need_bwd = not ps.pp_is_last + _, bwd_grad = _p2p(recv_bwd=need_bwd) + + old_inp = input_tensors.pop(0) + old_hid = output_hiddens.pop(0) + old_loss = losses.pop(0) + in_grad = _run_backward(old_inp, old_hid, old_loss, bwd_grad) + + send_bwd = in_grad if not ps.pp_is_first else None + if send_bwd is not None: + _p2p(send_bwd=send_bwd) + + return outputs + + +# ══════════════════════════════════════════════════════════════════════ +# Pipeline communication helpers +# ══════════════════════════════════════════════════════════════════════ +_PIPELINE_TENSOR_DTYPE = torch.bfloat16 + + +def _deallocate_output_tensor(tensor: torch.Tensor | None) -> None: + """Free a large output tensor after it has been sent to the next stage.""" + if tensor is not None: + tensor.data = torch.empty(1, device=tensor.device, dtype=tensor.dtype) + + +# ══════════════════════════════════════════════════════════════════════ +# Interleaved 1F1B (VPP) Schedule +# ══════════════════════════════════════════════════════════════════════ +def _build_schedule_table( + num_microbatches: int, num_chunks: int, group_size: int +) -> list[tuple[int, int]]: + """Build (microbatch_id, model_chunk_id) table for VPP scheduling.""" + table: list[tuple[int, int]] = [] + for start in range(0, num_microbatches, group_size): + end = min(start + group_size, num_microbatches) + for chunk in range(num_chunks): + for mb in range(start, end): + table.append((mb, chunk)) + return table + + +def _send_recv_pipeline( + send_fwd: torch.Tensor | None, + send_bwd: torch.Tensor | None, + recv_fwd: bool, + recv_bwd: bool, + ps: ParallelState, + tensor_shape: tuple[int, ...], + *, + fwd_recv_buf: torch.Tensor | None = None, + bwd_recv_buf: torch.Tensor | None = None, + batch_p2p: bool = True, + clone_recv: bool = False, +) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """P2P communication between pipeline stages.""" + _dbg = int(os.environ.get("MEGATRON_LITE_PP_DEBUG", "0")) + rank = dist.get_rank() + + ops: list[dist.P2POp] = [] + fwd_buf: torch.Tensor | None = None + bwd_buf: torch.Tensor | None = None + + p2p_group = ps.pp_group + + if send_fwd is not None: + t = send_fwd.to(_PIPELINE_TENSOR_DTYPE) + ops.append(dist.P2POp(dist.isend, t, ps.pp_next_rank, p2p_group)) + if recv_fwd: + fwd_buf = ( + fwd_recv_buf + if fwd_recv_buf is not None + else torch.empty(tensor_shape, dtype=_PIPELINE_TENSOR_DTYPE, device="cuda") + ) + ops.append(dist.P2POp(dist.irecv, fwd_buf, ps.pp_prev_rank, p2p_group)) + if send_bwd is not None: + t = send_bwd.to(_PIPELINE_TENSOR_DTYPE) + ops.append(dist.P2POp(dist.isend, t, ps.pp_prev_rank, p2p_group)) + if recv_bwd: + bwd_buf = ( + bwd_recv_buf + if bwd_recv_buf is not None + else torch.empty(tensor_shape, dtype=_PIPELINE_TENSOR_DTYPE, device="cuda") + ) + ops.append(dist.P2POp(dist.irecv, bwd_buf, ps.pp_next_rank, p2p_group)) + + if ops: + if _dbg: + desc = [] + if send_fwd is not None: + desc.append(f"send_fwd→{ps.pp_next_rank}({list(send_fwd.shape)})") + if recv_fwd: + desc.append(f"recv_fwd←{ps.pp_prev_rank}") + if send_bwd is not None: + desc.append(f"send_bwd→{ps.pp_prev_rank}") + if recv_bwd: + desc.append(f"recv_bwd←{ps.pp_next_rank}") + op_name = "batch_isend_irecv" if batch_p2p else "isend_irecv" + print(f"[P2P r{rank}] {op_name}: {' '.join(desc)}", flush=True) + if batch_p2p: + reqs = dist.batch_isend_irecv(ops) + else: + direct_tensors = [] + reqs = [] + if send_fwd is not None: + t = send_fwd.to(_PIPELINE_TENSOR_DTYPE) + direct_tensors.append(t) + reqs.append(dist.isend(t, ps.pp_next_rank, group=p2p_group)) + if recv_fwd: + reqs.append(dist.irecv(fwd_buf, ps.pp_prev_rank, group=p2p_group)) + if send_bwd is not None: + t = send_bwd.to(_PIPELINE_TENSOR_DTYPE) + direct_tensors.append(t) + reqs.append(dist.isend(t, ps.pp_prev_rank, group=p2p_group)) + if recv_bwd: + reqs.append(dist.irecv(bwd_buf, ps.pp_next_rank, group=p2p_group)) + for req in reqs: + req.wait() + if _dbg: + print(f"[P2P r{rank}] batch done", flush=True) + + if fwd_buf is not None: + if clone_recv: + fwd_buf = fwd_buf.clone() + fwd_buf.grad = None + fwd_buf.requires_grad_() + if bwd_buf is not None and clone_recv: + bwd_buf = bwd_buf.clone() + return fwd_buf, bwd_buf + + +def _pipeline_stage_barrier(ps: ParallelState) -> None: + if ps.pp_cpu_group is not None and ps.pp_size > 1: + dist.barrier(group=ps.pp_cpu_group) + + +def _set_virtual_pipeline_rank(chunk_id: int | None, num_chunks: int) -> None: + if chunk_id is None or num_chunks <= 1: + return + try: + from megatron.core import parallel_state as mpu # pyright: ignore[reportMissingImports] + except Exception: + return + if not mpu.is_initialized(): + return + vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() + if vpp_size is not None and vpp_size > 1: + mpu.set_virtual_pipeline_model_parallel_rank(chunk_id) + + +def _run_pipeline_chunk_forward( + forward_step_fn, + model, + batch, + input_tensor: torch.Tensor | None, + *, + is_first_stage: bool, + is_last_stage: bool, + num_microbatches: int, + pre_forward_hook=None, + loss_fn=None, +) -> dict: + _set_aux_loss_scale(pre_forward_hook, num_microbatches) + position_ids = _batch_get(batch, "position_ids") + packed_seq_params = _batch_get(batch, "packed_seq_params") + if is_first_stage: + return forward_step_fn(model, batch) + if is_last_stage: + out = model( + input_ids=_batch_input_ids(batch), + hidden_states=input_tensor, + position_ids=position_ids, + packed_seq_params=packed_seq_params, + labels=_batch_get(batch, "labels"), + loss_mask=_batch_get(batch, "loss_mask"), + temperature=_batch_get(batch, "temperature") or 1.0, + use_fused_kernels=bool(_batch_get(batch, "use_fused_kernels") or False), + calculate_entropy=bool(_batch_get(batch, "calculate_entropy") or False), + ) + _apply_external_loss(out, batch, loss_fn) + return out + return model( + hidden_states=input_tensor, position_ids=position_ids, packed_seq_params=packed_seq_params + ) + + +def _forward_only_pipeline_schedule( + forward_step_fn, + model_chunks: list, + data_iter, + num_microbatches: int, + ps: ParallelState, + tensor_shape: tuple[int, ...], + *, + pre_forward_hook=None, + loss_fn=None, +): + """Simple PP/VPP forward-only schedule used for log-prob inference.""" + num_chunks = len(model_chunks) + total_stages = ps.pp_size * num_chunks + outputs: list[dict] = [] + + for _mb in range(num_microbatches): + batch = next(data_iter) + _set_aux_loss_scale(pre_forward_hook, num_microbatches) + pending_activation: torch.Tensor | None = None + last_output: dict | None = None + for stage_id in range(total_stages): + stage_pp_rank = stage_id % ps.pp_size + is_local_stage = stage_pp_rank == ps.pp_rank + hidden: torch.Tensor | None = None + if is_local_stage: + chunk_id = stage_id // ps.pp_size + _set_virtual_pipeline_rank(chunk_id, num_chunks) + model = model_chunks[chunk_id] + is_first_stage = stage_id == 0 + is_last_stage = stage_id == total_stages - 1 + activation = None if is_first_stage else pending_activation + pending_activation = None + out = _run_pipeline_chunk_forward( + forward_step_fn, + model, + batch, + activation, + is_first_stage=is_first_stage, + is_last_stage=is_last_stage, + num_microbatches=num_microbatches, + pre_forward_hook=None, + loss_fn=loss_fn, + ) + if is_last_stage: + last_output = out + else: + hidden = out.get("hidden_states") + + if stage_id < total_stages - 1: + recv_next = (stage_id + 1) % ps.pp_size == ps.pp_rank + _pipeline_stage_barrier(ps) + fwd_buf, _ = _send_recv_pipeline( + hidden if is_local_stage else None, + None, + recv_next, + False, + ps, + tensor_shape, + batch_p2p=False, + clone_recv=True, + ) + if recv_next: + pending_activation = fwd_buf + + outputs.append(_compact_pipeline_output(last_output) if last_output is not None else {}) + + return outputs + + +def _interleaved_1f1b_schedule( + forward_step_fn, + model_chunks: list, + data_iter, + num_microbatches: int, + config, + ps: ParallelState, + tensor_shape: tuple[int, ...], + *, + grad_sync_fn=None, + pre_forward_hook=None, + loss_fn=None, +): + """ + Correct non-overlapped schedule for Virtual Pipeline Parallelism (VPP). + + Local chunks are laid out in global layer order as + ``chunk_id * pp_size + pp_rank``. The previous interleaved 1F1B bringup + schedule could ask the first physical PP stage to receive the next virtual + chunk before the last physical stage had produced it. This schedule keeps + the same VPP semantics but runs one micro-batch at a time in global stage + order, which is slower but deterministic and easy to validate against + Megatron. + """ + num_chunks = len(model_chunks) + total_stages = ps.pp_size * num_chunks + rank = dist.get_rank() + outputs: list[dict] = [] + + _dbg = int(os.environ.get("MEGATRON_LITE_PP_DEBUG", "0")) + if _dbg: + print( + f"[VPP r{rank}] entered simple schedule pp_rank={ps.pp_rank} " + f"microbatches={num_microbatches} chunks={num_chunks} total_stages={total_stages}", + flush=True, + ) + + for mb_id in range(num_microbatches): + batch = next(data_iter) + _set_aux_loss_scale(pre_forward_hook, num_microbatches) + saved: dict[ + int, tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, dict] + ] = {} + pending_activation: torch.Tensor | None = None + + # Forward in true virtual-stage order: + # chunk0/rank0 -> chunk0/rank1 -> ... -> chunkN/rankP. + for stage_id in range(total_stages): + stage_pp_rank = stage_id % ps.pp_size + is_local_stage = stage_pp_rank == ps.pp_rank + hidden: torch.Tensor | None = None + if is_local_stage: + chunk_id = stage_id // ps.pp_size + _set_virtual_pipeline_rank(chunk_id, num_chunks) + model = model_chunks[chunk_id] + is_first_stage = stage_id == 0 + is_last_stage = stage_id == total_stages - 1 + activation = None if is_first_stage else pending_activation + pending_activation = None + if _dbg: + activation_shape = None if activation is None else tuple(activation.shape) + print( + f"[VPP r{rank}] mb={mb_id} fwd stage={stage_id} " + f"chunk={chunk_id} activation_shape={activation_shape}", + flush=True, + ) + out = _run_pipeline_chunk_forward( + forward_step_fn, + model, + batch, + activation, + is_first_stage=is_first_stage, + is_last_stage=is_last_stage, + num_microbatches=num_microbatches, + pre_forward_hook=None, + loss_fn=loss_fn, + ) + + hidden = out.get("hidden_states") + if _dbg: + hidden_shape = None if hidden is None else tuple(hidden.shape) + print( + f"[VPP r{rank}] mb={mb_id} fwd stage={stage_id} " + f"chunk={chunk_id} hidden_shape={hidden_shape}", + flush=True, + ) + loss = out["loss"] / num_microbatches if is_last_stage and "loss" in out else None + saved[stage_id] = (activation, hidden, loss, out) + + if is_last_stage: + outputs.append(_compact_pipeline_output(out)) + + if stage_id < total_stages - 1: + recv_next = (stage_id + 1) % ps.pp_size == ps.pp_rank + if _dbg and (is_local_stage or recv_next): + print( + f"[VPP r{rank}] mb={mb_id} fwd boundary={stage_id} " + f"send={is_local_stage and hidden is not None} recv_next={recv_next}", + flush=True, + ) + _pipeline_stage_barrier(ps) + fwd_buf, _ = _send_recv_pipeline( + hidden if is_local_stage else None, + None, + recv_next, + False, + ps, + tensor_shape, + batch_p2p=False, + clone_recv=True, + ) + if recv_next: + pending_activation = fwd_buf + + if grad_sync_fn and mb_id == num_microbatches - 1: + grad_sync_fn() + + # Backward in the reverse virtual-stage order. + pending_grad: torch.Tensor | None = None + for stage_id in range(total_stages - 1, -1, -1): + stage_pp_rank = stage_id % ps.pp_size + is_local_stage = stage_pp_rank == ps.pp_rank + inp_grad: torch.Tensor | None = None + if is_local_stage: + chunk_id = stage_id // ps.pp_size + _set_virtual_pipeline_rank(chunk_id, num_chunks) + is_first_stage = stage_id == 0 + is_last_stage = stage_id == total_stages - 1 + inp, out_t, loss, _out = saved[stage_id] + grad = None if is_last_stage else pending_grad + pending_grad = None + if _dbg: + print(f"[VPP r{rank}] mb={mb_id} bwd stage={stage_id}", flush=True) + if is_last_stage: + if loss is not None: + loss.backward() + elif out_t is not None and out_t.requires_grad: + torch.autograd.backward(out_t, grad) + if not is_first_stage: + inp_grad = inp.grad if inp is not None else None + + if stage_id > 0: + recv_prev = (stage_id - 1) % ps.pp_size == ps.pp_rank + if _dbg and (is_local_stage or recv_prev): + print( + f"[VPP r{rank}] mb={mb_id} bwd boundary={stage_id} " + f"send={is_local_stage and inp_grad is not None} recv_prev={recv_prev}", + flush=True, + ) + _pipeline_stage_barrier(ps) + _, bwd_buf = _send_recv_pipeline( + None, + inp_grad if is_local_stage else None, + False, + recv_prev, + ps, + tensor_shape, + batch_p2p=False, + clone_recv=True, + ) + if recv_prev: + pending_grad = bwd_buf + + if _dbg: + print(f"[VPP r{rank}] mb={mb_id} complete", flush=True) + + return outputs + + +__all__ = ["forward_backward_pipelining"] diff --git a/experimental/lite/megatron/lite/primitive/parallel/pp.py b/experimental/lite/megatron/lite/primitive/parallel/pp.py new file mode 100644 index 00000000000..e5da34b8a37 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/parallel/pp.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Context/pipeline parallel sequence splitting utilities.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +from megatron.lite.primitive.utils import ensure_divisible + +if TYPE_CHECKING: + from megatron.lite.primitive.parallel.state import ParallelState + + +@dataclass +class PipelineChunkLayout: + layer_indices: list[int] = field(default_factory=list) + has_embed: bool = False + has_head: bool = False + + +def build_pipeline_chunk_layout( + num_hidden_layers: int, + ps: ParallelState, + vpp: int | None = None, + vpp_chunk_id: int | None = None, +) -> PipelineChunkLayout: + """Compute layer_indices, has_embed, has_head for this PP rank / VPP chunk.""" + if vpp_chunk_id is not None: + assert vpp is not None + layers_per_chunk = ensure_divisible(num_hidden_layers, ps.pp_size * vpp) + start = ps.pp_rank * layers_per_chunk + vpp_chunk_id * (ps.pp_size * layers_per_chunk) + layer_indices = list(range(start, start + layers_per_chunk)) + has_embed = ps.pp_is_first and vpp_chunk_id == 0 + has_head = ps.pp_is_last and vpp_chunk_id == vpp - 1 + elif vpp is not None: + layers_per_chunk = ensure_divisible(num_hidden_layers, ps.pp_size * vpp) + layer_indices = [] + for chunk in range(vpp): + start = ps.pp_rank * layers_per_chunk + chunk * (ps.pp_size * layers_per_chunk) + layer_indices.extend(range(start, start + layers_per_chunk)) + has_embed = ps.pp_is_first + has_head = ps.pp_is_last + else: + layers_per_stage = ensure_divisible(num_hidden_layers, ps.pp_size) + start = ps.pp_rank * layers_per_stage + layer_indices = list(range(start, start + layers_per_stage)) + has_embed = ps.pp_is_first + has_head = ps.pp_is_last + return PipelineChunkLayout(layer_indices=layer_indices, has_embed=has_embed, has_head=has_head) + + +__all__ = ["PipelineChunkLayout", "build_pipeline_chunk_layout"] diff --git a/experimental/lite/megatron/lite/primitive/parallel/sp.py b/experimental/lite/megatron/lite/primitive/parallel/sp.py new file mode 100644 index 00000000000..1c85c4b58cb --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/parallel/sp.py @@ -0,0 +1,45 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Sequence parallel scatter/gather helpers.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch # pyright: ignore[reportMissingImports] + +from megatron.lite.primitive.ops.sp_ops import ( + AllGatherDim0, + AllGatherDim0ForNonSPConsumer, + ScatterToSP, +) + +if TYPE_CHECKING: + from megatron.lite.primitive.parallel.state import ParallelState + + +def scatter_to_sequence_parallel(x: torch.Tensor, ps: ParallelState) -> torch.Tensor: + """Scatter [S, B, H] → [S/tp, B, H] for sequence parallel. No-op when tp=1.""" + if ps.tp_size == 1: + return x + return ScatterToSP.apply(x, ps.tp_size, ps.tp_rank, ps.tp_group) + + +def gather_from_sequence_parallel(x: torch.Tensor, ps: ParallelState) -> torch.Tensor: + """Gather [S/tp, B, H] → [S, B, H] from sequence parallel. No-op when tp=1.""" + if ps.tp_size == 1: + return x + return AllGatherDim0.apply(x, ps.tp_size, ps.tp_rank, ps.tp_group) + + +def gather_for_non_sp_head(x: torch.Tensor, ps: ParallelState) -> torch.Tensor: + """AllGather for non-SP consumer (e.g. vocab parallel head).""" + if ps.tp_size == 1: + return x + return AllGatherDim0ForNonSPConsumer.apply(x, ps.tp_size, ps.tp_rank, ps.tp_group) + + +__all__ = [ + "gather_for_non_sp_head", + "gather_from_sequence_parallel", + "scatter_to_sequence_parallel", +] diff --git a/experimental/lite/megatron/lite/primitive/parallel/state.py b/experimental/lite/megatron/lite/primitive/parallel/state.py new file mode 100644 index 00000000000..c282364a714 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/parallel/state.py @@ -0,0 +1,191 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""ParallelState and process group initialization.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch.distributed as dist # pyright: ignore[reportMissingImports] + +from megatron.lite.primitive.utils import ensure_divisible + + +@dataclass +class ParallelState: + tp_group: dist.ProcessGroup | None = None + ep_group: dist.ProcessGroup | None = None + etp_group: dist.ProcessGroup | None = None + cp_group: dist.ProcessGroup | None = None + pp_group: dist.ProcessGroup | None = None + pp_cpu_group: dist.ProcessGroup | None = None + dp_group: dist.ProcessGroup | None = None + dp_cp_group: dist.ProcessGroup | None = None + tp_ep_group: dist.ProcessGroup | None = None + ep_dp_group: dist.ProcessGroup | None = None + cp_global_ranks: list[int] | None = None + pp_global_ranks: list[int] | None = None + + tp_size: int = 1 + ep_size: int = 1 + etp_size: int = 1 + cp_size: int = 1 + pp_size: int = 1 + dp_size: int = 1 + dp_cp_size: int = 1 + expert_dp_size: int = 1 + + tp_rank: int = 0 + ep_rank: int = 0 + etp_rank: int = 0 + cp_rank: int = 0 + pp_rank: int = 0 + dp_rank: int = 0 + dp_cp_rank: int = 0 + expert_dp_rank: int = 0 + + pp_is_first: bool = True + pp_is_last: bool = True + pp_next_rank: int = -1 + pp_prev_rank: int = -1 + + +def init_parallel(config) -> ParallelState: + """ + Initialize all process groups using dual rank decomposition. + + Dense layers: world = TP × CP × PP × DP + Expert layers: world = ETP × EP × PP × expert_DP + """ + assert dist.is_initialized(), "Call torch.distributed.init_process_group first" + + world = dist.get_world_size() + rank = dist.get_rank() + tp, ep, etp, cp, pp = config.tp, config.ep, config.etp, config.cp, config.pp + + dense_dp = ensure_divisible(world, tp * cp * pp) + expert_dp = ensure_divisible(world, etp * ep * pp) + + ps = ParallelState() + ps.tp_size, ps.ep_size, ps.etp_size = tp, ep, etp + ps.cp_size, ps.pp_size, ps.dp_size = cp, pp, dense_dp + ps.expert_dp_size = expert_dp + ps.dp_cp_size = dense_dp * cp + + def _d(tp_i, cp_i, dp_i, pp_i): + return ((pp_i * dense_dp + dp_i) * cp + cp_i) * tp + tp_i + + def _e(etp_i, ep_i, edp_i, pp_i): + return ((pp_i * expert_dp + edp_i) * ep + ep_i) * etp + etp_i + + t = rank + my_tp = t % tp + t //= tp + my_cp = t % cp + t //= cp + my_ddp = t % dense_dp + t //= dense_dp + my_pp = t + + t = rank + my_etp = t % etp + t //= etp + my_ep = t % ep + t //= ep + my_edp = t % expert_dp + t //= expert_dp + assert t == my_pp, "PP rank must agree between dense and expert decompositions" + + ps.tp_rank, ps.cp_rank, ps.dp_rank, ps.pp_rank = my_tp, my_cp, my_ddp, my_pp + ps.dp_cp_rank = my_ddp * cp + my_cp + ps.ep_rank, ps.etp_rank, ps.expert_dp_rank = my_ep, my_etp, my_edp + ps.pp_is_first = my_pp == 0 + ps.pp_is_last = my_pp == pp - 1 + + for d in range(dense_dp): + for p in range(pp): + for c in range(cp): + ranks = [_d(t, c, d, p) for t in range(tp)] + g = dist.new_group(ranks) + if rank in ranks: + ps.tp_group = g + + for d in range(dense_dp): + for p in range(pp): + for t in range(tp): + ranks = [_d(t, c, d, p) for c in range(cp)] + g = dist.new_group(ranks) + if rank in ranks: + ps.cp_group = g + ps.cp_global_ranks = ranks + + for d in range(dense_dp): + for c in range(cp): + for t in range(tp): + ranks = [_d(t, c, d, p) for p in range(pp)] + g = dist.new_group(ranks) + try: + cpu_g = dist.new_group(ranks, backend="gloo") + except (RuntimeError, ValueError): + cpu_g = None + if rank in ranks: + ps.pp_group = g + ps.pp_cpu_group = cpu_g + ps.pp_global_ranks = ranks + + for p in range(pp): + for c in range(cp): + for t in range(tp): + ranks = [_d(t, c, d, p) for d in range(dense_dp)] + g = dist.new_group(ranks) + if rank in ranks: + ps.dp_group = g + + for p in range(pp): + for t in range(tp): + ranks = [_d(t, c, d, p) for d in range(dense_dp) for c in range(cp)] + g = dist.new_group(ranks) + if rank in ranks: + ps.dp_cp_group = g + + for d in range(expert_dp): + for p in range(pp): + for t in range(etp): + ranks = [_e(t, e, d, p) for e in range(ep)] + g = dist.new_group(ranks) + if rank in ranks: + ps.ep_group = g + + if etp > 1: + for d in range(expert_dp): + for p in range(pp): + for e in range(ep): + ranks = [_e(t, e, d, p) for t in range(etp)] + g = dist.new_group(ranks) + if rank in ranks: + ps.etp_group = g + + for d in range(expert_dp): + for p in range(pp): + ranks = [_e(t, e, d, p) for e in range(ep) for t in range(etp)] + g = dist.new_group(ranks) + if rank in ranks: + ps.tp_ep_group = g + + for p in range(pp): + for e in range(ep): + for t in range(etp): + ranks = [_e(t, e, d, p) for d in range(expert_dp)] + g = dist.new_group(ranks) + if rank in ranks: + ps.ep_dp_group = g + + pp_ranks = ps.pp_global_ranks + if pp_ranks is None: + raise RuntimeError("Pipeline ranks were not initialized.") + ps.pp_next_rank = pp_ranks[(my_pp + 1) % pp] if pp > 1 else rank + ps.pp_prev_rank = pp_ranks[(my_pp - 1) % pp] if pp > 1 else rank + + return ps + + +__all__ = ["ParallelState", "init_parallel"] diff --git a/experimental/lite/megatron/lite/primitive/parallel/thd.py b/experimental/lite/megatron/lite/primitive/parallel/thd.py new file mode 100644 index 00000000000..ac928bd8b5b --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/parallel/thd.py @@ -0,0 +1,457 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Packed sequence helpers for variable-length (THD) attention.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed as dist + + +@dataclass +class PackedSeqParams: + """Parameters for THD-format packed-sequence attention. + + Mirrors the fields consumed by TE's ``DotProductAttention.forward()`` + when ``qkv_format="thd"`` is used (total-tokens x heads x dim). + + Typical construction:: + + params = PackedSeqParams.from_cu_seqlens(batch.cu_seqlens, batch.max_seqlen) + """ + + qkv_format: str = "thd" + cu_seqlens_q: torch.Tensor | None = None + cu_seqlens_kv: torch.Tensor | None = None + max_seqlen_q: int | None = None + max_seqlen_kv: int | None = None + cu_seqlens_q_padded: torch.Tensor | None = None + cu_seqlens_kv_padded: torch.Tensor | None = None + local_cp_size: int | None = None + cp_group: Any | None = None + cp_rank: int | None = None + + @staticmethod + def from_cu_seqlens(cu_seqlens: torch.Tensor, max_seqlen: int) -> PackedSeqParams: + """Build from shared Q/KV cu_seqlens (self-attention).""" + return PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + cu_seqlens_q_padded=cu_seqlens, + cu_seqlens_kv_padded=cu_seqlens, + ) + + +@dataclass(frozen=True) +class PackedTHDBatch: + """Dense packed representation of a jagged no-padding batch.""" + + input_ids: torch.Tensor + labels: torch.Tensor | None + loss_mask: torch.Tensor | None + position_ids: torch.Tensor + packed_seq_params: Any + cu_seqlens_padded: torch.Tensor + lengths: torch.Tensor + padded_lengths: torch.Tensor + cp_size: int = 1 + cp_rank: int = 0 + cp_group: Any | None = None + + +def _make_packed_seq_params( + *, + cu_seqlens_padded: torch.Tensor, + max_seqlen: int, + cp_size: int = 1, + cp_rank: int = 0, + cp_group: Any | None = None, +): + extra_args = {} + if cp_size > 1: + extra_args["local_cp_size"] = cp_size + if cp_group is not None: + extra_args["cp_group"] = cp_group + try: + from megatron.core.packed_seq_params import PackedSeqParams as MCorePackedSeqParams + + params = MCorePackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens_padded, + cu_seqlens_kv=cu_seqlens_padded, + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + **extra_args, + ) + # MCore's PackedSeqParams does not carry rank, but local rolling needs it. + params.cp_rank = cp_rank + return params + except Exception: + return PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens_padded, + cu_seqlens_kv=cu_seqlens_padded, + max_seqlen_q=max_seqlen, + max_seqlen_kv=max_seqlen, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + cp_rank=cp_rank, + **extra_args, + ) + + +def _slice_along_dim(tensor: torch.Tensor, dim: int, start: int, end: int) -> torch.Tensor: + index = [slice(None)] * tensor.dim() + index[dim] = slice(start, end) + return tensor[tuple(index)] + + +def _assign_along_dim(dst: torch.Tensor, dim: int, start: int, src: torch.Tensor) -> None: + index = [slice(None)] * dst.dim() + index[dim] = slice(start, start + src.size(dim)) + dst[tuple(index)] = src + + +def _split_full_to_cp_local( + tensor: torch.Tensor, *, cu_seqlens_padded: torch.Tensor, cp_size: int, cp_rank: int, dim: int +) -> torch.Tensor: + if cp_size <= 1: + return tensor + + total_local = int(cu_seqlens_padded[-1].item()) // cp_size + out_shape = list(tensor.shape) + out_shape[dim] = total_local + local = torch.zeros(out_shape, dtype=tensor.dtype, device=tensor.device) + + for idx in range(int(cu_seqlens_padded.numel()) - 1): + full_start = int(cu_seqlens_padded[idx].item()) + full_end = int(cu_seqlens_padded[idx + 1].item()) + padded_len = full_end - full_start + if padded_len <= 0: + continue + chunk = padded_len // (2 * cp_size) + local_start = full_start // cp_size + + first = _slice_along_dim( + tensor, dim, full_start + cp_rank * chunk, full_start + (cp_rank + 1) * chunk + ) + second = _slice_along_dim( + tensor, dim, full_end - (cp_rank + 1) * chunk, full_end - cp_rank * chunk + ) + _assign_along_dim(local, dim, local_start, first) + _assign_along_dim(local, dim, local_start + chunk, second) + + return local + + +def _reconstruct_full_from_cp_parts( + parts: list[torch.Tensor], *, cu_seqlens_padded: torch.Tensor, cp_size: int, dim: int +) -> torch.Tensor: + if cp_size <= 1: + return parts[0] + + total_full = int(cu_seqlens_padded[-1].item()) + out_shape = list(parts[0].shape) + out_shape[dim] = total_full + full = torch.zeros(out_shape, dtype=parts[0].dtype, device=parts[0].device) + + for idx in range(int(cu_seqlens_padded.numel()) - 1): + full_start = int(cu_seqlens_padded[idx].item()) + full_end = int(cu_seqlens_padded[idx + 1].item()) + padded_len = full_end - full_start + if padded_len <= 0: + continue + chunk = padded_len // (2 * cp_size) + local_start = full_start // cp_size + + for rank, part in enumerate(parts): + first = _slice_along_dim(part, dim, local_start, local_start + chunk) + second = _slice_along_dim(part, dim, local_start + chunk, local_start + 2 * chunk) + _assign_along_dim(full, dim, full_start + rank * chunk, first) + _assign_along_dim(full, dim, full_end - (rank + 1) * chunk, second) + + return full + + +def reconstruct_packed_from_cp_parts( + parts: list[torch.Tensor], *, cu_seqlens_padded: torch.Tensor, cp_size: int, dim: int +) -> torch.Tensor: + """Reconstruct a full packed THD tensor from CP-local zigzag parts.""" + return _reconstruct_full_from_cp_parts( + parts, cu_seqlens_padded=cu_seqlens_padded, cp_size=cp_size, dim=dim + ) + + +def split_packed_to_cp_local( + tensor: torch.Tensor, *, cu_seqlens_padded: torch.Tensor, cp_size: int, cp_rank: int, dim: int +) -> torch.Tensor: + """Slice a full packed THD tensor back to one CP rank's zigzag shard.""" + return _split_full_to_cp_local( + tensor, cu_seqlens_padded=cu_seqlens_padded, cp_size=cp_size, cp_rank=cp_rank, dim=dim + ) + + +def _all_gather_cp_tensor( + tensor: torch.Tensor, *, cp_size: int, cp_group: Any +) -> list[torch.Tensor]: + if cp_size <= 1: + return [tensor] + if cp_group is None: + raise ValueError("CP THD gather requires cp_group.") + try: + from torch.distributed.nn.functional import all_gather + + return list(all_gather(tensor, group=cp_group)) + except Exception: + parts = [torch.empty_like(tensor) for _ in range(cp_size)] + dist.all_gather(parts, tensor, group=cp_group) + return parts + + +def _roll_packed_thd_left_local( + tensor: torch.Tensor, *, cu_seqlens_padded: torch.Tensor, dims: int = -1 +) -> tuple[torch.Tensor, torch.Tensor]: + dim = dims if dims >= 0 else tensor.dim() + dims + if dim < 0 or dim >= tensor.dim(): + raise ValueError(f"Invalid roll dim {dims} for tensor with {tensor.dim()} dims.") + + rolled = tensor.clone() + for idx in range(int(cu_seqlens_padded.numel()) - 1): + start = int(cu_seqlens_padded[idx].item()) + end = int(cu_seqlens_padded[idx + 1].item()) + if end <= start: + continue + + index = [slice(None)] * tensor.dim() + index[dim] = slice(start, end) + seq = torch.roll(tensor[tuple(index)], shifts=-1, dims=dim) + + zero_index = [slice(None)] * seq.dim() + zero_index[dim] = slice(-1, None) + seq[tuple(zero_index)] = 0 + rolled[tuple(index)] = seq + + return rolled, rolled.sum() + + +def roll_packed_thd_left( + tensor: torch.Tensor, + *, + cu_seqlens_padded: torch.Tensor | None = None, + packed_seq_params: Any | None = None, + dims: int = -1, +) -> tuple[torch.Tensor, torch.Tensor]: + """Roll a THD packed tensor left without crossing sequence boundaries.""" + + cp_size = 1 + cp_rank = 0 + cp_group = None + if packed_seq_params is not None: + cu_seqlens_padded = getattr(packed_seq_params, "cu_seqlens_q", None) + cp_size = int(getattr(packed_seq_params, "local_cp_size", None) or 1) + cp_rank = int(getattr(packed_seq_params, "cp_rank", 0) or 0) + cp_group = getattr(packed_seq_params, "cp_group", None) + if cu_seqlens_padded is None: + raise ValueError("THD packed roll requires cu_seqlens.") + + dim = dims if dims >= 0 else tensor.dim() + dims + if cp_size <= 1: + return _roll_packed_thd_left_local(tensor, cu_seqlens_padded=cu_seqlens_padded, dims=dim) + + parts = _all_gather_cp_tensor(tensor, cp_size=cp_size, cp_group=cp_group) + full = _reconstruct_full_from_cp_parts( + parts, cu_seqlens_padded=cu_seqlens_padded, cp_size=cp_size, dim=dim + ) + rolled_full, token_sum = _roll_packed_thd_left_local( + full, cu_seqlens_padded=cu_seqlens_padded, dims=dim + ) + local = _split_full_to_cp_local( + rolled_full, cu_seqlens_padded=cu_seqlens_padded, cp_size=cp_size, cp_rank=cp_rank, dim=dim + ) + return local, token_sum + + +def pack_nested_thd( + input_ids: torch.Tensor, + *, + tp_size: int = 1, + cp_size: int = 1, + cp_rank: int = 0, + cp_group: Any | None = None, + labels: torch.Tensor | None = None, + roll_labels: bool = False, + loss_mask: torch.Tensor | None = None, + roll_loss_mask: bool = False, +) -> PackedTHDBatch: + """Pack a jagged no-padding batch into Megatron Lite's THD model input. + + Mirrors VERL/Megatron's THD engine convention: each sequence is + padded to the tensor-parallel alignment, concatenated, then represented as + a single ``[1, local_padded_tokens]`` token row plus ``PackedSeqParams``. + For CP>1 the local row uses Megatron/TE zigzag chunking. + """ + + if cp_size < 1: + raise ValueError(f"cp_size must be >= 1, got {cp_size}") + if cp_rank < 0 or cp_rank >= cp_size: + raise ValueError(f"cp_rank must be in [0, {cp_size}), got {cp_rank}") + if not getattr(input_ids, "is_nested", False): + raise TypeError("pack_nested_thd expects a jagged NestedTensor input_ids.") + if labels is not None and not getattr(labels, "is_nested", False): + raise TypeError( + "pack_nested_thd expects jagged NestedTensor labels when labels are provided." + ) + if loss_mask is not None and not getattr(loss_mask, "is_nested", False): + raise TypeError( + "pack_nested_thd expects jagged NestedTensor loss_mask when loss_mask is provided." + ) + + align_size = max(int(tp_size), 1) * (2 * cp_size if cp_size > 1 else 1) + device = input_ids.device + offsets = input_ids.offsets().to(device=device) + lengths = offsets.diff().to(dtype=torch.int32) + pad_size = (align_size - lengths % align_size) % align_size + padded_lengths = lengths + pad_size + + cu_seqlens_padded = torch.zeros(lengths.numel() + 1, dtype=torch.int32, device=device) + cu_seqlens_padded[1:] = torch.cumsum(padded_lengths, dim=0) + total_padded = int(cu_seqlens_padded[-1].item()) + total_local = total_padded // cp_size + max_seqlen = int(padded_lengths.max().item()) if padded_lengths.numel() else 0 + + packed_input = torch.zeros(total_local, dtype=input_ids.dtype, device=device) + packed_labels = ( + torch.zeros(total_local, dtype=labels.dtype, device=device) if labels is not None else None + ) + packed_loss_mask = ( + torch.zeros(total_local, dtype=loss_mask.dtype, device=device) + if loss_mask is not None + else None + ) + # Megatron's rotary embedding slices position embeddings on the CP rank. + # Keep packed position ids full-length while CP-slicing tokens/labels/masks. + position_ids = torch.zeros(total_padded, dtype=torch.long, device=device) + + for idx, length_t in enumerate(lengths): + length = int(length_t.item()) + padded_length = int(padded_lengths[idx].item()) + full_start = int(cu_seqlens_padded[idx].item()) + local_start = full_start // cp_size + + seq_input = torch.zeros(padded_length, dtype=input_ids.dtype, device=device) + seq_input[:length] = input_ids[idx] + seq_labels = None + if labels is not None: + assert packed_labels is not None + seq_labels = torch.zeros(padded_length, dtype=labels.dtype, device=device) + seq_labels[:length] = labels[idx] + if roll_labels and length > 0: + seq_labels[:length] = torch.roll(seq_labels[:length], shifts=-1, dims=0) + seq_labels[length - 1] = 0 + seq_loss_mask = None + if loss_mask is not None: + assert packed_loss_mask is not None + seq_loss_mask = torch.zeros(padded_length, dtype=loss_mask.dtype, device=device) + seq_loss_mask[:length] = loss_mask[idx] + if roll_loss_mask and length > 0: + seq_loss_mask[:length] = torch.roll(seq_loss_mask[:length], shifts=-1, dims=0) + seq_loss_mask[length - 1] = 0 + seq_positions = torch.zeros(padded_length, dtype=torch.long, device=device) + seq_positions[:length] = torch.arange(length, dtype=torch.long, device=device) + + local_input = _split_full_to_cp_local( + seq_input, + cu_seqlens_padded=torch.tensor([0, padded_length], dtype=torch.int32, device=device), + cp_size=cp_size, + cp_rank=cp_rank, + dim=0, + ) + packed_input[local_start : local_start + local_input.numel()] = local_input + if seq_labels is not None: + local_labels = _split_full_to_cp_local( + seq_labels, + cu_seqlens_padded=torch.tensor( + [0, padded_length], dtype=torch.int32, device=device + ), + cp_size=cp_size, + cp_rank=cp_rank, + dim=0, + ) + assert packed_labels is not None + packed_labels[local_start : local_start + local_labels.numel()] = local_labels + if seq_loss_mask is not None: + local_loss_mask = _split_full_to_cp_local( + seq_loss_mask, + cu_seqlens_padded=torch.tensor( + [0, padded_length], dtype=torch.int32, device=device + ), + cp_size=cp_size, + cp_rank=cp_rank, + dim=0, + ) + assert packed_loss_mask is not None + packed_loss_mask[local_start : local_start + local_loss_mask.numel()] = local_loss_mask + position_ids[full_start : full_start + padded_length] = seq_positions + + return PackedTHDBatch( + input_ids=packed_input.unsqueeze(0), + labels=packed_labels.unsqueeze(0) if packed_labels is not None else None, + loss_mask=packed_loss_mask.unsqueeze(0) if packed_loss_mask is not None else None, + position_ids=position_ids.unsqueeze(0), + packed_seq_params=_make_packed_seq_params( + cu_seqlens_padded=cu_seqlens_padded, + max_seqlen=max_seqlen, + cp_size=cp_size, + cp_rank=cp_rank, + cp_group=cp_group, + ), + cu_seqlens_padded=cu_seqlens_padded, + lengths=lengths, + padded_lengths=padded_lengths, + cp_size=cp_size, + cp_rank=cp_rank, + cp_group=cp_group, + ) + + +def unpack_packed_thd_to_nested(output: torch.Tensor, batch: PackedTHDBatch) -> torch.Tensor: + """Unpack ``[1, total_padded, ...]`` THD model output back to jagged form.""" + + if output.dim() >= 2 and output.shape[0] == 1: + flat = output[0] + elif output.dim() >= 2 and output.shape[1] == 1: + flat = output[:, 0] + else: + flat = output + + if batch.cp_size > 1: + dim = 0 + parts = _all_gather_cp_tensor(flat, cp_size=batch.cp_size, cp_group=batch.cp_group) + flat = _reconstruct_full_from_cp_parts( + parts, cu_seqlens_padded=batch.cu_seqlens_padded, cp_size=batch.cp_size, dim=dim + ) + + pieces = [] + for idx, length_t in enumerate(batch.lengths): + length = int(length_t.item()) + start = int(batch.cu_seqlens_padded[idx].item()) + pieces.append(flat[start : start + length]) + return torch.nested.as_nested_tensor(pieces, layout=torch.jagged) + + +__all__ = [ + "PackedSeqParams", + "PackedTHDBatch", + "pack_nested_thd", + "reconstruct_packed_from_cp_parts", + "roll_packed_thd_left", + "split_packed_to_cp_local", + "unpack_packed_thd_to_nested", +] diff --git a/experimental/lite/megatron/lite/primitive/protocols.py b/experimental/lite/megatron/lite/primitive/protocols.py new file mode 100644 index 00000000000..d000864bc4a --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/protocols.py @@ -0,0 +1,24 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Train-side lightweight protocols and defaults.""" + +from __future__ import annotations + +from collections.abc import Callable + +from torch.distributed.tensor import Replicate # pyright: ignore[reportMissingImports] + +ExpertClassifierFn = Callable[[str], bool] +PlacementFn = Callable[[str], list] + + +def default_expert_classifier(name: str) -> bool: + """Default: params with 'experts' (but not 'router' or 'shared') are expert params.""" + return "experts" in name and "router" not in name and "shared" not in name + + +def default_placement_fn(name: str) -> list: + """Default: all Replicate (safe but no resharding benefit).""" + return [Replicate(), Replicate(), Replicate(), Replicate()] + + +__all__ = ["ExpertClassifierFn", "PlacementFn", "default_expert_classifier", "default_placement_fn"] diff --git a/experimental/lite/megatron/lite/primitive/recompute.py b/experimental/lite/megatron/lite/primitive/recompute.py new file mode 100644 index 00000000000..a83a209b0f7 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/recompute.py @@ -0,0 +1,381 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Model-agnostic activation recompute and offload wrappers.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +import torch # pyright: ignore[reportMissingImports] +import torch.nn as nn # pyright: ignore[reportMissingImports] + +# ── CheckpointWithoutOutput ─────────────────────────────────────────────────── +# Zero-copy C++ extension: makes dst's UntypedStorage point to src's data. +# Operates below TensorImpl level → ALL views/reshapes that share dst's StorageImpl +# (e.g. TE GroupedLinear's inp.reshape() saved for backward) see the restored data. +# Equivalent to MC's share_storage in megatron/core/tensor_parallel/random.py. +_SHARE_STORAGE_SRC = r""" +#include + +void share_storage(at::Tensor dst, at::Tensor src) { + auto* dst_impl = dst.storage().unsafeGetStorageImpl(); + auto* src_ref = new c10::Storage(src.storage()); + void* data = src_ref->data_ptr().get(); + size_t nbytes = src_ref->nbytes(); + c10::Device device = src_ref->device(); + c10::DataPtr shared( + data, + static_cast(src_ref), + [](void* ctx) { delete static_cast(ctx); }, + device); + dst_impl->set_data_ptr(std::move(shared)); + dst_impl->set_nbytes(nbytes); +} +""" + +_share_storage_ext = None + + +def _get_share_storage() -> Callable: + global _share_storage_ext + if _share_storage_ext is None: + from torch.utils.cpp_extension import load_inline + + _share_storage_ext = load_inline( + name="share_storage_ext", + cpp_sources=_SHARE_STORAGE_SRC, + functions=["share_storage"], + verbose=False, + ) + return _share_storage_ext.share_storage + + +class _CheckpointWithoutOutputFn(torch.autograd.Function): + """Autograd Function for CheckpointWithoutOutput. + + Forward: runs function with no_grad, saves inputs. + Backward: uses outputs/inputs set by CheckpointWithoutOutput._recompute(). + """ + + @staticmethod + def forward(ctx, run_function, ckpt_obj, *args): + ctx.run_function = run_function + ctx.preserve_rng_state = ckpt_obj.preserve_rng_state + if ckpt_obj.preserve_rng_state: + ctx.cpu_rng_state = torch.get_rng_state() + ctx.cuda_rng_state = torch.cuda.get_rng_state() + + ctx.tensor_indices = [i for i, a in enumerate(args) if isinstance(a, torch.Tensor)] + ctx.non_tensor_args = [ + (i, a) for i, a in enumerate(args) if not isinstance(a, torch.Tensor) + ] + ctx.num_args = len(args) + ctx.save_for_backward(*[a for a in args if isinstance(a, torch.Tensor)]) + + with torch.no_grad(): + outputs = run_function(*args) + + ckpt_obj.ctx = ctx + return outputs + + @staticmethod + def backward(ctx, *grad_outputs): + # inputs and outputs are set by CheckpointWithoutOutput._recompute() + # before this backward runs (via the hook registered on the downstream tensor). + inputs = ctx.inputs + outputs = ctx.outputs + torch.autograd.backward(outputs, grad_outputs) + ctx.outputs = None + ctx.inputs = None + grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else None for inp in inputs) + return (None, None) + grads + + +class CheckpointWithoutOutput: + """Checkpoint a function and discard its output to save memory. + + Equivalent to MC's CheckpointWithoutOutput from megatron/core/tensor_parallel/random.py. + The output tensor's storage is freed immediately after downstream computation; + it is recomputed just-in-time during backward via a hook on the downstream output. + + The C++ share_storage extension restores the output at the StorageImpl level so + that ALL aliases (including views saved by TE GroupedLinear's backward) see the data. + + Usage (mirrors MC's moe_act pattern):: + + ckpt = CheckpointWithoutOutput() + h = ckpt.checkpoint(activation_func, fc1_out, probs) + fc2_out = fc2(h, m_splits) + ckpt.discard_output_and_register_recompute(fc2_out) + """ + + def __init__(self, preserve_rng_state: bool = True): + self.preserve_rng_state = preserve_rng_state + self.run_function: Callable | None = None + self._cpu_rng: torch.Tensor | None = None + self._cuda_rng: torch.Tensor | None = None + self.ctx: Any | None = None + self.outputs: tuple[torch.Tensor, ...] | None = None + + def checkpoint(self, run_function: Callable, *args) -> Any: + self.run_function = run_function + if self.preserve_rng_state: + self._cpu_rng = torch.get_rng_state() + self._cuda_rng = torch.cuda.get_rng_state() + + outputs = _CheckpointWithoutOutputFn.apply(run_function, self, *args) + self.outputs = (outputs,) if isinstance(outputs, torch.Tensor) else tuple(outputs) + return outputs + + def _recompute(self, _) -> None: + if self.ctx is None: + return + + # Reconstruct args from saved context. + tensors = list(self.ctx.saved_tensors) + args: list = [None] * self.ctx.num_args + t_it = iter(tensors) + for i in self.ctx.tensor_indices: + t = next(t_it) + args[i] = t.detach().requires_grad_(t.requires_grad) + for i, val in self.ctx.non_tensor_args: + args[i] = val + + # Recompute with forward-time RNG states. + if self.preserve_rng_state: + saved_cpu = torch.get_rng_state() + saved_cuda = torch.cuda.get_rng_state() + torch.set_rng_state(self._cpu_rng) + torch.cuda.set_rng_state(self._cuda_rng) + with torch.enable_grad(): + outputs = self.run_function(*args) + if self.preserve_rng_state: + torch.set_rng_state(saved_cpu) + torch.cuda.set_rng_state(saved_cuda) + + if isinstance(outputs, torch.Tensor): + outputs = (outputs,) + + # Zero-copy: make original output's StorageImpl point to recomputed data. + share_storage = _get_share_storage() + for orig, new in zip(self.outputs, outputs, strict=False): + share_storage(orig, new) + + self.ctx.outputs = list(outputs) + self.ctx.inputs = args + self.run_function = None + self._cpu_rng = None + self._cuda_rng = None + self.outputs = None + self.ctx = None + + def reset(self) -> None: + """Reset state so the instance can be reused across forward passes.""" + self.run_function = None + self._cpu_rng = None + self._cuda_rng = None + self.ctx = None + self.outputs = None + + def discard_output_and_register_recompute(self, hook_tensor: torch.Tensor) -> None: + """Free output tensor storage; recompute when hook_tensor's grad is computed.""" + for out in self.outputs: + out.untyped_storage().resize_(0) + if hook_tensor.requires_grad: + hook_tensor.register_hook(self._recompute) + + +ModuleMap = dict[str, Callable[[nn.Module], nn.Module | None]] +"""Maps module name → lambda that extracts a sub-module from a layer.""" + + +def apply_recompute( + layers: nn.ModuleList, + module_names: list[str], + module_map: ModuleMap, + no_rng_modules: set[str] | None = None, +) -> None: + """Wrap specified sub-modules with activation checkpointing for recomputation.""" + if not module_names: + return + no_rng = no_rng_modules or set() + for layer in layers: + if "full" in module_names: + wrap_checkpoint(layer) + else: + for mod_name in module_names: + if mod_name in module_map: + submod = module_map[mod_name](layer) + if submod is not None: + wrap_checkpoint(submod, preserve_rng_state=mod_name not in no_rng) + + +def apply_offload(layers: nn.ModuleList, module_names: list[str], module_map: ModuleMap) -> None: + """Wrap specified sub-modules with activation offloading to CPU.""" + if not module_names: + return + try: + from torch.utils.checkpoint import CheckpointPolicy # noqa: F401 + except ImportError: + log_rank0("WARNING: torch.utils.checkpoint policy_fn not available, skipping offload") + return + for layer in layers: + for mod_name in module_names: + if mod_name in module_map: + submod = module_map[mod_name](layer) + if submod is not None: + wrap_offload(submod) + + +class CheckpointFunction(torch.autograd.Function): + """Reentrant activation checkpoint using custom autograd.Function. + + Adapted from Megatron-Core's CheckpointFunction. Key differences: + - No distribute_saved_activations (not needed for our TP implementation) + - No model-parallel RNG tracker (we use TE's built-in RNG management) + - Handles expert_bias restore for MoE router determinism + """ + + @staticmethod + def forward( + ctx: Any, run_function: Callable, preserve_rng_state: bool, *args: torch.Tensor + ) -> Any: + ctx.run_function = run_function + ctx.preserve_rng_state = preserve_rng_state + + # Save RNG states for deterministic recomputation. + if preserve_rng_state: + ctx.cpu_rng_state = torch.get_rng_state() + ctx.cuda_rng_state = torch.cuda.get_rng_state() + + # Run forward without gradient tracking — discard intermediate activations. + with torch.no_grad(): + outputs = run_function(*args) + + # Save inputs for recomputation in backward. + ctx.save_for_backward(*args) + return outputs + + @staticmethod + def backward(ctx: Any, *grad_outputs: torch.Tensor) -> tuple: + if not torch.autograd._is_checkpoint_valid(): + raise RuntimeError( + "Checkpointing is not compatible with .grad(), use .backward() instead" + ) + + inputs = ctx.saved_tensors + + # Fork RNG: restore forward-time states, then reset to current after recompute. + if ctx.preserve_rng_state: + current_cpu_rng = torch.get_rng_state() + current_cuda_rng = torch.cuda.get_rng_state() + torch.set_rng_state(ctx.cpu_rng_state) + torch.cuda.set_rng_state(ctx.cuda_rng_state) + + # Recompute forward pass with gradients enabled. + detached = tuple( + t.detach().requires_grad_(t.requires_grad) if isinstance(t, torch.Tensor) else t + for t in inputs + ) + with torch.enable_grad(): + outputs = ctx.run_function(*detached) + + # Restore RNG states. + if ctx.preserve_rng_state: + torch.set_rng_state(current_cpu_rng) + torch.cuda.set_rng_state(current_cuda_rng) + + if isinstance(outputs, torch.Tensor): + outputs = (outputs,) + + # Filter to outputs that need gradients. + outputs_with_grad = [] + grad_for_outputs = [] + for out, grad in zip(outputs, grad_outputs, strict=False): + if torch.is_tensor(out) and out.requires_grad: + outputs_with_grad.append(out) + grad_for_outputs.append(grad) + + if outputs_with_grad: + torch.autograd.backward(outputs_with_grad, grad_for_outputs) + + grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else None for inp in detached) + # None for run_function, None for preserve_rng_state, then grads for each input. + return (None, None) + grads + + +def wrap_checkpoint(module: nn.Module, *, preserve_rng_state: bool = True) -> None: + """Wrap a module's forward with reentrant activation checkpointing.""" + original_forward = module.forward + _routers = [m for m in module.modules() if hasattr(m, "expert_bias")] + + def _checkpointed_forward(*args, **kwargs): + # expert_bias is modified in-place by the router during forward. + # Save and restore it so the recomputation in backward sees the same values. + if _routers: + saved = [r.expert_bias.clone() for r in _routers] + call_count = [0] + + def _fwd(*a, **kw): + call_count[0] += 1 + if call_count[0] > 1: + for r, s in zip(_routers, saved, strict=False): + r.expert_bias.copy_(s) + return original_forward(*a, **kw) + + else: + _fwd = original_forward + + # CheckpointFunction.apply only accepts positional tensor args. + # Wrap kwargs into the function closure. + if kwargs: + + def _fn(*a): + return _fwd(*a, **kwargs) + + else: + _fn = _fwd + + return CheckpointFunction.apply(_fn, preserve_rng_state, *args) + + module.forward = _checkpointed_forward + + +def wrap_offload(module: nn.Module) -> None: + """Wrap a module's forward with activation offloading.""" + original_forward = module.forward + + def _offloaded_forward(*args, **kwargs): + return torch.utils.checkpoint.checkpoint( + original_forward, *args, use_reentrant=False, **kwargs + ) + + module.forward = _offloaded_forward + + +def log_rank0(msg: str) -> None: + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: + print(f"[megatron.lite] {msg}", flush=True) + + +def parse_recompute_spec(recompute: str | list[str] | None) -> list[str]: + """Parse a recompute spec into a list of module names.""" + if recompute is None or recompute == "none": + return [] + if recompute == "full": + return ["full"] + if isinstance(recompute, list): + return recompute + return recompute.split(",") + + +__all__ = [ + "CheckpointFunction", + "CheckpointWithoutOutput", + "ModuleMap", + "apply_offload", + "apply_recompute", + "parse_recompute_spec", + "wrap_checkpoint", + "wrap_offload", +] diff --git a/experimental/lite/megatron/lite/primitive/train_step.py b/experimental/lite/megatron/lite/primitive/train_step.py new file mode 100644 index 00000000000..630dc5f7701 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/train_step.py @@ -0,0 +1,182 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Reusable train-step primitives owned by Megatron Lite.""" + +from __future__ import annotations + +from collections.abc import Callable + +import torch +import torch.distributed as dist + +from megatron.lite.primitive.parallel import ParallelState +from megatron.lite.primitive.protocols import ExpertClassifierFn, default_expert_classifier + + +def run_microbatch_loop( + model, + data_iter, + num_microbatches: int, + forward_fn, + optimizer=None, + dist_opt: bool = False, + pre_forward_hook: Callable[[torch.Tensor], None] | None = None, + loss_fn: Callable | None = None, +): + """Run forward-backward over microbatches with loss accumulation. + + Args: + forward_fn: ``forward_fn(model, batch) -> dict`` with at least ``"loss"`` key + (when ``loss_fn`` is None) or model outputs (when ``loss_fn`` is provided). + pre_forward_hook: Optional callable ``hook(scale: torch.Tensor) -> None`` + invoked once per microbatch, right before ``forward_fn``. ``scale`` is + ``1.0 / num_microbatches`` (matches MC's ``schedules.forward_step``, + see `pipeline_parallel/schedules.py`). Used e.g. by MoE aux-loss + scale-setting; runtime stays model-agnostic by passing the hook + through from the model bundle's extras. + # TODO: once CP is supported, align with MC's + # `schedules:297`-style scale of `cp_group_size / num_microbatches` + # (currently assumes ``cp_group_size == 1``). + loss_fn: Optional external loss function. + ``loss_fn(model_output: dict, batch) -> (loss: Tensor, metrics: dict)``. + When provided, ``forward_fn`` output is passed to ``loss_fn`` instead of + reading ``out["loss"]`` directly. This enables RLHF policy/value losses. + """ + last_out = None + all_metrics: list[dict] = [] + for mb in range(num_microbatches): + batch = next(data_iter) + if pre_forward_hook is not None: + scale = torch.tensor(1.0 / num_microbatches, device="cuda") + pre_forward_hook(scale) + out = forward_fn(model, batch) + if dist_opt and optimizer is not None and mb == num_microbatches - 1: + optimizer.grad_sync_enabled = True + if loss_fn is not None: + loss, metrics = loss_fn(out, batch) + (loss / num_microbatches).backward() + out["loss"] = loss.detach() + all_metrics.append(metrics) + else: + (out["loss"] / num_microbatches).backward() + last_out = out + if last_out is not None and all_metrics: + last_out["_loss_fn_metrics"] = all_metrics + return last_out + + +def compute_and_clip_grad_norm( + model, + optimizer, + max_norm: float, + use_dist_opt: bool, + sp_params=None, + sp_group=None, + *, + report_global_norm: bool = False, + ps: ParallelState | None = None, + is_expert_param: ExpertClassifierFn = default_expert_classifier, +): + """SP AllReduce + finish grad sync + clip grad norm. Returns grad_norm.""" + if sp_params: + sp_grads = [p.grad for p in sp_params if p.grad is not None] + if sp_grads: + flat = torch.cat([g.view(-1) for g in sp_grads]) + dist.all_reduce(flat, op=dist.ReduceOp.SUM, group=sp_group) + offset = 0 + for g in sp_grads: + n = g.numel() + g.copy_(flat[offset : offset + n].view_as(g)) + offset += n + report_norm = None + if report_global_norm: + if ps is None: + raise ValueError("`ps` is required when `report_global_norm=True`.") + report_norm = compute_global_grad_norm(model, ps, is_expert_param=is_expert_param) + if use_dist_opt: + optimizer.finish_grad_sync() + return optimizer.clip_grad_norm() + local_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm) + return report_norm if report_norm is not None else local_norm + + +def compute_global_grad_norm( + model, ps: ParallelState, *, is_expert_param: ExpertClassifierFn = default_expert_classifier +) -> torch.Tensor: + """Compute benchmark global grad norm with dist-opt-aligned reduction order.""" + dense_sq = _bucketed_grad_sq_sum( + model, + include_param=lambda name: not is_expert_param(name), + replica_group=ps.dp_cp_group, + replica_size=ps.dp_cp_size, + ) + expert_sq = _bucketed_grad_sq_sum( + model, + include_param=is_expert_param, + replica_group=ps.ep_dp_group, + replica_size=ps.expert_dp_size, + ) + + if ps.tp_size > 1 and ps.tp_group is not None: + dist.all_reduce(dense_sq, group=ps.tp_group) + + if ps.ep_size > 1 and ps.ep_group is not None: + dist.all_reduce(expert_sq, group=ps.ep_group) + if ps.etp_size > 1 and ps.etp_group is not None: + dist.all_reduce(expert_sq, group=ps.etp_group) + + total_sq = dense_sq + expert_sq + if ps.pp_size > 1 and ps.pp_group is not None: + dist.all_reduce(total_sq, group=ps.pp_group) + return total_sq.sqrt() + + +def _bucketed_grad_sq_sum( + model, + *, + include_param: Callable[[str], bool], + replica_group, + replica_size: int, + max_bucket_bytes: int = 80 * 1024 * 1024, +) -> torch.Tensor: + """Accumulate squared norm after averaging replica grads within each bucket.""" + total_sq = torch.zeros(1, device="cuda") + bucket: list[torch.Tensor] = [] + bucket_bytes = 0 + + def flush_bucket() -> None: + nonlocal bucket_bytes, bucket, total_sq + if not bucket: + return + flat = torch.cat(bucket) + if replica_size > 1 and replica_group is not None: + dist.all_reduce(flat, group=replica_group) + flat.div_(replica_size) + total_sq += flat.float().norm().pow(2) + bucket = [] + bucket_bytes = 0 + + for name, param in model.named_parameters(): + if param.grad is None or not include_param(name): + continue + grad = param.grad.view(-1) + grad_bytes = grad.numel() * grad.element_size() + if bucket and bucket_bytes + grad_bytes > max_bucket_bytes: + flush_bucket() + bucket.append(grad) + bucket_bytes += grad_bytes + + flush_bucket() + return total_sq + + +def optimizer_step(optimizer) -> None: + """Execute optimizer step.""" + optimizer.step() + + +__all__ = [ + "compute_and_clip_grad_norm", + "compute_global_grad_norm", + "optimizer_step", + "run_microbatch_loop", +] diff --git a/experimental/lite/megatron/lite/primitive/utils.py b/experimental/lite/megatron/lite/primitive/utils.py new file mode 100644 index 00000000000..3ab4a5c33a7 --- /dev/null +++ b/experimental/lite/megatron/lite/primitive/utils.py @@ -0,0 +1,26 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import torch # pyright: ignore[reportMissingImports] + + +def build_fp8_recipe(train_config=None): + """Build the standard TE FP8 recipe (DelayedScaling, HYBRID format, H100).""" + from transformer_engine.common.recipe import DelayedScaling, Format + + return DelayedScaling(margin=0, fp8_format=Format.HYBRID) + + +def ensure_divisible(numerator: int, denominator: int, msg: str = "") -> int: + if numerator % denominator != 0: + detail = f" ({msg})" if msg else "" + raise ValueError(f"{numerator} is not divisible by {denominator}{detail}") + return numerator // denominator + + +def log_rank0(msg: str) -> None: + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: + print(f"[megatron.lite] {msg}", flush=True) + + +__all__ = ["build_fp8_recipe", "ensure_divisible", "log_rank0"] diff --git a/experimental/lite/megatron/lite/runtime/__init__.py b/experimental/lite/megatron/lite/runtime/__init__.py new file mode 100644 index 00000000000..9f8f1a8906c --- /dev/null +++ b/experimental/lite/megatron/lite/runtime/__init__.py @@ -0,0 +1,83 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Runtime entrypoints for Megatron Lite.""" + +from __future__ import annotations + +import importlib +from typing import TYPE_CHECKING + +from megatron.lite.runtime.contracts.config import RuntimeConfig + +if TYPE_CHECKING: + from megatron.lite.runtime.backends import Runtime + from megatron.lite.runtime.backends.bridge.config import BridgeConfig + from megatron.lite.runtime.backends.mlite.config import MegatronLiteConfig + from megatron.lite.runtime.contracts.data import ( + Batch, + ForwardResult, + ModelOutputs, + PackedBatch, + TrainBatch, + ) + from megatron.lite.runtime.contracts.handle import ModelHandle + + +def _runtime_registry() -> dict[str, str]: + from megatron.lite.runtime.backends import RUNTIME_REGISTRY + + return RUNTIME_REGISTRY + + +def register_runtime(name: str, module_path: str) -> None: + """Register a custom runtime backend. + + Args: + name: Backend name (used in ``RuntimeConfig.backend``). + module_path: Dotted module path providing a ``create(hf_path, cfg)`` function. + + Example:: + + from megatron.lite.runtime import register_runtime + register_runtime("my_backend", "my_package.my_runtime") + """ + _runtime_registry()[name] = module_path + + +def create_runtime(cfg: RuntimeConfig) -> Runtime: + """Create a Runtime instance for the given config.""" + mod = importlib.import_module(_runtime_registry()[cfg.backend]) + return mod.create(cfg.hf_path, cfg.backend_cfg) + + +def __getattr__(name: str): + _lazy = { + "Batch": "megatron.lite.runtime.contracts.data", + "BridgeConfig": "megatron.lite.runtime.backends.bridge.config", + "MegatronLiteConfig": "megatron.lite.runtime.backends.mlite.config", + "ForwardResult": "megatron.lite.runtime.contracts.data", + "ModelHandle": "megatron.lite.runtime.contracts.handle", + "ModelOutputs": "megatron.lite.runtime.contracts.data", + "PackedBatch": "megatron.lite.runtime.contracts.data", + "Runtime": "megatron.lite.runtime.backends", + "TrainBatch": "megatron.lite.runtime.contracts.data", + } + if name in _lazy: + mod = importlib.import_module(_lazy[name]) + return getattr(mod, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +__all__ = [ + "Batch", + "BridgeConfig", + "ForwardResult", + "MegatronLiteConfig", + "ModelHandle", + "ModelOutputs", + "PackedBatch", + "Runtime", + "RuntimeConfig", + "TrainBatch", + "create_runtime", + "register_runtime", +] diff --git a/experimental/lite/megatron/lite/runtime/backends/__init__.py b/experimental/lite/megatron/lite/runtime/backends/__init__.py new file mode 100644 index 00000000000..b4f0c729759 --- /dev/null +++ b/experimental/lite/megatron/lite/runtime/backends/__init__.py @@ -0,0 +1,178 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Runtime ABC and registry. + +Runtime API tiers +----------------- +L1 — **Pretrain Ready** (9 abstract methods, must implement): + build_model, save_checkpoint, load_checkpoint, + train_mode, eval_mode, + forward_backward, zero_grad, optimizer_step, lr_scheduler_step + +L2 — **RL Ready** (+ export_weights): + Enables RL frameworks to extract weights for the inference engine. + +L3 — **RL Best** (+ to): + Enables offloading model/optimizer/grad between training and rollout + phases to free GPU memory for the inference engine. + +A new backend only needs to implement L1 to work for pretraining. +Override ``export_weights`` and/or ``to`` to unlock higher tiers. +Check ``runtime.tier`` to see which level a backend supports. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Callable, Iterator +from typing import TYPE_CHECKING, Any, Literal + +if TYPE_CHECKING: + import torch + + from megatron.lite.runtime.contracts.data import ForwardResult + from megatron.lite.runtime.contracts.handle import ModelHandle + + +class Runtime(ABC): + """Base class for all runtime implementations. + + MegatronLiteRuntime and custom impls subclass this. + """ + + # ── L1: Pretrain Ready (必须实现) ──────────────────────────── + + # Model lifecycle + + @abstractmethod + def build_model(self, hf_path: str | None = None, cfg: Any = None, **kwargs) -> ModelHandle: + """Build model state for this runtime. + + Implementations should default to the ``hf_path`` / ``backend_cfg`` + captured by ``create_runtime(RuntimeConfig(...))``. Passing arguments to + ``build_model`` is supported as an advanced override path, but public + examples should prefer ``handle = rt.build_model()``. + """ + ... + + @abstractmethod + def save_checkpoint(self, handle: ModelHandle, path: str, **kwargs) -> None: ... + + @abstractmethod + def load_checkpoint(self, handle: ModelHandle, path: str, **kwargs) -> int: ... + + # Mode switching + + @abstractmethod + def train_mode(self, handle: ModelHandle) -> Any: ... + + @abstractmethod + def eval_mode(self, handle: ModelHandle) -> Any: ... + + # Training atoms + + @abstractmethod + def forward_backward( + self, + handle: ModelHandle, + data: Any, + loss_fn: Callable | None, + *, + num_microbatches: int = 1, + forward_only: bool = False, + ) -> ForwardResult: + """Forward + backward pass over data. + + Args: + num_microbatches: Number of microbatches to accumulate inside one + logical training step. + loss_fn: Optional external loss function with signature:: + + loss_fn(model_output: dict, batch) -> (loss: Tensor, metrics: dict) + + When ``loss_fn`` is None, the model computes loss internally + (standard pretrain/SFT path). When provided, ``model_output`` + is the dict returned by the model's forward (logits, log_probs, etc.) + and ``batch`` is the current microbatch. + """ + ... + + @abstractmethod + def zero_grad(self, handle: ModelHandle) -> None: ... + + @abstractmethod + def optimizer_step(self, handle: ModelHandle) -> tuple[bool, float, int | None]: + """Run optimizer step. + + Returns: + (update_successful, grad_norm, num_zeros_in_grad) + """ + ... + + @abstractmethod + def lr_scheduler_step(self, handle: ModelHandle) -> float | list[float]: ... + + # ── Parallel state queries ─────────────────────────────────── + + def is_mp_src_rank_with_outputs(self, handle: ModelHandle) -> bool: + """True if this rank is PP-last, TP-0, CP-0 (has full output).""" + return True # default: no parallelism, every rank has outputs + + # ── L2: RL Ready (覆盖即解锁) ─────────────────────────────── + + def export_weights(self, handle: ModelHandle, **kwargs) -> Iterator[tuple[str, torch.Tensor]]: + """Iterate over (name, tensor) pairs for HF-compatible weight export. + + Required by RL frameworks to send weights to the inference engine. + Override to unlock **RL Ready** tier. + """ + raise NotImplementedError( + f"{type(self).__name__} does not implement export_weights. " + "Implement it to unlock the RL Ready tier." + ) + + # ── L3: RL Best (覆盖即解锁) ──────────────────────────────── + + def to( + self, + handle: ModelHandle, + device: str, + *, + model: bool = True, + optimizer: bool = True, + grad: bool = True, + ) -> None: + """Move model / optimizer / gradients to *device*. + + Enables offloading between training and rollout phases so the + inference engine can reclaim GPU memory. + Override to unlock **RL Best** tier. + """ + raise NotImplementedError( + f"{type(self).__name__} does not implement to(). " + "Implement it to unlock the RL Best tier." + ) + + # ── Tier introspection ────────────────────────────────────── + + @property + def tier(self) -> Literal["pretrain", "rl_ready", "rl_best"]: + """Report the highest API tier this runtime supports.""" + cls = type(self) + has_export = cls.export_weights is not Runtime.export_weights + has_to = cls.to is not Runtime.to + if has_export and has_to: + return "rl_best" + if has_export: + return "rl_ready" + return "pretrain" + + +# --------------------------------------------------------------------------- +# Registry +# --------------------------------------------------------------------------- + +RUNTIME_REGISTRY: dict[str, str] = { + "bridge": "megatron.lite.runtime.backends.bridge", + "mbridge": "megatron.lite.runtime.backends.mbridge", + "mlite": "megatron.lite.runtime.backends.mlite", +} diff --git a/experimental/lite/megatron/lite/runtime/backends/bridge/__init__.py b/experimental/lite/megatron/lite/runtime/backends/bridge/__init__.py new file mode 100644 index 00000000000..71072edbdde --- /dev/null +++ b/experimental/lite/megatron/lite/runtime/backends/bridge/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Megatron-Bridge backend for the Megatron Lite runtime API.""" + +from __future__ import annotations + +from typing import Any + +from megatron.lite.runtime.backends import Runtime as RuntimeBase +from megatron.lite.runtime.backends.bridge.config import BridgeConfig +from megatron.lite.runtime.backends.bridge.runtime import BridgeRuntime + + +def create(hf_path: str, cfg: BridgeConfig | dict[str, Any]) -> RuntimeBase: + """Factory called by ``megatron.lite.runtime.create_runtime``.""" + return BridgeRuntime(hf_path, cfg) + + +__all__ = ["BridgeConfig", "BridgeRuntime", "create"] diff --git a/experimental/lite/megatron/lite/runtime/backends/bridge/config.py b/experimental/lite/megatron/lite/runtime/backends/bridge/config.py new file mode 100644 index 00000000000..e1df53439fa --- /dev/null +++ b/experimental/lite/megatron/lite/runtime/backends/bridge/config.py @@ -0,0 +1,73 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Megatron-Bridge backend configuration.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from megatron.lite.runtime.contracts.config import OptimizerConfig, ParallelConfig, pick_fields + + +@dataclass +class BridgeConfig: + """Config for ``BridgeRuntime``. + + The backend is intentionally thin: it lowers Megatron Lite's runtime + contract into Megatron-Bridge / Megatron-Core objects, while model-specific + mutations stay in examples as explicit benchmark hooks. + """ + + model_name: str = "auto" + + parallel: ParallelConfig = field(default_factory=ParallelConfig) + seed: int = 42 + param_offload: bool = False + optimizer_offload: bool = False + load_hf_weights: bool = True + build_optimizer: bool = True + + override_ddp_config: dict[str, Any] = field(default_factory=dict) + override_transformer_config: dict[str, Any] = field(default_factory=dict) + override_optimizer_config: dict[str, Any] = field(default_factory=dict) + + optimizer: OptimizerConfig = field(default_factory=OptimizerConfig) + + # Bench-only hook. Kept callable so examples can trim model configs without + # making the runtime know about benchmark profiles. + bridge_post_init: Any = None + + @classmethod + def from_dict(cls, cfg: dict[str, Any]) -> BridgeConfig: + """Construct ``BridgeConfig`` from a flat or nested mapping.""" + if "num_microbatches" in cfg: + raise ValueError( + "BridgeConfig does not accept `num_microbatches`; " + "pass it to Runtime.forward_backward(..., num_microbatches=...) instead." + ) + + parallel_src = cfg.get("parallel") + parallel_data = ( + pick_fields(ParallelConfig, parallel_src) if isinstance(parallel_src, dict) else {} + ) + parallel_data.update(pick_fields(ParallelConfig, cfg)) + parallel = ParallelConfig(**parallel_data) + + optimizer_src = cfg.get("optimizer", {}) + lr_src = cfg.get("lr_scheduler", {}) + optimizer_data: dict[str, Any] = {} + if isinstance(optimizer_src, dict): + optimizer_data.update(optimizer_src) + if isinstance(lr_src, dict): + optimizer_data.update(lr_src) + optimizer = OptimizerConfig(**pick_fields(OptimizerConfig, optimizer_data)) + + skip = {"parallel", "optimizer", "lr_scheduler"} + return cls( + **{k: v for k, v in pick_fields(cls, cfg).items() if k not in skip}, + parallel=parallel, + optimizer=optimizer, + ) + + +__all__ = ["BridgeConfig"] diff --git a/experimental/lite/megatron/lite/runtime/backends/bridge/runtime.py b/experimental/lite/megatron/lite/runtime/backends/bridge/runtime.py new file mode 100644 index 00000000000..7f4b7ce886e --- /dev/null +++ b/experimental/lite/megatron/lite/runtime/backends/bridge/runtime.py @@ -0,0 +1,811 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Runtime backend backed by Megatron-Bridge.""" + +from __future__ import annotations + +import logging +import os +from collections.abc import Iterator +from datetime import timedelta +from typing import Any + +import torch +import torch.distributed as dist + +from megatron.lite.primitive.optimizers.megatron_wrap import build_mc_optimizer_config +from megatron.lite.runtime.backends import Runtime as RuntimeBase +from megatron.lite.runtime.backends.bridge.config import BridgeConfig +from megatron.lite.runtime.contracts.data import Batch, ForwardResult, ModelOutputs +from megatron.lite.runtime.contracts.handle import ModelHandle +from megatron.lite.runtime.megatron_utils import ( + build_sharded_state_dict, + is_mp_src_rank_with_outputs, + load_model_to_gpu, + load_optimizer, + offload_model_to_cpu, + offload_optimizer, + register_training_hooks, +) + +logger = logging.getLogger(__name__) + + +class _MpuParallelState: + """Adapter exposing Megatron-Core mpu through ``ModelHandle`` properties.""" + + def __init__(self, mpu): + self._mpu = mpu + + @property + def dp_rank(self) -> int: + return self._mpu.get_data_parallel_rank() + + @property + def dp_size(self) -> int: + return self._mpu.get_data_parallel_world_size() + + @property + def dp_group(self): + return self._mpu.get_data_parallel_group() + + @property + def tp_rank(self) -> int: + return self._mpu.get_tensor_model_parallel_rank() + + @property + def tp_size(self) -> int: + return self._mpu.get_tensor_model_parallel_world_size() + + @property + def pp_rank(self) -> int: + return self._mpu.get_pipeline_model_parallel_rank() + + @property + def pp_size(self) -> int: + return self._mpu.get_pipeline_model_parallel_world_size() + + @property + def cp_rank(self) -> int: + return self._mpu.get_context_parallel_rank() + + @property + def cp_size(self) -> int: + return self._mpu.get_context_parallel_world_size() + + @property + def cp_group(self): + return self._mpu.get_context_parallel_group() + + +def _lower_transformer_overrides(cfg: BridgeConfig) -> dict[str, Any]: + overrides = {"attention_backend": "flash"} + overrides.update(cfg.override_transformer_config) + return overrides + + +def _bridge_hf_config(bridge): + hf_pretrained = getattr(bridge, "hf_pretrained", None) + if hf_pretrained is None: + return None + return getattr(hf_pretrained, "config", hf_pretrained) + + +def _lower_provider_value(key: str, value: Any) -> Any: + if key == "attention_backend" and isinstance(value, str): + from megatron.core.transformer.enums import AttnBackend + + return AttnBackend[value] + return value + + +def _configure_provider(provider, cfg: BridgeConfig) -> None: + from megatron.lite.primitive.deterministic import deterministic_requested + + p = cfg.parallel + provider.tensor_model_parallel_size = p.tp + provider.pipeline_model_parallel_size = p.pp + provider.context_parallel_size = p.cp + provider.expert_model_parallel_size = p.ep + if p.etp is not None: + provider.expert_tensor_parallel_size = p.etp + if p.vpp > 1: + provider.virtual_pipeline_model_parallel_size = p.vpp + provider.sequence_parallel = p.tp > 1 + provider.bf16 = True + provider.fp16 = False + provider.deterministic_mode = deterministic_requested() + if provider.deterministic_mode: + os.environ.setdefault("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + + for key, value in _lower_transformer_overrides(cfg).items(): + setattr(provider, key, _lower_provider_value(key, value)) + + +def _register_bridge_compat_aliases() -> None: + """Register local Megatron-Bridge aliases for supported checkpoint variants.""" + from megatron.bridge.models.conversion import model_bridge + from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry + from megatron.bridge.models.conversion.param_mapping import ( + AutoMapping, + GatedMLPMapping, + GDNConv1dMapping, + QKVMapping, + ReplicatedMapping, + RMSNorm2ZeroCenteredRMSNormMapping, + merge_gdn_linear_weights, + split_gdn_linear_weights, + ) + from megatron.bridge.models.qwen.qwen3_next_bridge import Qwen3NextBridge + from megatron.bridge.utils.common_utils import extract_expert_number_from_param + from megatron.core.models.gpt.gpt_model import GPTModel + + class Qwen35SplitGDNMapping(AutoMapping): + """Bridge mapping for Qwen3.5 split GDN in-proj weights.""" + + def __init__(self, megatron_param: str, qkv: str, z: str, b: str, a: str): + super().__init__( + megatron_param=megatron_param, hf_param={"qkv": qkv, "z": z, "b": b, "a": a} + ) + self._tp_mapping = AutoMapping(megatron_param, megatron_param) + + def hf_to_megatron( + self, hf_weights: dict[str, torch.Tensor], megatron_module + ) -> torch.Tensor: + if self.tp_rank == 0: + config = self._get_config(megatron_module) + qkvz = torch.cat([hf_weights["qkv"], hf_weights["z"]], dim=0) + ba = torch.cat([hf_weights["b"], hf_weights["a"]], dim=0) + merged = merge_gdn_linear_weights(config, qkvz, ba, tp_size=self.tp_size) + else: + merged = None + return self._tp_mapping.hf_to_megatron(merged, megatron_module) + + def megatron_to_hf(self, megatron_weights, megatron_module) -> dict[str, torch.Tensor]: + if megatron_weights is not None: + megatron_weights = self.maybe_dequantize(megatron_weights) + + if megatron_module is None: + config = self.broadcast_obj_from_pp_rank(None) + else: + config = self._get_config(megatron_module) + config = self.broadcast_obj_from_pp_rank(config) + + packed_dict = self._tp_mapping.megatron_to_hf(megatron_weights, megatron_module) + if not packed_dict: + return {} + + packed = next(iter(packed_dict.values())) + qkvz, ba = split_gdn_linear_weights(config, packed, tp_size=self.tp_size) + qk_dim = config.linear_key_head_dim * config.linear_num_key_heads + v_dim = config.linear_value_head_dim * config.linear_num_value_heads + qkv, z = qkvz.split([2 * qk_dim + v_dim, v_dim], dim=0) + b, a = ba.chunk(2, dim=0) + return { + self.hf_param["qkv"]: qkv, + self.hf_param["z"]: z, + self.hf_param["b"]: b, + self.hf_param["a"]: a, + } + + def resolve(self, captures): + megatron_param, hf_param = self._resolve_names(captures) + return type(self)( + megatron_param, hf_param["qkv"], hf_param["z"], hf_param["b"], hf_param["a"] + ) + + class Qwen35PackedExpertDownMapping(AutoMapping): + """Bridge mapping for Qwen3.5 packed expert down-projection weights.""" + + def __init__(self, megatron_param: str, hf_param: str, permute_dims=None): + super().__init__( + megatron_param=megatron_param, hf_param=hf_param, permute_dims=permute_dims + ) + self.allow_hf_name_mismatch = True + + def hf_to_megatron(self, hf_weights: torch.Tensor, megatron_module) -> torch.Tensor: + expert_number = extract_expert_number_from_param(self.megatron_param) + expert_weight = hf_weights[expert_number].contiguous() + return super().hf_to_megatron(expert_weight, megatron_module) + + def megatron_to_hf(self, megatron_weights, megatron_module) -> dict[str, torch.Tensor]: + converted = super().megatron_to_hf(megatron_weights, megatron_module) + return converted + + def _validate_patterns(self, *args, **kwargs): + pass + + class Qwen35RouterMapping(AutoMapping): + """Bridge mapping for Qwen3.5 router weights with bench expert truncation.""" + + def hf_to_megatron(self, hf_weights: torch.Tensor, megatron_module) -> torch.Tensor: + config = self._get_config(megatron_module) + num_experts = getattr(config, "num_moe_experts", None) + if num_experts is not None and hf_weights.shape[0] != num_experts: + hf_weights = hf_weights[:num_experts].contiguous() + return super().hf_to_megatron(hf_weights, megatron_module) + + class Qwen35PackedExpertGateUpMapping(AutoMapping): + """Bridge mapping for Qwen3.5 packed expert gate/up projection weights.""" + + def __init__(self, megatron_param: str, hf_param: str, permute_dims=None): + super().__init__( + megatron_param=megatron_param, hf_param=hf_param, permute_dims=permute_dims + ) + self.allow_hf_name_mismatch = True + GatedMLPMapping._validate_patterns = lambda *args, **kwargs: None + self._gated_mapping = GatedMLPMapping( + megatron_param=self.megatron_param, + gate=f"{self.hf_param}.gate", + up=f"{self.hf_param}.up", + ) + + def hf_to_megatron(self, hf_weights: torch.Tensor, megatron_module) -> torch.Tensor: + expert_number = extract_expert_number_from_param(self.megatron_param) + expert_weight = hf_weights[expert_number].contiguous() + gate, up = torch.chunk(expert_weight, 2, dim=0) + return self._gated_mapping.hf_to_megatron({"gate": gate, "up": up}, megatron_module) + + def megatron_to_hf(self, megatron_weights, megatron_module) -> dict[str, torch.Tensor]: + converted = self._gated_mapping.megatron_to_hf(megatron_weights, megatron_module) + if not converted: + return {} + + fused = {} + for name, tensor in converted.items(): + if not name.endswith(".gate"): + continue + base_name = name[: -len(".gate")] + up_tensor = converted.get(f"{base_name}.up") + if up_tensor is None: + continue + gate_tensor = tensor.contiguous() + up_tensor = up_tensor.contiguous() + fused[base_name] = torch.stack( + [gate_tensor, up_tensor], dim=0 if up_tensor.ndim == 2 else 1 + ) + return fused + + def _validate_patterns(self, *args, **kwargs): + pass + + class Qwen35MoEBridge(Qwen3NextBridge): + """Megatron-Bridge Qwen3-Next bridge adjusted for Qwen3.5 HF naming.""" + + def _text_config(self, hf_pretrained): + config = getattr(hf_pretrained, "config", hf_pretrained) + text_config = getattr(config, "text_config", config) + + if getattr(text_config, "intermediate_size", None) is None: + text_config.intermediate_size = 5120 + + rope_parameters = getattr(text_config, "rope_parameters", None) + if isinstance(rope_parameters, dict): + rope_theta = rope_parameters.get("rope_theta") + if rope_theta is not None: + text_config.rope_theta = rope_theta + partial_rotary_factor = rope_parameters.get("partial_rotary_factor") + if partial_rotary_factor is not None: + text_config.partial_rotary_factor = partial_rotary_factor + + if not hasattr(text_config, "tie_word_embeddings") and hasattr( + config, "tie_word_embeddings" + ): + text_config.tie_word_embeddings = config.tie_word_embeddings + + return text_config + + def provider_bridge(self, hf_pretrained): + text_config = self._text_config(hf_pretrained) + shim = type("_Qwen35TextConfigShim", (), {"config": text_config})() + provider = super().provider_bridge(shim) + + aux_loss = getattr(text_config, "router_aux_loss_coef", None) + if aux_loss is not None: + provider.moe_aux_loss_coeff = aux_loss + + return provider + + def mapping_registry(self): + prefix = "model.language_model" + param_mappings = { + "embedding.word_embeddings.weight": f"{prefix}.embed_tokens.weight", + "output_layer.weight": "lm_head.weight", + "decoder.final_layernorm.weight": f"{prefix}.norm.weight", + "decoder.layers.*.pre_mlp_layernorm.weight": f"{prefix}.layers.*.post_attention_layernorm.weight", + "decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": ( + f"{prefix}.layers.*.input_layernorm.weight" + ), + "decoder.layers.*.self_attention.q_layernorm.weight": f"{prefix}.layers.*.self_attn.q_norm.weight", + "decoder.layers.*.self_attention.k_layernorm.weight": f"{prefix}.layers.*.self_attn.k_norm.weight", + "decoder.layers.*.self_attention.linear_proj.weight": f"{prefix}.layers.*.self_attn.o_proj.weight", + "decoder.layers.*.self_attention.in_proj.layer_norm_weight": ( + f"{prefix}.layers.*.input_layernorm.weight" + ), + "decoder.layers.*.self_attention.out_proj.weight": f"{prefix}.layers.*.linear_attn.out_proj.weight", + "decoder.layers.*.self_attention.A_log": f"{prefix}.layers.*.linear_attn.A_log", + "decoder.layers.*.self_attention.dt_bias": f"{prefix}.layers.*.linear_attn.dt_bias", + } + + mapping_list = [ + AutoMapping(megatron_param=megatron_param, hf_param=hf_param) + for megatron_param, hf_param in param_mappings.items() + ] + AutoMapping.register_module_type("SharedExpertMLP", "column") + AutoMapping.register_module_type("GatedDeltaNet", "column") + + mapping_list.extend( + [ + QKVMapping( + megatron_param="decoder.layers.*.self_attention.linear_qkv.weight", + q=f"{prefix}.layers.*.self_attn.q_proj.weight", + k=f"{prefix}.layers.*.self_attn.k_proj.weight", + v=f"{prefix}.layers.*.self_attn.v_proj.weight", + ), + GDNConv1dMapping( + megatron_param="decoder.layers.*.self_attention.conv1d.weight", + hf_param=f"{prefix}.layers.*.linear_attn.conv1d.weight", + ), + Qwen35SplitGDNMapping( + megatron_param="decoder.layers.*.self_attention.in_proj.weight", + qkv=f"{prefix}.layers.*.linear_attn.in_proj_qkv.weight", + z=f"{prefix}.layers.*.linear_attn.in_proj_z.weight", + b=f"{prefix}.layers.*.linear_attn.in_proj_b.weight", + a=f"{prefix}.layers.*.linear_attn.in_proj_a.weight", + ), + Qwen35RouterMapping( + megatron_param="decoder.layers.*.mlp.router.weight", + hf_param=f"{prefix}.layers.*.mlp.gate.weight", + ), + Qwen35PackedExpertGateUpMapping( + megatron_param="decoder.layers.*.mlp.experts.linear_fc1.weight*", + hf_param=f"{prefix}.layers.*.mlp.experts.gate_up_proj", + ), + Qwen35PackedExpertDownMapping( + megatron_param="decoder.layers.*.mlp.experts.linear_fc2.weight*", + hf_param=f"{prefix}.layers.*.mlp.experts.down_proj", + ), + GatedMLPMapping( + megatron_param="decoder.layers.*.mlp.shared_experts.linear_fc1.weight", + gate=f"{prefix}.layers.*.mlp.shared_expert.gate_proj.weight", + up=f"{prefix}.layers.*.mlp.shared_expert.up_proj.weight", + ), + AutoMapping( + megatron_param="decoder.layers.*.mlp.shared_experts.linear_fc2.weight", + hf_param=f"{prefix}.layers.*.mlp.shared_expert.down_proj.weight", + ), + ReplicatedMapping( + megatron_param="decoder.layers.*.mlp.shared_experts.gate_weight", + hf_param=f"{prefix}.layers.*.mlp.shared_expert_gate.weight", + ), + RMSNorm2ZeroCenteredRMSNormMapping( + "decoder.layers.*.self_attention.out_norm.weight", + f"{prefix}.layers.*.linear_attn.norm.weight", + ), + ] + ) + + return MegatronMappingRegistry(*mapping_list) + + registry = getattr(model_bridge.get_model_bridge, "_exact_types", {}) + for source in ("Qwen3_5MoeForConditionalGeneration", "Qwen3_5MoeForCausalLM"): + if source not in registry: + model_bridge.register_bridge_implementation( + source=source, target=GPTModel, bridge_class=Qwen35MoEBridge + ) + + +def _build_bridge(hf_path: str, cfg: BridgeConfig): + """Build Megatron-Bridge AutoBridge lazily from an HF model path.""" + from megatron.bridge import AutoBridge + + _register_bridge_compat_aliases() + bridge = AutoBridge.from_hf_pretrained(hf_path, trust_remote_code=True) + hf_config = _bridge_hf_config(bridge) + if hf_config is not None and not hasattr(hf_config, "rope_theta"): + hf_config.rope_theta = hf_config.to_dict().get("rope_theta", 1000000.0) + + if callable(cfg.bridge_post_init): + cfg.bridge_post_init(bridge) + + return bridge + + +def _build_optimizer(model_list: list, cfg: BridgeConfig): + from megatron.core.optimizer import get_megatron_optimizer + + return get_megatron_optimizer( + config=build_mc_optimizer_config( + cfg.optimizer, override_optimizer_config=cfg.override_optimizer_config + ), + model_chunks=model_list, + ) + + +def _build_lr_scheduler(optimizer, cfg: BridgeConfig): + opt = cfg.optimizer + total_steps = opt.total_training_steps + if total_steps <= 0: + return None + + from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler + + warmup_steps = opt.lr_warmup_steps + if warmup_steps <= 0 and opt.lr_warmup_steps_ratio > 0: + warmup_steps = int(opt.lr_warmup_steps_ratio * total_steps) + warmup_steps = max(warmup_steps, 0) + decay_steps = opt.lr_decay_steps if opt.lr_decay_steps is not None else total_steps + + return OptimizerParamScheduler( + optimizer, + init_lr=opt.lr_warmup_init, + max_lr=opt.lr, + min_lr=opt.min_lr, + lr_warmup_steps=warmup_steps, + lr_decay_steps=decay_steps, + lr_decay_style=opt.lr_decay_style, + start_wd=opt.weight_decay, + end_wd=opt.weight_decay, + wd_incr_steps=total_steps, + wd_incr_style=opt.weight_decay_incr_style, + use_checkpoint_opt_param_scheduler=opt.use_checkpoint_opt_param_scheduler, + override_opt_param_scheduler=not opt.use_checkpoint_opt_param_scheduler, + wsd_decay_steps=opt.lr_wsd_decay_steps, + lr_wsd_decay_style=opt.lr_wsd_decay_style, + ) + + +def _build_ddp_config(cfg: BridgeConfig): + from megatron.core.distributed import DistributedDataParallelConfig + + ddp_kwargs = { + "use_distributed_optimizer": True, + "overlap_grad_reduce": False, + "grad_reduce_in_fp32": True, + } + ddp_kwargs.update(cfg.override_ddp_config) + return DistributedDataParallelConfig(**ddp_kwargs) + + +def _resolve_benchmark_protocol(cfg: BridgeConfig, bridge) -> Any | None: + """Best-effort protocol lookup for model stats used by bench examples.""" + from megatron.lite.model.registry import get_train_runtime_module, resolve_model_type_from_hf + + model_name = cfg.model_name + if model_name == "auto": + try: + model_name = resolve_model_type_from_hf(_bridge_hf_config(bridge)) + except ValueError: + return None + + try: + return get_train_runtime_module(model_name) + except ValueError: + return None + + +def _as_data_iter(data: Any): + if hasattr(data, "__next__"): + return data + if isinstance(data, list): + return iter(data) + return iter([data]) + + +class BridgeRuntime(RuntimeBase): + """Megatron-Bridge training backend using Megatron-Core optimizer state.""" + + def __init__(self, hf_path: str, cfg: BridgeConfig | dict[str, Any]): + self._hf_path = hf_path + self._cfg = cfg if isinstance(cfg, BridgeConfig) else BridgeConfig.from_dict(cfg) + self._offload_param = self._cfg.param_offload + self._offload_optimizer = self._cfg.optimizer_offload + + def build_model( + self, hf_path: str | None = None, cfg: BridgeConfig | dict[str, Any] | None = None, **kwargs + ) -> ModelHandle: + from megatron.core import parallel_state as mpu + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + from megatron.core.transformer.enums import ModelType + + if cfg is None: + rt_cfg = self._cfg + elif isinstance(cfg, BridgeConfig): + rt_cfg = cfg + else: + rt_cfg = BridgeConfig.from_dict(cfg) + hf_path = hf_path or self._hf_path + + if not dist.is_initialized(): + dist.init_process_group("nccl", timeout=timedelta(minutes=10)) + torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count()) + + p = rt_cfg.parallel + init_kwargs = dict( + tensor_model_parallel_size=p.tp, + pipeline_model_parallel_size=p.pp, + expert_model_parallel_size=p.ep, + context_parallel_size=p.cp, + ) + if p.etp is not None: + init_kwargs["expert_tensor_parallel_size"] = p.etp + if p.vpp > 1: + init_kwargs["virtual_pipeline_model_parallel_size"] = p.vpp + + if not mpu.model_parallel_is_initialized(): + mpu.initialize_model_parallel(**init_kwargs) + model_parallel_cuda_manual_seed(rt_cfg.seed) + + bridge = _build_bridge(hf_path, rt_cfg) + provider = bridge.to_megatron_provider( + load_weights=rt_cfg.load_hf_weights, hf_path=hf_path if rt_cfg.load_hf_weights else None + ) + _configure_provider(provider, rt_cfg) + if hasattr(provider, "finalize"): + provider.finalize() + + model_list = provider.provide_distributed_model( + model_type=ModelType.encoder_or_decoder, + wrap_with_ddp=True, + ddp_config=_build_ddp_config(rt_cfg), + bf16=True, + ) + + optimizer = _build_optimizer(model_list, rt_cfg) if rt_cfg.build_optimizer else None + lr_scheduler = _build_lr_scheduler(optimizer, rt_cfg) if optimizer is not None else None + register_training_hooks(model_list, optimizer) + + if self._offload_param: + offload_model_to_cpu(model_list) + if self._offload_optimizer and optimizer is not None: + offload_optimizer(optimizer) + + logger.info("BridgeRuntime: model built, tp=%d ep=%d pp=%d cp=%d", p.tp, p.ep, p.pp, p.cp) + + return ModelHandle( + model=model_list[0], + optimizer=optimizer, + lr_scheduler=lr_scheduler, + parallel_state=_MpuParallelState(mpu), + config=rt_cfg, + _extras={ + "bridge": bridge, + "provider": provider, + "model_list": model_list, + "mpu": mpu, + "model_cfg": _bridge_hf_config(bridge), + "protocol": _resolve_benchmark_protocol(rt_cfg, bridge), + "optimizer_backend": "distopt" if optimizer is not None else "none", + "world_size": dist.get_world_size(), + }, + ) + + def forward_backward( + self, + handle: ModelHandle, + data: Any, + loss_fn, + *, + num_microbatches: int = 1, + forward_only: bool = False, + ) -> ForwardResult: + from megatron.core import parallel_state as mpu + from megatron.core.pipeline_parallel.schedules import get_forward_backward_func + + if num_microbatches < 1: + raise ValueError("num_microbatches must be >= 1") + + model_list = handle._extras["model_list"] + data_iter = _as_data_iter(data) + last_loss: list[float | None] = [None] + + def _fwd_step(data_iterator, model): + sample = next(data_iterator) + if isinstance(sample, Batch): + sample = { + "input_ids": sample["input_ids"], + "labels": sample["labels"], + "position_ids": getattr(sample, "position_ids", None), + } + if not isinstance(sample, dict): + raise TypeError( + f"BridgeRuntime expected dict or Batch data, got {type(sample).__name__}." + ) + + output_tensor = model( + input_ids=sample["input_ids"], + position_ids=sample.get("position_ids"), + attention_mask=sample.get("attention_mask"), + labels=sample["labels"], + packed_seq_params=sample.get("packed_seq_params"), + ) + if isinstance(output_tensor, tuple): + output_tensor = output_tensor[0] + + def _mc_loss_fn(output_tensor, non_loss_data=False): + if loss_fn is not None: + loss, _metrics = loss_fn({"output_tensor": output_tensor}, sample) + else: + loss = output_tensor.mean() + last_loss[0] = float(loss.detach().item()) + return loss, {} + + return output_tensor, _mc_loss_fn + + vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() + if vpp_size is not None and vpp_size > 1: + batches = [next(data_iter) for _ in range(num_microbatches)] + batch_generator = [iter(batches) for _ in range(vpp_size)] + else: + batch_generator = data_iter + + get_forward_backward_func()( + forward_step_func=_fwd_step, + data_iterator=batch_generator, + model=model_list, + num_microbatches=num_microbatches, + forward_only=forward_only, + seq_length=1, + micro_batch_size=1, + ) + + if not forward_only: + from megatron.core.distributed.finalize_model_grads import finalize_model_grads + + finalize_model_grads(model_list) + + loss_val = last_loss[0] + if mpu.get_pipeline_model_parallel_world_size() > 1: + loss_t = torch.tensor([loss_val or 0.0], device="cuda") + dist.broadcast(loss_t, src=mpu.get_pipeline_model_parallel_last_rank()) + loss_val = float(loss_t.item()) + + result_loss = torch.tensor(loss_val or 0.0) + return ForwardResult( + model_output=ModelOutputs(loss=result_loss), + metrics={"loss": loss_val if loss_val is not None else 0.0}, + ) + + def zero_grad(self, handle: ModelHandle) -> None: + if handle._optimizer is not None: + handle._optimizer.zero_grad() + for model in handle._extras["model_list"]: + if handle._optimizer is None: + model.zero_grad(set_to_none=True) + if hasattr(model, "zero_grad_buffer"): + model.zero_grad_buffer() + + def optimizer_step(self, handle: ModelHandle) -> tuple[bool, float, int | None]: + if handle._optimizer is None: + return True, 0.0, 0 + update_successful, grad_norm, num_zeros = handle._optimizer.step() + return update_successful, float(grad_norm), num_zeros + + def lr_scheduler_step(self, handle: ModelHandle) -> float | list[float]: + if handle._lr_scheduler is not None: + handle._lr_scheduler.step(1) + return handle._optimizer.param_groups[0]["lr"] + return 0.0 + + def is_mp_src_rank_with_outputs(self, handle: ModelHandle) -> bool: + return is_mp_src_rank_with_outputs() + + def to( + self, + handle: ModelHandle, + device: str, + *, + model: bool = True, + optimizer: bool = True, + grad: bool = True, + ) -> None: + model_list = handle._extras["model_list"] + opt = handle._optimizer + if device == "cuda": + if model: + load_model_to_gpu(model_list, load_grad=grad) + if optimizer and opt is not None: + load_optimizer(opt) + elif device == "cpu": + if model: + offload_model_to_cpu(model_list) + if optimizer and opt is not None: + offload_optimizer(opt) + else: + raise ValueError(f"BridgeRuntime.to supports only 'cpu' or 'cuda', got {device!r}.") + + def train_mode(self, handle: ModelHandle): + return _BridgeTrainCtx(self, handle) + + def eval_mode(self, handle: ModelHandle): + return _BridgeEvalCtx(self, handle) + + def export_weights(self, handle: ModelHandle, **kwargs) -> Iterator[tuple[str, torch.Tensor]]: + bridge = handle._extras["bridge"] + model_list = handle._extras["model_list"] + load_model_to_gpu(model_list, load_grad=False) + return bridge.export_hf_weights(model_list, cpu=bool(kwargs.get("cpu", False))) + + def save_checkpoint(self, handle: ModelHandle, path: str, **kwargs) -> None: + model_list = handle._extras["model_list"] + opt = handle._optimizer + load_model_to_gpu(model_list, load_grad=True) + + from megatron.core import dist_checkpointing + + state = build_sharded_state_dict(model_list, opt, handle._lr_scheduler) + os.makedirs(path, exist_ok=True) + dist_checkpointing.save(state, path) + dist.barrier() + + if self._offload_param: + offload_model_to_cpu(model_list) + + def load_checkpoint(self, handle: ModelHandle, path: str, **kwargs) -> None: + model_list = handle._extras["model_list"] + opt = handle._optimizer + load_model_to_gpu(model_list, load_grad=True) + + from megatron.core import dist_checkpointing + + state = build_sharded_state_dict(model_list, opt, handle._lr_scheduler) + dist_checkpointing.load(state, path) + + if self._offload_param: + offload_model_to_cpu(model_list) + if self._offload_optimizer and opt is not None: + offload_optimizer(opt) + + +class _BridgeTrainCtx: + def __init__(self, runtime: BridgeRuntime, handle: ModelHandle): + self._runtime = runtime + self._handle = handle + + def __enter__(self): + if self._runtime._offload_param or self._runtime._offload_optimizer: + self._runtime.to( + self._handle, + "cuda", + model=self._runtime._offload_param, + optimizer=self._runtime._offload_optimizer, + grad=self._runtime._offload_param, + ) + for model in self._handle._extras["model_list"]: + model.train() + return self + + def __exit__(self, *exc): + self._runtime.zero_grad(self._handle) + if self._runtime._offload_param or self._runtime._offload_optimizer: + self._runtime.to( + self._handle, + "cpu", + model=self._runtime._offload_param, + optimizer=self._runtime._offload_optimizer, + grad=self._runtime._offload_param, + ) + return False + + +class _BridgeEvalCtx: + def __init__(self, runtime: BridgeRuntime, handle: ModelHandle): + self._runtime = runtime + self._handle = handle + self._prev_grad = torch.is_grad_enabled() + + def __enter__(self): + if self._runtime._offload_param: + self._runtime.to(self._handle, "cuda", model=True, optimizer=False, grad=False) + for model in self._handle._extras["model_list"]: + model.eval() + torch.set_grad_enabled(False) + return self + + def __exit__(self, *exc): + torch.set_grad_enabled(self._prev_grad) + if self._runtime._offload_param: + self._runtime.to(self._handle, "cpu", model=True, optimizer=False, grad=False) + return False + + +__all__ = ["BridgeRuntime"] diff --git a/experimental/lite/megatron/lite/runtime/backends/mbridge/__init__.py b/experimental/lite/megatron/lite/runtime/backends/mbridge/__init__.py new file mode 100644 index 00000000000..993aa2576c0 --- /dev/null +++ b/experimental/lite/megatron/lite/runtime/backends/mbridge/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""mbridge backend for the Megatron Lite runtime API.""" + +from __future__ import annotations + +from typing import Any + +from megatron.lite.runtime.backends import Runtime as RuntimeBase +from megatron.lite.runtime.backends.bridge.config import BridgeConfig +from megatron.lite.runtime.backends.mbridge.runtime import MBridgeRuntime + + +def create(hf_path: str, cfg: BridgeConfig | dict[str, Any]) -> RuntimeBase: + """Factory called by ``megatron.lite.runtime.create_runtime``.""" + return MBridgeRuntime(hf_path, cfg) + + +__all__ = ["BridgeConfig", "MBridgeRuntime", "create"] diff --git a/experimental/lite/megatron/lite/runtime/backends/mbridge/runtime.py b/experimental/lite/megatron/lite/runtime/backends/mbridge/runtime.py new file mode 100644 index 00000000000..faed2738dfa --- /dev/null +++ b/experimental/lite/megatron/lite/runtime/backends/mbridge/runtime.py @@ -0,0 +1,167 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Runtime backend backed by the legacy ``mbridge`` package.""" + +from __future__ import annotations + +import logging +from collections.abc import Iterator +from datetime import timedelta +from typing import Any + +import torch +import torch.distributed as dist + +from megatron.lite.runtime.backends.bridge.config import BridgeConfig +from megatron.lite.runtime.backends.bridge.runtime import ( + BridgeRuntime, + _build_lr_scheduler, + _build_optimizer, + _lower_transformer_overrides, + _MpuParallelState, +) +from megatron.lite.runtime.contracts.handle import ModelHandle +from megatron.lite.runtime.megatron_utils import ( + load_model_to_gpu, + offload_model_to_cpu, + offload_optimizer, + register_training_hooks, +) + +logger = logging.getLogger(__name__) + + +def _build_mbridge(hf_path: str, cfg: BridgeConfig): + """Build the legacy mbridge AutoBridge lazily from an HF model path.""" + from mbridge import AutoBridge + + from megatron.lite.primitive.deterministic import deterministic_requested + + bridge = AutoBridge.from_pretrained(hf_path, trust_remote_code=True) + bridge.set_extra_args(sequence_parallel=cfg.parallel.tp > 1) + + transformer_overrides = _lower_transformer_overrides(cfg) + if transformer_overrides: + bridge.set_extra_args(**transformer_overrides) + + if not hasattr(bridge.hf_config, "rope_theta"): + bridge.hf_config.rope_theta = bridge.hf_config.to_dict().get("rope_theta", 1000000.0) + + bridge.set_extra_args(bf16=True, fp16=False) + + if cfg.model_name == "qwen3_5": + bridge.set_extra_args( + deterministic_mode=deterministic_requested(), fused_single_qkv_rope=False + ) + + if callable(cfg.bridge_post_init): + cfg.bridge_post_init(bridge) + + return bridge + + +def _resolve_mbridge_benchmark_protocol(cfg: BridgeConfig, bridge) -> Any | None: + """Best-effort protocol lookup for model stats used by bench examples.""" + from megatron.lite.model.registry import get_train_runtime_module, resolve_model_type_from_hf + + model_name = cfg.model_name + if model_name == "auto": + try: + model_name = resolve_model_type_from_hf(bridge.hf_config) + except ValueError: + return None + + try: + return get_train_runtime_module(model_name) + except ValueError: + return None + + +class MBridgeRuntime(BridgeRuntime): + """mbridge training backend using Megatron-Core optimizer state.""" + + def build_model( + self, hf_path: str | None = None, cfg: BridgeConfig | dict[str, Any] | None = None, **kwargs + ) -> ModelHandle: + from megatron.core import parallel_state as mpu + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + from megatron.core.transformer.enums import ModelType + + if cfg is None: + rt_cfg = self._cfg + elif isinstance(cfg, BridgeConfig): + rt_cfg = cfg + else: + rt_cfg = BridgeConfig.from_dict(cfg) + hf_path = hf_path or self._hf_path + + if not dist.is_initialized(): + dist.init_process_group("nccl", timeout=timedelta(minutes=10)) + torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count()) + + p = rt_cfg.parallel + init_kwargs = dict( + tensor_model_parallel_size=p.tp, + pipeline_model_parallel_size=p.pp, + expert_model_parallel_size=p.ep, + context_parallel_size=p.cp, + ) + if p.etp is not None: + init_kwargs["expert_tensor_parallel_size"] = p.etp + if p.vpp > 1: + init_kwargs["virtual_pipeline_model_parallel_size"] = p.vpp + + if not mpu.model_parallel_is_initialized(): + mpu.initialize_model_parallel(**init_kwargs) + model_parallel_cuda_manual_seed(rt_cfg.seed) + + bridge = _build_mbridge(hf_path, rt_cfg) + + ddp_config = { + "use_distributed_optimizer": True, + "overlap_grad_reduce": False, + "grad_reduce_in_fp32": True, + } + ddp_config.update(rt_cfg.override_ddp_config) + + model_list = bridge.get_model( + model_type=ModelType.encoder_or_decoder, wrap_with_ddp=True, ddp_config=ddp_config + ) + if rt_cfg.load_hf_weights: + bridge.load_weights(model_list, hf_path, memory_efficient=True) + + optimizer = _build_optimizer(model_list, rt_cfg) if rt_cfg.build_optimizer else None + lr_scheduler = _build_lr_scheduler(optimizer, rt_cfg) if optimizer is not None else None + register_training_hooks(model_list, optimizer) + + if self._offload_param: + offload_model_to_cpu(model_list) + if self._offload_optimizer and optimizer is not None: + offload_optimizer(optimizer) + + logger.info("MBridgeRuntime: model built, tp=%d ep=%d pp=%d cp=%d", p.tp, p.ep, p.pp, p.cp) + + return ModelHandle( + model=model_list[0], + optimizer=optimizer, + lr_scheduler=lr_scheduler, + parallel_state=_MpuParallelState(mpu), + config=rt_cfg, + _extras={ + "bridge": bridge, + "model_list": model_list, + "mpu": mpu, + "model_cfg": bridge.hf_config, + "protocol": _resolve_mbridge_benchmark_protocol(rt_cfg, bridge), + "optimizer_backend": "distopt" if optimizer is not None else "none", + "world_size": dist.get_world_size(), + }, + ) + + def export_weights(self, handle: ModelHandle, **kwargs) -> Iterator[tuple[str, torch.Tensor]]: + bridge = handle._extras["bridge"] + model_list = handle._extras["model_list"] + load_model_to_gpu(model_list, load_grad=False) + return bridge.export_weights(model_list) + + +__all__ = ["MBridgeRuntime"] diff --git a/experimental/lite/megatron/lite/runtime/backends/mlite/__init__.py b/experimental/lite/megatron/lite/runtime/backends/mlite/__init__.py new file mode 100644 index 00000000000..cc514f2858e --- /dev/null +++ b/experimental/lite/megatron/lite/runtime/backends/mlite/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Megatron Lite backend — Megatron Lite's own training engine.""" + +from __future__ import annotations + +from typing import Any + +from megatron.lite.runtime.backends import Runtime as RuntimeBase +from megatron.lite.runtime.backends.mlite.config import MegatronLiteConfig +from megatron.lite.runtime.backends.mlite.runtime import MegatronLiteRuntime + + +def create(hf_path: str, cfg: MegatronLiteConfig | dict[str, Any]) -> RuntimeBase: + """Factory called by create_runtime. + + create_runtime escape hatch is checked inside MegatronLiteRuntime.build_model(). + """ + return MegatronLiteRuntime(hf_path, cfg) diff --git a/experimental/lite/megatron/lite/runtime/backends/mlite/config.py b/experimental/lite/megatron/lite/runtime/backends/mlite/config.py new file mode 100644 index 00000000000..edff0336517 --- /dev/null +++ b/experimental/lite/megatron/lite/runtime/backends/mlite/config.py @@ -0,0 +1,96 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Megatron Lite backend configuration.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from megatron.lite.runtime.contracts.config import OptimizerConfig, ParallelConfig, pick_fields + + +@dataclass(slots=True) +class DebugConfig: + """Megatron Lite backend debug flags. Not exposed to end users.""" + + param_update: bool = False + optimizer_state: bool = False + grad_phases: bool = False + router_summary: bool = False + moe_io: bool = False + attn_io: bool = False + + +@dataclass +class MegatronLiteConfig: + """Config for MegatronLiteRuntime (Megatron Lite's default 5D parallel runtime). + + Megatron Lite-specific training features live in ``impl_cfg`` (a plain dict — + each impl reads the keys it needs via its own typed ImplConfig). + """ + + # ── identity ── + model_name: str = "auto" + impl: str = "lite" + hf_path: str = "" + + # ── parallelism and optimizer ── + parallel: ParallelConfig = field(default_factory=ParallelConfig) + optimizer: OptimizerConfig = field(default_factory=OptimizerConfig) + + # ── common runtime/model fields ── + attention_backend_override: str | None = "flash" + router_aux_loss_coef: float | None = None + load_hf_weights: bool = True + + # ── impl-specific (each impl reads its own keys) ── + impl_cfg: dict[str, Any] = field(default_factory=dict) + + # ── debug ── + debug: DebugConfig = field(default_factory=DebugConfig) + + # ── bench-only hook: mutate model_cfg after build (e.g. expert truncation) ── + model_config_hook: Any = None + + @classmethod + def from_dict(cls, hf_path: str, cfg: dict[str, Any]) -> MegatronLiteConfig: + """Construct MegatronLiteConfig from a flat dict (legacy / OmegaConf path).""" + if "num_microbatches" in cfg: + raise ValueError( + "MegatronLiteConfig no longer accepts `num_microbatches`; " + "pass it to Runtime.forward_backward(..., num_microbatches=...) instead" + ) + parallel = ParallelConfig(**pick_fields(ParallelConfig, cfg)) + + opt_d = cfg.get("optimizer", {}) + optimizer = ( + OptimizerConfig(**pick_fields(OptimizerConfig, opt_d)) + if isinstance(opt_d, dict) + else OptimizerConfig() + ) + if isinstance(opt_d, dict) and isinstance(opt_d.get("override_optimizer_config"), dict): + optimizer.override_optimizer_config = dict(opt_d["override_optimizer_config"]) + + # impl_cfg: merge nested dict + top-level overrides + impl_cfg: dict[str, Any] = {} + nested = cfg.get("impl_cfg") + if isinstance(nested, dict): + impl_cfg.update(nested) + for k in list(impl_cfg): + if k in cfg: + impl_cfg[k] = cfg[k] + for k in ("recompute", "use_thd", "use_deepep", "precision_aware_opt"): + if k in cfg and k not in impl_cfg: + impl_cfg[k] = cfg[k] + + skip = {"parallel", "optimizer", "impl_cfg", "debug"} + return cls( + **{k: v for k, v in pick_fields(cls, cfg).items() if k not in skip}, + hf_path=hf_path, + parallel=parallel, + optimizer=optimizer, + impl_cfg=impl_cfg, + ) + + +__all__ = ["MegatronLiteConfig", "DebugConfig"] diff --git a/experimental/lite/megatron/lite/runtime/backends/mlite/runtime.py b/experimental/lite/megatron/lite/runtime/backends/mlite/runtime.py new file mode 100644 index 00000000000..c197447b70d --- /dev/null +++ b/experimental/lite/megatron/lite/runtime/backends/mlite/runtime.py @@ -0,0 +1,546 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""MegatronLiteRuntime — Megatron Lite's default training backend.""" + +from __future__ import annotations + +import os +from collections.abc import Callable, Iterator +from dataclasses import fields as dc_fields +from datetime import timedelta +from itertools import chain +from typing import Any + +import torch +import torch.distributed as dist + +from megatron.lite.runtime.backends import Runtime as RuntimeBase +from megatron.lite.runtime.backends.mlite.config import MegatronLiteConfig +from megatron.lite.runtime.contracts.data import ForwardResult, ModelOutputs +from megatron.lite.runtime.contracts.handle import ModelHandle + + +def _build_impl_cfg(proto, rt_cfg: MegatronLiteConfig): + """Construct typed impl config, backfilling hf_path + optimizer_config.""" + impl_cfg_kwargs = {**rt_cfg.impl_cfg, "parallel": rt_cfg.parallel} + init_fields = {f.name for f in dc_fields(proto.ImplConfig) if f.init} + if ( + "attention_backend_override" in init_fields + and impl_cfg_kwargs.get("attention_backend_override") is None + ): + impl_cfg_kwargs["attention_backend_override"] = rt_cfg.attention_backend_override + if "hf_path" in init_fields and impl_cfg_kwargs.get("hf_path") in (None, "") and rt_cfg.hf_path: + impl_cfg_kwargs["hf_path"] = rt_cfg.hf_path + # Thread the user-level OptimizerConfig so the protocol can pass it to + # optimizer primitives without reading runtime internals. + if ( + "optimizer_config" in init_fields + and impl_cfg_kwargs.get("optimizer_config") is None + and getattr(rt_cfg, "optimizer", None) is not None + ): + impl_cfg_kwargs["optimizer_config"] = rt_cfg.optimizer + return proto.ImplConfig(**impl_cfg_kwargs) + + +def _apply_attention_backend_env(backend: str | None, *, tag: str) -> None: + if backend is None: + return + + env_overrides = { + "auto": ("1", "1", "1"), + "flash": ("1", "0", "0"), + "fused": ("0", "1", "0"), + "unfused": ("0", "0", "1"), + "local": ("0", "0", "0"), + } + try: + flash, fused, unfused = env_overrides[backend] + except KeyError as exc: + raise ValueError( + "attention_backend_override must be one of {'auto', 'flash', 'fused', 'unfused', 'local'}" + ) from exc + + os.environ["NVTE_FLASH_ATTN"] = flash + os.environ["NVTE_FUSED_ATTN"] = fused + os.environ["NVTE_UNFUSED_ATTN"] = unfused + + +def _infer_pipeline_tensor_shape(batch: Any, model_cfg: Any, ps) -> tuple[int, int, int]: + if model_cfg is None or not hasattr(model_cfg, "hidden_size"): + raise ValueError("Megatron Lite pipeline runtime requires model_cfg.hidden_size.") + if not isinstance(batch, dict) or "input_ids" not in batch: + raise TypeError("Megatron Lite pipeline runtime requires dict batches with input_ids.") + + input_ids = batch["input_ids"] + if input_ids.dim() == 1: + batch_size = 1 + local_seq_len = int(input_ids.size(0)) + elif input_ids.dim() == 2: + batch_size = int(input_ids.size(0)) + local_seq_len = int(input_ids.size(1)) + else: + raise ValueError(f"Unsupported input_ids rank for pipeline runtime: {input_ids.dim()}.") + + if local_seq_len < 1: + raise ValueError("Pipeline tensor shape requires non-empty sequence.") + + tp_size = int(getattr(ps, "tp_size", 1) or 1) + if tp_size > 1: + if local_seq_len % tp_size != 0: + raise ValueError( + f"Pipeline tensor sequence length {local_seq_len} is not divisible by TP={tp_size}." + ) + # Megatron Lite Qwen3.5 scatters embeddings into Megatron sequence-parallel form + # before the first layer, so PP activations carry S / (CP * TP). + local_seq_len //= tp_size + + return (local_seq_len, batch_size, int(model_cfg.hidden_size)) + + +def _last_loss_output(outputs: list[dict]) -> dict: + for output in reversed(outputs): + if output.get("loss") is not None: + return output + return {} + + +def _checkpoint_module(model: Any) -> torch.nn.Module: + if isinstance(model, torch.nn.Module): + return model + if isinstance(model, list | tuple): + return torch.nn.ModuleList(model) + raise TypeError( + f"Checkpoint model must be an nn.Module or sequence of modules, got {type(model).__name__}." + ) + + +class MegatronLiteRuntime(RuntimeBase): + """Megatron Lite default training backend (Megatron-style 5D parallel).""" + + def __init__(self, hf_path: str, cfg: MegatronLiteConfig | dict[str, Any]): + self._hf_path = hf_path + self._cfg = ( + cfg + if isinstance(cfg, MegatronLiteConfig) + else MegatronLiteConfig.from_dict(hf_path, cfg) + ) + + # ── build_model ── + + def build_model( + self, + hf_path: str | None = None, + cfg: MegatronLiteConfig | dict[str, Any] | None = None, + **kwargs, + ) -> ModelHandle: + if cfg is not None and isinstance(cfg, dict): + rt_cfg = MegatronLiteConfig.from_dict(hf_path or self._hf_path, cfg) + elif cfg is not None and isinstance(cfg, MegatronLiteConfig): + rt_cfg = cfg + else: + rt_cfg = self._cfg + + # ── init distributed ── + if not dist.is_initialized(): + dist.init_process_group("nccl", timeout=timedelta(minutes=10)) + torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count()) + torch.cuda.manual_seed(42) + + # ── load model protocol module ── + proto = self._load_protocol(rt_cfg) + + _apply_attention_backend_env( + rt_cfg.attention_backend_override, tag=f"{rt_cfg.model_name}:{rt_cfg.impl}" + ) + + # ── escape hatch: model takes over ── + if hasattr(proto, "create_runtime"): + return proto.create_runtime(rt_cfg.hf_path, rt_cfg).build_model() + + # ── construct impl_cfg (parallel injected) ── + impl_cfg = _build_impl_cfg(proto, rt_cfg) + + # ── build model config ── + model_cfg = proto.build_model_config(rt_cfg.hf_path) + if callable(rt_cfg.model_config_hook): + model_cfg = rt_cfg.model_config_hook(model_cfg) + + # ── build model (model owns ps + optimizer + everything) ── + bundle = proto.build_model(model_cfg, impl_cfg=impl_cfg) + + # ── load HF weights (optional) ── + loaded_hf_weights = False + if rt_cfg.load_hf_weights and rt_cfg.hf_path and hasattr(proto, "load_hf_weights"): + for chunk in bundle.chunks: + proto.load_hf_weights(chunk, rt_cfg.hf_path, model_cfg, bundle.parallel_state) + loaded_hf_weights = True + + post_load_hook = bundle.extras.pop("post_model_load_hook", None) + if callable(post_load_hook): + post_load_updates = post_load_hook() + if post_load_updates is not None: + if not isinstance(post_load_updates, dict): + raise TypeError("post_model_load_hook must return a dict or None.") + if "optimizer" in post_load_updates: + bundle.optimizer = post_load_updates["optimizer"] + if "finalize_grads" in post_load_updates: + bundle.finalize_grads = post_load_updates["finalize_grads"] + extra_updates = post_load_updates.get("extras") + if extra_updates: + if not isinstance(extra_updates, dict): + raise TypeError("post_model_load_hook extras update must be a dict.") + bundle.extras.update(extra_updates) + + if loaded_hf_weights and bundle.optimizer is not None: + reload_model_params = getattr(bundle.optimizer, "reload_model_params", None) + if callable(reload_model_params): + reload_model_params() + + # ── forward_step default ── + forward_fn = bundle.forward_step or (lambda m, b: m(**b)) + + p = rt_cfg.parallel + model = bundle.chunks[0] if len(bundle.chunks) == 1 else bundle.chunks + return ModelHandle( + model=model, + optimizer=bundle.optimizer, + lr_scheduler=None, + parallel_state=bundle.parallel_state, + config=rt_cfg, + _extras={ + "model_chunks": bundle.chunks, + "model_cfg": model_cfg, + "forward_step": forward_fn, + "protocol": proto, + "finalize_grads": bundle.finalize_grads, + "world_size": dist.get_world_size(), + "cp_range": (p.cp, p.cp), + **bundle.extras, + }, + ) + + def _load_protocol(self, rt_cfg: MegatronLiteConfig): + """Load and return the model protocol module.""" + from megatron.lite.model.registry import TRAIN_RUNTIME_MODULES, resolve_runtime_model_name + + try: + runtime_key = resolve_runtime_model_name(rt_cfg.model_name, rt_cfg.impl) + except ValueError as exc: + raise ValueError( + f"No protocol registered for model={rt_cfg.model_name!r}, " + f"impl={rt_cfg.impl!r}. Register with register_model(...)." + ) from exc + + mod_path = TRAIN_RUNTIME_MODULES.get(runtime_key) + if mod_path is None: + raise ValueError(f"No protocol module for runtime key {runtime_key!r}") + + import importlib + + proto = importlib.import_module(mod_path) + + for fn_name in ("build_model_config", "build_model"): + if not callable(getattr(proto, fn_name, None)): + raise ValueError(f"Protocol module {mod_path} missing required function: {fn_name}") + if not hasattr(proto, "ImplConfig"): + raise ValueError(f"Protocol module {mod_path} missing ImplConfig class") + + return proto + + # ── Checkpoint ── + + def save_checkpoint(self, handle: ModelHandle, path: str, **kwargs) -> None: + from megatron.lite.primitive.ckpt import save_training_checkpoint + + step = kwargs.pop("step", None) + if step is None: + step = kwargs.pop("iteration", None) + if step is None: + step = kwargs.pop("global_step", 0) + use_dcp = bool(kwargs.pop("use_dcp", True)) + save_rng = bool(kwargs.pop("save_rng", True)) + get_placements, is_expert = _checkpoint_hooks(handle) + save_training_checkpoint( + _checkpoint_model(handle, use_dcp=use_dcp), + handle._optimizer, + int(step), + path, + _checkpoint_parallel_config(handle), + handle._parallel_state, + get_placements=kwargs.pop("get_placements", get_placements), + is_expert=kwargs.pop("is_expert", is_expert), + use_dcp=use_dcp, + save_rng=save_rng, + save_model=kwargs.pop("save_model", True), + save_optimizer=kwargs.pop("save_optimizer", True), + **kwargs, + ) + + def load_checkpoint(self, handle: ModelHandle, path: str, **kwargs) -> int: + from megatron.lite.primitive.ckpt import load_training_checkpoint + + use_dcp = bool(kwargs.pop("use_dcp", True)) + load_rng = bool(kwargs.pop("load_rng", True)) + update_legacy_format = bool( + kwargs.pop( + "load_parameter_state_update_legacy_format", + kwargs.pop("update_legacy_format", False), + ) + ) + get_placements, is_expert = _checkpoint_hooks(handle) + return load_training_checkpoint( + _checkpoint_model(handle, use_dcp=use_dcp), + handle._optimizer, + path, + _checkpoint_parallel_config(handle), + handle._parallel_state, + get_placements=kwargs.pop("get_placements", get_placements), + is_expert=kwargs.pop("is_expert", is_expert), + use_dcp=use_dcp, + load_rng=load_rng, + load_parameter_state_update_legacy_format=update_legacy_format, + load_model=kwargs.pop("load_model", True), + load_optimizer=kwargs.pop("load_optimizer", True), + **kwargs, + ) + + def export_weights(self, handle: ModelHandle, **kwargs) -> Iterator[tuple[str, torch.Tensor]]: + model_chunks = handle._extras.get("model_chunks", [handle._model]) + proto = handle._extras.get("protocol") + model_cfg = handle._extras.get("model_cfg") + ps = handle._parallel_state + + if proto and hasattr(proto, "export_hf_weights"): + yield from proto.export_hf_weights(model_chunks, model_cfg, ps, **kwargs) + else: + for chunk in model_chunks: + yield from chunk.named_parameters() + + # ── Memory ── + + def to( + self, + handle: ModelHandle, + device: str, + *, + model: bool = True, + optimizer: bool = True, + grad: bool = True, + ) -> None: + model_chunks = handle._extras.get("model_chunks", [handle._model]) + from megatron.lite.runtime.megatron_utils import ( + load_model_to_gpu, + load_optimizer, + offload_model_to_cpu, + offload_optimizer, + ) + + if device == "cpu": + if model: + offload_model_to_cpu(model_chunks) + if optimizer and handle._optimizer is not None: + offload_state = getattr(handle._optimizer, "offload_state_to_cpu", None) + if callable(offload_state): + offload_state() + else: + offload_optimizer(handle._optimizer) + elif device == "cuda": + if model: + load_model_to_gpu(model_chunks, load_grad=grad) + if optimizer and handle._optimizer is not None: + load_state = getattr(handle._optimizer, "load_state_to_device", None) + if callable(load_state): + load_state() + else: + load_optimizer(handle._optimizer) + + # ── Mode switching ── + + def train_mode(self, handle: ModelHandle): + return _TrainModeCtx(handle) + + def eval_mode(self, handle: ModelHandle): + return _EvalModeCtx(handle) + + # ── Training atoms ── + + def forward_backward( + self, + handle: ModelHandle, + data: Any, + loss_fn: Callable | None, + *, + num_microbatches: int = 1, + forward_only: bool = False, + ) -> ForwardResult: + from megatron.lite.primitive.train_step import run_microbatch_loop + + forward_step = handle._extras["forward_step"] + if num_microbatches < 1: + raise ValueError("num_microbatches must be >= 1") + + if hasattr(data, "__next__"): + data_iter = data + elif hasattr(data, "__iter__"): + data_iter = iter(data) + else: + data_iter = iter([data]) + + ps = handle._parallel_state + if ps.pp_size > 1: + from types import SimpleNamespace + + from megatron.lite.primitive.parallel.pipeline import forward_backward_pipelining + + first_batch = next(data_iter) + data_iter = chain([first_batch], data_iter) + tensor_shape = _infer_pipeline_tensor_shape( + first_batch, handle._extras.get("model_cfg"), ps + ) + outputs = forward_backward_pipelining( + forward_step, + handle._extras.get("model_chunks", [handle._model]), + data_iter, + SimpleNamespace(num_microbatches=num_microbatches), + ps, + tensor_shape=tensor_shape, + pre_forward_hook=handle._extras.get("pre_forward_hook"), + loss_fn=loss_fn, + forward_only=forward_only, + ) + out = _last_loss_output(outputs) + loss_obj = out.get("loss") if out else None + if isinstance(loss_obj, torch.Tensor): + loss_float = float(loss_obj.detach().item()) + elif loss_obj is not None: + loss_float = float(loss_obj) + else: + loss_float = 0.0 + loss_t = torch.tensor([loss_float], device="cuda") + if ps.pp_group is not None and ps.pp_global_ranks is not None: + dist.broadcast(loss_t, src=ps.pp_global_ranks[-1], group=ps.pp_group) + out = {"loss": loss_t.squeeze(0)} + else: + out = run_microbatch_loop( + handle._model, + data_iter, + num_microbatches, + forward_step, + optimizer=handle._optimizer if not forward_only else None, + dist_opt=not forward_only, + pre_forward_hook=handle._extras.get("pre_forward_hook"), + loss_fn=loss_fn, + ) + + if not forward_only: + finalize_grads = handle._extras.get("finalize_grads") + if finalize_grads is not None: + finalize_grads() + + loss_tensor = out.get("loss") if out else None + loss_val = ( + loss_tensor.item() + if isinstance(loss_tensor, torch.Tensor) + else float(loss_tensor or 0.0) + ) + metrics: dict = {"loss": loss_val} + for m in out.get("_loss_fn_metrics", []) if out else []: + for k, v in m.items(): + if k not in metrics: + metrics[k] = v + if ps.pp_size > 1: + for item in outputs: + for k, v in item.get("metrics", {}).items(): + if k not in metrics: + metrics[k] = v + metrics["_micro_outputs"] = outputs + + return ForwardResult( + model_output=ModelOutputs( + loss=loss_tensor, + vocab_parallel_logits=out.get("logits") if out else None, + log_probs=out.get("log_probs") if out else None, + routed_experts=out.get("routed_experts") if out else None, + ), + metrics=metrics, + ) + + def is_mp_src_rank_with_outputs(self, handle: ModelHandle) -> bool: + ps = handle._parallel_state + return ps.tp_rank == 0 and ps.cp_rank == 0 and ps.pp_rank == ps.pp_size - 1 + + def zero_grad(self, handle: ModelHandle) -> None: + for chunk in handle._extras.get("model_chunks", [handle._model]): + if hasattr(chunk, "zero_grad_buffer"): + chunk.zero_grad_buffer() + if handle._optimizer is not None: + handle._optimizer.zero_grad() + + def optimizer_step(self, handle: ModelHandle) -> tuple[bool, float, int | None]: + if handle._optimizer is None: + return True, 0.0, 0 + update_successful, grad_norm, num_zeros = handle._optimizer.step() + return update_successful, float(grad_norm), num_zeros + + def lr_scheduler_step(self, handle: ModelHandle) -> float | list[float]: + if handle._lr_scheduler is not None: + handle._lr_scheduler.step() + return handle._lr_scheduler.get_last_lr() + return 0.0 + + +# --------------------------------------------------------------------------- +# Context managers +# --------------------------------------------------------------------------- + + +class _TrainModeCtx: + def __init__(self, handle: ModelHandle): + self._handle = handle + + def __enter__(self): + for chunk in self._handle._extras.get("model_chunks", [self._handle._model]): + chunk.train() + return self + + def __exit__(self, *exc): + return False + + +class _EvalModeCtx: + def __init__(self, handle: ModelHandle): + self._handle = handle + self._prev_grad = torch.is_grad_enabled() + + def __enter__(self): + for chunk in self._handle._extras.get("model_chunks", [self._handle._model]): + chunk.eval() + torch.set_grad_enabled(False) + return self + + def __exit__(self, *exc): + torch.set_grad_enabled(self._prev_grad) + return False + + +def _checkpoint_parallel_config(handle: ModelHandle): + cfg = handle.config + if cfg is None: + return None + return getattr(cfg, "parallel", cfg) + + +def _checkpoint_model(handle: ModelHandle, *, use_dcp: bool): + model = handle._model + if not use_dcp or isinstance(model, torch.nn.Module): + return model + return torch.nn.ModuleList(handle._extras.get("model_chunks", model)) + + +def _checkpoint_hooks(handle: ModelHandle): + from megatron.lite.primitive.protocols import default_expert_classifier, default_placement_fn + + proto = handle._extras.get("protocol") + return ( + getattr(proto, "PLACEMENT_FN", default_placement_fn), + getattr(proto, "EXPERT_CLASSIFIER", default_expert_classifier), + ) diff --git a/experimental/lite/megatron/lite/runtime/contracts/__init__.py b/experimental/lite/megatron/lite/runtime/contracts/__init__.py new file mode 100644 index 00000000000..e8485baca29 --- /dev/null +++ b/experimental/lite/megatron/lite/runtime/contracts/__init__.py @@ -0,0 +1,65 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Lazy re-export hub for ``megatron.lite.runtime.contracts``. + +This preserves imports like ``from megatron.lite.runtime.contracts import X`` while +avoiding eager imports of heavyweight modules (for example ``torch`` from +``contracts.data``) during package import. +""" + +from __future__ import annotations + +import importlib +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from megatron.lite.runtime.backends.bridge.config import BridgeConfig + from megatron.lite.runtime.backends.mlite.config import DebugConfig, MegatronLiteConfig + from megatron.lite.runtime.contracts.config import ( + OptimizerConfig, + ParallelConfig, + RuntimeConfig, + ) + from megatron.lite.runtime.contracts.data import ( + Batch, + ForwardResult, + ModelOutputs, + PackedBatch, + TrainBatch, + ) + from megatron.lite.runtime.contracts.handle import ModelHandle + +__all__ = [ + "MegatronLiteConfig", + "Batch", + "BridgeConfig", + "DebugConfig", + "ForwardResult", + "ModelHandle", + "ModelOutputs", + "OptimizerConfig", + "PackedBatch", + "ParallelConfig", + "RuntimeConfig", + "TrainBatch", +] + + +def __getattr__(name: str): + _lazy = { + "Batch": "megatron.lite.runtime.contracts.data", + "BridgeConfig": "megatron.lite.runtime.backends.bridge.config", + "DebugConfig": "megatron.lite.runtime.backends.mlite.config", + "MegatronLiteConfig": "megatron.lite.runtime.backends.mlite.config", + "ForwardResult": "megatron.lite.runtime.contracts.data", + "ModelHandle": "megatron.lite.runtime.contracts.handle", + "ModelOutputs": "megatron.lite.runtime.contracts.data", + "OptimizerConfig": "megatron.lite.runtime.contracts.config", + "PackedBatch": "megatron.lite.runtime.contracts.data", + "ParallelConfig": "megatron.lite.runtime.contracts.config", + "RuntimeConfig": "megatron.lite.runtime.contracts.config", + "TrainBatch": "megatron.lite.runtime.contracts.data", + } + if name in _lazy: + mod = importlib.import_module(_lazy[name]) + return getattr(mod, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/experimental/lite/megatron/lite/runtime/contracts/config.py b/experimental/lite/megatron/lite/runtime/contracts/config.py new file mode 100644 index 00000000000..7ef5d29108e --- /dev/null +++ b/experimental/lite/megatron/lite/runtime/contracts/config.py @@ -0,0 +1,84 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Shared runtime configuration for Megatron Lite.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from dataclasses import fields as dc_fields +from typing import TYPE_CHECKING, Any + + +def pick_fields(cls, src: dict[str, Any]) -> dict[str, Any]: + """Extract fields of dataclass *cls* that exist in *src*.""" + return {f.name: src[f.name] for f in dc_fields(cls) if f.name in src} + + +if TYPE_CHECKING: + from megatron.lite.runtime.backends.bridge.config import BridgeConfig + from megatron.lite.runtime.backends.mlite.config import MegatronLiteConfig + + +@dataclass +class ParallelConfig: + """Parallel dimensions used by Megatron Lite.""" + + tp: int = 1 + etp: int | None = None + ep: int = 1 + pp: int = 1 + vpp: int = 1 + cp: int = 1 + + +@dataclass +class OptimizerConfig: + """Optimizer + LR scheduler config. Aligned with VERL McoreOptimizerConfig. + + Stable VERL fields use VERL default values. + Compatibility aliases are lowered into backend-specific override dicts + by the adapter layer before consumption. + """ + + # --- stable VERL fields --- + optimizer: str = "adam" + lr: float = 1e-3 + min_lr: float = 0.0 + clip_grad: float = 1.0 + weight_decay: float = 0.01 + lr_warmup_steps_ratio: float = 0.0 + total_training_steps: int = -1 + lr_warmup_steps: int = -1 + lr_warmup_init: float = 0.0 + lr_decay_steps: int | None = None + lr_decay_style: str = "linear" + weight_decay_incr_style: str = "constant" + lr_wsd_decay_style: str = "exponential" + lr_wsd_decay_steps: int | None = None + use_checkpoint_opt_param_scheduler: bool = False + + # --- compatibility aliases --- + adam_beta1: float | None = None + adam_beta2: float | None = None + adam_eps: float | None = None + offload_fraction: float | None = None + use_precision_aware_optimizer: bool | None = None + decoupled_weight_decay: bool | None = None + + +@dataclass +class RuntimeConfig: + """Top-level runtime configuration. + + Attributes: + backend: Runtime backend name. Use ``"mlite"`` for Megatron Lite or + ``"bridge"`` for Megatron-Bridge. + hf_path: Path to HuggingFace model directory. Required for real runs. + backend_cfg: Backend config or a compatible dict. + """ + + backend: str = "mlite" + hf_path: str = "" + backend_cfg: MegatronLiteConfig | BridgeConfig | dict[str, Any] = field(default_factory=dict) + + +__all__ = ["OptimizerConfig", "ParallelConfig", "RuntimeConfig"] diff --git a/experimental/lite/megatron/lite/runtime/contracts/data.py b/experimental/lite/megatron/lite/runtime/contracts/data.py new file mode 100644 index 00000000000..ff028a4fa20 --- /dev/null +++ b/experimental/lite/megatron/lite/runtime/contracts/data.py @@ -0,0 +1,123 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Data contracts — forward_backward input/output types.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch + + +class Batch: + """Protocol for data passed to Runtime.forward_backward. + + Sequences are packed without padding. ``sizes()`` gives per-sequence + lengths so the runtime can reconstruct ``cu_seqlens`` / ``position_ids`` + for THD attention. + + Subclass this to carry domain-specific fields (loss_mask, routed_experts, + etc.) while keeping the runtime interface uniform. + """ + + def __len__(self) -> int: + """Number of sequences in this batch.""" + raise NotImplementedError + + def sizes(self) -> torch.Tensor: + """Per-sequence token counts. Shape ``[num_seqs]``.""" + raise NotImplementedError + + def __getitem__(self, key: str) -> Any: + """Dict-like field access (``batch["input_ids"]``).""" + raise NotImplementedError + + +@dataclass(slots=True) +class PackedBatch(Batch): + """Variable-length packed batch — no padding. + + All token-level tensors are 1-D with length ``sum(seq_lens)``. + ``cu_seqlens`` and ``position_ids`` are derived automatically. + """ + + input_ids: torch.Tensor # [total_tokens] + labels: torch.Tensor # [total_tokens] + seq_lens: torch.Tensor # [num_seqs] + loss_mask: torch.Tensor | None = None # [total_tokens] + position_ids: torch.Tensor | None = None # [total_tokens], auto if None + routed_experts: torch.Tensor | None = None + extras: dict[str, Any] = field(default_factory=dict) + + def __len__(self) -> int: + return len(self.seq_lens) + + def sizes(self) -> torch.Tensor: + return self.seq_lens + + def __getitem__(self, key: str) -> Any: + if key in self.__slots__: + return getattr(self, key) + return self.extras[key] + + @property + def cu_seqlens(self) -> torch.Tensor: + """Cumulative sequence lengths for THD attention. Shape ``[num_seqs+1]``.""" + return torch.cat( + [ + torch.zeros(1, dtype=torch.int32, device=self.seq_lens.device), + self.seq_lens.cumsum(0).to(torch.int32), + ] + ) + + @property + def total_tokens(self) -> int: + return int(self.seq_lens.sum()) + + def make_position_ids(self) -> torch.Tensor: + """Generate per-token position_ids from seq_lens.""" + if self.position_ids is not None: + return self.position_ids + return torch.cat( + [torch.arange(s, device=self.seq_lens.device) for s in self.seq_lens.tolist()] + ) + + +@dataclass(slots=True) +class TrainBatch: + """Legacy fixed-shape batch (padded). Use PackedBatch for new code.""" + + input_ids: torch.Tensor + labels: torch.Tensor + loss_mask: torch.Tensor | None = None + position_ids: torch.Tensor | None = None + routed_experts: torch.Tensor | None = None + cp_size: int | None = None + extras: dict[str, Any] = field(default_factory=dict) + + +@dataclass(slots=True) +class ModelOutputs: + """Model forward output.""" + + loss: torch.Tensor | None = None + vocab_parallel_logits: torch.Tensor | None = None + log_probs: torch.Tensor | None = None + hidden_states: torch.Tensor | None = None + values: torch.Tensor | None = None + # MTP + mtp_logits: torch.Tensor | None = None + mtp_loss: torch.Tensor | None = None + # Router Replay: recorded routing decisions + routed_experts: torch.Tensor | None = None + + +@dataclass(slots=True) +class ForwardResult: + """Output of forward_backward.""" + + model_output: ModelOutputs = field(default_factory=ModelOutputs) + metrics: dict[str, Any] = field(default_factory=dict) + + +__all__ = ["Batch", "ForwardResult", "ModelOutputs", "PackedBatch", "TrainBatch"] diff --git a/experimental/lite/megatron/lite/runtime/contracts/handle.py b/experimental/lite/megatron/lite/runtime/contracts/handle.py new file mode 100644 index 00000000000..85c4f120037 --- /dev/null +++ b/experimental/lite/megatron/lite/runtime/contracts/handle.py @@ -0,0 +1,65 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""ModelHandle — opaque handle returned by Runtime.build_model().""" + +from __future__ import annotations + +from typing import Any + + +class ModelHandle: + """Opaque handle returned by Runtime.build_model(). + + Only documented properties on this class are part of the public contract. + Internal helpers may still use ``_model`` / ``_optimizer`` / ``_extras`` + while the higher-level runtime API is being stabilized. + """ + + def __init__( + self, + *, + model: Any, + optimizer: Any = None, + lr_scheduler: Any = None, + parallel_state: Any = None, + config: Any = None, + _extras: dict[str, Any] | None = None, + ): + self._model = model + self._optimizer = optimizer + self._lr_scheduler = lr_scheduler + self._parallel_state = parallel_state + self._config = config + self._extras = _extras or {} + + @property + def dp_rank(self) -> int: + ps = self._parallel_state + if ps is None: + return 0 + return getattr(ps, "dp_rank", 0) + + @property + def dp_size(self) -> int: + ps = self._parallel_state + if ps is None: + return 1 + return getattr(ps, "dp_size", 1) + + @property + def dp_group(self): + ps = self._parallel_state + if ps is None: + return None + return getattr(ps, "dp_group", None) + + @property + def cp_range(self) -> tuple[int, int]: + return self._extras.get("cp_range", (1, 1)) + + @property + def config(self) -> Any: + """Backend config captured when this handle was built.""" + return self._config + + +__all__ = ["ModelHandle"] diff --git a/experimental/lite/megatron/lite/runtime/megatron_utils.py b/experimental/lite/megatron/lite/runtime/megatron_utils.py new file mode 100644 index 00000000000..129868f08a5 --- /dev/null +++ b/experimental/lite/megatron/lite/runtime/megatron_utils.py @@ -0,0 +1,222 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Megatron-Core utilities aligned with VERL's megatron_utils. + +All functions here are ported from VERL and could theoretically be replaced +by ``from verl.utils.megatron_utils import ...`` if Megatron Lite ever depends on verl. + +Sources: + - verl/utils/megatron_utils.py (offload/load, register_megatron_training_hooks) + - verl/utils/megatron/optimizer.py (get_megatron_last_lr) +""" + +from __future__ import annotations + +import gc +from typing import Any + +import torch + +# ====================================================================== +# Rank utilities +# ====================================================================== + + +def is_mp_src_rank_with_outputs() -> bool: + """True on the rank that holds the final model output (loss). + + Only last PP stage, first TP rank, first CP rank has the output. + VERL: MegatronEngine.is_mp_src_rank_with_outputs + """ + from megatron.core import parallel_state as mpu + + return ( + mpu.get_tensor_model_parallel_rank() == 0 + and mpu.get_pipeline_model_parallel_rank() + == mpu.get_pipeline_model_parallel_world_size() - 1 + and mpu.get_context_parallel_rank() == 0 + ) + + +# ====================================================================== +# Training hooks — register_megatron_training_hooks +# ====================================================================== + + +def register_training_hooks(model_list: list, optimizer) -> None: + """Register megatron training callbacks on model config. + + Ref: megatron/training/training.py (core_v0.15.0rc7, L2039-L2057) + """ + from megatron.core.distributed import DistributedDataParallel as DDP + from megatron.core.distributed import finalize_model_grads + from megatron.core.utils import get_model_config + + for one_model in model_list: + config = get_model_config(one_model) + if optimizer is not None: + config.grad_scale_func = optimizer.scale_loss + config.finalize_model_grads_func = finalize_model_grads + + optimizer_config = getattr(optimizer, "config", None) + overlap_param_gather = getattr(optimizer_config, "overlap_param_gather", False) + overlap_grad_reduce = getattr(one_model.ddp_config, "overlap_grad_reduce", False) + align_grad_reduce = True + align_param_gather = getattr(one_model.ddp_config, "align_param_gather", False) + + if isinstance(model_list[0], DDP) and overlap_grad_reduce: + config.no_sync_func = [m.no_sync for m in model_list] + if len(model_list) == 1: + config.no_sync_func = config.no_sync_func[0] + if align_grad_reduce: + config.grad_sync_func = [m.start_grad_sync for m in model_list] + if len(model_list) == 1: + config.grad_sync_func = config.grad_sync_func[0] + if overlap_param_gather and align_param_gather: + config.param_sync_func = [m.start_param_sync for m in model_list] + if len(model_list) == 1: + config.param_sync_func = config.param_sync_func[0] + + +# ====================================================================== +# Model offload / load — offload_megatron_model_to_cpu / load_megatron_model_to_gpu +# ====================================================================== + + +def offload_model_to_cpu(model_list: list) -> None: + """Offload DDP model to CPU via buffer-resize (zero-copy on GPU side).""" + from megatron.core.distributed import DistributedDataParallel as DDP + + for model_chunk in model_list: + if isinstance(model_chunk, DDP): + all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers] + for buffers in all_buffers: + for buffer in buffers: + if buffer.param_data.storage().size() > 0: + buffer.param_data.cpu_data = buffer.param_data.data.cpu().pin_memory() + buffer.param_data_size = buffer.param_data.storage().size() + buffer.param_data.storage().resize_(0) + + if buffer.grad_data.storage().size() > 0: + buffer.grad_data_size = buffer.grad_data.storage().size() + buffer.grad_data.storage().resize_(0) + + for param in model_chunk.module.parameters(): + if not param.requires_grad and param.device.type != "cpu": + param.data = param.data.to("cpu", non_blocking=True) + else: + model_chunk.to("cpu") + + +def load_model_to_gpu(model_list: list, load_grad: bool = True) -> None: + """Load DDP model back to GPU from pinned CPU copy.""" + from megatron.core.distributed import DistributedDataParallel as DDP + + for model_chunk in model_list: + if isinstance(model_chunk, DDP): + all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers] + for buffers in all_buffers: + for buffer in buffers: + if load_grad and hasattr(buffer, "grad_data_size"): + current_size = buffer.grad_data.storage().size() + if current_size == 0 or current_size == buffer.grad_data_size: + buffer.grad_data.storage().resize_(buffer.grad_data_size) + buffer.grad_data.zero_() + else: + buffer.grad_data.zero_() + + if buffer.param_data.storage().size() == 0: + buffer.param_data.storage().resize_(buffer.param_data_size) + buffer.param_data.copy_(buffer.param_data.cpu_data, non_blocking=True) + + for param in model_chunk.module.parameters(): + if not param.requires_grad and param.device.type == "cpu": + param.data = param.data.to("cuda", non_blocking=True) + else: + model_chunk.to("cuda") + + +# ====================================================================== +# Optimizer offload / load — offload_megatron_optimizer / load_megatron_optimizer +# ====================================================================== + + +def offload_optimizer(optimizer) -> None: + """Offload optimizer states to CPU.""" + from megatron.core.optimizer import ChainedOptimizer + + for _opt in _iter_opts(optimizer, ChainedOptimizer): + if _opt.optimizer is not None: + hdo = _opt.optimizer + if all( + hasattr(hdo, a) for a in ("sub_optimizers", "inner_param_to_orig_param", "state") + ): + for sub_opt in hdo.sub_optimizers: + for param, state in sub_opt.state.items(): + for k, v in state.items(): + if not isinstance(v, torch.Tensor): + continue + orig_param = hdo.inner_param_to_orig_param.get(param, param) + hdo.state[orig_param][k] = state[k] = v.to("cpu") + else: + for v in _opt.optimizer.state.values(): + if "exp_avg" in v: + v["exp_avg"] = v["exp_avg"].to("cpu", non_blocking=True) + if "exp_avg_sq" in v: + v["exp_avg_sq"] = v["exp_avg_sq"].to("cpu", non_blocking=True) + + gc.collect() + torch.cuda.empty_cache() + + +def load_optimizer(optimizer) -> None: + """Load optimizer states back to GPU.""" + from megatron.core.optimizer import ChainedOptimizer + + for _opt in _iter_opts(optimizer, ChainedOptimizer): + if _opt.optimizer is not None: + if hasattr(_opt.optimizer, "_move_new_state_to_right_device"): + _opt.optimizer._move_new_state_to_right_device() + else: + for v in _opt.optimizer.state.values(): + if "exp_avg" in v: + v["exp_avg"] = v["exp_avg"].to("cuda", non_blocking=True) + if "exp_avg_sq" in v: + v["exp_avg_sq"] = v["exp_avg_sq"].to("cuda", non_blocking=True) + + gc.collect() + torch.cuda.empty_cache() + + +def _iter_opts(optimizer, chained_cls): + if isinstance(optimizer, chained_cls): + return optimizer.chained_optimizers + return [optimizer] + + +# ====================================================================== +# Checkpoint helpers +# ====================================================================== + + +def build_sharded_state_dict( + model_list: list, optimizer: Any = None, lr_scheduler: Any = None +) -> dict[str, Any]: + """Build sharded state dict for model + optimizer + lr_scheduler. + + Uses ``model0.`` / ``model1.`` prefix for VPP (multiple model chunks). + """ + sharded_state_dict: dict[str, Any] = {} + + for i, model_chunk in enumerate(model_list): + prefix = f"model{i}." if len(model_list) > 1 else "model." + chunk_sd = model_chunk.sharded_state_dict(prefix=prefix) + sharded_state_dict.update(chunk_sd) + + if optimizer is not None: + opt_sd = optimizer.sharded_state_dict(model_sharded_state_dict=sharded_state_dict) + sharded_state_dict.update(opt_sd) + + if lr_scheduler is not None: + sharded_state_dict["lr_scheduler"] = lr_scheduler.state_dict() + + return sharded_state_dict diff --git a/experimental/lite/skills/README.md b/experimental/lite/skills/README.md new file mode 100644 index 00000000000..51a4468714d --- /dev/null +++ b/experimental/lite/skills/README.md @@ -0,0 +1,167 @@ +# Megatron Lite Skills + +This directory defines agent-agnostic skills for maintaining Megatron Lite. +A skill is an operational contract for agents: what context to load, what +invariants to protect, how to make a change, and what evidence to leave behind. + +These files are not Codex, Claude, or tool-specific skills. Do not add +agent-specific frontmatter, prompt wrappers, or runtime assumptions here. + +## Function Model + +Treat a skill as a function: + +- `Schema` is the signature and summary. +- `imports` are skills that must be loaded before execution, like Python imports. +- `calls` are every skill this function explicitly calls in the body. +- the body is Python-like pseudocode with bounded exits. + +An agent should be able to route work from the `Schema` alone. It should read +the body only when executing that skill. + +## File Layout + +Each skill file has three regions: + +1. Before `MLITE_SKILL_SCHEMA_BEGIN`: a short human-facing title or note. Agents + should not depend on this region for routing. +2. Between `MLITE_SKILL_SCHEMA_BEGIN` and `MLITE_SKILL_SCHEMA_END`: the compact + schema used for retrieval, routing, imports, and call planning. +3. After `MLITE_SKILL_SCHEMA_END`: the skill body. This must be Python-like + pseudocode, not free-form prose. + +The pseudocode is a contract, not executable Python. It should still be precise: +clear inputs, explicit outputs, explicit skill calls, and finite exits. + +Skills are organized like Python modules. A skill is a single `.md` file; +directories are namespaces, not skill bodies. Do not create `skill-name/README.md` +for individual skills. +Map schema names to paths by replacing `.` with `/` and `_` with `-`, then +adding `.md`: `model_compose.config_mapping` lives at +`model-compose/config-mapping.md`. + +## Loading Model + +Agents should load skills progressively: + +1. Read this file. +2. Read exactly one leaf skill for the current work type. +3. Read linked Megatron Lite docs or source files only when the leaf skill asks + for them. + +Keep the active context small. A skill should point to durable source files and +validation commands instead of copying long design background. + +## Defined Skills + +`basic`: +- `basic.constitution` +- `basic.find_reference` +- `basic.align_precision` +- `basic.align_e2e_precision` +- `basic.lint_skill` +- `basic.construct_proxy_task` +- `basic.review_threshold` + +`primitive`: +- `primitive.contract` +- `primitive.principle` +- `primitive.select_for_compose` +- `primitive.design` +- `primitive.validate` +- `primitive.fuse` +- `primitive.process_group` +- `primitive.parallel.tp` +- `primitive.parallel.ep` +- `primitive.parallel.pp` +- `primitive.parallel.cp` +- `primitive.optimizer.fsdp` +- `primitive.optimizer.distopt` +- `primitive.module.moe` +- `primitive.module.gqa` +- `primitive.module.thd` + +Primitive work has two reusable meta layers: `primitive.principle` defines what +must be true, and `primitive.select_for_compose` decides when a primitive belongs +in a model composition. + +`model_compose`: +- `model_compose.config_mapping` +- `model_compose.weight_mapping` +- `model_compose.build_model` +- `model_compose.qwen` + +`application`: +- `application.runtime` +- `application.verl` +- `application.bench` + +`perf`: +- `perf.measure` +- `perf.memory` +- `perf.fusion` +- `perf.optimize` + +`insight`: +- `insight.progressive_model_support` +- `insight.progressive_primitive_design` +- `insight.scope_control` + +## Skill Contract + +Each skill file contains a delimited schema block near the top. Agents may read +only this block for routing, then read the body only when executing the skill. + +````text + +```python +schema = Skill( + "basic.align_precision", kind="state_machine", purpose="align precision", + imports=["basic.constitution"], calls=["basic.constitution", "basic.find_reference"], + inputs=["task", "files", "reference", "budget"], + outputs=["patch", "evidence", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + +```` + +Use the markers instead of fixed line counts. Keep the schema short enough to +scan without reading the body; prefer compact lists and short field names over +prose. + +After the schema, each skill body is Python-like pseudocode. The top-level +function should match the schema inputs and return one of the declared exits. +Use prose only inside short comments. + +Required sections: + +- one top-level function matching the schema name; +- explicit `return done(...)`, `return blocked(...)`, or + `return out_of_scope(...)`; +- explicit calls to skills as `namespace.skill(...)`. + +`imports` and `calls` are separate. `imports` says what an agent should preload. +`calls` says what the body invokes. If a loaded skill is directly invoked, list +it in `calls` too. + +Every loop must have a progress measure, a maximum attempt count, or a blocking +condition. A skill that can spin forever is invalid. + +## Skill Lint + +Every new or changed skill should pass `basic.lint_skill` before review. Lint is +structural: it checks that schema can be extracted, imports and calls resolve, +the file path matches the schema name, body calls match declared calls, exits +are explicit, and loops are bounded. Scenario dry-runs are optional but +preferred for complex skills; they use mocked inputs to verify the pseudocode +can reach a declared exit. + +## Boundaries + +- Skills describe how agents work; `docs/` describes Megatron Lite for users and + reviewers. +- Skills must not replace tests. They should name the right validation surface. +- Skills must not store task history, temporary job logs, or one-off debugging + transcripts. +- Skills should stay stable across agents. If a rule depends on one agent tool, + keep it outside this directory. diff --git a/experimental/lite/skills/application/bench.md b/experimental/lite/skills/application/bench.md new file mode 100644 index 00000000000..d8e5ba5c314 --- /dev/null +++ b/experimental/lite/skills/application/bench.md @@ -0,0 +1,31 @@ +# Bench Skill + +Run benchmark-style checks without weakening precision evidence. + +## Schema + + +```python +schema = Skill( + "application.bench", kind="state_machine", purpose="benchmark with controlled variables and evidence", + imports=["basic.constitution"], calls=["perf.measure"], + inputs=["task", "bench_config", "target", "budget"], + outputs=["bench", "evidence", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def bench(task, bench_config, target, budget): + controlled = bench_config.variables.freeze_all_except(bench_config.axis) + if controlled.has_unknown_unfrozen_axes(): + return blocked("bench has uncontrolled variables", risks=controlled.unknown_axes) + + measurement = perf.measure(task, target=target, workload=bench_config.workload, budget=budget.measure) + if not measurement.done: + return blocked("bench measurement failed", evidence=measurement) + + evidence = record_evidence(task, run=measurement.run, comparison=measurement.metrics, environment=budget.env) + risks = ["benchmarks are not correctness proof", "uncontrolled variables can dominate"] + return done(bench=measurement.metrics, evidence=evidence, risks=risks) +``` diff --git a/experimental/lite/skills/application/runtime.md b/experimental/lite/skills/application/runtime.md new file mode 100644 index 00000000000..f3e956701df --- /dev/null +++ b/experimental/lite/skills/application/runtime.md @@ -0,0 +1,30 @@ +# Runtime Skill + +Validate the MLite runtime path end to end. + +## Schema + + +```python +schema = Skill( + "application.runtime", kind="state_machine", purpose="validate MLite runtime build/train/save/load", + imports=["basic.constitution"], calls=["basic.align_precision"], + inputs=["task", "runtime_config", "model", "budget"], + outputs=["runtime", "evidence", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def runtime(task, runtime_config, model, budget): + if runtime_config.backend != "mlite": + return out_of_scope("not MLite runtime") + + run = execute_runtime_steps(["init", "build_model", "train_step", "save", "load"], runtime_config, model) + if not run.return_code == 0: + return blocked("runtime path failed", evidence=run) + + precision = basic.align_precision(task, target=model, variables=runtime_config.variables, budget=budget.precision) + evidence = record_evidence(task, run=run, comparison=precision, environment=budget.env) + return done(runtime=run, evidence=evidence, risks=["runtime success can hide model-local mismatch"]) +``` diff --git a/experimental/lite/skills/application/verl.md b/experimental/lite/skills/application/verl.md new file mode 100644 index 00000000000..e7257c2f139 --- /dev/null +++ b/experimental/lite/skills/application/verl.md @@ -0,0 +1,35 @@ +# VERL Skill + +Validate MLite through VERL SFT, RL, or GRPO workflows. + +## Schema + + +```python +schema = Skill( + "application.verl", kind="state_machine", purpose="validate VERL workflows with MLite", + imports=["basic.constitution"], calls=["basic.align_e2e_precision"], + inputs=["task", "verl_config", "model", "budget"], + outputs=["workflow", "evidence", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def verl(task, verl_config, model, budget): + if verl_config.algorithm not in ["SFT", "RL", "GRPO"]: + return out_of_scope("not a VERL SFT/RL/GRPO workflow") + + modes = {"SFT": ["SFT"], "RL": ["RL"], "GRPO": ["RL"]}[verl_config.algorithm] + e2e = basic.align_e2e_precision(task, target=model, modes=modes, budget=budget.e2e) + if not e2e.done: + return blocked("VERL e2e precision failed", evidence=e2e) + + evidence = record_evidence(task, run=e2e.evidence.run, comparison=e2e.evidence, environment=budget.env) + risks = [ + "SFT loss curve can pass while a lower-level bug remains shared", + "RL reward variance can hide precision drift", + "rollout backend can become the reference accidentally", + ] + return done(workflow=verl_config.algorithm, evidence=evidence, risks=risks) +``` diff --git a/experimental/lite/skills/basic/align-e2e-precision.md b/experimental/lite/skills/basic/align-e2e-precision.md new file mode 100644 index 00000000000..c05c7414c50 --- /dev/null +++ b/experimental/lite/skills/basic/align-e2e-precision.md @@ -0,0 +1,59 @@ +# Align E2E Precision Skill + +Use pretrain, SFT, or RL as the strongest but highest-cost precision evidence. + +## Schema + + +```python +schema = Skill( + "basic.align_e2e_precision", kind="state_machine", purpose="validate precision through pretrain/SFT/RL", + imports=["basic.constitution"], calls=["basic.constitution", "basic.find_reference", "basic.review_threshold"], + inputs=["task", "target", "modes", "budget"], + outputs=["alignment", "evidence", "override", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def align_e2e_precision(task, target, modes, budget): + if task.scope not in ["precision", "model", "runtime", "application"]: + return out_of_scope("not an end-to-end precision task") + if not budget.high_cost_approved: + return blocked("e2e precision requires high-cost approval") + + mode = choose_first_available(modes, order=["pretrain", "SFT", "RL"]) + if mode is None: + return blocked("no e2e mode selected") + + ref = basic.find_reference(task, layer=target.layer, candidates=task.references, budget=budget.reference) + if not ref.done: + return blocked("no usable e2e reference", evidence=ref) + + policy = basic.constitution(task, layer=target.layer, reference=ref.reference) + if not policy.done: + return blocked("e2e precision policy failed", evidence=policy) + + run = construct_e2e_run( + mode, + target, + reference=ref.reference, + freeze=["dataset", "tokenizer", "checkpoint", "schedule", "seed", "variables"], + compare=["loss_curve", "grad_norm", "reward_or_metric", "checkpoint_delta"], + ) + if run is None: + return blocked("cannot construct comparable e2e run") + + threshold = basic.review_threshold(task, ref.contract, requested_threshold=task.tolerance) + if not threshold.done: + return blocked("e2e precision threshold review failed", evidence=threshold) + compare_mode = threshold.compare_mode + + evidence = run_reference_and_target(run, compare_mode=compare_mode) + risks = ["high cost", "may preserve a shared bug", "can mask local mismatches"] + + if evidence.pass_: + return done(alignment="e2e_aligned", evidence=evidence, override=True, risks=risks) + + return done(alignment="e2e_mismatch", evidence=evidence, override=False, risks=risks) +``` diff --git a/experimental/lite/skills/basic/align-precision.md b/experimental/lite/skills/basic/align-precision.md new file mode 100644 index 00000000000..afbd29dccb3 --- /dev/null +++ b/experimental/lite/skills/basic/align-precision.md @@ -0,0 +1,99 @@ +# Align Precision Skill + +Align Megatron Lite behavior against a reference, or localize the first precision +break with controlled variables. + +## Schema + + +```python +schema = Skill( + "basic.align_precision", kind="state_machine", purpose="recursively align precision to reference", + imports=["basic.constitution"], + calls=[ + "basic.constitution", "basic.find_reference", "basic.construct_proxy_task", + "basic.review_threshold", "basic.align_precision", "basic.align_e2e_precision", + ], + inputs=["task", "target", "variables", "budget"], + outputs=["alignment", "evidence", "next"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def align_precision(task, target, variables, budget): + if task.scope not in ["precision", "primitive", "model"]: + return out_of_scope("not a precision-alignment task") + if budget.depth > budget.max_depth or budget.attempts > budget.max_attempts: + return blocked("precision recursion budget exhausted") + + ref = basic.find_reference(task, layer=target.layer, candidates=task.references, budget=budget.reference) + if not ref.done: + return blocked("no usable reference", evidence=ref) + + policy = basic.constitution(task, layer=target.layer, reference=ref.reference) + if not policy.done: + return blocked("precision policy failed", evidence=policy) + + proxy_result = basic.construct_proxy_task( + task, + target=target, + reference=ref.reference, + variables=variables, + budget=budget.proxy, + ) + if not proxy_result.done: + return blocked("cannot construct minimal proxy task") + proxy_task = proxy_result.proxy + + controlled = variables.freeze_all_except(task.variable_under_test) + if controlled.has_unknown_unfrozen_axes(): + return blocked("unknown variables remain uncontrolled", risks=controlled.unknown_axes) + axes = controlled.axes(one_at_a_time=True) + + deterministic = run_same_setting(proxy_task, runs=3) + if not deterministic.bitwise_equal(): + proxy_task = remove_known_nondeterminism(proxy_task) + deterministic = run_same_setting(proxy_task, runs=3) + if not deterministic.bitwise_equal() and policy.validation.requires_bitwise: + return blocked("reference or target is not deterministic", evidence=deterministic) + + threshold = basic.review_threshold(task, ref.contract, requested_threshold=task.tolerance) + if not threshold.done: + return blocked("precision threshold review failed", evidence=threshold) + mode = threshold.compare_mode + + full = compare(proxy_task, target, ref.reference, mode=mode) + if full.pass_: + return done(alignment="aligned", evidence=[deterministic, full], next=[]) + + evidence = [deterministic, full] + for axis in axes: + experiment = vary_one_axis(proxy_task, axis=axis) + evidence.append(compare(experiment, target, ref.reference, mode=mode)) + if evidence[-1].first_failure: + break + + for unit in decompose(target, order=["layer", "submodule", "primitive"]): + child = basic.align_precision( + task.narrow_to(unit), + target=unit, + variables=variables.freeze_except(unit.related_axes), + budget=budget.child(increment=["depth", "attempts"]), + ) + evidence.append(child) + if child.alignment in ["aligned", "localized_mismatch"] or child.blocked: + break + + if budget.e2e.high_cost_approved: + e2e = basic.align_e2e_precision(task, target=target, modes=task.e2e_modes, budget=budget.e2e) + evidence.append(e2e) + if e2e.alignment == "e2e_aligned": + return done( + alignment="e2e_aligned_override", + evidence=evidence, + next=["local mismatch overridden by high-cost e2e evidence", *e2e.risks], + ) + + return done(alignment="localized_mismatch", evidence=evidence, next=owner_of_first_failure(evidence)) +``` diff --git a/experimental/lite/skills/basic/constitution.md b/experimental/lite/skills/basic/constitution.md new file mode 100644 index 00000000000..8d9c006a447 --- /dev/null +++ b/experimental/lite/skills/basic/constitution.md @@ -0,0 +1,45 @@ +# Basic Constitution Skill + +Global constraints for all Megatron Lite skills. + +## Schema + + +```python +schema = Skill( + "basic.constitution", kind="constitution", purpose="set global MLite constraints", + imports=[], calls=[], + inputs=["task", "layer", "reference"], + outputs=["constraints", "validation", "stop"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def constitution(task, layer, reference): + if reference is None: + return blocked("no validation reference") + if task.requires_undefined_skill: + return out_of_scope("required procedural skill is undefined") + + constraints = [ + occam_razor("choose the smallest correct and reviewable design"), + modularity("primitives are replaceable unless explicitly fused"), + bounded_state_machine("every procedural skill has finite exits"), + ] + + validation = [ + reference_order(["Megatron", "HuggingFace", "Torch", "first principles"]), + require_bitwise_when_possible(reference), + minimal_primitive_check(scope=["single_gpu", "single_node"], reduce=["layers", "experts"]), + require_end_to_end_before_delivery(), + ] + + stop = [ + "missing reference", + "missing validation path", + "required procedural skill is undefined", + ] + + return done(constraints=constraints, validation=validation, stop=stop) +``` diff --git a/experimental/lite/skills/basic/construct-proxy-task.md b/experimental/lite/skills/basic/construct-proxy-task.md new file mode 100644 index 00000000000..3214f36d290 --- /dev/null +++ b/experimental/lite/skills/basic/construct-proxy-task.md @@ -0,0 +1,38 @@ +# Construct Proxy Task Skill + +Build the smallest task that can falsify a claim. + +## Schema + + +```python +schema = Skill( + "basic.construct_proxy_task", kind="state_machine", purpose="minimize validation while preserving claim", + imports=["basic.constitution"], calls=[], + inputs=["task", "target", "reference", "variables", "budget"], + outputs=["proxy", "controlled", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def construct_proxy_task(task, target, reference, variables, budget): + if reference is None: + return blocked("proxy task needs a reference") + + proxy = minimize( + target, + start=["single_gpu", "single_node"], + reduce=["layers", "experts", "sequence", "hidden", "batch"], + preserve=["operation", "shape_rules", "dtype_rules", task.variable_under_test], + ) + if proxy is None or not proxy.still_tests(task.claim): + return blocked("cannot build minimal proxy that preserves claim") + + controlled = variables.freeze_all_except(task.variable_under_test) + if controlled.has_unknown_unfrozen_axes(): + return blocked("proxy has uncontrolled variables", risks=controlled.unknown_axes) + + risks = ["proxy may miss bugs that require full scale"] + return done(proxy=proxy, controlled=controlled, risks=risks) +``` diff --git a/experimental/lite/skills/basic/find-reference.md b/experimental/lite/skills/basic/find-reference.md new file mode 100644 index 00000000000..526f5b1defb --- /dev/null +++ b/experimental/lite/skills/basic/find-reference.md @@ -0,0 +1,51 @@ +# Find Reference Skill + +Find the strongest checkable reference for a Megatron Lite task. + +## Schema + + +```python +schema = Skill( + "basic.find_reference", kind="state_machine", purpose="select a checkable validation reference", + imports=[], calls=[], + inputs=["task", "layer", "candidates", "budget"], + outputs=["reference", "contract", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def find_reference(task, layer, candidates, budget): + if task.scope not in ["primitive", "model", "runtime", "application", "precision"]: + return out_of_scope("unknown validation surface") + + ordered = [ + megatron_reference(task, layer), + huggingface_reference(task, layer), + torch_reference(task, layer), + first_principles_formula_or_distributed_invariant(task, layer), + *candidates, + ] + + risks = [] + for ref in ordered[:budget.max_candidates]: + if ref is None or not ref.exists(): + continue + + contract = extract_contract( + ref, + fields=["inputs", "outputs", "shape", "dtype", "seed", "variables", "tolerance"], + ) + if not contract.is_checkable(): + risks.append((ref, "reference exists but contract is incomplete")) + continue + if ref.requires_unavailable_assets(): + risks.append((ref, "reference requires unavailable assets")) + continue + + variables = freeze_variables(contract.variables, except_=task.variable_under_test) + return done(reference=ref, contract=contract.with_variables(variables), risks=risks) + + return blocked("no checkable reference", risks=risks) +``` diff --git a/experimental/lite/skills/basic/lint-skill.md b/experimental/lite/skills/basic/lint-skill.md new file mode 100644 index 00000000000..3700d05550d --- /dev/null +++ b/experimental/lite/skills/basic/lint-skill.md @@ -0,0 +1,58 @@ +# Lint Skill + +Validate a Megatron Lite skill file before review. + +## Schema + + +```python +schema = Skill( + "basic.lint_skill", kind="state_machine", purpose="validate skill structure and dry-runs", + imports=[], calls=[], + inputs=["skill_file", "registry", "scenarios", "budget"], + outputs=["lint", "dry_runs", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def lint_skill(skill_file, registry, scenarios, budget): + if not skill_file.path.startswith("experimental/lite/skills/"): + return out_of_scope("not an MLite skill") + + schema = extract_schema_block(skill_file) + if schema is None: + return blocked("missing schema markers") + + spec = parse_skill_schema(schema) + expected_path = module_name_to_path(spec.name) + if skill_file.path != expected_path: + return blocked("schema name does not match file path", evidence=[spec.name, skill_file.path]) + + lint = [ + require_single_md_file(skill_file), + reject_readme_skill_body(skill_file), + require_python_like_body(skill_file), + require_top_level_function(skill_file, name=spec.name.split(".")[-1], inputs=spec.inputs), + require_declared_exits(skill_file, exits=spec.exits), + require_bounded_loops(skill_file), + require_resolved_imports(spec.imports, registry), + require_resolved_calls(spec.calls, registry), + ] + body_calls = extract_skill_calls(skill_file.body, registry=registry) + lint.extend([ + require_body_calls_declared(body_calls, declared=spec.calls), + require_declared_calls_used(spec.calls, body_calls, allow_recursive=True), + ]) + if any(check.fail for check in lint): + return blocked("skill lint failed", lint=lint) + + dry_runs = [] + for scenario in scenarios[:budget.max_scenarios]: + dry_runs.append(trace_pseudocode(skill_file, scenario, max_steps=budget.max_steps)) + if dry_runs[-1].exit not in spec.exits: + return blocked("scenario reached undeclared exit", dry_runs=dry_runs) + + risks = ["dry-runs are contract checks, not executable unit tests"] + return done(lint=lint, dry_runs=dry_runs, risks=risks) +``` diff --git a/experimental/lite/skills/basic/review-threshold.md b/experimental/lite/skills/basic/review-threshold.md new file mode 100644 index 00000000000..cd296ea3dd2 --- /dev/null +++ b/experimental/lite/skills/basic/review-threshold.md @@ -0,0 +1,33 @@ +# Review Threshold Skill + +Choose the comparison mode and require human review for non-bitwise thresholds. + +## Schema + + +```python +schema = Skill( + "basic.review_threshold", kind="state_machine", purpose="select bitwise or reviewed tolerance", + imports=["basic.constitution"], calls=[], + inputs=["task", "reference_contract", "requested_threshold"], + outputs=["compare_mode", "review", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def review_threshold(task, reference_contract, requested_threshold): + if reference_contract is None: + return blocked("missing reference contract for threshold review") + + if reference_contract.bitwise_possible: + return done(compare_mode=bitwise(), review=False, risks=[]) + + threshold = requested_threshold or relative(0.01) + review = human_review_required( + reason="non-bitwise precision threshold; 1% relative is only a default", + threshold=threshold, + ) + risks = ["accepted threshold can hide real precision bugs"] + return done(compare_mode=threshold, review=review, risks=risks) +``` diff --git a/experimental/lite/skills/insight/progressive-model-support.md b/experimental/lite/skills/insight/progressive-model-support.md new file mode 100644 index 00000000000..2599564ea36 --- /dev/null +++ b/experimental/lite/skills/insight/progressive-model-support.md @@ -0,0 +1,201 @@ +# Progressive Model Support Skill + +Plan model support through either an HF-first path or a nearest-reference diff path. + +## Schema + + +```python +schema = Skill( + "insight.progressive_model_support", kind="state_machine", purpose="stage model support with HF and bridge/reference paths", + imports=["basic.constitution"], + calls=[ + "basic.find_reference", "basic.align_precision", "basic.align_e2e_precision", + "primitive.select_for_compose", + "model_compose.config_mapping", "model_compose.weight_mapping", "model_compose.build_model", + ], + inputs=["task", "model", "budget"], + outputs=["path", "stages", "evidence", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def progressive_model_support(task, model, budget): + if model.is_pure_new or not model.has_bridge_or_nearest_reference(): + result = support_pure_new_model_from_hf(task, model, budget) + else: + result = support_model_from_bridge_reference(task, model, budget) + + if not result.done: + return blocked("progressive model support failed", evidence=result) + return done(path=result.path, stages=result.stages, evidence=result.evidence, risks=result.risks) + + +def support_pure_new_model_from_hf(task, model, budget): + stages = ["hf_config", "hf_weights", "PP", "EP", "CP_if_needed", "TP_if_needed", "e2e"] + evidence = [] + + config = model_compose.config_mapping(task, model.hf.config, model.lite_config, budget.config) + evidence.append(config) + if not config.done: + return blocked("pure-new HF path stopped at config mapping", evidence=evidence) + + weights = model_compose.weight_mapping(task, model.hf.checkpoint, model.lite_model, budget.weights) + evidence.append(weights) + if not weights.done: + return blocked("pure-new HF path stopped at weight mapping", evidence=evidence) + + primitive_steps = [ + step("PP", required=True, features=["pipeline_parallel"]), + step("EP", required=model.has_moe, features=["expert_parallel", "moe", "deepep"]), + step("CP", required=model.needs_long_context or model.uses_thd, features=["context_parallel", "thd"]), + step("TP", required=model.needs_tensor_parallel_for_memory_or_perf, features=["tensor_parallel"]), + ] + partial = model.empty_lite_model(config=config.mapping, weights=weights.mapping) + + for primitive_step in primitive_steps: + if not primitive_step.required: + continue + selection = primitive.select_for_compose( + task.narrow_to(primitive_step.name), + model_spec=model.spec.require(primitive_step.features), + candidates=model.primitive_candidates(primitive_step.features), + budget=budget.primitive_step(primitive_step.name), + ) + evidence.append(selection) + if not selection.done: + return blocked("pure-new primitive selection failed", evidence=evidence) + + partial = model_compose.build_model( + task.narrow_to(primitive_step.name), + model_spec=partial.spec.with_primitives(selection.selection), + primitives=partial.primitives + selection.selection, + budget=budget.build_step(primitive_step.name), + ) + evidence.append(partial) + if not partial.done: + return blocked("pure-new incremental model build failed", evidence=evidence) + + precision = compare_model_precision_ladder( + task.narrow_to(primitive_step.name), + target=partial.model, + reference=model.hf, + variables=model.variables.freeze_all_except(primitive_step.features), + budget=budget.precision_step(primitive_step.name), + ) + evidence.append(precision) + if not precision.done: + return blocked("pure-new precision ladder failed", evidence=evidence) + + e2e = basic.align_e2e_precision(task, target=partial.model, modes=task.e2e_modes, budget=budget.e2e) + evidence.append(e2e) + if not e2e.done: + return blocked("pure-new model stopped at e2e", evidence=evidence) + + risks = ["HF-first path can miss bugs shared by HF conversion and Lite compose"] + return done(path="hf_first_pure_new", stages=stages, evidence=evidence, risks=risks) + + +def support_model_from_bridge_reference(task, model, budget): + stages = ["nearest_reference", "diff_primitives", "bridge_precision", "hf_cross_check", "e2e"] + evidence = [] + + reference = basic.find_reference( + task, + layer="model", + candidates=[model.bridge_reference, *model.nearest_existing_models, "mbridge", "megatron-bridge"], + budget=budget.reference, + ) + evidence.append(reference) + if not reference.done: + return blocked("no nearest model reference for diff path", evidence=evidence) + + diff = model.diff_against(reference.reference) + selection = primitive.select_for_compose( + task.narrow_to(diff), + model_spec=diff.model_spec, + candidates=diff.primitive_candidates, + budget=budget.diff_primitives, + ) + evidence.append(selection) + if not selection.done: + return blocked("diff primitive selection failed", evidence=evidence) + + config = model_compose.config_mapping(task, model.hf.config, model.lite_config, budget.config) + weights = model_compose.weight_mapping(task, model.hf.checkpoint, model.lite_model, budget.weights) + evidence.extend([config, weights]) + if not config.done or not weights.done: + return blocked("HF config or weight mapping failed in reference path", evidence=evidence) + + candidate = model_compose.build_model( + task, + model_spec=model.spec.apply_diff(diff, config.mapping), + primitives=reference.reference.primitives + selection.selection, + budget=budget.build, + ) + evidence.append(candidate) + if not candidate.done: + return blocked("reference-diff model build failed", evidence=evidence) + + bridge_precision = compare_model_precision_ladder( + task.narrow_to("bridge_reference"), + target=candidate.model, + reference=reference.reference, + variables=model.variables.freeze_all_except(diff.changed_features), + budget=budget.bridge_precision, + ) + evidence.append(bridge_precision) + if not bridge_precision.done: + return blocked("reference-diff precision failed", evidence=evidence) + + hf_precision = compare_model_precision_ladder( + task.narrow_to("hf_cross_check"), + target=candidate.model, + reference=model.hf, + variables=model.variables.freeze_all_except(diff.changed_features), + budget=budget.hf_precision, + ) + evidence.append(hf_precision) + if not hf_precision.done: + return blocked("HF cross-check precision failed", evidence=evidence) + + e2e = basic.align_e2e_precision(task, target=candidate.model, modes=task.e2e_modes, budget=budget.e2e) + evidence.append(e2e) + if not e2e.done: + return blocked("reference-diff model stopped at e2e", evidence=evidence) + + risks = ["nearest-reference path can inherit reference bugs", "HF cross-check is still required"] + return done(path="nearest_reference_diff", stages=stages, evidence=evidence, risks=risks) + + +def compare_model_precision_ladder(task, target, reference, variables, budget): + forward = basic.align_precision( + task.with_phase("forward").with_reference(reference).with_metrics(["bitwise_when_possible", "cos_sim"]), + target=target.forward_proxy(), + variables=variables.freeze_all_except("forward_math"), + budget=budget.forward, + ) + if not forward.done: + return blocked("forward precision failed", evidence=forward) + + backward = basic.align_precision( + task.with_phase("backward_grad").with_reference(reference).with_metrics(["bitwise_when_possible", "grad_cos_sim"]), + target=target.backward_proxy(), + variables=variables.freeze_all_except("backward_grad"), + budget=budget.backward, + ) + if not backward.done: + return blocked("backward gradient precision failed", evidence=[forward, backward]) + + grad_norm = compare_grad_norm( + target, + reference, + mode=["bitwise_when_possible", "relative_threshold", "cos_sim_supporting_signal"], + tolerance=budget.grad_norm_tolerance, + ) + if not grad_norm.pass_: + return blocked("grad norm precision failed", evidence=[forward, backward, grad_norm]) + + return done(alignment="forward_backward_grad_norm_aligned", evidence=[forward, backward, grad_norm], risks=[]) +``` diff --git a/experimental/lite/skills/insight/progressive-primitive-design.md b/experimental/lite/skills/insight/progressive-primitive-design.md new file mode 100644 index 00000000000..4ef2f5857bd --- /dev/null +++ b/experimental/lite/skills/insight/progressive-primitive-design.md @@ -0,0 +1,39 @@ +# Progressive Primitive Design Skill + +Plan a new primitive from first principles to performance. + +## Schema + + +```python +schema = Skill( + "insight.progressive_primitive_design", kind="state_machine", purpose="stage new primitive design", + imports=["basic.constitution"], calls=["primitive.design", "primitive.validate", "perf.optimize"], + inputs=["task", "primitive", "budget"], + outputs=["stages", "evidence", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def progressive_primitive_design(task, primitive, budget): + stages = ["principle", "single_device", "single_node", "distributed", "model_compose", "perf"] + evidence = [] + + design = primitive.design(task, primitive, primitive.requirements, budget.design) + evidence.append(design) + if not design.done: + return blocked("primitive design stopped at contract design", evidence=evidence) + + validation = primitive.validate(task, primitive=primitive, implementation=primitive.implementation, budget=budget.validation) + evidence.append(validation) + if not validation.done: + return blocked("primitive design stopped at validation", evidence=evidence) + + optimization = perf.optimize(task, target=primitive.implementation, constraints=primitive.perf_constraints, budget=budget.perf) + evidence.append(optimization) + if not optimization.done: + return blocked("primitive design stopped at performance", evidence=evidence) + + return done(stages=stages, evidence=evidence, risks=["performance stage must not rewrite correctness contract"]) +``` diff --git a/experimental/lite/skills/insight/scope-control.md b/experimental/lite/skills/insight/scope-control.md new file mode 100644 index 00000000000..f074c92e646 --- /dev/null +++ b/experimental/lite/skills/insight/scope-control.md @@ -0,0 +1,29 @@ +# Scope Control Skill + +Decide when to split work instead of expanding a skill or task. + +## Schema + + +```python +schema = Skill( + "insight.scope_control", kind="state_machine", purpose="prevent skill and task scope creep", + imports=["basic.constitution"], calls=["basic.lint_skill"], + inputs=["task", "change", "budget"], + outputs=["decision", "split", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def scope_control(task, change, budget): + if change.touches_multiple_namespaces() or change.exceeds(budget.max_files): + return done(decision="split", split=propose_child_tasks(change), risks=["large review surface"]) + if change.adds_new_procedure_without_schema(): + return blocked("new procedure needs a skill schema") + + lint = basic.lint_skill(change.skill_file, registry=budget.registry, scenarios=[], budget=budget.lint) + if not lint.done: + return blocked("scope-control lint failed", evidence=lint) + return done(decision="continue", split=[], risks=[]) +``` diff --git a/experimental/lite/skills/model-compose/build-model.md b/experimental/lite/skills/model-compose/build-model.md new file mode 100644 index 00000000000..c921ddfcbc1 --- /dev/null +++ b/experimental/lite/skills/model-compose/build-model.md @@ -0,0 +1,35 @@ +# Build Model Skill + +Compose a Megatron Lite model from validated primitives. + +## Schema + + +```python +schema = Skill( + "model_compose.build_model", kind="state_machine", purpose="compose MLite model from primitives", + imports=["basic.constitution"], calls=["primitive.select_for_compose", "primitive.validate", "basic.align_precision"], + inputs=["task", "model_spec", "primitives", "budget"], + outputs=["model", "validation", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def build_model(task, model_spec, primitives, budget): + selected = primitive.select_for_compose(task, model_spec=model_spec, candidates=primitives, budget=budget.selection) + if not selected.done: + return blocked("primitive selection failed before model compose", evidence=selected) + + for selected_primitive in selected.selection: + result = primitive.validate(task, primitive=selected_primitive, implementation=selected_primitive.impl, budget=budget.primitive) + if not result.done: + return blocked("primitive not validated before model compose", evidence=result) + + model = compose_layers(model_spec, selected.selection, boundary=["runtime", "model", "primitive"]) + precision = basic.align_precision(task, target=model, variables=model.variables, budget=budget.precision) + if not precision.done: + return blocked("model precision failed", evidence=precision) + + return done(model=model, validation=precision, risks=["composition can hide primitive boundary bugs"]) +``` diff --git a/experimental/lite/skills/model-compose/config-mapping.md b/experimental/lite/skills/model-compose/config-mapping.md new file mode 100644 index 00000000000..453695321c0 --- /dev/null +++ b/experimental/lite/skills/model-compose/config-mapping.md @@ -0,0 +1,34 @@ +# Config Mapping Skill + +Map external model configs into Megatron Lite configs. + +## Schema + + +```python +schema = Skill( + "model_compose.config_mapping", kind="state_machine", purpose="map model configs into MLite", + imports=["basic.constitution"], calls=["basic.find_reference"], + inputs=["task", "source_config", "target_config", "budget"], + outputs=["mapping", "evidence", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def config_mapping(task, source_config, target_config, budget): + reference = basic.find_reference(task, layer="config", candidates=[source_config], budget=budget.reference) + if not reference.done: + return blocked("config reference not found", evidence=reference) + + mapping = map_fields( + source_config, + target_config, + required=["hidden_size", "num_layers", "num_heads", "dtype", "moe", "rope"], + ) + if mapping.has_missing_required_fields(): + return blocked("config mapping incomplete", evidence=mapping.missing) + + evidence = record_evidence(task, run=mapping.check, comparison=mapping.comparison, environment=budget.env) + return done(mapping=mapping, evidence=evidence, risks=["default mismatch", "alias mismatch"]) +``` diff --git a/experimental/lite/skills/model-compose/qwen.md b/experimental/lite/skills/model-compose/qwen.md new file mode 100644 index 00000000000..7dc6eae8fd5 --- /dev/null +++ b/experimental/lite/skills/model-compose/qwen.md @@ -0,0 +1,37 @@ +# Qwen Compose Skill + +Compose Qwen3 MoE and Qwen3.5 Megatron Lite models. + +## Schema + + +```python +schema = Skill( + "model_compose.qwen", kind="state_machine", purpose="compose Qwen3 MoE/Qwen3.5 MLite models", + imports=["basic.constitution"], calls=["model_compose.config_mapping", "model_compose.weight_mapping", "model_compose.build_model"], + inputs=["task", "hf_model", "lite_model", "budget"], + outputs=["model", "mapping", "evidence", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def qwen(task, hf_model, lite_model, budget): + if hf_model.family not in ["qwen3_moe", "qwen3_5"]: + return out_of_scope("not a Qwen3 MoE/Qwen3.5 model") + + config = model_compose.config_mapping(task, hf_model.config, lite_model.config, budget.config) + if not config.done: + return blocked("Qwen config mapping failed", evidence=config) + + weights = model_compose.weight_mapping(task, hf_model.checkpoint, lite_model, budget.weights) + if not weights.done: + return blocked("Qwen weight mapping failed", evidence=weights) + + model = model_compose.build_model(task, lite_model.spec, lite_model.primitives, budget.build) + if not model.done: + return blocked("Qwen model build failed", evidence=model) + + risks = ["Qwen alias mismatch", "MoE router drift", "GQA head mapping drift"] + return done(model=model.model, mapping=[config.mapping, weights.mapping], evidence=[weights.evidence, model.validation], risks=risks) +``` diff --git a/experimental/lite/skills/model-compose/weight-mapping.md b/experimental/lite/skills/model-compose/weight-mapping.md new file mode 100644 index 00000000000..f5451539761 --- /dev/null +++ b/experimental/lite/skills/model-compose/weight-mapping.md @@ -0,0 +1,40 @@ +# Weight Mapping Skill + +Map checkpoint weights into Megatron Lite model state. + +## Schema + + +```python +schema = Skill( + "model_compose.weight_mapping", kind="state_machine", purpose="map checkpoint weights into MLite", + imports=["basic.constitution"], calls=["basic.find_reference", "basic.align_precision"], + inputs=["task", "source_checkpoint", "target_model", "budget"], + outputs=["mapping", "coverage", "evidence"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def weight_mapping(task, source_checkpoint, target_model, budget): + candidates = [ + source_checkpoint, + "mbridge", + "megatron-bridge", + *budget.reference.extra_checkpoint_tools, + ] + reference = basic.find_reference(task, layer="checkpoint", candidates=candidates, budget=budget.reference) + if not reference.done: + return blocked("checkpoint reference not found", evidence=reference) + + mapping = map_weight_names_shapes_dtypes(source_checkpoint, target_model, reference_tools=["mbridge", "megatron-bridge"]) + coverage = mapping.coverage() + if not coverage.complete: + return blocked("weight mapping incomplete", evidence=coverage.missing) + + precision = basic.align_precision(task, target=target_model, variables=mapping.variables, budget=budget.precision) + if not precision.done: + return blocked("mapped model precision failed", evidence=precision) + + return done(mapping=mapping, coverage=coverage, evidence=precision) +``` diff --git a/experimental/lite/skills/perf/fusion.md b/experimental/lite/skills/perf/fusion.md new file mode 100644 index 00000000000..1be55e041ce --- /dev/null +++ b/experimental/lite/skills/perf/fusion.md @@ -0,0 +1,31 @@ +# Fusion Skill + +Evaluate fusion opportunities without breaking primitive boundaries. + +## Schema + + +```python +schema = Skill( + "perf.fusion", kind="state_machine", purpose="evaluate safe fusion opportunities", + imports=["basic.constitution"], calls=["primitive.fuse", "perf.measure"], + inputs=["task", "candidates", "target", "budget"], + outputs=["fusion", "evidence", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def fusion(task, candidates, target, budget): + fused = propose_fusion(candidates, target) + decision = primitive.fuse(task, primitives=candidates, fused_design=fused, budget=budget.fuse) + if not decision.done: + return blocked("fusion decision failed", evidence=decision) + + measurement = perf.measure(task, target=fused, workload=budget.workload, budget=budget.measure) + if not measurement.done: + return blocked("fused target measurement failed", evidence=measurement) + + risks = [*decision.risks, "fusion must not replace precision validation"] + return done(fusion=fused, evidence=[decision, measurement], risks=risks) +``` diff --git a/experimental/lite/skills/perf/measure.md b/experimental/lite/skills/perf/measure.md new file mode 100644 index 00000000000..04a9a8d10cd --- /dev/null +++ b/experimental/lite/skills/perf/measure.md @@ -0,0 +1,30 @@ +# Measure Skill + +Measure runtime, memory, and quality metrics with a stable protocol. + +## Schema + + +```python +schema = Skill( + "perf.measure", kind="state_machine", purpose="measure performance without losing precision context", + imports=["basic.constitution"], calls=[], + inputs=["task", "target", "workload", "budget"], + outputs=["metrics", "run", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def measure(task, target, workload, budget): + if workload is None: + return blocked("missing workload") + + run = execute_with_protocol(target, workload, warmup=budget.warmup, repeats=budget.repeats) + metrics = collect_metrics(run, fields=["tokens_per_sec", "step_time", "memory", "loss_or_reward"]) + if metrics.has_missing_fields(): + return blocked("performance evidence incomplete", evidence=metrics) + + risks = ["performance numbers without precision evidence are not sufficient"] + return done(metrics=metrics, run=run, risks=risks) +``` diff --git a/experimental/lite/skills/perf/memory.md b/experimental/lite/skills/perf/memory.md new file mode 100644 index 00000000000..d634f485c54 --- /dev/null +++ b/experimental/lite/skills/perf/memory.md @@ -0,0 +1,40 @@ +# Memory Skill + +Analyze memory across parameters, optimizer state, activations, and offload. + +## Schema + + +```python +schema = Skill( + "perf.memory", kind="state_machine", purpose="analyze MLite memory behavior", + imports=["basic.constitution"], calls=["perf.measure"], + inputs=["task", "target", "memory_config", "budget"], + outputs=["memory", "evidence", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def memory(task, target, memory_config, budget): + measurement = perf.measure(task, target=target, workload=memory_config.workload, budget=budget.measure) + if not measurement.done: + return blocked("memory measurement failed", evidence=measurement) + + memory = split_memory(measurement.metrics.memory, buckets=["params", "grads", "optimizer", "activations", "offload"]) + model = estimate_memory( + target, + reference="https://developer.nvidia.cn/blog/explore-using-the-megatron-core-training-framework-to-improve-gpu-memory-efficiency-in-large-model-training/", + static_axes=["EP*PP changes expert/layer ownership", "FSDP changes param/grad/optimizer ownership"], + dynamic_axes=["TP*CP changes activation and attention working-set shape", "THD changes packed-token activation shape"], + ) + comparison = compare_estimate_to_measurement(model, measurement.metrics.memory) + if not comparison.within(memory_config.tolerance): + return blocked("memory estimate and measurement disagree", evidence=[model, measurement, comparison]) + + risks = [ + "offload can improve capacity while hiding transfer bottlenecks", + "static memory and dynamic peak must be inspected separately", + ] + return done(memory=memory, evidence=[measurement, model, comparison], risks=risks) +``` diff --git a/experimental/lite/skills/perf/optimize.md b/experimental/lite/skills/perf/optimize.md new file mode 100644 index 00000000000..df67b868e73 --- /dev/null +++ b/experimental/lite/skills/perf/optimize.md @@ -0,0 +1,61 @@ +# Optimize Skill + +Iterate on performance while preserving precision evidence. + +## Schema + + +```python +schema = Skill( + "perf.optimize", kind="state_machine", purpose="optimize performance under precision guardrails", + imports=["basic.constitution"], calls=["basic.align_precision", "perf.measure", "perf.fusion", "primitive.design"], + inputs=["task", "target", "candidates", "budget"], + outputs=["best", "evidence", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def optimize(task, target, candidates, budget): + best = None + evidence = [] + queue = candidates[:budget.max_candidates] + + for attempt in range(budget.max_attempts): + if not queue: + break + candidate = queue.pop(0) + + experiment = run_experiment(candidate, workload=budget.workload, knobs=budget.tunable_knobs) + profile = perf.measure(task, target=candidate, workload=experiment.workload, budget=budget.measure) + evidence.append((candidate, experiment, profile)) + if not profile.done: + continue + + opportunities = analyze_profile( + profile.metrics, + spaces=["compute_communication_overlap", "kernel_fusion", "schedule_tuning", "new_or_split_primitive"], + ) + if opportunities.require_fusion: + fusion = perf.fusion(task, candidates=opportunities.fusion_candidates, target=candidate, budget=budget.fusion) + evidence.append(fusion) + if fusion.done: + queue.append(fusion.fusion) + if opportunities.require_primitive_design: + design = primitive.design(task, opportunities.primitive, opportunities.requirements, budget.primitive_design) + evidence.append(design) + if design.done: + queue.append(design.design) + + tuned = tune_parameters(candidate, opportunities.knobs) + precision = basic.align_precision(task, target=tuned, variables=tuned.variables, budget=budget.precision) + if not precision.done: + continue + measurement = perf.measure(task, target=tuned, workload=budget.workload, budget=budget.measure) + evidence.append((tuned, precision, measurement)) + best = choose_better(best, tuned, measurement.metrics) + + if best is None: + return blocked("no optimized candidate preserved precision", evidence=evidence) + return done(best=best, evidence=evidence, risks=["optimization search can overfit benchmark", "profiling must be repeated after every primitive change"]) +``` diff --git a/experimental/lite/skills/primitive/checkpoint/distckpt.md b/experimental/lite/skills/primitive/checkpoint/distckpt.md new file mode 100644 index 00000000000..fef29da385b --- /dev/null +++ b/experimental/lite/skills/primitive/checkpoint/distckpt.md @@ -0,0 +1,45 @@ +# Distopt Distributed Checkpoint Skill + +Define, implement, use, and validate Megatron Core distributed checkpointing for MLite distopt continuity. + +## Schema + + +```python +schema = Skill( + "primitive.checkpoint.distckpt", kind="primitive", purpose="checkpoint DistributedOptimizer state with mcore dist_checkpointing", + imports=["basic.constitution"], calls=["primitive.contract", "primitive.validate", "primitive.optimizer.distopt"], + inputs=["task", "implementation", "config", "reference", "budget"], + outputs=["principle", "implementation_contract", "usage_contract", "validation", "risks"], + exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def distckpt(task, implementation, config, reference, budget): + contract = primitive.contract(implementation.distckpt, scope=task.scope, reference=reference) + if not contract.done: + return blocked("Distopt distckpt contract not satisfied", evidence=contract) + + principle = { + "semantics": "use mcore dist_checkpointing for DistributedOptimizer model and tensor optimizer state", + "invariants": ["save-load-continue matches uninterrupted training", "model sharded_state_dict keys are MLite-local"], + "reference": reference or "mcore DistributedOptimizer.sharded_state_dict roundtrip", + } + implementation_contract = { + "owned_files": ["primitive.ckpt.distckpt", "model protocol distopt compose sites"], + "state": ["model ShardedTensor metadata", "optimizer fp32 master params", "optimizer exp_avg", "optimizer exp_avg_sq", "optimizer step"], + "boundaries": ["do not attach sharded_state_dict to generic model classes", "do not change FSDP2 or mfsdp checkpoint paths"], + } + usage_contract = { + "choose_when": ["runtime use_dcp=True", "optimizer exposes sharded_state_dict", "model chunks expose sharded_state_dict"], + "avoid_when": ["use_dcp=False local checkpoint", "non-distopt optimizer", "cross-tool mcore-MLite checkpoint interchange"], + "compose_with": ["primitive.optimizer.distopt", "model-compose owned opt-in"], + } + validation = primitive.validate(task, primitive=implementation.distckpt, implementation=implementation, budget=budget) + risks = ["missing optimizer tensor state", "model key drift", "replica_id or shard offset mismatch"] + if not validation.done: + return blocked("Distopt distckpt validation failed", evidence=validation) + return done(principle=principle, implementation_contract=implementation_contract, usage_contract=usage_contract, validation=validation, risks=risks) +``` diff --git a/experimental/lite/skills/primitive/contract.md b/experimental/lite/skills/primitive/contract.md new file mode 100644 index 00000000000..7f270293f9f --- /dev/null +++ b/experimental/lite/skills/primitive/contract.md @@ -0,0 +1,61 @@ +# Primitive Contract Skill + +Define the required outputs for every MLite primitive skill. + +## Schema + + +```python +schema = Skill( + "primitive.contract", kind="constitution", purpose="define primitive skill outputs", + imports=["basic.constitution"], calls=[], + inputs=["primitive", "scope", "reference"], + outputs=["principle", "implementation_contract", "usage_contract", "validation", "risks"], + exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def contract(primitive, scope, reference): + if reference is None: + return blocked("primitive needs a reference or first-principles invariant") + + principle = require([ + "math_or_parallel_semantics", + "invariants", + "shape_dtype_rank_rules", + "what_must_match_reference", + ]) + implementation_contract = require([ + "owned_files_or_modules", + "public_api", + "state_and_config", + "process_groups_or_device_placement", + "forward_backward_update_details", + "failure_modes", + ]) + usage_contract = require([ + "config_keys", + "minimal_example", + "valid_combinations", + "selection_rules", + "unsupported_combinations", + ]) + validation = require([ + "single_gpu_or_single_node_proxy", + "controlled_variables", + "precision_reference", + "composition_test", + "e2e_path_if_applicable", + ]) + risks = ["silent mismatch", "dtype drift", "hidden coupling", "wrong selection rule"] + + return done( + principle=principle, + implementation_contract=implementation_contract, + usage_contract=usage_contract, + validation=validation, + risks=risks, + ) +``` diff --git a/experimental/lite/skills/primitive/design.md b/experimental/lite/skills/primitive/design.md new file mode 100644 index 00000000000..8246a688b9e --- /dev/null +++ b/experimental/lite/skills/primitive/design.md @@ -0,0 +1,41 @@ +# Primitive Design Skill + +Design a replaceable MLite primitive before implementation. + +## Schema + + +```python +schema = Skill( + "primitive.design", kind="state_machine", purpose="design a modular primitive", + imports=["basic.constitution"], calls=["primitive.principle", "primitive.contract", "basic.find_reference"], + inputs=["task", "primitive", "requirements", "budget"], + outputs=["design", "reference", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def design(task, primitive, requirements, budget): + ref = basic.find_reference(task, layer=primitive.layer, candidates=requirements.references, budget=budget.reference) + if not ref.done: + return blocked("no primitive reference", evidence=ref) + + principle = primitive.principle(primitive, reference=ref.reference, constraints=requirements.constraints) + if not principle.done: + return blocked("primitive principle failed", evidence=principle) + + contract = primitive.contract(primitive, scope=requirements.scope, reference=ref.reference) + if not contract.done: + return blocked("primitive contract failed", evidence=contract) + + design = { + "principle": principle.principle, + "implementation_details": define_owned_modules_and_dataflow(primitive), + "api": define_inputs_outputs_config_keys(primitive), + "composition": declare_valid_and_invalid_combinations(primitive), + "selection": decide_when_to_use_or_not_use(primitive), + "replaceability": require_no_hidden_dependency(primitive), + } + return done(design=design, reference=ref, risks=contract.risks) +``` diff --git a/experimental/lite/skills/primitive/fuse.md b/experimental/lite/skills/primitive/fuse.md new file mode 100644 index 00000000000..2c83f1b6238 --- /dev/null +++ b/experimental/lite/skills/primitive/fuse.md @@ -0,0 +1,40 @@ +# Primitive Fuse Skill + +Decide whether primitive coupling is allowed as an explicit fused primitive. + +## Schema + + +```python +schema = Skill( + "primitive.fuse", kind="state_machine", purpose="approve or reject primitive fusion", + imports=["basic.constitution"], calls=["primitive.validate", "basic.align_precision"], + inputs=["task", "primitives", "fused_design", "budget"], + outputs=["decision", "validation", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def fuse(task, primitives, fused_design, budget): + if not fused_design.names_all_coupled_primitives(primitives): + return blocked("fusion hides primitive coupling") + if not fused_design.has_independent_fallback_or_reference(): + return blocked("fusion needs an unfused reference") + + structure = primitive.validate( + task, + primitive=fused_design.primitive, + implementation=fused_design.implementation, + budget=budget.validate, + ) + if not structure.done: + return blocked("fused primitive structure not validated", evidence=structure) + + precision = basic.align_precision(task, target=fused_design, variables=fused_design.variables, budget=budget.precision) + if not precision.done: + return blocked("fused primitive precision not validated", evidence=precision) + + risks = ["fusion can mask individual primitive bugs", "fusion reduces replaceability"] + return done(decision="approved_fused_primitive", validation=[structure, precision], risks=risks) +``` diff --git a/experimental/lite/skills/primitive/module/gqa.md b/experimental/lite/skills/primitive/module/gqa.md new file mode 100644 index 00000000000..53fdd65c767 --- /dev/null +++ b/experimental/lite/skills/primitive/module/gqa.md @@ -0,0 +1,46 @@ +# GQA Primitive Skill + +Define, implement, use, and validate grouped-query attention primitives. + +## Schema + + +```python +schema = Skill( + "primitive.module.gqa", kind="primitive", purpose="define and validate GQA module primitive", + imports=["basic.constitution"], calls=["primitive.contract", "primitive.validate"], + inputs=["task", "implementation", "config", "reference", "budget"], + outputs=["principle", "implementation_contract", "usage_contract", "validation", "risks"], + exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def gqa(task, implementation, config, reference, budget): + contract = primitive.contract(implementation.gqa, scope=task.scope, reference=reference) + if not contract.done: + return blocked("GQA contract not satisfied", evidence=contract) + + principle = { + "semantics": "many query heads share fewer key/value heads", + "invariants": ["KV repeat mapping matches reference", "attention mask and rotary positions are unchanged"], + "reference": reference or "HuggingFace attention implementation", + } + implementation_contract = { + "details": ["qkv projection layout", "head mapping", "kv repeat", "attention call"], + "state": ["num_attention_heads", "num_key_value_heads", "head_dim"], + "boundaries": ["GQA owns head mapping; TP owns sharding of projection weights; THD owns packed sequence boundaries"], + } + usage_contract = { + "config": require_config_keys(config, ["num_attention_heads", "num_key_value_heads"]), + "choose_when": ["architecture declares grouped KV heads"], + "avoid_when": ["num_attention_heads not divisible by num_key_value_heads"], + "compose_with": ["TP projection sharding", "CP context sharding with explicit head/context mapping", "THD packed attention when use_thd=True"], + } + validation = primitive.validate(task, primitive=implementation.gqa, implementation=implementation, budget=budget) + risks = ["head repeat mismatch", "qkv layout bug", "attention mask drift"] + if not validation.done: + return blocked("GQA validation failed", evidence=validation) + return done(principle=principle, implementation_contract=implementation_contract, usage_contract=usage_contract, validation=validation, risks=risks) +``` diff --git a/experimental/lite/skills/primitive/module/moe.md b/experimental/lite/skills/primitive/module/moe.md new file mode 100644 index 00000000000..eb52ae650d8 --- /dev/null +++ b/experimental/lite/skills/primitive/module/moe.md @@ -0,0 +1,58 @@ +# MoE Primitive Skill + +Define, implement, use, and validate mixture-of-experts primitives. + +## Schema + + +```python +schema = Skill( + "primitive.module.moe", kind="primitive", purpose="define and validate MoE module primitive", + imports=["basic.constitution"], calls=["primitive.contract", "primitive.validate", "primitive.parallel.ep"], + inputs=["task", "implementation", "config", "reference", "budget"], + outputs=["principle", "implementation_contract", "usage_contract", "validation", "risks"], + exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def moe(task, implementation, config, reference, budget): + contract = primitive.contract(implementation.moe, scope=task.scope, reference=reference) + if not contract.done: + return blocked("MoE contract not satisfied", evidence=contract) + + principle = { + "semantics": "router selects experts and combines weighted expert outputs", + "invariants": ["router logits/topk match reference", "dispatch/combine preserves token order"], + "reference": reference or "HuggingFace/Megatron MoE layer or first-principles weighted sum", + } + implementation_contract = { + "details": ["router", "topk", "token dispatcher", "DeepEP dispatcher", "expert MLP", "combine weights"], + "state": ["expert params", "router dtype", "capacity/drop policy", "DeepEP dispatch metadata"], + "boundaries": ["module owns routing math; EP owns distributed expert placement; DeepEP owns dispatch/combine transport"], + } + usage_contract = { + "config": require_config_keys(config, ["num_experts", "top_k", "expert_parallel_size", "use_deepep"]), + "choose_when": ["model architecture has sparse experts", "expert count justifies EP or grouped GEMM"], + "avoid_when": ["router tie behavior cannot be stabilized", "DeepEP metadata cannot be validated against all-to-all"], + "compose_with": ["primitive.parallel.ep", "DeepEP dispatcher when EP>1", "primitive.parallel.tp for expert MLP if explicit"], + } + ep_validation = None + if config.expert_parallel_size > 1: + ep_validation = primitive.parallel.ep(task, implementation=implementation, config=config, reference=reference, budget=budget.ep) + if not ep_validation.done: + return blocked("MoE EP composition failed", evidence=ep_validation) + + validation = primitive.validate(task, primitive=implementation.moe, implementation=implementation, budget=budget) + risks = ["router tie sensitivity", "expert capacity mismatch", "DeepEP metadata mismatch", "hidden state flattening bug"] + if not validation.done: + return blocked("MoE validation failed", evidence=validation) + return done( + principle=principle, + implementation_contract=implementation_contract, + usage_contract=usage_contract, + validation=[ep_validation, validation], + risks=risks, + ) +``` diff --git a/experimental/lite/skills/primitive/module/thd.md b/experimental/lite/skills/primitive/module/thd.md new file mode 100644 index 00000000000..762923c6242 --- /dev/null +++ b/experimental/lite/skills/primitive/module/thd.md @@ -0,0 +1,59 @@ +# THD Primitive Skill + +Define, implement, use, and validate packed THD variable-length attention. + +## Schema + + +```python +schema = Skill( + "primitive.module.thd", kind="primitive", purpose="define and validate THD packed-sequence primitive", + imports=["basic.constitution"], calls=["primitive.contract", "primitive.validate", "primitive.parallel.cp"], + inputs=["task", "implementation", "config", "reference", "budget"], + outputs=["principle", "implementation_contract", "usage_contract", "validation", "risks"], + exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def thd(task, implementation, config, reference, budget): + contract = primitive.contract(implementation.thd, scope=task.scope, reference=reference) + if not contract.done: + return blocked("THD contract not satisfied", evidence=contract) + + principle = { + "semantics": "pack variable-length sequences as total_tokens x heads x dim for attention", + "invariants": ["cu_seqlens define sequence boundaries", "pack/unpack never crosses sequence boundaries"], + "reference": reference or "Megatron/Core TE THD packed-sequence attention contract", + } + implementation_contract = { + "details": ["PackedSeqParams", "PackedTHDBatch", "pack_nested_thd", "unpack_packed_thd_to_nested"], + "state": ["cu_seqlens", "cu_seqlens_padded", "max_seqlen", "qkv_format=thd"], + "boundaries": ["THD owns packing; CP owns zigzag rank partition when cp_size > 1"], + } + usage_contract = { + "config": require_config_keys(config, ["use_thd", "sequence_length", "context_parallel_size"]), + "choose_when": ["variable-length SFT/RL data", "padding waste dominates", "attention supports packed THD"], + "avoid_when": ["reference cannot expose cu_seqlens", "operator path lacks packed-sequence support"], + "compose_with": ["primitive.parallel.cp via zigzag THD slicing", "primitive.module.gqa attention path"], + } + + cp_validation = None + if config.context_parallel_size > 1 and not task.stack.contains("primitive.parallel.cp"): + cp_validation = primitive.parallel.cp(task.push("primitive.module.thd"), implementation=implementation, config=config, reference=reference, budget=budget.cp) + if not cp_validation.done: + return blocked("THD CP composition failed", evidence=cp_validation) + + validation = primitive.validate(task, primitive=implementation.thd, implementation=implementation, budget=budget) + risks = ["cu_seqlens mismatch", "padding alignment drift", "CP zigzag pack/unpack bug"] + if not validation.done: + return blocked("THD validation failed", evidence=validation) + return done( + principle=principle, + implementation_contract=implementation_contract, + usage_contract=usage_contract, + validation=[cp_validation, validation], + risks=risks, + ) +``` diff --git a/experimental/lite/skills/primitive/optimizer/distopt.md b/experimental/lite/skills/primitive/optimizer/distopt.md new file mode 100644 index 00000000000..25b3c41b687 --- /dev/null +++ b/experimental/lite/skills/primitive/optimizer/distopt.md @@ -0,0 +1,46 @@ +# Distributed Optimizer Skill + +Define, implement, use, and validate distributed optimizer sharding. + +## Schema + + +```python +schema = Skill( + "primitive.optimizer.distopt", kind="primitive", purpose="define and validate distributed optimizer primitive", + imports=["basic.constitution"], calls=["primitive.contract", "primitive.validate"], + inputs=["task", "implementation", "config", "reference", "budget"], + outputs=["principle", "implementation_contract", "usage_contract", "validation", "risks"], + exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def distopt(task, implementation, config, reference, budget): + contract = primitive.contract(implementation.distopt, scope=task.scope, reference=reference) + if not contract.done: + return blocked("DistOpt contract not satisfied", evidence=contract) + + principle = { + "semantics": "shard optimizer state and updates across data-parallel ranks", + "invariants": ["global update equals unsharded optimizer", "state partition is deterministic"], + "reference": reference or "single-rank optimizer with same grads", + } + implementation_contract = { + "details": ["state partition", "grad shard ownership", "update and param sync"], + "state": ["momentum/variance shard", "master param ownership", "offload policy"], + "boundaries": ["optimizer primitive owns update state; data-parallel grad sync is a separate runtime contract"], + } + usage_contract = { + "config": require_config_keys(config, ["use_distributed_optimizer", "data_parallel_size"]), + "choose_when": ["optimizer state memory dominates", "replicated optimizer state is too expensive"], + "avoid_when": ["single-rank proxy debugging", "FSDP already owns optimizer sharding"], + "compose_with": ["data-parallel gradient path", "TP/EP/PP with clear DP group"], + } + validation = primitive.validate(task, primitive=implementation.distopt, implementation=implementation, budget=budget) + risks = ["state shard drift", "param sync mismatch", "offload update device mismatch"] + if not validation.done: + return blocked("DistOpt validation failed", evidence=validation) + return done(principle=principle, implementation_contract=implementation_contract, usage_contract=usage_contract, validation=validation, risks=risks) +``` diff --git a/experimental/lite/skills/primitive/optimizer/fsdp.md b/experimental/lite/skills/primitive/optimizer/fsdp.md new file mode 100644 index 00000000000..e3c85ebe39d --- /dev/null +++ b/experimental/lite/skills/primitive/optimizer/fsdp.md @@ -0,0 +1,46 @@ +# FSDP Skill + +Define, implement, use, and validate fully sharded data parallelism. + +## Schema + + +```python +schema = Skill( + "primitive.optimizer.fsdp", kind="primitive", purpose="define and validate FSDP optimizer primitive", + imports=["basic.constitution"], calls=["primitive.contract", "primitive.validate"], + inputs=["task", "implementation", "config", "reference", "budget"], + outputs=["principle", "implementation_contract", "usage_contract", "validation", "risks"], + exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def fsdp(task, implementation, config, reference, budget): + contract = primitive.contract(implementation.fsdp, scope=task.scope, reference=reference) + if not contract.done: + return blocked("FSDP contract not satisfied", evidence=contract) + + principle = { + "semantics": "shard parameters, gradients, and optimizer state while materialized computation matches reference", + "invariants": ["materialized params equal reference", "optimizer update equals unsharded update"], + "reference": reference or "single-rank optimizer update with same params/grads", + } + implementation_contract = { + "details": ["wrap policy", "param shard", "all_gather before compute", "reduce_scatter grads"], + "state": ["optimizer state shard", "param offload", "optimizer offload", "update-state device"], + "boundaries": ["model modules expose params; optimizer primitive owns sharding/offload"], + } + usage_contract = { + "config": require_config_keys(config, ["fsdp_size", "param_offload", "optimizer_offload"]), + "choose_when": ["optimizer memory dominates", "param/optimizer offload is needed"], + "avoid_when": ["debugging model math; use unsharded optimizer first"], + "compose_with": ["TP/EP/PP through explicit param ownership", "DistOpt only with clear owner split"], + } + validation = primitive.validate(task, primitive=implementation.fsdp, implementation=implementation, budget=budget) + risks = ["offload device mismatch", "optimizer state drift", "materialization timing bug"] + if not validation.done: + return blocked("FSDP validation failed", evidence=validation) + return done(principle=principle, implementation_contract=implementation_contract, usage_contract=usage_contract, validation=validation, risks=risks) +``` diff --git a/experimental/lite/skills/primitive/parallel/cp.md b/experimental/lite/skills/primitive/parallel/cp.md new file mode 100644 index 00000000000..324568d1008 --- /dev/null +++ b/experimental/lite/skills/primitive/parallel/cp.md @@ -0,0 +1,52 @@ +# Context Parallel Skill + +Define, implement, use, and validate context parallelism. + +## Schema + + +```python +schema = Skill( + "primitive.parallel.cp", kind="primitive", purpose="define and validate CP primitive", + imports=["basic.constitution"], calls=["primitive.contract", "primitive.validate", "primitive.module.thd"], + inputs=["task", "implementation", "config", "reference", "budget"], + outputs=["principle", "implementation_contract", "usage_contract", "validation", "risks"], + exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def cp(task, implementation, config, reference, budget): + contract = primitive.contract(implementation.cp, scope=task.scope, reference=reference) + if not contract.done: + return blocked("CP contract not satisfied", evidence=contract) + + principle = { + "semantics": "partition sequence context while preserving attention result", + "invariants": ["CP=1 equals unsharded attention", "position and mask mapping are identical"], + "reference": reference or "unsharded attention formula", + } + implementation_contract = { + "details": ["cp_group", "zigzag_split_for_cp", "zigzag_reconstruct_from_cp_parts", "position ids", "mask partition"], + "collectives": ["context gather/scatter or ring attention path"], + "autograd": ["no non-differentiable fallback collectives"], + } + usage_contract = { + "config": require_config_keys(config, ["context_parallel_size", "sequence_length"]), + "choose_when": ["long context exceeds local memory", "attention implementation supports CP"], + "avoid_when": ["reference cannot validate positions/masks", "model attention variant lacks CP support"], + "compose_with": ["TP only with explicit head/context ownership", "PP by stage-local context policy", "THD packed sequences via zigzag THD slicing"], + } + thd_validation = None + if config.use_thd and not task.stack.contains("primitive.module.thd"): + thd_validation = primitive.module.thd(task.push("primitive.parallel.cp"), implementation=implementation, config=config, reference=reference, budget=budget.thd) + if not thd_validation.done: + return blocked("CP THD composition failed", evidence=thd_validation) + + validation = primitive.validate(task, primitive=implementation.cp, implementation=implementation, budget=budget) + risks = ["zigzag reconstruction mismatch", "position mismatch", "mask mismatch", "non-differentiable gather fallback"] + if not validation.done: + return blocked("CP validation failed", evidence=validation) + return done(principle=principle, implementation_contract=implementation_contract, usage_contract=usage_contract, validation=[thd_validation, validation], risks=risks) +``` diff --git a/experimental/lite/skills/primitive/parallel/ep.md b/experimental/lite/skills/primitive/parallel/ep.md new file mode 100644 index 00000000000..e03eb6d2fac --- /dev/null +++ b/experimental/lite/skills/primitive/parallel/ep.md @@ -0,0 +1,46 @@ +# Expert Parallel Skill + +Define, implement, use, and validate expert parallelism. + +## Schema + + +```python +schema = Skill( + "primitive.parallel.ep", kind="primitive", purpose="define and validate EP primitive", + imports=["basic.constitution"], calls=["primitive.contract", "primitive.validate"], + inputs=["task", "implementation", "config", "reference", "budget"], + outputs=["principle", "implementation_contract", "usage_contract", "validation", "risks"], + exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def ep(task, implementation, config, reference, budget): + contract = primitive.contract(implementation.ep, scope=task.scope, reference=reference) + if not contract.done: + return blocked("EP contract not satisfied", evidence=contract) + + principle = { + "semantics": "partition experts across ranks while router semantics stay global", + "invariants": ["EP=1 equals local MoE", "dispatch/combine preserve token order and weights"], + "reference": reference or "Megatron MoE expert parallel path", + } + implementation_contract = { + "details": ["ep_group", "expert placement", "token dispatcher", "combine weights"], + "collectives": ["all_to_all token exchange", "optional grouped GEMM locality"], + "state": ["expert ownership", "capacity/drop policy", "router dtype"], + } + usage_contract = { + "config": require_config_keys(config, ["expert_model_parallel_size", "num_experts"]), + "choose_when": ["num_experts exceeds local capacity", "MoE communication is acceptable"], + "avoid_when": ["router nondeterminism is unresolved", "expert count cannot map cleanly"], + "compose_with": ["TP only with explicit expert and tensor group nesting", "FSDP/DistOpt through optimizer owner"], + } + validation = primitive.validate(task, primitive=implementation.ep, implementation=implementation, budget=budget) + risks = ["router tie instability", "token drop mismatch", "all-to-all participant mismatch"] + if not validation.done: + return blocked("EP validation failed", evidence=validation) + return done(principle=principle, implementation_contract=implementation_contract, usage_contract=usage_contract, validation=validation, risks=risks) +``` diff --git a/experimental/lite/skills/primitive/parallel/pp.md b/experimental/lite/skills/primitive/parallel/pp.md new file mode 100644 index 00000000000..831da8578cd --- /dev/null +++ b/experimental/lite/skills/primitive/parallel/pp.md @@ -0,0 +1,46 @@ +# Pipeline Parallel Skill + +Define, implement, use, and validate pipeline parallelism. + +## Schema + + +```python +schema = Skill( + "primitive.parallel.pp", kind="primitive", purpose="define and validate PP primitive", + imports=["basic.constitution"], calls=["primitive.contract", "primitive.validate"], + inputs=["task", "implementation", "config", "reference", "budget"], + outputs=["principle", "implementation_contract", "usage_contract", "validation", "risks"], + exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def pp(task, implementation, config, reference, budget): + contract = primitive.contract(implementation.pp, scope=task.scope, reference=reference) + if not contract.done: + return blocked("PP contract not satisfied", evidence=contract) + + principle = { + "semantics": "partition ordered layers into pipeline stages", + "invariants": ["stage concatenation equals full layer order", "microbatch schedule preserves gradients"], + "reference": reference or "Megatron pipeline schedule", + } + implementation_contract = { + "details": ["stage assignment", "send/recv activation tensors", "microbatch schedule"], + "state": ["layer ownership", "activation shape", "loss stage"], + "boundaries": ["first/last stage embeddings and heads", "MTP or auxiliary heads if present"], + } + usage_contract = { + "config": require_config_keys(config, ["pipeline_model_parallel_size", "num_layers"]), + "choose_when": ["model depth exceeds single-rank memory", "stageable layer stack"], + "avoid_when": ["tiny proxy unless testing PP itself", "uneven stage ownership without explicit plan"], + "compose_with": ["TP/EP inside each stage", "FSDP/DistOpt outside stage groups"], + } + validation = primitive.validate(task, primitive=implementation.pp, implementation=implementation, budget=budget) + risks = ["stage boundary off by one", "activation shape mismatch", "schedule deadlock"] + if not validation.done: + return blocked("PP validation failed", evidence=validation) + return done(principle=principle, implementation_contract=implementation_contract, usage_contract=usage_contract, validation=validation, risks=risks) +``` diff --git a/experimental/lite/skills/primitive/parallel/tp.md b/experimental/lite/skills/primitive/parallel/tp.md new file mode 100644 index 00000000000..4ffcad1358d --- /dev/null +++ b/experimental/lite/skills/primitive/parallel/tp.md @@ -0,0 +1,46 @@ +# Tensor Parallel Skill + +Define, implement, use, and validate tensor parallelism. + +## Schema + + +```python +schema = Skill( + "primitive.parallel.tp", kind="primitive", purpose="define and validate TP primitive", + imports=["basic.constitution"], calls=["primitive.contract", "primitive.validate"], + inputs=["task", "implementation", "config", "reference", "budget"], + outputs=["principle", "implementation_contract", "usage_contract", "validation", "risks"], + exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def tp(task, implementation, config, reference, budget): + contract = primitive.contract(implementation.tp, scope=task.scope, reference=reference) + if not contract.done: + return blocked("TP contract not satisfied", evidence=contract) + + principle = { + "semantics": "split linear/vocab dimensions across tensor-parallel ranks", + "invariants": ["TP=1 equals unsharded reference", "row/column shard axes are explicit"], + "reference": reference or "Megatron tensor-parallel layers", + } + implementation_contract = { + "details": ["tp_group", "ColumnParallelLinear", "RowParallelLinear", "VocabParallelEmbedding"], + "collectives": ["all_gather output when needed", "reduce_scatter or all_reduce gradients"], + "state": ["sharded weight layout", "bias ownership", "vocab padding"], + } + usage_contract = { + "config": require_config_keys(config, ["tensor_model_parallel_size"]), + "choose_when": ["large matmul or vocab projection", "model has TP-compatible dimensions"], + "avoid_when": ["dimension not divisible and no padding contract", "debugging non-TP primitive"], + "compose_with": ["FSDP/DistOpt through optimizer owner", "PP", "EP with explicit group nesting", "CP/THD with explicit head/context ownership"], + } + validation = primitive.validate(task, primitive=implementation.tp, implementation=implementation, budget=budget) + risks = ["wrong shard axis", "missing collective", "dtype drift across shards"] + if not validation.done: + return blocked("TP validation failed", evidence=validation) + return done(principle=principle, implementation_contract=implementation_contract, usage_contract=usage_contract, validation=validation, risks=risks) +``` diff --git a/experimental/lite/skills/primitive/principle.md b/experimental/lite/skills/primitive/principle.md new file mode 100644 index 00000000000..a30753175dd --- /dev/null +++ b/experimental/lite/skills/primitive/principle.md @@ -0,0 +1,32 @@ +# Primitive Principle Skill + +State the principle and invariants of a primitive before implementation choices. + +## Schema + + +```python +schema = Skill( + "primitive.principle", kind="primitive", purpose="define primitive principle and invariants", + imports=["basic.constitution"], calls=[], + inputs=["primitive", "reference", "constraints"], + outputs=["principle", "invariants", "reference", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def principle(primitive, reference, constraints): + if reference is None and not constraints.has_first_principles: + return blocked("primitive principle needs a reference or first-principles invariant") + + principle = state_math_or_parallel_semantics(primitive, reference=reference) + invariants = [ + define_shape_dtype_rank_rules(primitive), + define_forward_backward_update_equivalence(primitive), + define_bitwise_or_threshold_contract(primitive, reference), + define_single_gpu_or_single_node_proxy(primitive), + ] + risks = ["weak principle leads to implementation-specific tests", "missing invariant hides shared bugs"] + return done(principle=principle, invariants=invariants, reference=reference, risks=risks) +``` diff --git a/experimental/lite/skills/primitive/process-group.md b/experimental/lite/skills/primitive/process-group.md new file mode 100644 index 00000000000..fc9661b41d6 --- /dev/null +++ b/experimental/lite/skills/primitive/process-group.md @@ -0,0 +1,32 @@ +# Process Group Skill + +Validate process-group ownership and collective participation. + +## Schema + + +```python +schema = Skill( + "primitive.process_group", kind="state_machine", purpose="validate process groups and collectives", + imports=["basic.constitution"], calls=[], + inputs=["task", "groups", "collectives", "budget"], + outputs=["mapping", "validation", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def process_group(task, groups, collectives, budget): + mapping = derive_rank_mapping(groups) + if mapping.has_overlap_without_owner(): + return blocked("rank belongs to overlapping groups without explicit owner") + if not collectives.have_same_participants(mapping): + return blocked("collective participants diverge", evidence=collectives.diff(mapping)) + + validation = [ + smoke_collective(groups, collectives, shape="tiny"), + vary_group_size_one_axis_at_a_time(groups, budget=budget), + ] + risks = ["NCCL hang from participant mismatch", "rank-order mismatch"] + return done(mapping=mapping, validation=validation, risks=risks) +``` diff --git a/experimental/lite/skills/primitive/select-for-compose.md b/experimental/lite/skills/primitive/select-for-compose.md new file mode 100644 index 00000000000..fa856d788d0 --- /dev/null +++ b/experimental/lite/skills/primitive/select-for-compose.md @@ -0,0 +1,44 @@ +# Primitive Select For Compose Skill + +Choose primitives for model composition without coupling their implementations. + +## Schema + + +```python +schema = Skill( + "primitive.select_for_compose", kind="state_machine", purpose="choose primitives for model compose", + imports=["basic.constitution"], calls=["primitive.principle"], + inputs=["task", "model_spec", "candidates", "budget"], + outputs=["selection", "rejected", "evidence", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def select_for_compose(task, model_spec, candidates, budget): + selection = [] + rejected = [] + evidence = [] + + for candidate in candidates[:budget.max_candidates]: + principle = primitive.principle(candidate, reference=candidate.reference, constraints=model_spec.constraints) + evidence.append(principle) + if not principle.done: + rejected.append((candidate, "principle missing")) + continue + if not candidate.supports(model_spec.required_features): + rejected.append((candidate, "missing required model feature")) + continue + if candidate.creates_hidden_coupling(selection): + rejected.append((candidate, "hidden primitive coupling")) + continue + selection.append(candidate) + + selection = minimize_complexity(selection, prefer=["single_gpu_proxy", "single_node_proxy", "existing_reference"]) + if not covers_required_features(selection, model_spec.required_features): + return blocked("primitive selection does not cover model spec", evidence=[evidence, rejected]) + + risks = ["selection can overfit one model family", "unsupported combinations must stay explicit"] + return done(selection=selection, rejected=rejected, evidence=evidence, risks=risks) +``` diff --git a/experimental/lite/skills/primitive/validate.md b/experimental/lite/skills/primitive/validate.md new file mode 100644 index 00000000000..34e4b97b96a --- /dev/null +++ b/experimental/lite/skills/primitive/validate.md @@ -0,0 +1,38 @@ +# Primitive Validate Skill + +Validate one primitive before using it in a model. + +## Schema + + +```python +schema = Skill( + "primitive.validate", kind="state_machine", purpose="validate primitive correctness and precision", + imports=["basic.constitution"], calls=["basic.construct_proxy_task", "basic.align_precision"], + inputs=["task", "primitive", "implementation", "budget"], + outputs=["validation", "evidence", "risks"], exits=["done", "blocked", "out_of_scope"], +) +``` + + +```python +def validate(task, primitive, implementation, budget): + proxy = basic.construct_proxy_task( + task, target=primitive, reference=implementation.reference, variables=implementation.variables, budget=budget.proxy + ) + if not proxy.done: + return blocked("primitive proxy task not constructed", evidence=proxy) + + precision = basic.align_precision(task, target=primitive, variables=implementation.variables, budget=budget.precision) + if not precision.done: + return blocked("primitive precision not validated", evidence=precision) + + validation = [ + "static_contract", + "single_gpu_or_node_proxy", + "controlled_variable_precision", + "composition_with_adjacent_primitives", + "usage_example_runs", + ] + return done(validation=validation, evidence=[proxy, precision], risks=precision.next) +``` diff --git a/experimental/lite/tests/README.md b/experimental/lite/tests/README.md new file mode 100644 index 00000000000..6917dab6034 --- /dev/null +++ b/experimental/lite/tests/README.md @@ -0,0 +1,47 @@ +# Megatron Lite Validation + +`experimental/lite/tests` separates MLite validation into two layers: + +- Unit tests: CPU/single-process contract tests for primitive, model, runtime, checkpoint sentinels, config, and helper behavior. Pure helper tests stub optional imports when no Transformer Engine runtime path is exercised; tests that need the real package explicitly skip when unavailable. +- Smoke tests: real `torch.distributed` tests for TP/EP/PP/CP/FSDP2/offload/checkpoint/distopt and tiny Qwen lite forward/backward behavior. Smoke runs are capped at one node and at most 8 GPUs. + +Run unit coverage: + +```bash +PYTHONPATH="$(pwd):$(pwd)/experimental/lite" pytest experimental/lite/tests/unit +``` + +Run smoke coverage on one node: + +```bash +PYTHONPATH="$(pwd):$(pwd)/experimental/lite" MLITE_RUN_SMOKE=1 MLITE_SMOKE_NPROC=8 \ + experimental/lite/tests/run_primitive_validation.sh +``` + +The smoke suite is skipped by default in regular `pytest` runs. Enable it with `--mlite-smoke` or `MLITE_RUN_SMOKE=1`. + +Current matrix: + +| Surface | Unit | Smoke | +| --- | --- | --- | +| TP/EP/PP/CP/SP topology | `unit/primitive/test_parallel_unit.py`, `unit/primitive/test_parallel_dimensions_independent_unit.py` | `smoke/primitive/test_parallel_topologies_smoke.py` | +| TP linear/vocab primitives | `unit/primitive/test_parallel_dimensions_independent_unit.py` | Qwen model smoke exercises TP linear surfaces | +| EP token dispatch | `unit/primitive/test_parallel_dimensions_independent_unit.py` | Qwen model smoke exercises router, dispatcher, and experts | +| THD packing helpers | `unit/primitive/test_parallel_unit.py` | CP topology smoke exercises distributed CP groups | +| GQA/attention split contract | `unit/primitive/test_attention_moe_unit.py` | Qwen model smoke exercises attention forward/backward | +| MoE router/aux-loss contract | `unit/primitive/test_attention_moe_unit.py` | Qwen model smoke exercises router, dispatcher, and experts | +| LoRA adapter primitives | `unit/primitive/test_module_primitives_independent_unit.py` | Qwen model smoke can enable adapters in follow-up coverage | +| MTP/MRoPE/Gated Delta helper contracts | `unit/primitive/test_module_primitives_independent_unit.py`, `unit/primitive/test_ops_data_trainstep_unit.py` | Qwen3.5 MoE model smoke exercises MRoPE/Gated DeltaNet paths | +| Loss/logprob/math ops | `unit/primitive/test_ops_data_trainstep_unit.py` | Qwen model smoke exercises loss plumbing | +| Data/recompute/train-step primitives | `unit/primitive/test_ops_data_trainstep_unit.py` | model/runtime smoke exercises training loop integration | +| DDP + distributed optimizer | `unit/primitive/test_checkpoint_unit.py`, `unit/primitive/test_checkpoint_runtime.py` | `smoke/primitive/test_distopt_checkpoint_smoke.py` | +| FSDP2 config/wrap/offload | `unit/primitive/test_fsdp2_unit.py` | `smoke/primitive/test_fsdp2_offload_checkpoint_smoke.py` | +| FSDP2 save/load resume | `unit/primitive/test_checkpoint_unit.py`, `unit/primitive/test_checkpoint_runtime.py` | `smoke/primitive/test_fsdp2_offload_checkpoint_smoke.py` | +| Checkpoint restore vs direct training | `unit/primitive/test_checkpoint_unit.py`, `unit/primitive/test_checkpoint_runtime.py` | FSDP2 and distopt checkpoint smokes cover distributed restore paths | +| Runtime backend registry/config | `unit/primitive/test_runtime_config_unit.py`, `unit/runtime/test_runtime_backend_unit.py` | covered through checkpoint/model handles | +| Runtime env/offload controls | `unit/runtime/test_runtime_backend_unit.py` | `smoke/primitive/test_fsdp2_offload_checkpoint_smoke.py` | +| Optimizer update-state offload fraction | `unit/primitive/test_runtime_config_unit.py` and single-process CUDA coverage in `unit/primitive/test_fsdp2_offload_gpu.py` | multi-rank offloaded grad clipping is checked against the non-offloaded baseline in `smoke/primitive/test_fsdp2_offload_checkpoint_smoke.py` | +| Qwen3 MoE lite config/build/forward | `unit/model/test_qwen_config_unit.py` | `smoke/model/test_qwen_lite_forward_smoke.py` | +| Qwen3.5 MoE lite config/build/forward | `unit/model/test_qwen_config_unit.py` | `smoke/model/test_qwen_lite_forward_smoke.py` | + +Classic FSDP is not a separate MLite primitive in the current source tree; MLite's native sharded optimizer coverage is FSDP2 plus Megatron DDP/distopt. diff --git a/experimental/lite/tests/conftest.py b/experimental/lite/tests/conftest.py new file mode 100644 index 00000000000..fa2840e238a --- /dev/null +++ b/experimental/lite/tests/conftest.py @@ -0,0 +1,71 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import os +import sys +import types +from pathlib import Path + +import pytest + +LITE_ROOT = Path(__file__).resolve().parents[1] +REPO_ROOT = Path(__file__).resolve().parents[3] +VERL_EXAMPLE_ROOT = LITE_ROOT / "examples" / "verl" +for root in (REPO_ROOT, LITE_ROOT, VERL_EXAMPLE_ROOT): + if str(root) not in sys.path: + sys.path.insert(0, str(root)) + + +def pytest_configure(config): + config.addinivalue_line("markers", "mlite: mark a test as Megatron Lite validation coverage") + config.addinivalue_line( + "markers", + "smoke: mark a Megatron Lite smoke test; skipped unless --mlite-smoke or MLITE_RUN_SMOKE=1 is set", + ) + config.addinivalue_line("markers", "gpu: mark a test as requiring CUDA") + config.addinivalue_line("markers", "distributed: mark a test as requiring torch.distributed") + + +def pytest_addoption(parser): + parser.addoption( + "--mlite-smoke", action="store_true", default=False, help="run Megatron Lite smoke tests" + ) + + +def pytest_collection_modifyitems(config, items): + run_smoke = config.getoption("--mlite-smoke") or os.getenv("MLITE_RUN_SMOKE") == "1" + if run_smoke: + return + skip_smoke = pytest.mark.skip(reason="set --mlite-smoke or MLITE_RUN_SMOKE=1 to run") + for item in items: + if "smoke" in item.keywords: + item.add_marker(skip_smoke) + + +@pytest.fixture +def transformer_engine_import_stub(monkeypatch): + def install() -> None: + try: + import transformer_engine.pytorch # noqa: F401 + + return + except ModuleNotFoundError as exc: + if exc.name not in {"transformer_engine", "transformer_engine.pytorch"}: + raise + + class _UnavailableTE: + def __init__(self, *args, **kwargs): + raise RuntimeError("Transformer Engine is not installed in this test environment.") + + root = types.ModuleType("transformer_engine") + root.__version__ = "0.0.0" + pytorch = types.ModuleType("transformer_engine.pytorch") + pytorch.DotProductAttention = _UnavailableTE + pytorch.LayerNormLinear = _UnavailableTE + pytorch.Linear = _UnavailableTE + pytorch.RMSNorm = _UnavailableTE + root.pytorch = pytorch + monkeypatch.setitem(sys.modules, "transformer_engine", root) + monkeypatch.setitem(sys.modules, "transformer_engine.pytorch", pytorch) + + return install diff --git a/experimental/lite/tests/run_primitive_validation.sh b/experimental/lite/tests/run_primitive_validation.sh new file mode 100755 index 00000000000..d2884a2f1a7 --- /dev/null +++ b/experimental/lite/tests/run_primitive_validation.sh @@ -0,0 +1,19 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)" +cd "${ROOT}" + +export PYTHONPATH="${ROOT}:${ROOT}/experimental/lite:${PYTHONPATH:-}" + +pytest experimental/lite/tests/unit "$@" + +if [[ "${MLITE_RUN_SMOKE:-0}" == "1" ]]; then + NPROC="${MLITE_SMOKE_NPROC:-${WORLD_SIZE:-1}}" + if (( NPROC < 1 || NPROC > 8 )); then + echo "MLITE smoke tests require 1 <= MLITE_SMOKE_NPROC <= 8, got ${NPROC}" >&2 + exit 2 + fi + torchrun --standalone --nproc_per_node="${NPROC}" \ + -m pytest --mlite-smoke experimental/lite/tests/smoke "$@" +fi diff --git a/experimental/lite/tests/smoke/model/test_qwen_lite_forward_smoke.py b/experimental/lite/tests/smoke/model/test_qwen_lite_forward_smoke.py new file mode 100644 index 00000000000..b00ee6b9646 --- /dev/null +++ b/experimental/lite/tests/smoke/model/test_qwen_lite_forward_smoke.py @@ -0,0 +1,156 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist + +pytestmark = [pytest.mark.mlite, pytest.mark.smoke, pytest.mark.gpu, pytest.mark.distributed] + + +def _qwen3_symbols(): + pytest.importorskip("transformer_engine.pytorch") + from megatron.lite.model.qwen3_moe.config import Qwen3MoEConfig + from megatron.lite.model.qwen3_moe.lite.model import Qwen3MoEModel + + return Qwen3MoEConfig, Qwen3MoEModel + + +def _qwen35_symbols(): + pytest.importorskip("transformer_engine.pytorch") + from megatron.lite.model.qwen3_5.config import Qwen35Config + from megatron.lite.model.qwen3_5.lite.model import Qwen35Model + + return Qwen35Config, Qwen35Model + + +@pytest.fixture(scope="module", autouse=True) +def _single_node_cuda_dist(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for Qwen lite model smoke tests.") + if int(os.environ.get("WORLD_SIZE", "1")) > 8: + pytest.skip("Megatron Lite smoke tests are capped at single-node 8 GPUs.") + + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("LOCAL_RANK", "0") + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29531") + + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + created_pg = False + if not dist.is_initialized(): + dist.init_process_group(backend="nccl", init_method="env://") + created_pg = True + yield + if created_pg and dist.is_initialized(): + dist.destroy_process_group() + + +def _tiny_qwen3_config(): + Qwen3MoEConfig, _Qwen3MoEModel = _qwen3_symbols() + return Qwen3MoEConfig( + num_hidden_layers=1, + hidden_size=16, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=4, + vocab_size=64, + num_experts=2, + num_experts_per_tok=1, + moe_intermediate_size=8, + max_position_embeddings=16, + layer_types=["full_attention"], + ) + + +def _tiny_qwen35_config(): + Qwen35Config, _Qwen35Model = _qwen35_symbols() + return Qwen35Config( + num_hidden_layers=1, + hidden_size=16, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=4, + vocab_size=64, + num_experts=2, + num_experts_per_tok=1, + moe_intermediate_size=8, + shared_expert_intermediate_size=8, + linear_num_key_heads=2, + linear_key_head_dim=4, + linear_num_value_heads=2, + linear_value_head_dim=4, + linear_conv_kernel_dim=2, + max_position_embeddings=16, + partial_rotary_factor=1.0, + mrope_section=[1, 1, 0], + layer_types=["full_attention"], + ) + + +def _parallel_state(): + from megatron.lite.primitive.parallel import init_parallel + from megatron.lite.runtime.contracts.config import ParallelConfig + + return init_parallel(ParallelConfig(tp=1, etp=1, ep=1, pp=1, cp=1)) + + +def _token_batch(vocab_size: int): + torch.manual_seed(9876 + dist.get_rank()) + input_ids = torch.randint(0, vocab_size, (2, 4), device="cuda") + labels = torch.randint(0, vocab_size, (2, 4), device="cuda") + return input_ids, labels + + +def _assert_loss_and_backward(output: dict, model: torch.nn.Module): + loss = output["loss"] + assert loss.ndim == 0 + assert torch.isfinite(loss) + loss.backward() + grad_params = [ + param for param in model.parameters() if param.requires_grad and param.grad is not None + ] + assert grad_params + assert all(torch.isfinite(param.grad.detach().float()).all() for param in grad_params) + + +def test_qwen3_moe_lite_tiny_forward_backward_smoke(): + _Qwen3MoEConfig, Qwen3MoEModel = _qwen3_symbols() + config = _tiny_qwen3_config() + model = Qwen3MoEModel(config, _parallel_state(), use_deepep=False).cuda().to(torch.bfloat16) + input_ids, labels = _token_batch(config.vocab_size) + + output = model(input_ids=input_ids, labels=labels, return_log_probs=True) + + assert output["hidden_states"].shape[-1] == config.hidden_size + assert output["log_probs"].shape == labels.shape + _assert_loss_and_backward(output, model) + + +def test_qwen35_lite_tiny_forward_backward_smoke(): + _Qwen35Config, Qwen35Model = _qwen35_symbols() + config = _tiny_qwen35_config() + train_config = SimpleNamespace( + tp=1, + ep=1, + etp=1, + pp=1, + cp=1, + vpp=None, + use_deepep=False, + fp8=False, + recompute_modules=[], + deterministic=True, + ) + model = Qwen35Model(config, train_config, _parallel_state()).cuda().to(torch.bfloat16) + input_ids, labels = _token_batch(config.vocab_size) + + output = model(input_ids=input_ids, labels=labels) + + assert output["hidden_states"].shape[-1] == config.hidden_size + assert output["log_probs"].shape == labels.shape + _assert_loss_and_backward(output, model) diff --git a/experimental/lite/tests/smoke/primitive/test_distopt_checkpoint_smoke.py b/experimental/lite/tests/smoke/primitive/test_distopt_checkpoint_smoke.py new file mode 100644 index 00000000000..e7732201191 --- /dev/null +++ b/experimental/lite/tests/smoke/primitive/test_distopt_checkpoint_smoke.py @@ -0,0 +1,291 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn +from torch.distributed.tensor import Replicate, Shard + +from megatron.core.dist_checkpointing import load_plain_tensors +from megatron.core.dist_checkpointing.dict_utils import diff +from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig +from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer +from megatron.core.transformer import TransformerConfig +from megatron.lite.primitive.ckpt import attach_model_sharded_state_dict +from megatron.lite.primitive.optimizers.megatron_wrap import build_mc_stack +from megatron.lite.primitive.parallel import ParallelState, init_parallel +from megatron.lite.runtime.backends.mlite.runtime import MegatronLiteRuntime +from megatron.lite.runtime.contracts.config import OptimizerConfig as LiteOptimizerConfig +from megatron.lite.runtime.contracts.config import ParallelConfig +from megatron.lite.runtime.contracts.handle import ModelHandle +from tests.unit_tests.test_utilities import Utils + +pytestmark = [pytest.mark.mlite, pytest.mark.smoke, pytest.mark.gpu, pytest.mark.distributed] + + +class TinyDense(nn.Module): + def __init__(self): + super().__init__() + self.fc1 = nn.Linear(8, 8, bias=False) + self.fc2 = nn.Linear(8, 4, bias=False) + + def forward(self, x): + return self.fc2(torch.relu(self.fc1(x))) + + +class TinyTopologyAwareState(nn.Module): + dense_shape = (8, 8) + expert_shape = (8, 4) + + def __init__(self, ps: ParallelState): + super().__init__() + self.dense_weight = nn.Parameter( + _local_shard(_global_tensor(self.dense_shape, 1.0), 0, ps.tp_rank, ps.tp_size) + .cuda() + .bfloat16() + ) + self.experts_weight = nn.Parameter( + _local_shard(_global_tensor(self.expert_shape, 101.0), 0, ps.etp_rank, ps.etp_size) + .cuda() + .bfloat16() + ) + + def forward(self, x): + return x + + +@pytest.fixture(scope="module", autouse=True) +def _single_node_cuda_distopt(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for distopt smoke tests.") + if int(os.environ.get("WORLD_SIZE", "1")) > 8: + pytest.skip("Megatron Lite smoke tests are capped at single-node 8 GPUs.") + + Utils.set_world_size( + int(os.environ.get("WORLD_SIZE", "1")), int(os.environ.get("LOCAL_RANK", "0")) + ) + Utils.initialize_model_parallel() + yield + Utils.destroy_model_parallel() + + +def _global_tensor(shape: tuple[int, ...], offset: float) -> torch.Tensor: + return torch.arange(offset, offset + int(torch.tensor(shape).prod().item())).reshape(shape) + + +def _local_shard(tensor: torch.Tensor, dim: int, rank: int, size: int) -> torch.Tensor: + if size <= 1: + return tensor.clone() + chunks = torch.chunk(tensor, size, dim=dim) + return chunks[rank].contiguous().clone() + + +def _topology_placements(name: str) -> list: + if _is_expert_param(name): + return [Replicate(), Replicate(), Replicate(), Shard(0)] + return [Replicate(), Replicate(), Replicate(), Shard(0)] + + +def _is_expert_param(name: str) -> bool: + return "experts" in name + + +def _build_sharded_model_and_distopt(parallel: ParallelConfig): + ps = init_parallel(parallel) + model = TinyTopologyAwareState(ps) + model_cfg = SimpleNamespace( + num_hidden_layers=1, + hidden_size=8, + num_attention_heads=2, + num_experts=2, + moe_intermediate_size=16, + add_bias_linear=False, + ) + engine_cfg = SimpleNamespace( + model_name="tiny_topology_state", + parallel=parallel, + optimizer=LiteOptimizerConfig(optimizer="adam", lr=1.0e-3, weight_decay=0.0), + deterministic=False, + ) + wrapped_chunks, optimizer = build_mc_stack( + [model], model_cfg=model_cfg, engine_cfg=engine_cfg, ps=ps, is_expert=_is_expert_param + ) + _seed_optimizer_state(optimizer) + attach_model_sharded_state_dict( + wrapped_chunks, ps, get_placements=_topology_placements, is_expert=_is_expert_param + ) + return wrapped_chunks, optimizer, ps + + +def _seed_optimizer_state(optimizer) -> None: + for inner_optimizer in _inner_optimizers(optimizer): + for group in inner_optimizer.param_groups: + for param in group["params"]: + state = inner_optimizer.state[param] + base = param.detach().float().abs() + state["exp_avg"] = (base + 0.125).to(dtype=param.dtype) + state["exp_avg_sq"] = (base + 0.25).to(dtype=param.dtype) + reload_model_params = getattr(optimizer, "reload_model_params", None) + if callable(reload_model_params): + reload_model_params() + + +def _inner_optimizers(optimizer): + chained = getattr(optimizer, "chained_optimizers", None) + if chained is not None: + for chained_optimizer in chained: + yield from _inner_optimizers(chained_optimizer) + return + try: + yield optimizer.optimizer + except AttributeError: + yield optimizer + + +def _distopt_handle(wrapped_chunks, optimizer, ps: ParallelState, parallel: ParallelConfig): + return ModelHandle( + model=wrapped_chunks, + optimizer=optimizer, + parallel_state=ps, + config=SimpleNamespace(parallel=parallel), + _extras={ + "model_chunks": wrapped_chunks, + "protocol": SimpleNamespace( + PLACEMENT_FN=_topology_placements, EXPERT_CLASSIFIER=_is_expert_param + ), + }, + ) + + +def _build_model_and_distopt(): + torch.manual_seed(2468) + model = TinyDense().bfloat16().cuda() + ddp_config = DistributedDataParallelConfig(use_distributed_optimizer=True) + wrapped = DistributedDataParallel( + TransformerConfig(num_attention_heads=1, num_layers=1), ddp_config, model + ) + optimizer = get_megatron_optimizer( + OptimizerConfig(optimizer="adam", lr=1.0e-3, bf16=True, use_distributed_optimizer=True), + [wrapped], + ) + attach_model_sharded_state_dict([wrapped], _single_node_parallel_state()) + return wrapped, optimizer + + +def _single_node_parallel_state() -> ParallelState: + rank = torch.distributed.get_rank() + world = torch.distributed.get_world_size() + return ParallelState(dp_size=world, dp_rank=rank, dp_cp_size=world, dp_cp_rank=rank) + + +def _shared_tmp_path(tmp_path) -> str: + payload = [str(tmp_path) if torch.distributed.get_rank() == 0 else None] + torch.distributed.broadcast_object_list(payload, src=0) + return payload[0] + + +def _train_step(model, optimizer, x: torch.Tensor): + output = model(x) + loss = output.float().square().mean() + loss.backward() + optimizer.step() + optimizer.zero_grad() + if hasattr(model, "zero_grad_buffer"): + model.zero_grad_buffer() + return loss.detach() + + +def _local_named_params(model) -> dict[str, torch.Tensor]: + return {name: param.detach().cpu().float().clone() for name, param in model.named_parameters()} + + +def _assert_model_close(lhs, rhs): + lhs_params = _local_named_params(lhs) + rhs_params = _local_named_params(rhs) + assert lhs_params.keys() == rhs_params.keys() + for name in lhs_params: + torch.testing.assert_close(lhs_params[name], rhs_params[name], atol=0.0, rtol=0.0) + + +def test_distopt_checkpoint_load_matches_uninterrupted_training_single_node(tmp_path): + model_for_ckpt, optimizer_for_ckpt = _build_model_and_distopt() + direct_model, direct_optimizer = _build_model_and_distopt() + loaded_model, loaded_optimizer = _build_model_and_distopt() + runtime = MegatronLiteRuntime.__new__(MegatronLiteRuntime) + + torch.manual_seed(1357) + x0 = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + x1 = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + + _train_step(model_for_ckpt, optimizer_for_ckpt, x0) + _train_step(direct_model, direct_optimizer, x0) + checkpoint_dir = _shared_tmp_path(tmp_path) + + runtime.save_checkpoint( + ModelHandle( + model=model_for_ckpt, + optimizer=optimizer_for_ckpt, + _extras={"model_chunks": [model_for_ckpt]}, + ), + checkpoint_dir, + step=1, + ) + assert ( + runtime.load_checkpoint( + ModelHandle( + model=loaded_model, + optimizer=loaded_optimizer, + _extras={"model_chunks": [loaded_model]}, + ), + checkpoint_dir, + ) + == 1 + ) + + _train_step(direct_model, direct_optimizer, x1) + _train_step(loaded_model, loaded_optimizer, x1) + _assert_model_close(direct_model, loaded_model) + + +def test_distopt_checkpoint_reshards_from_pp_ep_to_tp_pp_ep_etp(tmp_path): + if torch.distributed.get_world_size() < 8: + pytest.skip("TP2/PP2/EP2/ETP2 distopt reshard smoke requires 8 GPUs.") + + runtime = MegatronLiteRuntime.__new__(MegatronLiteRuntime) + checkpoint_root = _shared_tmp_path(tmp_path) + source_dir = os.path.join(checkpoint_root, "source") + reserialized_dir = os.path.join(checkpoint_root, "reserialized") + source_parallel = ParallelConfig(tp=1, ep=2, etp=1, pp=2, cp=1) + target_parallel = ParallelConfig(tp=2, ep=2, etp=2, pp=2, cp=1) + + source_chunks, source_optimizer, source_ps = _build_sharded_model_and_distopt(source_parallel) + runtime.save_checkpoint( + _distopt_handle(source_chunks, source_optimizer, source_ps, source_parallel), + source_dir, + step=3, + save_rng=False, + ) + + target_chunks, target_optimizer, target_ps = _build_sharded_model_and_distopt(target_parallel) + assert ( + runtime.load_checkpoint( + _distopt_handle(target_chunks, target_optimizer, target_ps, target_parallel), + source_dir, + load_rng=False, + ) + == 3 + ) + runtime.save_checkpoint( + _distopt_handle(target_chunks, target_optimizer, target_ps, target_parallel), + reserialized_dir, + step=3, + save_rng=False, + ) + + plain_source = load_plain_tensors(os.path.join(source_dir, "step_3")) + plain_reserialized = load_plain_tensors(os.path.join(reserialized_dir, "step_3")) + diffs = diff(plain_source, plain_reserialized) + assert not any(map(bool, diffs)), diffs diff --git a/experimental/lite/tests/smoke/primitive/test_fsdp2_offload_checkpoint_smoke.py b/experimental/lite/tests/smoke/primitive/test_fsdp2_offload_checkpoint_smoke.py new file mode 100644 index 00000000000..fbaebad32b5 --- /dev/null +++ b/experimental/lite/tests/smoke/primitive/test_fsdp2_offload_checkpoint_smoke.py @@ -0,0 +1,267 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.nn as nn + +from megatron.lite.primitive.optimizers.fsdp2 import ( + FSDP2Config, + build_fsdp2_adamw, + build_fsdp2_device_mesh, + fsdp2_available, + wrap_fsdp2, +) +from megatron.lite.primitive.optimizers.fsdp2.adamw import iter_torch_optimizers, to_local_tensor +from megatron.lite.primitive.parallel.state import ParallelState +from megatron.lite.runtime.backends.mlite.runtime import MegatronLiteRuntime +from megatron.lite.runtime.contracts.config import ParallelConfig +from megatron.lite.runtime.contracts.handle import ModelHandle + +pytestmark = [pytest.mark.mlite, pytest.mark.smoke, pytest.mark.gpu, pytest.mark.distributed] + + +class TinyUnit(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(8, 8) + self.act = nn.GELU() + + def forward(self, x): + return self.act(self.linear(x)) + + +class TinyModel(nn.Module): + def __init__(self): + super().__init__() + self.unit0 = TinyUnit() + self.unit1 = TinyUnit() + self.out = nn.Linear(8, 4) + + def forward(self, x): + return self.out(self.unit1(self.unit0(x))) + + +@pytest.fixture(scope="module", autouse=True) +def _single_node_cuda_dist(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for FSDP2 smoke tests.") + if not fsdp2_available(): + pytest.skip("Installed PyTorch does not expose FSDP2 fully_shard.") + if int(os.environ.get("WORLD_SIZE", "1")) > 8: + pytest.skip("Megatron Lite smoke tests are capped at single-node 8 GPUs.") + + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("LOCAL_RANK", "0") + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29521") + + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + created_pg = False + if not dist.is_initialized(): + dist.init_process_group(backend="nccl", init_method="env://") + created_pg = True + yield + if created_pg and dist.is_initialized(): + dist.destroy_process_group() + + +def _parallel_state() -> ParallelState: + rank = dist.get_rank() + world_size = dist.get_world_size() + return ParallelState( + dp_group=dist.group.WORLD, + dp_cp_group=dist.group.WORLD, + dp_size=world_size, + dp_cp_size=world_size, + dp_rank=rank, + dp_cp_rank=rank, + ) + + +def _shared_tmp_path(tmp_path) -> str: + payload = [str(tmp_path) if dist.get_rank() == 0 else None] + dist.broadcast_object_list(payload, src=0) + return payload[0] + + +def _checkpoint_config(): + return SimpleNamespace(parallel=ParallelConfig()) + + +def _build_fsdp2_model(dtype: torch.dtype = torch.bfloat16) -> tuple[nn.Module, ParallelState]: + torch.manual_seed(1234) + model = TinyModel().cuda().to(dtype=dtype) + ps = _parallel_state() + config = FSDP2Config(unit_modules=(TinyUnit,), reshard_after_forward=True) + mesh = build_fsdp2_device_mesh(ps, config) + return wrap_fsdp2(model, ps, config, mesh=mesh), ps + + +def _build_optimizer(model: nn.Module, ps: ParallelState, *, offload_fraction: float): + return build_fsdp2_adamw( + [model], + SimpleNamespace( + optimizer="adam", + lr=1.0e-3, + weight_decay=0.0, + adam_beta1=0.9, + adam_beta2=0.95, + adam_eps=1.0e-8, + clip_grad=1.0, + offload_fraction=offload_fraction, + ), + ps, + use_fp32_master=True, + ) + + +def _local_param_devices(model: nn.Module) -> set[str]: + return {to_local_tensor(param.detach()).device.type for param in model.parameters()} + + +def _optimizer_state_devices(optimizer) -> set[str]: + devices: set[str] = set() + for child in iter_torch_optimizers(optimizer.optimizer): + for param_state in getattr(child, "state", {}).values(): + if not isinstance(param_state, dict): + continue + for value in param_state.values(): + if isinstance(value, torch.Tensor): + devices.add(to_local_tensor(value).device.type) + return devices + + +def _local_named_params(model: nn.Module) -> dict[str, torch.Tensor]: + return { + name: to_local_tensor(param.detach()).cpu().clone() + for name, param in model.named_parameters() + } + + +def _train_step(model: nn.Module, optimizer, x: torch.Tensor, target: torch.Tensor): + optimizer.zero_grad() + loss = torch.nn.functional.mse_loss(model(x).float(), target.float()) + loss.backward() + success, grad_norm, _ = optimizer.step() + assert success + assert torch.isfinite(torch.tensor(grad_norm)) + return loss.detach(), float(grad_norm) + + +def _assert_grad_norm_exact(lhs: float, rhs: float) -> None: + assert lhs == rhs + + +def _assert_local_params_close(lhs: nn.Module, rhs: nn.Module): + lhs_params = _local_named_params(lhs) + rhs_params = _local_named_params(rhs) + assert lhs_params.keys() == rhs_params.keys() + for name in lhs_params: + torch.testing.assert_close(lhs_params[name], rhs_params[name], atol=0.0, rtol=0.0) + + +def test_fsdp2_runtime_model_and_optimizer_offload_roundtrip_single_node(): + model, ps = _build_fsdp2_model() + optimizer = _build_optimizer(model, ps, offload_fraction=0.0) + handle = ModelHandle( + model=model, optimizer=optimizer, parallel_state=ps, _extras={"model_chunks": [model]} + ) + runtime = MegatronLiteRuntime.__new__(MegatronLiteRuntime) + + assert _local_param_devices(model) == {"cuda"} + assert _optimizer_state_devices(optimizer) == {"cuda"} + + runtime.to(handle, "cpu", model=True, optimizer=True, grad=True) + assert _local_param_devices(model) == {"cpu"} + assert _optimizer_state_devices(optimizer) == {"cpu"} + + runtime.to(handle, "cuda", model=True, optimizer=True, grad=True) + assert _local_param_devices(model) == {"cuda"} + assert _optimizer_state_devices(optimizer) == {"cuda"} + + +def test_fsdp2_offload_fraction_matches_non_offloaded_grad_clip_single_node(): + if dist.get_world_size() < 2: + pytest.skip( + "multi-rank FSDP2 grad clipping equivalence requires WORLD_SIZE > 1." + ) + + baseline_model, baseline_ps = _build_fsdp2_model() + baseline_optimizer = _build_optimizer( + baseline_model, baseline_ps, offload_fraction=0.0 + ) + offload_model, offload_ps = _build_fsdp2_model() + offload_optimizer = _build_optimizer( + offload_model, offload_ps, offload_fraction=1.0 + ) + + assert _optimizer_state_devices(baseline_optimizer) == {"cuda"} + assert _optimizer_state_devices(offload_optimizer) == {"cpu"} + + torch.manual_seed(4321) + x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + target = torch.randn(4, 4, device="cuda", dtype=torch.bfloat16) + _loss, baseline_grad_norm = _train_step( + baseline_model, baseline_optimizer, x, target + ) + _loss, offload_grad_norm = _train_step(offload_model, offload_optimizer, x, target) + + _assert_grad_norm_exact(offload_grad_norm, baseline_grad_norm) + _assert_local_params_close(offload_model, baseline_model) + assert _local_param_devices(offload_model) == {"cuda"} + assert _optimizer_state_devices(offload_optimizer) == {"cpu"} + + +def test_fsdp2_checkpoint_load_matches_uninterrupted_training_single_node(tmp_path): + model_for_ckpt, ps = _build_fsdp2_model() + optimizer_for_ckpt = _build_optimizer(model_for_ckpt, ps, offload_fraction=0.0) + direct_model, direct_ps = _build_fsdp2_model() + direct_optimizer = _build_optimizer(direct_model, direct_ps, offload_fraction=0.0) + loaded_model, loaded_ps = _build_fsdp2_model() + loaded_optimizer = _build_optimizer(loaded_model, loaded_ps, offload_fraction=0.0) + runtime = MegatronLiteRuntime.__new__(MegatronLiteRuntime) + + torch.manual_seed(4321) + x0 = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + y0 = torch.randn(4, 4, device="cuda", dtype=torch.bfloat16) + x1 = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + y1 = torch.randn(4, 4, device="cuda", dtype=torch.bfloat16) + + _train_step(model_for_ckpt, optimizer_for_ckpt, x0, y0) + _train_step(direct_model, direct_optimizer, x0, y0) + checkpoint_dir = _shared_tmp_path(tmp_path) + + runtime.save_checkpoint( + ModelHandle( + model=model_for_ckpt, + optimizer=optimizer_for_ckpt, + parallel_state=ps, + config=_checkpoint_config(), + _extras={"model_chunks": [model_for_ckpt]}, + ), + checkpoint_dir, + step=1, + ) + assert ( + runtime.load_checkpoint( + ModelHandle( + model=loaded_model, + optimizer=loaded_optimizer, + parallel_state=loaded_ps, + config=_checkpoint_config(), + _extras={"model_chunks": [loaded_model]}, + ), + checkpoint_dir, + ) + == 1 + ) + + _train_step(direct_model, direct_optimizer, x1, y1) + _train_step(loaded_model, loaded_optimizer, x1, y1) + _assert_local_params_close(direct_model, loaded_model) diff --git a/experimental/lite/tests/smoke/primitive/test_parallel_topologies_smoke.py b/experimental/lite/tests/smoke/primitive/test_parallel_topologies_smoke.py new file mode 100644 index 00000000000..3ee1f4e54da --- /dev/null +++ b/experimental/lite/tests/smoke/primitive/test_parallel_topologies_smoke.py @@ -0,0 +1,96 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist + +from megatron.lite.primitive.parallel import init_parallel + +pytestmark = [pytest.mark.mlite, pytest.mark.smoke, pytest.mark.gpu, pytest.mark.distributed] + + +@pytest.fixture(scope="module", autouse=True) +def _single_node_cuda_dist(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for distributed primitive smoke tests.") + if int(os.environ.get("WORLD_SIZE", "1")) > 8: + pytest.skip("Megatron Lite smoke tests are capped at single-node 8 GPUs.") + + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("LOCAL_RANK", "0") + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29511") + + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + created_pg = False + if not dist.is_initialized(): + dist.init_process_group(backend="nccl", init_method="env://") + created_pg = True + yield + if created_pg and dist.is_initialized(): + dist.destroy_process_group() + + +def _assert_group_size(group, expected: int): + if expected == 1: + return + assert group is not None + assert dist.get_world_size(group) == expected + + +def _topologies(world_size: int): + yield "dp_only", SimpleNamespace(tp=1, ep=1, etp=1, cp=1, pp=1) + if world_size >= 2: + yield "tp2", SimpleNamespace(tp=2, ep=1, etp=1, cp=1, pp=1) + yield "cp2", SimpleNamespace(tp=1, ep=1, etp=1, cp=2, pp=1) + yield "pp2", SimpleNamespace(tp=1, ep=1, etp=1, cp=1, pp=2) + yield "ep2", SimpleNamespace(tp=1, ep=2, etp=1, cp=1, pp=1) + yield "etp2", SimpleNamespace(tp=1, ep=1, etp=2, cp=1, pp=1) + if world_size >= 4: + yield "tp2_ep2_pp2", SimpleNamespace(tp=2, ep=2, etp=1, cp=1, pp=2) + if world_size >= 8: + yield "tp2_ep2_cp2_pp2", SimpleNamespace(tp=2, ep=2, etp=1, cp=2, pp=2) + yield "tp2_ep2_etp2_pp2", SimpleNamespace(tp=2, ep=2, etp=2, cp=1, pp=2) + + +def test_parallel_state_builds_expected_primitive_groups(): + world_size = dist.get_world_size() + for name, cfg in _topologies(world_size): + if world_size % (cfg.tp * cfg.cp * cfg.pp) != 0: + continue + if world_size % (cfg.etp * cfg.ep * cfg.pp) != 0: + continue + + ps = init_parallel(cfg) + dense_dp = world_size // (cfg.tp * cfg.cp * cfg.pp) + expert_dp = world_size // (cfg.etp * cfg.ep * cfg.pp) + + assert ps.tp_size == cfg.tp + assert ps.cp_size == cfg.cp + assert ps.pp_size == cfg.pp + assert ps.ep_size == cfg.ep + assert ps.dp_size == dense_dp + assert ps.expert_dp_size == expert_dp + assert 0 <= ps.tp_rank < cfg.tp + assert 0 <= ps.cp_rank < cfg.cp + assert 0 <= ps.pp_rank < cfg.pp + assert 0 <= ps.ep_rank < cfg.ep + _assert_group_size(ps.tp_group, cfg.tp) + _assert_group_size(ps.cp_group, cfg.cp) + _assert_group_size(ps.pp_group, cfg.pp) + _assert_group_size(ps.ep_group, cfg.ep) + _assert_group_size(ps.dp_group, dense_dp) + _assert_group_size(ps.dp_cp_group, dense_dp * cfg.cp) + _assert_group_size(ps.ep_dp_group, expert_dp) + # tp_ep follows Megatron Core's expert tensor + expert model group, + # not the dense-layer TP group. + _assert_group_size(ps.tp_ep_group, cfg.etp * cfg.ep) + _assert_group_size(ps.etp_group, cfg.etp) + if cfg.pp > 1: + assert ps.pp_next_rank in ps.pp_global_ranks + assert ps.pp_prev_rank in ps.pp_global_ranks diff --git a/experimental/lite/tests/smoke/primitive/test_qwen3_moe_distopt_checkpoint_smoke.py b/experimental/lite/tests/smoke/primitive/test_qwen3_moe_distopt_checkpoint_smoke.py new file mode 100644 index 00000000000..38fcde9856b --- /dev/null +++ b/experimental/lite/tests/smoke/primitive/test_qwen3_moe_distopt_checkpoint_smoke.py @@ -0,0 +1,209 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import os +from types import SimpleNamespace +from typing import Any + +import pytest +import torch +import torch.distributed as dist + +from megatron.lite.primitive.deterministic import set_deterministic +from megatron.lite.runtime.backends.mlite.runtime import MegatronLiteRuntime +from megatron.lite.runtime.contracts.config import OptimizerConfig, ParallelConfig +from megatron.lite.runtime.contracts.handle import ModelHandle + +pytestmark = [pytest.mark.mlite, pytest.mark.smoke, pytest.mark.gpu, pytest.mark.distributed] + + +def _qwen3_moe_symbols(): + te = pytest.importorskip( + "transformer_engine.pytorch", + reason="Qwen3MoE distopt checkpoint smoke requires real Transformer Engine.", + ) + assert hasattr(te, "Linear"), "Qwen3MoE smoke requires real Transformer Engine Linear." + from megatron.lite.model.qwen3_moe.config import Qwen3MoEConfig + from megatron.lite.model.qwen3_moe.lite import protocol + + return Qwen3MoEConfig, protocol + + +@pytest.fixture(scope="module", autouse=True) +def _single_node_cuda_dist(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for Qwen3MoE distopt checkpoint smoke tests.") + if int(os.environ.get("WORLD_SIZE", "1")) > 8: + pytest.skip("Megatron Lite smoke tests are capped at single-node 8 GPUs.") + + os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + os.environ.setdefault("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("LOCAL_RANK", "0") + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29541") + + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + created_pg = False + if not dist.is_initialized(): + dist.init_process_group(backend="nccl", init_method="env://") + created_pg = True + yield + try: + from megatron.core import parallel_state as mpu + + if mpu.is_initialized(): + mpu.destroy_model_parallel() + finally: + if created_pg and dist.is_initialized(): + dist.destroy_process_group() + + +def _topology() -> ParallelConfig: + return ParallelConfig(tp=2, ep=2, etp=1, pp=2, cp=1) + + +def _tiny_qwen3_moe_config(): + Qwen3MoEConfig, _protocol = _qwen3_moe_symbols() + return Qwen3MoEConfig( + num_hidden_layers=2, + hidden_size=16, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=4, + vocab_size=64, + num_experts=4, + num_experts_per_tok=1, + moe_intermediate_size=8, + max_position_embeddings=16, + layer_types=["full_attention", "full_attention"], + ) + + +def _build_handle(model_seed: int) -> ModelHandle: + _Qwen3MoEConfig, protocol = _qwen3_moe_symbols() + torch.manual_seed(model_seed) + torch.cuda.manual_seed_all(model_seed) + + parallel = _topology() + model_cfg = _tiny_qwen3_moe_config() + impl_cfg = protocol.ImplConfig( + parallel=parallel, + optimizer="mc", + optimizer_config=OptimizerConfig( + optimizer="adam", lr=1.0e-3, weight_decay=0.0, clip_grad=1.0 + ), + use_deepep=False, + deterministic=True, + ) + bundle = protocol.build_model(model_cfg, impl_cfg=impl_cfg) + extras = dict(bundle.extras) + extras.update( + { + "model_chunks": bundle.chunks, + "forward_step": bundle.forward_step, + "finalize_grads": bundle.finalize_grads, + "protocol": protocol, + } + ) + return ModelHandle( + model=bundle.chunks, + optimizer=bundle.optimizer, + parallel_state=bundle.parallel_state, + config=SimpleNamespace(parallel=parallel), + _extras=extras, + ) + + +def _shared_tmp_path(tmp_path) -> str: + payload = [str(tmp_path) if dist.get_rank() == 0 else None] + dist.broadcast_object_list(payload, src=0) + return payload[0] + + +def _random_batch(vocab_size: int) -> dict[str, torch.Tensor | bool]: + return { + "input_ids": torch.randint(0, vocab_size, (2, 4), device="cuda"), + "labels": torch.randint(0, vocab_size, (2, 4), device="cuda"), + "return_log_probs": False, + } + + +def _clone_batch(batch: dict[str, Any]) -> dict[str, Any]: + return { + key: value.detach().clone() if torch.is_tensor(value) else value + for key, value in batch.items() + } + + +def _assert_batch_equal(actual: dict[str, Any], expected: dict[str, Any]) -> None: + assert actual.keys() == expected.keys() + for key, expected_value in expected.items(): + actual_value = actual[key] + if torch.is_tensor(expected_value): + assert torch.equal(actual_value, expected_value), key + else: + assert actual_value == expected_value + + +def _train_step(runtime: MegatronLiteRuntime, handle: ModelHandle, batch: dict[str, Any]) -> None: + runtime.zero_grad(handle) + runtime.forward_backward(handle, iter([batch]), None, num_microbatches=1) + runtime.optimizer_step(handle) + runtime.zero_grad(handle) + + +def _local_named_params(handle: ModelHandle) -> dict[str, torch.Tensor]: + params: dict[str, torch.Tensor] = {} + for chunk_idx, chunk in enumerate(handle._extras["model_chunks"]): + for name, param in chunk.named_parameters(): + params[f"{chunk_idx}.{name}"] = param.detach().cpu().float().clone() + return params + + +def _assert_params_bitwise_equal(lhs: ModelHandle, rhs: ModelHandle) -> None: + lhs_params = _local_named_params(lhs) + rhs_params = _local_named_params(rhs) + assert lhs_params.keys() == rhs_params.keys() + for name in lhs_params: + torch.testing.assert_close(lhs_params[name], rhs_params[name], atol=0.0, rtol=0.0) + + +def test_qwen3_moe_distopt_checkpoint_restores_rng_and_continues_bitwise_tp2_pp2_ep2(tmp_path): + if dist.get_world_size() != 8: + pytest.skip("Qwen3MoE tp2/pp2/ep2 distopt checkpoint smoke requires exactly 8 GPUs.") + + set_deterministic(2026) + model_cfg = _tiny_qwen3_moe_config() + runtime = MegatronLiteRuntime.__new__(MegatronLiteRuntime) + model_for_ckpt = _build_handle(model_seed=4242) + direct_model = _build_handle(model_seed=4242) + loaded_model = _build_handle(model_seed=4242) + + torch.manual_seed(1357 + dist.get_rank()) + torch.cuda.manual_seed_all(1357 + dist.get_rank()) + step0_batch = _random_batch(model_cfg.vocab_size) + + cpu_rng_before_step0 = torch.get_rng_state() + cuda_rng_before_step0 = torch.cuda.get_rng_state() + _train_step(runtime, model_for_ckpt, step0_batch) + + torch.set_rng_state(cpu_rng_before_step0) + torch.cuda.set_rng_state(cuda_rng_before_step0) + _train_step(runtime, direct_model, _clone_batch(step0_batch)) + _assert_params_bitwise_equal(model_for_ckpt, direct_model) + + checkpoint_dir = _shared_tmp_path(tmp_path) + runtime.save_checkpoint(model_for_ckpt, checkpoint_dir, step=1) + + direct_step1_batch = _random_batch(model_cfg.vocab_size) + expected_step1_batch = _clone_batch(direct_step1_batch) + _train_step(runtime, direct_model, direct_step1_batch) + + assert runtime.load_checkpoint(loaded_model, checkpoint_dir) == 1 + loaded_step1_batch = _random_batch(model_cfg.vocab_size) + _assert_batch_equal(loaded_step1_batch, expected_step1_batch) + _train_step(runtime, loaded_model, loaded_step1_batch) + + _assert_params_bitwise_equal(direct_model, loaded_model) diff --git a/experimental/lite/tests/unit/examples/test_bench_example.py b/experimental/lite/tests/unit/examples/test_bench_example.py new file mode 100644 index 00000000000..b54341786dc --- /dev/null +++ b/experimental/lite/tests/unit/examples/test_bench_example.py @@ -0,0 +1,308 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Unit tests for the benchmark example.""" + +from __future__ import annotations + +import json +import sys +from contextlib import nullcontext +from pathlib import Path + +from megatron.lite.runtime.contracts.config import ParallelConfig +from megatron.lite.runtime.contracts.data import ForwardResult +from megatron.lite.runtime.contracts.handle import ModelHandle + +_LITE_ROOT = str(Path(__file__).resolve().parents[3]) +sys.path = [path for path in sys.path if path != _LITE_ROOT] +sys.path.insert(0, _LITE_ROOT) + + +def test_bench_builds_mlite_runtime_config_with_model_hook(): + from examples.bench.bench import BenchCliConfig, build_runtime_config + from megatron.lite.model.qwen3_5.config import Qwen35Config + from megatron.lite.runtime.backends.mlite.config import MegatronLiteConfig + + cfg = BenchCliConfig( + backend="mlite", + hf_path="/tmp/hf", + model_name="qwen3_5", + use_thd=True, + truncate_layers=2, + disable_mtp=True, + ) + + runtime_cfg = build_runtime_config(cfg) + + assert runtime_cfg.backend == "mlite" + assert isinstance(runtime_cfg.backend_cfg, MegatronLiteConfig) + assert runtime_cfg.backend_cfg.impl_cfg["use_thd"] is True + assert callable(runtime_cfg.backend_cfg.model_config_hook) + + model_cfg = runtime_cfg.backend_cfg.model_config_hook(Qwen35Config()) + assert model_cfg.num_hidden_layers == 2 + assert len(model_cfg.layer_types) == 2 + assert model_cfg.num_nextn_predict_layers == 0 + + +def test_bench_mlite_deterministic_mounts_native_vision_not_mbridge(monkeypatch): + from examples.bench.bench import BenchCliConfig, build_runtime_config + + monkeypatch.setenv("MEGATRON_LITE_DETERMINISTIC", "1") + + runtime_cfg = build_runtime_config( + BenchCliConfig(backend="mlite", hf_path="/tmp/hf", model_name="qwen3_5") + ) + + impl_cfg = runtime_cfg.backend_cfg.impl_cfg + assert impl_cfg["mount_vision_model"] is True + assert ("mount_" + "mbridge_vision_model") not in impl_cfg + + +def test_bench_builds_bridge_dry_run_plan_without_bridge_import(): + from examples.bench.bench import BenchCliConfig, build_dry_run_plan + + plan = build_dry_run_plan( + BenchCliConfig( + backend="bridge", + hf_path="/tmp/hf", + model_name="qwen3_5", + truncate_layers=2, + override_transformer_json='{"attention_backend": "unfused"}', + dry_run=True, + ) + ) + + assert plan["dry_run"] is True + assert plan["runtime"]["backend"] == "bridge" + backend_cfg = plan["runtime"]["backend_cfg"] + assert backend_cfg["model_name"] == "qwen3_5" + assert backend_cfg["override_transformer_config"] == {"attention_backend": "unfused"} + assert backend_cfg["bridge_post_init"].startswith(" None: + pass + + def forward_backward(self, handle, data, loss_fn, *, num_microbatches: int = 1): + self.loss += 1 + return ForwardResult(metrics={"loss": float(self.loss)}) + + def optimizer_step(self, handle): + return True, 3.5, 0 + + def lr_scheduler_step(self, handle): + return 0.0 + + +def test_pretrain_session_runs_with_fake_runtime_on_cpu(): + from examples.bench.session import PretrainSessionConfig, run_pretrain_session + + handle = ModelHandle( + model=object(), + optimizer=object(), + parallel_state=None, + config=type( + "Cfg", (), {"model_name": "fake", "impl": "lite", "parallel": ParallelConfig()} + )(), + _extras={"optimizer_backend": "fake"}, + ) + + result = run_pretrain_session( + _FakeRuntime(), + handle, + PretrainSessionConfig(steps=3, warmup=1, device="cpu", seq_len=4), + data_iter=iter([{}, {}, {}]), + ) + + assert result.backend == "mlite" + assert result.seq_len == 4 + assert result.num_microbatches == 1 + assert len(result.step_traces) == 2 + assert [trace.loss for trace in result.step_traces] == [2.0, 3.0] + assert result.step_traces[0].grad_norm == 3.5 + + +def test_bench_main_writes_dry_run_output_json(tmp_path): + from examples.bench.bench import main + + output_path = tmp_path / "dry_run.json" + + artifact = main( + [ + "--backend", + "mlite", + "--hf-path", + "/tmp/hf", + "--model-name", + "qwen3_5", + "--truncate-layers", + "2", + "--disable-mtp", + "--dry-run", + "--output-json", + str(output_path), + ] + ) + + assert output_path.exists() + assert json.loads(output_path.read_text()) == artifact + + +def test_bench_main_writes_output_json_only_on_rank_zero(tmp_path, monkeypatch): + from examples.bench.bench import main + + output_path = tmp_path / "rank_one.json" + monkeypatch.setenv("RANK", "1") + + artifact = main( + [ + "--backend", + "mlite", + "--hf-path", + "/tmp/hf", + "--model-name", + "qwen3_5", + "--dry-run", + "--output-json", + str(output_path), + ] + ) + + assert artifact["dry_run"] is True + assert not output_path.exists() + + +def test_result_artifact_summary_and_trace_compare(tmp_path): + from examples.bench.results import compare_step_traces, load_result_artifact, result_summary + + baseline = { + "summary": { + "backend": "mlite", + "avg_step_ms": 10.0, + "tok_per_s": 3200.0, + "steps_measured": 2, + }, + "result": { + "step_traces": [ + {"step": 0, "loss": 1.0, "grad_norm": 2.0, "step_ms": 10.0}, + {"step": 1, "loss": 1.5, "grad_norm": 2.5, "step_ms": 10.0}, + ] + }, + } + candidate = { + "summary": { + "backend": "bridge", + "avg_step_ms": 11.0, + "tok_per_s": 2900.0, + "steps_measured": 2, + }, + "result": { + "step_traces": [ + {"step": 0, "loss": 1.00001, "grad_norm": 2.00001, "step_ms": 11.0}, + {"step": 1, "loss": 1.49999, "grad_norm": 2.49999, "step_ms": 11.0}, + ] + }, + } + baseline_path = tmp_path / "mlite.json" + baseline_path.write_text(json.dumps(baseline), encoding="utf-8") + + loaded = load_result_artifact(baseline_path) + + assert result_summary(loaded)["backend"] == "mlite" + assert compare_step_traces(baseline, candidate, atol=1e-3, rtol=0.0)["passed"] is True + + +def test_result_trace_compare_reports_metric_level_failures(): + from examples.bench.results import compare_step_traces + + baseline = { + "result": {"step_traces": [{"step": 0, "loss": 1.0, "grad_norm": 2.0, "step_ms": 10.0}]} + } + candidate = { + "result": {"step_traces": [{"step": 0, "loss": 1.00001, "grad_norm": 3.0, "step_ms": 10.0}]} + } + + comparison = compare_step_traces(baseline, candidate, atol=1e-3, rtol=0.0) + + assert comparison["passed"] is False + assert comparison["loss_passed"] is True + assert comparison["grad_norm_passed"] is False + + +def test_correctness_compare_requires_bitwise_fields(): + from examples.bench.results import compare_correctness_artifacts + + baseline = { + "eval_logits": {"sha256": "a", "shape": [1], "dtype": "torch.bfloat16"}, + "steps": [ + { + "loss": {"value": 1.0, "float_hex": (1.0).hex()}, + "logits": {"sha256": "b"}, + "grad_fingerprint": {"sha256": "c", "tensor_count": 1}, + "grad_norm": {"value": 2.0, "float_hex": (2.0).hex()}, + "update_successful": True, + "num_zeros": 0, + "post_step_weights": {"sha256": "d", "tensor_count": 1}, + } + ], + } + candidate = json.loads(json.dumps(baseline)) + + assert compare_correctness_artifacts(baseline, candidate)["passed"] is True + + candidate["steps"][0]["grad_norm"] = {"value": 2.5, "float_hex": (2.5).hex()} + comparison = compare_correctness_artifacts(baseline, candidate) + + assert comparison["passed"] is False + assert comparison["max_grad_norm_abs"] == 0.5 diff --git a/experimental/lite/tests/unit/model/test_qwen35_export.py b/experimental/lite/tests/unit/model/test_qwen35_export.py new file mode 100644 index 00000000000..698e6137e16 --- /dev/null +++ b/experimental/lite/tests/unit/model/test_qwen35_export.py @@ -0,0 +1,721 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from types import SimpleNamespace + +import torch +import torch.nn as nn + +from megatron.lite.model.qwen3_5.config import Qwen35Config +from megatron.lite.model.qwen3_5.lite.checkpoint import ( + Qwen35WeightSpec, + _merge_full_attn_qkvg, + _merge_gate_up_tp_shards, + _merge_linear_attn_conv1d_tp_shards, + _merge_linear_attn_in_proj_tp_shards, + export_hf_weights, +) +from megatron.lite.model.registry import TRAIN_RUNTIME_MODULES, resolve_runtime_model_name + + +def _tiny_config() -> Qwen35Config: + return Qwen35Config( + num_hidden_layers=1, + hidden_size=8, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=2, + vocab_size=16, + num_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=4, + shared_expert_intermediate_size=4, + linear_num_key_heads=2, + linear_key_head_dim=2, + linear_num_value_heads=2, + linear_value_head_dim=2, + linear_conv_kernel_dim=2, + layer_types=["full_attention"], + partial_rotary_factor=1.0, + ) + + +def _single_rank_parallel_state() -> SimpleNamespace: + return SimpleNamespace( + pp_size=1, tp_size=1, tp_group=None, ep_size=1, ep_group=None, etp_size=1, etp_group=None + ) + + +def test_qwen35_protocol_registers_vllm_export_entrypoint() -> None: + key = resolve_runtime_model_name("qwen3_5", "lite") + module = __import__(TRAIN_RUNTIME_MODULES[key], fromlist=["export_hf_weights"]) + + assert key == "qwen3_5" + assert callable(module.export_hf_weights) + + +def test_qwen35_export_uses_hf_checkpoint_names_without_module_prefix() -> None: + class TinyQwen35Module(nn.Module): + def __init__(self) -> None: + super().__init__() + self.embed = nn.Module() + self.embed.embedding = nn.Embedding(16, 8) + self.norm = nn.LayerNorm(8) + self.head = nn.Module() + self.head.col = nn.Module() + self.head.col.linear = nn.Linear(8, 16, bias=False) + + cfg = _tiny_config() + model = TinyQwen35Module() + + exported = dict(export_hf_weights(model, cfg, _single_rank_parallel_state())) + + assert set(exported) == { + "model.language_model.embed_tokens.weight", + "model.language_model.norm.weight", + "lm_head.weight", + } + assert all(not name.startswith("module.") for name in exported) + assert all(not name.startswith(("embed.", "norm.", "head.")) for name in exported) + + +def test_qwen35_export_dtype_cast_is_opt_in() -> None: + class TinyQwen35Module(nn.Module): + def __init__(self, config: Qwen35Config) -> None: + super().__init__() + self.norm = nn.LayerNorm(8) + self.layers = nn.ModuleList([nn.Module()]) + self.layers[0].moe = nn.Module() + self.layers[0].moe.experts = nn.Module() + self.layers[0].moe.experts.fc1 = nn.Module() + + rows = config.moe_intermediate_size * 2 + for expert_idx in range(config.num_experts): + tensor = torch.arange(rows * config.hidden_size, dtype=torch.float32).reshape( + rows, config.hidden_size + ) + tensor = tensor + expert_idx * 1000 + self.layers[0].moe.experts.fc1.register_parameter( + f"weight{expert_idx}", nn.Parameter(tensor) + ) + + cfg = _tiny_config() + model = TinyQwen35Module(cfg) + + default_export = dict(export_hf_weights(model, cfg, _single_rank_parallel_state())) + bf16_export = dict( + export_hf_weights(model, cfg, _single_rank_parallel_state(), export_dtype="bfloat16") + ) + + assert default_export["model.language_model.norm.weight"].dtype == torch.float32 + assert bf16_export["model.language_model.norm.weight"].dtype == torch.bfloat16 + assert ( + default_export["model.language_model.layers.0.mlp.experts.gate_up_proj"].dtype + == torch.float32 + ) + assert ( + bf16_export["model.language_model.layers.0.mlp.experts.gate_up_proj"].dtype + == torch.bfloat16 + ) + + +def test_qwen35_export_preserves_runtime_parameter_dtype_by_default() -> None: + class TinyQwen35Module(nn.Module): + def __init__(self, config: Qwen35Config) -> None: + super().__init__() + self.norm = nn.LayerNorm(8).to(torch.bfloat16) + self.layers = nn.ModuleList([nn.Module()]) + self.layers[0].moe = nn.Module() + self.layers[0].moe.experts = nn.Module() + self.layers[0].moe.experts.fc1 = nn.Module() + + rows = config.moe_intermediate_size * 2 + for expert_idx in range(config.num_experts): + tensor = torch.arange(rows * config.hidden_size, dtype=torch.bfloat16).reshape( + rows, config.hidden_size + ) + tensor = tensor + expert_idx * 1000 + self.layers[0].moe.experts.fc1.register_parameter( + f"weight{expert_idx}", nn.Parameter(tensor) + ) + + cfg = _tiny_config() + model = TinyQwen35Module(cfg) + + exported = dict(export_hf_weights(model, cfg, _single_rank_parallel_state())) + + assert exported["model.language_model.norm.weight"].dtype == torch.bfloat16 + assert ( + exported["model.language_model.layers.0.mlp.experts.gate_up_proj"].dtype == torch.bfloat16 + ) + + +def test_qwen35_export_batches_ep_expert_gather(monkeypatch) -> None: + class TinyQwen35Module(nn.Module): + def __init__(self, config: Qwen35Config) -> None: + super().__init__() + self.layers = nn.ModuleList([nn.Module()]) + self.layers[0].moe = nn.Module() + self.layers[0].moe.experts = nn.Module() + self.layers[0].moe.experts.fc1 = nn.Module() + + rows = config.moe_intermediate_size * 2 + for local_idx in range(config.num_experts // 2): + tensor = torch.arange(rows * config.hidden_size, dtype=torch.bfloat16).reshape( + rows, config.hidden_size + ) + tensor = tensor + local_idx * 1000 + self.layers[0].moe.experts.fc1.register_parameter( + f"weight{local_idx}", nn.Parameter(tensor) + ) + + cfg = _tiny_config() + model = TinyQwen35Module(cfg) + ps = SimpleNamespace( + pp_size=1, + tp_size=1, + tp_group=None, + ep_size=2, + ep_group=object(), + etp_size=1, + etp_group=None, + ) + gather_calls = [] + + def fake_all_gather(outputs, tensor, group=None): + del group + gather_calls.append(tensor.clone()) + outputs[0].copy_(tensor) + outputs[1].copy_(tensor + 2000) + + monkeypatch.setattr("megatron.lite.primitive.ckpt.hf_weights.dist.all_gather", fake_all_gather) + + exported = dict(export_hf_weights(model, cfg, ps)) + + assert len(gather_calls) == 1 + assert gather_calls[0].shape[0] == cfg.num_experts // ps.ep_size + local_tensors = [ + model.layers[0].moe.experts.fc1.weight0.detach(), + model.layers[0].moe.experts.fc1.weight1.detach(), + ] + expected = torch.stack( + [local_tensors[0], local_tensors[1], local_tensors[0] + 2000, local_tensors[1] + 2000], + dim=0, + ) + assert torch.equal(exported["model.language_model.layers.0.mlp.experts.gate_up_proj"], expected) + + +def test_qwen35_export_uses_packed_expert_group_names(monkeypatch) -> None: + class TinyQwen35Module(nn.Module): + def __init__(self, config: Qwen35Config) -> None: + super().__init__() + self.layers = nn.ModuleList([nn.Module()]) + self.layers[0].moe = nn.Module() + self.layers[0].moe.experts = nn.Module() + self.layers[0].moe.experts.fc1 = nn.Module() + + rows = config.moe_intermediate_size * 2 + for expert_idx in range(config.num_experts): + tensor = torch.arange(rows * config.hidden_size, dtype=torch.bfloat16).reshape( + rows, config.hidden_size + ) + tensor = tensor + expert_idx * 1000 + self.layers[0].moe.experts.fc1.register_parameter( + f"weight{expert_idx}", nn.Parameter(tensor) + ) + + cfg = _tiny_config() + seen_native_names = [] + original = Qwen35WeightSpec.native_to_hf + + def spy_native_to_hf(self, native_name, tensor): + seen_native_names.append(native_name) + return original(self, native_name, tensor) + + monkeypatch.setattr(Qwen35WeightSpec, "native_to_hf", spy_native_to_hf) + + exported = dict(export_hf_weights(TinyQwen35Module(cfg), cfg, _single_rank_parallel_state())) + + assert seen_native_names == ["layers.0.moe.experts.fc1.packed"] + assert set(exported) == {"model.language_model.layers.0.mlp.experts.gate_up_proj"} + + +def test_qwen35_export_rank0_only_still_participates_in_ep_gather(monkeypatch) -> None: + class TinyQwen35Module(nn.Module): + def __init__(self, config: Qwen35Config) -> None: + super().__init__() + self.layers = nn.ModuleList([nn.Module()]) + self.layers[0].moe = nn.Module() + self.layers[0].moe.experts = nn.Module() + self.layers[0].moe.experts.fc1 = nn.Module() + + rows = config.moe_intermediate_size * 2 + for local_idx in range(config.num_experts // 2): + tensor = torch.zeros(rows, config.hidden_size, dtype=torch.bfloat16) + local_idx + self.layers[0].moe.experts.fc1.register_parameter( + f"weight{local_idx}", nn.Parameter(tensor) + ) + + cfg = _tiny_config() + ps = SimpleNamespace( + pp_size=1, + tp_size=1, + tp_group=None, + ep_size=2, + ep_group=object(), + etp_size=1, + etp_group=None, + ) + gather_calls = [] + + def fake_all_gather(outputs, tensor, group=None): + del group + gather_calls.append(tensor.clone()) + outputs[0].copy_(tensor) + outputs[1].copy_(tensor + 2) + + monkeypatch.setattr("megatron.lite.primitive.ckpt.hf_weights.dist.is_initialized", lambda: True) + monkeypatch.setattr("megatron.lite.primitive.ckpt.hf_weights.dist.get_rank", lambda: 1) + monkeypatch.setattr("megatron.lite.primitive.ckpt.hf_weights.dist.all_gather", fake_all_gather) + + exported = list(export_hf_weights(TinyQwen35Module(cfg), cfg, ps, rank0_only=True)) + + assert exported == [] + assert len(gather_calls) == 1 + + +def test_qwen35_export_maps_top_level_and_layer_norm_names() -> None: + cfg = _tiny_config() + spec = Qwen35WeightSpec(cfg) + tensor = torch.arange(cfg.hidden_size) + + cases = { + "embed.embedding.weight": "model.language_model.embed_tokens.weight", + "norm.weight": "model.language_model.norm.weight", + "head.col.linear.weight": "lm_head.weight", + "layers.0.full_attn.qkv.linear.layer_norm_weight": ( + "model.language_model.layers.0.input_layernorm.weight" + ), + "layers.0.mlp_norm.weight": "model.language_model.layers.0.post_attention_layernorm.weight", + } + + for native_name, hf_name in cases.items(): + exported = dict(spec.native_to_hf(native_name, tensor)) + assert set(exported) == {hf_name} + assert torch.equal(exported[hf_name], tensor) + + +def test_qwen35_export_unpacks_full_attention_q_gate() -> None: + cfg = _tiny_config() + spec = Qwen35WeightSpec(cfg) + hidden = cfg.hidden_size + q_gate = torch.arange(cfg.num_attention_heads * 2 * cfg.head_dim * hidden).reshape(-1, hidden) + key = torch.arange( + q_gate.numel(), q_gate.numel() + cfg.num_key_value_heads * cfg.head_dim * hidden + ).reshape(-1, hidden) + value = torch.arange( + key[-1, -1] + 1, key[-1, -1] + 1 + cfg.num_key_value_heads * cfg.head_dim * hidden + ).reshape(-1, hidden) + + packed = _merge_full_attn_qkvg(q_gate, key, value, cfg=cfg) + exported = dict(spec.native_to_hf("layers.0.full_attn.qkv.linear.weight", packed)) + + assert set(exported) == { + "model.language_model.layers.0.self_attn.q_proj.weight", + "model.language_model.layers.0.self_attn.k_proj.weight", + "model.language_model.layers.0.self_attn.v_proj.weight", + } + assert torch.equal(exported["model.language_model.layers.0.self_attn.q_proj.weight"], q_gate) + assert torch.equal(exported["model.language_model.layers.0.self_attn.k_proj.weight"], key) + assert torch.equal(exported["model.language_model.layers.0.self_attn.v_proj.weight"], value) + + +def test_qwen35_export_maps_linear_attention_to_hf_checkpoint_names() -> None: + cfg = _tiny_config() + spec = Qwen35WeightSpec(cfg) + qk_dim = cfg.linear_num_key_heads * cfg.linear_key_head_dim + v_dim = cfg.linear_num_value_heads * cfg.linear_value_head_dim + rows = qk_dim * 2 + v_dim * 2 + cfg.linear_num_value_heads * 2 + tensor = torch.arange(rows * cfg.hidden_size).reshape(rows, cfg.hidden_size) + + exported = dict(spec.native_to_hf("layers.0.linear_attn.in_proj.linear.weight", tensor)) + + assert set(exported) == { + "model.language_model.layers.0.linear_attn.in_proj_qkv.weight", + "model.language_model.layers.0.linear_attn.in_proj_z.weight", + "model.language_model.layers.0.linear_attn.in_proj_b.weight", + "model.language_model.layers.0.linear_attn.in_proj_a.weight", + } + assert ( + exported["model.language_model.layers.0.linear_attn.in_proj_qkv.weight"].shape[0] + == qk_dim * 2 + v_dim + ) + assert exported["model.language_model.layers.0.linear_attn.in_proj_z.weight"].shape[0] == v_dim + assert ( + exported["model.language_model.layers.0.linear_attn.in_proj_b.weight"].shape[0] + == cfg.linear_num_value_heads + ) + assert ( + exported["model.language_model.layers.0.linear_attn.in_proj_a.weight"].shape[0] + == cfg.linear_num_value_heads + ) + + +def test_qwen35_export_reorders_linear_attention_tp_shards_before_hf_split() -> None: + cfg = _tiny_config() + qk_dim = cfg.linear_num_key_heads * cfg.linear_key_head_dim + v_dim = cfg.linear_num_value_heads * cfg.linear_value_head_dim + hidden = cfg.hidden_size + parts = [ + torch.arange(0, qk_dim * hidden).reshape(qk_dim, hidden), + torch.arange(100, 100 + qk_dim * hidden).reshape(qk_dim, hidden), + torch.arange(200, 200 + v_dim * hidden).reshape(v_dim, hidden), + torch.arange(300, 300 + v_dim * hidden).reshape(v_dim, hidden), + torch.arange(400, 400 + cfg.linear_num_value_heads * hidden).reshape( + cfg.linear_num_value_heads, hidden + ), + torch.arange(500, 500 + cfg.linear_num_value_heads * hidden).reshape( + cfg.linear_num_value_heads, hidden + ), + ] + full = torch.cat(parts, dim=0) + shards = [torch.cat([part.chunk(2, dim=0)[rank] for part in parts], dim=0) for rank in range(2)] + + merged = _merge_linear_attn_in_proj_tp_shards(shards, cfg=cfg) + + assert torch.equal(merged, full) + + +def test_qwen35_export_reorders_linear_attention_conv1d_tp_shards() -> None: + cfg = _tiny_config() + qk_dim = cfg.linear_num_key_heads * cfg.linear_key_head_dim + v_dim = cfg.linear_num_value_heads * cfg.linear_value_head_dim + trailing = (1, cfg.linear_conv_kernel_dim) + parts = [ + torch.arange(0, qk_dim * trailing[0] * trailing[1], dtype=torch.float32).reshape( + qk_dim, *trailing + ), + torch.arange(100, 100 + qk_dim * trailing[0] * trailing[1], dtype=torch.float32).reshape( + qk_dim, *trailing + ), + torch.arange(200, 200 + v_dim * trailing[0] * trailing[1], dtype=torch.float32).reshape( + v_dim, *trailing + ), + ] + full = torch.cat(parts, dim=0) + shards = [torch.cat([part.chunk(2, dim=0)[rank] for part in parts], dim=0) for rank in range(2)] + + merged = _merge_linear_attn_conv1d_tp_shards(shards, cfg=cfg) + + assert torch.equal(merged, full) + + +def test_qwen35_export_uses_mbridge_conv1d_tp_gather(monkeypatch) -> None: + class TinyQwen35Module(nn.Module): + def __init__(self, local_shard: torch.Tensor) -> None: + super().__init__() + self.layers = nn.ModuleList([nn.Module()]) + self.layers[0].linear_attn = nn.Module() + self.layers[0].linear_attn.conv1d = nn.Module() + self.layers[0].linear_attn.conv1d.register_parameter( + "weight", nn.Parameter(local_shard.clone()) + ) + + cfg = _tiny_config() + qk_dim = cfg.linear_num_key_heads * cfg.linear_key_head_dim + v_dim = cfg.linear_num_value_heads * cfg.linear_value_head_dim + trailing = (1, cfg.linear_conv_kernel_dim) + parts = [ + torch.arange(0, qk_dim * trailing[0] * trailing[1], dtype=torch.float32).reshape( + qk_dim, *trailing + ), + torch.arange(100, 100 + qk_dim * trailing[0] * trailing[1], dtype=torch.float32).reshape( + qk_dim, *trailing + ), + torch.arange(200, 200 + v_dim * trailing[0] * trailing[1], dtype=torch.float32).reshape( + v_dim, *trailing + ), + ] + full = torch.cat(parts, dim=0) + shards = [torch.cat([part.chunk(2, dim=0)[rank] for part in parts], dim=0) for rank in range(2)] + ps = SimpleNamespace( + pp_size=1, + tp_size=2, + tp_group=object(), + ep_size=1, + ep_group=None, + etp_size=1, + etp_group=None, + ) + gather_calls = [] + + def fake_all_gather(outputs, tensor, group=None): + assert group is ps.tp_group + gather_calls.append(tensor.clone()) + outputs[0].copy_(shards[0]) + outputs[1].copy_(shards[1]) + + monkeypatch.setattr( + "megatron.lite.model.qwen3_5.lite.checkpoint.dist.all_gather", fake_all_gather + ) + + exported = dict(export_hf_weights(TinyQwen35Module(shards[0]), cfg, ps)) + + assert len(gather_calls) == 1 + assert torch.equal(gather_calls[0], shards[0]) + assert torch.equal(exported["model.language_model.layers.0.linear_attn.conv1d.weight"], full) + + +def test_qwen35_export_reorders_shared_expert_gate_up_tp_shards() -> None: + gate = torch.arange(0, 32).reshape(4, 8) + up = torch.arange(100, 132).reshape(4, 8) + full = torch.cat([gate, up], dim=0) + shards = [ + torch.cat([gate.chunk(2, dim=0)[rank], up.chunk(2, dim=0)[rank]], dim=0) + for rank in range(2) + ] + + merged = _merge_gate_up_tp_shards(shards) + + assert torch.equal(merged, full) + + +def test_qwen35_export_restores_zero_centered_linear_attention_norm() -> None: + cfg = _tiny_config() + spec = Qwen35WeightSpec(cfg) + tensor = torch.tensor([-0.5, 0.0, 0.5]) + + exported = dict(spec.native_to_hf("layers.0.linear_attn.norm.weight", tensor)) + + assert set(exported) == {"model.language_model.layers.0.linear_attn.norm.weight"} + assert torch.equal( + exported["model.language_model.layers.0.linear_attn.norm.weight"], tensor + 1 + ) + + +def test_qwen35_export_maps_shared_expert_to_hf_checkpoint_names() -> None: + cfg = _tiny_config() + spec = Qwen35WeightSpec(cfg) + tensor = torch.arange(cfg.shared_expert_intermediate_size * 2 * cfg.hidden_size).reshape( + -1, cfg.hidden_size + ) + + exported = dict(spec.native_to_hf("layers.0.moe.shared_expert.gate_up.linear.weight", tensor)) + + assert set(exported) == { + "model.language_model.layers.0.mlp.shared_expert.gate_proj.weight", + "model.language_model.layers.0.mlp.shared_expert.up_proj.weight", + } + gate, up = tensor.chunk(2, dim=0) + assert torch.equal( + exported["model.language_model.layers.0.mlp.shared_expert.gate_proj.weight"], gate + ) + assert torch.equal( + exported["model.language_model.layers.0.mlp.shared_expert.up_proj.weight"], up + ) + + +def test_qwen35_export_packs_base_expert_fc1_to_hf_gate_up_proj() -> None: + cfg = _tiny_config() + spec = Qwen35WeightSpec(cfg) + base = torch.arange(cfg.moe_intermediate_size * 2 * cfg.hidden_size).reshape( + -1, cfg.hidden_size + ) + + exported = {} + expert_tensors = [] + for expert_idx in range(cfg.num_experts): + tensor = base + expert_idx * 1000 + expert_tensors.append(tensor) + exported.update( + dict(spec.native_to_hf(f"layers.0.moe.experts.fc1.weight{expert_idx}", tensor)) + ) + + assert set(exported) == {"model.language_model.layers.0.mlp.experts.gate_up_proj"} + assert torch.equal( + exported["model.language_model.layers.0.mlp.experts.gate_up_proj"], + torch.stack(expert_tensors, dim=0), + ) + + +def test_qwen35_export_matches_mbridge_qwen35_moe_packed_expert_contract() -> None: + cfg = _tiny_config() + spec = Qwen35WeightSpec(cfg) + rows = cfg.moe_intermediate_size * 2 + fc1_tensors = [ + torch.arange(rows * cfg.hidden_size, dtype=torch.bfloat16).reshape(rows, cfg.hidden_size) + + expert_idx * 1000 + for expert_idx in range(cfg.num_experts) + ] + fc2_tensors = [ + torch.arange(cfg.hidden_size * cfg.moe_intermediate_size, dtype=torch.bfloat16).reshape( + cfg.hidden_size, cfg.moe_intermediate_size + ) + + expert_idx * 1000 + for expert_idx in range(cfg.num_experts) + ] + + fc1_exported = {} + fc2_exported = {} + for expert_idx, (fc1, fc2) in enumerate(zip(fc1_tensors, fc2_tensors, strict=True)): + fc1_exported.update( + dict(spec.native_to_hf(f"layers.0.moe.experts.fc1.weight{expert_idx}", fc1)) + ) + fc2_exported.update( + dict(spec.native_to_hf(f"layers.0.moe.experts.fc2.weight{expert_idx}", fc2)) + ) + + assert set(fc1_exported) == {"model.language_model.layers.0.mlp.experts.gate_up_proj"} + assert set(fc2_exported) == {"model.language_model.layers.0.mlp.experts.down_proj"} + assert torch.equal( + fc1_exported["model.language_model.layers.0.mlp.experts.gate_up_proj"], + torch.stack(fc1_tensors, dim=0), + ) + assert torch.equal( + fc2_exported["model.language_model.layers.0.mlp.experts.down_proj"], + torch.stack(fc2_tensors, dim=0), + ) + + +def test_qwen35_export_vllm_target_uses_runtime_prefix_and_packed_expert_names() -> None: + cfg = _tiny_config() + spec = Qwen35WeightSpec(cfg, target="vllm") + dense = torch.arange(cfg.hidden_size) + + exported_embed = dict(spec.native_to_hf("embed.embedding.weight", dense)) + exported_norm = dict(spec.native_to_hf("norm.weight", dense)) + exported_head = dict(spec.native_to_hf("head.col.linear.weight", dense)) + exported_mlp_norm = dict(spec.native_to_hf("layers.0.mlp_norm.weight", dense)) + assert set(exported_embed) == {"language_model.model.embed_tokens.weight"} + assert set(exported_norm) == {"language_model.model.norm.weight"} + assert set(exported_head) == {"language_model.lm_head.weight"} + assert set(exported_mlp_norm) == { + "language_model.model.layers.0.post_attention_layernorm.weight" + } + assert torch.equal(exported_embed["language_model.model.embed_tokens.weight"], dense) + assert torch.equal(exported_norm["language_model.model.norm.weight"], dense) + assert torch.equal(exported_head["language_model.lm_head.weight"], dense) + assert torch.equal( + exported_mlp_norm["language_model.model.layers.0.post_attention_layernorm.weight"], dense + ) + + fc1_tensors = [ + torch.arange(cfg.moe_intermediate_size * 2 * cfg.hidden_size, dtype=torch.bfloat16).reshape( + -1, cfg.hidden_size + ) + + expert_idx * 1000 + for expert_idx in range(cfg.num_experts) + ] + fc2_tensors = [ + torch.arange(cfg.hidden_size * cfg.moe_intermediate_size, dtype=torch.bfloat16).reshape( + cfg.hidden_size, cfg.moe_intermediate_size + ) + + expert_idx * 2000 + for expert_idx in range(cfg.num_experts) + ] + + fc1_exported = {} + fc2_exported = {} + for expert_idx, (fc1, fc2) in enumerate(zip(fc1_tensors, fc2_tensors, strict=True)): + fc1_exported.update( + dict(spec.native_to_hf(f"layers.0.moe.experts.fc1.weight{expert_idx}", fc1)) + ) + fc2_exported.update( + dict(spec.native_to_hf(f"layers.0.moe.experts.fc2.weight{expert_idx}", fc2)) + ) + + assert set(fc1_exported) == {"language_model.model.layers.0.mlp.experts.gate_up_proj"} + assert set(fc2_exported) == {"language_model.model.layers.0.mlp.experts.down_proj"} + assert torch.equal( + fc1_exported["language_model.model.layers.0.mlp.experts.gate_up_proj"], + torch.stack(fc1_tensors, dim=0), + ) + assert torch.equal( + fc2_exported["language_model.model.layers.0.mlp.experts.down_proj"], + torch.stack(fc2_tensors, dim=0), + ) + + +def test_qwen35_export_vllm_target_packs_experts_with_runtime_prefix() -> None: + class TinyQwen35Module(nn.Module): + def __init__(self, config: Qwen35Config) -> None: + super().__init__() + self.layers = nn.ModuleList([nn.Module()]) + self.layers[0].moe = nn.Module() + self.layers[0].moe.experts = nn.Module() + self.layers[0].moe.experts.fc1 = nn.Module() + self.layers[0].moe.experts.fc2 = nn.Module() + + rows = config.moe_intermediate_size * 2 + for expert_idx in range(config.num_experts): + fc1 = torch.arange(rows * config.hidden_size, dtype=torch.bfloat16).reshape( + rows, config.hidden_size + ) + fc1 = fc1 + expert_idx * 1000 + fc2 = torch.arange( + config.hidden_size * config.moe_intermediate_size, dtype=torch.bfloat16 + ).reshape(config.hidden_size, config.moe_intermediate_size) + fc2 = fc2 + expert_idx * 2000 + self.layers[0].moe.experts.fc1.register_parameter( + f"weight{expert_idx}", nn.Parameter(fc1) + ) + self.layers[0].moe.experts.fc2.register_parameter( + f"weight{expert_idx}", nn.Parameter(fc2) + ) + + cfg = _tiny_config() + model = TinyQwen35Module(cfg) + + exported = dict(export_hf_weights(model, cfg, _single_rank_parallel_state(), target="vllm")) + + assert "model.language_model.layers.0.mlp.experts.gate_up_proj" not in exported + assert "model.language_model.layers.0.mlp.experts.down_proj" not in exported + assert "language_model.model.layers.0.mlp.experts.0.gate_proj.weight" not in exported + assert "language_model.model.layers.0.mlp.experts.0.down_proj.weight" not in exported + assert set(exported) == { + "language_model.model.layers.0.mlp.experts.gate_up_proj", + "language_model.model.layers.0.mlp.experts.down_proj", + } + + expected_fc1 = [] + expected_fc2 = [] + for expert_idx in range(cfg.num_experts): + fc1 = getattr(model.layers[0].moe.experts.fc1, f"weight{expert_idx}").detach() + fc2 = getattr(model.layers[0].moe.experts.fc2, f"weight{expert_idx}").detach() + expected_fc1.append(fc1) + expected_fc2.append(fc2) + + assert torch.equal( + exported["language_model.model.layers.0.mlp.experts.gate_up_proj"], + torch.stack(expected_fc1, dim=0), + ) + assert torch.equal( + exported["language_model.model.layers.0.mlp.experts.down_proj"], + torch.stack(expected_fc2, dim=0), + ) + + +def test_qwen35_export_packs_base_expert_fc2_and_expert_metadata() -> None: + cfg = _tiny_config() + spec = Qwen35WeightSpec(cfg) + base = torch.arange(cfg.hidden_size * cfg.moe_intermediate_size).reshape( + cfg.hidden_size, cfg.moe_intermediate_size + ) + native_name = "layers.0.moe.experts.fc2.weight2" + + exported = {} + expert_tensors = [] + for expert_idx in range(cfg.num_experts): + tensor = base + expert_idx * 1000 + expert_tensors.append(tensor) + exported.update( + dict(spec.native_to_hf(f"layers.0.moe.experts.fc2.weight{expert_idx}", tensor)) + ) + + assert set(exported) == {"model.language_model.layers.0.mlp.experts.down_proj"} + assert torch.equal( + exported["model.language_model.layers.0.mlp.experts.down_proj"], + torch.stack(expert_tensors, dim=0), + ) + assert spec.is_expert(native_name) + assert spec.expert_global_id(native_name) == 2 + assert spec.expert_local_name(native_name, 0) == "layers.0.moe.experts.fc2.weight0" + assert spec.tp_spec(native_name) == (1, 1) diff --git a/experimental/lite/tests/unit/model/test_qwen_config_unit.py b/experimental/lite/tests/unit/model/test_qwen_config_unit.py new file mode 100644 index 00000000000..4e660bbb646 --- /dev/null +++ b/experimental/lite/tests/unit/model/test_qwen_config_unit.py @@ -0,0 +1,146 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest + +from megatron.lite.model.qwen3_5.config import Qwen35Config +from megatron.lite.model.qwen3_moe.config import Qwen3MoEConfig +from megatron.lite.model.registry import resolve_model_type_from_hf, resolve_runtime_model_name + +pytestmark = pytest.mark.mlite + +LITE_ROOT = Path(__file__).resolve().parents[3] + + +def _tiny_qwen3_hf_dict() -> dict: + return { + "model_type": "qwen3_moe", + "hidden_size": 16, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "num_hidden_layers": 1, + "vocab_size": 64, + "num_experts": 2, + "num_experts_per_tok": 1, + "moe_intermediate_size": 8, + "rope_parameters": {"rope_theta": 12345.0}, + } + + +def _tiny_qwen35_text_config() -> dict: + return { + "hidden_size": 16, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 4, + "num_hidden_layers": 2, + "vocab_size": 64, + "num_experts": 2, + "num_experts_per_tok": 1, + "moe_intermediate_size": 8, + "shared_expert_intermediate_size": 8, + "linear_num_key_heads": 2, + "linear_key_head_dim": 4, + "linear_num_value_heads": 2, + "linear_value_head_dim": 4, + "linear_conv_kernel_dim": 2, + "num_nextn_predict_layers": 1, + "layer_types": ["linear_attention", "full_attention", "full_attention"], + "rope_parameters": {"partial_rotary_factor": 1.0, "mrope_section": [1, 1, 0]}, + } + + +def test_registry_resolves_qwen_lite_model_names(): + assert resolve_model_type_from_hf({"model_type": "qwen3_moe"}) == "qwen3" + assert resolve_model_type_from_hf({"model_type": "qwen3_5_moe"}) == "qwen3_5" + assert resolve_runtime_model_name("qwen3", "lite") == "qwen3" + assert resolve_runtime_model_name("qwen3_moe", "lite") == "qwen3_moe" + assert resolve_runtime_model_name("qwen3_5", "lite") == "qwen3_5" + + +def test_qwen3_config_from_hf_dict_derives_head_dim_and_rope_theta(): + cfg = Qwen3MoEConfig._from_hf_dict(_tiny_qwen3_hf_dict()) + + assert cfg.hidden_size == 16 + assert cfg.head_dim == 4 + assert cfg.layer_types == ["full_attention"] + assert cfg.rope_theta == 12345.0 + + +def test_qwen3_config_rejects_invalid_expert_topk(): + hf = _tiny_qwen3_hf_dict() + hf["num_experts_per_tok"] = 3 + + with pytest.raises(ValueError, match="num_experts_per_tok"): + Qwen3MoEConfig._from_hf_dict(hf) + + +def test_qwen35_config_from_text_config_splits_mtp_layer_types(): + cfg = Qwen35Config._from_hf_dict( + {"model_type": "qwen3_5_moe", "text_config": _tiny_qwen35_text_config()} + ) + + assert cfg.layer_types == ["linear_attention", "full_attention"] + assert cfg.mtp_layer_types == ["full_attention"] + assert cfg.rotary_dim == 4 + assert cfg.mrope_section == [1, 1, 0] + + +def test_qwen_lite_protocols_build_configs_from_hf_dicts(): + pytest.importorskip("transformer_engine.pytorch") + + from megatron.lite.model.qwen3_5.lite import protocol as qwen35_protocol + from megatron.lite.model.qwen3_moe.lite import protocol as qwen3_protocol + + qwen3_cfg = qwen3_protocol.build_model_config(_tiny_qwen3_hf_dict(), vocab_size=128) + qwen35_cfg = qwen35_protocol.build_model_config( + {"model_type": "qwen3_5_moe", "text_config": _tiny_qwen35_text_config()}, vocab_size=128 + ) + + assert qwen3_cfg.vocab_size == 128 + assert qwen35_cfg.vocab_size == 128 + assert qwen35_cfg.layer_type_at(0) == "linear_attention" + assert qwen35_cfg.layer_type_at(1) == "full_attention" + + +def test_qwen_lite_protocols_reexport_checkpoint_hook_names(): + protocol_paths = [ + LITE_ROOT / "megatron/lite/model/qwen3_moe/lite/protocol.py", + LITE_ROOT / "megatron/lite/model/qwen3_5/lite/protocol.py", + ] + + for path in protocol_paths: + tree = ast.parse(path.read_text()) + exported = _string_list_assignment(tree, "__all__") + checkpoint_imports = _checkpoint_import_names(tree) + + assert "EXPERT_CLASSIFIER" in exported + assert "PLACEMENT_FN" in exported + assert "EXPERT_CLASSIFIER" in checkpoint_imports + assert "PLACEMENT_FN" in checkpoint_imports + + +def _string_list_assignment(tree: ast.Module, name: str) -> set[str]: + for node in tree.body: + if not isinstance(node, ast.Assign): + continue + if not any(isinstance(target, ast.Name) and target.id == name for target in node.targets): + continue + if not isinstance(node.value, (ast.List, ast.Tuple)): + return set() + return {item.value for item in node.value.elts if isinstance(item, ast.Constant)} + return set() + + +def _checkpoint_import_names(tree: ast.Module) -> set[str]: + names: set[str] = set() + for node in tree.body: + if not isinstance(node, ast.ImportFrom): + continue + if node.module is None or not node.module.endswith(".lite.checkpoint"): + continue + names.update(alias.name for alias in node.names) + return names diff --git a/experimental/lite/tests/unit/primitive/test_attention_moe_unit.py b/experimental/lite/tests/unit/primitive/test_attention_moe_unit.py new file mode 100644 index 00000000000..c1e61aacba4 --- /dev/null +++ b/experimental/lite/tests/unit/primitive/test_attention_moe_unit.py @@ -0,0 +1,148 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = pytest.mark.mlite + + +def _split_grouped_qkvg(): + from megatron.lite.primitive.modules import split_grouped_qkvg + + return split_grouped_qkvg + + +def _moe_aux_scaler(): + from megatron.lite.primitive.modules.moe import MoEAuxLossAutoScaler + + return MoEAuxLossAutoScaler + + +def _router_and_parallel_state(monkeypatch): + from megatron.core.transformer.moe import moe_utils + + if not hasattr(moe_utils, "te_general_gemm"): + monkeypatch.setattr(moe_utils, "te_general_gemm", None, raising=False) + + from megatron.lite.primitive.modules.router import TopKRouter + from megatron.lite.primitive.parallel import ParallelState + + return TopKRouter, ParallelState + + +def _router_config(): + return SimpleNamespace( + hidden_size=4, num_experts=4, num_experts_per_tok=2, router_aux_loss_coef=0.1 + ) + + +def _walk_grad_fn_names(tensor: torch.Tensor) -> set[str]: + names: set[str] = set() + stack = [tensor.grad_fn] + while stack: + fn = stack.pop() + if fn is None: + continue + names.add(type(fn).__name__) + stack.extend(parent for parent, _idx in fn.next_functions) + return names + + +def test_gqa_split_grouped_qkvg_preserves_q_gate_kv_order(): + split_grouped_qkvg = _split_grouped_qkvg() + qkv = torch.arange(24).reshape(1, 24) + + query, gate, key, value = split_grouped_qkvg(qkv, num_heads=4, num_kv_heads=2, head_dim=2) + + assert query.shape == (1, 4, 2) + assert gate.shape == (1, 4, 2) + assert key.shape == (1, 2, 2) + assert value.shape == (1, 2, 2) + assert torch.equal(query, torch.tensor([[[0, 1], [2, 3], [12, 13], [14, 15]]])) + assert torch.equal(gate, torch.tensor([[[4, 5], [6, 7], [16, 17], [18, 19]]])) + assert torch.equal(key, torch.tensor([[[8, 9], [20, 21]]])) + assert torch.equal(value, torch.tensor([[[10, 11], [22, 23]]])) + + +def test_moe_aux_loss_auto_scaler_threads_scaled_aux_gradient(): + MoEAuxLossAutoScaler = _moe_aux_scaler() + MoEAuxLossAutoScaler.set_loss_scale(torch.tensor([0.25])) + output = torch.randn(3, requires_grad=True) + aux_loss = torch.tensor(2.0, requires_grad=True) + + scaled_output = MoEAuxLossAutoScaler.apply(output * 2.0, aux_loss) + scaled_output.sum().backward() + + torch.testing.assert_close(output.grad, torch.full_like(output, 2.0)) + torch.testing.assert_close(aux_loss.grad, torch.tensor(0.25)) + MoEAuxLossAutoScaler.main_loss_backward_scale = None + + +def test_topk_router_returns_finite_scores_and_valid_expert_indices(monkeypatch): + TopKRouter, ParallelState = _router_and_parallel_state(monkeypatch) + config = _router_config() + router = TopKRouter(config, ParallelState(), compute_aux_loss=False) + hidden = torch.randn(5, 4) + + scores, indices = router(hidden) + + assert scores.shape == (5, 2) + assert indices.shape == (5, 2) + assert scores.dtype == hidden.dtype + assert torch.isfinite(scores).all() + assert indices.min().item() >= 0 + assert indices.max().item() < config.num_experts + + +def test_topk_router_scores_are_normalized_and_deterministic_in_eval(monkeypatch): + TopKRouter, ParallelState = _router_and_parallel_state(monkeypatch) + config = _router_config() + router = TopKRouter(config, ParallelState(), compute_aux_loss=False) + hidden = torch.randn(5, config.hidden_size) + + router.eval() + scores_1, indices_1 = router(hidden) + scores_2, indices_2 = router(hidden) + + torch.testing.assert_close(scores_1.sum(dim=-1), torch.ones(hidden.size(0))) + torch.testing.assert_close(scores_1, scores_2, atol=0, rtol=0) + assert torch.equal(indices_1, indices_2) + + +def test_topk_router_does_not_attach_aux_scaler_in_eval(monkeypatch): + TopKRouter, ParallelState = _router_and_parallel_state(monkeypatch) + config = _router_config() + router = TopKRouter(config, ParallelState(), compute_aux_loss=True) + hidden = torch.randn(5, config.hidden_size) + + router.eval() + scores, _indices = router(hidden) + + assert not any("MoEAuxLoss" in name for name in _walk_grad_fn_names(scores)) + + +def test_topk_router_aux_loss_contributes_gate_gradient(monkeypatch): + TopKRouter, ParallelState = _router_and_parallel_state(monkeypatch) + config = _router_config() + router = TopKRouter(config, ParallelState(), compute_aux_loss=True) + hidden = torch.randn(8, config.hidden_size) + + router.train() + scores, _indices = router(hidden) + scores.sum().backward() + grad_with_aux = router.gate.weight.grad.detach().clone() + + router.zero_grad() + saved_coeff = router.aux_loss_coeff + router.aux_loss_coeff = 0.0 + scores_no_aux, _indices = router(hidden) + scores_no_aux.sum().backward() + grad_no_aux = router.gate.weight.grad.detach().clone() + router.aux_loss_coeff = saved_coeff + + assert torch.isfinite(grad_with_aux).all() + assert torch.isfinite(grad_no_aux).all() + assert (grad_with_aux - grad_no_aux).abs().sum().item() > 0.0 diff --git a/experimental/lite/tests/unit/primitive/test_checkpoint_runtime.py b/experimental/lite/tests/unit/primitive/test_checkpoint_runtime.py new file mode 100644 index 00000000000..de25fbdae05 --- /dev/null +++ b/experimental/lite/tests/unit/primitive/test_checkpoint_runtime.py @@ -0,0 +1,291 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import copy +import random +from types import SimpleNamespace +from unittest.mock import patch + +import numpy as np +import torch +import torch.nn as nn + +from megatron.lite.primitive.ckpt import save_training_checkpoint +from megatron.lite.runtime.backends.mlite.runtime import MegatronLiteRuntime +from megatron.lite.runtime.contracts.config import ParallelConfig +from megatron.lite.runtime.contracts.handle import ModelHandle + + +class TinyMLP(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.Sequential(nn.Linear(4, 8), nn.GELU(), nn.Linear(8, 2)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.layers(x) + + +def _step(model: nn.Module, optimizer: torch.optim.Optimizer, x: torch.Tensor, y: torch.Tensor): + optimizer.zero_grad(set_to_none=True) + loss = torch.nn.functional.mse_loss(model(x), y) + loss.backward() + optimizer.step() + return loss.detach() + + +def _clone_model_and_optimizer(model: nn.Module): + clone = copy.deepcopy(model) + optimizer = torch.optim.AdamW(clone.parameters(), lr=1.0e-3, weight_decay=0.0) + return clone, optimizer + + +def _assert_model_close(lhs: nn.Module, rhs: nn.Module): + for (lhs_name, lhs_param), (rhs_name, rhs_param) in zip( + lhs.named_parameters(), rhs.named_parameters(), strict=True + ): + assert lhs_name == rhs_name + torch.testing.assert_close(lhs_param, rhs_param, atol=0.0, rtol=0.0) + + +def test_runtime_local_checkpoint_load_matches_uninterrupted_training(tmp_path): + torch.manual_seed(2029) + base = TinyMLP() + ckpt_model, ckpt_optimizer = _clone_model_and_optimizer(base) + direct_model, direct_optimizer = _clone_model_and_optimizer(base) + loaded_model, loaded_optimizer = _clone_model_and_optimizer(base) + x0, y0 = torch.randn(3, 4), torch.randn(3, 2) + x1, y1 = torch.randn(3, 4), torch.randn(3, 2) + + _step(ckpt_model, ckpt_optimizer, x0, y0) + _step(direct_model, direct_optimizer, x0, y0) + + runtime = MegatronLiteRuntime.__new__(MegatronLiteRuntime) + runtime.save_checkpoint( + ModelHandle(model=ckpt_model, optimizer=ckpt_optimizer), + str(tmp_path), + step=1, + use_dcp=False, + ) + + assert ( + runtime.load_checkpoint( + ModelHandle(model=loaded_model, optimizer=loaded_optimizer), + str(tmp_path), + use_dcp=False, + ) + == 1 + ) + + _step(direct_model, direct_optimizer, x1, y1) + _step(loaded_model, loaded_optimizer, x1, y1) + + _assert_model_close(direct_model, loaded_model) + + +class DistOptLike: + """Small optimizer wrapper with the same checkpoint contract as distopt.""" + + def __init__(self, optimizer: torch.optim.Optimizer): + self.optimizer = optimizer + self.load_calls = 0 + self.parameter_save_calls = 0 + self.parameter_load_calls = 0 + self.update_legacy_format = None + + def zero_grad(self): + self.optimizer.zero_grad(set_to_none=True) + + def step(self): + self.optimizer.step() + return True, 0.0, 0 + + def state_dict(self): + state = self.optimizer.state_dict() + state["distopt_like_marker"] = {"load_calls": self.load_calls} + return state + + def load_state_dict(self, state): + marker = state.pop("distopt_like_marker") + self.load_calls = int(marker["load_calls"]) + 1 + self.optimizer.load_state_dict(state) + + def save_parameter_state(self, filename: str): + self.parameter_save_calls += 1 + torch.save({"parameter_save_calls": self.parameter_save_calls}, filename) + + def load_parameter_state(self, filename: str, *, update_legacy_format: bool = False): + state = torch.load(filename, weights_only=False) + self.parameter_load_calls = int(state["parameter_save_calls"]) + self.update_legacy_format = update_legacy_format + + +def test_runtime_local_checkpoint_uses_optimizer_parameter_state_contract(tmp_path): + torch.manual_seed(2030) + model = TinyMLP() + optimizer = DistOptLike(torch.optim.AdamW(model.parameters(), lr=1.0e-3)) + x, y = torch.randn(3, 4), torch.randn(3, 2) + optimizer.zero_grad() + torch.nn.functional.mse_loss(model(x), y).backward() + optimizer.step() + + runtime = MegatronLiteRuntime.__new__(MegatronLiteRuntime) + runtime.save_checkpoint( + ModelHandle(model=model, optimizer=optimizer), str(tmp_path), step=7, use_dcp=False + ) + + loaded_model = TinyMLP() + loaded_optimizer = DistOptLike(torch.optim.AdamW(loaded_model.parameters(), lr=1.0e-3)) + + assert ( + runtime.load_checkpoint( + ModelHandle(model=loaded_model, optimizer=loaded_optimizer), + str(tmp_path), + update_legacy_format=True, + use_dcp=False, + ) + == 7 + ) + assert loaded_optimizer.load_calls == 1 + assert loaded_optimizer.parameter_load_calls == 1 + assert loaded_optimizer.update_legacy_format is True + assert (tmp_path / "training_state.optimizer_parameter_state.pt").exists() + _assert_model_close(model, loaded_model) + + +def test_runtime_local_checkpoint_restores_rng_state(tmp_path): + model = TinyMLP() + runtime = MegatronLiteRuntime.__new__(MegatronLiteRuntime) + + random.seed(2031) + np.random.seed(2031) + torch.manual_seed(2031) + + runtime.save_checkpoint( + ModelHandle(model=model, optimizer=None), str(tmp_path), step=9, use_dcp=False + ) + + expected_python = random.random() + expected_numpy = np.random.random(4) + expected_torch = torch.rand(4) + + random.seed(9999) + np.random.seed(9999) + torch.manual_seed(9999) + + assert ( + runtime.load_checkpoint( + ModelHandle(model=model, optimizer=None), str(tmp_path), use_dcp=False + ) + == 9 + ) + assert random.random() == expected_python + np.testing.assert_allclose(np.random.random(4), expected_numpy, atol=0.0, rtol=0.0) + torch.testing.assert_close(torch.rand(4), expected_torch, atol=0.0, rtol=0.0) + + +def test_runtime_local_checkpoint_uses_rank_specific_files_when_distributed(tmp_path): + model = TinyMLP() + runtime = MegatronLiteRuntime.__new__(MegatronLiteRuntime) + + with ( + patch("megatron.lite.primitive.ckpt.dcp.dist.is_available", return_value=True), + patch("megatron.lite.primitive.ckpt.dcp.dist.is_initialized", return_value=True), + patch("megatron.lite.primitive.ckpt.dcp.dist.get_rank", return_value=3), + ): + runtime.save_checkpoint( + ModelHandle(model=model, optimizer=None), str(tmp_path), step=11, use_dcp=False + ) + assert (tmp_path / "training_state_rank_00003.pt").exists() + assert not (tmp_path / "training_state.pt").exists() + assert ( + runtime.load_checkpoint( + ModelHandle(model=model, optimizer=None), str(tmp_path), use_dcp=False + ) + == 11 + ) + + +def test_primitive_local_checkpoint_keeps_optimizer_checkpoints_local(tmp_path): + model = TinyMLP() + optimizer = torch.optim.AdamW(model.parameters(), lr=1.0e-3) + + with patch("megatron.lite.primitive.ckpt.dcp.dcp.save") as dcp_save_mock: + save_training_checkpoint(model, optimizer, 12, str(tmp_path), use_dcp=False) + + dcp_save_mock.assert_not_called() + assert (tmp_path / "training_state.pt").exists() + + +def test_primitive_explicit_dcp_saves_optimizer_rank_sidecar(tmp_path): + model = TinyMLP() + optimizer = torch.optim.AdamW(model.parameters(), lr=1.0e-3) + parallel = ParallelConfig(tp=1, ep=1, pp=1, cp=1) + + with ( + patch("megatron.lite.primitive.ckpt.dcp._build_meshes", return_value=(None, None)), + patch( + "megatron.lite.primitive.ckpt.dcp.DTensor.from_local", + side_effect=lambda tensor, *args, **kwargs: tensor, + ), + patch("megatron.lite.primitive.ckpt.dcp.dcp.save") as dcp_save_mock, + ): + save_training_checkpoint( + model, optimizer, 12, str(tmp_path), parallel, object(), use_dcp=True + ) + + dcp_save_mock.assert_called_once() + assert (tmp_path / "step_12" / "optimizer_rank_0.pt").exists() + + +def test_runtime_dcp_checkpoint_threads_parallel_config_and_protocol_hooks(tmp_path): + model = TinyMLP() + parallel = ParallelConfig(tp=2, ep=1, pp=1, cp=1) + ps = object() + + def placement_fn(name: str): + return ["placement", name] + + def expert_classifier(name: str): + return name.endswith("expert") + + proto = SimpleNamespace(PLACEMENT_FN=placement_fn, EXPERT_CLASSIFIER=expert_classifier) + handle = ModelHandle( + model=[model], + optimizer=None, + parallel_state=ps, + config=SimpleNamespace(parallel=parallel), + _extras={"model_chunks": [model], "protocol": proto}, + ) + runtime = MegatronLiteRuntime.__new__(MegatronLiteRuntime) + + with patch("megatron.lite.primitive.ckpt.save_training_checkpoint") as save_mock: + runtime.save_checkpoint(handle, str(tmp_path), global_step=13, use_dcp=True) + + save_args = save_mock.call_args.args + save_kwargs = save_mock.call_args.kwargs + assert isinstance(save_args[0], nn.ModuleList) + assert save_args[0][0] is model + assert save_args[2] == 13 + assert save_args[3] == str(tmp_path) + assert save_args[4] is parallel + assert save_args[5] is ps + assert save_kwargs["get_placements"] is placement_fn + assert save_kwargs["is_expert"] is expert_classifier + assert save_kwargs["use_dcp"] is True + assert save_kwargs["save_rng"] is True + + with patch( + "megatron.lite.primitive.ckpt.load_training_checkpoint", return_value=13 + ) as load_mock: + assert runtime.load_checkpoint(handle, str(tmp_path), use_dcp=True) == 13 + + load_args = load_mock.call_args.args + load_kwargs = load_mock.call_args.kwargs + assert isinstance(load_args[0], nn.ModuleList) + assert load_args[0][0] is model + assert load_args[3] is parallel + assert load_args[4] is ps + assert load_kwargs["get_placements"] is placement_fn + assert load_kwargs["is_expert"] is expert_classifier + assert load_kwargs["use_dcp"] is True + assert load_kwargs["load_rng"] is True diff --git a/experimental/lite/tests/unit/primitive/test_checkpoint_unit.py b/experimental/lite/tests/unit/primitive/test_checkpoint_unit.py new file mode 100644 index 00000000000..52008febbbc --- /dev/null +++ b/experimental/lite/tests/unit/primitive/test_checkpoint_unit.py @@ -0,0 +1,138 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import copy + +import pytest +import torch +import torch.nn as nn + +from megatron.lite.runtime.backends.mlite.runtime import MegatronLiteRuntime +from megatron.lite.runtime.contracts.handle import ModelHandle + +pytestmark = pytest.mark.mlite + + +class TinyMLP(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.Sequential(nn.Linear(4, 8), nn.GELU(), nn.Linear(8, 2)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.layers(x) + + +def _step(model: nn.Module, optimizer: torch.optim.Optimizer, x: torch.Tensor, y: torch.Tensor): + optimizer.zero_grad(set_to_none=True) + loss = torch.nn.functional.mse_loss(model(x), y) + loss.backward() + optimizer.step() + return loss.detach() + + +def _clone_model_and_optimizer(model: nn.Module): + clone = copy.deepcopy(model) + optimizer = torch.optim.AdamW(clone.parameters(), lr=1.0e-3, weight_decay=0.0) + return clone, optimizer + + +def _assert_model_close(lhs: nn.Module, rhs: nn.Module): + for (lhs_name, lhs_param), (rhs_name, rhs_param) in zip( + lhs.named_parameters(), rhs.named_parameters(), strict=True + ): + assert lhs_name == rhs_name + torch.testing.assert_close(lhs_param, rhs_param, atol=0.0, rtol=0.0) + + +def test_runtime_checkpoint_load_matches_uninterrupted_training(tmp_path): + torch.manual_seed(2029) + base = TinyMLP() + ckpt_model, ckpt_optimizer = _clone_model_and_optimizer(base) + direct_model, direct_optimizer = _clone_model_and_optimizer(base) + loaded_model, loaded_optimizer = _clone_model_and_optimizer(base) + x0, y0 = torch.randn(3, 4), torch.randn(3, 2) + x1, y1 = torch.randn(3, 4), torch.randn(3, 2) + + _step(ckpt_model, ckpt_optimizer, x0, y0) + _step(direct_model, direct_optimizer, x0, y0) + + runtime = MegatronLiteRuntime.__new__(MegatronLiteRuntime) + ckpt_handle = ModelHandle( + model=ckpt_model, optimizer=ckpt_optimizer, _extras={"model_chunks": [ckpt_model]} + ) + runtime.save_checkpoint(ckpt_handle, str(tmp_path), step=1, use_dcp=False) + + loaded_handle = ModelHandle( + model=loaded_model, optimizer=loaded_optimizer, _extras={"model_chunks": [loaded_model]} + ) + assert runtime.load_checkpoint(loaded_handle, str(tmp_path), use_dcp=False) == 1 + + _step(direct_model, direct_optimizer, x1, y1) + _step(loaded_model, loaded_optimizer, x1, y1) + + _assert_model_close(direct_model, loaded_model) + + +class DistOptLike: + """Small optimizer wrapper with the same checkpoint contract as distopt.""" + + def __init__(self, optimizer: torch.optim.Optimizer): + self.optimizer = optimizer + self.load_calls = 0 + self.parameter_save_calls = 0 + self.parameter_load_calls = 0 + + def zero_grad(self): + self.optimizer.zero_grad(set_to_none=True) + + def step(self): + self.optimizer.step() + return True, 0.0, 0 + + def state_dict(self): + state = self.optimizer.state_dict() + state["distopt_like_marker"] = {"load_calls": self.load_calls} + return state + + def load_state_dict(self, state): + marker = state.pop("distopt_like_marker") + self.load_calls = int(marker["load_calls"]) + 1 + self.optimizer.load_state_dict(state) + + def save_parameter_state(self, filename: str): + self.parameter_save_calls += 1 + torch.save({"parameter_save_calls": self.parameter_save_calls}, filename) + + def load_parameter_state(self, filename: str, *, update_legacy_format: bool = False): + state = torch.load(filename) + self.parameter_load_calls = int(state["parameter_save_calls"]) + + +def test_runtime_checkpoint_uses_optimizer_state_dict_contract(tmp_path): + torch.manual_seed(2030) + model = TinyMLP() + optimizer = DistOptLike(torch.optim.AdamW(model.parameters(), lr=1.0e-3)) + x, y = torch.randn(3, 4), torch.randn(3, 2) + optimizer.zero_grad() + torch.nn.functional.mse_loss(model(x), y).backward() + optimizer.step() + + runtime = MegatronLiteRuntime.__new__(MegatronLiteRuntime) + runtime.save_checkpoint( + ModelHandle(model=model, optimizer=optimizer, _extras={"model_chunks": [model]}), + str(tmp_path), + step=7, + use_dcp=False, + ) + + loaded_model = TinyMLP() + loaded_optimizer = DistOptLike(torch.optim.AdamW(loaded_model.parameters(), lr=1.0e-3)) + loaded_handle = ModelHandle( + model=loaded_model, optimizer=loaded_optimizer, _extras={"model_chunks": [loaded_model]} + ) + + assert runtime.load_checkpoint(loaded_handle, str(tmp_path), use_dcp=False) == 7 + assert loaded_optimizer.load_calls == 1 + assert loaded_optimizer.parameter_load_calls == 1 + assert (tmp_path / "training_state.optimizer_parameter_state.pt").exists() + _assert_model_close(model, loaded_model) diff --git a/experimental/lite/tests/unit/primitive/test_dist_opt_validation.py b/experimental/lite/tests/unit/primitive/test_dist_opt_validation.py new file mode 100644 index 00000000000..5af9fdc4e74 --- /dev/null +++ b/experimental/lite/tests/unit/primitive/test_dist_opt_validation.py @@ -0,0 +1,26 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import pytest + +from megatron.lite.primitive.optimizers.megatron_wrap import validate_mc_config, validate_mc_session +from megatron.lite.runtime.backends.mlite.config import MegatronLiteConfig +from megatron.lite.runtime.contracts.config import ParallelConfig + + +def _engine_cfg(*, model_name: str, pp: int = 1, vpp: int = 1) -> MegatronLiteConfig: + return MegatronLiteConfig(model_name=model_name, parallel=ParallelConfig(pp=pp, vpp=vpp)) + + +def test_dist_opt_validation_accepts_model_agnostic_config(): + validate_mc_config(_engine_cfg(model_name="synthetic_custom_model", pp=1, vpp=1)) + + +def test_dist_opt_validation_keeps_vpp_parallel_constraint(): + with pytest.raises(ValueError, match="dist_opt requires pp>1 when vpp>1"): + validate_mc_config(_engine_cfg(model_name="synthetic_custom_model", pp=1, vpp=2)) + + +def test_validate_mc_session_alias_matches_config_validator(): + assert validate_mc_session is validate_mc_config + validate_mc_session(_engine_cfg(model_name="another_synthetic_model", pp=2, vpp=2)) diff --git a/experimental/lite/tests/unit/primitive/test_fsdp2_offload_gpu.py b/experimental/lite/tests/unit/primitive/test_fsdp2_offload_gpu.py new file mode 100644 index 00000000000..abc42bf6469 --- /dev/null +++ b/experimental/lite/tests/unit/primitive/test_fsdp2_offload_gpu.py @@ -0,0 +1,161 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.nn as nn + +from megatron.lite.primitive.optimizers.fsdp2 import ( + FSDP2Config, + build_fsdp2_adamw, + build_fsdp2_device_mesh, + fsdp2_available, + wrap_fsdp2, +) +from megatron.lite.primitive.optimizers.fsdp2.adamw import iter_torch_optimizers, to_local_tensor +from megatron.lite.primitive.parallel.state import ParallelState +from megatron.lite.runtime.backends.mlite.runtime import MegatronLiteRuntime +from megatron.lite.runtime.contracts.handle import ModelHandle + + +class TinyUnit(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(8, 8) + self.act = nn.GELU() + + def forward(self, x): + return self.act(self.linear(x)) + + +class TinyModel(nn.Module): + def __init__(self): + super().__init__() + self.unit0 = TinyUnit() + self.unit1 = TinyUnit() + self.out = nn.Linear(8, 4) + + def forward(self, x): + return self.out(self.unit1(self.unit0(x))) + + +@pytest.fixture(scope="module", autouse=True) +def _single_node_cuda_dist(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for FSDP2 offload tests.") + if not fsdp2_available(): + pytest.skip("Installed PyTorch does not expose FSDP2 fully_shard.") + + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("LOCAL_RANK", "0") + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29501") + + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + created_pg = False + if not dist.is_initialized(): + dist.init_process_group(backend="nccl", init_method="env://") + created_pg = True + yield + if created_pg and dist.is_initialized(): + dist.destroy_process_group() + + +def _parallel_state() -> ParallelState: + rank = dist.get_rank() + world_size = dist.get_world_size() + return ParallelState( + dp_group=dist.group.WORLD, + dp_cp_group=dist.group.WORLD, + dp_size=world_size, + dp_cp_size=world_size, + dp_rank=rank, + dp_cp_rank=rank, + ) + + +def _build_fsdp2_model(dtype: torch.dtype = torch.bfloat16) -> tuple[nn.Module, ParallelState]: + torch.manual_seed(1234) + model = TinyModel().cuda().to(dtype=dtype) + ps = _parallel_state() + config = FSDP2Config(unit_modules=(TinyUnit,), reshard_after_forward=True) + mesh = build_fsdp2_device_mesh(ps, config) + return wrap_fsdp2(model, ps, config, mesh=mesh), ps + + +def _build_optimizer(model: nn.Module, ps: ParallelState, *, offload_fraction: float): + return build_fsdp2_adamw( + [model], + SimpleNamespace( + optimizer="adam", + lr=1.0e-3, + weight_decay=0.0, + adam_beta1=0.9, + adam_beta2=0.95, + adam_eps=1.0e-8, + clip_grad=1.0, + offload_fraction=offload_fraction, + ), + ps, + use_fp32_master=True, + ) + + +def _local_param_devices(model: nn.Module) -> set[str]: + return {to_local_tensor(param.detach()).device.type for param in model.parameters()} + + +def _optimizer_state_devices(optimizer) -> set[str]: + devices: set[str] = set() + for child in iter_torch_optimizers(optimizer.optimizer): + for param_state in getattr(child, "state", {}).values(): + if not isinstance(param_state, dict): + continue + for value in param_state.values(): + if isinstance(value, torch.Tensor): + devices.add(to_local_tensor(value).device.type) + return devices + + +def test_fsdp2_runtime_model_and_optimizer_offload_roundtrip_single_gpu(): + model, ps = _build_fsdp2_model() + optimizer = _build_optimizer(model, ps, offload_fraction=0.0) + handle = ModelHandle( + model=model, optimizer=optimizer, parallel_state=ps, _extras={"model_chunks": [model]} + ) + runtime = MegatronLiteRuntime.__new__(MegatronLiteRuntime) + + assert _local_param_devices(model) == {"cuda"} + assert _optimizer_state_devices(optimizer) == {"cuda"} + + runtime.to(handle, "cpu", model=True, optimizer=True, grad=True) + assert _local_param_devices(model) == {"cpu"} + assert _optimizer_state_devices(optimizer) == {"cpu"} + + runtime.to(handle, "cuda", model=True, optimizer=True, grad=True) + assert _local_param_devices(model) == {"cuda"} + assert _optimizer_state_devices(optimizer) == {"cuda"} + + +def test_fsdp2_offload_fraction_keeps_optimizer_update_state_on_cpu_single_gpu(): + model, ps = _build_fsdp2_model() + optimizer = _build_optimizer(model, ps, offload_fraction=1.0) + + assert _optimizer_state_devices(optimizer) == {"cpu"} + + x = torch.randn(4, 8, device="cuda", dtype=torch.bfloat16) + target = torch.randn(4, 4, device="cuda", dtype=torch.bfloat16) + optimizer.zero_grad() + loss = torch.nn.functional.mse_loss(model(x).float(), target.float()) + loss.backward() + success, grad_norm, _ = optimizer.step() + + assert success + assert torch.isfinite(torch.tensor(grad_norm)) + assert _local_param_devices(model) == {"cuda"} + assert _optimizer_state_devices(optimizer) == {"cpu"} diff --git a/experimental/lite/tests/unit/primitive/test_fsdp2_unit.py b/experimental/lite/tests/unit/primitive/test_fsdp2_unit.py new file mode 100644 index 00000000000..acdc7230836 --- /dev/null +++ b/experimental/lite/tests/unit/primitive/test_fsdp2_unit.py @@ -0,0 +1,469 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import copy +import importlib +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.nn as nn + +from megatron.lite.primitive.optimizers.fsdp2 import ( + FSDP2Config, + FSDP2Optimizer, + all_reduce_scalar_, + clip_grads_with_sharded_norm_, + fsdp2_available, +) +from megatron.lite.primitive.optimizers.fsdp2.adamw import build_adamw_optimizer +from megatron.lite.primitive.optimizers.fsdp2.wrap import build_fsdp2_shard_placement_fn +from megatron.lite.primitive.parallel.state import ParallelState + +pytestmark = pytest.mark.mlite + +fsdp2_wrap = importlib.import_module("megatron.lite.primitive.optimizers.fsdp2.wrap") +fsdp2_optimizer = importlib.import_module("megatron.lite.primitive.optimizers.fsdp2.optimizer") +fsdp2_grad_clip = importlib.import_module("megatron.lite.primitive.optimizers.fsdp2.grad_clip") + + +class ToyBlock(nn.Module): + def __init__(self): + super().__init__() + self.proj = nn.Linear(4, 4) + + def forward(self, x): + return self.proj(x) + + +class ToyModel(nn.Module): + def __init__(self): + super().__init__() + self.block = ToyBlock() + self.out = nn.Linear(4, 2) + + def forward(self, x): + return self.out(self.block(x)) + + +class TwoBlockModel(nn.Module): + def __init__(self): + super().__init__() + self.block0 = ToyBlock() + self.block1 = ToyBlock() + self.out = nn.Linear(4, 2) + + def forward(self, x): + return self.out(self.block1(self.block0(x))) + + +class NestedToyBlock(ToyBlock): + def __init__(self): + super().__init__() + self.inner = ToyBlock() + + +def test_fsdp2_config_validates_empty_wrap_surface(): + with pytest.raises(ValueError, match="wrap_root=True"): + FSDP2Config(wrap_root=False) + + +@pytest.mark.parametrize("field", ["mesh_dim_name", "device_type"]) +def test_fsdp2_config_rejects_empty_names(field: str): + with pytest.raises(ValueError, match=field): + FSDP2Config(**{field: ""}) + + +def test_fsdp2_config_normalizes_unit_and_leaf_modules(): + cfg = FSDP2Config(unit_modules=[nn.Linear], leaf_module_names=["embed"]) + + assert cfg.unit_modules == (nn.Linear,) + assert cfg.leaf_module_names == ("embed",) + assert isinstance(fsdp2_available(), bool) + + +def test_fsdp2_optimizer_offloads_dtensor_state_without_extra_knob(monkeypatch): + model = ToyModel() + torch_optimizer = torch.optim.AdamW(model.parameters(), lr=1.0e-3) + optimizer = FSDP2Optimizer(torch_optimizer, model.parameters()) + calls: list[bool] = [] + + def fake_move_optimizer_state_to_cpu(_optimizer, _offloaded_state, *, include_dtensor_state): + calls.append(include_dtensor_state) + + monkeypatch.setattr( + fsdp2_optimizer, "move_optimizer_state_to_cpu", fake_move_optimizer_state_to_cpu + ) + + optimizer.offload_state_to_cpu() + + assert calls == [True] + assert not hasattr(optimizer, "optimizer_offload_dtensor_state") + + +def test_fsdp2_shard_placement_prefers_first_divisible_dimension(): + placement_for_two = build_fsdp2_shard_placement_fn(2) + placement_for_three = build_fsdp2_shard_placement_fn(3) + + assert placement_for_two(nn.Parameter(torch.empty(3, 4))).dim == 1 + assert placement_for_three(nn.Parameter(torch.empty(3, 4))).dim == 0 + + +def test_fsdp2_shard_placement_rejects_invalid_group_size(): + with pytest.raises(ValueError, match="positive"): + build_fsdp2_shard_placement_fn(0) + + +def test_fsdp2_rejects_invalid_unit_path(): + with pytest.raises(ValueError, match="Invalid FSDP2 unit module path"): + fsdp2_wrap._resolve_unit_module_types(("Linear",)) + + +def test_fsdp2_rejects_non_module_unit_path(): + with pytest.raises(TypeError, match="does not resolve"): + fsdp2_wrap._resolve_unit_module_types(("math.sqrt",)) + + +def test_wrap_fsdp2_requires_distributed_when_mesh_is_not_provided(monkeypatch): + monkeypatch.setattr(fsdp2_wrap, "_load_fully_shard", lambda: lambda module, **kwargs: module) + + with pytest.raises(RuntimeError, match="torch.distributed"): + fsdp2_wrap.wrap_fsdp2(ToyModel(), ParallelState(), FSDP2Config()) + + +def test_wrap_fsdp2_wraps_units_then_root_and_preserves_param_attrs(monkeypatch): + model = ToyModel() + model.block.proj.weight.tensor_model_parallel = True + calls: list[nn.Module] = [] + + def fake_fully_shard(module, **kwargs): + calls.append(module) + for param in module.parameters(): + vars(param).clear() + module._fake_fsdp2_kwargs = kwargs + return module + + monkeypatch.setattr(fsdp2_wrap, "_load_fully_shard", lambda: fake_fully_shard) + + result = fsdp2_wrap.wrap_fsdp2( + model, + ParallelState(), + FSDP2Config(unit_modules=(ToyBlock,), reshard_after_forward=False), + mesh=SimpleNamespace(name="mesh"), + ) + + assert result is model + assert calls == [model.block, model] + assert model.block.proj.weight.tensor_model_parallel is True + assert model._fake_fsdp2_kwargs["reshard_after_forward"] is False + assert model._fake_fsdp2_kwargs["mesh"].name == "mesh" + + +def test_wrap_fsdp2_accepts_unit_module_import_paths(monkeypatch): + model = ToyModel() + calls: list[nn.Module] = [] + + def fake_fully_shard(module, **kwargs): + calls.append(module) + return module + + monkeypatch.setattr(fsdp2_wrap, "_load_fully_shard", lambda: fake_fully_shard) + + fsdp2_wrap.wrap_fsdp2( + model, + ParallelState(), + FSDP2Config(unit_modules=("torch.nn.modules.linear.Linear",), wrap_root=False), + mesh=SimpleNamespace(name="mesh"), + ) + + assert calls == [model.block.proj, model.out] + + +def test_wrap_fsdp2_uses_container_order_without_nested_unit_duplicates(monkeypatch): + model = nn.Module() + model.layers = nn.ModuleDict({"10": NestedToyBlock(), "2": ToyBlock(), "11": ToyBlock()}) + calls: list[nn.Module] = [] + + def fake_fully_shard(module, **kwargs): + calls.append(module) + module._fake_fsdp2_kwargs = kwargs + return module + + monkeypatch.setattr(fsdp2_wrap, "_load_fully_shard", lambda: fake_fully_shard) + + fsdp2_wrap.wrap_fsdp2( + model, + ParallelState(), + FSDP2Config(unit_modules=(ToyBlock,), reshard_after_forward=True), + mesh=SimpleNamespace(name="mesh"), + ) + + assert calls == [model.layers["10"], model.layers["2"], model.layers["11"], model] + assert model.layers["10"]._fake_fsdp2_kwargs["reshard_after_forward"] is True + assert not hasattr(model.layers["10"].inner, "_fake_fsdp2_kwargs") + assert model.layers["11"]._fake_fsdp2_kwargs["reshard_after_forward"] is False + + +def test_wrap_fsdp2_configures_default_forward_prefetch(monkeypatch): + model = TwoBlockModel() + calls: list[nn.Module] = [] + + def fake_fully_shard(module, **kwargs): + calls.append(module) + module._forward_prefetch = None + module._backward_prefetch = None + module._fake_fsdp2_kwargs = kwargs + + def set_forward_prefetch(modules, *, _module=module): + _module._forward_prefetch = list(modules) + + def set_backward_prefetch(modules, *, _module=module): + _module._backward_prefetch = list(modules) + + module.set_modules_to_forward_prefetch = set_forward_prefetch + module.set_modules_to_backward_prefetch = set_backward_prefetch + return module + + monkeypatch.setattr(fsdp2_wrap, "_load_fully_shard", lambda: fake_fully_shard) + + fsdp2_wrap.wrap_fsdp2( + model, + ParallelState(), + FSDP2Config(unit_modules=(ToyBlock,), reshard_after_forward=True), + mesh=SimpleNamespace(name="mesh"), + ) + + assert calls == [model.block0, model.block1, model] + assert model.block0._fake_fsdp2_kwargs["reshard_after_forward"] is True + assert model.block1._fake_fsdp2_kwargs["reshard_after_forward"] is False + assert model._fake_fsdp2_kwargs["reshard_after_forward"] is False + assert model._forward_prefetch == [model.block0] + assert model.block0._forward_prefetch == [model.block1] + assert model.block1._backward_prefetch is None + + +def test_wrap_fsdp2_prefetch_depths(monkeypatch): + model = nn.Sequential(ToyBlock(), ToyBlock(), ToyBlock()) + + def fake_fully_shard(module, **kwargs): + module._forward_prefetch = None + module._backward_prefetch = None + + def set_forward_prefetch(modules, *, _module=module): + _module._forward_prefetch = list(modules) + + def set_backward_prefetch(modules, *, _module=module): + _module._backward_prefetch = list(modules) + + module.set_modules_to_forward_prefetch = set_forward_prefetch + module.set_modules_to_backward_prefetch = set_backward_prefetch + return module + + monkeypatch.setattr(fsdp2_wrap, "_load_fully_shard", lambda: fake_fully_shard) + + fsdp2_wrap.wrap_fsdp2( + model, + ParallelState(), + FSDP2Config( + unit_modules=(ToyBlock,), + wrap_root=False, + forward_prefetch_depth=2, + backward_prefetch_depth=2, + ), + mesh=SimpleNamespace(name="mesh"), + ) + + assert model[0]._forward_prefetch == [model[1], model[2]] + assert model[1]._forward_prefetch == [model[2]] + assert model[2]._backward_prefetch == [model[1], model[0]] + + +def test_clip_grads_with_sharded_norm_scales_cpu_grads_once(): + p0 = nn.Parameter(torch.ones(2)) + p1 = nn.Parameter(torch.ones(2)) + p0.grad = torch.tensor([3.0, 4.0]) + p1.grad = torch.tensor([0.0, 12.0]) + + clip_grads_with_sharded_norm_([p0, p1], max_norm=6.5, total_norm=torch.tensor(13.0)) + + scale = 6.5 / (13.0 + 1.0e-6) + torch.testing.assert_close(p0.grad, torch.tensor([3.0, 4.0]) * scale) + torch.testing.assert_close(p1.grad, torch.tensor([0.0, 12.0]) * scale) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for NCCL scalar test.") +def test_all_reduce_scalar_moves_cpu_value_to_nccl_device(monkeypatch): + group = object() + reduced_devices: list[str] = [] + + def fake_all_reduce(value, *, op, group): + assert group is fake_all_reduce.group + assert op == dist.ReduceOp.SUM + reduced_devices.append(value.device.type) + value.add_(5.0) + + fake_all_reduce.group = group + monkeypatch.setattr(fsdp2_grad_clip.dist, "get_backend", lambda _group: "nccl") + monkeypatch.setattr(fsdp2_grad_clip.dist, "all_reduce", fake_all_reduce) + + value = torch.tensor(7.0) + all_reduce_scalar_(value, op=dist.ReduceOp.SUM, group=group) + + assert reduced_devices == ["cuda"] + assert value.device.type == "cpu" + torch.testing.assert_close(value, torch.tensor(12.0)) + + +def test_fsdp2_optimizer_uses_scalar_all_reduce_for_all_norm_groups(monkeypatch): + groups = SimpleNamespace( + dp_cp=object(), + tp=object(), + replicated=object(), + expert=object(), + pp=object(), + ) + reduced_groups: list[object] = [] + + def fake_all_reduce_scalar(value, *, op, group): + assert value.ndim == 0 + assert op == dist.ReduceOp.SUM + reduced_groups.append(group) + + monkeypatch.setattr(fsdp2_optimizer, "all_reduce_scalar_", fake_all_reduce_scalar) + monkeypatch.setattr(fsdp2_optimizer.dist, "is_initialized", lambda: True) + monkeypatch.setattr(fsdp2_optimizer.dist, "get_world_size", lambda _group: 2) + + sharded = nn.Parameter(torch.tensor([1.0])) + replicated = nn.Parameter(torch.tensor([1.0])) + expert = nn.Parameter(torch.tensor([1.0])) + tp_replicated = nn.Parameter(torch.tensor([1.0])) + sharded.grad = torch.tensor([2.0]) + replicated.grad = torch.tensor([3.0]) + expert.grad = torch.tensor([4.0]) + tp_replicated.grad = torch.tensor([5.0]) + + optimizer = FSDP2Optimizer( + torch.optim.SGD([sharded, replicated, expert, tp_replicated], lr=0.0), + [sharded, replicated, expert, tp_replicated], + ParallelState( + dp_cp_group=groups.dp_cp, + tp_group=groups.tp, + pp_group=groups.pp, + ), + clip_grad=100.0, + replicated_grad_params=[replicated], + replicated_grad_norm_group=groups.replicated, + expert_sharded_grad_params=[expert], + expert_sharded_grad_norm_group=groups.expert, + tp_replicated_grad_params=[tp_replicated], + ) + + assert optimizer.clip_grad_norm() == pytest.approx( + (2.0**2 + 3.0**2 + 4.0**2 + 5.0**2) ** 0.5 + ) + assert reduced_groups == [ + groups.dp_cp, + groups.tp, + groups.replicated, + groups.expert, + groups.pp, + ] + + +def test_fp32_adamw_state_dict_roundtrip_cpu(): + param = nn.Parameter(torch.tensor([1.0, -2.0], dtype=torch.bfloat16)) + optimizer = build_adamw_optimizer( + [{"params": [param], "weight_decay": 0.0}], + all_params=[param], + lr=0.1, + weight_decay=0.0, + betas=(0.9, 0.99), + eps=1.0e-8, + foreach=False, + use_fp32_master=True, + cpu_update=False, + model_param_dtypes={id(param): torch.bfloat16}, + opt=SimpleNamespace(), + ) + param.grad = torch.tensor([0.5, -0.25], dtype=torch.bfloat16) + optimizer.step() + state = optimizer.state_dict() + + loaded_param = nn.Parameter(torch.tensor([9.0, 9.0], dtype=torch.bfloat16)) + loaded_optimizer = build_adamw_optimizer( + [{"params": [loaded_param], "weight_decay": 0.0}], + all_params=[loaded_param], + lr=0.1, + weight_decay=0.0, + betas=(0.9, 0.99), + eps=1.0e-8, + foreach=False, + use_fp32_master=True, + cpu_update=False, + model_param_dtypes={id(loaded_param): torch.bfloat16}, + opt=SimpleNamespace(), + ) + loaded_optimizer.load_state_dict(state) + loaded_state = loaded_optimizer.state_dict() + + assert loaded_state["step_count"] == state["step_count"] + for key in ("master_params", "exp_avgs", "exp_avg_sqs", "steps"): + assert len(loaded_state[key]) == len(state[key]) + torch.testing.assert_close(loaded_state["master_params"][0], state["master_params"][0]) + torch.testing.assert_close(loaded_state["exp_avgs"][0], state["exp_avgs"][0]) + torch.testing.assert_close(loaded_state["exp_avg_sqs"][0], state["exp_avg_sqs"][0]) + assert loaded_state["steps"] == state["steps"] + + +@pytest.mark.parametrize("cpu_update", [False, True]) +def test_fp32_adamw_load_matches_uninterrupted_next_step_cpu(cpu_update: bool): + def build(initial_value: torch.Tensor): + param = nn.Parameter(initial_value.clone().to(dtype=torch.bfloat16)) + optimizer = build_adamw_optimizer( + [{"params": [param], "weight_decay": 0.0}], + all_params=[param], + lr=0.1, + weight_decay=0.0, + betas=(0.9, 0.99), + eps=1.0e-8, + foreach=False, + use_fp32_master=True, + cpu_update=cpu_update, + model_param_dtypes={id(param): torch.bfloat16}, + opt=SimpleNamespace(), + ) + return param, optimizer + + initial = torch.tensor([1.0, -2.0], dtype=torch.float32) + first_grad = torch.tensor([0.5, -0.25], dtype=torch.bfloat16) + second_grad = torch.tensor([-0.125, 0.375], dtype=torch.bfloat16) + + ckpt_param, ckpt_optimizer = build(initial) + direct_param, direct_optimizer = build(initial) + loaded_param, loaded_optimizer = build(initial) + + ckpt_param.grad = first_grad.clone() + ckpt_optimizer.step() + direct_param.grad = first_grad.clone() + direct_optimizer.step() + + saved_param = ckpt_param.detach().clone() + saved_state = copy.deepcopy(ckpt_optimizer.state_dict()) + + with torch.no_grad(): + loaded_param.copy_(saved_param) + loaded_optimizer.load_state_dict(saved_state) + + direct_param.grad = second_grad.clone() + direct_optimizer.step() + loaded_param.grad = second_grad.clone() + loaded_optimizer.step() + + torch.testing.assert_close(loaded_param, direct_param, atol=0.0, rtol=0.0) + direct_state = direct_optimizer.state_dict() + loaded_state = loaded_optimizer.state_dict() + assert loaded_state["step_count"] == direct_state["step_count"] + for key in ("master_params", "exp_avgs", "exp_avg_sqs"): + torch.testing.assert_close(loaded_state[key][0], direct_state[key][0], atol=0.0, rtol=0.0) + assert loaded_state["steps"] == direct_state["steps"] diff --git a/experimental/lite/tests/unit/primitive/test_module_primitives_independent_unit.py b/experimental/lite/tests/unit/primitive/test_module_primitives_independent_unit.py new file mode 100644 index 00000000000..9641cc26715 --- /dev/null +++ b/experimental/lite/tests/unit/primitive/test_module_primitives_independent_unit.py @@ -0,0 +1,131 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn + +from megatron.lite.primitive.modules.lora import ( + GroupedLinearLoRA, + LinearLoRA, + SharedGroupedLinearLoRA, + freeze_non_lora_params, + normalize_lora_config, + trainable_param_stats, +) + +pytestmark = pytest.mark.mlite + + +def test_lora_config_aliases_and_trainable_param_accounting(): + cfg = normalize_lora_config({"enabled": True, "rank": 2, "alpha": 6, "targets": ["qkv", "fc2"]}) + + assert cfg.enabled + assert cfg.scale == 3.0 + assert cfg.targets() == {"linear_qkv", "linear_fc2"} + assert cfg.targets_module("qkv") + assert cfg.targets_module("linear_fc2") + assert not normalize_lora_config({"enabled": False, "rank": 8}).enabled + with pytest.raises(TypeError, match="LoRA config"): + normalize_lora_config(object()) + + class TinyAdapterModel(nn.Module): + def __init__(self): + super().__init__() + self.base = nn.Linear(3, 2) + self.lora_adapter = nn.Linear(3, 2) + + model = TinyAdapterModel() + stats = freeze_non_lora_params(model) + + assert stats["lora_tensors"] == 2 + assert stats["frozen_tensors"] == 2 + assert not model.base.weight.requires_grad + assert model.lora_adapter.weight.requires_grad + assert trainable_param_stats(model) == { + "trainable_tensors": 2, + "trainable_numel": model.lora_adapter.weight.numel() + model.lora_adapter.bias.numel(), + } + + +def test_linear_lora_forward_backward_matches_low_rank_delta(): + layer = LinearLoRA(3, 2, rank=2, alpha=4, dropout=0.0) + with torch.no_grad(): + layer.lora_a.copy_(torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])) + layer.lora_b.copy_(torch.tensor([[1.0, 2.0], [3.0, 4.0]])) + + x = torch.tensor([[1.0, 2.0, 3.0]], requires_grad=True) + output = layer(x) + + torch.testing.assert_close(output, torch.tensor([[10.0, 22.0]])) + output.sum().backward() + torch.testing.assert_close(x.grad, torch.tensor([[8.0, 12.0, 0.0]])) + + +def test_grouped_lora_respects_per_expert_splits(): + layer = GroupedLinearLoRA(2, 2, 2, rank=1, alpha=1, dropout=0.0) + with torch.no_grad(): + layer.lora_a.copy_(torch.tensor([[[1.0, 0.0]], [[0.0, 1.0]]])) + layer.lora_b.copy_(torch.tensor([[[2.0], [3.0]], [[5.0], [7.0]]])) + + x = torch.tensor([[2.0, 9.0], [4.0, 1.0], [6.0, 3.0]]) + output = layer(x, [1, 2]) + + torch.testing.assert_close(output, torch.tensor([[4.0, 6.0], [5.0, 7.0], [15.0, 21.0]])) + with pytest.raises(ValueError, match="expected 2 splits"): + layer(x, [3]) + + +def test_shared_grouped_lora_uses_one_adapter_for_all_experts(): + layer = SharedGroupedLinearLoRA(2, 2, 2, rank=1, alpha=2, dropout=0.0) + with torch.no_grad(): + layer.lora_a.copy_(torch.tensor([[1.0, -1.0]])) + layer.lora_b.copy_(torch.tensor([[2.0], [3.0]])) + + output = layer(torch.tensor([[3.0, 1.0], [4.0, 7.0]]), [1, 1]) + + torch.testing.assert_close(output, torch.tensor([[8.0, 12.0], [-12.0, -18.0]])) + + +def test_mrope_interleaves_text_height_and_width_sections(): + from megatron.lite.primitive.modules.mrope import MultimodalRotaryEmbedding + + base = torch.arange(3 * 2 * 6, dtype=torch.float32).reshape(3, 2, 6) + + interleaved = MultimodalRotaryEmbedding._apply_interleaved_mrope(base, mrope_section=[1, 1, 1]) + + expected = base[0].clone() + expected[..., 1] = base[1, ..., 1] + expected[..., 2] = base[2, ..., 2] + torch.testing.assert_close(interleaved, expected) + + +def test_mtp_aux_loss_scaler_threads_independent_gradient(transformer_engine_import_stub): + transformer_engine_import_stub() + from megatron.lite.primitive.modules.mtp import MTPLossAutoScaler + + MTPLossAutoScaler.set_loss_scale(torch.tensor(0.125)) + output = torch.tensor([1.0, 2.0, 3.0], requires_grad=True) + mtp_loss = torch.tensor(4.0, requires_grad=True) + + MTPLossAutoScaler.apply(output * 3.0, mtp_loss).sum().backward() + + torch.testing.assert_close(output.grad, torch.full_like(output, 3.0)) + torch.testing.assert_close(mtp_loss.grad, torch.tensor(0.125)) + MTPLossAutoScaler.main_loss_backward_scale = 1.0 + + +def test_gated_delta_static_helpers_are_finite_and_shape_stable(transformer_engine_import_stub): + transformer_engine_import_stub() + from megatron.lite.primitive.modules.gated_delta_net import GatedDeltaNet + + alpha = torch.tensor([[[0.0, 1.0], [-1.0, 2.0]]]) + beta = torch.tensor([[[0.0, 2.0], [-2.0, 4.0]]]) + + g, beta_sigmoid = GatedDeltaNet._compute_g_and_beta(torch.zeros(2), torch.ones(2), alpha, beta) + + assert g.shape == alpha.shape + assert beta_sigmoid.shape == beta.shape + assert torch.isfinite(g).all() + assert torch.isfinite(beta_sigmoid).all() + assert torch.all(g < 0) diff --git a/experimental/lite/tests/unit/primitive/test_ops_data_trainstep_unit.py b/experimental/lite/tests/unit/primitive/test_ops_data_trainstep_unit.py new file mode 100644 index 00000000000..599595f35b8 --- /dev/null +++ b/experimental/lite/tests/unit/primitive/test_ops_data_trainstep_unit.py @@ -0,0 +1,193 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from megatron.lite.primitive.data import _resolve_thd_padding, fixed_batches +from megatron.lite.primitive.deterministic import deterministic_requested +from megatron.lite.primitive.ops.cross_entropy import vocab_parallel_cross_entropy +from megatron.lite.primitive.ops.gated_delta_rule import l2norm, torch_chunk_gated_delta_rule +from megatron.lite.primitive.ops.linear_cross_entropy import linear_cross_entropy +from megatron.lite.primitive.ops.logprob import vocab_parallel_log_probs_from_logits +from megatron.lite.primitive.recompute import apply_recompute, parse_recompute_spec +from megatron.lite.primitive.train_step import compute_and_clip_grad_norm, run_microbatch_loop +from megatron.lite.primitive.utils import ensure_divisible + +pytestmark = pytest.mark.mlite + + +def test_vocab_parallel_cross_entropy_matches_torch_cross_entropy(): + logits = torch.randn(2, 3, 5, dtype=torch.float64, requires_grad=True) + labels = torch.tensor([[0, 4, 2], [3, 1, 0]]) + + loss = vocab_parallel_cross_entropy(logits, labels) + expected = F.cross_entropy( + logits.float().reshape(-1, 5), labels.reshape(-1), reduction="none" + ).view_as(labels) + + torch.testing.assert_close(loss, expected) + loss.sum().backward() + assert logits.grad is not None + assert torch.isfinite(logits.grad).all() + + +def test_logprob_selects_labels_in_batch_sequence_or_sequence_batch_layout(): + logits = torch.randn(2, 3, 4) + labels_bsh = torch.tensor([[0, 1], [2, 3], [1, 0]]) + + selected = vocab_parallel_log_probs_from_logits(logits, labels_bsh) + expected = ( + torch.log_softmax(logits.float(), dim=-1) + .gather(-1, labels_bsh.transpose(0, 1).unsqueeze(-1)) + .squeeze(-1) + .transpose(0, 1) + .contiguous() + ) + + torch.testing.assert_close(selected, expected) + with pytest.raises(ValueError, match="Could not align"): + vocab_parallel_log_probs_from_logits(logits, torch.zeros(4, 2, dtype=torch.long)) + + +def test_linear_cross_entropy_fallback_matches_explicit_matmul(): + hidden = torch.tensor([[1.0, -2.0, 0.5], [0.25, 1.5, -0.5]]) + weight = torch.tensor( + [[0.5, -1.0, 0.25], [1.0, 0.0, -0.5], [-0.5, 0.75, 1.0], [0.25, 0.5, -1.5]] + ) + labels = torch.tensor([0, 3]) + + log_probs, entropy = linear_cross_entropy(hidden, weight, labels, temperature=2.0) + logits = hidden.matmul(weight.t()) / 2.0 + expected_loss = F.cross_entropy(logits, labels, reduction="none") + expected_entropy = torch.distributions.Categorical(logits=logits).entropy() + + torch.testing.assert_close(log_probs, -expected_loss) + torch.testing.assert_close(entropy, expected_entropy) + + +def test_gated_delta_rule_math_helpers_return_finite_stateful_outputs(): + x = torch.tensor([[3.0, 4.0], [0.0, 5.0]]) + normalized = l2norm(x, dim=-1, eps=0.0) + torch.testing.assert_close(normalized.norm(dim=-1), torch.ones(2)) + + query = torch.randn(1, 3, 1, 2) + key = torch.randn(1, 3, 1, 2) + value = torch.randn(1, 3, 1, 2) + g = -torch.rand(1, 3, 1) + beta = torch.rand(1, 3, 1) + + output, final_state = torch_chunk_gated_delta_rule( + query, + key, + value, + g, + beta, + chunk_size=2, + output_final_state=True, + use_qk_l2norm_in_kernel=False, + ) + + assert output.shape == value.shape + assert final_state is not None + assert final_state.shape == (1, 1, 2, 2) + assert torch.isfinite(output).all() + assert torch.isfinite(final_state).all() + + +def test_data_padding_env_and_fixed_batches_are_deterministic(monkeypatch): + monkeypatch.delenv("MEGATRON_LITE_THD_PAD_TO_ALIGNMENT", raising=False) + monkeypatch.delenv("MEGATRON_LITE_THD_PAD_MULTIPLE", raising=False) + assert _resolve_thd_padding(seq_len=7, cp_size=2) == (8, 4, True) + assert _resolve_thd_padding(seq_len=7, cp_size=1) == (7, 1, False) + + monkeypatch.setenv("MEGATRON_LITE_THD_PAD_TO_ALIGNMENT", "0") + monkeypatch.setenv("MEGATRON_LITE_THD_PAD_MULTIPLE", "8") + assert _resolve_thd_padding(seq_len=7, cp_size=2) == (7, 8, False) + monkeypatch.setenv("MEGATRON_LITE_THD_PAD_MULTIPLE", "bad") + with pytest.raises(ValueError, match="PAD_MULTIPLE"): + _resolve_thd_padding(seq_len=7, cp_size=2) + + batches_a = fixed_batches(11, seq_len=4, num_steps=2, batch_size=2, device="cpu", seed=123) + batches_b = fixed_batches(11, seq_len=4, num_steps=2, batch_size=2, device="cpu", seed=123) + for (ids_a, labels_a), (ids_b, labels_b) in zip(batches_a, batches_b, strict=True): + torch.testing.assert_close(ids_a, ids_b) + torch.testing.assert_close(labels_a, labels_b) + + +def test_deterministic_env_request_parser(monkeypatch): + monkeypatch.setenv("MEGATRON_LITE_DETERMINISTIC", "yes") + assert deterministic_requested() + monkeypatch.setenv("MEGATRON_LITE_DETERMINISTIC", "0") + assert not deterministic_requested() + + +def test_recompute_parser_and_wrapper_replays_forward_on_backward(): + assert parse_recompute_spec(None) == [] + assert parse_recompute_spec("none") == [] + assert parse_recompute_spec("full") == ["full"] + assert parse_recompute_spec("attn,mlp") == ["attn", "mlp"] + assert parse_recompute_spec(["attn"]) == ["attn"] + + class CountingModule(nn.Module): + def __init__(self): + super().__init__() + self.calls = 0 + + def forward(self, x): + self.calls += 1 + return x * x + + class Layer(nn.Module): + def __init__(self): + super().__init__() + self.inner = CountingModule() + + layer = Layer() + apply_recompute( + nn.ModuleList([layer]), + ["inner"], + {"inner": lambda module: module.inner}, + no_rng_modules={"inner"}, + ) + + x = torch.tensor([2.0, -3.0], requires_grad=True) + layer.inner(x).sum().backward() + + assert layer.inner.calls == 2 + torch.testing.assert_close(x.grad, torch.tensor([4.0, -6.0])) + + +def test_train_step_microbatch_loop_and_grad_clip_cpu_contract(): + model = nn.Linear(2, 1, bias=False) + with torch.no_grad(): + model.weight.fill_(0.5) + data = iter( + [ + {"x": torch.tensor([[1.0, 2.0]]), "y": torch.tensor([[1.0]])}, + {"x": torch.tensor([[3.0, 4.0]]), "y": torch.tensor([[2.0]])}, + ] + ) + + def forward_fn(module, batch): + return {"loss": F.mse_loss(module(batch["x"]), batch["y"])} + + output = run_microbatch_loop(model, data, 2, forward_fn) + + assert output is not None + assert output["loss"].ndim == 0 + assert model.weight.grad is not None + assert torch.isfinite(model.weight.grad).all() + + grad_norm = compute_and_clip_grad_norm(model, optimizer=None, max_norm=0.25, use_dist_opt=False) + + assert torch.isfinite(grad_norm) + assert model.weight.grad.norm() <= 0.25 + 1.0e-6 + + +def test_utils_ensure_divisible_returns_quotient_and_reports_context(): + assert ensure_divisible(12, 3, "tp") == 4 + with pytest.raises(ValueError, match=r"10 is not divisible by 4 \(tp\)"): + ensure_divisible(10, 4, "tp") diff --git a/experimental/lite/tests/unit/primitive/test_parallel_dimensions_independent_unit.py b/experimental/lite/tests/unit/primitive/test_parallel_dimensions_independent_unit.py new file mode 100644 index 00000000000..482d5449856 --- /dev/null +++ b/experimental/lite/tests/unit/primitive/test_parallel_dimensions_independent_unit.py @@ -0,0 +1,127 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from megatron.lite.primitive.parallel.cp import split_packed_for_cp +from megatron.lite.primitive.parallel.pipeline import _num_microbatches_from_config +from megatron.lite.primitive.parallel.pp import build_pipeline_chunk_layout +from megatron.lite.primitive.parallel.sp import ( + gather_for_non_sp_head, + gather_from_sequence_parallel, + scatter_to_sequence_parallel, +) +from megatron.lite.primitive.parallel.state import ParallelState + +pytestmark = pytest.mark.mlite + + +def test_tp_vocab_embedding_and_output_single_rank_contract(transformer_engine_import_stub): + transformer_engine_import_stub() + from megatron.lite.primitive.parallel.linear import ( + VocabParallelEmbedding, + VocabParallelOutput, + pad_vocab_for_tp, + ) + + ps = ParallelState(tp_size=1, tp_rank=0) + + assert pad_vocab_for_tp(151936, 2) == 151936 + assert pad_vocab_for_tp(129, 2) == 256 + + embedding = VocabParallelEmbedding(5, 3, ps, deterministic=True) + with torch.no_grad(): + values = torch.arange(embedding.local_vocab * 3, dtype=torch.float32).view( + embedding.local_vocab, 3 + ) + embedding.embedding.weight.copy_(values) + + input_ids = torch.tensor([[0, 4, 1]]) + output = embedding(input_ids) + + assert output.shape == (3, 1, 3) + torch.testing.assert_close(output[:, 0], values[input_ids[0]]) + + head = VocabParallelOutput(5, 3, ps) + logits = head(torch.ones(2, 1, 3, dtype=torch.bfloat16)) + gathered = head.gather(logits) + + assert logits.shape == (2, 1, 128) + assert gathered.shape == (2, 1, 5) + + +def test_sp_helpers_are_identity_for_single_rank_tp(): + ps = ParallelState(tp_size=1, tp_rank=0) + x = torch.randn(4, 2, 3, requires_grad=True) + + assert scatter_to_sequence_parallel(x, ps) is x + assert gather_from_sequence_parallel(x, ps) is x + assert gather_for_non_sp_head(x, ps) is x + + +def test_cp_packed_split_handles_each_sample_independently(): + input_ids = torch.arange(16) + position_ids = torch.arange(100, 116) + cu_seqlens = torch.tensor([0, 8, 16], dtype=torch.int32) + + rank0 = split_packed_for_cp( + input_ids, position_ids, cu_seqlens, max_seqlen=8, cp_rank=0, cp_size=2 + ) + rank1 = split_packed_for_cp( + input_ids, position_ids, cu_seqlens, max_seqlen=8, cp_rank=1, cp_size=2 + ) + + torch.testing.assert_close(rank0[0], torch.tensor([0, 1, 6, 7, 8, 9, 14, 15])) + torch.testing.assert_close(rank0[1], torch.tensor([100, 101, 106, 107, 108, 109, 114, 115])) + torch.testing.assert_close(rank0[2], torch.tensor([0, 4, 8], dtype=torch.int32)) + assert rank0[3] == 4 + + torch.testing.assert_close(rank1[0], torch.tensor([2, 3, 4, 5, 10, 11, 12, 13])) + torch.testing.assert_close(rank1[1], torch.tensor([102, 103, 104, 105, 110, 111, 112, 113])) + torch.testing.assert_close(rank1[2], torch.tensor([0, 4, 8], dtype=torch.int32)) + assert rank1[3] == 4 + + +def test_pp_layout_rejects_non_divisible_layer_counts(): + ps = ParallelState(pp_size=2, pp_rank=0, pp_is_first=True, pp_is_last=False) + + with pytest.raises(ValueError, match="not divisible"): + build_pipeline_chunk_layout(7, ps) + + +def test_dp_dimension_controls_dense_microbatch_contract(): + ps = ParallelState(dp_size=4, dp_rank=2, dp_cp_size=8, dp_cp_rank=5) + config = SimpleNamespace(gbs=32, mbs=2, num_microbatches=None) + + assert _num_microbatches_from_config(config, ps) == 4 + + config.num_microbatches = 7 + assert _num_microbatches_from_config(config, ps) == 7 + assert ps.dp_rank == 2 + assert ps.dp_cp_size == 8 + assert ps.dp_cp_rank == 5 + + +def test_ep_token_dispatcher_local_roundtrip_is_independent_of_deepep(): + from megatron.lite.primitive.modules.dispatcher import TokenDispatcher + + ps = ParallelState(ep_size=1, ep_rank=0) + dispatcher = TokenDispatcher(num_experts=3, hidden_size=2, ps=ps, use_deepep=False) + hidden = torch.tensor([[1.0, 10.0], [2.0, 20.0], [3.0, 30.0], [4.0, 40.0]], requires_grad=True) + topk_indices = torch.tensor([[0], [2], [1], [2]]) + topk_scores = torch.ones(4, 1) + + dispatched, tokens_per_expert, dispatched_probs = dispatcher.dispatch( + hidden, topk_scores, topk_indices + ) + combined = dispatcher.combine(dispatched) + + torch.testing.assert_close(tokens_per_expert, torch.tensor([1, 1, 2])) + torch.testing.assert_close(dispatched_probs, torch.ones(4)) + torch.testing.assert_close(combined, hidden) + + combined.sum().backward() + torch.testing.assert_close(hidden.grad, torch.ones_like(hidden)) diff --git a/experimental/lite/tests/unit/primitive/test_parallel_unit.py b/experimental/lite/tests/unit/primitive/test_parallel_unit.py new file mode 100644 index 00000000000..6eb42553542 --- /dev/null +++ b/experimental/lite/tests/unit/primitive/test_parallel_unit.py @@ -0,0 +1,86 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import pytest +import torch + +from megatron.lite.primitive.parallel import ( + ParallelState, + build_pipeline_chunk_layout, + reconstruct_packed_from_cp_parts, + roll_packed_thd_left, + split_packed_to_cp_local, + zigzag_position_ids_for_cp, + zigzag_reconstruct_from_cp_parts, + zigzag_slice_for_cp, + zigzag_split_for_cp, +) + +pytestmark = pytest.mark.mlite + + +def test_cp_zigzag_split_slice_and_reconstruct_match(): + tensor = torch.arange(16).reshape(1, 8, 2) + parts = [zigzag_split_for_cp(tensor, rank, cp_size=2, seq_dim=1) for rank in range(2)] + + assert torch.equal(parts[0], tensor[:, [0, 1, 6, 7], :]) + assert torch.equal(parts[1], tensor[:, [2, 3, 4, 5], :]) + assert torch.equal(zigzag_slice_for_cp(tensor, 0, cp_size=2, seq_dim=1), parts[0]) + assert torch.equal(zigzag_slice_for_cp(tensor, 1, cp_size=2, seq_dim=1), parts[1]) + assert torch.equal(zigzag_reconstruct_from_cp_parts(parts, seq_dim=1), tensor) + + +def test_cp_position_ids_follow_zigzag_order(): + assert torch.equal( + zigzag_position_ids_for_cp(8, cp_rank=0, cp_size=2, device=torch.device("cpu")), + torch.tensor([[0, 1, 6, 7]]), + ) + assert torch.equal( + zigzag_position_ids_for_cp(8, cp_rank=1, cp_size=2, device=torch.device("cpu")), + torch.tensor([[2, 3, 4, 5]]), + ) + + +def test_pp_layout_marks_stage_boundaries_and_vpp_chunks(): + rank0 = ParallelState(pp_size=2, pp_rank=0, pp_is_first=True, pp_is_last=False) + rank1 = ParallelState(pp_size=2, pp_rank=1, pp_is_first=False, pp_is_last=True) + + assert build_pipeline_chunk_layout(8, rank0).layer_indices == [0, 1, 2, 3] + assert build_pipeline_chunk_layout(8, rank0).has_embed is True + assert build_pipeline_chunk_layout(8, rank0).has_head is False + assert build_pipeline_chunk_layout(8, rank1).layer_indices == [4, 5, 6, 7] + assert build_pipeline_chunk_layout(8, rank1).has_embed is False + assert build_pipeline_chunk_layout(8, rank1).has_head is True + + vpp_rank0_chunk1 = build_pipeline_chunk_layout(8, rank0, vpp=2, vpp_chunk_id=1) + vpp_rank1_chunk1 = build_pipeline_chunk_layout(8, rank1, vpp=2, vpp_chunk_id=1) + assert vpp_rank0_chunk1.layer_indices == [4, 5] + assert vpp_rank0_chunk1.has_head is False + assert vpp_rank1_chunk1.layer_indices == [6, 7] + assert vpp_rank1_chunk1.has_head is True + + +def test_thd_roll_keeps_sequence_boundaries(): + cu_seqlens = torch.tensor([0, 4, 8], dtype=torch.int32) + rolled, token_sum = roll_packed_thd_left(torch.arange(8), cu_seqlens_padded=cu_seqlens, dims=0) + + assert torch.equal(rolled, torch.tensor([1, 2, 3, 0, 5, 6, 7, 0])) + assert token_sum.item() == 24 + + +def test_thd_cp_split_and_reconstruct_roundtrip(): + cu_seqlens = torch.tensor([0, 8], dtype=torch.int32) + tensor = torch.arange(8) + parts = [ + split_packed_to_cp_local( + tensor, cu_seqlens_padded=cu_seqlens, cp_size=2, cp_rank=rank, dim=0 + ) + for rank in range(2) + ] + + assert torch.equal(parts[0], torch.tensor([0, 1, 6, 7])) + assert torch.equal(parts[1], torch.tensor([2, 3, 4, 5])) + assert torch.equal( + reconstruct_packed_from_cp_parts(parts, cu_seqlens_padded=cu_seqlens, cp_size=2, dim=0), + tensor, + ) diff --git a/experimental/lite/tests/unit/primitive/test_runtime_config_unit.py b/experimental/lite/tests/unit/primitive/test_runtime_config_unit.py new file mode 100644 index 00000000000..6e91f51c1b3 --- /dev/null +++ b/experimental/lite/tests/unit/primitive/test_runtime_config_unit.py @@ -0,0 +1,77 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import pytest + +from megatron.lite.runtime import RuntimeConfig, create_runtime +from megatron.lite.runtime.backends.mlite.config import DebugConfig, MegatronLiteConfig +from megatron.lite.runtime.contracts.config import OptimizerConfig, ParallelConfig + +pytestmark = pytest.mark.mlite + + +def test_mlite_config_defaults_are_stable(): + cfg = MegatronLiteConfig(model_name="qwen3_moe") + + assert cfg.model_name == "qwen3_moe" + assert cfg.impl == "lite" + assert cfg.parallel.tp == 1 + assert cfg.parallel.ep == 1 + assert cfg.parallel.pp == 1 + assert cfg.parallel.cp == 1 + assert isinstance(cfg.optimizer, OptimizerConfig) + assert isinstance(cfg.debug, DebugConfig) + + +def test_mlite_config_from_dict_preserves_parallel_optimizer_and_impl_cfg(): + cfg = MegatronLiteConfig.from_dict( + "/models/qwen", + { + "model_name": "qwen3_moe", + "impl": "lite", + "tp": 2, + "ep": 4, + "pp": 2, + "cp": 2, + "optimizer": { + "lr": 1.0e-4, + "weight_decay": 0.1, + "adam_beta1": 0.9, + "offload_fraction": 1.0, + }, + "impl_cfg": {"attn_impl": "mcore", "moe_impl": "ml"}, + "use_thd": True, + "precision_aware_opt": True, + }, + ) + + assert cfg.hf_path == "/models/qwen" + assert cfg.parallel == ParallelConfig(tp=2, etp=None, ep=4, pp=2, vpp=1, cp=2) + assert cfg.optimizer.lr == 1.0e-4 + assert cfg.optimizer.weight_decay == 0.1 + assert cfg.optimizer.adam_beta1 == 0.9 + assert cfg.optimizer.offload_fraction == 1.0 + assert cfg.impl_cfg["attn_impl"] == "mcore" + assert cfg.impl_cfg["moe_impl"] == "ml" + assert cfg.impl_cfg["use_thd"] is True + assert cfg.impl_cfg["precision_aware_opt"] is True + + +def test_mlite_config_rejects_num_microbatches_in_backend_config(): + with pytest.raises(ValueError, match="num_microbatches"): + MegatronLiteConfig.from_dict( + "/models/qwen", {"model_name": "qwen3_moe", "num_microbatches": 2} + ) + + +def test_create_runtime_uses_mlite_backend_registry(): + runtime = create_runtime( + RuntimeConfig( + backend="mlite", + hf_path="/models/qwen", + backend_cfg={"model_name": "qwen3_moe", "load_hf_weights": False}, + ) + ) + + assert type(runtime).__name__ == "MegatronLiteRuntime" + assert runtime.tier == "rl_best" diff --git a/experimental/lite/tests/unit/primitive/test_training_checkpoint.py b/experimental/lite/tests/unit/primitive/test_training_checkpoint.py new file mode 100644 index 00000000000..f07cda7858d --- /dev/null +++ b/experimental/lite/tests/unit/primitive/test_training_checkpoint.py @@ -0,0 +1,370 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import copy +from types import SimpleNamespace + +import torch +from torch.distributed.tensor import Replicate, Shard + +from megatron.core.dist_checkpointing.strategies.torch import ( + _replace_state_dict_keys_with_sharded_keys, +) +from megatron.lite.primitive.ckpt import dcp +from megatron.lite.primitive.ckpt.distckpt import ( + _model_sharded_state_dict, + _rank_offsets_and_replica_id, + _single_or_all_model_state, + _synchronize_native_optimizer_steps, + attach_model_sharded_state_dict, +) +from megatron.lite.primitive.parallel import ParallelState +from megatron.lite.primitive.protocols import default_expert_classifier, default_placement_fn +from megatron.lite.runtime.backends.mlite.runtime import MegatronLiteRuntime +from megatron.lite.runtime.contracts.handle import ModelHandle + + +def _assert_state_equal(actual, expected) -> None: + if torch.is_tensor(expected): + assert torch.equal(actual, expected) + elif isinstance(expected, dict): + assert actual.keys() == expected.keys() + for key, value in expected.items(): + _assert_state_equal(actual[key], value) + elif isinstance(expected, list): + assert len(actual) == len(expected) + for actual_item, expected_item in zip(actual, expected, strict=True): + _assert_state_equal(actual_item, expected_item) + else: + assert actual == expected + + +def test_optimizer_checkpoint_roundtrips_rank_local_state(tmp_path) -> None: + model = torch.nn.Linear(4, 2) + optimizer = torch.optim.AdamW(model.parameters(), lr=0.01) + + loss = model(torch.ones(3, 4)).sum() + loss.backward() + optimizer.step() + + expected = copy.deepcopy(optimizer.state_dict()) + dcp._save_optimizer_checkpoint(optimizer, str(tmp_path)) + + for state in optimizer.state.values(): + for value in state.values(): + if torch.is_tensor(value): + value.zero_() + + dcp._load_optimizer_checkpoint(optimizer, str(tmp_path)) + + assert (tmp_path / "optimizer_rank_0.pt").exists() + _assert_state_equal(optimizer.state_dict(), expected) + + +class FakeDistOpt: + def __init__(self): + self.save_model_sd = None + self.load_model_sd = None + self.loaded_state = None + + def sharded_state_dict(self, model_sd, is_loading: bool = False, metadata=None): + assert metadata == DISTOPT_METADATA + if is_loading: + self.load_model_sd = model_sd + else: + self.save_model_sd = model_sd + return {"is_loading": is_loading} + + def load_state_dict(self, state): + self.loaded_state = state + + +class FakeWrapper(torch.nn.Module): + def __init__(self, module): + super().__init__() + self.module = module + self.wrapper_load_called = False + + def forward(self, *args, **kwargs): + return self.module(*args, **kwargs) + + def load_state_dict(self, *args, **kwargs): + self.wrapper_load_called = True + return super().load_state_dict(*args, **kwargs) + + +DISTOPT_METADATA = { + "distrib_optim_sharding_type": "fully_reshardable", + "distrib_optim_fully_reshardable_mem_efficient": False, + "chained_optim_avoid_prefix": True, +} + + +def test_distopt_checkpoint_dispatches_to_mcore_distckpt(monkeypatch, tmp_path) -> None: + model = torch.nn.Linear(4, 2) + optimizer = FakeDistOpt() + ps = ParallelState(pp_rank=1, tp_rank=2, dp_cp_rank=3) + attach_model_sharded_state_dict([model], ps) + saved = {} + + def fake_save(state_dict, checkpoint_dir, **kwargs): + saved["state_dict"] = state_dict + saved["checkpoint_dir"] = checkpoint_dir + saved["kwargs"] = kwargs + + monkeypatch.setattr("megatron.lite.primitive.ckpt.distckpt.dist_checkpointing.save", fake_save) + + dcp.save_training_checkpoint(model, optimizer, 5, str(tmp_path), use_dcp=True) + + model_sd = saved["state_dict"]["model"] + assert set(model_sd) == {"weight", "bias"} + assert model_sd["weight"].replica_id == (0, 2, 3) + assert optimizer.save_model_sd is model_sd + assert saved["state_dict"]["optimizer"] == {"is_loading": False} + assert saved["state_dict"]["step"] == 5 + assert saved["checkpoint_dir"] == str(tmp_path / "step_5") + assert saved["kwargs"]["validate_access_integrity"] is False + assert saved["kwargs"]["content_metadata"] == DISTOPT_METADATA + assert not (tmp_path / "step_5" / "optimizer_rank_0.pt").exists() + + +def test_distopt_checkpoint_offsets_cover_tp_pp_ep_etp_topology() -> None: + ps = ParallelState( + pp_size=2, + pp_rank=1, + tp_size=2, + tp_rank=1, + ep_size=2, + ep_rank=1, + etp_size=2, + etp_rank=1, + dp_size=2, + dp_rank=0, + cp_size=1, + cp_rank=0, + dp_cp_rank=0, + expert_dp_size=1, + expert_dp_rank=0, + ) + + dense_offsets, dense_replica = _rank_offsets_and_replica_id( + [Replicate(), Replicate(), Replicate(), Shard(0)], ps, expert=False + ) + expert_offsets, expert_replica = _rank_offsets_and_replica_id( + [Replicate(), Replicate(), Shard(0), Shard(0)], ps, expert=True + ) + + assert dense_offsets == ((0, 1, 2),) + assert dense_replica == (0, 0, 0) + assert expert_offsets == ((0, 3, 4),) + assert expert_replica == (0, 0, 0) + + +def test_distopt_replica_id_groups_sharded_axes_by_placement() -> None: + placements = [Replicate(), Replicate(), Replicate(), Shard(0)] + rank_offsets0, replica_id0 = _rank_offsets_and_replica_id( + placements, ParallelState(tp_size=2, tp_rank=0), expert=False + ) + rank_offsets1, replica_id1 = _rank_offsets_and_replica_id( + placements, ParallelState(tp_size=2, tp_rank=1), expert=False + ) + + assert rank_offsets0 == ((0, 0, 2),) + assert rank_offsets1 == ((0, 1, 2),) + assert replica_id0 == replica_id1 == (0, 0, 0) + + expert_offsets, expert_replica_id = _rank_offsets_and_replica_id( + [Replicate(), Replicate(), Shard(0), Shard(1)], + ParallelState(ep_size=2, ep_rank=1, etp_size=2, etp_rank=1), + expert=True, + ) + + assert expert_offsets == ((0, 1, 2), (1, 1, 2)) + assert expert_replica_id == (0, 0, 0) + + +def test_distopt_replica_id_does_not_treat_pp_as_a_replica_axis() -> None: + rank_offsets, replica_id = _rank_offsets_and_replica_id( + [Replicate(), Replicate(), Replicate(), Shard(0)], + ParallelState(pp_size=2, pp_rank=1, tp_size=2, tp_rank=1), + expert=False, + ) + + assert rank_offsets == ((0, 1, 2),) + assert replica_id == (0, 0, 0) + + _rank_offsets, replica_id = _rank_offsets_and_replica_id( + [Replicate(), Replicate(), Replicate(), Replicate()], + ParallelState(pp_size=2, pp_rank=1, tp_size=2, tp_rank=0), + expert=False, + ) + + assert replica_id == (0, 0, 0) + + +def test_distopt_pp_rank_one_model_keys_survive_torch_dist_main_replica_filter() -> None: + ps = ParallelState(pp_size=2, pp_rank=1, pp_is_first=False, pp_is_last=True) + model = torch.nn.Linear(4, 2) + attach_model_sharded_state_dict([model], ps) + + model_sd = _model_sharded_state_dict(model) + filtered_sd, _flat_mapping, _rename_mapping = _replace_state_dict_keys_with_sharded_keys( + model_sd, keep_only_main_replica=True + ) + + assert set(filtered_sd) == {"model_pp1.weight", "model_pp1.bias"} + + +def test_distopt_model_state_keys_are_pp_and_vpp_aware() -> None: + ps = ParallelState(pp_size=2, pp_rank=1, pp_is_first=False, pp_is_last=True) + single_chunk = torch.nn.Linear(4, 2) + attach_model_sharded_state_dict([single_chunk], ps) + + single_sd = _model_sharded_state_dict(single_chunk) + + assert set(single_sd) == {"model_pp1"} + assert set(single_sd["model_pp1"]) == {"weight", "bias"} + assert single_sd["model_pp1"]["weight"].key == "model_pp1.weight" + assert _single_or_all_model_state(single_sd) is single_sd + + chunks = [torch.nn.Linear(4, 2), torch.nn.Linear(4, 2)] + attach_model_sharded_state_dict(chunks, ps) + + vpp_sd = _model_sharded_state_dict(chunks) + + assert set(vpp_sd) == {"model_pp1_vpp0", "model_pp1_vpp1"} + assert set(vpp_sd["model_pp1_vpp0"]) == {"weight", "bias"} + assert set(vpp_sd["model_pp1_vpp1"]) == {"weight", "bias"} + assert vpp_sd["model_pp1_vpp0"]["weight"].key == "model_pp1_vpp0.weight" + assert vpp_sd["model_pp1_vpp1"]["weight"].key == "model_pp1_vpp1.weight" + assert _single_or_all_model_state(vpp_sd) is vpp_sd + + +def test_distopt_checkpoint_loads_from_mcore_distckpt(monkeypatch, tmp_path) -> None: + wrapped_module = torch.nn.Linear(4, 2) + model = FakeWrapper(wrapped_module) + optimizer = FakeDistOpt() + attach_model_sharded_state_dict([model], ParallelState()) + expected_weight = torch.full_like(wrapped_module.weight, 3.0) + expected_bias = torch.full_like(wrapped_module.bias, -2.0) + + def fake_load(sharded_state_dict, checkpoint_dir, **kwargs): + assert set(sharded_state_dict["model"]) == {"weight", "bias"} + assert optimizer.load_model_sd is sharded_state_dict["model"] + assert sharded_state_dict["optimizer"] == {"is_loading": True} + assert checkpoint_dir == str(tmp_path / "step_5") + assert kwargs["validate_access_integrity"] is False + return { + "step": 5, + "model": {"weight": expected_weight, "bias": expected_bias}, + "optimizer": {"loaded": True}, + } + + monkeypatch.setattr("megatron.lite.primitive.ckpt.distckpt.dist_checkpointing.load", fake_load) + + step = dcp.load_training_checkpoint(model, optimizer, str(tmp_path / "step_5"), use_dcp=True) + + assert step == 5 + assert not model.wrapper_load_called + torch.testing.assert_close(wrapped_module.weight, expected_weight) + torch.testing.assert_close(wrapped_module.bias, expected_bias) + assert optimizer.loaded_state == {"loaded": True} + + +def test_distopt_step_sync_traverses_multi_optimizer_chain_without_optimizer_property() -> None: + class FakeTorchOptimizer: + def __init__(self, steps): + self.state = { + object(): {"step": torch.tensor(step, dtype=torch.int64)} for step in steps + } + + class FakeDistOpt: + def __init__(self, steps): + self.optimizer = FakeTorchOptimizer(steps) + + class FakeChainedOptimizer: + def __init__(self): + self.chained_optimizers = [FakeDistOpt([1, 3]), FakeDistOpt([2, 4])] + + @property + def optimizer(self): + raise AssertionError( + "ChainedOptimizer has more than one optimizer when accessing self.optimizer" + ) + + chained = FakeChainedOptimizer() + + _synchronize_native_optimizer_steps(chained) + + for child in chained.chained_optimizers: + steps = [int(state["step"].item()) for state in child.optimizer.state.values()] + assert steps == [max(steps)] * len(steps) + + +def test_runtime_checkpoint_api_passes_current_training_checkpoint_signature( + monkeypatch, tmp_path +) -> None: + calls = {} + + def fake_save(model, optimizer, step, path, config, ps, **kwargs): + calls["save"] = (model, optimizer, step, path, config, ps, kwargs) + + def fake_load(model, optimizer, path, config, ps, **kwargs): + calls["load"] = (model, optimizer, path, config, ps, kwargs) + return 7 + + monkeypatch.setattr("megatron.lite.primitive.ckpt.save_training_checkpoint", fake_save) + monkeypatch.setattr("megatron.lite.primitive.ckpt.load_training_checkpoint", fake_load) + + runtime = MegatronLiteRuntime.__new__(MegatronLiteRuntime) + model = torch.nn.Linear(1, 1) + optimizer = object() + parallel = SimpleNamespace(tp=1, etp=1, ep=1, pp=1, cp=1) + ps = object() + handle = ModelHandle( + model=model, + optimizer=optimizer, + parallel_state=ps, + config=SimpleNamespace(parallel=parallel), + ) + + runtime.save_checkpoint( + handle, str(tmp_path), global_step=7, save_model=True, save_optimizer=False + ) + loaded_step = runtime.load_checkpoint( + handle, str(tmp_path), load_model=False, load_optimizer=True + ) + + assert calls["save"] == ( + model, + optimizer, + 7, + str(tmp_path), + parallel, + ps, + { + "get_placements": default_placement_fn, + "is_expert": default_expert_classifier, + "use_dcp": True, + "save_rng": True, + "save_model": True, + "save_optimizer": False, + }, + ) + assert calls["load"] == ( + model, + optimizer, + str(tmp_path), + parallel, + ps, + { + "get_placements": default_placement_fn, + "is_expert": default_expert_classifier, + "use_dcp": True, + "load_rng": True, + "load_parameter_state_update_legacy_format": False, + "load_model": False, + "load_optimizer": True, + }, + ) + assert loaded_step == 7 diff --git a/experimental/lite/tests/unit/runtime/test_bridge_backend.py b/experimental/lite/tests/unit/runtime/test_bridge_backend.py new file mode 100644 index 00000000000..90368fe0c56 --- /dev/null +++ b/experimental/lite/tests/unit/runtime/test_bridge_backend.py @@ -0,0 +1,190 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Unit tests for the Megatron-Bridge runtime backend surface.""" + +from __future__ import annotations + +import os +import sys +import types + +import pytest + + +def test_bridge_runtime_registered_and_lazy_constructible(): + from megatron.lite.runtime import RuntimeConfig, create_runtime + from megatron.lite.runtime.backends import RUNTIME_REGISTRY + from megatron.lite.runtime.backends.bridge.config import BridgeConfig + from megatron.lite.runtime.backends.bridge.runtime import BridgeRuntime + + assert RUNTIME_REGISTRY["bridge"] == "megatron.lite.runtime.backends.bridge" + + runtime = create_runtime( + RuntimeConfig( + backend="bridge", + hf_path="/tmp/hf-model", + backend_cfg=BridgeConfig(model_name="qwen3_5"), + ) + ) + + assert isinstance(runtime, BridgeRuntime) + assert runtime.tier == "rl_best" + + +def test_mbridge_runtime_registered_and_lazy_constructible(): + from megatron.lite.runtime import RuntimeConfig, create_runtime + from megatron.lite.runtime.backends import RUNTIME_REGISTRY + from megatron.lite.runtime.backends.bridge.config import BridgeConfig + from megatron.lite.runtime.backends.mbridge.runtime import MBridgeRuntime + + assert RUNTIME_REGISTRY["mbridge"] == "megatron.lite.runtime.backends.mbridge" + + runtime = create_runtime( + RuntimeConfig( + backend="mbridge", + hf_path="/tmp/hf-model", + backend_cfg=BridgeConfig(model_name="qwen3_5"), + ) + ) + + assert isinstance(runtime, MBridgeRuntime) + assert runtime.tier == "rl_best" + + +def test_bridge_config_from_dict_accepts_nested_and_flat_parallel_fields(): + from megatron.lite.runtime.backends.bridge.config import BridgeConfig + + cfg = BridgeConfig.from_dict( + { + "model_name": "qwen3_5", + "parallel": {"tp": 2, "pp": 1}, + "ep": 4, + "optimizer": {"lr": 2e-4, "weight_decay": 0.2}, + "lr_scheduler": {"total_training_steps": 16}, + "override_transformer_config": {"attention_backend": "unfused"}, + } + ) + + assert cfg.model_name == "qwen3_5" + assert cfg.parallel.tp == 2 + assert cfg.parallel.ep == 4 + assert cfg.optimizer.lr == 2e-4 + assert cfg.optimizer.weight_decay == 0.2 + assert cfg.optimizer.total_training_steps == 16 + assert cfg.override_transformer_config == {"attention_backend": "unfused"} + + +def test_bridge_config_from_dict_rejects_num_microbatches(): + from megatron.lite.runtime.backends.bridge.config import BridgeConfig + + with pytest.raises(ValueError, match="num_microbatches"): + BridgeConfig.from_dict({"num_microbatches": 2}) + + +def test_bridge_builds_mcore_ddp_config_object(monkeypatch): + from megatron.lite.runtime.backends.bridge.config import BridgeConfig + from megatron.lite.runtime.backends.bridge.runtime import _build_ddp_config + + class _FakeDDPConfig: + def __init__(self, **kwargs): + for key, value in kwargs.items(): + setattr(self, key, value) + + monkeypatch.setitem( + sys.modules, + "megatron.core.distributed", + types.SimpleNamespace(DistributedDataParallelConfig=_FakeDDPConfig), + ) + + ddp_config = _build_ddp_config( + BridgeConfig(override_ddp_config={"overlap_grad_reduce": True, "bucket_size": 1024}) + ) + + assert isinstance(ddp_config, _FakeDDPConfig) + assert ddp_config.use_distributed_optimizer is True + assert ddp_config.grad_reduce_in_fp32 is True + assert ddp_config.overlap_grad_reduce is True + assert ddp_config.bucket_size == 1024 + + +def test_bridge_deterministic_provider_sets_te_env(monkeypatch): + from megatron.lite.runtime.backends.bridge.config import BridgeConfig + from megatron.lite.runtime.backends.bridge.runtime import _configure_provider + + monkeypatch.setenv("MEGATRON_LITE_DETERMINISTIC", "1") + monkeypatch.delenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", raising=False) + provider = types.SimpleNamespace() + + _configure_provider(provider, BridgeConfig()) + + assert provider.deterministic_mode is True + assert "NVTE_ALLOW_NONDETERMINISTIC_ALGO" in os.environ + assert os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] == "0" + + +def test_bridge_registers_qwen35_moe_compat_aliases(monkeypatch): + from megatron.lite.runtime.backends.bridge.runtime import _register_bridge_compat_aliases + + registered = [] + + class _FakeMapping: + def __init__(self, *args, **kwargs): + self.args = args + self.kwargs = kwargs + + @classmethod + def register_module_type(cls, *args, **kwargs): + return None + + class _FakeQwen3NextBridge: + def provider_bridge(self, hf_pretrained): + return types.SimpleNamespace() + + class _FakeDispatcher: + _exact_types = {} + + model_bridge = types.SimpleNamespace( + get_model_bridge=_FakeDispatcher(), + register_bridge_implementation=lambda **kwargs: registered.append(kwargs), + ) + mapping_registry_mod = types.SimpleNamespace(MegatronMappingRegistry=lambda *items: list(items)) + param_mapping_mod = types.SimpleNamespace( + AutoMapping=_FakeMapping, + GDNConv1dMapping=_FakeMapping, + GatedMLPMapping=_FakeMapping, + QKVMapping=_FakeMapping, + ReplicatedMapping=_FakeMapping, + RMSNorm2ZeroCenteredRMSNormMapping=_FakeMapping, + merge_gdn_linear_weights=lambda *args, **kwargs: None, + split_gdn_linear_weights=lambda *args, **kwargs: (None, None), + ) + qwen_bridge_mod = types.SimpleNamespace(Qwen3NextBridge=_FakeQwen3NextBridge) + common_utils_mod = types.SimpleNamespace(extract_expert_number_from_param=lambda name: 0) + gpt_mod = types.SimpleNamespace(GPTModel=object) + + monkeypatch.setitem( + sys.modules, + "megatron.bridge.models.conversion", + types.SimpleNamespace(model_bridge=model_bridge), + ) + monkeypatch.setitem( + sys.modules, "megatron.bridge.models.conversion.mapping_registry", mapping_registry_mod + ) + monkeypatch.setitem( + sys.modules, "megatron.bridge.models.conversion.param_mapping", param_mapping_mod + ) + monkeypatch.setitem( + sys.modules, "megatron.bridge.models.qwen.qwen3_next_bridge", qwen_bridge_mod + ) + monkeypatch.setitem(sys.modules, "megatron.bridge.utils.common_utils", common_utils_mod) + monkeypatch.setitem(sys.modules, "megatron.core.models.gpt.gpt_model", gpt_mod) + + _register_bridge_compat_aliases() + + assert [item["source"] for item in registered] == [ + "Qwen3_5MoeForConditionalGeneration", + "Qwen3_5MoeForCausalLM", + ] + assert all(issubclass(item["bridge_class"], _FakeQwen3NextBridge) for item in registered) + assert all(item["bridge_class"] is not _FakeQwen3NextBridge for item in registered) + assert all("provider_bridge" in item["bridge_class"].__dict__ for item in registered) + assert all("mapping_registry" in item["bridge_class"].__dict__ for item in registered) diff --git a/experimental/lite/tests/unit/runtime/test_runtime_backend_unit.py b/experimental/lite/tests/unit/runtime/test_runtime_backend_unit.py new file mode 100644 index 00000000000..c71df0f4bff --- /dev/null +++ b/experimental/lite/tests/unit/runtime/test_runtime_backend_unit.py @@ -0,0 +1,295 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from __future__ import annotations + +import os +import subprocess +from dataclasses import dataclass +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +import torch.nn as nn + +from megatron.lite.runtime import create_runtime +from megatron.lite.runtime.backends.mlite.config import MegatronLiteConfig +from megatron.lite.runtime.backends.mlite.runtime import ( + MegatronLiteRuntime, + _apply_attention_backend_env, + _build_impl_cfg, +) +from megatron.lite.runtime.contracts.config import OptimizerConfig, ParallelConfig, RuntimeConfig +from megatron.lite.runtime.contracts.handle import ModelHandle + +pytestmark = pytest.mark.mlite + + +def test_runtime_config_defaults_to_mlite_backend(): + cfg = RuntimeConfig() + + assert cfg.backend == "mlite" + assert cfg.hf_path == "" + assert isinstance(cfg.backend_cfg, dict) + + +def test_runtime_config_accepts_mlite_backend_cfg(): + cfg = RuntimeConfig( + backend="mlite", + hf_path="/models/Qwen3", + backend_cfg={"model_name": "qwen3", "impl": "lite", "tp": 2, "ep": 4}, + ) + + assert cfg.backend == "mlite" + assert cfg.backend_cfg["model_name"] == "qwen3" + assert cfg.backend_cfg["tp"] == 2 + + +def test_mlite_config_defaults_and_parallel_fields(): + cfg = MegatronLiteConfig( + model_name="qwen3_moe", parallel=ParallelConfig(tp=4, etp=1, ep=8, pp=2, vpp=2, cp=2) + ) + + assert cfg.model_name == "qwen3_moe" + assert cfg.impl == "lite" + assert cfg.parallel.tp == 4 + assert cfg.parallel.ep == 8 + assert cfg.parallel.pp == 2 + assert cfg.parallel.cp == 2 + + +def test_mlite_config_impl_cfg_optimizer_and_load_gate(): + hook = lambda cfg: cfg # noqa: E731 + cfg = MegatronLiteConfig( + model_name="qwen3_moe", + impl_cfg={"recompute": "full", "use_deepep": True}, + optimizer=OptimizerConfig(lr=1e-4, weight_decay=0.1, adam_beta1=0.9), + load_hf_weights=False, + model_config_hook=hook, + ) + + assert cfg.impl_cfg["recompute"] == "full" + assert cfg.impl_cfg["use_deepep"] is True + assert cfg.optimizer.lr == 1e-4 + assert cfg.optimizer.adam_beta1 == 0.9 + assert cfg.load_hf_weights is False + assert cfg.model_config_hook is hook + + +def test_mlite_config_from_dict_accepts_optimizer_override_config(): + cfg = MegatronLiteConfig.from_dict( + "/models/Qwen3", + { + "optimizer": { + "override_optimizer_config": { + "fsdp2_use_fp32_master": False, + "offload_fraction": 1.0, + } + } + }, + ) + + assert cfg.optimizer.override_optimizer_config == { + "fsdp2_use_fp32_master": False, + "offload_fraction": 1.0, + } + + +def test_mlite_config_from_dict_rejects_num_microbatches(): + with pytest.raises(ValueError, match="num_microbatches"): + MegatronLiteConfig.from_dict( + "/models/Qwen3", {"model_name": "qwen3", "tp": 4, "num_microbatches": 2} + ) + + +@dataclass +class _FakeImplConfig: + parallel: object + hf_path: str = "" + optimizer_config: object = None + attention_backend_override: str | None = None + + +def test_build_impl_cfg_backfills_top_level_hf_path_and_runtime_fields(): + proto = type("Proto", (), {"ImplConfig": _FakeImplConfig}) + cfg = MegatronLiteConfig( + model_name="qwen3", hf_path="/models/top", attention_backend_override="local" + ) + + impl_cfg = _build_impl_cfg(proto, cfg) + + assert impl_cfg.parallel is cfg.parallel + assert impl_cfg.hf_path == "/models/top" + assert impl_cfg.optimizer_config is cfg.optimizer + assert impl_cfg.attention_backend_override == "local" + + +def test_build_impl_cfg_preserves_explicit_impl_hf_path(): + proto = type("Proto", (), {"ImplConfig": _FakeImplConfig}) + cfg = MegatronLiteConfig( + model_name="qwen3", hf_path="/models/top", impl_cfg={"hf_path": "/models/impl"} + ) + + impl_cfg = _build_impl_cfg(proto, cfg) + + assert impl_cfg.hf_path == "/models/impl" + + +@pytest.mark.parametrize( + ("backend", "expected"), + [ + ("auto", ("1", "1", "1")), + ("flash", ("1", "0", "0")), + ("fused", ("0", "1", "0")), + ("unfused", ("0", "0", "1")), + ("local", ("0", "0", "0")), + ], +) +def test_attention_backend_override_sets_expected_env(monkeypatch, backend, expected): + for name in ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN"): + monkeypatch.delenv(name, raising=False) + + _apply_attention_backend_env(backend, tag="unit") + + assert ( + os.environ["NVTE_FLASH_ATTN"], + os.environ["NVTE_FUSED_ATTN"], + os.environ["NVTE_UNFUSED_ATTN"], + ) == expected + + +def test_attention_backend_override_rejects_unknown_backend(): + with pytest.raises(ValueError, match="attention_backend_override"): + _apply_attention_backend_env("invalid", tag="unit") + + +class HookedOptimizer: + def __init__(self): + self.calls: list[str] = [] + + def offload_state_to_cpu(self): + self.calls.append("offload") + + def load_state_to_device(self): + self.calls.append("load") + + +def test_runtime_to_prefers_optimizer_specific_offload_hooks(): + optimizer = HookedOptimizer() + handle = ModelHandle(model=nn.Linear(2, 2), optimizer=optimizer, _extras={"model_chunks": []}) + runtime = MegatronLiteRuntime.__new__(MegatronLiteRuntime) + + runtime.to(handle, "cpu", model=False, optimizer=True, grad=False) + runtime.to(handle, "cuda", model=False, optimizer=True, grad=False) + + assert optimizer.calls == ["offload", "load"] + + +def test_model_handle_dp_defaults(): + handle = ModelHandle(model=MagicMock()) + + assert handle.dp_rank == 0 + assert handle.dp_size == 1 + assert handle.dp_group is None + + +def test_model_handle_dp_from_parallel_state(): + ps = MagicMock() + ps.dp_rank = 3 + ps.dp_size = 8 + ps.dp_group = "fake_group" + + handle = ModelHandle(model=MagicMock(), parallel_state=ps) + + assert handle.dp_rank == 3 + assert handle.dp_size == 8 + assert handle.dp_group == "fake_group" + + +def test_model_handle_cp_range_and_config_properties(): + cfg = {"tp": 8, "ep": 4} + default_handle = ModelHandle(model=MagicMock()) + configured_handle = ModelHandle(model=MagicMock(), config=cfg, _extras={"cp_range": (1, 8)}) + + assert default_handle.cp_range == (1, 1) + assert configured_handle.cp_range == (1, 8) + assert configured_handle.config is cfg + + +def test_runtime_dispatch_creates_mlite_backend(): + with patch("megatron.lite.runtime.backends.mlite.create") as mock_create: + backend = MagicMock() + mock_create.return_value = backend + + runtime = create_runtime( + RuntimeConfig( + backend="mlite", hf_path="/models/test", backend_cfg={"model_name": "qwen3"} + ) + ) + + assert runtime is backend + mock_create.assert_called_once_with("/models/test", {"model_name": "qwen3"}) + + +def test_runtime_dispatch_unknown_backend_raises(): + with pytest.raises(KeyError): + create_runtime(RuntimeConfig(backend="nonexistent")) + + +def _run_verl_sft_dry_run(script: Path, tmp_path: Path, **env_overrides: str) -> str: + env = { + **os.environ, + "MODEL_PATH": "/tmp/mlite-model", + "TRAIN_FILES": "/tmp/mlite-train.parquet", + "OUTPUT_ROOT": str(tmp_path), + "DRY_RUN": "1", + "NUM_GPUS": "1", + "NPROC_PER_NODE": "1", + "TP_SIZE": "1", + "PP_SIZE": "1", + "CP_SIZE": "1", + "EP_SIZE": "1", + "ETP_SIZE": "1", + **env_overrides, + } + completed = subprocess.run([str(script)], env=env, text=True, capture_output=True, check=True) + return completed.stdout + + +def test_verl_sft_script_maps_offload_env_to_backend_args(tmp_path): + script = ( + Path(__file__).resolve().parents[3] + / "examples" + / "verl" + / "scripts" + / "run_qwen3moe_sft.sh" + ) + + command = _run_verl_sft_dry_run( + script, + tmp_path, + PARAM_OFFLOAD="True", + OPTIMIZER_OFFLOAD="True", + OPTIMIZER_STATE_OFFLOAD_FRACTION="0.75", + ) + + assert "engine.param_offload=True" in command + assert "engine.optimizer_offload=True" in command + assert "+optim.override_optimizer_config.offload_fraction=0.75" in command + assert "+optim.override_optimizer_config.use_precision_aware_optimizer=True" in command + + +def test_verl_sft_script_does_not_emit_optimizer_state_offload_when_disabled(tmp_path): + script = ( + Path(__file__).resolve().parents[3] + / "examples" + / "verl" + / "scripts" + / "run_qwen3moe_sft.sh" + ) + + command = _run_verl_sft_dry_run( + script, tmp_path, PARAM_OFFLOAD="False", OPTIMIZER_OFFLOAD="False" + ) + + assert "engine.param_offload=False" in command + assert "engine.optimizer_offload=False" in command + assert "override_optimizer_config.offload_fraction" not in command diff --git a/experimental/lite/tests/unit/verl/test_mlite_engine_checkpoint.py b/experimental/lite/tests/unit/verl/test_mlite_engine_checkpoint.py new file mode 100644 index 00000000000..c68ff754011 --- /dev/null +++ b/experimental/lite/tests/unit/verl/test_mlite_engine_checkpoint.py @@ -0,0 +1,187 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from types import SimpleNamespace + +import pytest +import torch +from verl_mlite.engine.config import MegatronLiteEngineConfig +from verl_mlite.engine.mlite_engine import MegatronLiteEngine + + +class _Scheduler: + def __init__(self): + self.loaded_state = None + + def state_dict(self): + return {"step": 7, "lr": 0.25} + + def load_state_dict(self, state): + self.loaded_state = state + + +@pytest.fixture(autouse=True) +def _single_process_dist(monkeypatch): + monkeypatch.setattr("verl_mlite.engine.mlite_engine.dist.is_initialized", lambda: False) + + +def _optimizer_config() -> SimpleNamespace: + return SimpleNamespace( + optimizer="adam", + lr=1e-6, + min_lr=None, + min_lr_ratio=None, + clip_grad=1.0, + weight_decay=0.1, + lr_warmup_steps_ratio=0.0, + total_training_steps=10, + lr_warmup_steps=0, + override_optimizer_config={}, + ) + + +def _engine_config(**kwargs) -> MegatronLiteEngineConfig: + values = {"custom_backend_module": None, "impl_cfg": {"use_thd": True}} + values.update(kwargs) + return MegatronLiteEngineConfig(**values) + + +def _initialized_engine(*, checkpoint_config=None, param_offload=False): + engine = MegatronLiteEngine( + model_config=SimpleNamespace( + local_path="/tmp/qwen35", hf_config={"model_type": "qwen3_5_moe"}, mtp=None + ), + engine_config=_engine_config(param_offload=param_offload), + optimizer_config=_optimizer_config(), + checkpoint_config=checkpoint_config or {}, + ) + + def placement_fn(name): + return ["placement", name] + + def expert_classifier(name): + return name.endswith("expert") + + parallel = SimpleNamespace(tp=1, cp=1, pp=1) + parallel_state = SimpleNamespace(dp_rank=0) + module = torch.nn.Linear(2, 2) + optimizer = object() + scheduler = _Scheduler() + engine.module = module + engine.handle = SimpleNamespace( + _optimizer=optimizer, + _lr_scheduler=scheduler, + _config=SimpleNamespace(parallel=parallel), + _parallel_state=parallel_state, + _extras={ + "protocol": SimpleNamespace( + PLACEMENT_FN=placement_fn, EXPERT_CLASSIFIER=expert_classifier + ) + }, + ) + engine.runtime = object() + return ( + engine, + module, + optimizer, + scheduler, + parallel, + parallel_state, + placement_fn, + expert_classifier, + ) + + +def test_save_checkpoint_forwards_contents_scheduler_and_param_offload_reload( + tmp_path, monkeypatch +): + ( + engine, + module, + optimizer, + scheduler, + parallel, + parallel_state, + placement_fn, + expert_classifier, + ) = _initialized_engine(checkpoint_config={"save_contents": ["model"]}, param_offload=True) + to_calls = [] + save_calls = [] + sync_calls = [] + monkeypatch.setattr(engine, "to", lambda **kwargs: to_calls.append(kwargs)) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: sync_calls.append(True)) + monkeypatch.setattr( + "verl_mlite.engine.mlite_engine.save_training_checkpoint", + lambda *args, **kwargs: save_calls.append((args, kwargs)), + ) + + engine.save_checkpoint(str(tmp_path), global_step=13) + + assert to_calls == [ + {"device": "cuda", "model": True, "optimizer": False, "grad": False}, + {"device": "cpu", "model": True, "optimizer": False, "grad": False}, + ] + assert sync_calls == [True] + assert len(save_calls) == 1 + save_args, save_kwargs = save_calls[0] + assert save_args == (module, optimizer, 13, str(tmp_path), parallel, parallel_state) + assert save_kwargs["get_placements"] is placement_fn + assert save_kwargs["is_expert"] is expert_classifier + assert save_kwargs["save_model"] is True + assert save_kwargs["save_optimizer"] is False + assert ( + torch.load(tmp_path / "lr_scheduler.pt", map_location="cpu", weights_only=False) + == scheduler.state_dict() + ) + + +def test_save_checkpoint_skips_when_contents_exclude_model_and_optimizer(tmp_path, monkeypatch): + engine, *_ = _initialized_engine(checkpoint_config={"save_contents": ["extra"]}) + checkpoint_path = tmp_path / "ckpt" + save_calls = [] + monkeypatch.setattr( + "verl_mlite.engine.mlite_engine.save_training_checkpoint", + lambda *args, **kwargs: save_calls.append((args, kwargs)), + ) + + engine.save_checkpoint(str(checkpoint_path), global_step=13) + + assert save_calls == [] + assert not checkpoint_path.exists() + + +def test_load_checkpoint_restores_scheduler_and_param_offload_reload(tmp_path, monkeypatch): + ( + engine, + module, + optimizer, + scheduler, + parallel, + parallel_state, + placement_fn, + expert_classifier, + ) = _initialized_engine(param_offload=True) + torch.save({"step": 23, "lr": 0.125}, tmp_path / "lr_scheduler.pt") + to_calls = [] + load_calls = [] + sync_calls = [] + monkeypatch.setattr(engine, "to", lambda **kwargs: to_calls.append(kwargs)) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: sync_calls.append(True)) + monkeypatch.setattr( + "verl_mlite.engine.mlite_engine.load_training_checkpoint", + lambda *args, **kwargs: load_calls.append((args, kwargs)), + ) + + engine.load_checkpoint(str(tmp_path)) + + assert to_calls == [ + {"device": "cuda", "model": True, "optimizer": False, "grad": False}, + {"device": "cpu", "model": True, "optimizer": False, "grad": False}, + ] + assert sync_calls == [True] + assert scheduler.loaded_state == {"step": 23, "lr": 0.125} + assert len(load_calls) == 1 + load_args, load_kwargs = load_calls[0] + assert load_args == (module, optimizer, str(tmp_path), parallel, parallel_state) + assert load_kwargs["get_placements"] is placement_fn + assert load_kwargs["is_expert"] is expert_classifier + assert load_kwargs["load_model"] is True + assert load_kwargs["load_optimizer"] is True diff --git a/experimental/lite/tests/unit/verl/test_mlite_engine_config.py b/experimental/lite/tests/unit/verl/test_mlite_engine_config.py new file mode 100644 index 00000000000..fa6a9f71442 --- /dev/null +++ b/experimental/lite/tests/unit/verl/test_mlite_engine_config.py @@ -0,0 +1,113 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from types import SimpleNamespace + +from verl_mlite.engine.config import MegatronLiteEngineConfig +from verl_mlite.engine.mlite_engine import MegatronLiteEngine + + +def _optimizer_config(**override_optimizer_config) -> SimpleNamespace: + return SimpleNamespace( + optimizer="adam", + lr=1e-6, + min_lr=None, + min_lr_ratio=None, + clip_grad=1.0, + weight_decay=0.1, + lr_warmup_steps_ratio=0.0, + total_training_steps=10, + lr_warmup_steps=0, + lr_warmup_init=0.0, + lr_decay_steps=None, + lr_decay_style="constant", + weight_decay_incr_style="constant", + lr_wsd_decay_style="exponential", + lr_wsd_decay_steps=None, + use_checkpoint_opt_param_scheduler=False, + betas=(0.9, 0.95), + override_optimizer_config=override_optimizer_config, + ) + + +def _engine( + *, engine_config: MegatronLiteEngineConfig, optimizer_config: SimpleNamespace | None = None +) -> MegatronLiteEngine: + return MegatronLiteEngine( + model_config=SimpleNamespace( + local_path="/tmp/qwen35", hf_config={"model_type": "qwen3_5_moe"}, mtp=None + ), + engine_config=engine_config, + optimizer_config=optimizer_config or _optimizer_config(), + checkpoint_config={}, + ) + + +def _engine_config(**kwargs) -> MegatronLiteEngineConfig: + values = {"custom_backend_module": None, "impl_cfg": {"use_thd": True}} + values.update(kwargs) + return MegatronLiteEngineConfig(**values) + + +def test_optimizer_offload_enables_full_optimizer_state_offload_by_default() -> None: + engine = _engine( + engine_config=_engine_config(optimizer_offload=True), + optimizer_config=_optimizer_config( + use_precision_aware_optimizer=True, decoupled_weight_decay=True + ), + ) + + optimizer = engine._build_mlite_optimizer_config() + + assert optimizer.offload_fraction == 1.0 + assert optimizer.use_precision_aware_optimizer is True + assert optimizer.decoupled_weight_decay is True + assert optimizer.adam_beta1 == 0.9 + assert optimizer.adam_beta2 == 0.95 + + +def test_explicit_optimizer_offload_fraction_overrides_engine_default() -> None: + engine = _engine( + engine_config=_engine_config(optimizer_offload=True), + optimizer_config=_optimizer_config(offload_fraction=0.25), + ) + + optimizer = engine._build_mlite_optimizer_config() + + assert optimizer.offload_fraction == 0.25 + + +def test_optimizer_cpu_offload_alias_maps_to_full_offload_fraction() -> None: + engine = _engine( + engine_config=_engine_config(optimizer_offload=False), + optimizer_config=_optimizer_config(optimizer_cpu_offload=True), + ) + + optimizer = engine._build_mlite_optimizer_config() + + assert optimizer.offload_fraction == 1.0 + + +def test_mlite_config_threads_rl_parallel_and_impl_settings() -> None: + engine = _engine( + engine_config=_engine_config( + tp=2, + ep=8, + etp=1, + pp=1, + cp=1, + optimizer_offload=True, + attention_backend_override="flash", + impl_cfg={"use_thd": True, "deterministic": False}, + ) + ) + + config = engine._build_mlite_config() + + assert config.model_name == "qwen3_5" + assert config.impl == "lite" + assert config.parallel.tp == 2 + assert config.parallel.ep == 8 + assert config.parallel.etp == 1 + assert config.optimizer.offload_fraction == 1.0 + assert config.attention_backend_override == "flash" + assert config.impl_cfg["use_thd"] is True + assert config.impl_cfg["deterministic"] is False