feat: vendor MXFP8 MoE runtime - #1
Draft
S1ro1 wants to merge 1 commit into
Draft
Conversation
Member
|
Can the backward of the mxfp8 kernel be used as backward for our fused-mxf8-moe kernels? |
S1ro1
added a commit
to PrimeIntellect-ai/prime-rl
that referenced
this pull request
Aug 28, 2026
## Summary ### Model and checkpoint contract - Replace the separate standard, GPT-OSS, and Nemotron expert/router implementations with one canonical MoE tree. - Store routed expert weights as separate `gate_proj`, `up_proj`, and `down_proj` parameters. Non-gated experts omit `gate_proj`. - Nest every custom router at `mlp.router.gate`, routed experts at `mlp.experts`, and optional shared experts inside `mlp`. - Store the single selection-only router bias at `mlp.router.selection_bias`; model-specific checkpoint names are converted only at the HF boundary. - Keep shared expert projections canonical. Qwen3.5 alone nests its scalar sigmoid gate at `mlp.shared_expert.output_gate`, with conversion mapping the standalone Hugging Face key at the boundary. - Use the same activation protocol and `FeedForward` implementation for dense MLPs, shared experts, and grouped experts. - Make shared-expert construction exclusively parent-owned: model layers pass an explicit `FeedForward`/subclass or `None`, while `MoEArgs` and `MoE` never infer or construct one. - Keep existing HF source layouts at the conversion boundary. GPT-OSS alone adds the structural split/join required for its interleaved source tensor. - Require `.prime-v1` in reusable `<conversion_dir>/prime` caches so legacy Prime layouts fail before loading. ### MoE runtime contract - Use one routed execution sequence for every backend: dispatch tokens, run canonical grouped experts, combine tokens. - Select grouped expert compute through explicit BF16, DeepGEMM FP8, or MXFP8 implementations. `MoEArgs` remains architecture-only. - Select local/Torch or DeepEP token dispatch once during model setup. Dispatchers own routing metadata, communication, padding, chunking, and synchronization. - Keep `TokenDispatcher` as the structural interface while a shared base owns the canonical run/synchronize sequence; local, Torch, and DeepEP implementations own only their dispatch/combine details. - Keep DeepEP combine events, concatenate streams, and flattened dispatch state owned by each dispatcher; chunking uses the equal input shape shared by EP ranks and does not synthesize empty chunks. - Ignore dispatch configuration when expert parallelism resolves to one rank; local dispatch remains active and DeepEP has no optimizer side effects. - Keep grouped expert execution directly in `GroupedExperts.forward` and router confidence accounting directly in `TokenChoiceTopKRouter.forward`. - Shard every canonical expert parameter with one backend-independent `Shard(0)` placement. - Keep dense linear quantization, routed-expert compute, and EP transport independent. - Remove the TorchTitan module hooks, expert decorator, dynamically attached process group, backend attributes, tensor-subclass interception, and process-global MXFP8 patches from the reachable path. - Vendor the narrow MXFP8 grouped-GEMM and Torch EP transport implementation in [prime-kernels PR #1](PrimeIntellect-ai/prime-kernels#1), preserving the pinned torchao revision and BSD license. ## Public contracts Checkpoint tree: - `mlp.router.gate`: `nn.Linear` - `mlp.router.selection_bias`: optional persistent selection-only buffer - `mlp.experts.gate_proj`: stacked `nn.Parameter` for gated experts, absent for non-gated experts - `mlp.experts.up_proj`: stacked `nn.Parameter` - `mlp.experts.down_proj`: stacked `nn.Parameter` - `mlp.shared_expert.{gate_proj,up_proj,down_proj}`: `FeedForward` projections when present - `mlp.shared_expert.output_gate`: Qwen3.5-only scalar `nn.Linear` - expert matrices: `[num_experts, output_features, input_features]` Runtime configuration: ```toml [trainer.model.moe.compute] type = "bf16" # bf16 | deepgemm_fp8 | mxfp8 [trainer.model.moe.dispatch] type = "torch" # torch | deepep transport = "bf16" # bf16 | mxfp8 for torch dispatch ``` DeepEP owns its tuning fields: ```toml [trainer.model.moe.dispatch] type = "deepep" num_sms = 20 token_chunk_size = 4096 ``` This is intentionally breaking. `enable_grouped_gemm`, `enable_a2a`, `ep_comm_backend`, `deepep_num_sms`, and `deepep_token_chunk_size` are removed without aliases. ## Validation completed - Foundation conversion, detector, marker, state-loading, checkpoint, NIXL graph, and wire-format tests passed. - Focused CPU configuration and MoE contracts: 143 passed. - Focused MoE runtime configuration parametrizations after the EP=1 no-op change: 6 passed. - Affected custom-model CUDA suite: 24 passed, 3 skipped across AFMoE, GLM, GPT-OSS, Nemotron, Qwen3/Qwen3.5, and state loading. - Final H200 model suite after execution-path cleanup: 81 passed, 3 skipped (Slurm 2597). - Torch EP on two H200 GPUs passed with BF16 and DeepGEMM FP8, top-k 2, both score locations, shared experts, gated/non-gated/bias variants, uneven routing, a rank receiving zero routed tokens, and input/router/expert gradients. - Final Torch EP BF16 and DeepGEMM FP8 matrix passed after the cleanup (Slurm 2596). - DeepEP on two H200 GPUs passed with BF16 and DeepGEMM FP8, chunking on/off, both score locations, shared experts, a rank receiving zero routed tokens, and gradients. - Final DeepEP BF16 and DeepGEMM FP8 matrix passed after the cleanup (Slurm 2595). - Parent-owned shared-expert construction: 16 focused CPU tests and 64 H200 model tests passed, with 3 expected skips (Slurm 2600). - Full-graph MoE and shared `FeedForward` compile validation passed after the ownership cleanup (Slurm 2601). - Local DeepGEMM FP8 passed across all supported topology, activation, and bias combinations. - `torch.compile(fullgraph=True)` forward/backward passed for the full MoE and shared `FeedForward` across all supported topology/activation combinations. - Single-rank DCP and Torch EP + FSDP DCP model/optimizer resume passed. - The `mxfp8_moe` wheel build includes its public module, manifest, and license; the module imports against the pinned torchao build. - Unsupported H200 MXFP8 selection fails before model mutation. - Review audit: every queued DeepEP combine event is now synchronized, including chunked execution (two-H200 Slurm 2604). - Dispatcher review cleanup passed on two H200s: DeepEP BF16/DeepGEMM with chunking and a zero-receive rank (Slurm 2847), plus Torch EP topology/bias/zero-receive coverage (Slurm 2848). - The exact `reverse_text_moe` integration passed end-to-end on two H200s with an isolated conversion cache (Slurm 2607). - Conversion ranks now make the create decision on master and validate `.prime-v1` only after the final all-rank barrier. - Canonical router selection bias: 12 CPU MoE/conversion tests passed; CUDA checkpoint loading, Nemotron, and GLM model tests passed, with the existing GLM BF16 gradient comparison passing on isolated rerun after one threshold-edge miss (Slurm 2851/2852). - Ruff, Ruff formatting, compileall, and `git diff --check` passed. ## Deliberate exclusions - MultiLoRA is not adapted to the new runtime and may fail. - Fused-MoE trainer integration is removed and will be reconsidered only after the canonical runtime is finalized. - Routed-expert selective activation checkpointing is removed; the remaining selective targets are unchanged. ## Remaining before marking ready - Run numerical MXFP8 grouped-GEMM and MXFP8 Torch EP transport forward/backward on SM100 hardware. The available Slurm fleet currently has only H200 nodes. - Merge prime-kernels PR #1, publish wheels containing `mxfp8_moe`, and update the `pyproject.toml`/lock wheel pins. The currently pinned v0.8.0 wheel does not contain this module. <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **High Risk** > This changes the MoE forward/backward path, expert-parallel dispatch, checkpoint key layout, and trainer TOML in breaking ways across all custom MoE models. > > **Overview** > **Breaking trainer config:** MoE is configured via `[trainer.model.moe.compute]` (`bf16`, `deepgemm_fp8`, `mxfp8`) and `[trainer.model.moe.dispatch]` (`torch` with `transport`, or `deepep` with `num_sms` / `token_chunk_size`). Dense `[trainer.model.quantization]` no longer drives expert GEMMs or EP transport. Removed flags include `ep_comm_backend`, `deepep_*`, `moe_use_grouped_mm`, `moe_fused_kernel`, and quantization `enable_grouped_gemm` / `enable_a2a`. > > **Unified MoE execution:** All custom models share one `MoE` path—`TokenChoiceTopKRouter`, stacked `GroupedExperts` (`gate_proj` / `up_proj` / `down_proj`), and a pluggable `TokenDispatcher` (`LocalTokenDispatcher`, `TorchTokenDispatcher`, `MXFP8TorchTokenDispatcher`, `DeepEPTokenDispatcher`). `configure_moe_runtime` wires grouped GEMM backends and dispatch at setup; EP weight sharding is only `ExpertWeightParallel`. DeepEP dispatch/combine, grouped-GEMM permutation, and chunk pipelining live in the new dispatcher layer instead of ad hoc `MoE` hooks. > > **Checkpoint and model surface:** PrimeRL checkpoints rename routed weights from `w1/w2/w3` to `gate/down/up_proj`, move selection bias to `mlp.router.selection_bias`, and add GPT-OSS HF↔prime conversion. `LatentMoE` / Nemotron-specific expert classes are folded into the canonical tree; dense layers use `FeedForward`. Reusable `prime` conversion caches must include a `.prime-v1` marker. > > **Removed / docs:** Fused `flash_moe` trainer integration and `bench_fused_moe.py` are dropped; `routed_experts` selective activation checkpointing is removed. Docs and examples (e.g. GLM-5) are updated for the new tables. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit ffd5be7. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY -->
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
2 participants
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.
Summary
Validation