Support EagleDraftModel.from_pretrained for finetuning

[fynnsu]

Prereq

Depends on #332, which should be merged first.

Purpose

We'd like to support finetuning existing Eagle3 models. This pr adds the --from-pretrained option, which loads a pretrained model from HF / local path.

Setup flow

Fresh model setup, i.e. from_training_args pathway

1. model = Eagle3DraftModel.init
   • Initialize all modules/parameters/buffers
   • Use torch.zeros (w/ correct shape + dtype) as placeholders for t2d/d2t
   • Use torch.nan (w/ correct shape + dtype) as placeholders for lm_head, embed_tokens, and verifier_lm_head
2. model.load_vocab_mappings(t2d, d2t)
   • t2d/d2t can be None, which will cause an early return
   • Verify shapes match and then load t2d/d2t
3. model.load_verifier_weights()
   • Loads embed_tokens, lm_head (into both lm_head and verifier_lm_head), verifier_norm from verifier model
   • Checks if embed_tokens and lm_head are still NaN (from init) before loading. Otherwise they are skipped
   • verifier_lm_head is always loaded
   • verifier_norm is skipped if the weight isn't found (with a warning). Note we don't use NaN init for this module

Continues below

Finetuning setup, i.e. from_pretrained pathway

1. model = Eagle3DraftModel.init call internally by PreTrainedModel.from_pretrained under meta device context
2. PreTrainedModel.from_pretrained also loads model weights for us
3. model.load_vocab_mappings(t2d, d2t) called, same as above
4. model.load_verifier_weights() called
   • lm_head and embed_token loading should be skipped because the values have already been set by PreTrainedModel.from_pretrained
   • verifier_lm_head and verifier_norm loaded (overriding existing values but should be the same)

Continues below

Joint (Fresh + Finetuning) next steps (in Trainer.setup_model)

If distributed:
    cache state dict on rank0
    Apply `fully_shard` to model layers and model (don't move to meta device, don't re-init weights)
    
    if checkpoint exists:
        load checkpoint (distributed)
    else:
        broadcast cached state dict from rank0 to all ranks
else:
    move model to local device
    if checkpoint exists:
        load checkpoint (single device)

Implementation

This uses the .from_pretrained method from the base PreTrainedModel class to resolve local vs remote model and handle downloading checkpoints.

As part of this pr, I also clean up the model initialize process, which previously also including loading verifier weights and t2d/d2t tensors. These will now be loaded after the init function. This was required to support .from_pretrained as the model init is run under a meta device context.

This means the init instead sets up placeholder buffers for verifier parameters/vocab mapping. These are being set intentionally in a way that makes it easy to confirm they are overwritten when training starts (e.g. by initializing some values to NaN so that failing to overwrite them will result in immediate NaN outputs from the model).

I also updated the fully shard handling to ensure values are set on all ranks correctly when using distributed training.

Testing

Added comprehensive test coverage for loading pathway combinations (e.g. loading from checkpoint, pretrained, fresh init, w/ and w/o vocab mappings, single gpu and distributed, etc.).

Note some tests require single or multi-gpu. These are skipped if requirements are not met. It would be good to ensure these are being run correctly at regular intervals.

[github-actions]

📦 Build Artifacts Available
The build artifacts (`.whl` and `.tar.gz`) have been successfully generated and are available for download: https://github.com/vllm-project/speculators/actions/runs/22909838668/artifacts/5852352190.
They will be retained for up to 30 days.
Commit: c09fab9

[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

[VincentG1234]

Tested locally and everything behaved smoothly as expected. 👍

One small observation: I noticed that the current tests are mainly unit tests. For a feature like this (loading from pretrained + finetuning), an E2E sanity check like the test here could be useful and robust to guard against regressions in the full workflow. That kind of test tends to catch subtle issues that unit tests may miss.

Of course, it could also make CI heavier since it involves running a short training job, so this is just a suggestion.

[fynnsu]

@VincentG1234

Tested locally and everything behaved smoothly as expected. 👍

One small observation: I noticed that the current tests are mainly unit tests. For a feature like this (loading from pretrained + finetuning), an E2E sanity check like the test here could be useful and robust to guard against regressions in the full workflow. That kind of test tends to catch subtle issues that unit tests may miss.

Of course, it could also make CI heavier since it involves running a short training job, so this is just a suggestion.

Yeah this make sense, I will add something like that. We have a separate nightly job that runs our E2E tests, so it won't block prs but will be a useful signal if there's a regression.

[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]

Removed merge commit and rebased.