Repository navigation
Conversation
S1ro1
force-pushed
the
feat/nvfp4-main-refresh
branch
from
September 9, 2026 15:03
f6fcd3e to
c2ac26b
Compare
S1ro1
force-pushed
the
feat/nvfp4-main-refresh
branch
from
September 9, 2026 19:45
5129093 to
35683f3
Compare
S1ro1
force-pushed
the
feat/nvfp4-main-refresh
branch
from
September 9, 2026 21:02
35683f3 to
b8ff1b0
Compare
S1ro1
force-pushed
the
feat/nvfp4-main-refresh
branch
from
September 9, 2026 21:34
b8ff1b0 to
3c9ebb0
Compare
S1ro1
added a commit
that referenced
this pull request
Sep 9, 2026
Select the model layers that use the configured MoE compute backend with `model.moe.compute.apply_to`: - `"all"` (default): all routed-expert groups. - `"85%"`: the first 85% of model layers, rounded down; 48 layers selects indices 0–39. - `[0, 1, 2, 3]`: explicit zero-based model layer indices. `[]` selects none. Non-MoE blocks in hybrid models count toward indices and percentages. Other expert groups use BF16 compute and transport while retaining the configured dispatch backend and expert parallelism. Out-of-range indices and invalid percentages fail early. Dense linear quantization is unchanged. ```toml [trainer.model.moe.compute] type = "mxfp8" apply_to = "85%" ``` BF16, DeepGEMM FP8 and MXFP8 configs inherit the common selector. NVFP4 inherits it in #3511. Inference precision remains configured separately. First PR in the stack: main → #3518 → #3515 (filter flag) → #3513 (vLLM upgrade) → #3511 (NVFP4). Validation: 62 focused MoE/config tests passed on this main-based branch, including all three forms through CLI/config round trips, percentage rounding, hybrid-layer indexing, and invalid indices. Ruff lint/format and whitespace checks passed. GB200 mixed DeepGEMM-FP8/BF16 forward/backward passed with an exact BF16 fallback output and unchanged parameters/checkpoint keys. EP2 with layer-index selection verified MXFP8 transport for the selected group, BF16 transport for the excluded group, both forwards, and BF16 fallback backward. MXFP8 backward has a separately reproduced CuTe TMA-alignment compiler failure in the existing kernel; that kernel is unchanged here. <sub>Stack created with <a href="https://github.com/github/gh-stack">GitHub Stacks CLI</a> • <a href="https://gh.io/stacks-feedback">Give Feedback 💬</a></sub> <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Medium Risk** > Per-layer MoE precision and transport wiring affects forward/backward numerics and expert-parallel paths; misaligned trainer vs inference precision could skew RL logprobs if rollout config is not matched manually. > > **Overview** > Adds **`model.moe.compute.apply_to`** so BF16, DeepGEMM FP8, and MXFP8 routed-expert backends can target **all layers**, a **leading percentage** (e.g. `"85%"`, rounded down), or **explicit zero-based layer indices** (including `[]` for none). Shared **`MoEComputeConfigBase`** validates inputs and **`resolve_layers`** maps selections to layer sets; out-of-range indices fail at resolve time. > > **`configure_moe_runtime`** maps selection onto language-model **`layers`** (including VLM via **`get_language_model`**): selected MoE modules get the configured grouped GEMM; others **fall back to BF16** compute and matching token dispatch alignment. **MXFP8 torch transport** is used only when that layer’s grouped GEMM is actually **`MXFP8GroupedGemm`**, so excluded layers no longer take MXFP8 all-to-all when compute is partial. Startup logging reports selected vs total MoE layers. > > Docs describe hybrid-layer counting, dense-quant independence, and RL rollout/inference precision alignment. Unit tests cover config/CLI round-trips and BF16 fallback without loading MXFP8 backends when nothing is selected. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit f367d48. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY --> --------- Co-authored-by: matejsirovatka <matejsirovatka@users.noreply.github.com>
S1ro1
force-pushed
the
feat/nvfp4-main-refresh
branch
from
September 9, 2026 22:11
3c9ebb0 to
92bbac9
Compare
S1ro1
changed the title
S1ro1
force-pushed
the
feat/nvfp4-main-refresh
branch
from
September 17, 2026 14:24
ada2172 to
52cb529
Compare
S1ro1
force-pushed
the
feat/nvfp4-main-refresh
branch
from
September 21, 2026 14:14
52cb529 to
19f78e2
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add routed-expert NVFP4 forward compute through
model.moe.compute.type = "nvfp4", using the prime-kernels grouped GEMM with BF16 output. The defaultdequant_bf16backward reconstructs the packed forward operands;bf16backward retains the original operands. Reuse existing MoE dispatch, expert parallelism, activation checkpointing, andapply_tolayer selection.four_over_sixdefaults tofalse. Enabling it applies adaptive 4/6 quantization to weights and activation inputs using FlashInfer's default recipe. Enable the matching sampler withFLASHINFER_NVFP4_4OVER6=1. Partial-graph compilation is supported; full-graph compilation and LoRA are unsupported.Include SM103 kernel build/runtime support and FA4 auto-selection on SM103 for B300. Pin prime-kernels
2802a7aand its published x86_64 wheel. ARM64 retains its existing NVFP4 wheels. The x86_64 HiSparse/NVFP4 vLLM wheel, including the reload-scale fix, is inherited from #3513. Transformers stays at 5.6.2.Stack: main → #3513 → #3511, rebased on
a1822f7a0. The debug algorithm and FP8 indexer constexpr fix are inherited from main. The zero-advantage filtering PR #3515 and all its changes are removed from the stack. Cluster-specific configs, diagnostic scripts, and the FP32 expert-output experiment are excluded.Validation on the complete rebased stack:
uv sync --all-extras --all-packagesanduv sync --frozen --all-extras --all-packagessucceeded with the published wheel and matching hash.uv run --no-sync pytest tests/unit/test_configs.py tests/unit/test_parsers.py tests/unit/train/models/test_moe.py tests/unit/inference/test_serving_tokens.py -q: 381 passed.uv run --no-sync pre-commit run --from-ref origin/main --to-ref HEAD,uv lock --check, and whitespace checks passed.The retained kernels and vLLM fixes were exercised in the earlier B300 experiments. No fresh full RL or GPU run was launched after this rebase.