feat: support torchspec training - #798
Conversation
Signed-off-by: Dogac Eldenk <dogacel@gmail.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: 885b66f8ed
ℹ️ 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".
| if not hasattr(self.model_runner.model, "set_eagle3_layers_to_capture"): | ||
| raise ValueError( | ||
| "The target model does not support auxiliary hidden-state capture" |
There was a problem hiding this comment.
This capability check passes for every BaseCausalLM subclass because set_eagle3_layers_to_capture is inherited from the base class, even when the underlying model ignores layers_to_capture; for example, Qwen2Model.forward always returns None for auxiliary states. In that case the FULL-capture fallback in LogitsProcessor uses the final hidden states as hidden_states, so Mooncake silently receives the final states for both _hs and _lhs, corrupting TorchSpec training data instead of rejecting the unsupported model.
Useful? React with 👍 / 👎.
| if config.global_rank == 0: | ||
| from tokenspeed.runtime.execution.spec_training_mooncake import ( | ||
| SpecTrainingMooncakeExporter, | ||
| ) | ||
|
|
||
| self.spec_training_exporter = SpecTrainingMooncakeExporter() |
There was a problem hiding this comment.
With attention data parallelism greater than one, each DP replica processes a different request batch, but this condition creates an exporter only on global rank 0. Consequently requests assigned to every other DP group are never published, while output_processor.py still advertises Mooncake store keys for those completed requests; the gate needs to select one TP/CP leader per DP replica rather than only the global leader.
Useful? React with 👍 / 👎.
| parser.add_argument( | ||
| "--enable-spec-training-mooncake", | ||
| action="store_true", | ||
| default=ServerArgs.enable_spec_training_mooncake, | ||
| help="Capture full-sequence auxiliary and final hidden states during " | ||
| "offline prefill and publish them to Mooncake for TorchSpec.", |
There was a problem hiding this comment.
Add user-facing documentation covering how to enable this new server mode, its required Mooncake environment variables, supported models and parallelism, request requirements, and the returned key convention; currently the CLI help is the only guidance for a feature that otherwise requires several undiscoverable settings.
AGENTS.md reference: AGENTS.md:L13-L13
Useful? React with 👍 / 👎.
|
This PR has been inactive for 14 days and is marked as stale. It will be closed in 3 days if there is no further activity. |
1 participant
Summary
Add hidden state capture & speculative decoding training via TorchSpec support (PR) to Tokenspeed.
Test Plan
Currently tested on Qwen3-8B for simplicity over 1000 hidden states, both offline and online training is done.
2 Qwen3-8B instances with TP=2 (stress testing the distributed path) can handle around 75 samples per second.
TODO: Will report actual runs with the 1000 step trained EAGLE3 model with Tokenspeed to validate trained model is actually compatible with tokenspeed.