raise error for nan hidden states

[shanjiaz]

Purpose

We had a failure early this week due to hiddenstatesconnector sending nan hidden states. However, our training kept going. We should explicitly fair when getting garbage hidden states

Tests

Tested locally and ran the e2e smoke tests.

Now failes directly:

                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/shanjiaz/speculators/src/speculators/train/data.py", line 327, in _maybe_generate_hs
    check_hidden_states(loaded_hs, dataset_item["input_ids"].tolist())
  File "/home/shanjiaz/speculators/src/speculators/data_generation/offline.py", line 17, in check_hidden_states
    raise ValueError("Hidden states contain NaN values")
ValueError: Hidden states contain NaN values
================================================== 1 failed, 16 warnings in 180.86s (0:03:00) ==================================================

Checklist

I have filled in:

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan/results, such as providing test command and pasting the results.
• (Optional) The necessary documentation update.
• I (a human) have written or reviewed the code in this pr to the best of my ability.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: aeda6a2c-7645-444a-8e3b-9174216c4fbe

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

check_safetensors_file is moved from scripts/data_generation_offline.py into src/speculators/data_generation/offline.py as a shared helper. The training data pipeline (train/data.py) imports and calls this helper after hidden-state generation, and re-raises ValueError containing "NaN" instead of silently skipping.

Changes

Centralize safetensors validation helper

Layer / File(s)	Summary
Shared check_safetensors_file implementation src/speculators/data_generation/offline.py	Adds check_safetensors_file with safe_open-based validation: checks token_ids equality, hidden_states NaN presence, and sequence-length match against the expected token list, raising ValueError on any failure.
Training pipeline validation and NaN re-raise src/speculators/train/data.py	Imports check_safetensors_file and calls it against input_ids immediately after hidden-state file generation. Adds special handling to re-raise ValueError messages containing "NaN" rather than falling through to the generic warn-and-skip path. Updates the hidden_states shape comment to reference num_layers instead of a fixed value.
Remove local helper from generation script scripts/data_generation_offline.py	Replaces the locally defined check_safetensors_file function and its safetensors import with an import of the same function from speculators.data_generation.offline.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title directly summarizes the main change: raising errors for NaN hidden states, which aligns with the primary objective across all three modified files.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	The PR description clearly explains the issue being addressed (NaN hidden states causing silent failures) and describes the solution (explicit error handling to fail fast). It includes test results demonstrating the new behavior.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fail-if-nan

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[mergify]

The quality checks have failed. Please run make style and make quality under
the root directory to address the lint failures. You will need to install the
dev optional install to get the required linting packages:
https://github.com/vllm-project/speculators/blob/main/CONTRIBUTING.md

[coderabbitai]

🧹 Nitpick comments (1)

src/speculators/train/data.py (1)

327-330: ⚡ Quick win

Consider optimizing to avoid opening the file twice.

The safetensors file is opened twice in sequence: once at line 325 by _maybe_load_hs_file to load tensors, and again here by check_safetensors_file to validate. This adds I/O overhead during hidden state generation.

Consider refactoring check_safetensors_file to optionally return the loaded tensors, allowing a single file open operation to serve both purposes.

🤖 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 `@src/speculators/train/data.py` around lines 327 - 330, The safetensors file
is being opened twice in succession: once by _maybe_load_hs_file to load tensors
and again by check_safetensors_file to validate, creating unnecessary I/O
overhead. Refactor the check_safetensors_file function to optionally return the
loaded tensors in addition to performing validation, allowing you to reuse the
already-loaded tensors from check_safetensors_file instead of calling
_maybe_load_hs_file separately. This consolidates the file open operation into a
single call while maintaining the validation logic.

🤖 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.

Nitpick comments:
In `@src/speculators/train/data.py`:
- Around line 327-330: The safetensors file is being opened twice in succession:
once by _maybe_load_hs_file to load tensors and again by check_safetensors_file
to validate, creating unnecessary I/O overhead. Refactor the
check_safetensors_file function to optionally return the loaded tensors in
addition to performing validation, allowing you to reuse the already-loaded
tensors from check_safetensors_file instead of calling _maybe_load_hs_file
separately. This consolidates the file open operation into a single call while
maintaining the validation logic.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 30c6f624-a91f-4ac5-958d-9a061c6e1736

📥 Commits

Reviewing files that changed from the base of the PR and between b939395 and 7b22bc8.

📒 Files selected for processing (3)

• scripts/data_generation_offline.py
• src/speculators/data_generation/offline.py
• src/speculators/train/data.py

[fynnsu]

One comment, otherwise looks good.

[mergify]

The quality checks have failed. Please run make style and make quality under
the root directory to address the lint failures. You will need to install the
dev optional install to get the required linting packages:
https://github.com/vllm-project/speculators/blob/main/CONTRIBUTING.md

[github-actions]

Link Check Results (DOCS)

All links are now valid - this issue has been resolved.

---

Marked as resolved: 9e20311

[github-actions]

Link Check Results (REPO)

All links are now valid - this issue has been resolved.

---

Marked as resolved: 9e20311