[WS2] feat: add TP=1 logprob comparison harness - #262
Conversation
|
Warning Review limit reached
Next review available in: 29 minutes You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. 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 single-GPU logprob/LSE diagnostics for PyTorch, Triton, and SM90 CUDA backends. Added comparison APIs, a JSON CLI harness, validation tests, and maintainer documentation. ChangesSingle-GPU logprob comparison
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant ComparisonAPI
participant PyTorchReference
participant CandidateBackend
participant JSONReport
CLI->>ComparisonAPI: submit seeded inputs and selected backends
ComparisonAPI->>PyTorchReference: compute reference logprob and LSE
ComparisonAPI->>CandidateBackend: execute selected diagnostic backend
CandidateBackend-->>ComparisonAPI: return logprob and LSE tensors
ComparisonAPI->>JSONReport: calculate drift and provenance
JSONReport-->>CLI: return serialized comparison report
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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
🧹 Nitpick comments (1)
rl_engine/testing/logprob_comparison.py (1)
132-146: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winDisable autograd during diagnostic execution.
If
inputs.logits.requires_gradis true, the reference creates two autograd graphs, and a PyTorch candidate can create another graph. The laterdetach()calls occur after full-vocabulary intermediates are retained. Run diagnostic calls undertorch.no_grad().Proposed change
- reference_logp, reference_lse = _run_ws1_reference( - inputs.logits, effective_targets, inputs.ignore_index - ) + with torch.no_grad(): + reference_logp, reference_lse = _run_ws1_reference( + inputs.logits, effective_targets, inputs.ignore_index + ) @@ - logp, lse = _run_candidate( - candidate, - inputs.logits, - effective_targets, - inputs.ignore_index, - ) + with torch.no_grad(): + logp, lse = _run_candidate( + candidate, + inputs.logits, + effective_targets, + inputs.ignore_index, + )🤖 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/testing/logprob_comparison.py` around lines 132 - 146, The diagnostic execution creates autograd graphs during _run_ws1_reference and _run_candidate calls which retain full-vocabulary intermediates in memory, even though detach() is applied later. Wrap the _validate_inputs call, the _run_ws1_reference invocation, and the candidate iteration loop (containing the _run_candidate calls) in a torch.no_grad() context manager to disable autograd tracking entirely during these diagnostic operations.
🤖 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/testing/logprob_comparison.py`:
- Around line 178-185: Run Black and isort formatters on the specified Python
files to match repository formatting standards. Apply isort to the imports and
Black to the code formatting at the following locations:
rl_engine/testing/logprob_comparison.py lines 178-185 (apply both isort to the
NativeBatchInvariantLogpOp import and Black to the op() and forward_with_lse()
call formatting), scripts/compare_logprob.py lines 19-22 (apply isort to the
package import), tests/test_logprob_comparison.py lines 15-17 (apply both isort
and Black to the import), and tests/test_logprob_comparison.py lines 197-199
(apply Black to the function call formatting).
- Around line 216-224: Update _candidate_provenance so candidate.provenance is
merged before the canonical requested_backend, actual_backend, tp_world,
communication, and lse_source fields. Ensure these canonical fields remain
authoritative and cannot be overwritten in the serialized report.
---
Nitpick comments:
In `@rl_engine/testing/logprob_comparison.py`:
- Around line 132-146: The diagnostic execution creates autograd graphs during
_run_ws1_reference and _run_candidate calls which retain full-vocabulary
intermediates in memory, even though detach() is applied later. Wrap the
_validate_inputs call, the _run_ws1_reference invocation, and the candidate
iteration loop (containing the _run_candidate calls) in a torch.no_grad()
context manager to disable autograd tracking entirely during these diagnostic
operations.
🪄 Autofix
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: fae2870a-cf06-4b8c-9dbb-30a35b083445
📒 Files selected for processing (9)
docs/design/ws2-logprob-single-gpu-harness.mddocs/design/ws2-logprob-sm90-validation.mdrl_engine/kernels/ops/cuda/loss/batch_invariant_logp.pyrl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.pyrl_engine/kernels/ops/triton/loss/batch_invariant_logp.pyrl_engine/testing/__init__.pyrl_engine/testing/logprob_comparison.pyscripts/compare_logprob.pytests/test_logprob_comparison.py
1 participant
Summary
Add the TP=1 logprob comparison harness requested by PR2 of #241.
The harness registers the existing WS1 batch-invariant PyTorch logprob path as the reference and compares the supported single-GPU backends against it before tensor-parallel communication is introduced. It reports direct vocabulary-LSE drift and active-token-only selected-logprob drift, while recording enough backend provenance to detect accidental fallback.
The existing production operator contract remains unchanged. The new
forward_with_lsemethods are diagnostic entry points used by the comparison harness and tests.Implements PR2 of #241.
Scope
This PR covers the single-GPU registration and regression guard described in PR2:
tp_world=1andcommunication=nonein the report.This PR does not implement vocab sharding, collective communication, fixed-order cross-rank LSE merging, CP reconstruction, or distributed artifact generation. Those remain part of the later PRs in #241.
Changes
Single-GPU comparison harness
Add
rl_engine/testing/logprob_comparison.pywith:pytorch,triton, andcuda-sm90.ignore_indexusage.batch_invariant_logp.Diagnostic LSE entry points
Add a diagnostic-only method to each supported backend:
The normal production call remains:
The diagnostic path exposes the LSE computed by the backend itself. The harness does not reconstruct LSE from selected logprobs, which keeps the LSE comparison independent and useful for later TP work.
For an explicit
cuda-sm90request, the diagnostic path requires the compiled SM90 extension and compatible Hopper inputs. It does not use the production operator's fallback behavior.Command-line comparison tool
Add
scripts/compare_logprob.pyfor reproducible local and GPU comparisons.Example:
The command writes a structured JSON report to stdout. RL-Kernel diagnostic logs are routed to stderr so redirected stdout remains valid machine-readable JSON.
Comparison contract
For each logical token row, the compared values are:
LSE drift is measured over every logical token row. Selected-logprob drift is measured only where the active-token mask is true. Each drift report contains:
The report also records:
Tests
Add focused coverage for:
ignore_indexusage.Validation
Windows CPU
Result:
The skipped cases require CUDA/Triton backends.
WSL Triton
Focused Triton validation:
The skipped case requires a compiled CUDA SM90 extension.
NVIDIA H800 / SM90
Validated on:
The editable CUDA extension built successfully with the SM90 kernel enabled:
Test results:
Observed BF16 SM90 drift against the PyTorch reference:
[2, 8, 1024]4.76837158203125e-074.76837158203125e-07[2, 16, 151936]9.5367431640625e-079.5367431640625e-07Both comparisons used
tp_world=1,communication=none, and the requestedcuda-sm90implementation without fallback.Additional checks:
Notes for review
forward_with_lseexists to expose backend-native diagnostics without changing production callers.Summary by CodeRabbit
New Features
Documentation
Tests