fixed s_vals divergence across GPUs - #69
Open
Mustafa-cuda-dev wants to merge 1 commit into
Open
Conversation
Contributor
WalkthroughChangesGRAIL proof computation
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant compute_proofs
participant forward_single_layer
participant _project_fixed_point
participant verifier.create_commitments_batch
participant torch.log_softmax
compute_proofs->>forward_single_layer: produce hidden states and logits
compute_proofs->>_project_fixed_point: project hidden states with randomness
_project_fixed_point-->>compute_proofs: return projected values
compute_proofs->>verifier.create_commitments_batch: create commitments
compute_proofs->>torch.log_softmax: calculate chunked logprobs
torch.log_softmax-->>compute_proofs: return selected token logprobs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
grail/environments/proofs.py (1)
203-443: 📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift
compute_proofsis ~240 lines, far beyond the 25-line guideline.The function now spans forward-pass fallback, deterministic-projection selection, commitment generation, and chunked logprob computation. As per coding guidelines, "functions exceeding 25 lines require refactoring justification"; each phase (projection selection, commitment call, logprob computation) is a good candidate for extraction into a private helper, which would also make the pieces independently testable.
🤖 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 `@grail/environments/proofs.py` around lines 203 - 443, Refactor compute_proofs into focused private helpers so it no longer contains the full forward-pass fallback, projection selection, commitment generation, and chunked logprob logic. Extract the projection/commitment path and logprob calculation at minimum, while preserving existing behavior, inputs, outputs, and error handling; keep compute_proofs responsible for orchestration and result assembly.Source: Coding guidelines
🧹 Nitpick comments (5)
grail/environments/proofs.py (5)
47-49: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winReplace bare
assertwith explicit validation.
assert r_vec.dim() in (1, 2), ...and the follow-up singleton-dim assert are stripped when Python runs with-O, silently disabling the input-shape guard in production.🛡️ Suggested fix
- assert r_vec.dim() in (1, 2), "r_vec must be 1D or 2D (batch x hidden)" - if r_vec.dim() == 2: - assert r_vec.size(0) == 1 or r_vec.size(1) == 1, "r_vec 2D must have a singleton dimension of size 1" + if r_vec.dim() not in (1, 2): + raise ValueError("r_vec must be 1D or 2D (batch x hidden)") + if r_vec.dim() == 2 and r_vec.size(0) != 1 and r_vec.size(1) != 1: + raise ValueError("r_vec 2D must have a singleton dimension of size 1")🤖 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 `@grail/environments/proofs.py` around lines 47 - 49, Replace both shape-checking assert statements in the r_vec validation logic with explicit runtime validation that raises an appropriate exception when r_vec is not 1D/2D or a 2D tensor lacks a singleton dimension. Preserve the existing validation messages and allow valid shapes to continue unchanged.Source: Coding guidelines
379-421: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winChunked logprob GPU/CPU paths are near-duplicate.
Both branches replicate the same chunk-index bookkeeping,
log_softmax, gather, andLOGPROB_CHUNK = 512constant, differing only in device placement of the token tensor. Extract a shared helper parameterized by device to avoid drift between the two copies on this correctness-sensitive path.♻️ Suggested direction
+LOGPROB_CHUNK = 512 + + +def _chunked_completion_logprobs( + logits: torch.Tensor, valid_token_ids: list[int], device: torch.device +) -> list[float]: + token_tensor = torch.tensor(valid_token_ids, dtype=torch.long, device=device) + n_valid = len(valid_token_ids) + chunk_logprobs: list[float] = [] + for c_start in range(0, n_valid, LOGPROB_CHUNK): + c_end = min(c_start + LOGPROB_CHUNK, n_valid) + log_probs_chunk = torch.log_softmax(logits[c_start:c_end].float(), dim=-1) + selected = log_probs_chunk[ + torch.arange(c_end - c_start, device=device), token_tensor[c_start:c_end] + ] + chunk_logprobs.extend(selected.tolist()) + return chunk_logprobs🤖 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 `@grail/environments/proofs.py` around lines 379 - 421, Refactor the duplicated chunked logprob computation into a shared helper near the surrounding proof logic, parameterized by the target device for the token tensor and index operations. Have both the gpu_logprobs/logits.is_cuda and CPU branches call this helper while preserving the existing chunk size, slicing, log_softmax, token gathering, and returned chunk_logprobs behavior.
45-116: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoff
_project_fixed_pointexceeds the 25-line function-length guideline.The function is ~70 lines with three near-duplicate chunking branches (1D r_vec, 2D-batch r_vec, 2D-transposed r_vec) both chunked and non-chunked. As per coding guidelines, "functions exceeding 25 lines require refactoring justification." Consider factoring the shape-dispatch matmul into a small helper reused by both the chunked and non-chunked paths.
🤖 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 `@grail/environments/proofs.py` around lines 45 - 116, Refactor _project_fixed_point to stay within the 25-line guideline by extracting the shape-dispatch and chunked/non-chunked matmul logic into a small helper. Reuse that helper for 1D vectors, singleton-row 2D vectors, and transposed 2D batches, preserving the current output shapes and numerical behavior while leaving quantization and de-quantization in _project_fixed_point.Source: Coding guidelines
298-298: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winHoist per-call config/mode resolution out of the per-sequence loop.
_get_env_config()(line 298) and thetorch.use_deterministic_algorithms(True)/finallyrestore (lines 304-327) are invoked once per sequence in the batch, not once percompute_proofscall. Besides the redundantos.getenvparsing, an invalid env value re-logs a warning for every sequence, and toggling global deterministic-algorithm state repeatedly on a hot path adds unnecessary overhead. Resolvequant_bits/deterministic_modeonce before thefor idx, all_token_ids in enumerate(...)loop.♻️ Suggested fix
verifier = GRAILVerifier(hidden_dim=hidden_dim) r_vec = verifier.generate_r_vec(randomness_hex) + quant_bits, deterministic_mode = _get_env_config()- # deterministic projection - quant_bits, deterministic_mode = _get_env_config() - s_vals = None + # deterministic projection + s_vals: torch.Tensor | None = NoneAlso applies to: 304-327
🤖 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 `@grail/environments/proofs.py` at line 298, Update compute_proofs to call _get_env_config once before the for idx, all_token_ids loop, then reuse quant_bits and deterministic_mode for every sequence. Move the torch.use_deterministic_algorithms(True) setup and its finally-based restoration outside the loop so global deterministic state is toggled once per compute_proofs call while preserving restoration on exit.Source: Coding guidelines
304-327: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the
fp32_deterministicshape dispatch out of the try block.grail/environments/proofs.py:303-327Thetrystill wraps the 1D/2D branching and squeeze logic; move that into a helper so thetryonly covers thematmulcall that can raise. If this path is expected to stay deterministic on CUDA, keep the fallback explicit: withoutCUBLAS_WORKSPACE_CONFIGset before CUDA starts, theRuntimeErrorpath can silently degrade to ordinary float32.🤖 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 `@grail/environments/proofs.py` around lines 304 - 327, Extract the r_vec dimensionality dispatch and squeeze logic from the deterministic-algorithm try block into a helper dedicated to the fp32_deterministic projection path. Keep the try/finally narrowly scoped around the matmul operation that may raise, preserve the existing shapes and device conversion, and retain an explicit fallback when deterministic CUDA matmul fails, including the required CUBLAS_WORKSPACE_CONFIG setup before CUDA initialization.Source: Coding guidelines
🤖 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 `@grail/environments/proofs.py`:
- Around line 297-347: The deterministic projection path in the proof-generation
flow requires projected_s_vals support that
GRAILVerifier.create_commitments_batch currently lacks, causing the default
fixed_point mode to raise. Update GRAILVerifier.create_commitments_batch and its
implementation to accept and use the optional projected_s_vals argument, while
preserving the existing h_layer/r_vec behavior when it is not provided.
---
Outside diff comments:
In `@grail/environments/proofs.py`:
- Around line 203-443: Refactor compute_proofs into focused private helpers so
it no longer contains the full forward-pass fallback, projection selection,
commitment generation, and chunked logprob logic. Extract the
projection/commitment path and logprob calculation at minimum, while preserving
existing behavior, inputs, outputs, and error handling; keep compute_proofs
responsible for orchestration and result assembly.
---
Nitpick comments:
In `@grail/environments/proofs.py`:
- Around line 47-49: Replace both shape-checking assert statements in the r_vec
validation logic with explicit runtime validation that raises an appropriate
exception when r_vec is not 1D/2D or a 2D tensor lacks a singleton dimension.
Preserve the existing validation messages and allow valid shapes to continue
unchanged.
- Around line 379-421: Refactor the duplicated chunked logprob computation into
a shared helper near the surrounding proof logic, parameterized by the target
device for the token tensor and index operations. Have both the
gpu_logprobs/logits.is_cuda and CPU branches call this helper while preserving
the existing chunk size, slicing, log_softmax, token gathering, and returned
chunk_logprobs behavior.
- Around line 45-116: Refactor _project_fixed_point to stay within the 25-line
guideline by extracting the shape-dispatch and chunked/non-chunked matmul logic
into a small helper. Reuse that helper for 1D vectors, singleton-row 2D vectors,
and transposed 2D batches, preserving the current output shapes and numerical
behavior while leaving quantization and de-quantization in _project_fixed_point.
- Line 298: Update compute_proofs to call _get_env_config once before the for
idx, all_token_ids loop, then reuse quant_bits and deterministic_mode for every
sequence. Move the torch.use_deterministic_algorithms(True) setup and its
finally-based restoration outside the loop so global deterministic state is
toggled once per compute_proofs call while preserving restoration on exit.
- Around line 304-327: Extract the r_vec dimensionality dispatch and squeeze
logic from the deterministic-algorithm try block into a helper dedicated to the
fp32_deterministic projection path. Keep the try/finally narrowly scoped around
the matmul operation that may raise, preserve the existing shapes and device
conversion, and retain an explicit fallback when deterministic CUDA matmul
fails, including the required CUBLAS_WORKSPACE_CONFIG setup before CUDA
initialization.
🪄 Autofix (Beta)
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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 9d7c100b-f7db-408d-b9a5-9067fc2df897
📒 Files selected for processing (1)
grail/environments/proofs.py
Comment on lines
+297
to
+347
| # deterministic projection | ||
| quant_bits, deterministic_mode = _get_env_config() | ||
| s_vals = None | ||
|
|
||
| if deterministic_mode == "fixed_point": | ||
| s_vals = _project_fixed_point(h_layer, r_vec, quant_bits) | ||
| elif deterministic_mode == "fp32_deterministic": | ||
| prev_deterministic = torch.are_deterministic_algorithms_enabled() | ||
| try: | ||
| torch.use_deterministic_algorithms(True) | ||
| if r_vec.dim() == 1: | ||
| s_vals = torch.matmul(h_layer, r_vec.to(h_layer.device)) | ||
| elif r_vec.dim() == 2: | ||
| if r_vec.size(0) == 1: | ||
| s_vals = torch.matmul(h_layer, r_vec.to(h_layer.device).t()) | ||
| else: | ||
| s_vals = torch.matmul(h_layer, r_vec.to(h_layer.device)) | ||
| else: | ||
| s_vals = torch.matmul(h_layer, r_vec.to(h_layer.device).view(-1)) | ||
|
|
||
| if r_vec.dim() == 1 and s_vals.dim() == 2: | ||
| s_vals = s_vals.squeeze(-1) | ||
| except RuntimeError as e: | ||
| logger.warning( | ||
| "Deterministic fp32 algorithms not fully supported on this platform: %s. " | ||
| "Falling back to standard float32 projection.", | ||
| e, | ||
| ) | ||
| s_vals = None | ||
| finally: | ||
| torch.use_deterministic_algorithms(prev_deterministic) | ||
|
|
||
| # commitments | ||
| if s_vals is not None: | ||
| # check signature compatibility | ||
| has_projected_arg = False | ||
| try: | ||
| sig = inspect.signature(verifier.create_commitments_batch) | ||
| if "projected_s_vals" in sig.parameters: | ||
| has_projected_arg = True | ||
| except (ValueError, TypeError): | ||
| has_projected_arg = False | ||
|
|
||
| if has_projected_arg: | ||
| commitments = verifier.create_commitments_batch(h_layer, r_vec, projected_s_vals=s_vals) | ||
| else: | ||
| raise RuntimeError( | ||
| "Verifier must support 'projected_s_vals' argument. Please update GRAILVerifier." | ||
| ) | ||
| else: | ||
| commitments = verifier.create_commitments_batch(h_layer, r_vec) |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔴 Critical | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n "def create_commitments_batch" -A 8 grail/protocol/grail_verifier.py
rg -n "projected_s_vals" -g '*.py'Repository: one-covenant/grail
Length of output: 601
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== proofs.py around env config and compute_proofs =="
sed -n '1,80p' grail/environments/proofs.py
echo
sed -n '290,350p' grail/environments/proofs.py
echo
echo "== verifier signature and usages =="
sed -n '248,290p' grail/protocol/grail_verifier.py
echo
rg -n "projected_s_vals|create_commitments_batch\(" grail -g '*.py'Repository: one-covenant/grail
Length of output: 8208
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("grail/environments/proofs.py")
text = p.read_text()
for needle in ["GRAIL_PROOF_DETERMINISTIC_MODE", "fixed_point", "fp32_deterministic", "projected_s_vals", "create_commitments_batch"]:
print(f"FOUND {needle}: {needle in text}")
PYRepository: one-covenant/grail
Length of output: 320
Default deterministic mode breaks proof generation
GRAIL_PROOF_DETERMINISTIC_MODE defaults to "fixed_point", so s_vals is always set here. But GRAILVerifier.create_commitments_batch() still only accepts (h_layer, r_vec), so this branch always raises RuntimeError instead of creating commitments. Either add projected_s_vals support to the verifier or stop requiring it in this path.
🤖 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 `@grail/environments/proofs.py` around lines 297 - 347, The deterministic
projection path in the proof-generation flow requires projected_s_vals support
that GRAILVerifier.create_commitments_batch currently lacks, causing the default
fixed_point mode to raise. Update GRAILVerifier.create_commitments_batch and its
implementation to accept and use the optional projected_s_vals argument, while
preserving the existing h_layer/r_vec behavior when it is not provided.
Author
There was a problem hiding this comment.
Fixed in #70 — added projected_s_vals parameter to verifier.
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
repo="one-covenant/grail"
echo "== Pull request `#70` metadata =="
gh pr view 70 --repo "$repo" --json number,state,mergedAt,url,title
echo
echo "== Relevant verifier changes in `#70` =="
gh pr diff 70 --repo "$repo" -- grail/protocol/grail_verifier.pyLength of output: 6612
The initial query confirmed that #70 is open, but its diff filter was unsupported. I’ll retrieve the changed-file patch directly to confirm the parameter is actually consumed.
#!/bin/bash
set -euo pipefail
gh api --paginate \
"repos/one-covenant/grail/pulls/70/files?per_page=100" \
--jq '.[] | select(.filename == "grail/protocol/grail_verifier.py") | .patch'
---
<details>
<summary>🧬 Code Graph Analysis Results</summary>
### Relevant snippets
#### `grail/environments/proofs.py` (lines **296-346**)
```py
# deterministic projection
quant_bits, deterministic_mode = _get_env_config()
s_vals = None
if deterministic_mode == "fixed_point":
s_vals = _project_fixed_point(h_layer, r_vec, quant_bits)
elif deterministic_mode == "fp32_deterministic":
prev_deterministic = torch.are_deterministic_algorithms_enabled()
try:
torch.use_deterministic_algorithms(True)
if r_vec.dim() == 1:
s_vals = torch.matmul(h_layer, r_vec.to(h_layer.device))
elif r_vec.dim() == 2:
if r_vec.size(0) == 1:
s_vals = torch.matmul(h_layer, r_vec.to(h_layer.device).t())
else:
s_vals = torch.matmul(h_layer, r_vec.to(h_layer.device))
else:
s_vals = torch.matmul(h_layer, r_vec.to(h_layer.device).view(-1))
if r_vec.dim() == 1 and s_vals.dim() == 2:
s_vals = s_vals.squeeze(-1)
except RuntimeError as e:
logger.warning(
"Deterministic fp32 algorithms not fully supported on this platform: %s. "
"Falling back to standard float32 projection.",
e,
)
s_vals = None
finally:
torch.use_deterministic_algorithms(prev_deterministic)
# commitments
if s_vals is not None:
# check signature compatibility
has_projected_arg = False
try:
sig = inspect.signature(verifier.create_commitments_batch)
if "projected_s_vals" in sig.parameters:
has_projected_arg = True
except (ValueError, TypeError):
has_projected_arg = False
if has_projected_arg:
commitments = verifier.create_commitments_batch(h_layer, r_vec, projected_s_vals=s_vals)
else:
raise RuntimeError(
"Verifier must support 'projected_s_vals' argument. Please update GRAILVerifier."
)
else:
commitments = verifier.create_commitments_batch(h_layer, r_vec)grail/environments/proofs.py (lines 45-116)
def _project_fixed_point(h_layer: torch.Tensor, r_vec: torch.Tensor, bits: int) -> torch.Tensor:
"""Project using fixed-point integer arithmetic via float64."""
...(Summarized)
- Parameters:
h_layer: hidden states tensor (expected shape[seq_len, hidden_dim])r_vec: coefficient vector tensor (expected dim 1 or 2; asserts/special-cases singleton dimension)bits: quantization bit width (used to set fixed-point integer range)
- Returns: projected values
s_valsastorch.Tensor(de-quantized back to float32; may squeeze to match expected shape whenr_vec.dim()==1) - Implementation details: quantizes
h_layerandr_vecto int types, converts to float64 for chunkedtorch.matmul, then de-quantizes using derived scales; usestorch.no_grad(), selects chunk size based onhidden_dim.
grail/protocol/grail_verifier.py (lines 186-219)
def generate_r_vec(self, randomness_hex: str) -> torch.Tensor:
"""Generate small bounded coefficient vector from randomness.
Uses tiny coefficients in [-127, 127] to reduce sensitivity to
floating-point variations while maintaining cryptographic security.
Args:
randomness_hex: Hex string of beacon randomness
Returns:
Tensor of shape [topk] with int8 coefficients in [-R, R]
"""
...(Included fully as provided in snippet)
grail/protocol/grail_verifier.py (lines 256-298)
def create_commitments_batch(self, h_layer: torch.Tensor, r_vec: torch.Tensor) -> list[dict]:
"""Create commitments for all positions at once (vectorized).
Produces bit-identical results to calling create_commitment() in a loop.
Args:
h_layer: Hidden states [seq_len, hidden_dim] on any device.
r_vec: Coefficient vector [topk] (int8, typically on CPU).
Returns:
List of commitment dicts, one per position.
"""
seq_len = h_layer.size(0)
# Step 1: Batched top-k selection
abs_h = h_layer.abs() # [seq_len, hidden_dim]
_, topk_indices = torch.topk(abs_h, k=self.topk, dim=1) # [seq_len, topk]
del abs_h
# Sort by index position for order-invariant sketch (v3)
topk_indices, _ = torch.sort(topk_indices, dim=1)
signed_values = torch.gather(h_layer, dim=1, index=topk_indices) # [seq_len, topk]
# Step 2: Vectorized log-magnitude bucketing (float64 for precision match)
buckets = log_magnitude_bucket_vectorized(signed_values, self.num_buckets)
# [seq_len, topk] int64
del signed_values
# Step 3: Batched sketch via matrix-vector product
# Max |dot| = topk * 7 * 127 = 14,224 (topk=16) — fits in int32, but CUDA
# doesn't support int matmul ("addmv_impl_cuda" not implemented for 'Int').
# Use float32 — exact for integers up to 2^24 (our max is 14,224).
buckets_f = buckets.to(torch.float32)
r_vec_f = r_vec.to(torch.float32).to(buckets_f.device)
sketches = (buckets_f @ r_vec_f).to(torch.int64) # [seq_len]
del buckets, buckets_f
# Use Python % for exact match with scalar: int(sketch.item()) % PRIME_Q
sketches_list = sketches.tolist()
sketch_vals = [s % PRIME_Q for s in sketches_list]
# Step 4: Package as list of dicts
return [{"sketch": sketch_vals[pos]} for pos in range(seq_len)]
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
1 participant
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.
s_vals were different on different GPUs (A100, H100, RTX 4090) with same inputs.
Fixed it by using fixed-point quantization + FP64 for accumulation. Now they match exactly.
Tested on A100, H100, RTX 4090 — all give same s_vals now.
Performance: ~35ms on A100, ~14.5ms on H100, ~52ms on RTX 4090.
Memory: ~19GB with FP16 weights — fits on RTX 4090.
Works with 7B and 13B models.
Summary by CodeRabbit