Repository navigation
Fix out-of-bounds row access in inference layer_norm/rms_norm kernels for hidden sizes <= 128 - #8830
jayden1711 wants to merge 2 commits into
Conversation
… for hidden sizes <= 128 Fixes deepspeedai#8828 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: 31b0247e75
ℹ️ 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".
| const int block_offset = (tb.group_index().x * (maxThreads / threadsPerGroup) * elems_per_row) + | ||
| (tb.thread_index().y * elems_per_row); | ||
| // Rows past the end only take part in the reductions; they skip all global loads/stores. | ||
| const int row = tb.group_index().x * (maxThreads / threadsPerGroup) + tb.thread_index().y; |
There was a problem hiding this comment.
Under inference-v2, elems_per_row <= 128 with more than one packed block and a non-multiple row count (for example fp16 H=128 with 17 rows) still corrupts memory: I checked deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/layer_norm_cuda.cu and deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm_cuda.cu, and their identical launchers create 16 groups in the final block but their kernels receive no rows value or guard, so row indices 17–31 load and store past the tensors. This is reachable through public CUDARMSNorm, including EXAONE4's per-head Q/K normalization in model.py:101-130; apply this guard to both sibling implementations and add v2 coverage rather than leaving the same production OOB in place.
AGENTS.md reference: AGENTS.md:L19-L19
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed: the inference-v2 core-op kernels have the same unguarded sub-block launch.
- Kernels:
rms_norm/pre_rms_normindeepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm_cuda.cu, andfused_ln/fused_residual_lnincuda_layer_norm/layer_norm_cuda.cu. - Overrun: for elems_per_row <= 128 they pack several rows per block with no row bound, so when rows > rows-per-block and is not a multiple of it, the last block reads and writes past the end.
- Reachability: through the public
CUDARMSNorm, for example EXAONE4's per-head Q/K RMSNorm (model_implementations/exaone4/model.py). There,channels = head_sizeand Q/K are normalized as[tokens * n_heads, head_size]. - Evidence on master (5f266e2): compute-sanitizer memcheck through
CUDARMSNormreports out-of-bounds accesses for fp16/bf16 hidden 128 at 17 and 33 rows, and for hidden 64 at 33 rows; 16 rows is clean.
This PR is intentionally limited to the v1 / kernel-injection kernels, so I'll keep the v2 kernels out of scope here and open a separate issue for them, with the same row-guard fix offered.
| @pytest.mark.parametrize("seq_len", [1, 17, 33, 128]) | ||
| @pytest.mark.parametrize("channels", [32, 64, 384, 512, 768, 1024, 2048, 8192, 14432]) |
There was a problem hiding this comment.
On a normal pytest run, these added RMS shapes and the analogous base layer-norm shapes do not detect the defect being fixed: the tests only compare values inside the returned logical tensors, while the surplus groups write after those tensors, so the parent implementation passes all of these RMS cases unless it happens to corrupt another allocation (the reported baseline is 360/360). Consequently fused_ln, rms_norm, or pre_rms_norm can lose its valid_row guard while CI remains green; use a canary/guard allocation or a sanitizer-backed regression that reliably fails on the old kernels instead of only broadening the parameter matrix.
AGENTS.md reference: AGENTS.md:L76-L76
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Agreed. The new seq_len/channels cases pass on master under plain pytest; only compute-sanitizer memcheck catches the overrun, and memcheck isn't part of the default CI.
I've added canary tests that fail under plain pytest:
- New tests:
test_layer_norm_no_write_past_last_rowandtest_rms_norm_no_write_past_last_row. - Helper:
run_with_output_canariesininference_test_utils.py. - Coverage: all five bindings (
layer_norm,_layer_norm_residual,layer_norm_residual_store_pre_ln_res,rms_norm,pre_rms_norm), rows 17 and 33, at fp32 hidden 32/64 and fp16/bf16 hidden 64/128. - How it works:
- The outputs are allocated inside the bindings, so the helper uses a private
torch.cuda.MemPoolto place each output directly in front of a caller-owned canary tensor filled with a finite sentinel (1000.0). - It asserts that the layout holds, and then that the canary is unchanged after the call.
- The tests skip when
MemPoolisn't available.
- The outputs are allocated inside the bindings, so the helper uses a private
Results:
- Master (unguarded kernels): 40 failed / 10 passed in each of 20 runs. The 10 passes are rows 17 at the 8-thread sizes (fp32 32, fp16 64), which fit in one 32-row block and don't overrun.
- This branch: 50 passed in each of 20 runs.
- Full files on this branch: test_layer_norm.py 1134 passed, test_rms_norm.py 380 passed.
Adds a MemPool-based canary that places each kernel's output directly in front of a sentinel-filled guard tensor and asserts it stays untouched. Catches every overrunning case deterministically (40 failed/10 passed on master across 20 runs; 50/50 on this branch). Signed-off-by: Jayden Vasquez <jaydenvasquez1711@gmail.com>
This branch has not been deployed
1 participant
Fixes #8828
Summary
For hidden sizes <= 128 the inference norm launchers pack
256 / threadsPerGrouprows into one thread block, but the kernels had no row bound. Whenrowsis larger than that and not a multiple of it, the last block read its inputs and wrote its outputs for rows past the end of the tensors (e.g. fp32 hidden 64 or fp16/bf16 hidden 128 with 17 or 33 rows; see #8828).This PR passes
rowstofused_ln,fused_residual_ln(used by both residual launchers) andrms_norm/pre_rms_norm, computes the row index from the same expression the kernels already use forblock_offset, and masks the global loads and stores of rows past the end. Those threads still take part in the reductions (masked loads return zeros), so the warp/block reductions are unchanged; there is no early return. Launch shapes (grid/block) are unchanged, and the launch macros just passrowsthrough.Why this is not a duplicate
gh pr list --search "8828 in:body"and searches for the launcher names).Tests
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.
transformer_inferenceJIT-built separately for master and this branch (TORCH_CUDA_ARCH_LIST=12.0,MAX_JOBS=4).Test changes (
tests/unit/ops/transformer/inference/):test_layer_norm.py,test_rms_norm.py:seq_len[1, 128]->[1, 17, 33, 128]andchannelsgains32, 64in all five parametrizations, so fp32/fp16/bf16 at 4-, 8- and 16-thread sizes run with partially filled blocks.test_layer_norm.py: newtest_layer_norm_residual_cudacalling the CUDA residual bindings directly (_layer_norm_residual,layer_norm_residual_store_pre_ln_res) for rows 17/33 at fp32 32/64, fp16 64/128 and bf16 128, sinceLayerNormOp.layer_norm_residual*are PyTorch implementations.Commands (from
tests/):pytest -m inference_ops unit/ops/transformer/inference/test_layer_norm.pyand... test_rms_norm.py.5f266e2)test_layer_norm.pytest_rms_norm.pyThe 8 failures on master are
test_layer_norm_residual_cuda[store_pre_ln_res=True]at exactly the partially filled shapes (fp32 32 / rows 33, fp32 64 / rows 17 and 33, fp16 64 / rows 33, fp16 and bf16 128 / rows 17 and 33): the overrun of the normalized output lands in the adjacentres_output. The other new cases (including all oftest_rms_norm.py) pass on master too even though they overrun, because the stray accesses happen to hit unused memory; the memcheck table below is what shows the overrun.memcheck (
PYTORCH_NO_CUDA_MEMORY_CACHING=1 compute-sanitizer --tool memcheck --padding 65536, one process per case, direct calls to the five bindings; number = invalid accesses reported, 0 = clean; rows 9/13/17/33):Counts are for the first faulting launch in each process (the CUDA context is lost afterwards). For
layer_norm,_layer_norm_residualandrms_normthe count is 240/248; forlayer_norm_residual_store_pre_ln_resandpre_rms_normit is 238/246.Speed (CUDA events, 50 warm-up calls, then 200 calls captured in one CUDA graph and replayed 30 times; median per-call GPU time, mean of two passes alternating master/this branch, fresh process per configuration):
The two passes agreed to within 1% for every entry. The differences are within about +/-5% (at most ~0.04 us per call). At 1024 rows layer_norm and rms_norm are 2-4% slower in both passes, at 4096 rows the two are equal, and at 17 rows this branch is 3-5% faster. Called from Python (eager, including the binding and output allocation) every configuration takes 2.5-2.7 us on both branches, so the kernel-time differences are not visible at that level.
Compatibility matrix: only the inference (v1, kernel-injection) layer_norm/rms_norm CUDA kernels change. Validated: fp32/fp16/bf16 via the unit tests above and memcheck, on sm_120. Not applicable: ZeRO stages, offload, training. Not validated: other GPUs/architectures, ROCm, end-to-end model inference.
Formatting:
pre-commit run --files <changed files>(yapf reformatted the new test once; all hooks pass afterwards, clang-format 23.1.1).AI assistance was used to prepare this change (code, tests and this description). I reviewed every changed line and ran the tests above.
(Reminder for the committer: DeepSpeed requires DCO sign-off,
git commit -s.)