| <!-- SUPERSEDED --> |
| > **Superseded by `plan.md`** (the reconciled SpecForge Architecture Plan). This is the |
| > original "redesign-from-scratch" draft, kept for its detailed component sketches |
| > (CheckpointManager, Evaluator, Pydantic config, MLA draft, SGLang export, train-with-decode). |
| > Where this draft and `plan.md` disagree (notably the "no Mooncake" stance), `plan.md` wins. |
|
|
| # SpecForge Redesign Plan |
|
|
| > **Status**: Draft (train-with-decode promoted to Phase 5) |
| > **Last updated**: 2026-05-31 |
|
|
| --- |
|
|
| ## 1. Context |
|
|
| SpecForge trains speculative decoding draft models (EAGLE3, DFlash) aligned with SGLang |
| serving. It works well for single-cluster torchrun workflows, but has accumulated structural |
| debt that makes it hard to add new architectures, new training modes, and production tooling. |
|
|
| TorchSpec (a sibling project) solves many of these gaps via Ray + Mooncake disaggregation, |
| but at the cost of heavy dependencies and operational complexity. This plan takes a different |
| approach: **keep SpecForge's simple torchrun-native design, but restructure internals so |
| new capabilities compose cleanly.** |
|
|
| ### What we keep verbatim |
|
|
| - `specforge/core/loss.py` — Triton `LogSoftmaxLoss` and `_compute_loss`. |
| - `specforge/optimizer.py` — `BF16Optimizer` (FP32 master weights + AdamW + grad clip). |
| - `specforge/lr_scheduler.py` — `CosineAnnealingWarmupLR` and the `TwoStageScheduler` family. |
| - `specforge/tracker.py` — `Tracker` ABC + `TRACKER_REGISTRY` (wandb/tensorboard/swanlab/mlflow). |
| - `specforge/distributed.py` — `init_distributed`, device meshes, yunchang USP integration. |
| - `specforge/core/eagle3_adapters.py` — `BackendAdapter` / `StepState` / `UspAdapter`. |
| - SGLang target backend code in `specforge/modeling/target/sglang_backend/`. |
| - All 30+ existing draft model configs. |
|
|
| ### What we throw away |
|
|
| - `scripts/train_eagle3.py` and `scripts/train_dflash.py` as god scripts (become thin shims). |
| - Hardcoded `_model_mapping` / `_config_mapping` dicts in `specforge/modeling/auto.py`. |
| - `configs/deepseek-v3-671b-eagle3.json` that claims `model_type: llama` for a DeepSeek target. |
| - The argparse-flags + per-arch JSON config split. |
| - `QwenVLOnlineEagle3Model` (VLM handled by target engine + data pipeline, not a separate model class). |
|
|
| ### Workloads in scope |
|
|
| These are the concrete training workflows the redesign must support. All four share the same |
| trainer / strategy / draft surface; they differ only in which `HiddenStateStream` and |
| `TargetEngine` get composed at the top. |
|
|
| | # | Workload | Description | Engine | Stream | |
| |---|---|---|---|---| |
| | W1 | **Offline** | Pre-computed hidden states from disk; trainer runs DP-only | none (`TargetHead` for logits) | `OfflineStream` | |
| | W2 | **In-process online** | Target on the same GPUs as draft (TP-collective); current default | `SGLangEagle3TargetModel` / `HFEagle3TargetModel` / `CustomEagle3TargetModel` | `OnlineStream` | |
| | W3 | **Disaggregated online** | Target on a separate SGLang server cluster; trainer talks HTTP | `SGLangServerEngine` | `RemoteStream` | |
| | W4 | **Train-with-decode** *(new in this revision — promoted from Phase 6)* | One long-lived SGLang server simultaneously **(a)** generates training data via prefill+aux and **(b)** serves real spec-decoding traffic. Trainer pushes draft weights into the same server every N steps so production traffic immediately benefits. | `SGLangServerEngine` (decode mode + weight push) | `OnlineStream` over static jsonl, or new `ServingTrafficStream` over a serving-traffic buffer | |
|
|
| W4 is what makes the no-Ray bet non-trivial: TorchSpec gets W4 "for free" because every |
| SglEngine actor already supports both modes. SpecForge needs to add three things explicitly |
| — `TargetEngine.update_draft_weights`, a decode-mode flag on `SGLangServerEngine`, and a |
| periodic weight-sync hook in `Trainer` — but **no actor topology change**: it stays one |
| torchrun-native trainer process talking HTTP to one always-on SGLang server. See §4.10. |
|
|
| --- |
|
|
| ## 2. Current State Gap Analysis |
|
|
| This section grounds the design in what is concretely wrong or missing in the current |
| codebase. Two flavors of gap: **(A) missing capabilities** — things the design assumes but |
| no code exists for; **(B) structural problems** — code that exists but the shape is wrong. |
|
|
| ### 2.1 Missing capabilities |
|
|
| | Gap | Evidence in current tree | Fixed by | |
| |---|---|---| |
| | **MLA-aware EAGLE3 draft** | Zero references to `DeepseekV3Config` / `Eagle3Deepseek*` outside test data. `configs/deepseek-v3-671b-eagle3.json` is a *Llama* draft mislabeled as DeepSeek (`model_type: "llama"`, `architectures: ["LlamaForCausalLMEagle3"]`). | Phase 1 #2 — port `deepseek_eagle.py` from TorchSpec, register via `@register_draft`. | |
| | **Backbone-agnostic DFlash** | `DFlashDraftModel` extends `Qwen3PreTrainedModel`; uses `Qwen3MLP`, `Qwen3DFlashAttention`, `Qwen3RMSNorm` directly (`specforge/modeling/draft/dflash.py:212`). All DFlash configs are Qwen3-only. | Phase 6 #25 — parameterize MLP/Norm/RoPE; add `llama_dflash.py`. | |
| | **Remote / disaggregated target** | `modeling/target/sglang_backend/` only does in-process SGLang (target loaded on same node as trainer). No HTTP client. Online mode for 671B targets is infeasible. | Phase 4 #17 — `SGLangServerEngine` over HTTP. | |
| | **Eval during training (proper)** | No `EvalCache`, no `simulated_acc_len`, no per-position-accuracy aggregation, no best-checkpoint tracking — grep is clean. | Phase 2 #9 — `Evaluator` + `EvalCache`. | |
| | **Checkpoint rotation / best tracking** | `train_eagle3.py:552:def save_checkpoints(...)` just writes every N steps. No `max_checkpoints`, no `best_checkpointed_iteration.txt`, no `meta.json`. | Phase 2 #8 — `CheckpointManager`. | |
| | **Resume that actually advances the stream** | No `stream.seek()`. Current `--resume` reloads weights but the dataloader yields from sample 0, silently re-training on the prefix. | §4.2 trainer pseudo-code + Phase 2 #5 — `HiddenStateStream.seek()`. | |
| | **Correct gradient accumulation** | Zero `no_sync()` calls anywhere in the codebase. FSDP all-reduces on every micro-step, defeating the point of `--accumulation-steps`. | Phase 2 #10 — `model.no_sync()` on non-sync micro-steps. | |
| | **Plugin registry for drafts** | Hardcoded `_model_mapping = {LlamaConfig: LlamaForCausalLMEagle3}` (`specforge/modeling/auto.py:35`). Code already has TODO: "should support lazy model mapping via registry". | Phase 1 #1 — `@register_draft` decorator. | |
| | **Train-with-decode (W4)** | No `update_draft_weights` on any target abstraction. `SGLangEagle3TargetModel.from_pretrained` hardcodes `disable_cuda_graph=True` (training-data-gen only); no notion of a long-lived dual-purpose server. No serving-traffic stream — `core/eagle3.py` only knows about static jsonl batches. No periodic weight push from trainer. | Phase 5 #21–#24 — `update_draft_weights` on `TargetEngine`, decode-mode flag on `SGLangServerEngine`, `ServingTrafficStream`, `Trainer._maybe_sync_draft_weights`. See §4.10. | |
| | **SGLang export with MLA weight map** | No `export/to_sglang.py`. MLA draft weights (`q_a_proj`, `kv_a_proj_with_mqa`, `kv_b_proj`) need explicit rename to whatever SGLang's spec-decoder loader expects. | Phase 3 #15 + `docs/export_weight_map_mla.md`. | |
| | **FSDP2 readiness** | `apply_fsdp` is FSDP1-only, inlined in scripts. No seam for FSDP2 swap. Locks the trainer to a pattern that composes poorly with `torch.compile`. | §4.9 — versioned `apply_fsdp` seam in Phase 2 (FSDP2 impl deferred to Phase 7 #31). | |
| | **Pydantic config + CLI** | `specforge/args.py` is 219 lines of argparse; architecture lives in JSON (`--draft-model-config`); training flags live in CLI. Two sources of truth, no validation. | Phase 3 #13/14 — one validated YAML per run. | |
| | **VLM unification** | Separate `QwenVLOnlineEagle3Model` class for VLM targets — parallel hierarchy that doesn't fit the abstraction. | Phase 4 #19 — typed `MediaInputs` in `TargetEngine`. | |
|
|
| ### 2.2 Structural / workflow problems |
|
|
| Code exists for these, but the shape is wrong: duplication, missing abstractions, |
| maintenance debt that compounds with each new arch. |
|
|
| | Problem | Concrete evidence | Fixed by | |
| |---|---|---| |
| | **God-script duplication** | `train_eagle3.py` = 1012 lines / 60 `add_argument`s; `train_dflash.py` = 562 lines / 44 `add_argument`s. Same argparse, same distributed init, same checkpoint save, same logging — copy-pasted twice. Total ~1574 lines across two top-level scripts. | Phase 2 #6 — single `Trainer`, strategy dispatch. Reduces to one ~70-line loop + strategy classes. | |
| | **Target-model duplication in modeling/** | `modeling/target/eagle3_target_model.py` (873 lines) and `modeling/target/dflash_target_model.py` (315 lines) are parallel hierarchies with no shared base. | Phase 2 #4 — single `TargetEngine` ABC; three concrete impls. | |
| | **Online ≠ class, offline ≠ class** | `core/eagle3.py` is the "online" wrapper (606 lines); offline is a different path through `OfflineEagle3Dataset` + `TargetHead`. The trainer must branch by mode. | Phase 2 #5 — both become `HiddenStateStream` implementations. Trainer takes one iterator regardless. | |
| | **Arch dispatch by 4-file shotgun** | Adding a new arch today touches `modeling/auto.py:35`, `modeling/auto.py:88`, `modeling/auto.py:134`, plus `modeling/draft/__init__.py`. | Phase 1 #1 — one file, one decorator. | |
| | **30 hand-maintained shell scripts** | `examples/run_*_online.sh` × 30. Defaults drift between them; new features need 30 edits. | Phase 3 #14 — 30 YAML files derived from one schema; shell scripts collapse to `specforge train --config ...`. | |
| | **No data prefetch overlap** | `core/eagle3.py` runs target forward synchronously inside the training step. Draft backward blocks on target forward. | Phase 2 #5 — `OnlineStream.prefetch_factor` overlaps producer and consumer. | |
| | **Misleading config in repo** | `configs/deepseek-v3-671b-eagle3.json` declares `model_type: "llama"` for a DeepSeek target — silently wrong if read as "MLA Eagle3 for V3." | §4.7 — rename to `..._llama_draft.json`, add real MLA config, deprecation warning. | |
| | **No numerical-equivalence gate for refactors** | Nothing prevents Phase 2's `Trainer` extraction from quietly changing loss curves. | §10.1 — `atol/rtol` gate at fixed steps on every PR touching `core/`, `training/`, `models/drafts/`. | |
| | **No FSDP all-reduce verification** | No profiler check that accumulation actually skips sync. `--accumulation-steps` flag exists but the sync isn't suppressed (no `no_sync()`) — same number of all-reduces as without accumulation. | §10.3 — one-time profiler check; one all-reduce per `optimizer.step()`. | |
| | **`Eagle3DraftModel.backbone()` ABC return shape** | Returns `Tensor`; KV cache mutates `past_key_values` in place (HF `DynamicCache` pattern). DFlash respects this. A naive MLA port that returns `(out, k, v)` would break the ABC. | §4.7 — explicit "use `DynamicCache.update`, return `Tensor`" rule for MLA. | |
|
|
| ### 2.3 The two highest-leverage chunks |
|
|
| If only two things ship this quarter, do these: |
|
|
| **Chunk 1 — Phase 1 (Week 1-2): `@register_draft` + MLA draft + Kimi-K2.5 TP/SP smoke test.** |
| Unblocks the actual feature ask (DFlash + MLA Eagle3 on SGLang). Additive — old |
| `_model_mapping` dicts stay as a fallback. ~3 new files, ~10 lines of edits to `auto.py`. |
| The TP/SP smoke test (§4.7) is the de-risking step that turns Kimi-K2.5 from "should |
| work" into a verified deliverable. |
| |
| **Chunk 2 — Phase 2 (Week 3-6): `TargetEngine` + `HiddenStateStream` + `Trainer` together.** |
| These ship as a unit. Extracting one without the others either leaves the trainer talking |
| directly to target models (Phase 4 has to undo it) or leaves the stream abstraction |
| without a consumer (can't be tested end-to-end). The Phase 2 acceptance gate is that the |
| legacy shells, now calling new internals, match the old loss curves to numerical- |
| equivalence tolerance (§10.1). Pass that, and Phase 4 becomes drop-in plugins |
| (`SGLangServerEngine`, `RemoteStream`) without re-plumbing. |
| |
| Everything else (export tool, CLI, VLM cleanup, FSDP2, WSD, Mooncake) is debt cleanup |
| that compounds — important but not blocking the stated goal. |
| |
| --- |
| |
| ## 3. Design Principles |
| |
| 1. **One trainer, many drafts, many targets.** Config-driven dispatch via a plugin registry. |
| New arch = new file, not a new script. |
| |
| 2. **The target is an interface, not a process.** Same trainer whether the target is in-process |
| HF, a separate SGLang server, or a remote cluster. |
| |
| 3. **Three data modes behind one abstraction.** Offline-cached, online-local, online-remote — |
| all consumed via the same `HiddenStateStream` iterator. But respect the asymmetry: online |
| streams own GPU resources and backpressure; offline streams are pure readers. |
| |
| 4. **Strategies, not forks.** EAGLE3 (TTT unroll) and DFlash (block-causal) are two |
| `DraftTrainStrategy` implementations sharing the trainer. |
| |
| 5. **Keep what works.** Don't rewrite battle-tested code for symmetry. |
| |
| --- |
| |
| ## 4. Target Architecture |
| |
| ### 4.1 Module layout |
| |
| ``` |
| specforge/ |
| ├── config/ # Structured configuration |
| │ ├── schema.py # Pydantic models (Config, ModelConfig, DatasetConfig, ...) |
| │ ├── loader.py # YAML load + merge + CLI override + validation |
| │ └── draft_configs/ # Draft model JSONs (moved from top-level configs/) |
| │ ├── llama3_8b_eagle3.json |
| │ ├── qwen3_8b_eagle3.json |
| │ ├── qwen3_8b_eagle3_mla.json # NEW |
| │ ├── kimi_k25_eagle3_mla.json # NEW |
| │ ├── deepseek_v3_671b_eagle3.json # NEW — real MLA config |
| │ └── ... |
| │ |
| ├── core/ # Core algorithms (preserved) |
| │ ├── eagle3.py # Eagle3Model (renamed from OnlineEagle3Model) |
| │ ├── dflash.py # DFlashModel (renamed from OnlineDFlashModel) |
| │ ├── loss.py # Triton LogSoftmaxLoss (UNCHANGED) |
| │ └── adapters.py # BackendAdapter / UspAdapter (UNCHANGED) |
| │ |
| ├── models/ |
| │ ├── drafts/ # Draft model plugin registry (see §4.2 for the registry + naming contract) |
| │ │ ├── __init__.py # DRAFT_REGISTRY + @register_draft decorator |
| │ │ ├── base.py # Eagle3DraftModel ABC (from modeling/draft/base.py) |
| │ │ ├── <model>_<strategy>.py # one file per arch — ${model_name}_${strategy} |
| │ │ │ # e.g. llama_eagle3, deepseek_eagle3, qwen3_dflash |
| │ │ └── auto.py # AutoEagle3DraftModel (registry-backed, no hardcoded dicts) |
| │ │ |
| │ └── targets/ # Target engine abstraction |
| │ ├── base.py # TargetEngine ABC + TargetOutput dataclass |
| │ ├── hf_engine.py # In-process HF target (lazy import) |
| │ ├── sglang_engine.py # In-process SGLang target (lazy import) |
| │ ├── sglang_server_engine.py # SGLang-as-service over HTTP (NEW) |
| │ ├── custom_engine.py # Custom TP backend (existing custom_backend/) |
| │ └── target_head.py # TargetHead for offline logits (existing) |
| │ |
| ├── data/ |
| │ ├── streams/ # HiddenStateStream abstraction |
| │ │ ├── base.py # HiddenStateStream protocol |
| │ │ ├── online.py # In-process target generates hidden states |
| │ │ ├── offline.py # Pre-computed hidden states from disk |
| │ │ └── remote.py # Target on separate SGLang server (NEW) |
| │ ├── template.py # TEMPLATE_REGISTRY (UNCHANGED) |
| │ ├── parse.py # Parsers + KimiK25Parser, MiniMaxParser (NEW) |
| │ ├── preprocessing.py # build_eagle3_dataset etc. (UNCHANGED) |
| │ ├── collator.py # DataCollatorWithPadding (extracted from utils.py) |
| │ └── cache.py # Tokenization cache + eval cache (NEW) |
| │ |
| ├── training/ |
| │ ├── trainer.py # Trainer: unified training loop |
| │ ├── strategies/ |
| │ │ ├── __init__.py # DraftTrainStrategy protocol |
| │ │ ├── eagle3_ttt.py # Eagle3 TTT unroll + forward-KL |
| │ │ └── dflash_block.py # DFlash block-causal + anchor sampling |
| │ ├── optimizer.py # BF16Optimizer (UNCHANGED) |
| │ ├── lr_scheduler.py # CosineWarmup + WSD scheduler (NEW) |
| │ ├── checkpoint.py # CheckpointManager (NEW) |
| │ ├── fsdp.py # apply_fsdp (FSDP1 SHARD_GRAD_OP + future FSDP2) |
| │ └── distributed.py # init_distributed etc. (moved from specforge/distributed.py) |
| │ |
| ├── eval/ # Evaluation system (NEW) |
| │ ├── evaluator.py # Evaluator: eval loop + metric aggregation |
| │ ├── cache.py # EvalCache: MD5-keyed disk cache |
| │ └── metrics.py # simulated_acc_len, avg_loss, avg_acc |
| │ |
| ├── export/ # Model export tools (NEW) |
| │ ├── to_hf.py # FSDP checkpoint → HF format + vocab pruning |
| │ └── to_sglang.py # Checkpoint → SGLang spec-decoder layout |
| │ |
| ├── tracker.py # UNCHANGED |
| ├── utils.py # UNCHANGED |
| │ |
| ├── cli.py # `specforge train|prepare|export|eval` (NEW) |
| └── __init__.py |
| |
| scripts/ |
| ├── legacy/ # Old scripts as thin shims |
| │ ├── train_eagle3.py |
| │ └── train_dflash.py |
| ├── prepare_data.py # UNCHANGED |
| ├── prepare_hidden_states.py # UNCHANGED |
| └── regenerate_train_data.py # UNCHANGED |
| ``` |
| |
| ### 4.2 Key abstractions |
| |
| #### Draft model registry |
| |
| This is the single source of truth for the `models/drafts/` layout sketched in §4.1. |
| One file per architecture, named `${model_name}_${strategy}` (strategy ∈ `eagle3`, |
| `dflash`, `domino`, `dspark`) — e.g. `llama_eagle3.py`, `deepseek_eagle3.py`, |
| `qwen3_dflash.py`; the registry key matches the filename stem. |
| |
| ```python |
| # models/drafts/__init__.py |
| DRAFT_REGISTRY: dict[str, type] = {} |
| |
| def register_draft(name: str): |
| """Decorator. New arch = new file + @register_draft("deepseek_v3_eagle3").""" |
| def wrapper(cls): |
| DRAFT_REGISTRY[name] = cls |
| return cls |
| return wrapper |
| |
| # models/drafts/llama_eagle3.py |
| @register_draft("llama_eagle3") |
| class LlamaForCausalLMEagle3(Eagle3DraftModel): |
| ... |
| |
| # models/drafts/deepseek_eagle3.py — MLA (NEW) |
| @register_draft("deepseek_v3_eagle3") |
| class Eagle3DeepseekV2ForCausalLM(Eagle3DraftModel): |
| config_class = DeepseekV3Config |
| ... |
| |
| # models/drafts/qwen3_dflash.py — DFlash (existing) |
| @register_draft("qwen3_dflash") |
| class DFlashDraftModel(Eagle3DraftModel): |
| ... |
| ``` |
| |
| The `Eagle3DraftModel` ABC is preserved exactly as-is. Its interface is already well-designed: |
| |
| ```python |
| class Eagle3DraftModel(PreTrainedModel, ABC): |
| def embed_input_ids(self, input_ids: Tensor) -> Tensor: ... |
| def project_hidden_states(self, hidden_states: Tensor) -> Tensor: ... |
| def backbone(self, input_embeds, hidden_states, cache_hidden, |
| attention_mask, position_ids, past_key_values=None, |
| use_cache=True) -> Tensor: ... |
| def compute_logits(self, hidden_states: Tensor) -> Tensor: ... |
| # Concrete: load_embedding, freeze_embedding, load_vocab_mapping, t2d/d2t |
| ``` |
| |
| `AutoDraftModelConfig.from_file()` is updated to look up `DRAFT_REGISTRY` first, fall back to |
| HF config type dispatch for backward compatibility. |
| |
| #### Target engine |
| |
| ```python |
| # models/targets/base.py |
| @dataclass |
| class TargetOutput: |
| aux_hidden_states: torch.Tensor # [batch, seq, hidden * num_aux_layers] — raw concat (NOT projected) |
| target_logits: torch.Tensor # [batch, seq, vocab] — pre-softmax logits in draft vocab space |
| loss_mask: torch.Tensor # [batch, seq] |
| input_ids: torch.Tensor # [batch, seq] |
| attention_mask: torch.Tensor # [batch, seq] |
| last_hidden_states: Optional[torch.Tensor] = None |
| |
| class TargetEngine(ABC): |
| @abstractmethod |
| def generate_train_data( |
| self, |
| input_ids: Tensor, |
| attention_mask: Tensor, |
| loss_mask: Tensor, |
| media: Optional[MediaInputs] = None, # typed VLM payload (pixel_values, image_grid_thw, ...) |
| ) -> TargetOutput: ... |
| |
| @property |
| @abstractmethod |
| def aux_layer_ids(self) -> list[int]: ... |
| |
| def set_aux_hidden_states_layers(self, layers: list[int]) -> None: ... |
| ``` |
| |
| Naming: |
| - `target_logits` (not `target`) — the previous overload of "target" with "target model" was confusing. |
| - `aux_hidden_states` — raw concatenated multi-layer hidden states from the target. |
| The projection (3·hidden → hidden) **stays inside the draft model** (`project_hidden_states`), |
| so streams and engines never need to know the draft's hidden_size. |
| |
| This is a rename + mild generalization of the existing `Eagle3TargetModel` + |
| `Eagle3TargetOutput`. The existing SGLang/HF/Custom implementations become concrete engines |
| with no API changes. VLM is handled via `**media_kwargs` — no separate model class needed. |
| |
| Lazy imports ensure SGLang is never imported unless `backend="sglang"`: |
| |
| ```python |
| def get_target_engine(backend: str, **kwargs) -> TargetEngine: |
| if backend == "sglang": |
| from specforge.models.targets.sglang_engine import SGLangTargetEngine |
| return SGLangTargetEngine(**kwargs) |
| elif backend == "sglang_server": |
| from specforge.models.targets.sglang_server_engine import SGLangServerEngine |
| return SGLangServerEngine(**kwargs) |
| elif backend == "hf": |
| from specforge.models.targets.hf_engine import HFTargetEngine |
| return HFTargetEngine(**kwargs) |
| elif backend == "custom": |
| from specforge.models.targets.custom_engine import CustomTargetEngine |
| return CustomTargetEngine(**kwargs) |
| raise ValueError(f"Unknown target backend: {backend}") |
| ``` |
| |
| #### HiddenStateStream |
|
|
| ```python |
| # data/streams/base.py |
| class HiddenStateStream(ABC): |
| """Produces TrainBatch instances for the draft trainer.""" |
| |
| def setup(self, draft_config) -> None: |
| """Inform stream of draft requirements (aux layers, vocab mapping, etc.).""" |
| pass |
| |
| @abstractmethod |
| def __iter__(self) -> Iterator[TrainBatch]: ... |
| |
| def seek(self, step: int) -> None: |
| """Resume from step N. Required for checkpoint resume to behave correctly. |
| |
| Default behavior (provided by base class): advance the iterator by `step` |
| batches. Implementations with random-access storage (e.g. OfflineStream) |
| should override to seek directly without consuming. |
| """ |
| |
| def teardown(self) -> None: |
| """Release GPU / network resources.""" |
| pass |
| |
| @dataclass |
| class TrainBatch: |
| input_ids: torch.Tensor |
| attention_mask: torch.Tensor |
| target_logits: torch.Tensor # target logits on draft vocab (pre-softmax) |
| loss_mask: torch.Tensor |
| aux_hidden_states: torch.Tensor # raw concat of target aux layers — draft projects internally |
| position_ids: Optional[torch.Tensor] = None |
| ``` |
|
|
| Three implementations: |
|
|
| | Stream | Source | GPU cost | Notes | |
| |--------|--------|----------|-------| |
| | `OnlineStream` | In-process `TargetEngine` | High (target on same GPUs) | Current "online" mode; supports `prefetch_factor` to overlap target inference with draft training | |
| | `OfflineStream` | Pre-computed .pt files | None | Current "offline" mode; overrides `seek()` for O(1) resume | |
| | `RemoteStream` | SGLang server over HTTP | None locally | NEW — target on separate node(s); supports `prefetch_factor` and request batching | |
|
|
| `OnlineStream` owns the target engine lifecycle and handles TP→DP batch sharding |
| (currently done manually in `train_eagle3.py:get_dp_data_shard_from_tp`). Exposes |
| `prefetch_factor: int = 2` so the next batch's target forward overlaps the current |
| batch's draft backward. |
|
|
| `OfflineStream` wraps the existing `OfflineEagle3Dataset` + `TargetHead` path. Overrides |
| `seek()` to jump directly to a sample index without iterating intermediate files. |
|
|
| `RemoteStream` sends tokenized batches to an SGLang server, receives hidden states back. |
| This covers the "671B target on dedicated GPUs" case without Ray/Mooncake. Should also |
| expose `prefetch_factor` (network roundtrip dominates without it) and a `max_in_flight` |
| bound so back-pressure is explicit. |
|
|
| #### DraftTrainStrategy |
|
|
| ```python |
| # training/strategies/__init__.py |
| class DraftTrainStrategy(ABC): |
| """Encapsulates one training algorithm's forward + loss computation.""" |
| |
| @abstractmethod |
| def forward_and_loss( |
| self, |
| draft_model: Eagle3DraftModel, |
| batch: TrainBatch, |
| ) -> tuple[torch.Tensor, dict[str, float]]: |
| """Returns (loss, metrics_dict). Target data is already in batch.""" |
| ... |
| |
| @abstractmethod |
| def build_model(self, draft_model, **kwargs) -> nn.Module: |
| """Wrap draft_model in the strategy-specific training wrapper.""" |
| ... |
| |
| def fsdp_wrap_policy(self) -> Optional[Callable]: |
| """Optional: return an FSDP auto-wrap policy tailored to this strategy's wrapper. |
| |
| The trainer applies this if non-None; otherwise falls back to a default |
| transformer-block wrap policy. Lets strategies declare wrapping intent |
| instead of the trainer second-guessing. |
| """ |
| return None |
| ``` |
|
|
| Two implementations: |
|
|
| - **`Eagle3TTTStrategy`**: wraps draft in `Eagle3Model` (current `core/eagle3.py`), runs |
| TTT unroll, forward-KL loss via `LogSoftmaxLoss`. The existing `Eagle3Model.forward()` |
| signature maps directly. |
|
|
| - **`DFlashBlockStrategy`**: wraps draft in `DFlashModel` (current `core/dflash.py`), does |
| anchor sampling + block-causal CE. The existing `OnlineDFlashModel.forward()` maps directly. |
|
|
| The strategy does NOT touch the target engine — batch already contains materialized tensors. |
| This is a hard boundary. |
|
|
| #### Trainer |
|
|
| ```python |
| # training/trainer.py |
| class Trainer: |
| def __init__(self, config: Config): |
| self.config = config |
| self.strategy = self._build_strategy() |
| self.checkpoint_mgr = CheckpointManager(config.training) |
| self.evaluator = Evaluator(config.eval) if config.eval.enabled else None |
| self.tracker = create_tracker(config.logging) |
| |
| def train(self): |
| # 1. Distributed setup |
| init_distributed(...) |
| |
| # 2. Build draft model from registry |
| draft_config = AutoDraftModelConfig.from_file(self.config.model.draft_model_config) |
| draft_model = DRAFT_REGISTRY[draft_config.architectures[0]].from_config(draft_config) |
| |
| # 3. Build strategy-specific wrapper |
| model = self.strategy.build_model(draft_model, ...) |
| |
| # 4. Apply FSDP (strategy may override wrap policy) |
| model = apply_fsdp( |
| model, |
| self.config.training, |
| wrap_policy=self.strategy.fsdp_wrap_policy(), |
| ) |
| |
| # 5. Build optimizer |
| optimizer = BF16Optimizer(model, ...) |
| |
| # 6. Build data stream |
| stream = self._build_stream() |
| stream.setup(draft_config) |
| |
| # 7. Resume if needed — note: stream.seek() is required for correctness. |
| # enumerate(stream, start=start_step) only changes the counter; without |
| # seek(), the stream still yields batch 0 first, silently re-training on it. |
| start_step = 0 |
| if self.config.training.resume: |
| start_step = self.checkpoint_mgr.load(model, optimizer) |
| stream.seek(start_step) |
| |
| # 8. Training loop |
| accum = self.config.training.accumulation_steps |
| for step, batch in enumerate(stream, start=start_step): |
| if step >= self.config.training.max_steps: |
| break |
| |
| # Forward + loss (strategy-specific) |
| loss, metrics = self.strategy.forward_and_loss(model, batch) |
| scaled_loss = loss / accum |
| |
| # Skip gradient sync on all but the final micro-step of an accumulation |
| # window. With FSDP this is the difference between one all-reduce per |
| # micro-batch and one per optimizer step — non-trivial on large models. |
| is_sync_step = ((step + 1) % accum == 0) |
| sync_ctx = nullcontext() if is_sync_step else model.no_sync() |
| with sync_ctx: |
| scaled_loss.backward() |
| |
| if is_sync_step: |
| optimizer.step() |
| |
| # Logging |
| if (step + 1) % self.config.training.log_interval == 0: |
| self.tracker.log(metrics, step=step) |
| |
| # Eval |
| if self.evaluator and (step + 1) % self.config.eval.interval == 0: |
| eval_metrics = self.evaluator.run( |
| lambda b: self.strategy.forward_and_loss(model, b), |
| self._eval_stream, |
| ) |
| self.checkpoint_mgr.update_best(step, eval_metrics) |
| self.tracker.log(eval_metrics, step=step) |
| |
| # Save |
| if (step + 1) % self.config.training.save_interval == 0: |
| self.checkpoint_mgr.save(step, model, optimizer) |
| |
| stream.teardown() |
| self.tracker.close() |
| ``` |
|
|
| This is ~70 lines of logic. Everything else is behind interfaces. |
|
|
| Two non-obvious correctness points the pseudo-code makes explicit: |
|
|
| 1. **Resume requires `stream.seek(start_step)`.** Without it, `enumerate(stream, start=N)` |
| only relabels the counter — the iterator still yields batch 0 first, so resume |
| silently re-trains on the prefix. |
| |
| 2. **Gradient accumulation needs `model.no_sync()` on non-sync steps.** Otherwise FSDP |
| all-reduces on every micro-batch, defeating the point of accumulation. |
|
|
| ### 4.3 Checkpoint manager |
|
|
| ```python |
| # training/checkpoint.py |
| class CheckpointManager: |
| def __init__(self, config: TrainingConfig): |
| self.checkpoint_dir = Path(config.checkpoint_dir) |
| self.max_checkpoints = config.max_checkpoints # 0 = keep all |
| self.best_score = -float("inf") |
| |
| def save(self, step: int, model, optimizer, meta: dict = None): |
| """Save model + optimizer + LR state + RNG + meta.json.""" |
| step_dir = self.checkpoint_dir / f"iter_{step + 1:07d}" |
| # torch.distributed.checkpoint.save for model, optimizer |
| # rank-0: rng.pt, meta.json (step, timestamp, global_step, world_size) |
| self._update_latest(step + 1) |
| self._rotate() |
| |
| def load(self, model, optimizer=None, continual=False) -> int: |
| """Load latest or specified checkpoint. Returns start_step.""" |
| step_id = self._read_latest() |
| # torch.distributed.checkpoint.load |
| # If continual: skip optimizer/LR/RNG, only restore weights + rebuild FP32 master |
| return step_id |
| |
| def update_best(self, step: int, eval_metrics: dict): |
| """Track best checkpoint by simulated_acc_len.""" |
| score = eval_metrics.get("simulated_acc_len", eval_metrics.get("avg_acc", 0)) |
| if score > self.best_score: |
| self.best_score = score |
| self._write_best(step + 1, eval_metrics) |
| |
| def _rotate(self): |
| """Keep only max_checkpoints newest iter_* directories.""" |
| if self.max_checkpoints <= 0: |
| return |
| dirs = sorted(self.checkpoint_dir.glob("iter_*")) |
| for d in dirs[:-self.max_checkpoints]: |
| shutil.rmtree(d) |
| |
| def _update_latest(self, step_id: int): |
| (self.checkpoint_dir / "latest_checkpointed_iteration.txt").write_text(str(step_id)) |
| |
| def _write_best(self, step_id: int, metrics: dict): |
| (self.checkpoint_dir / "best_checkpointed_iteration.txt").write_text(str(step_id)) |
| (self.checkpoint_dir / "best_meta.json").write_text(json.dumps(metrics, indent=2)) |
| ``` |
|
|
| ### 4.4 Evaluation system |
|
|
| ```python |
| # eval/evaluator.py |
| class Evaluator: |
| def __init__(self, config: EvalConfig): |
| self.cache = EvalCache(config.cache_dir) |
| self.micro_batch_size = config.micro_batch_size |
| |
| def run(self, forward_fn, eval_stream: HiddenStateStream) -> dict: |
| """Run full eval pass, return aggregated metrics. |
| |
| Per-position accuracy must be averaged *across all batches first*, then fed |
| into the geometric sum. Treating each batch's per-position vector as if it |
| were positions makes simulated_acc_len batch-size-dependent. |
| """ |
| total_loss_x_tokens = 0.0 |
| total_tokens = 0 |
| # Sum and count per draft position (TTT step), aggregated across the whole pass. |
| per_pos_acc_sum: Optional[torch.Tensor] = None # shape [ttt_length] |
| per_pos_acc_count: Optional[torch.Tensor] = None # shape [ttt_length] |
| |
| for batch in self._iter_micro_batches(eval_stream): |
| with torch.no_grad(): |
| loss, metrics = forward_fn(batch) |
| total_loss_x_tokens += metrics["loss"] * metrics["num_tokens"] |
| total_tokens += metrics["num_tokens"] |
| |
| ppa = metrics.get("per_position_acc") # [ttt_length], weighted by num_tokens |
| ppc = metrics.get("per_position_count") # [ttt_length], token counts per position |
| if ppa is not None: |
| if per_pos_acc_sum is None: |
| per_pos_acc_sum = torch.zeros_like(ppa) |
| per_pos_acc_count = torch.zeros_like(ppc) |
| per_pos_acc_sum += ppa * ppc # accumulate weighted sum |
| per_pos_acc_count += ppc |
| |
| # Aggregate first, then geometric-sum. |
| per_position_acc = (per_pos_acc_sum / per_pos_acc_count.clamp_min(1)).tolist() |
| return { |
| "eval/avg_loss": total_loss_x_tokens / max(total_tokens, 1), |
| "eval/avg_acc": float(per_position_acc[0]) if per_position_acc else 0.0, |
| "eval/simulated_acc_len": self._simulated_acc_len(per_position_acc), |
| } |
| |
| @staticmethod |
| def _simulated_acc_len(per_position_acc: list[float]) -> float: |
| """E[accepted tokens] = acc_0 + acc_0*acc_1 + acc_0*acc_1*acc_2 + ... |
| |
| `per_position_acc` is the *aggregated* per-position accuracy across the full |
| eval set, length = ttt_length. Not a list of per-batch vectors. |
| """ |
| cumulative = 1.0 |
| total = 0.0 |
| for acc in per_position_acc: |
| cumulative *= acc |
| total += cumulative |
| return total |
| |
| # eval/cache.py |
| class EvalCache: |
| """Disk cache for pre-computed eval hidden states. |
| |
| Cache key must cover everything that would change the cached tensors: |
| eval data, target model identity & revision, tokenizer, chat template, |
| aux layer ids, and sequence length. Missing any of these silently serves |
| stale data after a target swap or template change. |
| """ |
| |
| def __init__(self, cache_dir: str): |
| self.cache_dir = Path(cache_dir) |
| |
| def cache_key( |
| self, |
| eval_path: str, |
| target_path: str, |
| target_revision: str, |
| tokenizer_path: str, |
| chat_template: str, |
| aux_layer_ids: list[int], |
| max_seq_len: int, |
| ) -> str: |
| content = "|".join([ |
| eval_path, |
| target_path, |
| target_revision or "", |
| tokenizer_path, |
| chat_template, |
| ",".join(map(str, aux_layer_ids)), |
| str(max_seq_len), |
| ]) |
| return hashlib.md5(content.encode()).hexdigest()[:12] |
| |
| def try_load(self, key: str) -> Optional[list]: |
| path = self.cache_dir / "eval_cache" / key |
| if path.exists(): |
| return [torch.load(f) for f in sorted(path.glob("rank_*.pt"))] |
| return None |
| |
| def save(self, key: str, rank: int, data: list): |
| path = self.cache_dir / "eval_cache" / key |
| path.mkdir(parents=True, exist_ok=True) |
| torch.save(data, path / f"rank_{rank:04d}.pt") |
| ``` |
|
|
| Note: `data/cache.py` (tokenization cache) and `eval/cache.py` (eval hidden-state cache) |
| are deliberately separate — they key on different things and live at different lifecycle |
| points. Don't merge them. |
|
|
| ### 4.5 Structured config |
|
|
| ```python |
| # config/schema.py |
| from pydantic import BaseModel, Field |
| from typing import Optional, Literal |
| |
| class ModelConfig(BaseModel): |
| target_model_path: str |
| draft_model_config: str # path to JSON |
| target_backend: Literal["sglang", "hf", "custom", "sglang_server"] = "sglang" |
| trust_remote_code: bool = False |
| embedding_key: str = "model.embed_tokens.weight" |
| lm_head_key: str = "lm_head.weight" |
| is_vlm: bool = False |
| |
| class DatasetConfig(BaseModel): |
| train_data_path: str |
| eval_data_path: str = "" |
| train_hidden_states_path: str = "" # non-empty = offline mode |
| chat_template: str = "llama3" |
| max_length: int = 2048 |
| train_only_last_turn: bool = False |
| num_proc: int = 8 |
| |
| class TrainingConfig(BaseModel): |
| strategy: Literal["eagle3", "dflash"] = "eagle3" |
| num_epochs: int = 1 |
| max_steps: int = 10000 |
| batch_size: int = 1 |
| learning_rate: float = 3e-4 |
| warmup_ratio: float = 0.015 |
| max_grad_norm: float = 0.5 |
| accumulation_steps: int = 1 |
| ttt_length: int = 7 # Eagle3-specific |
| block_size: int = 16 # DFlash-specific |
| num_anchors: int = 512 # DFlash-specific |
| loss_decay_gamma: Optional[float] = None |
| attention_backend: Literal["sdpa", "flex_attention", "fa", "usp"] = "sdpa" |
| fsdp_strategy: Literal["NO_SHARD", "SHARD_GRAD_OP", "FULL_SHARD", "HYBRID_SHARD"] = "SHARD_GRAD_OP" |
| fsdp_version: Literal[1, 2] = 1 # 2 = FSDP2 (PT 2.4+); see §4.9 |
| compile_model: bool = False |
| tp_size: int = 1 |
| sp_ulysses_size: int = 1 |
| sp_ring_size: int = 1 |
| save_interval: int = 500 |
| log_interval: int = 10 |
| max_checkpoints: int = 5 # 0 = keep all |
| checkpoint_dir: str = "checkpoints" |
| resume: bool = False |
| seed: int = 42 |
| |
| class EvalConfig(BaseModel): |
| enabled: bool = False |
| interval: int = 500 |
| micro_batch_size: int = 4 |
| cache_dir: str = "eval_cache" |
| |
| class LRConfig(BaseModel): |
| decay_style: Literal["cosine", "WSD"] = "cosine" |
| wsd_decay_steps: int = 0 |
| wsd_decay_style: Literal["linear", "cosine", "exponential"] = "cosine" |
| |
| class LoggingConfig(BaseModel): |
| report_to: Literal["wandb", "tensorboard", "swanlab", "mlflow", "none"] = "wandb" |
| wandb_project: str = "specforge" |
| wandb_run_name: str = "" |
| |
| class SGLangConfig(BaseModel): |
| attention_backend: str = "flashinfer" |
| # ... other SGLang server args |
| |
| class Config(BaseModel): |
| model: ModelConfig |
| dataset: DatasetConfig |
| training: TrainingConfig = Field(default_factory=TrainingConfig) |
| eval: EvalConfig = Field(default_factory=EvalConfig) |
| lr: LRConfig = Field(default_factory=LRConfig) |
| logging: LoggingConfig = Field(default_factory=LoggingConfig) |
| sglang: SGLangConfig = Field(default_factory=SGLangConfig) |
| output_dir: str = "output" |
| cache_dir: str = "cache" |
| ``` |
|
|
| YAML example for MLA Eagle3 training: |
|
|
| ```yaml |
| model: |
| target_model_path: Qwen/Qwen3-8B |
| draft_model_config: specforge/config/draft_configs/qwen3_8b_eagle3_mla.json |
| target_backend: sglang |
| |
| dataset: |
| train_data_path: data/train.jsonl |
| eval_data_path: data/eval.jsonl |
| chat_template: qwen3 |
| max_length: 4096 |
| |
| training: |
| strategy: eagle3 |
| max_steps: 20000 |
| batch_size: 2 |
| learning_rate: 3e-4 |
| ttt_length: 7 |
| attention_backend: flex_attention |
| accumulation_steps: 2 |
| tp_size: 4 |
| max_checkpoints: 3 |
| |
| eval: |
| enabled: true |
| interval: 1000 |
| |
| output_dir: runs/qwen3_8b_mla_eagle3 |
| ``` |
|
|
| CLI override: `specforge train --config config.yaml --training.learning_rate=1e-4` |
|
|
| ### 4.6 WSD learning rate scheduler |
|
|
| > **Deferred to Phase 6** (was Phase 2 in the original draft). Cosine warmup is sufficient |
| > for typical draft training runs; WSD is most useful for continued pretraining and long |
| > stable-phase runs, which aren't on the immediate roadmap. |
|
|
| Add to `specforge/training/lr_scheduler.py`: |
|
|
| ```python |
| class WSDScheduler(_LRScheduler): |
| """Warmup → Stable → Decay schedule. |
| |
| Stable phase holds max LR until (total_steps - wsd_decay_steps). |
| Decay phase applies linear/cosine/exponential decay to min_lr. |
| """ |
| def __init__(self, optimizer, total_steps, warmup_steps, min_lr=0.0, |
| wsd_decay_steps=0, wsd_decay_style="cosine", last_epoch=-1): |
| ... |
| ``` |
|
|
| Wired via `LRConfig.decay_style == "WSD"` in the trainer. |
|
|
| ### 4.7 MLA Eagle3 draft model |
|
|
| Port from TorchSpec's `deepseek_eagle.py`. Key design decisions: |
|
|
| **KV cache in TTT context**: During Eagle3 TTT unroll, cache expanded K/V (not compressed |
| latent). This matches TorchSpec's approach and avoids requiring MLA-specific cache logic in |
| `core/eagle3.py`. The MLA compression only happens during projection (forward pass), not in |
| the cache. |
|
|
| **Cache integration**: use the HF `DynamicCache` pattern (mutate `past_key_values` in place |
| via `past_key_values.update(k, v, layer_idx, cache_kwargs)`) — same as the existing |
| `DFlashDraftModel` and `Eagle3DraftModel.backbone()` ABC, which returns `Tensor` (hidden |
| states) and never a `(Tensor, k, v)` tuple. Returning a tuple would break the base ABC. |
|
|
| ```python |
| # models/drafts/deepseek_eagle3.py |
| @register_draft("deepseek_v3_eagle3") |
| class Eagle3DeepseekV2ForCausalLM(Eagle3DraftModel): |
| config_class = DeepseekV3Config |
| |
| class DeepSeekMLAAttention(nn.Module): |
| """MLA attention with Q/KV LoRA, decoupled RoPE.""" |
| def __init__(self, config: DeepseekV3Config): |
| # Q path: q_a_proj → RMSNorm → q_b_proj (if q_lora_rank) |
| # KV path: kv_a_proj_with_mqa → split(kv_compressed, k_rope_raw) |
| # → kv_a_layernorm → kv_b_proj → split(k_nope, value) |
| # RoPE: interleaved rotation on qk_rope_head_dim slice only |
| # o_proj: num_heads * v_head_dim → hidden_size |
| ... |
| |
| def forward(self, hidden_states, past_key_values=None, |
| attention_mask=None, position_ids=None, cache_position=None, |
| use_cache=False): |
| # 1. Project Q (with optional LoRA) |
| # 2. Project KV (compressed → expand via kv_b_proj) |
| # 3. Split k_nope from kv_b_proj, k_rope_raw from kv_a_proj |
| # 4. Apply interleaved RoPE to q_rope and k_rope slices |
| # 5. Expand k_rope (MQA → MHA) across heads |
| # 6. Concat [k_nope, k_rope] → full expanded key |
| # 7. If use_cache and past_key_values is not None: |
| # k, v = past_key_values.update(k, v, self.layer_idx, cache_kwargs) |
| # 8. Attention (SDPA or FlexAttention) |
| # 9. Return attn_output — k/v stored in the mutable DynamicCache |
| ... |
| ``` |
|
|
| Config fields consumed from `DeepseekV3Config`: |
| - `q_lora_rank` — Q low-rank dim (None = dense Q) |
| - `kv_lora_rank` — KV compressed dim |
| - `qk_nope_head_dim` — per-head key dim without RoPE |
| - `qk_rope_head_dim` — per-head key dim with RoPE |
| - `v_head_dim` — per-head value dim (may differ from key dim) |
| - `rope_scaling` — YaRN parameters |
|
|
| **Memory budget for expanded TTT cache**. Per token per draft layer: |
|
|
| key_bytes_per_token = num_heads * (qk_nope_head_dim + qk_rope_head_dim) * dtype_size |
| value_bytes_per_token = num_heads * v_head_dim * dtype_size |
| |
| For DeepSeek-V3 numbers (128 heads, 128 nope + 64 rope, 128 v_dim, bf16 = 2B): |
| - per token: 128*(128+64)*2 = 49 152 B key + 128*128*2 = 32 768 B value ≈ 80 KB |
| - at seq=4096, ttt=7, one draft layer, one sample: ~2.3 GB |
| - the draft typically has 1 layer, so this scales with `batch_size` not with `num_layers` |
| |
| So per-GPU peak from the TTT cache ≈ `batch_per_gpu * 2.3 GB`. Bound `batch_size` |
| accordingly, or revisit the compressed-cache option if this blocks longer-context training. |
|
|
| **TP/SP for MLA — validate before Phase 1 sign-off**. MLA has asymmetric per-head dims |
| (`qk_nope`, `qk_rope`, `v_head_dim`) — Yunchang USP and any TP slicing assume per-head |
| homogeneity in places. A small smoke test on Kimi-K2.5 with `tp_size=2`, |
| `sp_ulysses_size=2` is on the Phase 1 critical path, not a Phase 4 concern. |
|
|
| **Migration for the old `configs/deepseek-v3-671b-eagle3.json`**. That file declares |
| `model_type: "llama"` and trains a Llama-style draft against a DeepSeek target — it's not |
| an MLA draft despite the filename. Phase 1 plan: |
| - Move it to `specforge/config/draft_configs/deepseek_v3_671b_eagle3_llama_draft.json` |
| (preserves backward compat for anyone training against it). |
| - Add the new MLA config at `specforge/config/draft_configs/deepseek_v3_671b_eagle3.json`. |
| - Emit a one-time deprecation warning when the old path is loaded. |
|
|
| ### 4.8 SGLang export tool |
|
|
| ```python |
| # export/to_sglang.py |
| def export_to_sglang(checkpoint_dir: str, output_dir: str, target_model_path: str, |
| prune_vocab: bool = False, prune_dataset: str = None): |
| """Convert FSDP training checkpoint to SGLang-loadable spec-decoder format. |
| |
| Handles: |
| - FSDP state dict → flat state dict |
| - Weight key renaming (training names → SGLang spec-decoder expected names) |
| - MLA weight naming: q_a_proj, kv_a_proj_with_mqa, kv_b_proj etc. |
| - Optional vocab pruning based on dataset token frequency |
| - Config generation (draft config JSON + tokenizer copy) |
| """ |
| ... |
| ``` |
|
|
| **Weight-name compatibility is the riskiest single piece of the rewrite** — it's silent |
| when it goes wrong (loader picks up zeros for missing keys, or refuses to load). Before |
| Phase 3 starts, produce an explicit two-column map per draft arch: |
|
|
| | Trainer key | SGLang spec-decoder loader key | |
| |---|---| |
| | `model.layers.{i}.self_attn.q_a_proj.weight` | _e.g._ `draft_model.layers.{i}.self_attn.q_a_proj.weight` | |
| | `model.layers.{i}.self_attn.q_a_layernorm.weight` | ... | |
| | `model.layers.{i}.self_attn.q_b_proj.weight` | ... | |
| | `model.layers.{i}.self_attn.kv_a_proj_with_mqa.weight` | ... | |
| | `model.layers.{i}.self_attn.kv_a_layernorm.weight` | ... | |
| | `model.layers.{i}.self_attn.kv_b_proj.weight` | ... | |
| | `model.layers.{i}.self_attn.o_proj.weight` | ... | |
| | `model.embed_tokens.weight` | ... | |
| | `lm_head.weight` | ... | |
| | `t2d`, `d2t` (vocab mapping) | ... | |
|
|
| Filling the RHS column requires reading SGLang's current spec-decoding draft loader (the |
| `Eagle3*` or `LightseekSpec*` loader, whichever is current in sgl-project/sglang) and |
| should be a documented artifact (`docs/export_weight_map_mla.md`), not implicit in code. |
|
|
| ### 4.9 FSDP seam (FSDP1 now, FSDP2-ready) |
|
|
| The plan ships on FSDP1 (`SHARD_GRAD_OP`) but `apply_fsdp` is a stable seam from day one, |
| gated by `TrainingConfig.fsdp_version`: |
|
|
| ```python |
| # training/fsdp.py |
| def apply_fsdp(model, training_config, wrap_policy=None): |
| if training_config.fsdp_version == 1: |
| return _apply_fsdp1(model, training_config, wrap_policy) |
| elif training_config.fsdp_version == 2: |
| return _apply_fsdp2(model, training_config, wrap_policy) |
| raise ValueError(...) |
| ``` |
|
|
| Rationale: FSDP2 (PT 2.4+) composes much better with `torch.compile` and per-parameter |
| sharding (`fully_shard` on individual submodules). The compute-friendly default for the |
| next 12 months is FSDP2. Pinning to FSDP1 in the Trainer interface means rewriting again |
| once `compile_model: True` becomes the norm. Implementing `_apply_fsdp2` is deferred to |
| Phase 7 #31 — the seam is what matters now. |
|
|
| ### 4.10 Train-with-decode mode (Phase 5) |
|
|
| W4 from §1 ("Workloads in scope") is a real workload: the same long-lived SGLang server |
| simultaneously generates training data and serves real spec-decoding traffic, with the |
| trainer pushing freshly-trained draft weights into that server every N steps. The TorchSpec |
| implementation does this via Ray actors (`SglEngine.update_weights_from_disk` + |
| `controller/loop._maybe_sync_draft_weights`). SpecForge gets the same workload by extending |
| the existing `TargetEngine` / `Stream` / `Trainer` interfaces — **no Ray, no Mooncake**. |
|
|
| Three new primitives, all additive to Phase 4 abstractions: |
|
|
| #### (a) `TargetEngine.update_draft_weights` — weight push API |
|
|
| ```python |
| # models/targets/base.py |
| class TargetEngine(ABC): |
| def update_draft_weights( |
| self, |
| weights_path: str | os.PathLike, |
| *, |
| blocking: bool = True, |
| load_format: Literal["pt", "safetensors"] = "safetensors", |
| ) -> dict[str, Any]: |
| """Push new draft weights into a long-lived serving target. |
| |
| Default impl raises NotImplementedError. In-process engines (sglang/hf/custom) |
| do not need this — they're torn down with the trainer. Implemented by |
| `SGLangServerEngine`, which forwards to SGLang's |
| `/update_weights_from_disk` endpoint (already present in sglang ≥0.4). |
| """ |
| raise NotImplementedError(f"{type(self).__name__} does not support live draft weight updates") |
| ``` |
|
|
| `SGLangServerEngine` adds: |
|
|
| ```python |
| class SGLangServerEngine(TargetEngine): |
| def update_draft_weights(self, weights_path, *, blocking=True, load_format="safetensors"): |
| return self._post_json( |
| "/update_weights_from_disk", |
| {"model_path": str(weights_path), "load_format": load_format}, |
| ) |
| |
| def serving_metrics(self) -> dict[str, float]: |
| """Pull serving-side acceptance rate / queue depth from SGLang. |
| |
| Mapped to wandb keys like serving/spec_accept_rate, serving/spec_accept_length. |
| """ |
| return self._get_json("/spec_decoding_stats") |
| ``` |
|
|
| The shared filesystem assumption is the same one TorchSpec already lives with: trainer |
| saves to a path readable by the SGLang server (NFS, or rank-0 saves and broadcasts the |
| path; the server reads from that path). |
|
|
| #### (b) `ServingTrafficStream` — train on real serving prompts (optional) |
|
|
| For pure W4 you can keep using static jsonl with `OnlineStream`/`RemoteStream` — the |
| training data has nothing to do with the serving traffic, just the engine is shared. But |
| **true on-policy training** wants to sample prompts from real serving traffic. New stream: |
|
|
| ```python |
| # data/streams/serving_traffic.py |
| class ServingTrafficStream(RemoteStream): |
| """RemoteStream variant that pulls prompts from a serving-traffic buffer |
| instead of a static dataset. |
| |
| Buffer backends in scope: |
| - file: jsonl tail (simplest; rotated by serving) |
| - redis: list with LPUSH (serving) / BRPOP (this stream) |
| - kafka: topic with consumer group |
| |
| Out of scope: anything that requires Ray. |
| """ |
| def __init__( |
| self, |
| buffer_uri: str, # "file:///var/log/sgl/prompts.jsonl" |
| # "redis://host:6379/0/serving_prompts" |
| target_engine: SGLangServerEngine, |
| sample_rate: float = 1.0, # PII-safe rate limiting |
| cold_start_jsonl: Optional[str] = None, # bootstrap when buffer is empty |
| **kwargs, |
| ): ... |
| ``` |
|
|
| The privacy / PII / sample-rate decisions live with whoever runs the serving cluster, |
| **not in the trainer**. Stream just consumes whatever it's pointed at. |
|
|
| #### (c) `Trainer._maybe_sync_draft_weights` — periodic push hook |
|
|
| ```python |
| # training/trainer.py — addition to the loop in §4.2 |
| class Trainer: |
| def _maybe_sync_draft_weights(self, step: int) -> None: |
| """Save draft state and call target_engine.update_draft_weights. |
| |
| Mirrors TorchSpec's _maybe_sync_draft_weights (controller/loop.py:42-73) |
| but stays in-trainer. No Ray needed because there's exactly one engine |
| to update (one trainer ↔ one serving SGLang). |
| """ |
| cfg = self.config.weight_sync |
| if not cfg.enabled or self._target_engine is None: |
| return |
| if (step + 1) % cfg.interval != 0: |
| return |
| |
| with rank_0_priority(): |
| tmp = Path(cfg.tmp_dir or self.config.checkpoint_dir) / "draft_weight_sync" |
| tmp.mkdir(parents=True, exist_ok=True) |
| self._save_draft_for_serving(tmp) # FSDP full-state -> safetensors |
| |
| # Only rank 0 actually issues the HTTP call; barrier so other ranks wait. |
| if dist.get_rank() == 0: |
| self._target_engine.update_draft_weights(tmp, load_format="safetensors") |
| dist.barrier() |
| ``` |
|
|
| Wired into the existing `_maybe_log / _maybe_eval / _maybe_save` chain. ~30 LOC. |
|
|
| #### (d) Decode-mode `SGLangServerEngine` — server-side config |
|
|
| The serving SGLang must be launched with both prefill+aux capture (for training data |
| generation) **and** spec decoding (for serving traffic) enabled simultaneously. This is |
| already supported by sglang ≥0.4 via the `--enable-aux-hidden-states` + |
| `--enable-spec-training-mooncake=False` (we don't need mooncake) + |
| `--enable-cuda-graph=True` combination. `SGLangServerEngine` gets a `decode_mode: bool` |
| config flag that simply governs which set of server-args gets validated; it doesn't |
| launch the server itself (operations responsibility). |
|
|
| #### Acceptance-rate as a first-class metric |
|
|
| In W1-W3 the relevant metric is `simulated_acc_len` (computed offline from logged |
| per-position accuracy). In W4, the **real serving acceptance rate** is available from |
| the same SGLang server. The `Trainer.train()` loop calls `target_engine.serving_metrics()` |
| on the same cadence as eval and merges into the tracker dict: |
|
|
| ```python |
| if self.config.weight_sync.report_serving_metrics and step % self.config.eval.interval == 0: |
| sm = self._target_engine.serving_metrics() |
| self.tracker.log({ |
| "serving/spec_accept_rate": sm["accept_rate"], |
| "serving/spec_accept_length": sm["accept_length"], |
| }, step=step) |
| ``` |
|
|
| **Why this is sufficient (no Ray, no Mooncake)**: |
|
|
| - One trainer ↔ one server is a hard scope boundary (§8 non-goal). The dispatcher / |
| multi-engine load balancing that TorchSpec's `AsyncInferenceManager` provides isn't |
| needed at this scope. |
| - Weight push is HTTP `/update_weights_from_disk`; SGLang already implements this for |
| serving updates, no new server work. |
| - Training data path stays identical to W2/W3 — `RemoteStream` / |
| `ServingTrafficStream` produce `TrainBatch` exactly the same way. |
| - The trainer has no idea decode mode is on. It just pushes weights and reads metrics. |
|
|
| If profiling later shows that the HTTP weight-push latency itself is a problem (saving + |
| loading 100MB-1GB of weights every 500 steps), the optimization is to upgrade |
| `update_draft_weights` to accept an in-memory state_dict + zero-copy transfer, again |
| without changing the trainer. Same shape as the L2 transport upgrade discussed in §6. |
| |
| --- |
| |
| ## 5. Feature List (Prioritized) |
| |
| **Ordering note (from review):** the original draft put the unified `Trainer` in Phase 2 |
| and the `TargetEngine` / `HiddenStateStream` abstractions in Phase 4. That meant Phase 2 |
| would build the trainer against the *old* `Eagle3TargetModel`/`DFlashTargetModel` |
| interface and then Phase 4 would re-plumb it. Reordered so the interfaces (and a single |
| in-process implementation of each) land in Phase 2 alongside the trainer; Phase 4 then |
| adds new implementations (`SGLangServerEngine`, `RemoteStream`) of stable interfaces |
| rather than re-plumbing. |
| |
| ### Phase 1: MLA Draft + Registry (Week 1-2) |
| |
| | # | Feature | Why | Effort | |
| |---|---------|-----|--------| |
| | 1 | Draft plugin registry (`@register_draft`) | Every new arch today touches 4 files. This is the single biggest extensibility blocker. | S | |
| | 2 | MLA Eagle3 draft (`Eagle3DeepseekV2ForCausalLM`) | The gap for DeepSeek-V3 / Kimi-K2. Port from TorchSpec. Includes MLA-aware KV cache handling via HF `DynamicCache` (no ABC break). | M | |
| | 3 | MLA draft configs | `qwen3_8b_eagle3_mla.json`, `kimi_k25_eagle3_mla.json`, `deepseek_v3_671b_eagle3.json` (real MLA, not Llama pretending). | S | |
| | 3a | **MLA + TP/SP smoke test** | Validate Yunchang USP and TP slicing on Kimi-K2.5 with asymmetric MLA head dims before declaring Phase 1 done. | S | |
|
|
| **Deliverable**: Train an MLA Eagle3 draft for Qwen3-8B / Kimi-K2.5 using existing `train_eagle3.py`. |
| Existing scripts unchanged — registry lives alongside the old `_model_mapping` dicts. |
|
|
| ### Phase 2: Interfaces + Trainer + Eval + Checkpoints (Week 3-6) |
|
|
| | # | Feature | Why | Effort | |
| |---|---------|-----|--------| |
| | 4 | `TargetEngine` protocol + in-process impls | Define the interface up front so the trainer is built against it from day 1. Adapt existing `Eagle3TargetModel` / `DFlashTargetModel` as the initial concrete impls. | M | |
| | 5 | `HiddenStateStream` protocol + `OnlineStream` + `OfflineStream` | Same reason — stable abstraction before trainer code is written. `RemoteStream` deferred to Phase 4. | M | |
| | 6 | `Trainer` class + `DraftTrainStrategy` protocol | Collapse `train_eagle3.py` and `train_dflash.py`. Stop the script-per-arch sprawl. Trainer consumes `TrainBatch` only — no direct target-model calls. | M | |
| | 7 | `Eagle3TTTStrategy` + `DFlashBlockStrategy` | Extract from existing scripts into strategy implementations. Each declares its own `fsdp_wrap_policy()`. | M | |
| | 8 | `CheckpointManager` | Rotation (`max_checkpoints`), best tracking, `latest_checkpointed_iteration.txt`, `meta.json`. Must integrate with `stream.seek()` on resume. | S | |
| | 9 | `Evaluator` + `EvalCache` | Online eval during training: `simulated_acc_len`, `avg_loss`, `avg_acc`. Per-position acc aggregated across batches before geometric sum. Cache key covers eval/target/tokenizer/template/aux/seqlen. | M | |
| | 10 | Gradient accumulation correctness | `model.no_sync()` on non-sync micro-steps, explicit `zero_grad`, `accumulation_steps` in config. Verify by step-time profiling. | S | |
| | 11 | `target_layer_ids` as first-class contract | EAGLE3 hardcodes 3 aux layers; DFlash uses a list. Generalize so every draft declares what it needs and the stream materializes accordingly. | S | |
| | 12 | FSDP seam (FSDP1 default, FSDP2-ready) | `apply_fsdp` dispatches on `fsdp_version`. FSDP2 impl deferred — but the seam is stable now so the trainer doesn't get rewritten when FSDP2 lands. | S | |
|
|
| **Deliverable**: `Trainer` (no CLI yet — legacy scripts call into it) works for both Eagle3 |
| and DFlash, using the new interfaces with single in-process target/stream impls. Eval |
| metrics logged during training. Best checkpoint auto-saved. Numerical equivalence (§10) |
| checked against the pre-refactor scripts. |
|
|
| ### Phase 3: Config + Export (Week 7-8) |
|
|
| | # | Feature | Why | Effort | |
| |---|---------|-----|--------| |
| | 13 | Pydantic config schema | Replace argparse + JSON split with one validated YAML per run. | S-M | |
| | 14 | CLI (`specforge train\|prepare\|export\|eval`) | Single entry point. Legacy scripts now become shims that build a `Config` and call `Trainer`. | S | |
| | 15 | SGLang export tool | Checkpoint → SGLang spec-decoder layout. **Weight-name map (§4.8) is a documented artifact, finalized before this phase starts.** | M | |
| | 16 | HF export tool + vocab pruning | FSDP checkpoint → HF format. Optional vocab pruning by dataset frequency. | S-M | |
|
|
| **Deliverable**: Complete train → eval → export → serve pipeline in one tool. Old MLA |
| draft trained in Phase 1 successfully loads in a SGLang server. |
|
|
| ### Phase 4: Remote Target + VLM + Parsers (Week 9-11) |
|
|
| | # | Feature | Why | Effort | |
| |---|---------|-----|--------| |
| | 17 | `SGLangServerEngine` | Target as HTTP service. Unlocks online training against 671B targets on separate GPUs. New impl of the Phase-2 `TargetEngine` interface. | M-H | |
| | 18 | `RemoteStream` | Stream backend that talks to `SGLangServerEngine`. `prefetch_factor` + `max_in_flight` back-pressure. | M | |
| | 19 | VLM unification | Delete `QwenVLOnlineEagle3Model`. VLM handled via typed `MediaInputs` in `TargetEngine` + data pipeline. | M | |
| | 20 | Additional parsers | `KimiK25Parser`, `MiniMaxParser` from TorchSpec. | S | |
|
|
| **Deliverable**: Train Eagle3 for DeepSeek-V3 with the target running on a separate SGLang |
| server cluster. |
|
|
| ### Phase 5: Train-with-decode (Week 12-14) |
|
|
| W4 from §1 — promoted from Phase 6 because it's a real workload, not a future option. All |
| features build on the Phase 4 `SGLangServerEngine` + `RemoteStream` foundation; **no |
| trainer / strategy / draft changes**. |
|
|
| | # | Feature | Why | Effort | |
| |---|---------|-----|--------| |
| | 21 | `TargetEngine.update_draft_weights` | Weight push API on the ABC; default raises NotImplementedError. `SGLangServerEngine` impl forwards to SGLang's `/update_weights_from_disk`. | S | |
| | 22 | Decode-mode `SGLangServerEngine` | `decode_mode: bool` config + `serving_metrics()` endpoint. Validates that the server was launched with prefill+aux **and** spec decoding both enabled. | S-M | |
| | 23 | `Trainer._maybe_sync_draft_weights` | Periodic push hook in the main loop. ~30 LOC. Save draft state on rank 0 → engine.update_draft_weights → barrier. | S | |
| | 24 | `ServingTrafficStream` | Optional stream that pulls prompts from a serving-traffic buffer (file / Redis / Kafka). Subclass of `RemoteStream`; cold-start fallback to a static jsonl. | M | |
| | 24a | Acceptance-rate metrics | Trainer reads `serving/spec_accept_rate` + `serving/spec_accept_length` from the engine and logs alongside `eval/simulated_acc_len`. Closes the gap between offline proxy and production reality. | S | |
|
|
| **Deliverable**: A long-lived SGLang server simultaneously serves real spec-decoding |
| traffic and produces training data; the trainer pushes draft weights every N steps and |
| production acceptance rate trends up over the run, monotonically across `weight_sync_interval`. |
| End-to-end smoke test: 1k-step run on a small target, acceptance-rate slope > 0. |
|
|
| ### Phase 6: Polish (Week 15-16) |
|
|
| | # | Feature | Why | Effort | |
| |---|---------|-----|--------| |
| | 25 | Backbone-agnostic DFlash | Factor `DFlashDraftModel` so MLP/Norm/RoPE are parameterized. Currently hardcoded to Qwen3. | S-M | |
| | 26 | `torch.compile` support | Optional `compile_model` flag for draft model compilation. | S | |
| | 27 | `defer_tokenization` + dynamic loss mask | Tokenize at fetch time, not preprocessing time. Enables dynamic loss masking. | M | |
| | 28 | Tokenization disk cache | Cache tokenized datasets keyed by data path + template + max_length. | S | |
| | 29 | WSD learning rate scheduler | Warmup-Stable-Decay. Useful for long runs / continued pretraining; defer unless a concrete need surfaces. | S | |
| | 30 | Remove legacy `scripts/legacy/` shims | After two minor releases with `DeprecationWarning`. Concrete target: SpecForge 0.X. | S | |
| |
| ### Phase 7: Optional / Future |
| |
| | # | Feature | Why | Effort | |
| |---|---------|-----|--------| |
| | 31 | FSDP2 implementation | Fill in `_apply_fsdp2` behind the Phase-2 seam. Required if `torch.compile` becomes default. | M | |
| | 32 | Mooncake streaming backend | GPU↔GPU tensor transport for multi-node disaggregation. Only if profiling shows HTTP/gRPC is the bottleneck (T1 trigger, §6). | H | |
| | 33 | In-memory `update_draft_weights` | Bypass save-to-disk in W4 weight sync; transfer state_dict directly over a binary channel. Only if 100MB-1GB / N-step push becomes a measurable slowdown. | M | |
| | 34 | Multi-job inference pool sharing | Single SGLang cluster amortized across ≥5 concurrent training jobs (T2 trigger, §6). At this point the actor topology becomes worth the complexity. | H | |
| | 35 | FA4 attention backend | FlashAttention 4 support for draft attention. | S-M | |
|
|
| --- |
|
|
| ## 6. Tradeoffs Worth Flagging |
|
|
| ### Ray + Mooncake vs HTTP/gRPC |
|
|
| TorchSpec gets a lot from Ray + Mooncake, but they're large dependencies with significant |
| operational complexity. For SpecForge, start with HTTP/gRPC to a SGLang server for cross-node |
| training — this covers most use cases (target on dedicated GPUs, trainer on others) at a |
| fraction of the surface area. |
|
|
| **Train-with-decode (W4) is in scope without Ray.** The TorchSpec value-add from actor |
| topology is concentrated in *multi-job inference pool sharing* and *zero-overhead async |
| producer-consumer pipelines*. For one-trainer-↔-one-server (which is W4's scope per §8), |
| HTTP `/update_weights_from_disk` plus a `ServingTrafficStream` over a Redis/Kafka/file |
| buffer is sufficient. See §4.10. Ray actor topology only starts paying off at the T2 |
| trigger below. |
|
|
| **Triggers that would force a re-evaluation:** |
|
|
| - **T1 — transport bottleneck**: profile shows `last_hidden_states` over JSON/HTTP |
| consumes >30% of step time, with target ≥70B. Response: upgrade `RemoteStream` |
| transport to a binary protocol (gRPC + protobuf, or Mooncake-as-transport). |
| **L1/L3 unchanged.** This is feature #32 in Phase 7. |
|
|
| - **T2 — multi-job inference pool sharing**: ≥5 concurrent training jobs hammering the |
| same 671B target, each tearing the engine up/down is wasteful. Response: introduce |
| Ray actor pool with `AsyncInferenceManager`-style dispatcher. This is the only |
| trigger that genuinely justifies adopting Ray; until then, complexity is pure cost. |
| Feature #34 in Phase 7. |
|
|
| - **T3 — Mooncake transport**: only meaningful after T1; even then, evaluate gRPC |
| binary first. Mooncake's RDMA edge matters for >100GB/s aggregate, which neither W3 |
| nor W4 typically hit. Feature #32 in Phase 7. |
|
|
| The bet: 90% of the time we're in W1-W3 territory and the refactor's interface boundaries |
| do real work. W4 lands in Phase 5 with ~200 LOC of additive code on those boundaries. |
| TorchSpec's full topology is reserved for the day T2 actually fires. |
|
|
| ### Plugin registry vs HF AutoModel pattern |
|
|
| HF's "drop in a config JSON" is a nice affordance. Keep both: decorator registry is primary |
| for dispatch, HF type dispatch as fallback for configs that specify a known `transformers` |
| config class. `AutoDraftModelConfig.from_file()` checks `DRAFT_REGISTRY` by architecture name |
| first, falls back to config type mapping. |
|
|
| ### Strategy pattern is an indirection |
|
|
| If only EAGLE3 and DFlash ever exist, two scripts are fine. The bet is that Medusa, MTP, or |
| hybrid variants will want the same trainer infrastructure — if so, strategies pay off fast. |
| If not, the cost is one extra level of dispatch that doesn't hurt readability. |
|
|
| ### Pydantic vs OmegaConf |
|
|
| Pydantic gives better validation and serialization. OmegaConf gives effortless YAML merge + |
| CLI override (`--training.lr=1e-4`). Options: |
| 1. **Pydantic + custom CLI overlay** (parse `--key=value` args, `model_validate` from merged dict) — our choice. |
| 2. OmegaConf + dataclass (TorchSpec's approach). |
| 3. Pydantic + typer/click for CLI. |
|
|
| We go with (1) because SpecForge already uses Pydantic (`ChatTemplate`) and validation |
| matters more than merge ergonomics at this stage. The CLI overlay is ~30 lines of code. |
|
|
| ### MLA cache: compressed vs expanded |
|
|
| Two options for KV cache during Eagle3 TTT unroll: |
| - **Compressed**: Cache `kv_compressed` (low-rank) + `k_rope_raw`. Saves memory but requires |
| MLA-specific cache logic in `core/eagle3.py`. |
| - **Expanded**: Cache full `(K, V)` after projection. More memory but `core/eagle3.py` is |
| architecture-agnostic. |
|
|
| We choose **expanded** (matching TorchSpec). The TTT unroll is typically 5-7 steps with |
| short sequences — memory savings from compressed cache are marginal, and keeping |
| `core/eagle3.py` architecture-agnostic is worth more. |
|
|
| --- |
|
|
| ## 7. Migration Path |
|
|
| Each phase ships independently and can be validated without breaking existing users. |
|
|
| ### Phase 1 plan |
|
|
| 1. Add `models/drafts/__init__.py` with `DRAFT_REGISTRY` and `@register_draft`. |
| 2. Move `LlamaForCausalLMEagle3` to `models/drafts/llama_eagle3.py`, decorate with |
| `@register_draft("LlamaForCausalLMEagle3")`. |
| 3. Port `deepseek_eagle.py` from TorchSpec → `models/drafts/deepseek_eagle3.py`, decorate |
| with `@register_draft("Eagle3DeepseekV2ForCausalLM")`. |
| 4. Update `AutoDraftModelConfig.from_file()` to check `DRAFT_REGISTRY` first. |
| 5. Update `AutoEagle3DraftModel.from_config()` to check `DRAFT_REGISTRY` first. |
| 6. Add MLA draft configs (`qwen3_8b_eagle3_mla.json`, `kimi_k25_eagle3_mla.json`). |
| 7. Keep old `_model_mapping` / `_config_mapping` as fallback — remove later. |
| 8. Test: train MLA Eagle3 draft using existing `train_eagle3.py` — zero changes to script. |
|
|
| ### Phase 2 plan |
|
|
| 1. Define `TargetEngine` protocol in `models/targets/base.py`. |
| 2. Adapt existing `SGLangEagle3TargetModel` → `SGLangTargetEngine` (in-process). |
| 3. Adapt existing `HFEagle3TargetModel` → `HFTargetEngine`. |
| 4. Adapt existing `CustomEagle3TargetModel` → `CustomTargetEngine`. |
| 5. Define `HiddenStateStream` protocol + `OnlineStream` + `OfflineStream`. `seek()` method |
| required for resume correctness. `OnlineStream` exposes `prefetch_factor`. |
| 6. Create `training/trainer.py` with the `Trainer` class. Trainer consumes `TrainBatch` |
| only, never a target model directly. |
| 7. Extract `Eagle3TTTStrategy` from `train_eagle3.py` logic. Implement `fsdp_wrap_policy()`. |
| 8. Extract `DFlashBlockStrategy` from `train_dflash.py` logic. Implement `fsdp_wrap_policy()`. |
| 9. Add the `apply_fsdp` seam dispatching on `fsdp_version` (FSDP1 only for now). |
| 10. Implement `CheckpointManager` — `save` / `load` / `update_best` / `_rotate`. |
| 11. Implement `Evaluator` (per-position acc aggregation across batches) + `EvalCache` |
| (full-key MD5). |
| 12. Move old scripts to `scripts/legacy/`, create shims that instantiate `Trainer`. |
| 13. **Numerical equivalence test (§10)**: legacy shim must match the pre-refactor script's |
| loss curve to within tolerance on a fixed seed at steps 100/500/1000. Gate. |
| |
| ### Phase 3 plan |
|
|
| 1. Define `config/schema.py` with Pydantic models (Literal-typed enums). |
| 2. Implement `config/loader.py` (YAML load + CLI override). |
| 3. Implement `cli.py` (`specforge train|prepare|export|eval`). |
| 4. Finalize the SGLang weight-name map (`docs/export_weight_map_mla.md`) by reading the |
| current SGLang spec-decoding loader. Block #5 on this. |
| 5. Implement `export/to_sglang.py` and `export/to_hf.py`. |
| 6. Mechanically migrate existing JSON configs to YAML with the new schema. |
| 7. Test: `specforge train --config x.yaml` end-to-end. MLA draft from Phase 1 successfully |
| loads in a SGLang server with non-trivial acceptance rate. |
|
|
| ### Phase 4 plan |
|
|
| 1. Implement `SGLangServerEngine` (HTTP client to SGLang server). Reuses the Phase-2 |
| `TargetEngine` interface — no trainer changes. |
| 2. Implement `RemoteStream` (`prefetch_factor`, `max_in_flight`). Reuses the Phase-2 |
| `HiddenStateStream` interface — no trainer changes. |
| 3. Add typed `MediaInputs` to `TargetEngine.generate_train_data`. |
| 4. Delete `QwenVLOnlineEagle3Model`; wire VLM through `MediaInputs` + data pipeline. |
| 5. Port `KimiK25Parser`, `MiniMaxParser` from TorchSpec. |
| 6. Test: online training with SGLang server on separate node; results within tolerance of |
| in-process training (same target weights). |
|
|
| ### Phase 5 plan |
|
|
| 1. Add `update_draft_weights(weights_path, *, blocking, load_format)` to `TargetEngine` |
| ABC with `NotImplementedError` default. |
| 2. Implement `SGLangServerEngine.update_draft_weights` → POST `/update_weights_from_disk`. |
| Implement `serving_metrics()` → GET `/spec_decoding_stats`. |
| 3. Add `decode_mode: bool` to `SGLangServerEngine`. Validate that the operator's server |
| was launched with prefill+aux **and** spec decoding both enabled (warn loudly otherwise). |
| 4. Add `Trainer._maybe_sync_draft_weights(step)` hook between `_maybe_eval` and |
| `_maybe_save`. Save FSDP full-state on rank 0 → call engine update on rank 0 → barrier. |
| 5. Implement `ServingTrafficStream(RemoteStream)` with file/Redis/Kafka buffer backends |
| and a cold-start jsonl fallback. PII / sample-rate are operator concerns, not |
| trainer concerns. |
| 6. Wire `serving/spec_accept_rate` and `serving/spec_accept_length` through the existing |
| `Tracker` (every `eval_interval` steps). |
| 7. **End-to-end test**: 1k-step run with a small target (Qwen3-8B) co-located with a |
| spec-decoding load generator. Acceptance-rate slope across `weight_sync_interval` |
| buckets must be > 0; final acceptance rate must exceed cold-start by ≥ X% |
| (X TBD — set baseline from Phase 4 numbers). |
| 8. **Weight-sync correctness test (§10.5)**: before/after a sync, the served draft's |
| token-level outputs match the trainer's draft on a 32-prompt fixed eval set within |
| acceptance-rate tolerance. |
|
|
| --- |
|
|
| ## 8. Non-Goals (Explicit) |
|
|
| - **No Ray dependency.** SpecForge stays torchrun-native, including for train-with-decode (W4). |
| - **No Mooncake in Phase 1-6.** HTTP/gRPC to SGLang is sufficient for both disaggregated |
| training (W3) and train-with-decode (W4). Mooncake is gated behind T1/T3 in §6. |
| - **No vLLM target backend.** SGLang is primary; HF is the reference. Adding vLLM is possible |
| via the `TargetEngine` interface but not prioritized. |
| - **No multi-engine load balancing / multi-job inference pool sharing.** Exactly **one** |
| target engine instance per training job — including in train-with-decode (one trainer |
| ↔ one server). Scaling within a single engine is via SGLang's own scaling (multiple |
| TP workers, replicas behind a single endpoint). Multi-job sharing is gated behind T2 |
| in §6 and lands as Phase 7 #34. |
| - **No on-policy serving-prompt mining without a buffer.** `ServingTrafficStream` reads |
| from an external buffer (Redis/Kafka/file) populated by the serving cluster; SpecForge |
| does **not** intercept inflight serving requests directly. |
|
|
| --- |
|
|
| ## 9. Success Criteria |
|
|
| | Phase | Criterion | |
| |-------|-----------| |
| | 1 | MLA Eagle3 draft trains on Qwen3-8B and converges to comparable loss as Llama draft. Kimi-K2.5 TP/SP smoke test passes. | |
| | 2 | Legacy shims (using new `Trainer` internally) match pre-refactor scripts to numerical-equivalence tolerance (§10). Eval metrics agree with manual eval. | |
| | 3 | `specforge train --config x.yaml` end-to-end. MLA draft from Phase 1 loads in SGLang and produces non-trivial acceptance rate. | |
| | 4 | Online training with 671B target on separate SGLang server; per-step loss within tolerance of in-process baseline using identical target weights. | |
| | 5 | **Train-with-decode**: 1k-step run; one long-lived SGLang server simultaneously serves spec-decoding traffic and produces training data; trainer pushes draft weights every `weight_sync_interval`; serving acceptance rate trends up monotonically across sync buckets. Weight-sync correctness gate §10.5 passes. | |
| | 6 | DFlash on Llama backbone; tokenization cache cuts data prep time by >50%; legacy shims removed. | |
|
|
| --- |
|
|
| ## 10. Testing Strategy (gating, not optional) |
|
|
| Every phase that changes the training path must pass a numerical-equivalence gate against |
| the immediately prior commit. Treat these as CI requirements, not nice-to-haves — |
| behavior-preserving refactors are the most common place silent regressions hide. |
|
|
| ### 10.1 Numerical-equivalence gate (Phase 2, 4 critical) |
|
|
| For a fixed seed and a fixed micro-batch sequence (3 batches × 4 samples is enough), the |
| new code path must match the old: |
|
|
| | Metric | Tolerance | Steps to check | |
| |---|---|---| |
| | Per-step training loss | `atol=1e-4, rtol=1e-4` (BF16 master copy) | 0, 1, 100, 500, 1000 | |
| | Per-position eval acc (vector of length `ttt_length`) | `atol=1e-3` | end of eval at step 1000 | |
| | `simulated_acc_len` | `atol=1e-3` | end of eval at step 1000 | |
| | Model state dict (selected keys) | `atol=1e-4` | after step 100 | |
|
|
| The gate run uses: Llama Eagle3 draft, Qwen3-8B target, offline mode, sharegpt eval slice. |
| Cheap enough to run on every PR that touches `core/`, `training/`, or `models/drafts/`. |
|
|
| ### 10.2 Smoke tests (every phase) |
|
|
| - One config per draft arch (`llama_eagle3`, `deepseek_v3_eagle3`, `qwen3_dflash`) trains |
| for 20 steps without crashing under TP=1, TP=2, and TP=2+SP=2. |
| - Checkpoint save + resume produces the same loss curve as no-resume (validates |
| `stream.seek()`). |
| - Eval cache miss + hit produce identical metrics (validates the cache key set). |
|
|
| ### 10.3 Distributed correctness (Phase 1 + 2) |
|
|
| - MLA + Yunchang USP with `qk_nope_head_dim != qk_rope_head_dim != v_head_dim` on |
| Kimi-K2.5. This is the Phase 1 risk item and must be on the critical path, not the |
| Phase 4 wishlist. |
| - Gradient accumulation: confirm one `all_reduce` per `optimizer.step()`, not per |
| `backward()`. Use a NCCL communicator log or `torch.profiler` once. |
|
|
| ### 10.4 Export-loop test (Phase 3) |
|
|
| - Train MLA draft for 100 steps, export to SGLang, load in SGLang server, run 32 generation |
| requests. Acceptance rate > 0 (i.e. the loader actually consumed the weights, not zeros). |
| This catches weight-name-map regressions at the integration boundary. |
|
|
| ### 10.5 Weight-sync correctness gate (Phase 5) |
|
|
| Every PR that touches `Trainer._maybe_sync_draft_weights`, `SGLangServerEngine.update_draft_weights`, |
| or the FSDP-state-save path must pass a parity gate: |
|
|
| | Step | Action | Pass condition | |
| |---|---|---| |
| | 1 | Train 50 steps; capture trainer-side draft state (full FSDP state dict, broadcast to rank 0). | — | |
| | 2 | Call `Trainer._maybe_sync_draft_weights(50)`. | Engine returns 200; serving SGLang reports `update_weights/success=True`. | |
| | 3 | Run 32 fixed-seed prompts through both: (a) trainer's local draft (offline-style forward), (b) the now-updated serving SGLang. | Logits agree to `atol=1e-2, rtol=1e-2` (allowing for kernel diffs between training and serving stacks). | |
| | 4 | Compare serving acceptance rate over 256 prompts before vs after sync. | Acceptance rate **non-decreasing**; if 50 training steps moved the loss meaningfully, acceptance rate strictly increases. | |
|
|
| The gate fails if any of: (i) the HTTP push silently 200-no-ops, (ii) the serving stack |
| loaded a corrupt or partial state dict, (iii) FSDP state-dict gather missed a parameter. |
| This is the train-with-decode analog of the §10.4 export-loop test. |
|
|
| ### 10.6 Long-run weight-sync soak (Phase 5, optional) |
|
|
| - 1k-step run with `weight_sync_interval=100`. Bucket serving acceptance rate by |
| 100-step windows; the slope across buckets must be monotonically non-decreasing |
| (allowing one regression bucket per 10 — to absorb rare unlucky prompt batches). |
| - Watches for: (a) silent weight push failures (acceptance rate would plateau), |
| (b) draft drift due to incorrect FSDP gather, (c) serving-side OOM that quietly |
| rolls back to old weights. |
|
|
