Repository navigation
Fix inference layer_norm/rms_norm launching no kernel for fp32 hidden sizes 65-128 - #8827
jayden1711 wants to merge 2 commits into
Conversation
…ng threadsPerGroup==32 launch) Fixes deepspeedai#8826 Signed-off-by: Jayden Vasquez <jaydenvasquez1711@gmail.com>
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: f8386f9296
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| } else if (threadsPerGroup == 16) { | ||
| LAUNCH_FUSED_LN(1, 16, maxThreads); | ||
| } else if (threadsPerGroup == 32) { | ||
| LAUNCH_FUSED_LN(1, 32, maxThreads); |
There was a problem hiding this comment.
When float32 hidden size is 96 or 128 and rows > 8 is not divisible by 8 (for example an input shaped (1, 9, 128)), this branch uses the scheduler above to launch block=(32, 8) with grid=ceil(rows / 8). fused_ln receives no row count and unconditionally processes every threadIdx.y, so the last block reads and writes up to seven rows beyond the input/output allocations; the residual and RMS branches added here have the same layout. The new tests only exercise row counts 1, 32, 128, and 4096, so make these launches tail-safe before enabling them.
AGENTS.md reference: AGENTS.md:L19-L19
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed with compute-sanitizer memcheck: at fp32 H=96/128 with rows 9/13/17/33, the new 32-thread launches overran the last block (70-224 errors per kernel). Fixed in the latest commit by launching one row per block (grid(rows), block(32)), as the >128 path does; memcheck now reports 0 errors across all 40 cases (5 bindings x H 96/128 x rows 9/13/17/33). Tests now include seq_len 9 and a direct test of the CUDA residual bindings.
Note: the existing 16-thread (and 8-thread) branches on master have the same pre-existing overrun when rows exceed the per-block row count and aren't a multiple of it. I'll file that separately.
… in the kernels) Addresses review: partially filled multi-row blocks read/wrote past the last row. Adds seq_len 9 and a direct CUDA residual test. Signed-off-by: Jayden Vasquez <jaydenvasquez1711@gmail.com>
This branch has not been deployed
1 participant
Fixes #8826
Summary
For float32, one 16-byte load holds 4 elements. So hidden sizes 65..128 take the
elems_per_row <= 128sub-block schedule withthreadsPerGroup = next_pow2(ceil(H / 4)) = 32. The dispatch in that branch only handled 1/2/4/8/16 threads, so no kernel was launched and the output tensor (at::empty_like) was returned uninitialized. This PR adds the missingthreadsPerGroup == 32branch to the four sub-block launchers incsrc/transformer/inference/csrc/:launch_fused_ln,launch_fused_residual_lnandlaunch_fused_residual_ln_store_pre_ln_resinlayer_norm.culaunch_rms_norminrms_norm.cureduce::partitioned_block<..., 32>already supports 32 threads (a full-warp reduce). fp16/bf16 (8 elements per load) never reach 32 threads in this branch, so their behaviour is unchanged. +8 lines, no other changes.End-to-end evidence
deepspeed.init_inference(model, dtype=torch.float32, replace_with_kernel_inject=True)ongoogle/bert_uncased_L-2_H-128_A-2(BERT-tiny, hidden 128), withlast_hidden_statecompared against a float64 Hugging Face run. Script and outputs: https://gist.github.com/jayden1711/c81e4ba679caa804ab6cf1a91a9019b9Tests
I extended the existing parametrization in
tests/unit/ops/transformer/inference/test_layer_norm.pyandtest_rms_norm.pyfromchannels in [384, ..., 14432]to also include 64, 96 and 128. The existing tests only used channels >= 384, so they could not reach this branch.Commands (from
tests/, matching thenv-inferenceworkflow's marker):5f266e2)test_layer_norm.py(extended)test_rms_norm.py(extended)All 24 failures on master are float32 at channels 96 and 128 (
test_layer_norm,test_rms_norm,test_pre_norm). All the original channel sizes pass before and after, and so does channel 64 (it maps to 16 threads).Note on the residual layer-norm tests:
test_layer_norm_residualandtest_layer_norm_residual_store_pre_ln_respass at 96/128 even on master. That is becauseLayerNormOp.layer_norm_residualandlayer_norm_residual_store_pre_ln_resinop_binding/layer_norm.pyare pure PyTorch (F.layer_norm), so those tests never calllaunch_fused_residual_ln*. I checked the two residual launchers directly through the CUDA bindings (inference_module._layer_norm_residualandinference_module.layer_norm_residual_store_pre_ln_res), fp32, shape (32, 128, H). Before each call, the allocator was pre-filled with NaN blocks so that an unwritten output shows up as NaN:launch_fused_residual_lnis also used by the injected MLP path (mlp_after_attninpt_binding.cpp), which is how BERT-tiny reaches it. I did not change the residual tests to call the CUDA bindings. If you want them covered by CI, that would be a small follow-up in this PR or a separate one.The InferenceBuilder op was JIT-built separately from master and from this branch (separate
TORCH_EXTENSIONS_DIR,TORCH_CUDA_ARCH_LIST=12.0,MAX_JOBS=4), and each run logged whichtransformer_inference.soit loaded.Hardware: 1x NVIDIA GeForce RTX 5080 Laptop GPU (sm_120), driver 595.99.02, CUDA 12.8, torch 2.11.0+cu128, Python 3.12.3. Only this GPU was tested.
Compatibility matrix: this only touches the inference (v1, kernel-injection) layer_norm/rms_norm CUDA kernels. Validated: fp32/fp16/bf16 through the extended unit tests, and fp32 end-to-end on BERT-tiny. Not applicable: ZeRO stages, offload, training. Not validated: other GPUs/architectures and ROCm.
Duplicate and sibling-path checks
gh issue view 8826 --comments: no comments.gh pr list --state open --search "8826 in:body": none.layer_norm threadsPerGroup,inference layer norm,rms_norm,layer_norm.cuandthreadsPerGroup: the only related open PR is Return early from inference norm launchers when rows == 0 #8743 ("Return early from inference norm launchers when rows == 0", for Inference layer_norm / rms_norm kill the process (SIGFPE) on an empty batch #8742). It touches the same launchers and test files but fixes a different bug (empty batch) and does not add the 32-thread branch. Whichever lands second may need a trivial rebase.deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/layer_norm_cuda.cuandcuda_rms_norm/rms_norm_cuda.cuhave the same 1/2/4/8/16-only dispatch and list float32 as supported. They look affected too (by code reading; not run). I left them out to keep this PR to the reported v1 path and can add the same branch there if you want.csrc/quantization/quantize.cuhas a similar ladder, but I did not check whether it can reach 32 threads.Notes
pre-commit run --files <changed files>. All hooks pass, including yapf, flake8, codespell, check-license and clang-format (clang-format 23.1.1 was put on PATH, since the hook uses the system binary).AI assistance was used to prepare this change (code, tests and this description). I reviewed every changed line and ran the tests above.