Repository navigation
Conversation
S1ro1
marked this pull request as ready for review
September 9, 2026 19:47
This was referenced Sep 9, 2026
S1ro1
marked this pull request as draft
September 9, 2026 19:47
S1ro1
force-pushed
the
fix/restore-zero-advantage-filter
branch
from
September 9, 2026 21:01
1885c2a to
dbcd9e4
Compare
S1ro1
force-pushed
the
fix/restore-zero-advantage-filter
branch
from
September 9, 2026 21:33
dbcd9e4 to
01fecc4
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
fix/restore-zero-advantage-filter
branch
from
September 9, 2026 22:11
01fecc4 to
0f52a95
Compare
S1ro1
force-pushed
the
fix/restore-zero-advantage-filter
branch
from
September 17, 2026 14:24
0f52a95 to
7cf4dbb
Compare
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.
Removed from the NVFP4 stack as requested. The remaining stack is main → #3513 (vLLM) → #3511 (NVFP4 training), rebased on main
a1822f7a0.The proposed
orchestrator.train.filter_zero_advantagesflag and this PR's associated code/documentation changes are absent from both remaining PRs. Main's existing filtering behavior is retained.