DeepSpeed

583fc9c3 - Add an opt-in fused Triton RoPE for Hugging Face split-half models (#8663)

Add an opt-in fused Triton RoPE for Hugging Face split-half models (#8663) ## Summary Adds an opt-in fused Triton kernel for Hugging Face's split-half rotary position embedding, `apply_rotary_pos_emb`, as used by Llama-, Mistral-, Qwen- and DeepSeek-V3-style models: ```python from deepspeed.ops.triton_ops.fused_rotary_pos_emb import replace_rotary_pos_emb patched = replace_rotary_pos_emb(model) # number of modeling modules patched; 1 for Qwen3-30B-A3B ``` The eager function, `q * cos + rotate_half(q) * sin`, runs several elementwise kernels, each reading and writing the whole query or key tensor, and its backward does the same again. The fused kernel reads and writes each tensor once per pass and applies each position's `cos`/`sin` to all of that position's heads. Nothing changes unless `replace_rotary_pos_emb` is called. ## Semantics and safety - **Process-wide patch.** Attention layers look `apply_rotary_pos_emb` up in their modeling module, so the installer replaces it there. The replacement applies to every model of that architecture in the process, not only to `model`. Calling it again is a no-op, and `restore_rotary_pos_emb()` undoes it. - **Which modules.** A modeling module is patched only if all of these hold: - one of the model's submodules is defined in it; - it is listed in `SUPPORTED_ROTARY_MODULES` (Llama, Mistral, Mixtral, Qwen2, Qwen2-MoE, Qwen3, Qwen3-MoE, DeepSeek-V3); - its `apply_rotary_pos_emb` and `rotate_half` still have the split-half expression's code: the same bytecode, names, constants and defaults, whatever their docstrings. Constants are compared because bytecode refers to them by index: a `rotate_half` that splits at a quarter, or along another axis, otherwise has the same bytecode. Several architectures define functions of the same name that rotate part of the head dimension or add casts (for example Phi-3 and OLMo-2). A module whose functions differ is left alone with a warning. This includes transformers releases whose `apply_rotary_pos_emb` still takes `position_ids`, such as 4.51. - **Fallback.** For any input the kernel does not cover, the patched function runs the original eager function and warns once. Uncovered inputs are: - dtypes other than bf16/fp16, or mixed dtypes; - tensors not on one CUDA device; - an odd head dimension, one over 512, or one without unit stride; - shapes whose batch, sequence or head dimension do not match `cos`/`sin`; - `unsqueeze_dim` other than 1 or 2; - `cos`/`sin` that require grad. Without Triton, or on ROCm, it always runs eager. - **Direct call.** `fused_apply_rotary_pos_emb` applies the kernel directly and raises on unsupported inputs instead. - **Launch device.** Forward and backward select the inputs' CUDA device for the Triton launch, use that device's current stream, and restore the caller's current device. Inputs need not be on the device that was current when the function was called. - **Autograd.** Gradients are first order, for `q` and `k`; `cos` and `sin` are constants. Double backward raises. ## Numerics and layout Every product and sum is rounded to the input dtype where the eager expression rounds it. The kernel is compiled with `enable_fp_fusion=False`, because an FMA contraction would skip a product's rounding. Outputs and the gradients for `q` and `k` equal eager's element for element under `torch.equal`, which does not distinguish `+0.0` from `-0.0`. Layout: - For dense inputs, such as attention's transposed projection, outputs keep the input strides, as eager's do. - Other non-contiguous inputs can give outputs laid out differently from eager's, with equal values. - The gradients for `q` and `k` take the outputs' layout, whereas eager's follow the incoming gradient. Weight gradients computed from them further back can therefore differ in the last bits (see the attention block below). ## Validation Three levels, reported separately. **1. Operator.** Earlier validation passed the original 32 cases on one H100. The review follow-up adds two cross-device cases (BF16/FP16); current regression results are listed below. **2. Full attention block, one H100.** Qwen3-30B-A3B attention layer (32 query heads, 4 KV heads, head dim 128). The unit is forward, recompute and backward under reentrant activation checkpointing; each arm ran twice in mirrored order. | Sequence | Eager RoPE | Fused RoPE | Saved per layer | RoPE segment | |---|---|---|---|---| | 8192 | 9.27 ms | 7.80 ms | 1.46 ms | 1.57 → 0.19 ms | | 16384 | 26.62 ms | 23.66 ms | 2.96 ms | 3.13 → 0.39 ms | - The block output and the block's input gradient equal eager's. - The `q_proj` and input-norm weight gradients differ from eager's by a relative 4e-5. A `torch.compile`d RoPE moves them too (8e-6 and 1e-5). This comes from the layout of the gradient the weight GEMMs receive, not from RoPE's arithmetic. - The fused kernel is about 0.2–0.4 ms per layer faster than a `torch.compile`d RoPE configured to reproduce eager rounding. **3. End to end, together with the narrow-row RMSNorm change.** AutoEP training of Qwen3-30B-A3B: - 16× H100, expert parallel 16, 48 layers, one sequence per GPU, BF16, activation checkpointing. - Both arms use other opt-in AutoEP changes, including the fused row weighting (#8644), fused RMSNorm (#8645) and chunked loss (#8653). - "After" adds this PR and a narrow-row path for the fused RMSNorm (a follow-up to #8645), and nothing else. - Both arms ran on the same 16 GPUs, two runs each. | Sequence | Before | After | Change | |---|---|---|---| | 8192 | 1713.8 ms | 1654.2 ms | −3.5% | | 16384 | 3631.7 ms | 3510.5 ms | −3.3% | - At 8192 the result separates: both "after" runs are below both "before" runs. At 16384 the run ranges overlap, so that result is indicative. - These numbers cover both changes and cannot be split between them. In the attention-block measurement scaled to 48 layers, RoPE accounts for 70 of the 84 ms the two save at 8192, and 142 of 179 ms at 16384. End to end, about 70% of that combined estimate shows up. - Peak allocated memory is unchanged. - The step-0 loss differs by 8e-4. That comes from the narrow-row RMSNorm's reduction order, since RoPE's forward equals eager's. Final-step losses differ by at most 0.0023, within the 0.0016–0.0068 spread between runs of the same arm. This is a short-run check, not a convergence result. The attention-block and end-to-end numbers were measured with this PR's kernel and autograd function. Earlier follow-ups hardened the installer check, tests and documentation. The multi-GPU review follow-up adds an input-device guard to the shared forward/backward launcher; kernel arithmetic and launch configuration are unchanged. The existing performance measurements were not rerun for this correctness fix. ## Tests `tests/unit/v1/ops/triton_ops/test_fused_rotary_pos_emb.py`, 34 tests. They need CUDA and Triton and are skipped otherwise; the two cross-device cases additionally require two visible CUDA devices. Comparisons with eager use `torch.equal`. Coverage: - forward and backward in bf16/fp16 across shapes, for attention's transposed projection and for contiguous inputs, including output and gradient strides; - `unsqueeze_dim=2`, a table broadcast over the batch, non-contiguous incoming gradients, an unused output leaving no gradient, an empty sequence, and double backward raising; - every rejected input on the direct call, and the eager fallback of the installed function; - an installed two-layer Qwen3 (eager attention, small shapes) whose logits and every parameter gradient equal eager's; - installer idempotence and restore, and an unlisted model; - refusal of a changed expression, and of two `rotate_half` variants that differ from the split-half one only in constants; - acceptance of functions that differ only in their docstrings. - BF16/FP16 inputs and a non-default stream on one device while another device is current: compare outputs and both input gradients against eager, and check caller device and stream restoration. The rejected-input cases are parametrized by name and built inside the test, so collection creates no tensors on any device. Results: - 1× NVIDIA H100 80GB HBM3, Python 3.12, torch 2.10 (CUDA 13.0), Triton 3.6.0, transformers 5.12.0: 32 passed. - The same tests against the previous installer check, which compared bytecode, names and defaults but not constants: 30 passed. The two constant-variant tests failed because that check patched the module. - CPU review regression: 34 skipped under `pytest -n 4 --forked`. - Multi-GPU review regression on two NVIDIA H100 80GB HBM3 GPUs, torch 2.10 (CUDA 13.0), Triton 3.6.0, transformers 5.12.0: the full suite passes all 34 cases under `pytest --forked`, with no skips. The two BF16/FP16 cross-device cases also pass plain pytest and `pytest -n 4 --forked`. - With the same new tests and the previous runtime, both cross-device cases fail with the reported `Pointer argument ... cannot be accessed from Triton` error. This confirms the device guard fixes the reported issue. - Code check without a GPU, on transformers 5.12 and 5.17 under Python 3.13 and 3.14: all eight listed modules still match. The functions of Phi-3, OLMo-2, GPT-NeoX, Qwen3-Next, GLM-4-MoE, ERNIE-4.5-MoE and Cohere, none of them listed, are refused. - `pre-commit run --files` on the changed files: passed. ## CI - `cpu-torch-latest` collects the file and skips it. - `python ci/tests_fetcher.py --base upstream/master --explain` selects exactly this file. `modal-torch-latest` runs it on the GPU, without `--forked`, when the PR enters the merge queue. - `aws-torch-latest` runs `tests/unit/v1` with `pytest --forked`. The cases now use `DistributedTest(world_size=1, init_distributed=False)` to run in clean workers despite CUDA state created during collection. This file no longer probes CUDA availability in its collection-time skip condition; process-wide installer patches are restored in the worker that installed them. Documentation: a new "Fused rotary position embedding (experimental)" section in `docs/code-docs/source/autoep.rst`, before the AutoEP constraints. ## Not covered - Architectures outside `SUPPORTED_ROTARY_MODULES`. For example, Gemma's functions also match the split-half expression but are not listed or tested. - Transformers releases whose `apply_rotary_pos_emb` takes `position_ids`, which are left unpatched. - Partial, interleaved or other RoPE variants, fp32, ROCm, and `torch.compile`/`torch.func` transforms. - Matching eager's layout for non-dense inputs; the values are equal. - Splitting the end-to-end gain between this PR and the narrow-row RMSNorm change, and long-run convergence. --------- Signed-off-by: yh0903 <helloyu0903@gmail.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>