[WS2][PR5][Logp] Add deterministic DP-aware GRPO loss on the TP-aware logprob path - #291
[WS2][PR5][Logp] Add deterministic DP-aware GRPO loss on the TP-aware logprob path#291KJLdefeated wants to merge 10 commits into
Conversation
Implements PR 1 of issue #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 #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 #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 #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.
📝 WalkthroughWalkthroughThis PR adds validated WS2 contracts, deterministic tensor-parallel logprob, distributed GRPO loss execution, contract-aware kernel dispatch, documentation, CI commands, and comprehensive single-rank and distributed tests. ChangesWS2 Contract and Execution
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant KernelRegistry
participant DistributedGRPOLossOp
participant VocabParallelLogprobOp
participant TPDPCollectives
KernelRegistry->>DistributedGRPOLossOp: resolve loss backend from contract
DistributedGRPOLossOp->>TPDPCollectives: validate cross-rank contract agreement
DistributedGRPOLossOp->>VocabParallelLogprobOp: compute sharded policy and reference logprob
VocabParallelLogprobOp->>TPDPCollectives: merge tiled vocabulary statistics
DistributedGRPOLossOp->>TPDPCollectives: gather rewards and sequence reductions
TPDPCollectives->>DistributedGRPOLossOp: return global advantages and reductions
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 💡 2📝 Generate docstrings 💡
🛠️ 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: 8
🧹 Nitpick comments (9)
tests/test_distributed_grpo_loss.py (2)
508-515: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDrop the unused
contractbinding.Ruff flags
contractas unused (RUF059).- result, tensors, contract = _run_single_rank() + result, tensors, _ = _run_single_rank()🤖 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_grpo_loss.py` around lines 508 - 515, Remove the unused contract binding in test_forward_and_backward_run by unpacking only the result and tensors returned from _run_single_rank, while preserving the existing assertions and test behavior.Source: Linters/SAST tools
185-189: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAssert that
NUM_SEQUENCESdividesdp.
_dp_boundsuses integer division. If a futureMESH_CONFIGSentry adds a DP degree that does not divideNUM_SEQUENCES, the bounds silently drop the trailing sequences and every bitwise comparison then runs on a smaller batch than the baseline. An assertion turns that into an immediate failure.♻️ Proposed change
def _dp_bounds(dp: int) -> tuple[tuple[int, int], ...]: """Contiguous sequence partition in DP-rank order.""" + assert NUM_SEQUENCES % dp == 0, f"DP={dp} does not divide NUM_SEQUENCES={NUM_SEQUENCES}" seqs = NUM_SEQUENCES // dp return tuple((d * seqs, (d + 1) * seqs) for d in range(dp))🤖 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_grpo_loss.py` around lines 185 - 189, Update _dp_bounds to assert that NUM_SEQUENCES is evenly divisible by dp before calculating seqs, so unsupported DP configurations fail immediately while preserving the existing contiguous bounds calculation.rl_engine/kernels/ops/pytorch/loss/distributed_grpo_loss.py (2)
443-447: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSort
__all__to clear the Ruff RUF022 warning.__all__ = [ "BACKEND_ID", - "GRPOLossResult", "DistributedGRPOLossOp", + "GRPOLossResult", ]🤖 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/distributed_grpo_loss.py` around lines 443 - 447, Sort the entries in __all__ alphabetically to resolve Ruff RUF022, keeping the same exported symbols: BACKEND_ID, DistributedGRPOLossOp, and GRPOLossResult.Source: Linters/SAST tools
355-359: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueBuild the active mask once per contract instead of once per call.
torch.tensor(...)converts a Python tuple ofnum_tokensbooleans on everyapplycall, then copies it to the device. For long sequences this dominates the CPU time of the operator and adds a host-to-device copy on the critical path. Cache the tensor on the op keyed by the contract identity, or derive it fromcontract.logprob.maskonce at contract construction.🤖 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/distributed_grpo_loss.py` around lines 355 - 359, Update the active-mask handling in the loss operator around the apply path so the boolean tensor is created once per contract rather than on every call. Cache or attach the device-appropriate mask using the contract identity, reuse it across calls, and preserve its dtype and device behavior.rl_engine/kernels/registry.py (2)
641-749: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the fallback semantics in
get_loss_opas well.
get_logprob_opcarries a comment at Lines 585-588 that explains why policy-only skips do not incrementcapability_rejections.get_loss_opimplements the same rule at Lines 713-716 with no comment. Add the same note so a future edit does not treat a policy skip as a fallback here.🤖 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/registry.py` around lines 641 - 749, Add an inline comment in get_loss_op immediately before or around the _loss_policy_mismatch handling to document that policy-only skips must not increment capability_rejections, matching the existing explanation in get_logprob_op. Do not change the fallback logic or rejection behavior.
751-793: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsolidate the two identical policy-mismatch helpers.
_loss_policy_mismatchand_logprob_policy_mismatchhave identical bodies. They differ only in the capability type annotation, and both capability types exposeimplementation_kindandbackend_id. Keep one helper so a future policy keyword cannot be added to one path and forgotten in the other.♻️ Proposed refactor
`@staticmethod` - def _loss_policy_mismatch( + def _policy_mismatch( requested_backend: str, - capability: LossBackendCapability, + capability: LossBackendCapability | LogprobBackendCapability, ) -> str | None: policy = requested_backend.lower() if policy == "auto": return None if policy in IMPLEMENTATION_KINDS: if capability.implementation_kind == policy: return None return ( f"implementation_kind={capability.implementation_kind} does not satisfy " f"requested_backend={policy}" ) if capability.backend_id == requested_backend: return None return ( f"backend_id={capability.backend_id} does not match " f"requested_backend={requested_backend}" ) - - `@staticmethod` - def _logprob_policy_mismatch( - requested_backend: str, - capability: LogprobBackendCapability, - ) -> str | None: - policy = requested_backend.lower() - if policy == "auto": - return None - if policy in IMPLEMENTATION_KINDS: - if capability.implementation_kind == policy: - return None - return ( - f"implementation_kind={capability.implementation_kind} does not satisfy " - f"requested_backend={policy}" - ) - if capability.backend_id == requested_backend: - return None - return ( - f"backend_id={capability.backend_id} does not match " - f"requested_backend={requested_backend}" - )Update both call sites at Line 597 and Line 713.
🤖 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/registry.py` around lines 751 - 793, Consolidate _loss_policy_mismatch and _logprob_policy_mismatch into a single shared policy-mismatch helper, using a capability type that supports implementation_kind and backend_id. Update both existing call sites near the loss and logprob selection paths to use the shared helper, and remove the duplicate method while preserving all current matching and error-message behavior.tests/test_vocab_parallel_logp.py (1)
367-388: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winThe
preflightscenario depends on rank 0 surviving_tile_size.The scenario gives rank 0
num_vocab_tiles = 64so the ranks disagree. This only reaches_preflight_cross_rank_agreementbecause_tile_sizeruns first and happens to accept 64:TP_PADDED_VOCAB=1024is divisible by 64, and the shard bounds are multiples ofTP_TILE=32, so they are also 16-aligned.If
TP_PADDED_VOCAB,TP_NUM_TILES, or_tp_boundschange,_tile_sizewill reject 64 on rank 0 only. Rank 0 then raises before the collective while rank 1 blocks insideall_gather_objectuntil_SPAWN_TIMEOUT_S(300 s) expires. The test would fail as a timeout instead of a clear assertion.Consider asserting the invariant in the scenario setup, for example that
TP_PADDED_VOCAB % (TP_NUM_TILES * 2) == 0and that every bound is aligned to the doubled tile size, so a future geometry change fails fast.🤖 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 367 - 388, The preflight scenario assumes the doubled tile count remains valid for _tile_size, but does not verify that invariant. In the scenario setup around TP_NUM_TILES, TP_PADDED_VOCAB, and bounds, add fail-fast assertions that the padded vocabulary is divisible by TP_NUM_TILES * 2 and every shard boundary is aligned to the corresponding doubled tile size, before invoking op.rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py (2)
284-307: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueConsider the memory cost of saving fp32
z_maskedfor backward.
forwardupcastslocal_logitsto fp32 and saves the full[n, local_vocab]tensor. For a bf16 input this retains 4x the input bytes for the whole backward window, in addition to the caller's own copy oflocal_logits.backwardderivespfromz_maskedandlseonly, so the cast could be recomputed from the saved input instead.This is a deliberate tradeoff for a reference backend. Consider saving
local_logitspluspadding_colsand re-derivingz_maskedinbackwardif activation memory becomes a limit at largelocal_vocab.🤖 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 284 - 307, The forward path currently saves the full fp32 z_masked tensor; reduce activation memory by saving the original local_logits together with padding_cols instead, then recast and apply the padding mask in backward before deriving probabilities from z_masked and lse. Update ctx.save_for_backward and the backward symbol that consumes the saved tensors while preserving the existing numerical behavior.
384-400: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winReduce the per-call host work and device syncs on the default path.
Two costs land on every call, and
validatedefaults toTrue:
- Line 384:
torch.tensor(contract.mask.active_mask, ...)converts a Python tuple ofnum_tokensbools into a device tensor on each call.LogprobContractis frozen, so this value never changes for a given contract and device.- Line 388 (
_validate_active_targets) and Line 396 each call.item(), which forces a device sync. That is two syncs per invocation on the training path.Cache the mask tensor per
(contract, device)and consider documenting that callers should passvalidate=Falsein steady-state training after a warmup validation.🤖 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 384 - 400, Reduce default-path overhead in the logprob call around _VocabParallelLogprobFunction.apply by caching the immutable active_mask tensor per LogprobContract and device instead of recreating it from the Python tuple on every invocation. Avoid the two per-call device synchronizations from _validate_active_targets and the non-finite lse check where possible, while preserving validation behavior; document that steady-state callers may use validate=False after warmup 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 `@docs/operators/batch-invariant-logp.md`:
- Around line 78-80: Update the usage example around
kernel_registry.get_logprob_op to explicitly import LogprobContract, complete
the contract reference comment, and revise the result.op comment to identify it
as the selected dispatched operator rather than the contract module.
In `@rl_engine/kernels/loss_contract.py`:
- Around line 624-631: Update cross_rank_fingerprint to preserve the
rank-invariant logprob["mask"] fields instead of removing the entire mask block.
Exclude only num_tokens and active_mask_sha256, while retaining ignore_index and
any other globally consistent mask fields; leave the existing sharding filtering
unchanged.
- Line 126: Suppress Ruff’s S105 false positive for the TWO_PASS enum
discriminator using the repository’s standard inline or file-level noqa
mechanism, while preserving the constant name and value and avoiding broader
lint-rule suppression.
In `@rl_engine/kernels/ops/pytorch/loss/distributed_grpo_loss.py`:
- Around line 410-417: Update the zero-active-token guard around total_active
and _normalized so no normalization occurs when total_active is zero, regardless
of validate. Preserve the LossContractError for validate=True, and for
validate=False return the established zero-loss behavior or otherwise bypass
policy_loss and kl normalization without producing inf, NaN, or gradients from
division by zero.
In `@rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py`:
- Line 190: Update _local_tile_stats to unpack z_masked.shape without assigning
the unused n variable, while preserving the local_vocab value used by the
function.
In `@tests/test_distributed_grpo_loss.py`:
- Around line 820-836: Run Black on tests/test_distributed_grpo_loss.py to apply
the required formatting around the assertions, ensuring the file passes the CI
linting checks.
In `@tests/test_vocab_parallel_logp.py`:
- Around line 471-481: Update the process cleanup in the test’s finally block to
record each process terminated after the bounded join, then restrict the
exitcode assertions to processes not recorded as terminated. Preserve result
validation while avoiding assertions for terminated workers or workers whose
exitcode remains unavailable.
- Around line 179-184: Remove the unused initial contract assignment in
test_matches_ws1_batch_invariant_logp_within_contract_tolerance, keeping only
the later _contract(real_vocab=PADDED_VOCAB, padded_vocab=PADDED_VOCAB)
assignment used by the test.
---
Nitpick comments:
In `@rl_engine/kernels/ops/pytorch/loss/distributed_grpo_loss.py`:
- Around line 443-447: Sort the entries in __all__ alphabetically to resolve
Ruff RUF022, keeping the same exported symbols: BACKEND_ID,
DistributedGRPOLossOp, and GRPOLossResult.
- Around line 355-359: Update the active-mask handling in the loss operator
around the apply path so the boolean tensor is created once per contract rather
than on every call. Cache or attach the device-appropriate mask using the
contract identity, reuse it across calls, and preserve its dtype and device
behavior.
In `@rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py`:
- Around line 284-307: The forward path currently saves the full fp32 z_masked
tensor; reduce activation memory by saving the original local_logits together
with padding_cols instead, then recast and apply the padding mask in backward
before deriving probabilities from z_masked and lse. Update
ctx.save_for_backward and the backward symbol that consumes the saved tensors
while preserving the existing numerical behavior.
- Around line 384-400: Reduce default-path overhead in the logprob call around
_VocabParallelLogprobFunction.apply by caching the immutable active_mask tensor
per LogprobContract and device instead of recreating it from the Python tuple on
every invocation. Avoid the two per-call device synchronizations from
_validate_active_targets and the non-finite lse check where possible, while
preserving validation behavior; document that steady-state callers may use
validate=False after warmup validation.
In `@rl_engine/kernels/registry.py`:
- Around line 641-749: Add an inline comment in get_loss_op immediately before
or around the _loss_policy_mismatch handling to document that policy-only skips
must not increment capability_rejections, matching the existing explanation in
get_logprob_op. Do not change the fallback logic or rejection behavior.
- Around line 751-793: Consolidate _loss_policy_mismatch and
_logprob_policy_mismatch into a single shared policy-mismatch helper, using a
capability type that supports implementation_kind and backend_id. Update both
existing call sites near the loss and logprob selection paths to use the shared
helper, and remove the duplicate method while preserving all current matching
and error-message behavior.
In `@tests/test_distributed_grpo_loss.py`:
- Around line 508-515: Remove the unused contract binding in
test_forward_and_backward_run by unpacking only the result and tensors returned
from _run_single_rank, while preserving the existing assertions and test
behavior.
- Around line 185-189: Update _dp_bounds to assert that NUM_SEQUENCES is evenly
divisible by dp before calculating seqs, so unsupported DP configurations fail
immediately while preserving the existing contiguous bounds calculation.
In `@tests/test_vocab_parallel_logp.py`:
- Around line 367-388: The preflight scenario assumes the doubled tile count
remains valid for _tile_size, but does not verify that invariant. In the
scenario setup around TP_NUM_TILES, TP_PADDED_VOCAB, and bounds, add fail-fast
assertions that the padded vocabulary is divisible by TP_NUM_TILES * 2 and every
shard boundary is aligned to the corresponding doubled tile size, before
invoking op.
🪄 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: b7cca0bc-d787-41ec-ae7d-1d2e010e7361
📒 Files selected for processing (13)
.github/workflows/ci.ymldocs/design/runtime-dispatch.mddocs/operators/batch-invariant-logp.mddocs/operators/grpo-loss.mdrl_engine/kernels/logprob_contract.pyrl_engine/kernels/loss_contract.pyrl_engine/kernels/ops/pytorch/loss/distributed_grpo_loss.pyrl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.pyrl_engine/kernels/registry.pytests/test_distributed_grpo_loss.pytests/test_grpo_loss_contract.pytests/test_logprob_contract.pytests/test_vocab_parallel_logp.py
| result = kernel_registry.get_logprob_op(contract) # LogprobContract from | ||
| op = result.op # rl_engine.kernels.logprob_contract | ||
| logp, lse = op(local_logits, target_ids, contract=contract, tp_group=tp_group) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Correct the imports and inline comments in the usage example.
Line 78 contains an incomplete comment. Line 79 associates result.op with rl_engine.kernels.logprob_contract, but that module defines the contract, not the dispatched operator. Add an explicit LogprobContract import and describe result.op as the selected operator.
Proposed documentation fix
from rl_engine.kernels.registry import kernel_registry
+from rl_engine.kernels.logprob_contract import LogprobContract
-result = kernel_registry.get_logprob_op(contract) # LogprobContract from
-op = result.op # rl_engine.kernels.logprob_contract
+result = kernel_registry.get_logprob_op(contract) # contract: LogprobContract
+op = result.op # selected TP-aware logprob operator🤖 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/batch-invariant-logp.md` around lines 78 - 80, Update the
usage example around kernel_registry.get_logprob_op to explicitly import
LogprobContract, complete the contract reference comment, and revise the
result.op comment to identify it as the selected dispatched operator rather than
the contract module.
| claim. The two-pass form subtracts the mean before squaring. | ||
| """ | ||
|
|
||
| TWO_PASS = "two_pass" |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Silence the Ruff S105 false positive on TWO_PASS.
Ruff flags TWO_PASS = "two_pass" as a hardcoded password (S105) because the name contains PASS. The value is an enum discriminator, not a secret. If the lint gate treats S rules as errors, this blocks CI.
🔧 Proposed fix
- TWO_PASS = "two_pass"
+ TWO_PASS = "two_pass" # noqa: S105 - enum discriminator, not a secret📝 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.
| TWO_PASS = "two_pass" | |
| TWO_PASS = "two_pass" # noqa: S105 - enum discriminator, not a secret |
🧰 Tools
🪛 Ruff (0.16.1)
[error] 126-126: Possible hardcoded password assigned to: "TWO_PASS"
(S105)
🤖 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/loss_contract.py` at line 126, Suppress Ruff’s S105 false
positive for the TWO_PASS enum discriminator using the repository’s standard
inline or file-level noqa mechanism, while preserving the constant name and
value and avoiding broader lint-rule suppression.
Source: Linters/SAST tools
| payload = self.to_dict() | ||
| logprob = payload["logprob"] | ||
| logprob.pop("mask", None) | ||
| logprob["sharding"] = { | ||
| key: value | ||
| for key, value in logprob["sharding"].items() | ||
| if key not in {"tp_rank", "cp_rank", "local_vocab_start", "local_vocab_end"} | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Keep rank-invariant mask fields in the cross-rank fingerprint.
cross_rank_fingerprint drops the entire nested logprob["mask"] block. Only num_tokens and active_mask_sha256 legitimately differ per DP rank. ignore_index is global and is part of the numerical identity, so dropping it lets two ranks disagree on ignore_index and still pass the preflight that this method exists to provide. docs/operators/grpo-loss.md Lines 170-177 state that the preflight catches exactly this class of disagreement.
Retain the rank-invariant subset instead of removing the whole block.
🔧 Proposed fix
payload = self.to_dict()
logprob = payload["logprob"]
- logprob.pop("mask", None)
+ # num_tokens and the mask digest are rank-local under DP sharding;
+ # ignore_index is global and must stay in the agreement check.
+ mask = logprob.pop("mask", None)
+ if mask is not None:
+ logprob["mask"] = {"ignore_index": mask["ignore_index"]}
logprob["sharding"] = {📝 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.
| payload = self.to_dict() | |
| logprob = payload["logprob"] | |
| logprob.pop("mask", None) | |
| logprob["sharding"] = { | |
| key: value | |
| for key, value in logprob["sharding"].items() | |
| if key not in {"tp_rank", "cp_rank", "local_vocab_start", "local_vocab_end"} | |
| } | |
| payload = self.to_dict() | |
| logprob = payload["logprob"] | |
| # num_tokens and the mask digest are rank-local under DP sharding; | |
| # ignore_index is global and must stay in the agreement check. | |
| mask = logprob.pop("mask", None) | |
| if mask is not None: | |
| logprob["mask"] = {"ignore_index": mask["ignore_index"]} | |
| logprob["sharding"] = { | |
| key: value | |
| for key, value in logprob["sharding"].items() | |
| if key not in {"tp_rank", "cp_rank", "local_vocab_start", "local_vocab_end"} | |
| } |
🤖 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/loss_contract.py` around lines 624 - 631, Update
cross_rank_fingerprint to preserve the rank-invariant logprob["mask"] fields
instead of removing the entire mask block. Exclude only num_tokens and
active_mask_sha256, while retaining ignore_index and any other globally
consistent mask fields; leave the existing sharding filtering unchanged.
| total_active = int(per_sequence_counts.sum().item()) | ||
| if validate and total_active == 0: | ||
| raise LossContractError( | ||
| "the global batch holds no active tokens; the loss normalizer would divide by zero" | ||
| ) | ||
|
|
||
| policy_loss = _normalized(per_sequence_policy, per_sequence_counts, contract) | ||
| kl = _normalized(per_sequence_kl, per_sequence_counts, contract) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Guard the zero-denominator path when validate=False.
total_active == 0 raises only when validate is true. With validate=False, _normalized divides by zero for GLOBAL_ACTIVE_TOKENS and by live.sum()==0 for PER_SEQUENCE_THEN_MEAN, so the loss becomes inf or NaN and the gradient propagates that silently. The count is already materialized on the host at Line 410, so the check costs nothing extra.
🛡️ Proposed fix
- total_active = int(per_sequence_counts.sum().item())
- if validate and total_active == 0:
+ total_active = int(per_sequence_counts.sum().item())
+ if total_active == 0 and contract.reduction.token_normalizer is not TokenNormalizer.FIXED_CONSTANT:
raise LossContractError(
"the global batch holds no active tokens; the loss normalizer would divide by zero"
)📝 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.
| total_active = int(per_sequence_counts.sum().item()) | |
| if validate and total_active == 0: | |
| raise LossContractError( | |
| "the global batch holds no active tokens; the loss normalizer would divide by zero" | |
| ) | |
| policy_loss = _normalized(per_sequence_policy, per_sequence_counts, contract) | |
| kl = _normalized(per_sequence_kl, per_sequence_counts, contract) | |
| total_active = int(per_sequence_counts.sum().item()) | |
| if total_active == 0 and contract.reduction.token_normalizer is not TokenNormalizer.FIXED_CONSTANT: | |
| raise LossContractError( | |
| "the global batch holds no active tokens; the loss normalizer would divide by zero" | |
| ) | |
| policy_loss = _normalized(per_sequence_policy, per_sequence_counts, contract) | |
| kl = _normalized(per_sequence_kl, per_sequence_counts, contract) |
🤖 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/distributed_grpo_loss.py` around lines 410
- 417, Update the zero-active-token guard around total_active and _normalized so
no normalization occurs when total_active is zero, regardless of validate.
Preserve the LossContractError for validate=True, and for validate=False return
the established zero-loss behavior or otherwise bypass policy_loss and kl
normalization without producing inf, NaN, or gradients from division by zero.
| the identity partial ``(-inf, 0)`` without evaluating ``exp(-inf - (-inf))``. | ||
| """ | ||
|
|
||
| n, local_vocab = z_masked.shape |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Remove the unused unpacked variable n.
Ruff flags n as never used in _local_tile_stats (RUF059). Only local_vocab is read.
🔧 Proposed fix
- n, local_vocab = z_masked.shape
+ local_vocab = z_masked.shape[1]📝 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.
| n, local_vocab = z_masked.shape | |
| local_vocab = z_masked.shape[1] |
🧰 Tools
🪛 Ruff (0.16.1)
[warning] 190-190: Unpacked variable n is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
🤖 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` at line 190,
Update _local_tile_stats to unpack z_masked.shape without assigning the unused n
variable, while preserving the local_vocab value used by the function.
Source: Linters/SAST tools
| assert actual["global_active"] == baseline["global_active"] | ||
| assert actual["per_sequence_counts"] == baseline["per_sequence_counts"] | ||
| assert actual["advantage_bits"] == baseline["advantage_bits"] | ||
| # The sensitive comparison: per-sequence totals, before the scalar | ||
| # average rounds a reordering away. | ||
| assert actual["per_sequence_policy_bits"] == baseline["per_sequence_policy_bits"], ( | ||
| f"{label} per-sequence policy totals differ from the single-rank baseline" | ||
| ) | ||
| assert actual["per_sequence_kl_bits"] == baseline["per_sequence_kl_bits"], ( | ||
| f"{label} per-sequence KL totals differ from the single-rank baseline" | ||
| ) | ||
| assert actual["loss_bits"] == baseline["loss_bits"], f"{label} loss differs" | ||
| assert actual["policy_bits"] == baseline["policy_bits"] | ||
| assert actual["kl_bits"] == baseline["kl_bits"] | ||
| assert actual["grad_row"] == baseline["grad_row"], ( | ||
| f"{label} gradient differs from the single-rank baseline" | ||
| ) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Run Black on this file to fix the failing lint job.
The CI linting job reformatted this range. Run black tests/test_distributed_grpo_loss.py or pre-commit run --all-files, then commit the result.
🤖 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_grpo_loss.py` around lines 820 - 836, Run Black on
tests/test_distributed_grpo_loss.py to apply the required formatting around the
assertions, ensuring the file passes the CI linting checks.
Source: Pipeline failures
| 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() |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Delete the dead contract assignment.
Line 181 assigns contract, then Line 183 overwrites it before any use. The first value is never read. It also suggests the test covers a padded-vocab comparison, which it does not: the effective contract sets real_vocab == padded_vocab.
🔧 Proposed fix
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.
| 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() | |
| 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() |
🤖 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 - 184, Remove the unused
initial contract assignment in
test_matches_ws1_batch_invariant_logp_within_contract_tolerance, keeping only
the later _contract(real_vocab=PADDED_VOCAB, padded_vocab=PADDED_VOCAB)
assignment used by the test.
| finally: | ||
| for process in processes: | ||
| process.join(timeout=30) | ||
| if process.is_alive(): | ||
| process.terminate() | ||
| results.sort(key=lambda item: item["rank"]) | ||
| for result in results: | ||
| assert result["ok"], result.get("traceback") | ||
| for process in processes: | ||
| assert process.exitcode == 0 | ||
| return results |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
The exitcode assertion can fail spuriously after a bounded join.
The finally block joins each process with timeout=30 and calls terminate() if the process is still alive. Line 480 then asserts process.exitcode == 0.
A worker puts its result on the queue before it calls destroy_process_group(), so the parent can reach the join while the child is still tearing down NCCL. If teardown exceeds 30 s, terminate() sets a negative exitcode and Line 480 fails even though every result was correct. If the process is alive and not yet reaped, exitcode is None and the assertion fails on None == 0.
Record which processes were terminated and assert the exit codes only for the processes that exited on their own.
🤖 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 471 - 481, Update the process
cleanup in the test’s finally block to record each process terminated after the
bounded join, then restrict the exitcode assertions to processes not recorded as
terminated. Preserve result validation while avoiding assertions for terminated
workers or workers whose exitcode remains unavailable.
2 participants
[WS2] Add deterministic DP-aware GRPO loss on the TP-aware logprob path
Overview
This pull request implements PR 5 of issue #241. Loss, per-sequence totals, and gradients are bit-identical across TP=1/2/4 × DP=1/2/4.
Stacked on #265; only the top commits are new.
The GRPO objective is elementwise — ratio, clipping, reference KL and the group-relative advantage all act per token or per sequence — so the only place the parallel layout can change the answer is the final sum over tokens. That sum is what this PR pins down. TP is delegated wholesale to the #265 logprob reference: by the time the objective sees a selected logprob, the vocabulary has already been reduced.
Key Changes
Loss Contract (rl_engine/kernels/loss_contract.py)
GRPOLossContractnestsLogprobContractrather than restating it, and cross-checks the two at construction: owned token count, CP degree, and determinism scope must agreeLossReductionSpec.token_normalizermakes the GRPO normalizer ambiguity explicit instead of implicit —global_active_tokens(default; matchesNativeGRPOLossOp's masked mean at DP=1),per_sequence_then_mean(original GRPO),fixed_constant(Dr.GRPO). They differ by more than a scale factor once sequence lengths vary, so the choice is part of the numerical identity and travels in the fingerprintLossShardingSpec.sequence_shard_boundsis a contiguous[0, num_sequences)partition in DP-rank order, mirroringShardingSpec.vocab_shard_bounds. A rank owns each of its sequences whole, which is what makes the reduction degree-independentReference Op (rl_engine/kernels/ops/pytorch/loss/distributed_grpo_loss.py)
DistributedGRPOLossOp: fp32 scalars replicated across DP and TP;GRPOLossResultunpacks as(loss, policy_loss, kl)so it can stand in for the single-GPU tuplepadded_seq_lenslots → one sequence total (entirely local, since the rank owns the whole sequence), thennum_sequencestotals → the scalar. No partial-sequence state ever crosses a rank boundary, so there is nothing to merge and no alignment rule to get wrongall_reduceis excluded on purpose — its combine order follows the collective's topology, not the declared sequence order[num_sequences]vector and keeps its own slice. This is what lets an advantage group straddle DP ranks with no extra machineryall_gathersevers the graph, so a rank writes its own slice into the vector from the live tensor; the other ranks' slices are constants with respect to its logits, which is exactly correct(fingerprint, backend id, num_vocab_tiles)on both the DP and TP axes before any other collective. The TP-axis check is load-bearing: the loss is replicated across TP, so two TP siblings disagreeing onbetawould compute different losses for one sharded model, and the logprob path's own preflight cannot see that becausebetais not part of the logprob contractSummary by CodeRabbit
New Features
Documentation
Tests