[WS2][PR4][Logp] feat: add distributed TP/CP logprob drift report - #289
Open
hihaluemen wants to merge 28 commits into
Open
[WS2][PR4][Logp] feat: add distributed TP/CP logprob drift report#289hihaluemen wants to merge 28 commits into
hihaluemen wants to merge 28 commits into
Conversation
Implements PR 1 of issue RL-Align#241: a typed contract for vocab-parallel selected-token logprob, mirroring the WS2 attention contract pattern. - rl_engine/kernels/logprob_contract.py: LogprobContract, ShardingSpec (per-rank vocab shard bounds, padded-vs-real vocab, TP/CP rank metadata, owner_rank resolution), MaskSpec (active-token mask, ignore_index), ReductionSpec (fp32 (max, sumexp) merge in fixed global vocab-shard index order, all-gather transport, CP declared a non-merge axis), and LogprobBackendCapability. - KernelRegistry.get_logprob_op(contract): contract-aware dispatch that only selects backends with a declared capability; incompatible or undeclared candidates are rejected with explicit reasons and never used as a silent fallback. Existing WS1 batch-invariant logp backends are declared truthfully as single-shard references, so strict WS2 requests fail loudly until the deterministic vocab-parallel TP reference (PR 3) lands. Legacy get_op() behavior is unchanged. - Design doc, runtime-dispatch and operator doc updates, and CPU-safe contract/dispatch tests covering the Qwen3-8B TP=2 BF16 target and the TP=1/2/4 sweep shapes. Tolerance values remain owned by RL-Align#108.
- docs: correct the TP-invariance claim — fixed merge order gives determinism per TP degree; cross-degree bitwise equality additionally requires a TP-degree-independent local tile decomposition (PR 3 obligation), otherwise RL-Align#108 tolerances apply - contract: store backend_id stripped so id-based dispatch matches; summarize the active mask in to_dict() provenance instead of copying every per-token boolean; sort __all__ per RUF022 - registry: add public register_logprob_backend() seam for PR 3 and tests; delegate _platform() to _platform_for_device(None); reuse _get_or_create_backend() in get_op so WS2 and legacy dispatch share one cache/blacklist code path - tests: use the registration seam instead of poking private state, pin _even_bounds' last bound for non-divisible vocabularies, assert candidate-list decoupling in both directions, cover registration replace semantics and backend_id normalization
- docs: state that cross-TP bitwise equality needs a global tile-level merge structure independent of TP partitioning (per-shard tiles alone leave different grouping at shard boundaries), and that padded columns are masked to -inf before the local (max, sumexp) partials - registry: scope logprob capabilities per platform so the same backend enum can declare different support on cuda/rocm/cpu; validate the platform argument of register_logprob_backend against known platforms - contract: derive IMPLEMENTATION_KINDS from RESERVED_DISPATCH_POLICIES and use it for the kind check; wrap non-iterable roles/dtypes in LogprobContractError for consistent error handling - tests: cover per-platform capability scoping, unknown-platform rejection, and non-iterable roles/dtypes
…typed contract Address external review: the cross-TP bitwise guarantee lived only in prose, so a fixed-topology-deterministic backend could pass dispatch as fully conformant. - DeterminismScope (fixed_topology | cross_tp_bitwise): requested via ReductionSpec (default cross_tp_bitwise, the RL-Align#241 PR 3 target), declared per backend via determinism_scopes, enforced by dispatch; replaces the deterministic_tp_merge bool - MaskMode (explicit_active_mask | ignore_index) replaces supports_inactive_tokens: the contract permits inactive targets that do not hold ignore_index, so ignore-index-only backends are rejected for contracts with inactive tokens - LogprobOutputSpec pins the output surface: fp32 selected logprob and fp32 vocab LSE, replicated across the TP group - implementation_kind is now a tier (reference | production); determinism is no longer conflated with it, and requesting "deterministic" as a policy raises a loud error pointing at determinism_scope - fallback provenance: policy evaluation now precedes capability checks, so a candidate excluded by the caller's own policy never counts as a fallback even when it also lacks capabilities - docs: define the (-inf, 0) identity partial for padding-only or all--inf shards; document that requested_backend="auto" is not distributed-safe and specify the preflight fingerprint agreement - LogprobContract.cross_rank_fingerprint(): rank-independent identity for that preflight; provenance now records active_mask_sha256 so masks with equal active counts remain distinguishable
Fold the normative reduction semantics (padded-column masking, fp32 (max, sumexp) merge formulas, the (-inf, 0) identity partial, and the cross-TP tile-structure requirement) into the ReductionSpec and DeterminismScope docstrings, and repoint the runtime-dispatch and batch-invariant-logp doc references at the module. The contract summary moves to the PR description.
Shrink class docstrings toward the attention-contract one-liner style and cut design-rationale comments; the normative reduction semantics stay in the ReductionSpec and DeterminismScope docstrings.
The guard added per review rejects requested_backend="auto" whenever tp_world_size > 1, so TP-sharded dispatch tests now name an explicit policy and auto-policy tests use TP=1 contracts. Add coverage for the guard itself and document the restriction in get_logprob_op. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012PyjQEqDJwy9Cos4Sb9QBK
hihaluemen
requested review from
EthanZero2Hero,
Flink-ddd,
bitborne and
inaniloquentee
as code owners
August 11, 2026 14:31
📝 WalkthroughWalkthroughAdded WS2 tensor-parallel logprob contracts, deterministic vocabulary-parallel execution, capability-aware dispatch, direct LSE diagnostics, distributed comparison tooling, documentation, tests, and CI coverage. ChangesWS2 logprob
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant KernelRegistry
participant VocabParallelLogprobOp
participant TPGroup
participant ComparisonReport
Caller->>KernelRegistry: request contract-aware logprob backend
KernelRegistry->>VocabParallelLogprobOp: materialize selected backend
VocabParallelLogprobOp->>TPGroup: reduce vocabulary tile statistics
TPGroup-->>VocabParallelLogprobOp: return replicated LSE and logprob outputs
VocabParallelLogprobOp-->>ComparisonReport: provide outputs and provenance
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: 5
🧹 Nitpick comments (4)
tests/test_logprob_comparison.py (1)
204-224: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd a
timeoutto this subprocess call.
subprocess.runhas no timeout here. If the child process blocks, the test blocks until the CI job limit. The distributed CLI test attests/test_distributed_logprob_comparison.pyline 206 already setstimeout=120. Match that behavior.The Ruff S603 and ast-grep injection hints on this call are false positives. The argument list is literal, and no shell is used.
♻️ Proposed change
check=True, capture_output=True, text=True, + timeout=120, )🤖 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_logprob_comparison.py` around lines 204 - 224, Add timeout=120 to the subprocess.run call in the test using the literal argument list, matching the distributed CLI test’s timeout behavior; leave the existing check, output capture, and shell-free invocation unchanged.Source: Linters/SAST tools
tests/test_distributed_logprob_comparison.py (1)
172-208: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSurface child stderr when this four-process run fails.
check=TrueraisesCalledProcessError, and that message omits the captured stderr.torchrunwrites per-rank tracebacks to stderr. Without them, a CI failure of this test reports only a return code.♻️ Proposed change
- check=True, + check=False, capture_output=True, text=True, timeout=120, env=environment, ) + assert result.returncode == 0, result.stderr🤖 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_distributed_logprob_comparison.py` around lines 172 - 208, Update the subprocess invocation in the distributed run test to surface captured child stderr when torch.distributed.run fails. Preserve check=True and the existing capture behavior, but catch the resulting CalledProcessError and include its stderr in the test failure output.rl_engine/testing/distributed_logprob_comparison.py (1)
17-17: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueImport
Sequencefromcollections.abc.The project supports Python 3.10+, and the sibling module already uses this 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/testing/distributed_logprob_comparison.py` at line 17, Update the typing imports in distributed_logprob_comparison.py to import Sequence from collections.abc instead of typing, while retaining Any from typing and leaving the module’s usage unchanged.rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py (1)
263-295: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winEach new
forward_with_lseentry point re-implements the validation already present inapply. Both backends added a diagnostic path with its own copy of the device, dimensionality, shape, and target-range checks, so the two copies in each file can diverge.
rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py#L263-L295: call the new_validate_inputshelper fromapply(lines 210-243) instead of keeping the inline copy.rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py#L178-L202: extract the checks into one static validator and call it from bothapply(lines 138-161) andforward_with_lse.🤖 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/batch_invariant_logp.py` around lines 263 - 295, The Triton validation is duplicated instead of reused, and the CUDA backend also needs a shared validator. In rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py:263-295, update apply to call the existing _validate_inputs helper and remove its inline checks. In rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py:178-202, extract the validation into one static validator and call it from both apply and forward_with_lse.
🤖 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/batch-invariant-logp.md`:
- Around line 57-62: Update the VocabParallelLogprobOp validation statement to
make bit-identical results conditional rather than unconditional: specify that
it applies when num_vocab_tiles is fixed and all shard bounds are tile-aligned,
and acknowledge that CUDA/NCCL BF16 validation remains pending.
In `@rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py`:
- Around line 378-401: Make local validation in the vocab-parallel logprob entry
point collective when contract.sharding.tp_world_size > 1: capture failures from
_tile_size, _validate_invocation, and _validate_active_targets, exchange
per-rank success or error status with the existing
_preflight_cross_rank_agreement all_gather_object mechanism, and raise the
corresponding LogprobContractError on every rank before any later collective or
_VocabParallelLogprobFunction.apply call. Preserve single-rank behavior and
avoid entering subsequent collectives when any rank reports validation failure.
In `@rl_engine/testing/distributed_logprob_comparison.py`:
- Around line 621-627: Update run_distributed_logprob_case to accept a
keyword-only collective_timeout_s parameter defaulting to 300, expose it through
the --collective-timeout CLI option, and pass
datetime.timedelta(seconds=collective_timeout_s) as timeout to
dist.init_process_group in the initialization block. Ensure the CLI value is
forwarded to run_distributed_logprob_case for both gloo and nccl runs.
- Around line 386-405: Update the relative-error calculation in the drift
comparison flow around selected_diff and selected_ref to clamp the denominator
to a representable epsilon, preventing max_rel overflow for near-zero
references. Also update both JSON serialization write sites using json.dumps to
pass allow_nan=False so non-finite report values fail loudly instead of emitting
invalid JSON.
In `@tests/test_vocab_parallel_logp.py`:
- Around line 179-192: Remove the initial overwritten contract assignment in
test_matches_ws1_batch_invariant_logp_within_contract_tolerance, keeping only
the real_vocab=PADDED_VOCAB and padded_vocab=PADDED_VOCAB contract used by the
test.
---
Nitpick comments:
In `@rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py`:
- Around line 263-295: The Triton validation is duplicated instead of reused,
and the CUDA backend also needs a shared validator. In
rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py:263-295, update apply
to call the existing _validate_inputs helper and remove its inline checks. In
rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py:178-202, extract the
validation into one static validator and call it from both apply and
forward_with_lse.
In `@rl_engine/testing/distributed_logprob_comparison.py`:
- Line 17: Update the typing imports in distributed_logprob_comparison.py to
import Sequence from collections.abc instead of typing, while retaining Any from
typing and leaving the module’s usage unchanged.
In `@tests/test_distributed_logprob_comparison.py`:
- Around line 172-208: Update the subprocess invocation in the distributed run
test to surface captured child stderr when torch.distributed.run fails. Preserve
check=True and the existing capture behavior, but catch the resulting
CalledProcessError and include its stderr in the test failure output.
In `@tests/test_logprob_comparison.py`:
- Around line 204-224: Add timeout=120 to the subprocess.run call in the test
using the literal argument list, matching the distributed CLI test’s timeout
behavior; leave the existing check, output capture, and shell-free invocation
unchanged.
🪄 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: 4e5c280f-d98e-48fb-a4a9-61cfe9df7c9f
📒 Files selected for processing (17)
.github/workflows/ci.ymldocs/design/runtime-dispatch.mddocs/operators/batch-invariant-logp.mdrl_engine/kernels/logprob_contract.pyrl_engine/kernels/ops/cuda/loss/batch_invariant_logp.pyrl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.pyrl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.pyrl_engine/kernels/ops/triton/loss/batch_invariant_logp.pyrl_engine/kernels/registry.pyrl_engine/testing/__init__.pyrl_engine/testing/distributed_logprob_comparison.pyrl_engine/testing/logprob_comparison.pyrl_engine/testing/logprob_drift.pytests/test_distributed_logprob_comparison.pytests/test_logprob_comparison.pytests/test_logprob_contract.pytests/test_vocab_parallel_logp.py
Comment on lines
+378
to
+401
| if not isinstance(contract, LogprobContract): | ||
| raise LogprobContractError("contract must be a LogprobContract") | ||
| tile = _tile_size(contract, num_vocab_tiles) | ||
| _validate_invocation(local_logits, target_ids, contract, tp_group) | ||
|
|
||
| target_1d = target_ids.reshape(-1).to(device=local_logits.device, dtype=torch.long) | ||
| active_mask = torch.tensor( | ||
| contract.mask.active_mask, dtype=torch.bool, device=local_logits.device | ||
| ) | ||
| if validate: | ||
| _validate_active_targets(target_1d, active_mask, contract.sharding.real_vocab_size) | ||
| if contract.sharding.tp_world_size > 1: | ||
| _preflight_cross_rank_agreement(contract, tp_group, num_vocab_tiles) | ||
|
|
||
| selected_logp, lse = _VocabParallelLogprobFunction.apply( | ||
| local_logits, target_1d, active_mask, contract, tp_group, tile | ||
| ) | ||
|
|
||
| if validate and bool((~torch.isfinite(lse) & active_mask).any().item()): | ||
| raise LogprobContractError( | ||
| "non-finite logsumexp on an active row: logits over the real " | ||
| "vocabulary must be finite for every active token" | ||
| ) | ||
| return selected_logp, lse |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Rank-local validation failures can strand the other TP ranks in a collective.
_tile_size, _validate_invocation, and _validate_active_targets all raise from local state only. If one rank fails one of these checks and the other ranks pass, the failing rank returns from apply while the remaining ranks proceed into _preflight_cross_rank_agreement (all_gather_object) or into the forward all_gather. Those ranks then block until the process-group timeout.
The current tests only cover symmetric failures: _case_tile_misaligned_bounds perturbs the bounds on every rank, and the preflight scenario still lets every rank enter all_gather_object. An asymmetric failure, for example a wrong local shard width on a single rank, is not covered and hangs.
Consider converting local validation into a collective decision when tp_world_size > 1: gather a per-rank status (or the validation error text) with the same all_gather_object payload already used by the preflight, then raise on every rank.
🛡️ Sketch: fold local validation results into the existing preflight exchange
def _preflight_cross_rank_agreement(
- contract: LogprobContract, tp_group: Any, num_vocab_tiles: int
+ contract: LogprobContract, tp_group: Any, num_vocab_tiles: int, local_error: str | None = None
) -> None:
dist = _require_distributed_initialized()
payload = (contract.cross_rank_fingerprint(), BACKEND_ID, int(num_vocab_tiles))
world = dist.get_world_size(group=tp_group)
- gathered: list[Any] = [None] * world
- dist.all_gather_object(gathered, payload, group=tp_group)
+ gathered: list[Any] = [None] * world
+ dist.all_gather_object(gathered, (payload, local_error), group=tp_group)
+ failures = [(rank, err) for rank, (_, err) in enumerate(gathered) if err]
+ if failures:
+ raise LogprobContractError(f"rank-local validation failed on ranks: {failures}")🤖 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/vocab_parallel_logp.py` around lines 378 -
401, Make local validation in the vocab-parallel logprob entry point collective
when contract.sharding.tp_world_size > 1: capture failures from _tile_size,
_validate_invocation, and _validate_active_targets, exchange per-rank success or
error status with the existing _preflight_cross_rank_agreement all_gather_object
mechanism, and raise the corresponding LogprobContractError on every rank before
any later collective or _VocabParallelLogprobFunction.apply call. Preserve
single-rank behavior and avoid entering subsequent collectives when any rank
reports validation failure.
Comment on lines
+179
to
+192
| def test_matches_ws1_batch_invariant_logp_within_contract_tolerance(self): | ||
| tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"] | ||
| contract = _contract(padded_vocab=REAL_VOCAB + 5) | ||
| # Use a real==padded contract so the WS1 op sees identical logits. | ||
| contract = _contract(real_vocab=PADDED_VOCAB, padded_vocab=PADDED_VOCAB) | ||
| logits, targets = _inputs() | ||
| logp, _ = VocabParallelLogprobOp()( | ||
| logits, targets, contract=contract, num_vocab_tiles=NUM_TILES | ||
| ) | ||
| ws1 = NativeBatchInvariantLogpOp().apply(logits, targets) | ||
| active = torch.tensor(ACTIVE) | ||
| assert torch.allclose( | ||
| logp[active], ws1[active], atol=tolerance["atol"], rtol=tolerance["rtol"] | ||
| ) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Remove the overwritten contract assignment.
Line 181 builds a contract with padded_vocab=REAL_VOCAB + 5, and line 183 immediately replaces it. The first assignment has no effect and suggests a padded-vocab case is covered here when it is not.
🧹 Proposed cleanup
tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"]
- contract = _contract(padded_vocab=REAL_VOCAB + 5)
# Use a real==padded contract so the WS1 op sees identical logits.
contract = _contract(real_vocab=PADDED_VOCAB, padded_vocab=PADDED_VOCAB)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| def test_matches_ws1_batch_invariant_logp_within_contract_tolerance(self): | |
| tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"] | |
| contract = _contract(padded_vocab=REAL_VOCAB + 5) | |
| # Use a real==padded contract so the WS1 op sees identical logits. | |
| contract = _contract(real_vocab=PADDED_VOCAB, padded_vocab=PADDED_VOCAB) | |
| logits, targets = _inputs() | |
| logp, _ = VocabParallelLogprobOp()( | |
| logits, targets, contract=contract, num_vocab_tiles=NUM_TILES | |
| ) | |
| ws1 = NativeBatchInvariantLogpOp().apply(logits, targets) | |
| active = torch.tensor(ACTIVE) | |
| assert torch.allclose( | |
| logp[active], ws1[active], atol=tolerance["atol"], rtol=tolerance["rtol"] | |
| ) | |
| def test_matches_ws1_batch_invariant_logp_within_contract_tolerance(self): | |
| tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"] | |
| # Use a real==padded contract so the WS1 op sees identical logits. | |
| contract = _contract(real_vocab=PADDED_VOCAB, padded_vocab=PADDED_VOCAB) | |
| logits, targets = _inputs() | |
| logp, _ = VocabParallelLogprobOp()( | |
| logits, targets, contract=contract, num_vocab_tiles=NUM_TILES | |
| ) | |
| ws1 = NativeBatchInvariantLogpOp().apply(logits, targets) | |
| active = torch.tensor(ACTIVE) | |
| assert torch.allclose( | |
| logp[active], ws1[active], atol=tolerance["atol"], rtol=tolerance["rtol"] | |
| ) |
🤖 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_vocab_parallel_logp.py` around lines 179 - 192, Remove the initial
overwritten contract assignment in
test_matches_ws1_batch_invariant_logp_within_contract_tolerance, keeping only
the real_vocab=PADDED_VOCAB and padded_vocab=PADDED_VOCAB contract used by the
test.
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 `@rl_engine/testing/distributed_logprob_comparison.py`:
- Line 630: Move the existing try/finally scope in the process-group setup flow
so it begins before dist.init_process_group and includes initialization, device
setup, rank_topology, _create_tp_group, and execution. Preserve the cleanup in
the matching finally so failures during any setup step destroy the process group
and reset state.
🪄 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: 6116c6d3-092f-483b-a355-074f4179aa9b
📒 Files selected for processing (3)
docs/operators/batch-invariant-logp.mdrl_engine/testing/distributed_logprob_comparison.pytests/test_distributed_logprob_comparison.py
🚧 Files skipped from review as they are similar to previous changes (2)
- docs/operators/batch-invariant-logp.md
- tests/test_distributed_logprob_comparison.py
6 tasks
KJLdefeated
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add the cross-config distributed logprob drift runner requested by PR4 of #241.
The runner materializes one explicit TP/CP topology per
torchruninvocation, executes the deterministic vocab-parallel logprob reference introduced by PR3, and compares its BF16 selected logprob and vocabulary LSE against an independent full-vocabulary FP32 PyTorch oracle. It emits per-rank and aggregate artifacts with topology, backend, reduction, tolerance, and worst-token provenance.The distributed logprob implementation itself remains owned by PR3. This PR adds the scoped integration, validation, and reporting layer.
Implements PR4 of #241.
Dependencies and current branch state
This PR depends on the three preceding issue #241 changes:
All three prerequisite PRs are still open. To start PR4 validation without waiting for them to merge, this branch was built on a temporary integration of their current reviewed heads. As a result, a PR opened against the current
mainwill temporarily show prerequisite changes in addition to the PR4 changes.The PR can be reviewed as a draft while those dependencies are open. The PR4-specific commits are:
After #259, #262, and #265 merge, these PR4 changes will be replayed onto the latest
mainand the branch will be updated so that the final merge diff contains PR4 only. This PR should not merge before its prerequisites.Known PR3 follow-ups
Review of the stacked PR3 content identified two items that should be resolved in #265 before this PR is rebased for merge:
test_matches_ws1_batch_invariant_logp_within_contract_tolerancecontains an immediately overwritten contract assignment that should be removed.The operator documentation in this branch also qualifies the cross-TP bitwise guarantee with its required fixed-tile and tile-aligned-shard conditions. PR4 will sync the final #265 resolution before leaving draft state.
Scope
This PR covers the cross-config integration and distributed drift report described in PR4:
TP=1/2/4 x CP=1/2 x BF16topology matrix.pytorch-vocab-parallel-logp-ws2backend from PR3.torchruncommand required to reproduce every artifact.CP is represented in topology and report reconstruction, but it is not a numerical merge axis for vocab-parallel logprob. This PR does not implement a new logprob reducer or fused kernel, GRPO loss reduction, PP/DP validation, or rollout-versus-training integration from WS3.
Changes
Distributed topology runner
Add
rl_engine/testing/distributed_logprob_comparison.pywith:DistributedLogprobCasefor TP/CP topology, dtype, vocabulary, tiles, shape, mask, and seed.torchrunentry point.Each invocation materializes exactly one topology.
WORLD_SIZEmust equalTP * CP; a mismatched launch fails before the process group is used.Testing-module boundary
Keep the distributed entry point separate from
rl_engine/testing/logprob_comparison.pyintentionally. The PR2 module is a single-process, multi-backend comparison CLI, while the PR4 module is atorchrunentry point that materializes one explicit distributed backend and owns process groups, TP/CP topology, rank aggregation, and artifact writing. Both remain kernel-local underrl_engine/testing/; no project-wide script is added.This separation avoids changing the already reviewed PR2 CLI and avoids combining two different execution models into a single module of more than one thousand lines. If maintainers prefer one consolidated logprob testing entry point, the two modules can be combined behind explicit single-GPU and distributed subcommands during review.
Independent FP32 oracle
Every rank reconstructs the same seeded logical FP32 logits, targets, and active mask. The candidate path receives only its BF16 token/vocabulary shard, while the oracle computes directly from the complete real vocabulary:
The oracle does not use the PR3 sharded merge implementation, so the candidate is not compared against itself.
Drift and replication checks
For LSE and active-token dlogp, report:
The worst-drift record also includes the global token position, target id, target-owner TP rank, candidate value, and reference value. Tolerances come from the existing #108 tolerance contract.
After the PR3 TP merge, every rank in the same TP group should hold the same LSE and selected logprob. The runner gathers those outputs and checks them with
torch.equal. Aggregate statistics then use onlytp_rank=0from each CP shard, avoiding duplicate counting of TP replicas.Fail-closed backend materialization
The distributed runner requires an explicit backend and records both the requested and actual implementation. It fails when:
autoor an empty backend is requested.Structured distributed artifacts
Global rank zero writes a #116-style JSON artifact containing:
A failed numerical comparison still writes the artifact, then exits nonzero on global rank zero.
Operator documentation
Document the distributed topology contract, oracle, report fields, launch commands, and validation status in
docs/operators/batch-invariant-logp.md.Comparison contract
TP is the only numerical merge axis. CP partitions independent token rows and is used only when reconstructing the full logical report:
LSE drift includes every logical token row. Selected-logprob drift includes only active response/action tokens. Prompt and ignored rows do not contribute to dlogp statistics.
Tests
Add
tests/test_distributed_logprob_comparison.pywith coverage for:TP=1/2/4 x CP=1/2topology plan.autobackend and non-tileable vocabulary cases.TP=2, CP=2subprocess run.The CPU-safe PR2 and PR4 comparison suites are wired into
.github/workflows/ci.yml, including the four-process Gloo case.Validation
WSL focused PR1-PR4 suite
Result:
This includes the real four-process CPU/Gloo
TP=2, CP=2test. The skipped tests require CUDA, Triton, or a compiled CUDA extension.Batch-invariant and vocab-parallel regression
Result:
The skipped tests are GPU/backend-specific in the current WSL environment.
Static checks
Pending CUDA/NCCL validation
The distributed control flow and artifacts have been validated with CPU/Gloo. The scoped BF16 numerical gate still needs a real CUDA/NCCL run on H800/H100-class hardware with Qwen3 vocabulary size
151936:Print the exact single-node launch matrix:
Example four-GPU target case:
The full matrix requires eight ranks for
TP=4, CP=2. Multi-node launch compatibility is supported by standardtorchrunrendezvous arguments, but multi-node NCCL has not yet been validated and is not required if an eight-GPU single-node host is available.Notes for review
torchrunlifecycle and topology/artifact responsibilities differ from the PR2 single-process comparison CLI. This boundary can be consolidated if maintainers prefer one entry point.