[WS2] Add deterministic vocab-parallel TP logprob reference - #265
[WS2] Add deterministic vocab-parallel TP logprob reference#265KJLdefeated wants to merge 9 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.
📝 WalkthroughWalkthroughThe PR adds a WS2 TP-aware logprob contract, a deterministic PyTorch vocab-parallel implementation, strict capability-based registry dispatch, supporting documentation, distributed 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
Caller->>KernelRegistry: get_logprob_op(contract)
KernelRegistry->>VocabParallelLogprobOp: select compatible backend
Caller->>VocabParallelLogprobOp: apply(local_logits, target_ids, contract)
VocabParallelLogprobOp->>TPGroup: all-gather tiled reduction partials
TPGroup-->>VocabParallelLogprobOp: deterministic global reductions
VocabParallelLogprobOp-->>Caller: selected logprob and vocabulary LSE
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 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (8)
rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py (2)
346-347: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the empty
__init__.
VocabParallelLogprobOpholds no instance state. The explicit__init__withpassadds nothing over the default.♻️ Proposed removal
op_class = "logprob" is_batch_invariant = True - def __init__(self) -> None: - pass - def __call__(🤖 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 346 - 347, Remove the empty __init__ method from VocabParallelLogprobOp and rely on Python’s default constructor, since the class holds no instance state.
309-337: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winBackward assumes
grad_logpandgrad_lseare identical on every TP rank.The backward pass is fully elementwise and launches no collective. That is correct only when the incoming output gradients are bitwise identical on every rank, which holds because
selected_logpandlseare replicated. If a caller applies a rank-dependent reduction to the outputs, each rank produces a differentgrad, and the parameter gradients silently diverge across the TP group. Nothing in the op detects this.State the requirement in the docstring, and consider adding a debug-only preflight that all-gathers a digest of
grad_logpandgrad_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/pytorch/loss/vocab_parallel_logp.py` around lines 309 - 337, Document in the relevant operation docstring that backward requires grad_logp and grad_lse to be identical across all tensor-parallel ranks because no collective synchronizes them. Add an optional debug-only preflight in backward, using the project’s existing TP communication and validation utilities if available, to all-gather and compare digests of each non-None incoming gradient; report a clear error when ranks disagree while leaving the normal elementwise gradient path unchanged.tests/test_vocab_parallel_logp.py (1)
179-192: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRemove the dead
contractassignment at line 181.Line 181 builds a contract with
padded_vocab=REAL_VOCAB + 5, and line 183 overwrites it before any use. Only the line 183 contract is exercised. Delete line 181 so the test states one intent.♻️ 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)🤖 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 unused first contract assignment in test_matches_ws1_batch_invariant_logp_within_contract_tolerance, keeping only the subsequent _contract(real_vocab=PADDED_VOCAB, padded_vocab=PADDED_VOCAB) assignment used by the test.rl_engine/kernels/logprob_contract.py (1)
120-143: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRename the
fieldparameter to avoid shadowingdataclasses.field.The module imports
fieldfromdataclassesat line 26. The helpers_enum_value,_positive_int,_non_negative_int, and_plain_intbindfieldto astrparameter. The same shadowing exists inLogprobBackendCapability._validated_world_sizesat line 553. There is no current defect, because none of these functions calldataclasses.field. A future edit inside one of them would silently use the string.♻️ Proposed rename
-def _enum_value(enum_type: type[_EnumT], value: Any, field: str) -> _EnumT: +def _enum_value(enum_type: type[_EnumT], value: Any, field_name: str) -> _EnumT: try: return enum_type(value) except (TypeError, ValueError) as exc: allowed = ", ".join(item.value for item in enum_type) - raise LogprobContractError(f"{field} must be one of: {allowed}; got {value!r}") from exc + raise LogprobContractError(f"{field_name} must be one of: {allowed}; got {value!r}") from exc -def _positive_int(value: Any, field: str) -> int: +def _positive_int(value: Any, field_name: str) -> int: if isinstance(value, bool) or not isinstance(value, int) or value <= 0: - raise LogprobContractError(f"{field} must be a positive integer; got {value!r}") + raise LogprobContractError(f"{field_name} must be a positive integer; got {value!r}") return value🤖 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/logprob_contract.py` around lines 120 - 143, Rename the field parameter in _enum_value, _positive_int, _non_negative_int, _plain_int, and LogprobBackendCapability._validated_world_sizes to avoid shadowing the imported dataclasses.field, and update all corresponding references and call sites within those implementations.rl_engine/kernels/registry.py (4)
635-640: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winLog the traceback for instantiation failures.
The broad
except Exceptionis appropriate here, because a failing backend constructor must not abort dispatch. Ruff flags it as BLE001; suppress the rule with a reason instead of narrowing the catch. Also record the traceback.logger.error(f"...: {exc}")discards the stack, and backend constructors fail deep inside Triton or CUDA initialization.♻️ Proposed logging change
try: op = op_class() - except Exception as exc: - logger.error(f"Failed to instantiate {backend.name}: {exc}") + except Exception as exc: # noqa: BLE001 - any backend init failure must fall through + logger.error("Failed to instantiate %s: %s", backend.name, exc, exc_info=True) self._failed_backends.add(backend.name) return None🤖 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 635 - 640, Update the backend instantiation path around op_class() in the registry dispatch flow to keep the broad Exception catch, add a Ruff BLE001 suppression with a short reason on that block, and change the logger.error call so it records the full traceback instead of only the exception message. Preserve the existing behavior of marking backend.name as failed and returning None when instantiation fails.Source: Linters/SAST tools
541-542: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win
get_logprob_opcannot target a non-default platform.
get_logprob_opresolves the platform throughself._platform(), which is_platform_for_device(None)and therefore always the process default fromdevice_ctx. Legacyget_opaccepts adeviceargument and resolves the platform per call. A WS2 caller that runs the contract on a device other than the process default receives the wrong candidate list, andprovenance["platform"]records the wrong platform.Add an optional
deviceparameter that mirrorsget_op.♻️ Proposed signature change
def get_logprob_op( self, contract: LogprobContract, *, requested_backend: str = "auto", + device: torch.device | str | None = None, ) -> LogprobDispatchResult:- platform = self._platform() + platform = self._platform_for_device(device)🤖 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 541 - 542, Update get_logprob_op to accept an optional device parameter matching get_op, resolve the platform through the provided device rather than always using self._platform(), and use that resolved platform for candidate selection and provenance["platform"]. Preserve default-platform behavior when device is omitted.
600-620: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueA backend id that collides with a policy keyword is unreachable.
_logprob_policy_mismatchcheckspolicy in IMPLEMENTATION_KINDSbefore the exactbackend_idcomparison. If a backend registersbackend_id="reference"orbackend_id="auto", the caller can never select it by id, because the keyword branch consumes the string first. Theautocase is worse: it matches every candidate.Reject those reserved values in
register_logprob_backendso the conflict fails at registration instead of silently changing selection.♻️ Proposed guard in `register_logprob_backend`
if not isinstance(capability, LogprobBackendCapability): raise LogprobContractError("capability must be a LogprobBackendCapability") + reserved = {"auto"} | set(IMPLEMENTATION_KINDS) + if capability.backend_id.lower() in reserved: + raise LogprobContractError( + f"backend_id={capability.backend_id!r} collides with a dispatch policy " + f"keyword; expected an id outside {sorted(reserved)}" + )🤖 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 600 - 620, Update register_logprob_backend to reject backend registrations whose backend_id is the reserved "auto" value or any value in IMPLEMENTATION_KINDS, raising the established registration validation error before storing the backend. Keep _logprob_policy_mismatch unchanged and preserve valid backend registration behavior.
372-378: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winConsider the default ordering of the reference backend.
prepend=TrueplacesPYTORCH_VOCAB_PARALLEL_LOGPat the front of every platform candidate list. Its capability is a superset of the WS1 descriptors, sorequested_backend="auto"selects the pure-PyTorch reference even for contracts that the Triton or CUDA production backends can serve (TP=1, ignore-index masking, no LSE export,fixed_topology).implementation_kind="reference"is recorded, but nothing inget_logprob_opprefersproductionunderauto.If the intent is that
autoprefers production kernels, register the reference withprepend=False, or rank candidates byimplementation_kindinget_logprob_op. If the intent is that WS2 dispatch always resolves to the reference for now, state that in the comment above the loop.🤖 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 372 - 378, The WS2 reference backend is prepended ahead of production candidates, causing automatic dispatch to select the pure-PyTorch implementation when production kernels are compatible. Update the registration in the loop over self._priority_map to append this backend instead, or adjust get_logprob_op to rank production candidates before reference candidates under requested_backend="auto"; if reference-first dispatch is intentional, document that policy above the loop.
🤖 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 @.github/workflows/ci.yml:
- Around line 82-83: Add a GPU-backed CI job alongside the existing CPU-safe
“Run WS2 Vocab-Parallel Logprob Tests” step, using a runner and PyTorch
installation that support CUDA/NCCL. Configure it to execute the distributed
TP=2 and TP=4 cases in tests/test_vocab_parallel_logp.py, ensuring CI validates
cross-TP bitwise equality rather than skipping those tests.
In `@docs/operators/batch-invariant-logp.md`:
- Around line 57-72: Update the Tensor Parallel documentation around
VocabParallelLogprobOp to state that the bit-identical TP=1/2/4 guarantee
requires the same num_vocab_tiles across compared runs, agreement on that value
across ranks, and tile-aligned shard bounds. Also add punctuation or line breaks
so the operator name, source path, and claim are clearly separated.
In `@rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py`:
- Around line 193-202: Update the tile partial handling in the visible loop and
`_merge_tile_partials` so NaN maxima are not treated like all-`-inf` tiles:
distinguish the exact all-`-inf` case from non-finite NaN values, preserve zero
partials only for all-`-inf`, and allow NaN to propagate into `lse` for the
existing `validate=True` check.
In `@tests/test_vocab_parallel_logp.py`:
- Around line 471-481: Update the process cleanup in the finally block and the
queue.Empty handling path to call process.join() again after terminate() for any
worker still alive. Ensure each process is reaped before the later exitcode
assertions in the result-collection flow.
---
Nitpick comments:
In `@rl_engine/kernels/logprob_contract.py`:
- Around line 120-143: Rename the field parameter in _enum_value, _positive_int,
_non_negative_int, _plain_int, and
LogprobBackendCapability._validated_world_sizes to avoid shadowing the imported
dataclasses.field, and update all corresponding references and call sites within
those implementations.
In `@rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py`:
- Around line 346-347: Remove the empty __init__ method from
VocabParallelLogprobOp and rely on Python’s default constructor, since the class
holds no instance state.
- Around line 309-337: Document in the relevant operation docstring that
backward requires grad_logp and grad_lse to be identical across all
tensor-parallel ranks because no collective synchronizes them. Add an optional
debug-only preflight in backward, using the project’s existing TP communication
and validation utilities if available, to all-gather and compare digests of each
non-None incoming gradient; report a clear error when ranks disagree while
leaving the normal elementwise gradient path unchanged.
In `@rl_engine/kernels/registry.py`:
- Around line 635-640: Update the backend instantiation path around op_class()
in the registry dispatch flow to keep the broad Exception catch, add a Ruff
BLE001 suppression with a short reason on that block, and change the
logger.error call so it records the full traceback instead of only the exception
message. Preserve the existing behavior of marking backend.name as failed and
returning None when instantiation fails.
- Around line 541-542: Update get_logprob_op to accept an optional device
parameter matching get_op, resolve the platform through the provided device
rather than always using self._platform(), and use that resolved platform for
candidate selection and provenance["platform"]. Preserve default-platform
behavior when device is omitted.
- Around line 600-620: Update register_logprob_backend to reject backend
registrations whose backend_id is the reserved "auto" value or any value in
IMPLEMENTATION_KINDS, raising the established registration validation error
before storing the backend. Keep _logprob_policy_mismatch unchanged and preserve
valid backend registration behavior.
- Around line 372-378: The WS2 reference backend is prepended ahead of
production candidates, causing automatic dispatch to select the pure-PyTorch
implementation when production kernels are compatible. Update the registration
in the loop over self._priority_map to append this backend instead, or adjust
get_logprob_op to rank production candidates before reference candidates under
requested_backend="auto"; if reference-first dispatch is intentional, document
that policy above the loop.
In `@tests/test_vocab_parallel_logp.py`:
- Around line 179-192: Remove the unused first contract assignment in
test_matches_ws1_batch_invariant_logp_within_contract_tolerance, keeping only
the subsequent _contract(real_vocab=PADDED_VOCAB, padded_vocab=PADDED_VOCAB)
assignment used by the test.
🪄 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: c54a23c6-918c-4648-a9fe-62b2f0661e24
📒 Files selected for processing (8)
.github/workflows/ci.ymldocs/design/runtime-dispatch.mddocs/operators/batch-invariant-logp.mdrl_engine/kernels/logprob_contract.pyrl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.pyrl_engine/kernels/registry.pytests/test_logprob_contract.pytests/test_vocab_parallel_logp.py
| - name: Run WS2 Vocab-Parallel Logprob Tests (CPU-safe) | ||
| run: python -m pytest tests/test_vocab_parallel_logp.py -v |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Run multi-GPU logprob tests in CI.
This job installs the CPU-only PyTorch wheel and uses ubuntu-latest. The TP=2 and TP=4 NCCL tests skip on this runner. Add a GPU-backed CI job that runs the distributed test cases. Otherwise, CI does not validate the cross-TP bitwise-equality requirement.
🤖 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 @.github/workflows/ci.yml around lines 82 - 83, Add a GPU-backed CI job
alongside the existing CPU-safe “Run WS2 Vocab-Parallel Logprob Tests” step,
using a runner and PyTorch installation that support CUDA/NCCL. Configure it to
execute the distributed TP=2 and TP=4 cases in
tests/test_vocab_parallel_logp.py, ensuring CI validates cross-TP bitwise
equality rather than skipping those tests.
| ## Tensor Parallel | ||
|
|
||
| `VocabParallelLogprobOp` | ||
| (`rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py`) | ||
| **TP=1, TP=2, and TP=4 produce bit-identical results.** | ||
|
|
||
| 1. Split the padded vocabulary into `num_vocab_tiles` fixed tiles. | ||
| 2. Each rank computes fp32 `(max, sumexp)` for the tiles it owns. Every tile | ||
| is reduced as the same contiguous `[n, tile]` shape, on any rank. | ||
| 3. All tile partials are shared with `all_gather`. The collective only moves | ||
| bytes; it never does math, so it cannot round anything. | ||
| 4. Every rank merges all tiles in the same fixed order, over the same | ||
| `[n, num_vocab_tiles]` shape. `LSE = M + log(sum(s_t * exp(m_t - M)))`. | ||
| 5. The target logit is copied from the rank that owns it (never summed). | ||
| 6. `logp = target_logit - LSE`. Inactive rows become `0.0`. | ||
|
|
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
State the num_vocab_tiles precondition with the bit-identical claim.
Line 61 asserts that TP=1, TP=2, and TP=4 produce bit-identical results. The implementation provides that guarantee only while num_vocab_tiles is held fixed across the compared TP degrees and agreed on by every rank. The module docstring of vocab_parallel_logp.py states this explicitly, and _preflight_cross_rank_agreement enforces the cross-rank half. A reader of this section alone can change num_vocab_tiles between two runs and lose the property. Also add the shard-alignment requirement, since _tile_size rejects non-tile-aligned bounds.
Lines 59 to 61 also render as a single paragraph, so the operator name, the file path, and the claim run together without punctuation.
📝 Proposed wording
## Tensor Parallel
-`VocabParallelLogprobOp`
-(`rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py`)
-**TP=1, TP=2, and TP=4 produce bit-identical results.**
+`VocabParallelLogprobOp`
+(`rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py`) computes the
+selected-token logprob across vocab-parallel TP ranks.
+**TP=1, TP=2, and TP=4 produce bit-identical results**, provided
+`num_vocab_tiles` is the same value at every TP degree and on every rank, and
+every shard boundary is tile-aligned. Both conditions fail loudly.📝 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.
| ## Tensor Parallel | |
| `VocabParallelLogprobOp` | |
| (`rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py`) | |
| **TP=1, TP=2, and TP=4 produce bit-identical results.** | |
| 1. Split the padded vocabulary into `num_vocab_tiles` fixed tiles. | |
| 2. Each rank computes fp32 `(max, sumexp)` for the tiles it owns. Every tile | |
| is reduced as the same contiguous `[n, tile]` shape, on any rank. | |
| 3. All tile partials are shared with `all_gather`. The collective only moves | |
| bytes; it never does math, so it cannot round anything. | |
| 4. Every rank merges all tiles in the same fixed order, over the same | |
| `[n, num_vocab_tiles]` shape. `LSE = M + log(sum(s_t * exp(m_t - M)))`. | |
| 5. The target logit is copied from the rank that owns it (never summed). | |
| 6. `logp = target_logit - LSE`. Inactive rows become `0.0`. | |
| ## Tensor Parallel | |
| `VocabParallelLogprobOp` | |
| (`rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py`) computes the | |
| selected-token logprob across vocab-parallel TP ranks. | |
| **TP=1, TP=2, and TP=4 produce bit-identical results**, provided | |
| `num_vocab_tiles` is the same value at every TP degree and on every rank, and | |
| every shard boundary is tile-aligned. Both conditions fail loudly. | |
| 1. Split the padded vocabulary into `num_vocab_tiles` fixed tiles. | |
| 2. Each rank computes fp32 `(max, sumexp)` for the tiles it owns. Every tile | |
| is reduced as the same contiguous `[n, tile]` shape, on any rank. | |
| 3. All tile partials are shared with `all_gather`. The collective only moves | |
| bytes; it never does math, so it cannot round anything. | |
| 4. Every rank merges all tiles in the same fixed order, over the same | |
| `[n, num_vocab_tiles]` shape. `LSE = M + log(sum(s_t * exp(m_t - M)))`. | |
| 5. The target logit is copied from the rank that owns it (never summed). | |
| 6. `logp = target_logit - LSE`. Inactive rows become `0.0`. |
🤖 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 57 - 72, Update the
Tensor Parallel documentation around VocabParallelLogprobOp to state that the
bit-identical TP=1/2/4 guarantee requires the same num_vocab_tiles across
compared runs, agreement on that value across ranks, and tile-aligned shard
bounds. Also add punctuation or line breaks so the operator name, source path,
and claim are clearly separated.
| for tile_index in range(local_vocab // tile): | ||
| block = z_masked[:, tile_index * tile : (tile_index + 1) * tile].contiguous() | ||
| m_t = block.max(dim=-1).values | ||
| finite = m_t > float("-inf") | ||
| m_safe = torch.where(finite, m_t, torch.zeros_like(m_t)) | ||
| s_t = (block - m_safe.unsqueeze(-1)).exp().sum(dim=-1) | ||
| s_t = torch.where(finite, s_t, torch.zeros_like(s_t)) | ||
| m_parts.append(m_t) | ||
| s_parts.append(s_t) | ||
| return torch.stack(m_parts, dim=1), torch.stack(s_parts, dim=1) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
A NaN logit is silently dropped instead of propagating.
m_t = block.max(dim=-1).values returns NaN if the tile holds a NaN. Line 196 evaluates finite = m_t > float("-inf"), which is False for NaN. Line 199 then forces s_t = 0, so the whole tile becomes the identity partial and the NaN never reaches lse. The module docstring states that logits over the real vocabulary must be finite and that validate=True fails loudly on a non-finite lse. That guard cannot fire here, because the NaN is removed before lse is formed. The result is a finite, wrong logprob on a corrupted row.
Distinguish "all -inf" from "not a number":
🐛 Proposed fix
m_t = block.max(dim=-1).values
- finite = m_t > float("-inf")
+ # NaN must stay NaN; only an all-(-inf) tile takes the identity partial.
+ finite = torch.isfinite(m_t) | m_t.isnan()
m_safe = torch.where(finite, m_t, torch.zeros_like(m_t))
s_t = (block - m_safe.unsqueeze(-1)).exp().sum(dim=-1)
s_t = torch.where(finite, s_t, torch.zeros_like(s_t))With that change a NaN tile yields s_t = NaN, the NaN reaches lse, and the validate=True check at line 396 rejects the row. Apply the same reasoning to _merge_tile_partials at line 274.
🤖 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 193 -
202, Update the tile partial handling in the visible loop and
`_merge_tile_partials` so NaN maxima are not treated like all-`-inf` tiles:
distinguish the exact all-`-inf` case from non-finite NaN values, preserve zero
partials only for all-`-inf`, and allow NaN to propagate into `lse` for the
existing `validate=True` check.
| 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
Join each process after terminate() before reading exitcode.
The finally block calls process.join(timeout=30) and then process.terminate() for a process that is still alive. It does not join again. Process.exitcode is None until the process is reaped, so line 480 can report assert None == 0 for a hung worker instead of the timeout or traceback that caused the hang. Line 469 has the same gap in the queue.Empty path.
🐛 Proposed fix
finally:
for process in processes:
process.join(timeout=30)
if process.is_alive():
process.terminate()
+ process.join(timeout=30)📝 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.
| 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 | |
| finally: | |
| for process in processes: | |
| process.join(timeout=30) | |
| if process.is_alive(): | |
| process.terminate() | |
| process.join(timeout=30) | |
| 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 |
🤖 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 finally block and the queue.Empty handling path to call
process.join() again after terminate() for any worker still alive. Ensure each
process is reaped before the later exitcode assertions in the result-collection
flow.
2 participants
Overview
This pull request implements PR 3 of issue #241 . Outputs, LSE, and gradients are bit-identical across TP=1/2/4.
Stacked on #259; only the top commit is new.
Key Changes
Reference Op (rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py)
VocabParallelLogprobOp: fp32 (selected_logp, lse), replicated per LogprobOutputSpec, differentiable through both outputsTesting
test_invalid_invocations_fail_loudlynum_vocab_tiles, out-of-real-vocab active target, all--infactive rowTestSingleRankTestBackwardTestCrossTPBitwiseselected_logp,lse, and the local gradient shard — across even/uneven tile-aligned shards and fp32/bf16, plus rerun stability under collectives and identical output bits on every rankTestCrossTPGuardsnum_vocab_tilesdisagreement and misaligned shard bounds must abort loudly on every rank, so a disagreement never strands some ranks inside a collective26 Tests Passed. Each cross-TP worker computes the same batch twice — once sharded
across the TP group, once at TP=1 on the full unsharded logits in-process — and compares
the two bit for bit, forward and backward.
Summary by CodeRabbit
New Features
Documentation
Tests