Reduce graph breaks and recompiles

[fynnsu]

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

Purpose

Late in epoch graph recompiles can cause training to crash after a substantial number of batches have been processed. This pr works to decrease these occurrences and also reduce the number of graph breaks in the compiled code.

Related Issue

I think this should resolve #544.

Tests

Will add some tests to this pr.

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: b8ab62d1-8d20-43bf-9bfa-941b835c54e8

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

The PR modifies the data collation pipeline to pad the lengths tensor to a minimum size of 4 elements, replacing direct tensor construction with a padded version using torch.nn.functional.pad.

Changes

Data Collation Tensor Padding

Layer / File(s)	Summary
Lengths tensor minimum size padding src/speculators/train/data.py	After computing new_lengths, the collated_data["lengths"] tensor is now explicitly created as torch.long and padded to ensure at least 4 elements, instead of using the unpadded tensor directly.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~3 minutes

Suggested labels

bug, training

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

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.
Title check	❓ Inconclusive	The PR title 'Reduce graph breaks and recompiles' is vague and does not clearly reflect the actual change of padding the lengths tensor to minimum size 4.	Consider revising the title to be more specific about the change, such as 'Pad lengths tensor to minimum size 4 to avoid specialization' or similar, to better reflect the actual code modification.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The code change addresses issue #544 by padding lengths tensors to prevent inductor pattern-matching failures, aligning with the root cause mitigation objective.
Out of Scope Changes check	✅ Passed	The change is narrowly scoped to the collation function and directly addresses the lengths tensor padding requirement from issue #544.
Description check	✅ Passed	The pull request description clearly relates to the changeset. It explains the purpose (reducing late-epoch graph recompiles and graph breaks), links to a related issue, and describes the technical approach.

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

---

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

[coderabbitai]

🧹 Nitpick comments (1)

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

469-472: ⚡ Quick win

Convert length tensor scalars to Python ints before rebuilding lengths.

In src/speculators/train/data.py (create_collate_fn), iterating for length in lengths: yields 0-d tensor scalars; mixing them with cum_length (Python int) and collecting them into new_lengths makes torch.tensor(new_lengths, ...) less idiomatic/stable. Convert with length.item() and keep cum_length/the condition as pure ints.

Proposed diff

         lengths = collated_data["lengths"]
         new_lengths = []
         cum_length = 0
         for length in lengths:
-            if length + cum_length >= max_len:
+            length_i = int(length.item())
+            if length_i + cum_length >= max_len:
                 new_lengths.append(max_len - cum_length)
                 break
-            new_lengths.append(length)
-            cum_length += length
+            new_lengths.append(length_i)
+            cum_length += length_i
         lengths = torch.tensor(new_lengths, dtype=torch.long)
         # Pad to at least size 4 to avoid torch.compile recompiling on sizes 0/1
         pad_to = max(4, lengths.shape[0])
         collated_data["lengths"] = F.pad(lengths, (0, pad_to - lengths.shape[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 `@src/speculators/train/data.py` around lines 469 - 472, In create_collate_fn,
when iterating the tensor variable lengths you should convert each 0-dim tensor
to a Python int (use length.item()) before comparing/adding to cum_length and
appending to new_lengths; ensure cum_length stays an int and the condition uses
ints so new_lengths is a list of ints, then rebuild lengths with
torch.tensor(new_lengths, dtype=torch.long) and apply the existing padding logic
(pad_to, F.pad) as before.

🤖 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 469-472: In create_collate_fn, when iterating the tensor variable
lengths you should convert each 0-dim tensor to a Python int (use length.item())
before comparing/adding to cum_length and appending to new_lengths; ensure
cum_length stays an int and the condition uses ints so new_lengths is a list of
ints, then rebuild lengths with torch.tensor(new_lengths, dtype=torch.long) and
apply the existing padding logic (pad_to, F.pad) as before.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: d4014969-132a-4a0f-b139-3ddfff652ab0

📥 Commits

Reviewing files that changed from the base of the PR and between 1bc3788 and 9ac1d14.

📒 Files selected for processing (1)

• src/speculators/train/data.py

[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

[fynnsu]

Rebased

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @fynnsu.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[fynnsu]

E2E (vllm@nightly, transformers<5.0, this branch): https://github.com/neuralmagic/llm-compressor-testing/actions/runs/27716366270/job/81990214563

[shanjiaz]

Nice!