[WS2] feat: add TP=1 logprob comparison harness - #262
Conversation
|
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:
📝 WalkthroughWalkthroughAdded ChangesSingle-GPU logprob comparison
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant ComparisonHarness
participant ReferenceBackend
participant CandidateBackend
participant JSONReport
CLI->>ComparisonHarness: create seeded inputs
ComparisonHarness->>ReferenceBackend: compute logprob and LSE
ComparisonHarness->>CandidateBackend: execute selected backend
CandidateBackend-->>ComparisonHarness: return logprob and LSE
ComparisonHarness->>JSONReport: calculate drift and provenance
JSONReport-->>CLI: emit serialized JSON
Possibly related issues
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: 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
KJLdefeated
left a comment
KJLdefeated
left a comment
There was a problem hiding this comment.
Overall looks good to me. After addressed the request, I am happy to approve.
There was a problem hiding this comment.
To let future maintenance easier, I think it is not necessary to create two docs for this. I suggest that you can compress the context of logprob harness and put into existing docs/operators/batch-invariant-logp.md.
There was a problem hiding this comment.
Same as this one, compress the context and put to docs/operators/batch-invariant-logp.md.
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rl_engine/testing/logprob_comparison.py (1)
116-124: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winPreserve backend execution failures as failures.
LogprobBackendUnavailableis used by callers to skip unavailable backends. The wrapper convertsRuntimeError,NotImplementedError, andOSErrorfromforward_with_lse, so kernel bugs, unsupported inputs, CUDA/OOM errors, and similar execution failures can be treated as skip conditions. Introduce a dedicated unsupported-capability/missing-feature exception and let execution failures propagate.🤖 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 116 - 124, Update run so LogprobBackendUnavailable is raised only for a dedicated unsupported-capability or missing-feature exception from the exact backend. Remove the conversion of RuntimeError, NotImplementedError, and OSError; allow those execution failures from diagnostic to propagate unchanged.
🧹 Nitpick comments (1)
rl_engine/testing/logprob_comparison.py (1)
304-309: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winRedirect only the intended logging handlers.
logging.FileHandleris a subclass oflogging.StreamHandler, so this condition can redirect file logging tosys.stderr. The loop also does not affect handlers reached through logger propagation. Inspectrl_engine.utils.logger.loggerand target the intended stream handler explicitly.🤖 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 304 - 309, Update _route_rl_kernel_logs_to_stderr to target only the intended console/stream handler, excluding logging.FileHandler, and preserve file-handler streams. Inspect the configured handlers on logger, including propagated handlers if the logger configuration relies on propagation, and redirect only the explicitly identified stream handler to sys.stderr.
🤖 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_logprob_comparison.py`:
- Around line 201-205: Update test_cli_runs_directly_from_testing_module to
derive the logprob_comparison.py subprocess path from Path(__file__).resolve(),
ensuring the script is found regardless of pytest’s working directory while
preserving the existing subprocess validation.
---
Outside diff comments:
In `@rl_engine/testing/logprob_comparison.py`:
- Around line 116-124: Update run so LogprobBackendUnavailable is raised only
for a dedicated unsupported-capability or missing-feature exception from the
exact backend. Remove the conversion of RuntimeError, NotImplementedError, and
OSError; allow those execution failures from diagnostic to propagate unchanged.
---
Nitpick comments:
In `@rl_engine/testing/logprob_comparison.py`:
- Around line 304-309: Update _route_rl_kernel_logs_to_stderr to target only the
intended console/stream handler, excluding logging.FileHandler, and preserve
file-handler streams. Inspect the configured handlers on logger, including
propagated handlers if the logger configuration relies on propagation, and
redirect only the explicitly identified stream handler to sys.stderr.
🪄 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: 7ee39833-edf9-4862-8e0e-ba514ba135bd
📒 Files selected for processing (3)
docs/operators/batch-invariant-logp.mdrl_engine/testing/logprob_comparison.pytests/test_logprob_comparison.py
@KJLdefeated Addressed, thanks for the suggestions. I consolidated both harness documents into docs/operators/batch-invariant-logp.md and moved the kernel-specific CLI into rl_engine/testing/logprob_comparison.py. |
2 participants
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.Kernel-local command-line entry
Make
rl_engine/testing/logprob_comparison.pydirectly executable for reproducible localand GPU comparisons. This keeps the kernel-specific diagnostic next to its harness instead
of adding a project-wide script.
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.
Operator documentation
Document the TP=1 harness contract, exact backend selection, CLI usage, and SM90 validation
in
docs/operators/batch-invariant-logp.md. The content is consolidated into the existingoperator documentation rather than maintained as separate WS2 design documents.
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
Current focused and complete operator validation:
The skipped cases require a compiled CUDA SM90 extension on Hopper hardware. The focused
suite includes a subprocess test of the direct module CLI and machine-readable JSON output.
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