[WS1] Close embedding and LM-head invariance coverage - #243
Conversation
Signed-off-by: inaniloquentee <3051000145@qq.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughThe PR adds fixed-K row-wise LM-head projection, Hopper SM90 embedding and LM-head kernels, explicit-device dispatch, operator harness registrations, and tests for batch-layout, padding, permutation, gradient, and LM-head-to-log-probability invariance. ChangesBatch-invariant forward path
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant DeepSpeedTrainingWorker
participant KernelRegistry
participant SM90LMHeadOp
participant SM90Extension
participant NativeLMHeadOp
DeepSpeedTrainingWorker->>KernelRegistry: resolve linear_logp for worker device
KernelRegistry->>SM90LMHeadOp: select Hopper backend when available
SM90LMHeadOp->>SM90Extension: compute LM-head forward
SM90LMHeadOp-->>NativeLMHeadOp: fallback for ineligible inputs
Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rl_engine/kernels/ops/pytorch/linear/lm_head.py`:
- Around line 118-119: Update the empty-input branch in the LM-head
implementation around flat_hidden to return an empty matmul result connected to
hidden and weight instead of hidden.new_empty, preserving autograd for zero-row
inputs. Also add the equivalent zero-row branch to the reference snippets in
tests/test_lm_head.py lines 94-100 and docs/operators/lm_head.md lines 57-61 so
they avoid torch.stack([]); the implementation site is
rl_engine/kernels/ops/pytorch/linear/lm_head.py lines 118-119.
In `@tests/test_issue151_embedding_lm_head_invariance.py`:
- Around line 240-271: Update
test_lm_head_linear_logp_handoff_is_layout_invariant_for_rl_batches to reuse the
existing CUDA/bfloat16 parameterization, covering both CPU/float32 and
CUDA/bfloat16 execution. Before calling NativeLinearLogpOp, alter the hidden
rows selected by ~mask while leaving valid rows unchanged, then retain the
reference and layout-invariance assertions so padded hidden values cannot affect
results.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: daceb832-b7f6-4a13-8c93-19c721132703
📒 Files selected for processing (12)
docs/operators/lm_head.mdrl_engine/executors/deepspeed_trainer.pyrl_engine/kernels/gtest/operator_inputs.pyrl_engine/kernels/gtest/operator_specs.pyrl_engine/kernels/ops/pytorch/linear/lm_head.pyrl_engine/kernels/registry.pyrl_engine/tests/test_dispatch.pytests/test_deepspeed_training_worker.pytests/test_issue151_embedding_lm_head_invariance.pytests/test_lm_head.pytests/test_op_checks.pytests/test_operator_inputs.py
Signed-off-by: inaniloquentee <3051000145@qq.com>
|
please resolve conflicts first, thanks! |
|
@inaniloquentee I only see the PyTorch implementation; I don’t see the corresponding CUDA and Triton implementations. If they haven't been implemented yet, could you please add them? Thank you. |
Signed-off-by: inaniloquentee <3051000145@qq.com> # Conflicts: # rl_engine/kernels/gtest/operator_specs.py # rl_engine/kernels/registry.py
Add single-card Hopper/H200 CUDA forward paths for embedding gather and LM-head projection. The LM-head kernel assigns one CTA per output logit and reduces the full hidden dimension without Split-K, preserving batch-invariant reduction order. Signed-off-by: inaniloquentee <3051000145@qq.com>
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
csrc/cuda/embedding_lm_head_sm90.cu (1)
82-112: 🧹 Nitpick | 🔵 TrivialConfirm the no-reuse, one-CTA-per-logit LM-head design is an accepted throughput tradeoff.
Each block independently re-reads the full
hidden_row/weight_rowfrom global memory with zero data reuse across thevocab_sizeblocks that share the same token, so total memory traffic scales asnum_tokens * vocab_size * hidden_size. This matches the PR's stated intent ("no Split-K... to preserve batch-invariant reduction order"), so flagging for awareness only — for large vocabularies this reference kernel will be significantly memory-bandwidth bound relative to a tiled GEMM.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@csrc/cuda/embedding_lm_head_sm90.cu` around lines 82 - 112, Confirm that the one-CTA-per-logit implementation in lm_head_sm90_forward_kernel is an intentional throughput tradeoff: preserve its independent full hidden_row/weight_row reads and deterministic reduction order, and document or acknowledge that this reference design may be memory-bandwidth bound for large vocabularies rather than refactoring it for cross-logit data reuse.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@csrc/cuda/embedding_lm_head_sm90.cu`:
- Around line 166-207: Resolve the conflicting out-of-range ID handling between
the host path around embedding_sm90_forward and the kernel’s bounds check:
either remove the synchronizing min/max validation and preserve/document the
kernel’s zero-for-invalid-ID behavior, or retain host-side validation and remove
the kernel masking branch. Ensure only one policy remains and eliminate the
unreachable or unnecessary path.
In `@rl_engine/kernels/ops/cuda/linear/embedding.py`:
- Around line 34-52: Update backward in the custom embedding autograd
implementation to replace grad_weight.index_add_ with a deterministic
accumulation path for duplicate token_ids on CUDA, such as sorting ids followed
by fixed-order segmented reduction or a dedicated deterministic kernel. Preserve
filtering of invalid ids, output shape/device/dtype, and the existing
no-gradient path while ensuring torch.use_deterministic_algorithms(True)
succeeds and matches PyTorch reference gradients.
In `@rl_engine/kernels/ops/cuda/linear/lm_head.py`:
- Around line 44-61: Update SM90LMHeadOp.backward to compute grad_hidden and
grad_weight through the deterministic det_gemm_da and det_gemm_db wrappers
instead of torch.matmul, preserving the existing reshaping, float conversion,
dtype conversion, and ctx.needs_input_grad guards. Use the wrappers with the
equivalent dC @ B^T and A^T @ dC operands, while leaving grad_bias and the
returned tuple unchanged.
---
Nitpick comments:
In `@csrc/cuda/embedding_lm_head_sm90.cu`:
- Around line 82-112: Confirm that the one-CTA-per-logit implementation in
lm_head_sm90_forward_kernel is an intentional throughput tradeoff: preserve its
independent full hidden_row/weight_row reads and deterministic reduction order,
and document or acknowledge that this reference design may be memory-bandwidth
bound for large vocabularies rather than refactoring it for cross-logit data
reuse.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: e244f52c-c9ab-4bce-8b8f-a39c6622e67f
📒 Files selected for processing (17)
csrc/cuda/embedding_lm_head_sm90.cucsrc/ops.cppdocs/operators/embedding.mddocs/operators/lm_head.mdrl_engine/_C.pyirl_engine/kernels/gtest/operator_inputs.pyrl_engine/kernels/gtest/operator_specs.pyrl_engine/kernels/ops/cuda/linear/__init__.pyrl_engine/kernels/ops/cuda/linear/embedding.pyrl_engine/kernels/ops/cuda/linear/lm_head.pyrl_engine/kernels/registry.pysetup.pytests/test_embedding.pytests/test_kernel_registry.pytests/test_lm_head.pytests/test_operator_inputs.pytests/test_sm90_linear_wrappers.py
🚧 Files skipped from review as they are similar to previous changes (6)
- tests/test_operator_inputs.py
- rl_engine/kernels/gtest/operator_specs.py
- rl_engine/kernels/gtest/operator_inputs.py
- tests/test_lm_head.py
- docs/operators/lm_head.md
- rl_engine/kernels/registry.py
|
cc @maxiaosong1124 @KJLdefeated @frank-2077 PTAL, we need merge this pr tonight. Thank you. |
Signed-off-by: inaniloquentee <3051000145@qq.com>
KJLdefeated
left a comment
KJLdefeated
left a comment
There was a problem hiding this comment.
Overall look good and clean. Left some comment. After addressed, I am happy to approve.
| "embedding_sm90 supports fp32, fp16, and bf16 weights"); | ||
|
|
||
| c10::cuda::CUDAGuard device_guard(weight.device()); | ||
| check_sm90_device(); |
There was a problem hiding this comment.
check_sm90_device() calls cudaGetDeviceProperties on every forward. The Python _can_use_sm90 already gates on _is_hopper, then the C++ re-checks prop.major == 9 and hard-TORCH_CHECKs. Maybe this is not necessary?
|
|
||
| namespace { | ||
|
|
||
| constexpr int kThreads = 256; |
There was a problem hiding this comment.
Non-blocking. One CTA per logit and kThreads = 256 would be slow for Qwen3 prefill (say 4096 tokens × 151936 vocab) that's ~622M blocks, each reducing 4096 elements with 256 threads (16 elements/thread). But this is fine for reference kernel.
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
csrc/cuda/embedding_lm_head_sm90.cu (1)
134-141: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueCache the SM90 capability probe per device.
check_sm90_devicecallscudaGetDevicePropertieson every embedding/LM-head forward. The fullcudaDevicePropfetch is heavier than needed on a per-step hot path;cudaDeviceGetAttribute(cudaDevAttrComputeCapabilityMajor, ...)or a static per-device cache keeps the check but removes the repeated cost.♻️ Proposed refactor
void check_sm90_device() { int device = 0; C10_CUDA_CHECK(cudaGetDevice(&device)); - cudaDeviceProp prop{}; - C10_CUDA_CHECK(cudaGetDeviceProperties(&prop, device)); - TORCH_CHECK(prop.major == 9, "SM90 embedding/lm_head kernels require Hopper, got sm_", - prop.major, prop.minor); + static std::array<int, 2> cached_cc[8] = {}; // or a small map keyed by device + int major = 0; + int minor = 0; + C10_CUDA_CHECK(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device)); + C10_CUDA_CHECK(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device)); + TORCH_CHECK(major == 9, "SM90 embedding/lm_head kernels require Hopper, got sm_", + major, minor); }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@csrc/cuda/embedding_lm_head_sm90.cu` around lines 134 - 141, Update check_sm90_device to avoid calling cudaGetDeviceProperties on every invocation: use cudaDeviceGetAttribute for the compute capability or cache the per-device SM90 result, while preserving the existing device lookup, CUDA error checks, and Hopper-only TORCH_CHECK validation.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rl_engine/kernels/ops/cuda/linear/embedding.py`:
- Around line 50-64: Replace the per-row accumulation loop in the embedding
backward helper with a padded [num_unique, max_count, H] segment buffer and fold
each segment using a vectorized reduction, preserving left-to-right accumulation
order while reducing launches to max(counts). Reuse the existing sorting/segment
structure and handle invalid IDs through the discarded sorted segment rather
than a separate validity synchronization; avoid introducing additional device
synchronizations beyond the existing counts.max().item() path.
In `@tests/test_sm90_linear_wrappers.py`:
- Around line 198-206: Update expected_weight in the relevant gradient test to
mirror det_gemm_db’s hidden^T @ dy orientation, then transpose the result to the
expected weight shape before comparison. Keep the existing dtype conversions and
torch.equal assertion unchanged.
---
Nitpick comments:
In `@csrc/cuda/embedding_lm_head_sm90.cu`:
- Around line 134-141: Update check_sm90_device to avoid calling
cudaGetDeviceProperties on every invocation: use cudaDeviceGetAttribute for the
compute capability or cache the per-device SM90 result, while preserving the
existing device lookup, CUDA error checks, and Hopper-only TORCH_CHECK
validation.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 99987810-767d-4463-8203-4fe39da02fc1
📒 Files selected for processing (6)
csrc/cuda/embedding_lm_head_sm90.cudocs/operators/embedding.mddocs/operators/lm_head.mdrl_engine/kernels/ops/cuda/linear/embedding.pyrl_engine/kernels/ops/cuda/linear/lm_head.pytests/test_sm90_linear_wrappers.py
🚧 Files skipped from review as they are similar to previous changes (1)
- rl_engine/kernels/ops/cuda/linear/lm_head.py
Signed-off-by: inaniloquentee <3051000145@qq.com>
maxiaosong1124
left a comment
maxiaosong1124
left a comment
There was a problem hiding this comment.
LGTM now! Thanks!
Thanks! Addressed the review, cached the SM90 device capability check, reduced embedding backward launches with a fixed-order segmented fold, and fixed the LM-head backward test expectation. CI is green now, including H100 SM90. |
Flink-ddd
left a comment
Flink-ddd
left a comment
There was a problem hiding this comment.
LGTM, Thank you for work.
4 participants
Summary
Relationship to existing PRs
Closes #151
Validation
Summary by CodeRabbit
forward_fp32variants.embeddingandlm_headwith device-aware backend selection.