Multi node data generation offline

[fynnsu]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

To speed up offline hidden states generation, this pr adds support for partitioning the data generation process across multiple nodes.

Description

This pr adds --world-size and --rank args to the script. When set, they will partition the dataset so the first N//k (or plus 1) samples get processed by the first node, the second batch by the second node, etc, where N is the total dataset size and k is the world size.

The script should handle partitioning when N is not divisible by k by increasing the size of the first N % k nodes batches by 1.

The changes shouldn't interfere with the ability to resume data generation that has been stopped.

Related Issue

Tests

Added tests for the partitioning / finding existing samples logic.

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: 3d8a5875-0a28-45d4-bdb0-859528ad72f0

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

This PR adds multi-node dataset partitioning support for offline hidden-state generation. A new helper module extracts index discovery and rank-based partitioning logic, the main script integrates these helpers via new --world-size and --rank CLI arguments, and documentation explains usage across distributed setups.

Changes

Multi-Node Offline Data Generation

Layer / File(s)	Summary
Index management helper module src/speculators/data_generation/offline.py	New module provides get_existing_hidden_state_indices() to scan output paths for existing hidden-state files and get_indices_to_process() to partition dataset indices across ranks, accounting for max_samples caps, pre-existing files, and balanced remainder distribution.
Unit tests for index helpers tests/unit/data_generation/test_offline.py	Comprehensive test suite validates single/multi-node partitioning, max_samples limits, exclusion of existing indices, empty-rank cases, and full coverage. File discovery tests verify numeric suffix detection, invalid filename rejection, and sorted output.
Script CLI and integration scripts/data_generation_offline.py	Script adds --world-size and --rank arguments and imports the new helpers, replacing local index-selection logic. The call to get_indices_to_process() now passes world-size and rank to determine per-node work assignment.
User-facing documentation docs/cli/data_generation_offline.md, docs/user_guide/tutorials/train_eagle3_offline.md	CLI reference and tutorial updated with multi-node argument documentation and example commands for distributed dataset generation across multiple nodes.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Suggested labels

data-generation, enhancement

Suggested reviewers

• shanjiaz
• dsikka
• rahul-tuli

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title directly summarizes the main feature added: multi-node data generation offline, which matches the core objective of partitioning data generation across multiple nodes.
Description check	✅ Passed	The description clearly explains the purpose (speed up offline hidden states generation), describes the new --world-size and --rank arguments, explains partitioning logic, mentions resumption capability preservation, and confirms tests were added.
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.

_{✏️ 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 multi_node_datagen

---

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]

Actionable comments posted: 4

🧹 Nitpick comments (2)

tests/unit/data_generation/test_offline.py (1)

12-77: ⚡ Quick win

Add edge-case tests for max_samples=0 and out-of-range existing indices.

Current suite misses two high-value boundaries that affect correctness of the new partitioning path.

Suggested additions

 class TestGetIndicesToProcess:
+    def test_max_samples_zero_returns_empty(self):
+        result = get_indices_to_process(10, 0, [], world_size=1, rank=0)
+        assert result == []
+
+    def test_existing_out_of_range_does_not_mark_all_processed(self):
+        result = get_indices_to_process(3, None, [10, 11, 12], world_size=1, rank=0)
+        assert result == [0, 1, 2]

As per coding guidelines, "Ensure PyTest tests are clear, comprehensive, and cover edge cases specific to speculative decoding ... Verify that new code paths introduced in the PR are covered."

🤖 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/unit/data_generation/test_offline.py` around lines 12 - 77, Add two
pytest cases to cover max_samples=0 and out-of-range existing indices for
get_indices_to_process: (1) test that when max_samples=0 the function returns an
empty list for any world_size/rank (e.g., call get_indices_to_process(10, 0, [],
world_size=1, rank=0) and assert []), and (2) test that out-of-range values in
the existing list are ignored (e.g., pass existing containing negative and
>=num_samples values like [-1, 10, 999] to get_indices_to_process(5, None,
existing, world_size=1, rank=0) and assert it returns list(range(5))); add these
to the TestGetIndicesToProcess class so the new partitioning path is exercised
and boundaries validated.

docs/user_guide/tutorials/train_eagle3_offline.md (1)

157-171: ⚡ Quick win

Clarify shared storage requirement for multi-node generation.

The documentation mentions "shared (or later merged) output directory" but doesn't explain the implications of non-shared storage or how to merge outputs afterward. If nodes don't share the output directory during generation, each node will see different existing files and the partitioning logic may produce overlapping or incorrect work distribution.

Consider adding a note that clarifies:

• Whether shared storage (e.g., NFS) is required during generation, or
• If separate storage is used, how outputs should be merged and whether re-running with merged directories is necessary

📚 Suggested clarification

 **Use multiple nodes:**
 
-If you have access to multiple machines, each with its own vLLM server, you can split the dataset across them with `--world-size` and `--rank`. Each node generates a contiguous, non-overlapping chunk of the data into a shared (or later merged) output directory. See the [data_generation_offline.py cli reference](/cli/data_generation_offline.md) for details.
+If you have access to multiple machines, each with its own vLLM server, you can split the dataset across them with `--world-size` and `--rank`. Each node generates a contiguous, non-overlapping chunk of the data. 
+
+**Important:** All nodes should write to a shared output directory (e.g., mounted via NFS) so that each node can detect already-completed samples and resume correctly if interrupted. If using separate local storage per node, you must manually merge the output directories after generation completes.
+
+See the [data_generation_offline.py cli reference](/cli/data_generation_offline.md) for details.

As per coding guidelines, ensure documentation is clear and complete.

🤖 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/user_guide/tutorials/train_eagle3_offline.md` around lines 157 - 171,
Update the docs/user_guide/tutorials/train_eagle3_offline.md section that
describes multi-node generation to explicitly state the shared-storage
requirement and the correct workflow when nodes do not share storage: clarify
that if the nodes have a shared output directory (e.g., NFS) the partitioning
via --world-size and --rank will produce non-overlapping chunks safely,
otherwise each node will write separate outputs that must be merged later;
describe a recommended merge procedure (concatenate or deduplicate per-sample
files into a single output directory) and note that after merging you should
treat the merged directory as the canonical preprocessed-data for any downstream
run (or re-run data_generation_offline.py with consistent --world-size/--rank if
necessary). Reference the data_generation_offline.py script and the
--world-size/--rank flags in the note so readers know where partitioning
behavior is controlled.

🤖 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 `@scripts/data_generation_offline.py`:
- Around line 170-187: Add validation after parsing arguments in parse_args() to
ensure the --rank value is within [0, world-size): check that args.rank is >= 0
and args.rank < args.world_size and raise parser.error (or SystemExit with a
clear message) if not; update any callers (e.g., main()) to use the validated
args returned from parse_args() so invalid rank values are caught early and
produce a clear error message referencing the offending values.
- Around line 170-187: The help strings for parser.add_argument entries for
"--world-size" and "--rank" concatenate two adjacent string literals without a
space after the colon, producing "IMPORTANT: thisis..."; update the help text
for the parser.add_argument calls for "--world-size" and "--rank" (the help
tuples) by inserting a space after the colon so the fragments become "IMPORTANT:
this is ..." (you can add the trailing space at the end of the first literal or
a leading space at the second literal).

In `@src/speculators/data_generation/offline.py`:
- Around line 55-60: Validate inputs before computing chunk_size: check that
world_size is an integer > 0 and that rank is an integer in the range 0 <= rank
< world_size; if not, raise a ValueError with a clear message. Add these checks
immediately before the chunk_size calculation (the block using variables
world_size, rank, chunk_size, remainder, start, end) so you avoid
ZeroDivisionError and silent invalid slices and fail fast with informative
errors.
- Around line 49-53: The early-return and target calculation are wrong: treat
max_samples<=0 as "no limit", compute target before checking completion, and
normalize/filter existing entries to valid, deduplicated sample indices;
specifically, normalize max_samples to None if <=0, compute target =
min(max_samples, num_samples) only when max_samples is set else target =
num_samples, replace existing with a filtered set containing only indices 0 <= i
< num_samples to remove stale/out-of-range entries, then perform the completion
check using len(existing_set) >= target and return []; update the block around
variables existing, num_samples, max_samples and target so the check uses the
filtered set and correct target.

---

Nitpick comments:
In `@docs/user_guide/tutorials/train_eagle3_offline.md`:
- Around line 157-171: Update the
docs/user_guide/tutorials/train_eagle3_offline.md section that describes
multi-node generation to explicitly state the shared-storage requirement and the
correct workflow when nodes do not share storage: clarify that if the nodes have
a shared output directory (e.g., NFS) the partitioning via --world-size and
--rank will produce non-overlapping chunks safely, otherwise each node will
write separate outputs that must be merged later; describe a recommended merge
procedure (concatenate or deduplicate per-sample files into a single output
directory) and note that after merging you should treat the merged directory as
the canonical preprocessed-data for any downstream run (or re-run
data_generation_offline.py with consistent --world-size/--rank if necessary).
Reference the data_generation_offline.py script and the --world-size/--rank
flags in the note so readers know where partitioning behavior is controlled.

In `@tests/unit/data_generation/test_offline.py`:
- Around line 12-77: Add two pytest cases to cover max_samples=0 and
out-of-range existing indices for get_indices_to_process: (1) test that when
max_samples=0 the function returns an empty list for any world_size/rank (e.g.,
call get_indices_to_process(10, 0, [], world_size=1, rank=0) and assert []), and
(2) test that out-of-range values in the existing list are ignored (e.g., pass
existing containing negative and >=num_samples values like [-1, 10, 999] to
get_indices_to_process(5, None, existing, world_size=1, rank=0) and assert it
returns list(range(5))); add these to the TestGetIndicesToProcess class so the
new partitioning path is exercised and boundaries validated.

🪄 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: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 9246024a-996d-4a09-a5da-475b75f20cf4

📥 Commits

Reviewing files that changed from the base of the PR and between 4f80f5d and 6115c89.

📒 Files selected for processing (6)

• docs/cli/data_generation_offline.md
• docs/user_guide/tutorials/train_eagle3_offline.md
• scripts/data_generation_offline.py
• src/speculators/data_generation/offline.py
• tests/unit/data_generation/__init__.py
• tests/unit/data_generation/test_offline.py

[shanjiaz]

Took a brief look, do we have capacity to run tests for this tho?