[FEAT][kernels] Add fused ratio-clip-aggregate loss primitive - #255
[FEAT][kernels] Add fused ratio-clip-aggregate loss primitive#255Chen-BUPT wants to merge 5 commits into
Conversation
|
Warning Review limit reached
Next review available in: 52 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughAdded a dispatchable ChangesRatio Clip Aggregate
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant GRPOLoss
participant RatioKL
participant RatioClipAggregate
GRPOLoss->>RatioKL: compute ratios and KL terms
GRPOLoss->>RatioClipAggregate: pass ratios, advantages, masks, and penalties
RatioClipAggregate->>GRPOLoss: return total loss and component metrics
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 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: 4
🧹 Nitpick comments (8)
docs/operators/grpo-loss.md (1)
16-18: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a language identifier to the fenced block.
Line 16 triggers markdownlint rule MD040. Use
textfor the pipeline diagram.Proposed fix
-``` +```text logits --[ratio_kl]--> (ratio, kl) --[ratio_clip_aggregate + group adv]--> loss🤖 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 `@docs/operators/grpo-loss.md` around lines 16 - 18, Update the fenced block containing the GRPO loss pipeline diagram to specify the text language identifier, using ```text while preserving the diagram content.Source: Linters/SAST tools
rl_engine/kernels/ops/triton/loss/ratio_clip_aggregate.py (3)
248-251: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the duplicated dtype checks.
_validate_ratio_clip_inputsat line 245 already raisesTypeErrorfor non-floatingratio,advantages, andpenalty_terms. These branches are unreachable.♻️ Proposed cleanup
per_token_advantages = _validate_ratio_clip_inputs( ratio, advantages, mask, penalty_terms, clip_low, clip_high ) - if not ratio.is_floating_point() or not advantages.is_floating_point(): - raise TypeError("ratio and advantages must be floating-point tensors.") - if penalty_terms is not None and not penalty_terms.is_floating_point(): - raise TypeError("penalty_terms must be a floating-point tensor.")🤖 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 `@rl_engine/kernels/ops/triton/loss/ratio_clip_aggregate.py` around lines 248 - 251, Remove the duplicated floating-point dtype validation branches from the surrounding function, keeping the existing _validate_ratio_clip_inputs checks as the sole validation for ratio, advantages, and penalty_terms. Preserve all other input validation and computation behavior.
14-14: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePromote the shared validator to a public helper.
The Triton backend imports
_validate_ratio_clip_inputsacross package boundaries. The underscore prefix marks it private to the PyTorch module, so the shared contract is not expressed in the public surface. Rename it tovalidate_ratio_clip_inputs, or move it to a shared contract module that both backends import.🤖 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 `@rl_engine/kernels/ops/triton/loss/ratio_clip_aggregate.py` at line 14, Promote _validate_ratio_clip_inputs to the public validate_ratio_clip_inputs helper, or relocate it to a shared contract module used by both backends. Update its definition and all imports/call sites, including the Triton ratio-clip aggregation path, so no cross-package import relies on the underscore-prefixed private symbol.
338-349: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winDo not alias the saved
ratiotensor as the gradient output pointer.When
has_penaltyis false,grad_penaltypoints intoratio, and the no-gradient cases passoutputsto the other grad pointers. Thetl.constexprguards avoid touching these pointers today, but any later unconditionaltl.storewould silently overwrite saved tensors.Pass an explicit empty placeholder, such as
ratio.new_empty(0), for unused gradient outputs.🤖 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 `@rl_engine/kernels/ops/triton/loss/ratio_clip_aggregate.py` around lines 338 - 349, The gradient-output setup must not alias saved tensors when gradients are unused. Update the fallback assignments in the backward path around grad_total, grad_policy, and grad_mean_penalty to use explicit empty placeholders such as ratio.new_empty(0), while preserving contiguous provided gradients; ensure grad_penalty no longer relies on ratio as an unused output pointer when has_penalty is false.rl_engine/executors/deepspeed_trainer.py (1)
168-176: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMove the clipping bounds and the penalty coefficient into the config.
clip_low,clip_high, andpenalty_coefare hardcoded at 0.2, 0.2, and 0.01.DeepSpeedTrainingConfigalready carries the training hyperparameters. Expose these three values there so runs can change them without editing the worker.🤖 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 `@rl_engine/executors/deepspeed_trainer.py` around lines 168 - 176, Move the hardcoded clip_low, clip_high, and penalty_coef values from the _ratio_clip_aggregate call into DeepSpeedTrainingConfig, adding corresponding configurable fields with the existing values as defaults. Update the trainer to read these fields from its config while preserving the current loss calculation behavior.rl_engine/kernels/ops/pytorch/loss/grpo_loss.py (2)
119-128: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
_masked_meanandexpand_advantagesare no longer used byapply.The delegated call replaces both helpers. Remove them, or keep them only if other callers or tests still use them.
🤖 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 `@rl_engine/kernels/ops/pytorch/loss/grpo_loss.py` around lines 119 - 128, Remove the now-unused _masked_mean and expand_advantages helpers from the loss implementation unless they have other callers or required tests; verify references before deleting and keep apply’s _ratio_clip_aggregate delegation unchanged.
119-128: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOrphaned reduction helpers in both GRPO implementations. Delegating the clipped surrogate and the masked reductions to
ratio_clip_aggregateleaves_masked_meanandexpand_advantagesuncalled in both GRPO ops.
rl_engine/kernels/ops/pytorch/loss/grpo_loss.py#L119-L128: removeexpand_advantages(lines 88-96) and_masked_mean(lines 165-171), or keep them only for confirmed external callers.rl_engine/kernels/ops/triton/loss/grpo_loss.py#L208-L217: removeexpand_advantages(lines 169-177) and_masked_mean(lines 179-185) in the same way.
expand_advantagesis a public static method, so check for callers before deleting it.🤖 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 `@rl_engine/kernels/ops/pytorch/loss/grpo_loss.py` around lines 119 - 128, Remove the now-unused _masked_mean and expand_advantages helpers from rl_engine/kernels/ops/pytorch/loss/grpo_loss.py (lines 88-96 and 165-171) and rl_engine/kernels/ops/triton/loss/grpo_loss.py (lines 169-177 and 179-185), unless expand_advantages has confirmed external callers; if it does, retain it. The _ratio_clip_aggregate call sites at rl_engine/kernels/ops/pytorch/loss/grpo_loss.py:119-128 and rl_engine/kernels/ops/triton/loss/grpo_loss.py:208-217 require no direct changes.tests/test_ratio_clip_aggregate.py (1)
293-328: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd a
clip_fractionnon-differentiability assertion.The Triton forward calls
ctx.mark_non_differentiable(clip_fraction). No test asserts that contract, so a regression that makes the metric differentiable would pass. Add a check thatclip_fraction.requires_gradis false for both backends.🤖 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 `@tests/test_ratio_clip_aggregate.py` around lines 293 - 328, Update test_triton_matches_native_at_single_and_staged_reduction_boundary to assert that the clip_fraction output has requires_grad set to false for both NativeRatioClipAggregateOp and TritonRatioClipAggregateOp. Identify the clip_fraction result by its corresponding output position or named result, and retain the existing numerical comparisons and determinism test unchanged.
🤖 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 `@docs/operators/ratio-clip-aggregate.md`:
- Around line 105-108: Align the benchmark command near the documented results
with the stated configuration by explicitly passing 100 warmups and 500
iterations, or revise the surrounding documentation to match the command’s
defaults. Keep the hardware, software, and input-condition details unchanged.
In `@rl_engine/kernels/ops/triton/loss/ratio_clip_aggregate.py`:
- Around line 317-332: Update the output unpacking in the custom autograd
forward method so the first three returned scalars—total, policy, and
mean_penalty—are independent clones rather than storage-sharing views of
outputs; keep clip_fraction unchanged and preserve the existing return order and
non-differentiable handling.
- Around line 265-282: Reduce the single-pass threshold used by
`_SINGLE_PASS_MAX` to a smaller safe tile size such as 4096 or 8192, so
`_ratio_clip_single_pass_kernel` does not launch with excessively large `BLOCK`
values. Ensure inputs above the new threshold continue through the existing
staged path while preserving deterministic coverage.
In `@tests/test_ratio_clip_aggregate.py`:
- Around line 8-20: Make the TritonRatioClipAggregateOp import conditional on
_HAS_TRITON so test collection succeeds without Triton, assigning it a safe
absent value otherwise. Update test_registry_dispatches_ratio_clip_aggregate to
require TritonRatioClipAggregateOp is not None in the Triton assertion branch,
while preserving native-test execution when Triton is unavailable.
---
Nitpick comments:
In `@docs/operators/grpo-loss.md`:
- Around line 16-18: Update the fenced block containing the GRPO loss pipeline
diagram to specify the text language identifier, using ```text while preserving
the diagram content.
In `@rl_engine/executors/deepspeed_trainer.py`:
- Around line 168-176: Move the hardcoded clip_low, clip_high, and penalty_coef
values from the _ratio_clip_aggregate call into DeepSpeedTrainingConfig, adding
corresponding configurable fields with the existing values as defaults. Update
the trainer to read these fields from its config while preserving the current
loss calculation behavior.
In `@rl_engine/kernels/ops/pytorch/loss/grpo_loss.py`:
- Around line 119-128: Remove the now-unused _masked_mean and expand_advantages
helpers from the loss implementation unless they have other callers or required
tests; verify references before deleting and keep apply’s _ratio_clip_aggregate
delegation unchanged.
- Around line 119-128: Remove the now-unused _masked_mean and expand_advantages
helpers from rl_engine/kernels/ops/pytorch/loss/grpo_loss.py (lines 88-96 and
165-171) and rl_engine/kernels/ops/triton/loss/grpo_loss.py (lines 169-177 and
179-185), unless expand_advantages has confirmed external callers; if it does,
retain it. The _ratio_clip_aggregate call sites at
rl_engine/kernels/ops/pytorch/loss/grpo_loss.py:119-128 and
rl_engine/kernels/ops/triton/loss/grpo_loss.py:208-217 require no direct
changes.
In `@rl_engine/kernels/ops/triton/loss/ratio_clip_aggregate.py`:
- Around line 248-251: Remove the duplicated floating-point dtype validation
branches from the surrounding function, keeping the existing
_validate_ratio_clip_inputs checks as the sole validation for ratio, advantages,
and penalty_terms. Preserve all other input validation and computation behavior.
- Line 14: Promote _validate_ratio_clip_inputs to the public
validate_ratio_clip_inputs helper, or relocate it to a shared contract module
used by both backends. Update its definition and all imports/call sites,
including the Triton ratio-clip aggregation path, so no cross-package import
relies on the underscore-prefixed private symbol.
- Around line 338-349: The gradient-output setup must not alias saved tensors
when gradients are unused. Update the fallback assignments in the backward path
around grad_total, grad_policy, and grad_mean_penalty to use explicit empty
placeholders such as ratio.new_empty(0), while preserving contiguous provided
gradients; ensure grad_penalty no longer relies on ratio as an unused output
pointer when has_penalty is false.
In `@tests/test_ratio_clip_aggregate.py`:
- Around line 293-328: Update
test_triton_matches_native_at_single_and_staged_reduction_boundary to assert
that the clip_fraction output has requires_grad set to false for both
NativeRatioClipAggregateOp and TritonRatioClipAggregateOp. Identify the
clip_fraction result by its corresponding output position or named result, and
retain the existing numerical comparisons and determinism test unchanged.
🪄 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: f8fcc558-80ad-45b5-8335-1e4c5898d835
📒 Files selected for processing (12)
benchmarks/benchmark_ratio_clip_aggregate.pydocs/.nav.ymldocs/operators/README.mddocs/operators/grpo-loss.mddocs/operators/ratio-clip-aggregate.mdrl_engine/executors/deepspeed_trainer.pyrl_engine/kernels/ops/pytorch/loss/grpo_loss.pyrl_engine/kernels/ops/pytorch/loss/ratio_clip_aggregate.pyrl_engine/kernels/ops/triton/loss/grpo_loss.pyrl_engine/kernels/ops/triton/loss/ratio_clip_aggregate.pyrl_engine/kernels/registry.pytests/test_ratio_clip_aggregate.py
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@tests/test_ratio_clip_aggregate.py`:
- Around line 10-16: Update the Triton import guard around
TritonRatioClipAggregateOp to catch only ModuleNotFoundError for missing
top-level modules named "triton" or "triton.language"; set the fallback state
only for those cases, and re-raise all other import failures so broken rl_engine
Triton implementation imports fail test collection.
🪄 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: 9113f1de-511d-4da2-85b2-24ceef25961a
📒 Files selected for processing (5)
docs/operators/grpo-loss.mddocs/operators/ratio-clip-aggregate.mdrl_engine/kernels/ops/pytorch/loss/ratio_clip_aggregate.pyrl_engine/kernels/ops/triton/loss/ratio_clip_aggregate.pytests/test_ratio_clip_aggregate.py
🚧 Files skipped from review as they are similar to previous changes (4)
- docs/operators/ratio-clip-aggregate.md
- rl_engine/kernels/ops/pytorch/loss/ratio_clip_aggregate.py
- rl_engine/kernels/ops/triton/loss/ratio_clip_aggregate.py
- docs/operators/grpo-loss.md
2 participants
Closes #252.
What changed
ratio_clip_aggregatebackends for PPO/GRPO clipped-surrogate loss.[B, T]tensor.atomicAddor CPU synchronization.Performance
Measured on NVIDIA RTX PRO 5000 Blackwell, PyTorch 2.11, CUDA 12.9, Triton 3.6, FP32 inputs, 90% mask density, 100 warmups, and 500 iterations.
At 4.19M tokens, measured forward intermediates decrease from 64.0 MiB to 0.3 MiB.
Validation
pytest -q tests/test_ratio_clip_aggregate.py tests/test_grpo_loss.py tests/test_deepspeed_training_worker.py— 58 passedpytest -q tests/test_kernel_registry.py tests/test_rl_kernel_loss_step.py tests/test_ratio_kl.py tests/test_grpo_single_gpu_example.py— 29 passed, 1 skippedruff checkandruff format --checkon every changed Python filemkdocs build --strict -f mkdocs.yamlpython benchmarks/benchmark_ratio_clip_aggregate.py --iterations 500 --warmup 100The disposable full-suite container did not include the repository's deterministic-attention CUDA extension, so its extension-dependent tests were not a valid signal; the affected operator and integration suites above are green.
Summary by CodeRabbit
New Features
Documentation
Tests
Benchmarks