feat(train): add --gradient-checkpointing for DFlash/DSpark

[WindChimeRan]

Summary

Adds opt-in --gradient-checkpointing for DFlash and DSpark. It recomputes draft-layer activations during backward, reducing peak memory at the cost of step time.

The flag remains off by default. EAGLE-3, P-EAGLE, and MTP are unchanged.

H100 results

DFlash with a Qwen3-8B verifier, --total-seq-len 8192, and current defaults (--num-layers 5 --block-size 16, CE, Muon) on one H100 80GB. Values are medians of three interleaved runs after warmup.

max-anchors	Step time, off → on	Peak allocated, off → on
512 (default)	365 → 395 ms (+8.2%)	33.0 → 25.5 GB (-22.6%)
1024	507 → 562 ms (+10.9%)	44.1 → 31.2 GB (-29.1%)

The added time is in backward recomputation; forward and optimizer times were effectively unchanged.

Benchmark command and environment

tree      007c7c1 (content-identical to current head b9f5e1e)
stack     Python 3.12.9, torch 2.13.0+cu130, CUDA 13.0, transformers 5.14.1
protocol  synthetic batches; 10 warmup + 4 measured steps; 3 interleaved repeats
command   python scripts/benchmark.py run --synthetic --warmup-steps 10 \
            --measured-steps 4 -- --speculator-type dflash \
            --verifier-name-or-path Qwen/Qwen3-8B [--gradient-checkpointing]

Implementation

• Adds the training flag and CLI documentation.
• Enables non-reentrant checkpointing for DFlashDraftModel (inherited by DSpark).
• Records the setting in benchmark metadata.
• Adds a regression test for recomputation and gradient equivalence.

Non-reentrant checkpointing is required because draft layers pass inputs as keyword arguments; the reentrant wrapper forwards only positional arguments and would break draft-layer gradients.

Validation

• Targeted tests: 490 passed, 5 skipped
• Ruff and mypy: no findings introduced relative to main
• End-to-end training with checkpoint creation

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 27b1cc72-2180-43ea-a2a5-81b6c833797f

📥 Commits

Reviewing files that changed from the base of the PR and between 1735da6 and 01760b1.

📒 Files selected for processing (6)

• docs/cli/train.md
• scripts/benchmark.py
• scripts/train.py
• src/speculators/models/dflash/core.py
• src/speculators/train/config/schema.py
• tests/unit/models/test_dflash_gradient_checkpointing.py

---

📝 Walkthrough

Walkthrough

Changes

The training configuration now exposes disabled-by-default gradient checkpointing. DFlash models advertise support, training and benchmark scripts apply the setting, benchmark metadata records it, and unit tests verify recomputation and gradient equivalence.

Gradient checkpointing

Layer / File(s)	Summary
Checkpointing configuration and model support src/speculators/train/config/schema.py, src/speculators/models/dflash/core.py	TrainerArgs adds the gradient_checkpointing option. DFlashDraftModel declares checkpointing support.
Training and benchmark wiring scripts/train.py, scripts/benchmark.py	Training and benchmark setup enable non-reentrant checkpointing on the draft model. Benchmark results record the configured value.
Checkpointing validation and CLI documentation tests/unit/models/test_dflash_gradient_checkpointing.py, docs/cli/train.md	The test compares checkpointed and baseline gradients and verifies recomputation. The CLI documentation describes the new flag.

Possibly related issues

• vllm-project/speculators issue 831 — Both changes add gradient-checkpointing support and model capability declarations, but this PR uses the training configuration path.

Possibly related PRs

• vllm-project/speculators#872 — Adds related gradient-checkpointing support through the shared model hierarchy.
• vllm-project/speculators#852 — Introduces benchmark metadata that this PR extends with the checkpointing setting.
• vllm-project/speculators#824 — Shares DFlash training and model files, although it implements Liger kernel support.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.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
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.
Title check	✅ Passed	The title clearly and concisely identifies the new training flag for DFlash and DSpark.
Description check	✅ Passed	The description directly explains the feature, implementation, performance results, and validation for the changeset.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require approval from approved reviewers list	👀 reviews

🔴 Require approval from approved reviewers list

Waiting for any of

• approved-reviews-by = dsikka
• approved-reviews-by = fynnsu
• approved-reviews-by = orestis-z
• approved-reviews-by = rahul-tuli
• approved-reviews-by = shanjiaz

This rule is failing.

All pull requests must have at least one approving review from a member of the approved reviewers list before merging.

• any of:
  • approved-reviews-by = dsikka
  • approved-reviews-by = fynnsu
  • approved-reviews-by = orestis-z
  • approved-reviews-by = rahul-tuli
  • approved-reviews-by = shanjiaz

[mgoin]

I feel like --max-anchors 3072 is an extreme case to exaggerate the memory usage. I'd be curious to see the impact with a more standard 512 or 1024

Still I think this is useful to have as an option for users if not difficult to maintain