[EAGLE] Configurable number of TTT steps

[benchislett]

What does this PR do?

Type of change: new CLI option for existing option

• Added num_ttt_steps CLI flag
• Changed num_ttt_steps default from 4 to 3 for consistency. Num_spec_tokens == 3 or == 7 are most common in practice, so rounding down to 3 and allowing users to increment higher on-demand. Will also improve training efficiency for the OOTB experience.

Usage

Users can now pass --num_ttt_steps 7 to launch_train.sh when training an EAGLE3 model for extended speculation lengths.

Testing

N/A

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: ✅
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: N/A
• Did you update Changelog?: N/A

Summary by CodeRabbit

• New Features
  • Added ability to configure train-time-test steps for speculative decoding training via command-line argument.
  • Updated default train-time-test steps value from 4 to 3.

[coderabbitai]

📝 Walkthrough

Walkthrough

The pull request adds support for a configurable num_ttt_steps parameter across the speculative decoding training pipeline, allowing users to specify the number of train-time-test steps via command-line arguments with a default value of 3.

Changes

Cohort / File(s)	Summary
Shell script argument handling examples/speculative_decoding/launch_train.sh	Added --num_ttt_steps command-line argument parsing with default value of 3 and propagation to the acceleration launch command.
Python argument definition examples/speculative_decoding/main.py	Added num_ttt_steps integer argument to EagleArguments class with default value 3 and propagated it to the Eagle configuration dictionary.
Configuration defaults modelopt/torch/speculative/config.py	Updated default value of eagle_ttt_steps field in EagleConfig from 4 to 3 to align with new standard.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 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 (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title '[EAGLE] Configurable number of TTT steps' directly and clearly describes the main change: making the number of TTT (test-time training) steps configurable via a new CLI flag and configuration parameter.
Security Anti-Patterns	✅ Passed	PR introduces CLI flag for train-time-test steps without introducing any of six CRITICAL security anti-patterns.

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

✨ Finishing Touches

• 📝 Generate docstrings (stacked PR)
• 📝 Generate docstrings (commit on current branch)

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch bchislett/configurable-ttt-steps

📝 Coding Plan

• Generate coding plan for human review comments

---

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

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@examples/speculative_decoding/launch_train.sh`:
- Around line 89-92: NUM_TTT_STEPS is taken directly from user input and later
interpolated into a shell command; validate and harden it in launch_train.sh
before any sh -c invocation by enforcing it is a positive integer within an
acceptable range (e.g., digits only, min/max bounds) and rejecting or defaulting
on invalid values; locate the parsing block that assigns NUM_TTT_STEPS and add
validation (or sanitize with a strict numeric cast) so only safe numeric values
are used, and avoid passing raw user strings into sh -c (prefer passing as
positional arguments or using printf %q if quoting is required).

In `@examples/speculative_decoding/main.py`:
- Around line 133-136: num_ttt_steps may be non-positive and is forwarded into
the Eagle config unchecked; add a validation check where num_ttt_steps is
consumed (before constructing or passing into the Eagle config) that ensures it
is an integer >= 1 (optionally enforce an upper bound if desired), and raise a
clear ValueError if the check fails; e.g., validate the num_ttt_steps attribute
from the dataclass/config and refuse to build the Eagle config (or normalize)
when the value is <= 0 so the downstream Eagle config creation code (the place
where num_ttt_steps is passed to the Eagle configuration) never receives invalid
values.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: dda22186-a514-429c-86b4-1b48b23a6da7

📥 Commits

Reviewing files that changed from the base of the PR and between 1070d89 and 1a5276d.

📒 Files selected for processing (3)

• examples/speculative_decoding/launch_train.sh
• examples/speculative_decoding/main.py
• modelopt/torch/speculative/config.py

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 70.09%. Comparing base (1070d89) to head (1a5276d).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1042      +/-   ##
==========================================
- Coverage   70.10%   70.09%   -0.02%     
==========================================
  Files         221      221              
  Lines       25541    25541              
==========================================
- Hits        17905    17902       -3     
- Misses       7636     7639       +3     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.