feat: staged training from a published draft, plus two draft-loading fixes - #168
Conversation
…rmers 5.x
Transformers 5.x reads `model.all_tied_weights_keys` while loading a
checkpoint, to keep tied parameters out of the missing-key set. The attribute
is populated by `post_init()`, which no draft model calls, so every draft --
Eagle3 Llama, Eagle3 DeepSeek MLA, DFlash and DSpark alike -- fails
`from_pretrained` with
AttributeError: 'LlamaForCausalLMEagle3' object has no attribute
'all_tied_weights_keys'
Declare the mapping as empty on the two draft base classes. Empty is the
correct value rather than a placeholder: a draft has its own (often pruned)
draft vocabulary, so its LM head can never share storage with the input
embedding.
Adds `tests/test_draft_from_pretrained.py`, which round-trips every draft
family through `save_pretrained`/`from_pretrained` and pins the non-tying
invariant the empty mapping depends on.
Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
`AutoEagle3DraftModel` builds `DSparkDraftModel` for any `DSparkConfig`, but
`AutoDraftModelConfig` only recognised the serving-side architecture names
`Qwen3DSparkModel` and `K3DSparkModel`. A DSpark draft config that names the
class this repo actually instantiates was rejected with
ValueError: Architecture DSparkDraftModel not supported
even though it is exported from the package root and instantiable directly.
Register the name alongside the serving alias so both spellings reach the same
model.
Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
Adds `model.initial_draft_model_path`. When set, rank 0 loads the named draft into the freshly constructed model before the target embedding is applied, which is what makes staged or continual training off an already-published draft possible. It is distinct from `training.load_path`: that resumes a training checkpoint with its optimizer and scheduler state, this seeds weights from a serving artifact and starts a new run. The path may name either `model.safetensors` or the directory holding it, and anything else is rejected rather than ignored -- a typo here would otherwise start a full run from random weights. Loading is strict in both directions for the same reason. `tools/convert_to_hf.py` publishes drafts under serving key names (`midlayer.` -> `layers.0.`, `context_proj.` -> `fc.`, ...), so the keys are mapped back first, via a new `keymap.to_internal_keys`. The forward mapping is not injective across draft families -- an Eagle3 draft's own `norm.*` is also DFlash's export name for `final_norm.*` -- so the reverse rename is applied per key and only where the key is not already a parameter of the model being loaded, which lets the model itself disambiguate. The path joins `_ALWAYS_LOCAL_PATH_KEYS` so it reaches Ray actors absolutized, since they do not share the launcher's working directory. Tests cover path resolution, loading published Eagle3 and DFlash drafts plus a native-keyed one, both strictness directions, the reverse key mapping, the config plumbing, and `init_model` end to end -- including that the target model's embedding is applied on top of the loaded checkpoint rather than under it. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ae5f589e61
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| loaded_from = checkpoint.load_initial_draft_weights( | ||
| draft_model, initial_draft_model_path | ||
| ) |
There was a problem hiding this comment.
When initial_draft_model_path points at a vocab-pruned Eagle3 draft (draft_vocab_size < vocab_size, e.g. configs/draft_models/qwen3_8b_eagle3.json), this load brings in both the pruned lm_head.weight and the checkpoint's saved d2t/t2d ordering. After actor initialization, however, train_entry.py still recomputes the mapping for the current dataset and calls set_vocab_buffers, which only overwrites the buffers and does not permute the already-loaded lm_head. If the staged run's token mapping differs from the published artifact, target probabilities are sliced with the new t2d while the draft head rows remain in the old order, silently corrupting the loss/logits; validate/reuse the artifact mapping or remap the head before training.
Useful? React with 👍 / 👎.
Most of them restated the code or the commit message. Kept only the ones a reader cannot derive: why the empty tied-weight mapping is not dead code, why the reverse key rename consults the model, and why load_embedding has to run after the initial checkpoint. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6aec6cf9a0
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| @dataclass | ||
| class ModelConfig: | ||
| draft_model_config: Optional[str] = None | ||
| initial_draft_model_path: Optional[str] = None |
There was a problem hiding this comment.
Because this option is added to the shared ModelConfig and flattened for every trainer, a DFlash or DSpark run can now accept model.initial_draft_model_path, but TrainerActor dispatches those configs to DFlashTrainer/DSparkTrainer and only Eagle3Trainer reads the field. In that scenario the run proceeds from freshly initialized draft weights even though the user requested a published draft warm-start, so either reject the option for non-Eagle3 configs or consume it in the DFlash/DSpark init path.
Useful? React with 👍 / 👎.
| initial_draft_model_path = getattr(self.args, "initial_draft_model_path", None) | ||
| if initial_draft_model_path: | ||
| loaded_from = checkpoint.load_initial_draft_weights( | ||
| draft_model, initial_draft_model_path | ||
| ) |
There was a problem hiding this comment.
When the initial draft path is missing or the safetensors state dict mismatches in a multi-rank run, rank 0 raises from this load inside the enclosing dist.get_rank() == 0 block while the other ranks skip it and continue to the subsequent dist.barrier(group=get_gloo_group()). That turns the explicit bad-path failure into a distributed init hang until timeout; validate/propagate the failure to all ranks before the barrier or make every rank fail consistently.
Useful? React with 👍 / 👎.
A vocab-pruned draft's `lm_head` row j holds the j-th target token its mapping selected, and `compute_target_p_padded` prunes the target side with `target_lm_head_weight[t2d]` in that same ascending-token-id order. The two are only aligned while both sides use one mapping. `train_entry` recomputes the mapping from the current dataset after the actors have initialised, then calls `set_vocab_buffers`, which replaced the buffers and nothing else. A draft seeded from a published artifact -- by the new `initial_draft_model_path`, or by `training.load_path` restoring buffers through DCP -- therefore kept its old head ordering while the target side moved to the new token set, corrupting the loss with no error and no log line. `set_vocab_buffers` now refuses to replace a mapping that arrived with loaded weights unless the new one selects the same tokens. A freshly built draft is unaffected: its `t2d` is still the all-pass sentinel, which is what `has_vocab_pruning` keys off. For the staged case where the new corpus legitimately has different token statistics, `model.keep_initial_vocab_mapping` trains against the mapping the weights were built with and skips the recompute entirely. It is rejected at config load without `initial_draft_model_path` or `load_path`, since there would be no mapping to keep and the head would train against an all-pass t2d. Remapping the head instead is not offered: tokens the new mapping adds have no trained row to permute into place. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f8cdd36768
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| return | ||
| if not ( | ||
| OmegaConf.select(config, "model.initial_draft_model_path") | ||
| or OmegaConf.select(config, "training.load_path") |
There was a problem hiding this comment.
Here training.load_path is treated as proof that a saved t2d/d2t mapping will be restored, but checkpoint.load() later returns None after only logging when the path, tracker, model dir, or model load is missing/incompatible. In that scenario with model.keep_initial_vocab_mapping=true and a pruned Eagle3 draft, train_entry.py skips computing/sending a dataset mapping, so the draft keeps its initial all-true t2d while lm_head is still only draft_vocab_size, leading to a late loss/logit shape failure instead of a valid mapping. Please make this mode require an actually loadable model checkpoint, or only skip mapping after the mapping was restored.
Useful? React with 👍 / 👎.
2 participants
Three independent, small changes to draft-model loading. They are batched because each is a few
lines on its own; they can be reviewed commit by commit.
1.
all_tied_weights_keys— drafts cannot be loaded at all on Transformers 5.xTransformers 5.x reads
model.all_tied_weights_keyswhile loading a checkpoint, to keep tiedparameters out of the missing-key set. The attribute is populated by
post_init(), which no draftmodel calls, so
AutoEagle3DraftModel.from_pretrainedcurrently fails for every draft family:Declared as empty on the two draft base classes (
Eagle3DraftModel,DFlashDraftModel) rather thanper model, since the cause is shared. Empty is the correct value, not a placeholder: a draft has its
own — often pruned — draft vocabulary, so its LM head can never share storage with the input
embedding. The new test pins that invariant alongside the round trip.
2.
DSparkDraftModelis not accepted as an architecture nameAutoEagle3DraftModelbuildsDSparkDraftModelfor anyDSparkConfig, butAutoDraftModelConfigonly recognised the serving-side names
Qwen3DSparkModelandK3DSparkModel, so a config namingthe class this repo actually instantiates was rejected with
ValueError: Architecture DSparkDraftModel not supported— despite the class being exported from the package root.3.
model.initial_draft_model_path— start training from a published draftNew config field. When set, rank 0 loads the named draft into the freshly constructed model before
the target embedding is applied, which is what makes staged or continual training off an
already-published draft possible. It is distinct from
training.load_path: that resumes a trainingcheckpoint with its optimizer and scheduler state, this seeds weights from a serving artifact and
starts a new run.
Two details worth a reviewer's attention:
tools/convert_to_hf.pypublishes drafts under serving key names(
midlayer.→layers.0.,context_proj.→fc., …), so a published draft does not load into atraining-side model as-is. The keys are mapped back through a new
keymap.to_internal_keys. Theforward mapping is not injective across draft families — an Eagle3 draft's own
norm.*is alsoDFlash's export name for
final_norm.*, and the same holds forfc.*vscontext_proj.*— sothe reverse rename is applied per key and only where the key is not already a parameter of the
model being loaded. That lets the model itself disambiguate instead of guessing from the config.
that does not line up with the configured draft, would otherwise leave part of the model randomly
initialised and cost a full training run to notice.
The helpers live in
torchspec/training/checkpoint.pynext to the other checkpoint loading, and thepath joins
_ALWAYS_LOCAL_PATH_KEYSso it reaches Ray actors absolutized — they do not share thelauncher's working directory.
Only the Eagle3 trainer is wired up; DFlash and DSpark would each need the same three lines and are
left alone here.
Testing
ruff check .andruff format --check .clean.New coverage, 24 tests:
tests/test_draft_from_pretrained.py—save_pretrained/from_pretrainedround trip across allfour draft families, plus the non-tying invariant. Every one of these fails on
mainwith theAttributeErrorabove.tests/test_initial_draft_model_path.py— path resolution (file, directory,~, and the threerejection cases), loading published Eagle3 and DFlash drafts as well as a native-keyed one, both
strictness directions, the reverse key mapping including the
norm/final_normcollision, andthe config plumbing.
tests/test_eagle3_trainer.py—init_modelend to end with FSDP2 and the optimizer stubbed,including that the target model's embedding is applied on top of the loaded checkpoint rather
than under it.
tests/test_dspark.py— the new architecture name resolves and buildsDSparkDraftModel.Full-suite comparison against a pristine
mainworktree in the same sandbox: identical failuresets (149 pre-existing failures and 5 collection errors, all from missing CUDA/mooncake/
ray.dagin the sandbox), with 24 more tests passing on this branch.
tests/test_eagle3_trainer.pyis one ofthe modules that cannot be collected without a CUDA-linked mooncake, so its three new tests were run
with
mooncakestubbed; they should be re-run on a real GPU box before merge, along with the restof the suite.
4. Swapping the vocabulary mapping under a seeded draft head
Found while reviewing (3), and it predates it. A vocab-pruned draft's
lm_headrow j holds thej-th target token its mapping selected, and
compute_target_p_paddedprunes the target side withtarget_lm_head_weight[t2d]in that same ascending-token-id order — the two are aligned only whileboth use one mapping.
train_entryrecomputes the mapping from the current dataset after the actors have initialisedand then calls
set_vocab_buffers, which replaced the buffers and nothing else. So a draft seededfrom a published artifact kept its old head ordering while the target side moved to the new token
set: wrong rows, no error, no log line. This reaches
training.load_pathtoo, since DCP restoresthe buffers with the rest of the state, so it is not new to this PR — the new field just makes it
easy to hit.
set_vocab_buffersnow refuses to replace a mapping that came in with loaded weights unless thenew one selects the same tokens. A freshly built draft is unaffected: its
t2dis still theall-pass sentinel that
has_vocab_pruningkeys off. For the staged case where the new corpusgenuinely has different token statistics,
model.keep_initial_vocab_mappingtrains against themapping the weights were built with and skips the recompute; it is rejected at config load without
initial_draft_model_pathorload_path. Remapping the head instead is not offered — tokens thenew mapping adds have no trained row to permute into place.