Repository navigation
Conversation
yh0903
marked this pull request as ready for review
September 25, 2026 18:49
yh0903
marked this pull request as draft
September 25, 2026 19:35
meefs
pushed a commit
to meefs/DeepSpeed
that referenced
this pull request
Oct 2, 2026
…speedai#8671) ## Summary AutoEP's native collective path (`comm_backend="comm"`, NCCL all-to-all) puts each rank's received rows into expert-major order before the grouped GEMM, and back afterwards, with `permute_by_local_expert` and `unpermute_by_local_expert`. The `autoep_size=1` path uses the same two functions. Both moved every routed row more often than needed: - **`permute_by_local_expert`:** - appended a zero row to the received rows with `torch.vstack`, a full copy, before indexing the padded expert-major buffer; - its backward was an accumulating `index_put`, i.e. a sort plus atomic adds. - **`unpermute_by_local_expert`:** zero-filled an `[N + 1, H]` buffer, then scattered into it with an elementwise `index_put`. On CUDA with Triton, both directions are now gathers: | Direction | Forward | Backward | |---|---|---| | permute | a Triton row gather that writes zero rows where `generate_permute_indices` marked padding (`-1`) | `index_select` through the inverse permutation | | unpermute | `index_select` through the inverse permutation | the same padded row gather | **Scope.** - This is a local, on-rank change to the expert reorder. - The all-to-alls send and receive the same rows in the same order, so communication volume is unchanged, and no communication/computation overlap is added. - `generate_permute_indices` is unchanged, and the DeepEP backend does not call these functions. - There is no switch: the gathers replace the previous code for 2-D CUDA tensors when Triton is available, because their results are equal to it (see Numerics). Other tensors keep the previous code. ## Permutation contract `generate_permute_indices` sends each received row to exactly one slot of the padded expert-major buffer and marks every other slot `-1`. The gathers rely on that one-to-one mapping: - **Permute forward:** writes each output row once: the row mapped to its slot, or zeros at a padding slot. - **Inverse permutation:** a scatter of slot positions into scratch of `N + N_padded` entries. Each padding slot goes to its own entry past `N`, so no entry is written twice, and every row's entry is written. - **Unpermute forward:** reads each row's slot through the inverse; padding rows are never read. - **Backward:** each direction's backward is the other gather. A row's gradient is that of the one slot it occupied, not a sum, and padding slots get zero gradient. Both backwards are autograd functions themselves, so double backward goes through the same pair. - **Top-k:** each received row is one (token, expert) assignment. Summing a token's top-k expert outputs, in the combine, and accumulating its top-k gradients, in the backward of the routed gather `x[token_indices // top_k]`, stay in the surrounding code, unchanged. These helpers are therefore not a general gather/scatter. `permute_rows`, `unpermute_rows` and `inverse_permutation` require the valid entries to be a complete one-to-one permutation of the rows, and do not check it; the docstrings say so. Both callers meet the requirement, because the counts sum to the number of rows: - with expert parallelism, each all-to-all receives exactly the row sums of the exchanged counts; - with `autoep_size=1`, the `T × top_k` routed rows are sorted by expert, and the counts are per expert. The previous code tolerated counts that sum to fewer rows (the rows left out came back as zeros); the gathers do not. The row gather reads its index by linear offset, so it now makes the index contiguous first, as it already did the source rows. ## Numerics The gathers only move data. Outputs, gradients and double-backward gradients equal the previous implementation's under `torch.equal`, which does not distinguish `+0.0` from `-0.0`. We also compared bit patterns on H100, in bf16, fp16 and fp32, with inputs and upstream gradients of which a quarter were `-0.0`: - Both outputs and the unpermute gradient are bitwise identical to the previous implementation's. - The permute gradient differs only where the upstream gradient is `-0.0`. The previous backward accumulated it into a zero-filled buffer, giving `+0.0`, whereas the gather passes `-0.0` through. ## Validation **Unit tests** (below): the earlier H100 validation passed these 27 tests and the existing `tests/unit/v1/moe/test_autoep_unit.py` (217 tests). The review follow-up additionally validates all 27 reorder cases with plain, forked and xdist-forked execution. **Local token movement**, one rank's work for one MoE layer, on one H100: - Shape: T tokens, top-8 of 128 experts, EP16 (8 local experts), hidden 2048. - The all-to-alls are replaced by the identity and the experts by a scalar multiply, so only the local work is timed. - Each iteration runs a no-grad forward, a grad forward and a backward, as under reentrant activation checkpointing. - The whole local chain is timed; only the reorder differs between the two columns. | T | Before | After | Saving × 48 layers | |---|---|---|---| | 8192 | 7.04 ms | 5.12 ms | −92 ms | | 16384 | 13.53 ms | 9.58 ms | −190 ms | At T=8192, per grad forward, the permute goes from 0.44 to 0.27 ms and the unpermute from 0.71 to 0.21 ms. **End to end**, Qwen3-30B-A3B training with AutoEP: - 16× H100, expert parallel 16, 48 layers, micro-batch 1, BF16, ZeRO stage 1, reentrant activation checkpointing, and a fixed routing, so both arms route identically. - `comm_backend="comm"` with `combine_impl="fused_weighted_sum"` and `async_split_plan=true`. - Both arms also carry other opt-in kernels under review: fused RMSNorm (deepspeedai#8645, deepspeedai#8664), fused RoPE (deepspeedai#8663), chunked loss (deepspeedai#8653) and BF16 gradient release (deepspeedai#8654). The arms differ only in this PR. - Both arms ran on the same 16 GPUs in mirrored order, two runs each. | Sequence | Before | After | Change | |---|---|---|---| | 8192 | 2649.4 ms | 2554.3 ms | −95.2 ms (−3.6%) | | 16384 | 5434.6 ms | 5244.5 ms | −190.2 ms (−3.5%) | - Each arm's two runs agree within 5 ms. - The end-to-end saving matches the local estimate scaled to 48 layers (−92 and −190 ms). - Peak memory and the step-0 loss are identical in both arms. These numbers come from one model, two sequence lengths, one GPU type and one parallel layout. The saving scales with the rows each rank reorders, so its share of the step depends on the configuration. Small shapes are not guaranteed to improve in every direction: a review microbenchmark on RTX PRO 4500 found slower forwards at 8,192 rows × hidden 1,025, while both backwards improved. This PR does not add a shape-dependent dispatch heuristic. The local and end-to-end numbers were measured with this PR's first commit. The later commits add `index.contiguous()`, a no-op for the contiguous indices both callers pass, and change docstrings and tests. The forked-runner review follow-up changes only test execution, not the runtime implementation. ## Tests `tests/unit/v1/moe/test_autoep_permute_rows.py`, 27 tests: 26 need CUDA and Triton, and 1 runs on the host. All compare against the previous zero-row-and-index implementation with `torch.equal`: - forward and backward of both reorders: - bf16 and fp32; - hidden 2048, 100 and 1; - source-major `[ep, E_local]` counts with and without empty experts, and aggregated `[E_local]` counts; - fp16 and bf16 rows read through a column-strided view, with column-major upstream gradients, in both directions, including the output strides; - double backward of both reorders; - a round trip with zero padding; - zero tokens; - the row gather with both padding encodings, through a contiguous and a strided index; - host tensors keeping the previous code. Results: - 1× NVIDIA H100 80GB HBM3, Python 3.12, torch 2.10 (CUDA 13.0), Triton 3.6.0: 27 passed, and 217 together with `tests/unit/v1/moe/test_autoep_unit.py`. - The same 27 tests against this PR's first commit, the code the timings used: 26 passed. The strided-index case failed, because that commit read the index as if contiguous. - CPU, torch 2.14: 1 passed, 26 skipped. - Forked-runner review regression on H100, torch 2.10 (CUDA 13.0), Triton 3.6.0: plain pytest, `pytest --forked` and `pytest -n 4 --forked` each pass all 27 cases, with no skips and no bad-fork errors. In the same environment, the previous tests fail under `--forked` (26 failed, 1 host case passed), reproducing the reported issue. - CPU review regression, torch 2.14: `--forked` and `-n 4 --forked` each give 1 passed, 26 skipped. - `pre-commit run --files` on the changed files: passed. ## CI - `cpu-torch-latest` runs the host test and skips the rest. - The test file is under `tests/unit/v1`, in the scope of both AWS and Modal GPU CI; Modal executes the selected GPU tests when the PR enters the merge queue. - `aws-torch-latest` uses `pytest -n 4 --forked`. The GPU cases now use `DistributedTest` with `world_size=1`, launching clean workers even when CUDA was touched while importing DeepSpeed during collection. These local operator tests do not need a process group (`init_distributed=False`). CUDA availability is checked by the runner during test execution, not by this file's collection-time skip condition. The host case remains an ordinary test, and all 27 cases and their numerical assertions are unchanged. ## Not covered - ROCm: PyTorch reports ROCm devices as `cuda`, so the gathers would be used there, as the existing Triton index generation is. Neither was tested on ROCm here. - Inputs that break the permutation contract. It is not checked; both callers meet it by construction. - `torch.compile` of the MoE layer itself. AutoEP's regional compilation keeps that layer eager, and the path already launched Triton kernels there. - Models, sequence lengths, hardware and parallel layouts other than the ones above. --------- Signed-off-by: yh0903 <helloyu0903@gmail.com> Co-authored-by: pengdurice <pengduhit@gmail.com>
banxingmjj
pushed a commit
to openanolis/DeepSpeed
that referenced
this pull request
Oct 10, 2026
deepspeedai#8645) ## Summary Adds `deepspeed/ops/triton_ops/fused_rms_norm.py`: - `fused_rms_norm(hidden, weight, eps)`: Triton forward and backward kernels for the RMSNorm expression that Llama-style Hugging Face models share, `weight * (x.float() * rsqrt(mean(x.float()**2) + eps)).to(x.dtype)`. - `replace_rms_norm(model) -> int`: an explicit, opt-in installer that runs the kernels in place of the forward of each supported RMSNorm module (the RMSNorm classes of Llama, Mistral, Mixtral, Phi-3, Qwen2, Qwen2-MoE, Qwen3, Qwen3-MoE, DeepSeek-V2 and DeepSeek-V3) and returns how many modules it replaced. For inputs the kernels do not take, a replaced module runs its own eager forward. Nothing changes unless `replace_rms_norm` (or `fused_rms_norm`) is called. In the AutoEP Qwen3-30B-A3B training setup described under Performance, it cut the median step time by **9.0%**, with peak memory unchanged. ## Module copying and checkpoints The installer uses a class descriptor and a per-instance opt-in marker, rather than storing a bound `forward` method on each module. The dispatcher is process-wide, but instances not explicitly selected keep their original eager forward and remain eligible for a later opt-in. Shallow/deep copies and `torch.nn.DataParallel` replicas retain their live opt-in and bind to their own receiver, weights and epsilon. HF module classes, Parameters and state-dict keys are unchanged. Full-model `torch.save(model, ...)` / `torch.load(..., weights_only=False)` clears the opt-in: the marker serializes as `False`, not as an enabled flag. The loaded model remains eager until `replace_rms_norm(model)` is called again, even when another instance has already installed the class dispatcher in the loading process. Legacy boolean markers also cannot enable the new dispatcher. The original class/instance-patch guards, idempotence and unsupported-input fallback are preserved. This follow-up preserves the earlier master merge and receiver-binding fix; it does not include the narrow-row scheduling changes from deepspeedai#8664. ## Semantics The kernels keep the eager order: FP32 variance and normalization, one rounding of the normalized value to the input dtype, then the multiply by the weight (gamma). The backward rounds where eager autograd rounds: `grad_out * weight` is rounded to the input dtype before use, and the gamma-gradient product `grad_out * normalized` is rounded to the input dtype before its FP32 column sum. Results are still not bitwise identical to eager, because the FP32 reductions sum in a different order. ## Which modules are replaced Instances of the exact classes listed in `SUPPORTED_RMS_NORM_CLASSES` that still run their class's own `forward` and are at most 2048 wide. The list is explicit because a class name ending in `RMSNorm`, a 1-D `weight` and a float `variance_epsilon` do not determine the math. In transformers 5.17, 140 `*RMSNorm` classes pass that rule, and 8 of them multiply by the weight *before* the cast: `GptOssRMSNorm`, `Olmo2RMSNorm`, `Olmo3RMSNorm`, `FlexOlmoRMSNorm`, `OlmoHybridRMSNorm`, `HeliumRMSNorm`, `AfmoeRMSNorm` and `OpenAIPrivacyFilterRMSNorm`. `Olmo2RMSNorm` changed order between releases: it casts before gamma in 4.51.3 and after gamma in 5.17. The kernels would run on those modules without error and silently change their numerics, so falling back to eager cannot catch this; the selection has to. The installer therefore: - replaces only exact instances of the listed classes. CPU tests hold the currently listed classes to the kernels' expression using `torch.equal` (which does not distinguish signed zero), so a class that changes its cast order upstream fails CI. Adding a class also requires updating the collection-safe class-name metadata in the tests; - leaves subclasses alone, as well as modules whose `forward` was patched or wrapped on the class or on the instance, whether by another kernel installer or by an earlier call (so a second call replaces nothing). A wrapper made with `functools.wraps` carries the original's name, so it is recognized by the `__wrapped__` attribute that `functools.wraps` sets; - leaves norms wider than 2048 eager, because the kernels hold a whole row in one block. On Qwen3-30B-A3B it replaces all 193 RMSNorm modules: `input_layernorm`, `post_attention_layernorm`, `q_norm` and `k_norm` in each of the 48 layers, plus the final norm. On Qwen3-235B-A22B (hidden size 4096) it replaces the 188 128-wide `q_norm` and `k_norm` modules and leaves the 189 4096-wide norms eager. ## Where the kernels run | | | |---|---| | Kernels | Input and weight are bfloat16 or float16 CUDA tensors of the same dtype, and Triton is available on CUDA (not ROCm). | | Everything else | The module runs its class's own eager forward, so the result is exactly eager's, and a warning is logged once. This includes CPU tensors, float32 inputs and float32 weights under autocast, so the installer can run before the model moves to the GPU. | | Layout | Any leading shape. Non-contiguous inputs and weights are copied before the kernels run. The kernels' output and input gradient are contiguous, whereas eager follows the input and upstream-gradient strides. | | Autograd | First-order gradients for input and weight; `eps` is a constant. On the kernel path, double backward raises (`once_differentiable`) instead of silently dropping this op's second derivative. `torch.compile` and `torch.func` transforms are not covered. | | `fused_rms_norm` | Applies the kernels directly and raises on unsupported inputs instead of falling back. | ## Correctness evidence Three separate questions, each with its own evidence. ### 1. Does the operator compute the expression? (unit tests, random inputs) The reference is the eager expression in the same dtype. The metric is the distance in units in the last place (ULP), computed from ordered bit patterns: - forward: max ≤ 2 ULP, and ≥ 99% of elements within 1 ULP - input and gamma gradients: max ≤ 8 ULP, and ≥ 95% within 1 ULP (≤ 16 ULP and ≥ 90% at input scales of 1e-4 and 1e4) Coverage: bfloat16 and float16; widths 128, 130, 256 and 2048, plus an empty input; head-major non-contiguous input; strided weight; and the installer on every listed class, against that class's own eager forward, with a non-default epsilon and the head-major upstream gradient that attention produces. Every random input comes from a seeded generator. ### 2. Does it hold on real training tensors? (frozen replay, Qwen3-30B-A3B) 21 RMSNorm boundaries were captured from one eager Qwen3-30B-A3B training step (48 layers, EP16, sequence length 1024): all four norm kinds at layers 0, 11, 23, 35 and 47, plus the final norm, with the real inputs, gamma, epsilon, upstream gradients and strides. The eager replay reproduces the captured outputs exactly (relative L2 of 0 in all 21 cases). Fused against eager, on one H100: | | Result | |---|---| | Gamma gradient | Bitwise identical in 20 of 21 cases; one element differs in the remaining case | | Input gradient | 15–136 differing elements per case (out of 0.5M–4.2M), relative L2 ≤ 5e-5 | | Forward | 0–14 differing elements per case | | Elements outside \|a − b\| ≤ 2⁻⁶ + 2⁻⁶·\|b\| (BF16 policy) | 0 in every case | The forward differences come from the FP32 summation order of the variance, which a `torch.compile`d RMSNorm does not reproduce either. No replay was run for the other listed classes; for them, the evidence is the unit tests above plus each class computing the same expression bitwise. ### 3. What does it do to a short training run? (sensitivity, not equivalence) These are paired 24-step runs in the setup described under Performance, with 4 repeats per arm: - step 0: mean loss differs by about 4e-4 (forward rounding) - step 1: 17.100 eager vs 17.044 fused (−0.056), where eager repeats agree to within 0.004 - max |Δ mean loss| over the 24 steps: fused 0.058; eager vs eager 0.013–0.019 A control changes only the eager RMSNorm's variance reduction, computing it in FP64 before the usual FP32 steps, which gives a more accurate eager norm. It moved step 1 by +0.043 and reached a max |Δ| of 0.049. So a 0.02 threshold on short-run loss also rejects a more accurate eager norm: ULP-level changes in the norm move the first AdamW steps by more than 0.02, and the threshold cannot decide on its own whether a kernel is correct. Runs this short cannot establish long-run training equivalence, and this PR does not claim it. ## Performance Measured setup, identical in both arms except for the RMSNorm implementation: - Qwen3-30B-A3B, 48 layers, sequence length 8192, micro-batch 1 per GPU - AutoEP with the DeepEP backend, EP16 on 2 nodes × 8 H100 80GB, with both arms on the same allocation - BF16, ZeRO stage 1 with FP32 gradient accumulation, AdamW, reentrant activation checkpointing, Python GC disabled during training - a fixed routing schedule (identical route histograms in both arms); HF router logits on, so both arms also build the zero-coefficient load-balancing loss - both arms include two changes that are not part of this PR: the DeepEP cached-dispatch layout fix (deepspeedai#8636) and the opt-in fused AutoEP row weighting (deepspeedai#8644) - 24 steps, of which the first 12 are excluded; step time is the synchronized training step, taking the maximum across ranks - one discarded preflight run per arm, then 4 runs per arm in ABBA BAAB order | | Median step (ms) | Per-repeat medians (ms) | |---|---|---| | Eager RMSNorm | 1934.3 | 1881.2, 1910.7, 1977.1, 1958.0 | | Fused RMSNorm | 1760.4 | 1814.3, 1706.5, 1693.1, 1869.3 | That is −174 ms (−9.0%). All four paired repeats are faster (−3.6%, −10.7%, −14.4%, −4.5%), and the 95% t-interval over those four deltas is [−16.5%, −0.1%]. Peak memory is unchanged (58.3 GiB allocated, 71.2 GiB reserved). An earlier screen on a different allocation, run before the gamma-gradient rounding fix, measured −8.96%. Scope: one model, one parallel configuration and one sequence length, on H100. The interval is wide, so read this as "about 9% in this configuration", not as a general speedup. Not measured: the other listed model families, other widths, sequence lengths or EP sizes, other GPUs, ZeRO-3, `torch.compile`, and non-reentrant checkpointing. The measured kernels are byte-identical to this PR's. Since those runs, the installer covers more classes and is stricter about which modules it takes (it still replaces the same 193 modules on this model), replaced modules fall back to eager for inputs the kernels do not take, `fused_rms_norm` copies a non-contiguous weight first (a no-op for the contiguous weights measured), and the backward became `once_differentiable` (first-order gradients unchanged). On supported inputs a replaced module still makes one input check per call and then runs the same kernels. The receiver-binding and checkpoint opt-in follow-ups change only module dispatch and checkpoint behavior, not the kernels, autograd or direct operator implementation. The existing performance and replay evidence is reused; there is no new timing or full E2E run, and the historical percentage is not an independently retimed measurement of the new dispatcher. ## Tests `tests/unit/v1/ops/triton_ops/test_fused_rms_norm.py`: 70 cases, including the original numerical and receiver/checkpoint coverage: - GPU (31 original numerical/installer cases): forward and backward against eager (dtypes × widths, empty input, head-major input, extreme magnitudes, strided weight); `fused_rms_norm` rejecting widths above 2048, float32 and mixed dtypes; double backward raising; the installer on each listed class against that class's own eager forward (count, the module's own epsilon, gradients reaching the module weights, a second call replacing nothing); a replaced module taking the kernels for supported input and returning exactly eager's result for float32 input. - CPU (30 original cases): device and dtype guards of `fused_rms_norm`; each listed class matching the kernels' expression under `torch.equal`, on inputs where the two cast orders differ; each listed class replaced and still returning exactly its eager result for CPU input; the 2048-wide boundary; a local name-alike class, `Olmo2RMSNorm`, `GptOssRMSNorm`, a `Qwen3MoeRMSNorm` subclass, and class-patched, class-wrapped (`functools.wraps`) and instance-patched forwards all left unreplaced. - Receiver/checkpoint cases (9): shallow/deep copies retaining opt-in with their own epsilon on CPU/CUDA; full-model round trips on CPU/CUDA that retain eager layout until explicit reinstallation and check independently derived forward/input/weight gradients; loading in a fresh interpreter with and without a previously installed dispatcher; and actual two-GPU `DataParallel` training against the eager expression over two SGD steps. Constant rows make the extra BF16 DP gradient reduction exact; the existing forward/gradient ULP thresholds are unchanged. The existing conditional-kernel test also verifies that an unselected sibling stays eager and can opt in later. Results: - Current checkpoint/forked-test follow-up on two NVIDIA H100 80GB HBM3 GPUs, torch 2.10/CUDA 13.0, Triton 3.6.0 and transformers 5.12.0: **70 passed, 0 skipped** under `pytest --forked`. The 9 receiver/checkpoint/DataParallel cases also pass `pytest -n 4 --forked`. - Collection regression on the same hardware: the new file does not import the RMSNorm Triton module or initialize CUDA during collection, and a child forked afterwards allocates CUDA successfully. The preceding test file initializes CUDA during collection and its forked child fails with the reported bad-fork error. - With the preceding class-dispatcher runtime and the same new tests, both CPU/CUDA same-process checkpoint cases and the fresh-process/preinstalled-dispatcher case fail as expected (**3 failed, 0 skipped**). - CPU follow-up, torch 2.14 and transformers 5.12.0: **35 passed, 35 CUDA cases skipped** with `pytest --forked`. Against the preceding class-dispatcher runtime, the same-process checkpoint and fresh-process/preinstalled-dispatcher cases both fail as expected. - A local macOS / torch 2.14 `-n 4 --forked` attempt aborted in the eager reference's native autograd engine (SIGABRT); single-runner CPU validation and the Linux CI-style receiver/checkpoint matrix above passed. No assertions were removed or platform safety checks disabled to bypass that abort. - Current-master targeted receiver/installer run on NVIDIA H100 80GB HBM3: **35 passed, 0 skipped** (the 29 existing installer cases plus all six new regressions), with torch 2.10/CUDA 13.0, Triton 3.6.0 and transformers 5.12.0. The DataParallel case used two GPUs and checked two SGD steps against the independent eager expression. This was targeted validation, not a rerun of the full 67-case suite or E2E. - Old-runtime negative control with the same new tests: **all six fail as expected, 0 skipped**. Only the tests are changed in this arm, so it retains the actual prior instance-bound runtime and the same kernel/autograd code. - 1× NVIDIA H100 80GB HBM3, torch 2.10 (CUDA 13.0), Triton 3.6.0, transformers 5.12.0: 61 passed, 0 skipped. - The same tests against the first version of this PR (same kernels, only `Qwen3MoeRMSNorm` listed, so 34 collected): 14 failed, 20 passed. A strided weight put the forward up to 34,883 ULP off; double backward did not raise; the seven name-alike, subclass, and patched or wrapped modules were matched; CPU and float32 inputs raised instead of running eager; a 2048-wide norm on the CPU was rejected; and a second call replaced the norms again. Every kernel accuracy test passed in both runs. - CPU only, torch 2.14: 30 passed and 31 skipped with transformers 5.17; 27 passed and 34 skipped with transformers 4.51.3, the `requirements-dev.txt` minimum, where `DeepseekV2RMSNorm` and `GptOssRMSNorm` do not exist. - `pre-commit run --files` on the changed files: passed. ## CI - Published follow-up head: DCO and test collection passed. Full GPU CI was skipped on the PR event; the targeted H100 execution above is the new device-validation evidence. - `cpu-torch-latest` runs the 35 CPU cases, with transformers from git main on PRs. CUDA cases defer Triton imports and device-availability checks to execution, and use `DistributedTest(world_size=1, init_distributed=False)` for fresh workers under forked CI. Parametrization and collection-time skipping use only metadata. The DataParallel regression requires two CUDA devices and skips on a single-GPU runner. - For `modal-torch-latest` (L40S), `python ci/tests_fetcher.py --base upstream/master --explain` selects exactly `tests/unit/v1/ops/triton_ops/test_fused_rms_norm.py`. That workflow only selects tests on PR events and runs them on the GPU when the PR enters the merge queue, so until then the GPU results are the H100 run above. ## Merge dependency As requested in the review, deepspeedai#8737 should land before this PR. It addresses a pre-existing Transformers 4.51.3 / Hybrid Engine OPT circular import when loading a checkpoint before importing DeepSpeed. This follow-up does not fold in or work around that separate import-order bug. ## Not covered None of these apply to the measured configuration (Qwen3-30B-A3B, BF16, H100, ZeRO-1, reentrant checkpointing, contiguous weights, first-order gradients). - Norms wider than 2048 stay eager. That includes the hidden-size norms of most models of 7B parameters and up (hidden sizes of 3584 for Qwen2.5-7B, 4096 for Llama-3-8B and Mistral-7B); in those models only narrower norms, such as Qwen3's 128-wide query and key norms, are replaced. Wider rows need a kernel change and new measurements. - RMSNorm classes that are not listed, including those with different math (the Gemma family's `(1 + weight)`, GPT-OSS, OLMo 2 and 3) and classes from `trust_remote_code` modeling files. - Replay and training runs for the listed classes other than Qwen3-MoE. - ZeRO-3 partitioned weights, `torch.compile` or AutoEP regional compile around the norms, and `torch.func` transforms. - Non-reentrant activation checkpointing, and float16 end to end (float16 is covered by the unit tests only). - Long-run convergence: only 24-step runs were measured. --------- Signed-off-by: yh0903 <helloyu0903@gmail.com> Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
…ward Norms at most 128 wide, such as the query and key norms, ran one row per program, and the backward read each row twice: once for the input gradient and once for the weight-gradient partials. They now run several rows per program, and one backward kernel produces both gradients. Wider norms keep the row kernels. Signed-off-by: yh0903 <helloyu0903@gmail.com>
A width-1 norm has an analytically zero input gradient, so its computed value is rounding noise and a ULP comparison with eager measures nothing. Signed-off-by: yh0903 <helloyu0903@gmail.com>
The width-one input derivative is nonzero when epsilon is positive; test it at a measurable scale against an FP64 oracle instead of assuming zero. Cover masked columns below width 16 with the existing ULP limits and document narrow-row reduction order without claiming bitwise eager parity. Signed-off-by: yh0903 <helloyu0903@gmail.com>
Add width 129 to the existing BF16/FP16 forward/backward matrix while retaining the merged parent's eager reference and numerical thresholds. This pins the narrow/wide scheduling boundary after integrating the child kernels with the parent module, checkpoint and forked-test fixes. Signed-off-by: yh0903 <helloyu0903@gmail.com>
yh0903
force-pushed
the
yh0903/fused-rmsnorm-narrow-rows
branch
from
October 12, 2026 20:10
3827f93 to
45afc05
Compare
yh0903
changed the title
yh0903
marked this pull request as ready for review
October 12, 2026 20:11
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 45afc059cc
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Use the same row-aware power-of-two tile cap in forward and backward so short inputs do not instantiate a full 4096-element tile. Keep filled tiles and the wider row kernels unchanged.\n\nAdd host scheduling cases and real BF16/FP16 launch and gradient regressions, reusing the FP64 width-one oracle and existing numerical limits. Document the short-input scheduling bound. Signed-off-by: yh0903 <helloyu0903@gmail.com>
There was a problem hiding this comment.
🔵 Needs a closer look
The GPU-specific Triton kernel and reduction changes require final human review despite strong targeted test coverage.
0 open findings
What changed in this PR
Optimizes the opt-in fused RMSNorm from #8645 for narrow dimensions while preserving the existing wide path.
Changes:
- Adds multi-row forward and combined backward kernels for widths ≤128.
- Caps short-input tile sizes and expands numerical, scheduling, and determinism tests.
- Documents scheduling and width-one gradient behavior.
| File | Description |
|---|---|
deepspeed/ops/triton_ops/fused_rms_norm.py |
Implements narrow multi-row kernels and dispatch. |
tests/unit/v1/ops/triton_ops/test_fused_rms_norm.py |
Adds boundary, accuracy, launch, and determinism coverage. |
docs/code-docs/source/autoep.rst |
Documents narrow scheduling and numerical behavior. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
tohtana
requested changes
Oct 14, 2026
tohtana
left a comment
tohtana
left a comment
Collaborator
There was a problem hiding this comment.
Hi @yh0903,
Thank you for the optimization!
One test failed when I ran with RTX PRO 4500. The failure is the new BF16 (2, 300, 32, 128) case. One of 2,457,600 dx values is 119 ULP from eager, above the asserted maximum of 8. What about relaxing the tolerance like <8 ULP or 1e-7?
Also, the comparison using the GPU, showed 11–16% slower narrow forward. We don't need to address this now, but we might need to clarify the permance impact in a wider variety of environments.
Allow each narrow input-gradient element to satisfy the existing 8 ULP limit or an absolute error of at most 1e-7, while retaining the 95% within-one-ULP requirement. Keep forward, weight-gradient and wide-path thresholds unchanged. Cover mixed error bounds, excessive outliers and non-finite values. Document the reviewer-reported RTX PRO 4500 narrow-forward slowdown without changing kernel arithmetic or claiming a universal speedup. Signed-off-by: yh0903 <helloyu0903@gmail.com>
Contributor
Author
Thanks for checking this @tohtana . Addressed in the latest commit: the narrow-shape |
Contributor
|
quick question: is this available in other libs, e.g. liger? if so, is using them a viable choice as well? |
Contributor
Author
|
@pengdurice Yes—Liger already provides fused RMSNorm, including a multi-row/block path and joint This PR extends the existing in-tree opt-in path without adding a Liger runtime dependency; the optimization itself is not unique. Our measurements compare against the existing baselines, not Liger, so they don't establish an advantage over Liger. I'm open to reusing Liger or contributing useful scheduling changes there rather than maintaining duplicate kernels. We'd first need to verify the casting/gradient and backward-buffer contracts and compare the relevant narrow shapes. I haven't done that comparison yet, so I wouldn't claim this separate implementation is necessary on performance grounds. |
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.
Summary
Builds on the opt-in fused RMSNorm implementation merged in #8645. This branch has been rebased onto current upstream
master: the PR diff contains only narrow-row scheduling/backward changes, boundary tests and documentation, not the former parent commits.The merged parent's per-instance opt-in, checkpoint opt-out, copy/DataParallel behavior, eager fallback and collection-safe CUDA test framework are preserved.
The parent fused RMSNorm kernels run one row per program. For the 128-wide query and key norms of models like Qwen3, that leaves most of each program idle. The backward also reads
hiddenand the incoming gradient twice, once in the dx kernel and once in the weight-gradient partial kernel.For norms at most 128 wide, this change:
Wider norms, including the hidden-size norms, keep the existing row kernels unchanged.
The short-input cap addresses the Codex review in this existing PR, rather than opening a duplicate optimization PR. Both passes use
min(max_rows_per_tile, next_power_of_2(n_rows)); a single width-128 or width-one row usesBLOCK_M=1, not 32 or 256 padded rows. Filled tiles keep their original scheduling. AI assistance was used for this follow-up, with author review before submission.Numerics
The expression and its rounding points are the same as the parent row kernels'. Multi-row scheduling can change the FP32 reduction order for variance,
grad_rstd, and weight-gradient partials. Bitwise equality with eager is not claimed. Forward, gamma-gradient and wide-path numerical thresholds are unchanged. The narrow-shape input-gradient comparison accepts each element at at most 8 ULP or absolute error at most 1e-7, while still requiring at least 95% within one ULP. This bounds near-zero cancellation errors without increasing the ULP limit for other elements. Identical fused runs give repeatable input and weight gradients at widths 128 and 2048, which is a separate determinism property.At width one and
eps > 0, the pre-gamma input derivative iseps / (x**2 + eps)**(3/2), not strictly zero. A tiny epsilon can make that value hard to distinguish from low-precision cancellation. The dedicated width-one test useseps=0.25and an FP64 reference so the nonzero gradient is measurable, without relaxing the existing ULP thresholds.Tests
The existing
test_fused_rms_norm.pyforward/backward test gains narrow shapes:(4099, 128);(2, 300, 32, 128);A width-one test checks the nonzero input derivative against FP64 and forward/gradient parity with eager. The repeated-run test checks fused input and weight gradients at widths 128 and 2048. Existing tests cover head-major non-contiguous inputs, transposed incoming gradients, and non-contiguous weights.
The first width outside the narrow cutoff, 129, is additionally checked in BF16/FP16 under the unchanged wide-path accuracy contract. The shape, FP64 width-one and determinism cases now run inside the merged parent's
DistributedTestCUDA class, with Triton imports and availability checks deferred to execution.Earlier validation passed the six BF16/FP16 sub-16 and width-one cases on one NVIDIA H100 80GB HBM3 (torch 2.10/CUDA 13.0, Triton 3.6).
Pre-cap rebased validation on two NVIDIA H100 80GB HBM3 GPUs, torch 2.10/CUDA 13.0, Triton 3.6.0 and transformers 5.12.0: 90 passed, 0 skipped under
pytest --forked. This includes the narrow numerical/masked-width/FP64/determinism tests and the merged parent's guards, copying, checkpoint loading/reinstallation and two-GPU DataParallel training.A separate public-operator comparison with the merged parent at widths 129, 130, 256 and 2048 passes for BF16 and FP16: all sampled outputs, input gradients and weight gradients have identical bit patterns in these 8 comparisons. The wide kernel implementations and parent dispatch/checkpoint logic are also unchanged by AST comparison.
Short-row follow-up in the same pinned two-H100 environment: 56 focused tests passed, 0 skipped, including 21 explicit scheduling cases and BF16/FP16 regressions that record both real Triton launches while checking forward/dx/dgamma under the unchanged ULP limits. These cover empty, single-row, non-power-of-two, tile-boundary and filled inputs, the FP64 width-one oracle, and affected copying, checkpoint round-trip and two-GPU DataParallel paths.
The two new real-launch regressions fail against the previous runtime: a single width-one row launches
BLOCK_M=256in both passes instead of 1. Separate public-operator controls at(8192, 128),(257, 16),(7, 129)and(7, 2048)pass in BF16/FP16, with identical output/dx/dgamma bit patterns in all 8 before/after comparisons.Focused GPU selector, from
tests/:PYTHONPATH=.. pytest --forked unit/v1/ops/triton_ops/test_fused_rms_norm.py \ -k 'narrow_tile or matches_hf_forward_and_backward or width_one or copy_uses_its_own_epsilon or model_round_trip or data_parallel'The pre-numerical-review file on CPU (macOS, torch 2.10.0, transformers 5.12.0) with
DS_ACCELERATOR=cpuandpytest --forked: 56 passed, 57 CUDA cases skipped. Existing support limits and opt-in/fallback behavior are unchanged; ZeRO/offload/TP/PP and other GPU architectures were not newly exercised by this bounded follow-up. No full E2E or new performance matrix was rerun, and no measured latency claim for short decoding batches is made.RTX numerical-review follow-up: the reviewer reported one of 2,457,600 BF16
dxvalues at 119 ULP in(2, 300, 32, 128)on RTX PRO 4500 and suggested an absolute alternative. Only the narrow commondxassertion now has the bounded alternative above; kernels, scheduling, forward/dgamma checks and the separate FP64 width-one oracle are unchanged. Five CPU assertion regressions cover mixed per-element ULP/absolute acceptance, strict default behavior, errors beyond both limits, too many small outliers, and NaN/Inf rejection. A synthetic BF16 119-ULP example has absolute error 1.39e-8; it is not the reviewer's tensor.The current whole file passes on CPU under
pytest --forked: 61 passed, 57 CUDA cases skipped. The focused H100 run passes 9 tests, 0 skipped under--forked: the five assertion cases, BF16/FP16 at the reported shape and both short-tile correctness/launch regressions. The allocation has 2x H100 80GB HBM3, driver 580.159.04, torch 2.10.0.2/CUDA 13.0, Triton 3.6.0 and transformers 5.12.0. On that H100 environment the reported shape has maxdxdistance 1 ULP in BF16 and 2 ULP in FP16; neither reproduces the RTX outlier. RTX PRO 4500 has not been rerun by the author for this follow-up, so confirmation of the reported case on that hardware remains with review. Changed-files pre-commit passed.Focused reproduction, from
tests/:PYTHONPATH=.. pytest --forked unit/v1/ops/triton_ops/test_fused_rms_norm.py \ -k 'dx_tolerance or (matches_hf_forward_and_backward and shape6) or narrow_tiles'Performance
Hardware-dependent tradeoff: the reviewer measured narrow forward 11-16% slower on RTX PRO 4500 than the parent row kernels. This is a reviewer-reported forward-only observation, not an author-reproduced timing or a measured E2E regression. The H100 training benefits below do not imply a universal speedup; the path remains opt-in. No new performance tuning or benchmark matrix was added in this review follow-up.
Incremental local effect of this PR: measured on one H100, Qwen3-30B-A3B attention layer under reentrant checkpointing (forward, recompute and backward), with fused RoPE already enabled in both arms.
Combined E2E effect: Qwen3-30B-A3B training with AutoEP, 16x H100, expert parallel 16, 48 layers, one sequence per GPU, activation checkpointing. The "after" arm adds both this narrow-row RMSNorm change and fused RoPE (#8663) on an otherwise identical stack, in the same allocation. The total change cannot be attributed to this PR alone:
The narrow kernel arithmetic and filled-tile scheduling used by these measured workloads are unchanged; only underfilled short inputs use the smaller row tile. Rebase/conflict resolution preserves the merged parent rather than reinstating the former instance-bound forward or serialized opt-in behavior. The performance numbers above are reused historical measurements, not a retiming of the rebased stack or the short-row follow-up.
CI
ci/tests_fetcher.pyagainst current master selects exactlytests/unit/v1/ops/triton_ops/test_fused_rms_norm.py. The file retains the parent's fork-safe worker execution and CPU coverage; the DataParallel case additionally requires two CUDA devices. Full GPU CI follows the repository's merge-queue workflow.