[dev] Add experimental Megatron Lite as agentic exploration

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"]

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.

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.

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.

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.

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)`.

experimental/lite/examples/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""Megatron Lite examples package."""

experimental/lite/examples/bench/.gitignore

@@ -0,0 +1 @@
+outputs/