Merge branch 'yuya/perf/cuda-graph-memory' into yuya/perf/moe-paged-stash

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

# Conflicts:
#	tests/unit_tests/moe/test_backend_config.py

Author: yaoyu-33

.agents/contributor-skills/build-and-dependency/SKILL.md

@@ -84,21 +84,27 @@ The media extras (`vlm-media`, `diffusion-media`, `media`) bundle FFmpeg and are
 deliberately **excluded from `all`** and from the container image — add them
 explicitly for video/image decode.
 
-### Option 3: pip
+### Option 3: uv pip
 
 Full install (matches `uv sync --extra all`):
 
 ```bash
-pip install -e ".[all]"
+uv venv
+source .venv/bin/activate
+uv pip install -e ".[all]"
 ```
 
-Login-node / submitter-only install — lightweight package for SLURM, k8s, or
-NeMo-Run job submission without local CUDA deps:
+To add NeMo Run submission support to the base package:
 
 ```bash
-pip install nemo-automodel[cli]
+uv venv
+source .venv/bin/activate
+uv pip install "nemo-automodel[cli]"
 ```
 
+The `cli` extra is additive: it adds `nemo-run` but does not remove the base
+package's core training dependencies, including PyTorch.
+
 ## Package Management
 
 Always use `uv`. Do not introduce `pip install` commands in scripts or docs.
@@ -120,31 +126,20 @@ export HF_HOME="/path/to/hf_cache" # Hugging Face cache directory
 
 ## CLI Usage
 
-The entry point is `automodel` (defined at `nemo_automodel._cli.app:main`).
+The entry point is `automodel` (defined at `nemo_automodel.cli.app:main`).
 
-Pattern: `automodel <command> <domain> -c <config.yaml>`
+Pattern: `uv run automodel <config.yaml> [--nproc-per-node N] [--key.subkey value ...]`
 
 ```bash
-# LLM
-automodel finetune llm -c examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml
-automodel pretrain llm -c config.yaml
-automodel kd llm -c config.yaml
-automodel benchmark llm -c config.yaml
-
-# VLM
-automodel finetune vlm -c config.yaml
-
-# Diffusion
-automodel finetune diffusion -c config.yaml
-
-# Retrieval
-automodel finetune retrieval -c config.yaml
+# The YAML's recipe field selects LLM, VLM, diffusion, or retrieval behavior.
+uv run automodel examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml --nproc-per-node 8
 ```
 
 Override any config value from the CLI:
 
 ```bash
-automodel finetune llm -c config.yaml --model.name_or_path meta-llama/Llama-3.2-1B
+uv run automodel examples/llm_finetune/llama3_2/llama3_2_1b_squad.yaml \
+    --model.pretrained_model_name_or_path meta-llama/Llama-3.2-1B
 ```
 
 ## Common Pitfalls

.github/workflows/cicd-main.yml

@@ -687,6 +687,8 @@ jobs:
           verbose: true
           flags: ${{ matrix.flag }}
           base_sha: ${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').base.sha }}
+          override_pr: ${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').number }}
+          override_commit: ${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').head.sha }}
 
       - name: Upload artifacts
         uses: actions/upload-artifact@v6

.github/workflows/claude-review.yml

@@ -12,7 +12,7 @@ permissions:
 
 jobs:
   claude-review:
-    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_claude_review.yml@v1.7.0
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_claude_review.yml@209ac7913b0419a5ccbac47b02d00fbea4939243 #v1.8.4
     with:
       model: ${{ vars.CLAUDE_MODEL }}
       prompt: |
@@ -107,6 +107,16 @@ jobs:
           - ambiguous public signatures: boolean-trap parameters, multiple
             same-typed positional runtime arguments that should be keyword-only,
             or long anonymous tuple returns that need a named typed result.
+        - Model-specific ownership is mandatory. Model-specific logic may live
+          only under `nemo_automodel/components/models/<model>/`. Treat logic as
+          model-specific when it names or branches on a model or family, hard-codes
+          its module/parameter paths or weight layout, or implements its
+          architecture, initialization, or parallelization policy. If changed code
+          adds or expands such logic anywhere else, including another directory
+          under `nemo_automodel/components/` or a recipe, report a critical finding
+          and require relocation to the owning model package. Shared components may
+          expose generic typed contracts and invoke model-owned hooks, but must not
+          encode model identities or model-specific policy.
         - Tensor contract docstrings. For every new function or method, or any
           function/method whose signature or tensor-handling body is materially
           changed, that accepts a tensor input, require its docstring to document
@@ -118,12 +128,20 @@ jobs:
           `nn.Parameter`, optional tensors, variadic tensors, and tensors nested in
           tuples/lists/mappings/dataclasses as tensor inputs. Require the Google-style
           `Args` section and, when values are returned, `Returns` section to document:
-          - every tensor input and output's semantic shape and axis order, with
-            symbolic dimensions defined at first use (for example, `B` = batch, `S`
-            = sequence, `H` = hidden). Flag vague descriptions such as "embedding
-            tensor" or unexplained `[B, S, H]` notation. If arbitrary ranks or
-            leading dimensions are accepted, state that explicitly and identify the
-            constrained axes (for example, `[..., H]` with `H` = hidden size);
+          - every tensor input and output's semantic shape and axis order. Follow
+            the common format `hidden_states: Tensor of shape [batch, sequence,
+            hidden].` Use clear, conventional dimension names such as `batch`,
+            `sequence`, `tokens`, `hidden`, `heads`, `channels`, and `vocab` without
+            explaining them. Define non-obvious or model-specific dimension names at
+            first use. Flag vague descriptions such as "embedding tensor" or
+            layouts that do not follow this format. If arbitrary ranks or leading
+            dimensions are accepted, state that explicitly and identify the
+            constrained axes (for example, `hidden_states: Tensor of shape [...,
+            hidden], with arbitrary leading dimensions.`). When reporting a layout
+            documentation finding, include a concrete compliant rewrite in the
+            review; for example, `fused_qkv: Tensor of shape [batch, sequence, 3,
+            heads, head_dim], where axis 2 stores query, key, and value in that
+            order.`;
           - nonstandard logical layout such as packed, padded, ragged, flattened,
             interleaved, fused-QKV/GateUp, THD, or channels-first/channels-last, and
             the before/after layout when the API transforms it;
@@ -163,6 +181,29 @@ jobs:
           that class's nested `ModelCapabilities` dataclass (static pattern) or
           inside the `get_capabilities(cls, config)` classmethod (dynamic
           dispatch pattern).
+        - Checkpoint-free model initialization. When a PR adds or materially
+          changes model parameters, persistent buffers, initialization methods,
+          routing state, or a random-init recipe path, verify that construction
+          followed by the production initialization path leaves every tensor that
+          can affect the first forward/backward in a valid state without loading a
+          checkpoint:
+          - flag parameters or persistent buffers allocated with `torch.empty`, or
+            initialized to sentinel/placeholder values, unless every supported
+            model variant and stage (including auxiliary heads and MTP blocks)
+            overwrites them before use. A later checkpoint load is not a substitute
+            when checkpoint-free or random-init training is supported;
+          - require routing/index tables to satisfy the downstream dispatcher's
+            structural invariants before first use: indices must be in range and,
+            when the dispatcher requires it (for example DeepEP top-k routing),
+            expert IDs for each token must be distinct. An all-zero placeholder is
+            valid only when it satisfies the actual routing contract;
+          - require a focused test that disables checkpoint/base-model loading,
+            invokes the production initializer, and asserts finite initialized
+            parameters plus the relevant buffer/routing invariants. When this state
+            affects execution or gradients, also require a finite first
+            forward/backward on the smallest representative backend/topology; apply
+            the distributed test matrix below when a single-process backend cannot
+            exercise the contract.
         - Low-precision dtype hazards in RoPE / precision-sensitive buffers. Flag
           when a module registers a floating-point buffer used to build rotary
           tables — `inv_freq`, `freqs_cis`, or precomputed `cos`/`sin` — and the

README.md

@@ -198,7 +198,8 @@ uv venv
 # Choose ONE:
 uv sync --frozen  # LLM recipes (default)
 # uv sync --frozen --extra vlm --extra vlm-media  # VLM recipes (Qwen/Mistral/Omni need vlm-media for video/vision; fixes: ImportError: qwen_vl_utils is not installed)
-# uv sync --frozen --extra cuda  # Optional CUDA deps (e.g., Transformer Engine, bitsandbytes)
+# uv sync --frozen --extra cuda  # Optional CUDA deps (e.g., Transformer Engine, Mamba SSM)
+# uv sync --frozen --extra cuda_source  # Optional bitsandbytes dependency
 # uv sync --frozen --extra all  # Most optional deps (includes `vlm` and `cuda`; NOTE: excludes media — add --extra media for video/image decode)
 # uv sync --frozen --all-extras  # Everything (includes `fa`, `moe`, `media`, etc.)
 
@@ -224,7 +225,7 @@ uv run automodel examples/llm_finetune/llama3_2/llama3_2_1b_hellaswag.yaml --npr
 ```
 
 > [!TIP]
-> **Login-node / CI installs:** If you only need to submit jobs (SLURM, k8s, NeMo-Run) and don't need to train locally, install the lightweight CLI package: `pip install nemo-automodel[cli]`
+> **NeMo-Run submission:** The `cli` extra adds NeMo Run to the base package: `uv pip install "nemo-automodel[cli]"`. It is additive; the base package still installs its core training dependencies, including PyTorch.
 
 
 ## LLM Pre-training