From 34fc8b8649d2aafbd2d42d5faace0418f4b2ee3e Mon Sep 17 00:00:00 2001 From: ali-88123 <1940747290@qq.com> Date: Thu, 30 Jul 2026 22:32:45 +0800 Subject: [PATCH 1/2] refactor(drafter): promote DFly to a standalone DFlash-family drafter + add e2e multi-step TV loss --- LICENSE | 6 +- PR_DESCRIPTION.md | 130 ++++++++ ...spark_treeflash_qwen3_4b_draft_config.json | 27 -- angelspec/config/dfly_hy3_draft_config.json | 8 +- .../config/dfly_qwen3_4b_draft_config.json | 8 +- .../config/dfly_qwen3_8b_draft_config.json | 8 +- angelspec/config/train_config.py | 8 +- angelspec/config/utils.py | 3 +- angelspec/controller/eval.py | 5 +- angelspec/controller/inference_manager.py | 6 +- angelspec/controller/training_controller.py | 6 +- angelspec/data/dataset.py | 5 +- angelspec/data/parse.py | 16 +- angelspec/data/preprocessing.py | 3 +- angelspec/data/template.py | 6 +- .../mooncake_hidden_states_connector.py | 12 +- .../inference/engine/score_worker_ext.py | 4 +- angelspec/inference/engine/sgl_engine.py | 6 +- angelspec/models/__init__.py | 2 + angelspec/models/dflash.py | 175 +++++----- angelspec/models/dfly.py | 40 +++ angelspec/models/draft/__init__.py | 2 - angelspec/models/draft/auto.py | 21 +- angelspec/models/draft/base.py | 6 +- angelspec/models/draft/dflare.py | 4 +- angelspec/models/draft/dflash.py | 4 +- angelspec/models/draft/dfly.py | 146 +++++++-- angelspec/models/draft/dspark.py | 304 +----------------- angelspec/models/draft/llama3_eagle.py | 26 +- angelspec/models/draft/mtp.py | 14 +- .../models/draft/treeflash_dspark_dflare.py | 55 ---- angelspec/models/dspark.py | 69 +--- angelspec/models/eagle3.py | 8 +- angelspec/models/ops/flex_attention.py | 6 +- angelspec/train_entry.py | 43 +-- angelspec/training/data_fetcher.py | 6 +- angelspec/training/dflash_trainer.py | 102 +++--- angelspec/training/dfly_trainer.py | 41 +++ angelspec/training/dspark_trainer.py | 46 +-- angelspec/training/eagle3_trainer.py | 6 +- angelspec/training/mtp_trainer.py | 17 +- angelspec/training/optimizer.py | 5 +- angelspec/training/trainer.py | 8 +- angelspec/training/trainer_actor.py | 11 +- angelspec/transfer/mooncake/buffers.py | 2 +- angelspec/transfer/mooncake/store.py | 3 +- angelspec/utils/profiling.py | 6 +- angelspec/utils/usp.py | 4 +- configs/vllm_hy3_dfly.yaml | 16 +- configs/vllm_qwen3_8b_dfly.yaml | 18 +- docs/concepts/dflash.md | 3 +- docs/concepts/dfly.md | 25 +- docs/concepts/draft_model_family.md | 3 +- docs/concepts/dspark.md | 9 +- tests/test_dflash.py | 11 +- tests/test_dfly.py | 44 +-- tests/test_mtp.py | 18 +- tests/test_treeflash.py | 213 ------------ tools/eval_accept_rate.py | 28 +- tools/generate_data.py | 16 +- 60 files changed, 773 insertions(+), 1080 deletions(-) create mode 100644 PR_DESCRIPTION.md delete mode 100644 angelspec/config/dflare_dspark_treeflash_qwen3_4b_draft_config.json create mode 100644 angelspec/models/dfly.py delete mode 100644 angelspec/models/draft/treeflash_dspark_dflare.py create mode 100644 angelspec/training/dfly_trainer.py delete mode 100644 tests/test_treeflash.py diff --git a/LICENSE b/LICENSE index ab8621c..0f1c12e 100644 --- a/LICENSE +++ b/LICENSE @@ -1,8 +1,8 @@ -Tencent is pleased to support the open source community by making AngelSpec available. +Tencent is pleased to support the open source community by making AngelSpec available. -Copyright (C) 2026 Tencent. All rights reserved. +Copyright (C) 2026 Tencent. All rights reserved. -The open-source software included in this distribution may have been modified by Tencent (“Tencent Modifications”). All Tencent Modifications are Copyright (C) Tencent. +The open-source software included in this distribution may have been modified by Tencent (“Tencent Modifications”). All Tencent Modifications are Copyright (C) Tencent. AngelSpec is licensed under the Apache-2.0 except for the third-party components listed below. diff --git a/PR_DESCRIPTION.md b/PR_DESCRIPTION.md new file mode 100644 index 0000000..d9162d5 --- /dev/null +++ b/PR_DESCRIPTION.md @@ -0,0 +1,130 @@ +# [Refactor] Promote DFly to a first-class DFlash-family drafter and add an end-to-end multi-step TV loss + +> **Title:** `refactor(drafter): promote DFly to a standalone DFlash-family drafter + add e2e multi-step TV loss` + + +## Summary + +This PR promotes **DFly** from an architecture *variant* +that piggy-backed on the DSpark code path into a standalone, first-class member +of the DFlash drafter family. It gets its own config, model, training wrapper, +and trainer, and no longer depends on DSpark in any way. + +Alongside the refactor, this PR: + +- Removes the dead `treeflash_dspark_dflare` drafter and its test. +- Slims down DSpark by moving the shared TreeFlash hidden-states correction (and + related knobs) out of `dspark.py` and into `dfly.py`, where it now belongs. +- Adds an optional **end-to-end multi-step TV loss** (`γ`-step MTP) to the + DFlash composable loss, replacing the now-unused KL/LK temperature knobs. + +The net effect is a cleaner architecture graph (DFly no longer "rides" DSpark), +less coupling between drafters, and a large reduction in DSpark's surface area +(~+367 / −933 lines overall). + +## Motivation + +Previously, DFly was selected via `DSparkConfig` + `model_arch == "dfly"` and was +dispatched through `DSparkTrainer` / `DSparkModel`. This meant: + +- DFly's behavior was implicit and hard to discover (hidden behind a string flag). +- DSpark carried a lot of machinery (hidden-states correction, position-adaptive + alpha, etc.) that only DFly actually used. +- The `auto` dispatch logic had brittle special-case branches keyed on + `model_arch` string comparison. + +Making DFly a proper config/model/trainer triple removes the cross-routing, +makes the dispatch type-based, and lets each drafter own only what it needs. + +## Changes + +### New — DFly as a first-class drafter + +- **`angelspec/models/dfly.py`** (new): `DFlyModel` training wrapper. Subclasses + `DFlashModel` and overrides `_compute_draft_logits` to apply the optional + TreeFlash hidden-states correction (formula (1)) before the LM head. +- **`angelspec/training/dfly_trainer.py`** (new): `DFlyTrainer`, a thin subclass + of `DFlashTrainer` that only overrides the two model-build seams + (`_build_draft_model` / `_build_training_wrapper`). Reads the `dflash_*` + hyperparameter namespace. +- **`angelspec/models/draft/dfly.py`**: now defines its own `DFlyConfig` + (extends `DFlashConfig`, `model_type = "qwen3"`) and owns the + `HiddenStatesCorrection` module / `build_hidden_correction` helper (moved here + from `dspark.py`). `DFlyDraftModel.config_class` is now `DFlyConfig`. + +### Dispatch / registration + +- **`angelspec/models/draft/auto.py`**: register `DFlyConfig → DFlyDraftModel` + and architecture `"Qwen3DFlyModel" → DFlyConfig`. Removed the + `model_arch == "dfly"` and `model_arch == "dflare"` (TreeFlash) special-case + branches. +- **`angelspec/training/trainer_actor.py`**: add a `DFlyConfig` dispatch branch. + Since both `DSparkConfig` and `DFlyConfig` subclass `DFlashConfig`, they are + checked before the `DFlashConfig` branch. +- **`angelspec/models/__init__.py`** / **`angelspec/models/draft/__init__.py`**: + export `DFlyModel`; drop the `TreeflashDSparkDFlareDraftModel` export. + +### DSpark slim-down + +- **`angelspec/models/draft/dspark.py`**: removed the hidden-states correction, + `PositionAdaptiveAlpha`, position-adaptive Markov knobs, and related + parameters — DSpark now only carries the Markov head and confidence head. +- **`angelspec/models/dspark.py`**: corresponding wrapper cleanup. + +### End-to-end multi-step TV loss + +- **`angelspec/models/dflash.py`**: add `_compute_e2e_tv_loss`, an independent + γ-step MTP TV term added on top of the total loss (not mutually exclusive with + KL/LK), gated on `e2e_tv_loss_weight > 0` and the presence of target + `last_hidden_states`. Emits `e2e_tv_loss` in `loss_components`. + + L_e2e = 1 - (1/γ) * Σ_{j=1..γ} Π_{i=1..j} α_i + +- **`angelspec/config/train_config.py`**: add `dflash_e2e_tv_loss_weight` + (default `0.0`, disabled) and `DatasetConfig.num_proc` (default `64`); remove + the now-unused `dflash_kl_temperature`, `dflash_kl_topk_renormalize`, and + `dflash_lk_temperature`. + +### Removals + +- **`angelspec/models/draft/treeflash_dspark_dflare.py`** (deleted). +- **`tests/test_treeflash.py`** (deleted). +- **`angelspec/config/dflare_dspark_treeflash_qwen3_4b_draft_config.json`** (deleted). + +### Configs + +- **`angelspec/config/dfly_*_draft_config.json`**: switch from + `architectures: ["DSparkDraftModel"]` / `model_type: "dspark"` / + `model_arch: "dfly"` to `architectures: ["Qwen3DFlyModel"]` / + `model_type: "qwen3"`; drop the DSpark-only `markov_rank` / + `enable_confidence_head` / `confidence_head_with_markov` keys. +- **`configs/vllm_qwen3_8b_dfly.yaml`**, **`configs/vllm_hy3_dfly.yaml`**: fix + `lm_head_key` to `lm_head.weight`, migrate `dspark_*` hyperparameters to the + `dflash_*` namespace, and set up the two-stage loss schedule (cold-start + lk-loss → final `e2e_tv_loss`). + +### Docs & tests + +- **`docs/concepts/dfly.md`**: document DFly as an independent `DFlyConfig` / + `Qwen3DFlyModel` / `DFlyTrainer` drafter; update the comparison table. +- **`docs/concepts/dspark.md`**, **`dflash.md`**, **`draft_model_family.md`**: + minor updates reflecting the moved correction module and the new loss term. +- **`tests/test_dfly.py`**: updated to build via `DFlyConfig` and exercise the + `DFlyModel` wrapper; asserts a plain `DSparkConfig` still routes to + `DSparkDraftModel` (no cross-routing). + +## Compatibility / migration notes + +- **Breaking config change:** existing DFly checkpoints/configs using + `architectures: ["DSparkDraftModel"]` + `model_arch: "dfly"` must be migrated + to `architectures: ["Qwen3DFlyModel"]` (see updated `dfly_*_draft_config.json`). +- The removed `dflash_kl_temperature` / `dflash_kl_topk_renormalize` / + `dflash_lk_temperature` training args are no longer accepted. +- `treeflash_dspark_dflare` is gone; any references must be removed. + +## Testing + +- `tests/test_dfly.py` covers auto-dispatch, model structure (shared-KV layers, + re-added `context_proj`, inherited fusion), the zero-init identity of the + hidden correction, a tiny forward through `DFlyModel` (finite loss, correct + `loss_components`, correction actually affects the loss), and state-dict keys. diff --git a/angelspec/config/dflare_dspark_treeflash_qwen3_4b_draft_config.json b/angelspec/config/dflare_dspark_treeflash_qwen3_4b_draft_config.json deleted file mode 100644 index 7006345..0000000 --- a/angelspec/config/dflare_dspark_treeflash_qwen3_4b_draft_config.json +++ /dev/null @@ -1,27 +0,0 @@ -{ - "architectures": ["DSparkDraftModel"], - "model_type": "dspark", - "model_arch": "dflare", - "hidden_size": 2560, - "intermediate_size": 9728, - "head_dim": 128, - "num_hidden_layers": 5, - "num_attention_heads": 32, - "num_key_value_heads": 8, - "vocab_size": 151936, - "rms_norm_eps": 1e-6, - "max_position_embeddings": 40960, - "rope_theta": 1000000.0, - "num_target_layers": 8, - "target_hidden_size": 2560, - "target_num_hidden_layers": 36, - "target_layer_ids": [1, 5, 10, 15, 20, 25, 30, 33], - "mask_token_id": 151669, - "tie_word_embeddings": false, - "markov_rank": 256, - "markov_head_type": "vanilla", - "enable_confidence_head": false, - "confidence_head_with_markov": false, - "enable_hidden_correction": true, - "hidden_correction_intermediate_size": null -} diff --git a/angelspec/config/dfly_hy3_draft_config.json b/angelspec/config/dfly_hy3_draft_config.json index 32e6599..ee51d74 100644 --- a/angelspec/config/dfly_hy3_draft_config.json +++ b/angelspec/config/dfly_hy3_draft_config.json @@ -1,7 +1,6 @@ { - "architectures": ["DSparkDraftModel"], - "model_type": "dspark", - "model_arch": "dfly", + "architectures": ["Qwen3DFlyModel"], + "model_type": "qwen3", "hidden_size": 4096, "intermediate_size": 13312, "head_dim": 128, @@ -18,9 +17,6 @@ "target_layer_ids": [1, 20, 39, 58, 77], "mask_token_id": 120009, "tie_word_embeddings": false, - "markov_rank": 0, - "enable_confidence_head": false, - "confidence_head_with_markov": false, "enable_hidden_correction": true, "hidden_correction_intermediate_size": null } diff --git a/angelspec/config/dfly_qwen3_4b_draft_config.json b/angelspec/config/dfly_qwen3_4b_draft_config.json index 169ce80..adc9528 100644 --- a/angelspec/config/dfly_qwen3_4b_draft_config.json +++ b/angelspec/config/dfly_qwen3_4b_draft_config.json @@ -1,7 +1,6 @@ { - "architectures": ["DSparkDraftModel"], - "model_type": "dspark", - "model_arch": "dfly", + "architectures": ["Qwen3DFlyModel"], + "model_type": "qwen3", "hidden_size": 2560, "intermediate_size": 9728, "head_dim": 128, @@ -18,9 +17,6 @@ "target_layer_ids": [1, 9, 17, 25, 33], "mask_token_id": 151669, "tie_word_embeddings": false, - "markov_rank": 0, - "enable_confidence_head": false, - "confidence_head_with_markov": false, "enable_hidden_correction": true, "hidden_correction_intermediate_size": null } diff --git a/angelspec/config/dfly_qwen3_8b_draft_config.json b/angelspec/config/dfly_qwen3_8b_draft_config.json index dbad53b..c5d7f5e 100644 --- a/angelspec/config/dfly_qwen3_8b_draft_config.json +++ b/angelspec/config/dfly_qwen3_8b_draft_config.json @@ -1,7 +1,6 @@ { - "architectures": ["DSparkDraftModel"], - "model_type": "dspark", - "model_arch": "dfly", + "architectures": ["Qwen3DFlyModel"], + "model_type": "qwen3", "hidden_size": 4096, "intermediate_size": 12288, "head_dim": 128, @@ -18,9 +17,6 @@ "target_layer_ids": [1, 9, 17, 25, 33], "mask_token_id": 151669, "tie_word_embeddings": false, - "markov_rank": 0, - "enable_confidence_head": false, - "confidence_head_with_markov": false, "enable_hidden_correction": true, "hidden_correction_intermediate_size": null } diff --git a/angelspec/config/train_config.py b/angelspec/config/train_config.py index e4aa975..d5c0d8c 100644 --- a/angelspec/config/train_config.py +++ b/angelspec/config/train_config.py @@ -49,6 +49,7 @@ class DatasetConfig: prompt_key: str = "conversations" shuffle_dataset: bool = True train_data_path: str = "" + num_proc: int = 64 @dataclass @@ -224,13 +225,14 @@ class TrainingConfig: dflash_ce_loss_alpha: float = 1.0 dflash_l1_loss_alpha: float = 0.0 dflash_kl_loss_weight: float = 0.0 - dflash_kl_temperature: float = 1.0 dflash_kl_topk: int = 10 - dflash_kl_topk_renormalize: bool = True dflash_lk_loss_weight: float = 0.0 dflash_lk_loss_type: str = "hybrid" # "alpha" or "hybrid" dflash_lk_eta: float = 3.0 - dflash_lk_temperature: float = 1.0 + # End-to-end multi-step TV loss (γ-step MTP; γ=block_size). Independent term + # added on top of the total loss (not mutually exclusive with KL/LK). Needs + # target last_hidden_states. 0 disables. + dflash_e2e_tv_loss_weight: float = 0.0 # Gated-sum layer-selection run only (fusion_type=gated_sum in the draft config). # Optional sparsity penalty weight on the layer gate: adds # `weight * H(softmax(gate))` to the loss to push the gate toward a peakier diff --git a/angelspec/config/utils.py b/angelspec/config/utils.py index 630cd4f..fc2f494 100644 --- a/angelspec/config/utils.py +++ b/angelspec/config/utils.py @@ -100,7 +100,8 @@ def generate_draft_model_config( warnings.warn( "No template config provided for draft model. " "Auto-generating config entirely from target model. " - "Consider providing a template via draft_model_config for full control." + "Consider providing a template via draft_model_config for full control.", + stacklevel=2, ) draft_config = { "architectures": ["LlamaForCausalLMEagle3"], diff --git a/angelspec/controller/eval.py b/angelspec/controller/eval.py index 5104105..616f4ca 100644 --- a/angelspec/controller/eval.py +++ b/angelspec/controller/eval.py @@ -30,7 +30,10 @@ import wandb from tqdm import tqdm -from angelspec.training.checkpoint import _read_checkpoint_metadata, _write_checkpoint_metadata +from angelspec.training.checkpoint import ( + _read_checkpoint_metadata, + _write_checkpoint_metadata, +) from angelspec.utils.logging import logger EVAL_CACHE_IDLE_TIMEOUT = 300.0 diff --git a/angelspec/controller/inference_manager.py b/angelspec/controller/inference_manager.py index 1d94b9a..0b3bf21 100644 --- a/angelspec/controller/inference_manager.py +++ b/angelspec/controller/inference_manager.py @@ -421,9 +421,9 @@ def _prepare_engine_inputs(self, entries: list[InferenceInput]) -> dict: if self._defer_tokenization: input_ids_ref = None packed_loss_mask_list = None - assert all(e.formatted_prompt is not None for e in entries), ( - "formatted_prompt is required when defer_tokenization is True" - ) + assert all( + e.formatted_prompt is not None for e in entries + ), "formatted_prompt is required when defer_tokenization is True" formatted_prompts = [e.formatted_prompt for e in entries] else: input_ids_ref = ray.put([e.input_ids for e in entries]) diff --git a/angelspec/controller/training_controller.py b/angelspec/controller/training_controller.py index 4a34965..4f58c56 100644 --- a/angelspec/controller/training_controller.py +++ b/angelspec/controller/training_controller.py @@ -477,9 +477,9 @@ def compute_vocab_mapping(self, target_vocab_size: int, draft_vocab_size: int) - from angelspec.data.preprocessing import generate_vocab_mapping assert self._stored_dataset is not None, "No stored dataset for vocab mapping" - assert "input_ids" in self._stored_dataset[0], ( - "compute_vocab_mapping requires input_ids in dataset. Set defer_tokenization=False to enable tokenization." - ) + assert ( + "input_ids" in self._stored_dataset[0] + ), "compute_vocab_mapping requires input_ids in dataset. Set defer_tokenization=False to enable tokenization." return generate_vocab_mapping( prompts=self._stored_dataset, target_vocab_size=target_vocab_size, diff --git a/angelspec/data/dataset.py b/angelspec/data/dataset.py index 4861d8c..18de4e6 100644 --- a/angelspec/data/dataset.py +++ b/angelspec/data/dataset.py @@ -27,7 +27,10 @@ from tqdm import tqdm from angelspec.data.parse import create_parser, has_thinking_content -from angelspec.data.preprocessing import _normalize_conversation, preprocess_conversations +from angelspec.data.preprocessing import ( + _normalize_conversation, + preprocess_conversations, +) from angelspec.data.template import TEMPLATE_REGISTRY from angelspec.data.utils import ( estimate_row_count, diff --git a/angelspec/data/parse.py b/angelspec/data/parse.py index 263621f..90b47cf 100644 --- a/angelspec/data/parse.py +++ b/angelspec/data/parse.py @@ -178,7 +178,8 @@ def format(self, conversation: "Conversation", **kwargs) -> str: if conversation[0]["role"] == "system": warnings.warn( - "The first message is from system, we will use the system prompt from the data and ignore the system prompt from the template" + "The first message is from system, we will use the system prompt from the data and ignore the system prompt from the template", + stacklevel=2, ) messages.append({"role": "system", "content": conversation[0]["content"]}) conversation = conversation[1:] @@ -191,19 +192,22 @@ def format(self, conversation: "Conversation", **kwargs) -> str: if j == 0: if role != "user": warnings.warn( - f"Conversation must start with a 'user' role, but found '{role}'. Conversation truncated." + f"Conversation must start with a 'user' role, but found '{role}'. Conversation truncated.", + stacklevel=2, ) break else: prev_role = conversation[j - 1]["role"] if role == "tool" and prev_role not in ["assistant", "tool"]: warnings.warn( - f"A 'tool' message must follow an 'assistant' or 'tool' message, but was preceded by '{prev_role}'. Conversation truncated." + f"A 'tool' message must follow an 'assistant' or 'tool' message, but was preceded by '{prev_role}'. Conversation truncated.", + stacklevel=2, ) break if role == "assistant" and prev_role not in ["user", "tool"]: warnings.warn( - f"An 'assistant' message must follow a 'user' or 'tool' message, but was preceded by '{prev_role}'. Conversation truncated." + f"An 'assistant' message must follow a 'user' or 'tool' message, but was preceded by '{prev_role}'. Conversation truncated.", + stacklevel=2, ) break messages.append(sentence) @@ -216,7 +220,9 @@ def format(self, conversation: "Conversation", **kwargs) -> str: try: return self._apply_chat_template(messages, **kwargs) except (ValueError, TypeError): - warnings.warn("Tokenizer does not have a chat_template, using fallback rendering.") + warnings.warn( + "Tokenizer does not have a chat_template, using fallback rendering.", stacklevel=2 + ) add_generation_prompt = kwargs.get("add_generation_prompt", False) parts = [] bos_token = getattr(self.tokenizer, "bos_token", None) diff --git a/angelspec/data/preprocessing.py b/angelspec/data/preprocessing.py index b1730bb..5abe685 100644 --- a/angelspec/data/preprocessing.py +++ b/angelspec/data/preprocessing.py @@ -283,7 +283,8 @@ def process_token_dict_to_mappings( else: warnings.warn( f"Unique tokens ({len(token_dict)}) exceed draft vocab size ({draft_vocab_size}). " - f"{len(token_dict) - draft_vocab_size} tokens will be dropped from the vocab mapping." + f"{len(token_dict) - draft_vocab_size} tokens will be dropped from the vocab mapping.", + stacklevel=2, ) total_frequency = sum(token_dict.values()) top_N = token_dict.most_common(draft_vocab_size) diff --git a/angelspec/data/template.py b/angelspec/data/template.py index 441d181..1928274 100644 --- a/angelspec/data/template.py +++ b/angelspec/data/template.py @@ -41,9 +41,9 @@ def __init__(self): self.templates = {} def register(self, name: str, template: ChatTemplate, override: bool = False): - assert override or name not in self.templates, ( - f"Chat template for the model type {name} has already been registered" - ) + assert ( + override or name not in self.templates + ), f"Chat template for the model type {name} has already been registered" self.templates[name] = template def get(self, name: str) -> ChatTemplate: diff --git a/angelspec/inference/engine/mooncake_hidden_states_connector.py b/angelspec/inference/engine/mooncake_hidden_states_connector.py index 8444a0e..38b9c6e 100644 --- a/angelspec/inference/engine/mooncake_hidden_states_connector.py +++ b/angelspec/inference/engine/mooncake_hidden_states_connector.py @@ -153,9 +153,9 @@ def __init__( self.cache_layers: list[str] = [] self._cache_layer_group_id: int = self._find_cache_layer_group_id(kv_cache_config) - assert self._vllm_config.speculative_config is not None, ( - "MooncakeHiddenStatesConnector requires 'extract_hidden_states' speculative method" - ) + assert ( + self._vllm_config.speculative_config is not None + ), "MooncakeHiddenStatesConnector requires 'extract_hidden_states' speculative method" spec_config = self._vllm_config.speculative_config.draft_model_config.hf_config self._layer_ids = list(getattr(spec_config, "eagle_aux_hidden_state_layer_ids", [])) self.num_hidden_states = len(self._layer_ids) @@ -271,9 +271,9 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self._vllm_config, CacheOnlyAttentionLayer, list(kv_caches.keys()) ) self.cache_layers = list(layers.keys()) - assert len(self.cache_layers) == 1, ( - f"Expected 1 CacheOnlyAttentionLayer, got {len(self.cache_layers)}" - ) + assert ( + len(self.cache_layers) == 1 + ), f"Expected 1 CacheOnlyAttentionLayer, got {len(self.cache_layers)}" def save_kv_layer( self, diff --git a/angelspec/inference/engine/score_worker_ext.py b/angelspec/inference/engine/score_worker_ext.py index 44c3277..05bd8cb 100644 --- a/angelspec/inference/engine/score_worker_ext.py +++ b/angelspec/inference/engine/score_worker_ext.py @@ -121,7 +121,9 @@ def score_packed( from vllm.config import CUDAGraphMode from vllm.forward_context import set_forward_context - from angelspec.models.ops.flex_attention import compile_friendly_create_block_mask + from angelspec.models.ops.flex_attention import ( + compile_friendly_create_block_mask, + ) from angelspec.models.ops.tree_mask import create_tree_mask_mod layers = self._score_attention_layers() # fail-closed guard diff --git a/angelspec/inference/engine/sgl_engine.py b/angelspec/inference/engine/sgl_engine.py index 81ab376..703d932 100644 --- a/angelspec/inference/engine/sgl_engine.py +++ b/angelspec/inference/engine/sgl_engine.py @@ -290,9 +290,9 @@ def init( else: engine_kwargs["disable_cuda_graph"] = True - assert pre_allocated_port is not None, ( - f"SglEngine rank {self.rank}: pre_allocated_port is required (ports must be pre-allocated by the factory)" - ) + assert ( + pre_allocated_port is not None + ), f"SglEngine rank {self.rank}: pre_allocated_port is required (ports must be pre-allocated by the factory)" engine_kwargs["port"] = pre_allocated_port engine_kwargs["nccl_port"] = pre_allocated_port + 1 diff --git a/angelspec/models/__init__.py b/angelspec/models/__init__.py index d768b6d..1102555 100644 --- a/angelspec/models/__init__.py +++ b/angelspec/models/__init__.py @@ -19,6 +19,7 @@ # SOFTWARE. from angelspec.models.dflash import DFlashModel +from angelspec.models.dfly import DFlyModel from angelspec.models.dspark import DSparkModel from angelspec.models.eagle3 import ( Eagle3Model, @@ -32,6 +33,7 @@ __all__ = [ "Eagle3Model", "DFlashModel", + "DFlyModel", "DSparkModel", "MTPModel", "compute_lazy_target_padded", diff --git a/angelspec/models/dflash.py b/angelspec/models/dflash.py index 05d4956..80a908a 100644 --- a/angelspec/models/dflash.py +++ b/angelspec/models/dflash.py @@ -32,7 +32,7 @@ import torch.nn.functional as F from angelspec.models.ops.flex_attention import compile_friendly_create_block_mask -from angelspec.models.ops.loss import lk_tv_kl_per_pos +from angelspec.models.ops.loss import _kl_variant_b, lk_tv_kl_per_pos from angelspec.utils.logging import logger _VALID_DFLASH_LOSS_OBJECTIVES = {"decay", "dpace"} @@ -138,13 +138,11 @@ def __init__( ce_loss_alpha: float = 1.0, l1_loss_alpha: float = 0.0, kl_loss_weight: float = 0.0, - kl_temperature: float = 1.0, kl_topk: int = 10, - kl_topk_renormalize: bool = True, lk_loss_weight: float = 0.0, lk_loss_type: str = "hybrid", lk_eta: float = 3.0, - lk_temperature: float = 1.0, + e2e_tv_loss_weight: float = 0.0, ): super().__init__() loss_objective = loss_objective.lower() @@ -175,13 +173,13 @@ def __init__( # are convex-mix coefficients in [0, 1] against CE; LK and KL are mutually # exclusive (LK takes precedence when both are set). self.kl_loss_weight = float(kl_loss_weight) - self.kl_temperature = float(kl_temperature) self.kl_topk = int(kl_topk) - self.kl_topk_renormalize = bool(kl_topk_renormalize) self.lk_loss_weight = float(lk_loss_weight) self.lk_loss_type = str(lk_loss_type) self.lk_eta = float(lk_eta) - self.lk_temperature = float(lk_temperature) + # End-to-end multi-step TV loss (independent term, added to the total; + # not mutually exclusive with KL/LK). 0 => off. Fixed T=1. + self.e2e_tv_loss_weight = float(e2e_tv_loss_weight) def _sample_anchor_positions( self, @@ -363,11 +361,11 @@ def _draft_backbone( ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: """Shared DFlash backbone (steps 1-6): context features → anchor sampling → noise embedding → position ids → block-causal mask → draft - forward. Both ``DFlashModel.forward`` and DSpark/TreeFlash subclasses build - the draft hidden states this exact way; only the label/loss tail differs. + forward. ``DFlashModel.forward`` and DSpark/DFly subclasses build the + draft hidden states this way; only the label/loss tail differs. Doc-aware (``ctx_doc_ids`` / ``base_position_ids``) and anchor-injection - arguments are threaded through so packing and parity tests keep working. + args are threaded through for packing and parity tests. Returns: draft_hidden: [B, n_blocks*block_size, D] pre-loss draft hidden states @@ -435,66 +433,21 @@ def _compute_l1_loss( ) -> torch.Tensor: """L1 distribution-distillation loss (DSpark ``l1_per_token``). - Per-position L1 distance between the student and teacher full-vocab - next-token distributions: ``Σ_i |softmax(student)_i - softmax(target)_i|`` - (== 2·TV). Float32 for stable O(V) softmax sums; no temperature. - - Args: - student_logits: [N, V] student logits. - teacher_logits: [N, V] teacher logits (detached). - Returns: - [N] per-position L1 distance. + Per-position L1 distance ``Σ_i |softmax(student)_i - softmax(teacher)_i|`` + between the full-vocab next-token distributions, which equals ``2·TV``. + Returns [N]. """ - student_probs = torch.softmax(student_logits.float(), dim=-1) - target_probs = torch.softmax(teacher_logits.float(), dim=-1) - return (student_probs - target_probs).abs().sum(dim=-1) + tv, _ = lk_tv_kl_per_pos(student_logits, teacher_logits, form="tv") + return 2.0 * tv def _compute_topk_kl_loss_variant_b( self, student_logits: torch.Tensor, teacher_logits: torch.Tensor, - temperature: float = 1.0, topk: int = 10, - renormalize_teacher_topk: bool = True, ) -> torch.Tensor: - """Top-K KL divergence (Variant B) for DFlash distillation. - - Teacher probs + student log-probs over the FULL vocab, then gather the - teacher's top-k. With ``renormalize_teacher_topk`` the gathered teacher - probs are renormalized to sum to 1 (student log-probs still full-vocab) — - a non-negative top-k objective that still penalizes mass outside the - teacher top-k. Otherwise falls back to truncated full-vocab KL. - - Returns [N] per-position KL (scaled by T^2). - """ - v = student_logits.shape[-1] - - teacher_logits_scaled = teacher_logits.float() / temperature - student_logits_scaled = student_logits.float() / temperature - - teacher_probs = torch.softmax(teacher_logits_scaled, dim=-1) - student_log_probs = torch.log_softmax(student_logits_scaled, dim=-1) - - if 0 < topk < v: - _, topk_indices = torch.topk(teacher_logits_scaled, topk, dim=-1) - - teacher_probs_topk = torch.gather(teacher_probs, -1, topk_indices) - student_log_probs_topk = torch.gather(student_log_probs, -1, topk_indices) - if renormalize_teacher_topk: - teacher_probs_topk = teacher_probs_topk / teacher_probs_topk.sum( - dim=-1, keepdim=True - ).clamp_min(1e-12) - teacher_log_probs_topk = torch.log(teacher_probs_topk.clamp_min(1e-10)) - - kl_per_position = ( - teacher_probs_topk * (teacher_log_probs_topk - student_log_probs_topk) - ).sum(dim=-1) - else: - kl_per_position = F.kl_div(student_log_probs, teacher_probs, reduction="none").sum( - dim=-1 - ) - - return kl_per_position * (temperature**2) + """Top-K KL divergence (Variant B) for DFlash distillation. Returns [N].""" + return _kl_variant_b(student_logits, teacher_logits, topk) def _compute_lk_loss( self, @@ -502,7 +455,6 @@ def _compute_lk_loss( teacher_logits: torch.Tensor, loss_type: str = "hybrid", eta: float = 3.0, - temperature: float = 1.0, ) -> torch.Tensor: """LK (acceptance-rate) distillation loss for DFlash. @@ -511,27 +463,62 @@ def _compute_lk_loss( lambda = exp(-eta * sg[alpha]), alpha = sum_i min(p_i, q_i). ``p`` = teacher (detached), ``q`` = student, both full-vocab. Returns [N] - per-position LK loss (no temperature^2 scaling). + per-position LK loss. """ if loss_type == "alpha": - teacher_logits_scaled = teacher_logits.float() / temperature - student_logits_scaled = student_logits.float() / temperature - teacher_probs = torch.softmax(teacher_logits_scaled, dim=-1) - student_probs = torch.log_softmax(student_logits_scaled, dim=-1).exp() - alpha = torch.minimum(teacher_probs, student_probs).sum(dim=-1) # [N] - return -torch.log(alpha.clamp_min(1e-10)) + # alpha = Σ_i min(p_i, q_i) == 1 − TV(p, q); reuse the shared TV term. + tv, _ = lk_tv_kl_per_pos(student_logits, teacher_logits, form="tv") + alpha = (1.0 - tv).clamp_min(1e-10) # [N] + return -torch.log(alpha) if loss_type == "hybrid": - ell, _tv = lk_tv_kl_per_pos( - student_logits / temperature, - teacher_logits / temperature, - form="lk", - eta=eta, - ) + ell, _tv = lk_tv_kl_per_pos(student_logits, teacher_logits, form="lk", eta=eta) return ell raise ValueError(f"Unknown lk_loss_type={loss_type!r}; expected 'alpha' or 'hybrid'.") + @staticmethod + def _compute_e2e_tv_loss( + student_logits_pb: torch.Tensor, + teacher_logits_pb: torch.Tensor, + valid_mask_pb: torch.Tensor, + ): + """End-to-end multi-step TV loss (γ-step accepted-length objective):: + + α_i = 1 - TV(p_i, q_i) = Σ_v min(p_i,v, q_i,v) ∈ (0, 1] + L_e2e = 1 - (1/γ) * Σ_{j=1..γ} Π_{i=1..j} α_i + + γ = block_size. The prefix product couples steps inside a block, giving + intrinsic per-step weighting, so this term ignores decay / flat_weights. + Inputs are ``[B, n_blocks, block_size, V]`` logits and a + ``[B, n_blocks, block_size]`` validity mask (T=1). Returns + ``(e2e_tv_loss, accept_length)`` (the latter detached, for logging). + """ + # fp32 for the O(V) min-sum and the chain of up-to-γ products. + t = torch.softmax(teacher_logits_pb.float(), dim=-1) + s = torch.softmax(student_logits_pb.float(), dim=-1) + + # α = Σ_v min(p, q); gradient flows through the student branch of min. + alpha = torch.minimum(t, s).sum(dim=-1) # [B, nb, γ] + + # Set α:=1 on invalid slots so cumprod treats them as identity. + m = valid_mask_pb.float() + alpha_effective = alpha * m + (1.0 - m) + prefix_prod = torch.cumprod(alpha_effective, dim=-1) # [B, nb, γ] + + gamma_valid = m.sum(dim=-1).clamp(min=1.0) # [B, nb] + accept_length_pb = (prefix_prod * m).sum(dim=-1) # [B, nb] + e2e_per_block = 1.0 - accept_length_pb / gamma_valid + + block_has_valid = (m.sum(dim=-1) > 0).float() + denom = block_has_valid.sum().clamp(min=1.0) + e2e_tv_loss = (e2e_per_block * block_has_valid).sum() / denom + + with torch.no_grad(): + accept_length = (accept_length_pb * block_has_valid).sum() / denom + + return e2e_tv_loss, accept_length.detach() + # ------------------------------------------------------------------ # Subclass extension hooks (no-ops for base DFlash). DSpark / TreeFlash # override these to inject hidden-state correction + Markov logit bias and @@ -570,10 +557,8 @@ def _compute_extra_loss( ) -> Tuple[torch.Tensor, dict]: """Add subclass-specific loss terms on top of the DFlash objective. - Returns ``(loss, extra_components)`` where ``extra_components`` is a dict - of detached scalars merged into ``loss_components`` for logging (so a new - term is just a dict key + an entry in the trainer's - ``_extra_loss_component_keys`` — no tuple/forward changes). DFlash adds + Returns ``(loss, extra_components)``, where ``extra_components`` holds + detached scalars merged into ``loss_components`` for logging. DFlash adds nothing; DSpark adds e.g. ``{"confidence_loss": ...}``.""" return loss, {} @@ -747,15 +732,15 @@ def forward( anchor_token_ids = torch.gather(input_ids, 1, anchor_positions.clamp(0, seq_len - 1)) prev_token_ids = torch.cat([anchor_token_ids.unsqueeze(-1), target_ids[:, :, :-1]], dim=-1) - # Chunked draft-logit projection (memory: full-vocab logits + their CE - # gradient are ~half the training-step peak). Only the production path — - # decay objective, no distillation, no subclass teacher head — is chunked; - # everything else keeps the single full projection below unchanged. + # Chunked draft-logit projection (full-vocab logits + their CE gradient + # are ~half the training-step peak). Only the production path — decay + # objective, no distillation, no subclass teacher head — is chunked. chunk = _dflash_loss_chunk() distill_active = ( self.l1_loss_alpha > 0 or (self.lk_loss_weight > 0.0 and last_hidden_states is not None) or (self.kl_loss_weight > 0.0 and last_hidden_states is not None) + or (self.e2e_tv_loss_weight > 0.0 and last_hidden_states is not None) or self._extra_distill_needed() ) if chunk > 0 and self.loss_objective == "decay" and not distill_active: @@ -860,12 +845,14 @@ def forward( base_loss = loss kl_loss = torch.zeros((), device=device, dtype=base_loss.dtype) lk_loss = torch.zeros((), device=device, dtype=base_loss.dtype) + e2e_tv_loss = torch.zeros((), device=device, dtype=base_loss.dtype) lk_active = self.lk_loss_weight > 0.0 and last_hidden_states is not None kl_active = ( (not lk_active) and self.kl_loss_weight > 0.0 and last_hidden_states is not None ) - want_teacher = lk_active or kl_active or self._extra_distill_needed() + e2e_tv_active = self.e2e_tv_loss_weight > 0.0 and last_hidden_states is not None + want_teacher = lk_active or kl_active or e2e_tv_active or self._extra_distill_needed() teacher_logits_flat = None if want_teacher and last_hidden_states is not None: @@ -894,7 +881,6 @@ def forward( teacher_logits=teacher_logits_flat, loss_type=self.lk_loss_type, eta=self.lk_eta, - temperature=self.lk_temperature, ) lk_loss = ( (lk_per_position * flat_weights.float()).sum() / valid_token_count.float() @@ -909,9 +895,7 @@ def forward( kl_per_position = self._compute_topk_kl_loss_variant_b( student_logits=flat_logits, teacher_logits=teacher_logits_flat, - temperature=self.kl_temperature, topk=self.kl_topk, - renormalize_teacher_topk=self.kl_topk_renormalize, ) kl_loss = ( (kl_per_position * flat_weights.float()).sum() / valid_token_count.float() @@ -923,6 +907,20 @@ def forward( else distill_w * kl_loss + (1.0 - distill_w) * base_loss ) + # 9c'. Independent e2e multi-step TV term, added on top of the total + # (not mutually exclusive with KL/LK). Bypasses flat_weights/decay. + if e2e_tv_active and teacher_logits_flat is not None: + vocab_size_e2e = flat_logits.size(-1) + e2e_tv_loss, _accept_len = self._compute_e2e_tv_loss( + student_logits_pb=flat_logits.view(bsz, n_blocks, self.block_size, vocab_size_e2e), + teacher_logits_pb=teacher_logits_flat.view( + bsz, n_blocks, self.block_size, vocab_size_e2e + ), + valid_mask_pb=weight_mask, + ) + e2e_tv_loss = e2e_tv_loss.to(base_loss.dtype) + loss = loss + self.e2e_tv_loss_weight * e2e_tv_loss + # 9d. Subclass extra-loss hook (DSpark confidence head; no-op for DFlash). loss, extra_components = self._compute_extra_loss( loss, @@ -971,6 +969,7 @@ def forward( "ce_loss": ce_component.detach(), "kl_loss": kl_loss.detach(), "lk_loss": lk_loss.detach(), + "e2e_tv_loss": e2e_tv_loss.detach(), } if l1_per_token is not None: loss_components["l1_loss"] = ( diff --git a/angelspec/models/dfly.py b/angelspec/models/dfly.py new file mode 100644 index 0000000..df45001 --- /dev/null +++ b/angelspec/models/dfly.py @@ -0,0 +1,40 @@ +"""DFly training wrapper — DFlash backbone + TreeFlash hidden-states correction. + +Like :class:`DFlashModel`, but applies the optional TreeFlash correction +(formula (1)) to the drafter output before the LM head, conditioning the token +distribution on the previous token. The correction lives on +``draft_model.hidden_correction`` and is a no-op (``None``) when disabled. +""" + +import torch +import torch.nn.functional as F + +from angelspec.models.dflash import DFlashModel + + +class DFlyModel(DFlashModel): + """DFly training wrapper (DFlash backbone + hidden-states correction).""" + + def _compute_draft_logits( + self, + draft_hidden: torch.Tensor, + lm_head_weight: torch.Tensor, + prev_token_ids: torch.Tensor, + n_blocks: int, + ) -> torch.Tensor: + """Apply hidden-states correction (if present), then project to logits. + + ``prev_token_ids`` is ``[B, n_blocks, block_size]`` — the ground-truth + token preceding each draft slot's target (aligned with ``draft_hidden``). + """ + # Correction (TreeFlash formula (1)) BEFORE the LM head, conditioning the + # token distribution on the previous token. + if getattr(self.draft_model, "hidden_correction", None) is not None: + bsz = draft_hidden.size(0) + prev_embeds = self.draft_model.embed_tokens(prev_token_ids) + prev_embeds = prev_embeds.view(bsz, -1, prev_embeds.size(-1)) + draft_hidden = self.draft_model.hidden_correction(draft_hidden, prev_embeds) + + if hasattr(self.draft_model, "lm_head"): + return self.draft_model.lm_head(draft_hidden) + return F.linear(draft_hidden, lm_head_weight) diff --git a/angelspec/models/draft/__init__.py b/angelspec/models/draft/__init__.py index 36ab30a..4ffc89a 100644 --- a/angelspec/models/draft/__init__.py +++ b/angelspec/models/draft/__init__.py @@ -25,7 +25,6 @@ from angelspec.models.draft.dfly import DFlyDraftModel from angelspec.models.draft.dspark import DSparkConfig, DSparkDraftModel from angelspec.models.draft.llama3_eagle import LlamaForCausalLMEagle3 -from angelspec.models.draft.treeflash_dspark_dflare import TreeflashDSparkDFlareDraftModel __all__ = [ "AutoDraftModelConfig", @@ -38,5 +37,4 @@ "DSparkConfig", "DSparkDraftModel", "LlamaForCausalLMEagle3", - "TreeflashDSparkDFlareDraftModel", ] diff --git a/angelspec/models/draft/auto.py b/angelspec/models/draft/auto.py index 36a7894..ff069cd 100644 --- a/angelspec/models/draft/auto.py +++ b/angelspec/models/draft/auto.py @@ -28,6 +28,7 @@ from angelspec.models.draft.deepseek_eagle import Eagle3DeepseekV2ForCausalLM from angelspec.models.draft.dflash import DFlashConfig, DFlashDraftModel +from angelspec.models.draft.dfly import DFlyConfig, DFlyDraftModel from angelspec.models.draft.dspark import DSparkConfig, DSparkDraftModel from angelspec.models.draft.llama3_eagle import LlamaForCausalLMEagle3 from angelspec.models.draft.mtp import MTPConfig, MTPDraftModel @@ -40,6 +41,7 @@ class AutoEagle3DraftModel(AutoModelForCausalLMBase): DeepseekV3Config: Eagle3DeepseekV2ForCausalLM, DFlashConfig: DFlashDraftModel, DSparkConfig: DSparkDraftModel, + DFlyConfig: DFlyDraftModel, MTPConfig: MTPDraftModel, } @@ -55,24 +57,6 @@ def from_config(cls, config: PretrainedConfig, torch_dtype=None, **config_kwargs from angelspec.models.draft.dflare import DFlareDraftModel _model_cls = DFlareDraftModel - # DSpark + ``model_arch == "dflare"`` → TreeFlash (DFlare backbone + DSpark - # heads + hidden-states correction). Same rationale as the DFlare branch - # above: rebuild the exact architecture so its extra modules survive. - if _model_cls is DSparkDraftModel and getattr(config, "model_arch", "dflash") == "dflare": - from angelspec.models.draft.treeflash_dspark_dflare import ( - TreeflashDSparkDFlareDraftModel, - ) - - _model_cls = TreeflashDSparkDFlareDraftModel - # DSpark + ``model_arch == "dfly"`` → DFlareV2 (DFlash shared-KV layers + - # DFlash FC context with a DFlare fusion residual + hidden-states - # correction). Same rationale: rebuild the exact architecture so its - # ``context_proj`` / ``layer_fusion_weights`` / ``hidden_correction`` - # survive a round-trip. - if _model_cls is DSparkDraftModel and getattr(config, "model_arch", "dflash") == "dfly": - from angelspec.models.draft.dfly import DFlyDraftModel - - _model_cls = DFlyDraftModel model = _model_cls(config, **config_kwargs) if torch_dtype is not None: @@ -109,6 +93,7 @@ class AutoDraftModelConfig: "Eagle3DeepseekV2ForCausalLM": DeepseekV3Config, "DFlashDraftModel": DFlashConfig, "DSparkDraftModel": DSparkConfig, + "Qwen3DFlyModel": DFlyConfig, "MTPDraftModel": MTPConfig, } diff --git a/angelspec/models/draft/base.py b/angelspec/models/draft/base.py index 4c4f737..f04bdcc 100644 --- a/angelspec/models/draft/base.py +++ b/angelspec/models/draft/base.py @@ -183,9 +183,9 @@ def get_lm_head_params(self) -> Tuple[torch.Tensor, torch.Tensor, float]: @torch.no_grad() def set_vocab_buffers(self, d2t: torch.Tensor, t2d: torch.Tensor) -> None: """Set the t2d/d2t vocab mapping buffers directly from tensors.""" - assert hasattr(self, "t2d") and hasattr(self, "d2t"), ( - "t2d and d2t buffers are not found in the draft model" - ) + assert hasattr(self, "t2d") and hasattr( + self, "d2t" + ), "t2d and d2t buffers are not found in the draft model" self.t2d.copy_(t2d) self.d2t.copy_(d2t) diff --git a/angelspec/models/draft/dflare.py b/angelspec/models/draft/dflare.py index c90bc3e..ee44e62 100644 --- a/angelspec/models/draft/dflare.py +++ b/angelspec/models/draft/dflare.py @@ -127,7 +127,9 @@ def forward( k = (k * cos_k) + (_rotate_half(k) * sin_k) if block_mask is not None: - from angelspec.models.ops.flex_attention import compile_friendly_flex_attention + from angelspec.models.ops.flex_attention import ( + compile_friendly_flex_attention, + ) attn_output = compile_friendly_flex_attention( query=q, diff --git a/angelspec/models/draft/dflash.py b/angelspec/models/draft/dflash.py index e97532b..fe851bb 100644 --- a/angelspec/models/draft/dflash.py +++ b/angelspec/models/draft/dflash.py @@ -262,7 +262,9 @@ def forward( k = (k * cos_k) + (_rotate_half(k) * sin_k) if block_mask is not None: - from angelspec.models.ops.flex_attention import compile_friendly_flex_attention + from angelspec.models.ops.flex_attention import ( + compile_friendly_flex_attention, + ) # Use enable_gqa=True to let FlexAttention handle GQA internally # instead of materializing expanded KV via _repeat_kv diff --git a/angelspec/models/draft/dfly.py b/angelspec/models/draft/dfly.py index 2842521..49a56dc 100644 --- a/angelspec/models/draft/dfly.py +++ b/angelspec/models/draft/dfly.py @@ -1,20 +1,10 @@ -"""DFlareV2 (``dfly``) draft model. - -DFlareV2 combines DFlash's shared FC context with DFlare's per-draft-layer -target fusion: +"""DFly draft model: DFlash shared-FC context + DFlare per-layer fusion residual. layer_context_i = RMSNorm(FC(concat(target_hidden)) + fusion_i(target_hidden)) -The resulting context has shape ``[B, S, hidden_size]`` and is consumed by -DFlash decoder layers, where target context and draft hidden states share the -same K/V projections. - -Selected via ``DSparkConfig`` + ``model_arch == "dfly"`` (dispatched in -``AutoEagle3DraftModel.from_config`` / ``DSparkTrainer.init_model``). The -optional hidden-states correction is the shared TreeFlash module from -``dspark.py`` and is applied at train time by the ``DSparkModel`` wrapper's -``_compute_draft_logits`` hook (which reads ``draft_model.hidden_correction`` -generically) — this drafter only carries the module. +Selected via ``DFlyConfig`` (architecture ``"Qwen3DFlyModel"``). The optional +hidden-states correction (TreeFlash formula (1)) is applied at train time by the +``DFlyModel`` wrapper's ``_compute_draft_logits`` hook. """ from typing import Optional @@ -24,35 +14,128 @@ import torch.nn.functional as F from angelspec.models.draft.dflare import DFlareDraftModel -from angelspec.models.draft.dflash import DFlashDecoderLayer -from angelspec.models.draft.dspark import DSparkConfig, build_hidden_correction +from angelspec.models.draft.dflash import ( + DFlashConfig, + DFlashDecoderLayer, + DFlashRMSNorm, +) + + +class DFlyConfig(DFlashConfig): + """DFly config: :class:`DFlashConfig` plus TreeFlash hidden-correction knobs.""" + + model_type = "qwen3" + + def __init__( + self, + enable_hidden_correction: bool = True, + hidden_correction_intermediate_size: Optional[int] = None, + **kwargs, + ): + super().__init__(**kwargs) + self.enable_hidden_correction = enable_hidden_correction + self.hidden_correction_intermediate_size = hidden_correction_intermediate_size + + +class HiddenStatesCorrection(nn.Module): + """TreeFlash hidden-states correction (formula (1)), residual + zero-init:: + + h'_{t+i} = h_{t+i} + SwiGLU( norm(h_{t+i}) :: norm(e_{t+i-1}) ) + + ``e_{t+i-1}`` is the previous (teacher-forced) token embedding, ``::`` is + feature-dim concat. The output projection is zero-initialized, so the + correction starts at 0 and the model degenerates back to DFlash. + """ + + def __init__( + self, + hidden_size: int, + embed_size: int, + intermediate_size: int, + rms_norm_eps: float = 1e-6, + ): + super().__init__() + self.hidden_size = int(hidden_size) + self.embed_size = int(embed_size) + self.intermediate_size = int(intermediate_size) + + # Normalize each stream independently before concat so the hidden state + # and the differently-scaled token embedding contribute comparably. + self.hidden_norm = DFlashRMSNorm(self.hidden_size, eps=rms_norm_eps) + self.embed_norm = DFlashRMSNorm(self.embed_size, eps=rms_norm_eps) + + in_features = self.hidden_size + self.embed_size + self.gate_proj = nn.Linear(in_features, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(in_features, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + + # Residual zero-init: correction starts at 0 -> exactly recovers DFlash. + nn.init.zeros_(self.down_proj.weight) + + def forward( + self, hidden_states: torch.Tensor, prev_token_embeds: torch.Tensor + ) -> torch.Tensor: + """Apply the residual correction. + + Args: + hidden_states: ``[..., hidden_size]`` drafter output hidden states. + prev_token_embeds: ``[..., embed_size]`` previous-token embeddings + (same leading shape as ``hidden_states``). + + Returns: + Corrected hidden states, same shape/dtype as ``hidden_states``. + """ + h_norm = self.hidden_norm(hidden_states) + e_norm = self.embed_norm(prev_token_embeds.to(hidden_states.dtype)) + gate_in = torch.cat([h_norm, e_norm], dim=-1) + delta = self.down_proj(F.silu(self.gate_proj(gate_in)) * self.up_proj(gate_in)) + delta = delta.to(hidden_states.dtype) + + return hidden_states + delta + + +def build_hidden_correction(config) -> Optional[nn.Module]: + """Build the hidden-states correction module, or ``None`` if disabled. + + The previous-token embedding dim equals the draft ``hidden_size`` (the + drafter reuses the target token embedding). The SwiGLU intermediate width + defaults to ``hidden_size`` unless ``hidden_correction_intermediate_size`` + is set. + """ + if not bool(getattr(config, "enable_hidden_correction", False)): + return None + + hidden_size = int(config.hidden_size) + intermediate = getattr(config, "hidden_correction_intermediate_size", None) + intermediate = int(intermediate) if intermediate else hidden_size + return HiddenStatesCorrection( + hidden_size=hidden_size, + embed_size=hidden_size, + intermediate_size=intermediate, + rms_norm_eps=getattr(config, "rms_norm_eps", 1e-6), + ) class DFlyDraftModel(DFlareDraftModel): - """DFlash layers with DFlash FC context plus DFlare fusion residual. + """DFlash decoder layers + DFlash FC context + DFlare fusion residual. - Rides the DSpark path (``DSparkConfig`` / ``DSparkTrainer`` / ``DSparkModel``) - so the shared ``hidden_correction`` runs through the wrapper hook; it builds - no Markov or confidence head (the hooks read those via ``getattr(..., None)`` - and no-op when absent). + Carries the optional ``hidden_correction`` module, applied by the + ``DFlyModel`` wrapper via the shared ``_compute_draft_logits`` hook. """ - config_class = DSparkConfig + config_class = DFlyConfig def __init__(self, config): super().__init__(config) - # DFlare builds layers with separate context/draft K/V projections and - # deletes DFlash's ``context_proj`` (keeping the reinitialized - # ``context_norm`` + ``layer_fusion_weights``). DFlareV2 intentionally - # restores DFlash layers so both sources share the same k_proj/v_proj, - # and re-adds the FC ``context_proj`` while retaining the DFlare fusion. + # Restore DFlash layers so context and draft share k_proj/v_proj, while + # keeping DFlare's per-layer fusion (context_norm + layer_fusion_weights). self.layers = nn.ModuleList([DFlashDecoderLayer(config) for _ in range(self.num_layers)]) target_hidden_size = getattr(config, "target_hidden_size", config.hidden_size) if target_hidden_size != config.hidden_size: raise ValueError( - "DFlareV2 residual fusion requires target_hidden_size == hidden_size, " + "DFly residual fusion requires target_hidden_size == hidden_size, " f"got target_hidden_size={target_hidden_size} and hidden_size={config.hidden_size}" ) @@ -62,8 +145,7 @@ def __init__(self, config): bias=False, ) - # Shared TreeFlash hidden-states correction (dspark.build_hidden_correction); - # ``None`` when ``enable_hidden_correction`` is unset. Applied by DSparkModel. + # TreeFlash correction (formula (1)); ``None`` unless enabled. Applied by DFlyModel. self.hidden_correction = build_hidden_correction(config) def _project_base_context(self, context_feature: torch.Tensor) -> torch.Tensor: @@ -99,7 +181,7 @@ def forward( """Run the draft model with FC + fusion-residual context per layer.""" if context_feature.ndim != 4: raise ValueError( - f"DFlareV2 context_feature must have shape [B, S, T, D], got {tuple(context_feature.shape)}" + f"DFly context_feature must have shape [B, S, T, D], got {tuple(context_feature.shape)}" ) if context_feature.shape[2] != self.num_target_layers: raise ValueError( @@ -125,4 +207,4 @@ def forward( return self.final_norm(draft_hidden) -__all__ = ["DFlyDraftModel"] +__all__ = ["DFlyConfig", "DFlyDraftModel", "HiddenStatesCorrection", "build_hidden_correction"] diff --git a/angelspec/models/draft/dspark.py b/angelspec/models/draft/dspark.py index 1a6ac69..6110953 100644 --- a/angelspec/models/draft/dspark.py +++ b/angelspec/models/draft/dspark.py @@ -1,30 +1,25 @@ -""" -DSpark draft model: DFlash backbone + EAGLE-style Markov and confidence heads. +"""DSpark draft model: DFlash backbone + EAGLE-style Markov and confidence heads. -DSpark shares DFlash's block-diffusion drafter (dual-source KV injection, anchor -sampling, MASK-token noise stream) and adds two heads on top: +Shares DFlash's block-diffusion drafter and adds two heads: - - Markov head: a low-rank learned bigram bias added to the draft logits, - conditioned on the (teacher-forced) previous token. Improves the per-token - distribution without touching the backbone. - - Confidence head (AcceptRatePredictor): predicts a per-draft-position - acceptance probability, trained against the empirical draft-vs-target - accept rate (used at inference time for adaptive block length). + - Markov head: low-rank learned bigram logit bias on the (teacher-forced) + previous token; refines the per-token distribution without touching the + backbone. + - Confidence head (AcceptRatePredictor): predicts per-position acceptance + probability, trained against the empirical accept rate (used at inference + for adaptive block length). """ from typing import Optional import torch import torch.nn as nn -import torch.nn.functional as F -from angelspec.models.draft.dflash import DFlashConfig, DFlashDraftModel, DFlashRMSNorm +from angelspec.models.draft.dflash import DFlashConfig, DFlashDraftModel class DSparkConfig(DFlashConfig): - """ - Configuration for the DSpark draft model. Extends :class:`DFlashConfig`. - """ + """Config for the DSpark draft model. Extends :class:`DFlashConfig`.""" model_type = "dspark" @@ -32,124 +27,15 @@ def __init__( self, markov_rank: int = 256, markov_head_type: str = "vanilla", - markov_pos_adaptive: bool = False, - markov_alpha_max: float = 1.0, - markov_alpha_start: float = 0.1, - markov_alpha_end: float = 1.0, - markov_smooth_lambda: float = 0.0, enable_confidence_head: bool = True, confidence_head_with_markov: bool = True, - enable_hidden_correction: bool = True, - hidden_correction_intermediate_size: Optional[int] = None, - hidden_correction_pos_adaptive: bool = False, - hidden_correction_alpha_max: float = 0.8, - hidden_correction_alpha_start: float = 0.1, - hidden_correction_alpha_end: float = 0.5, - hidden_correction_smooth_lambda: float = 0.0, - block_size: Optional[int] = None, **kwargs, ): super().__init__(**kwargs) self.markov_rank = markov_rank self.markov_head_type = markov_head_type - # Position-adaptive transition strength for the Markov logit bias - # (``logits_i += alpha_i * bias_i``), sharing the same design as the - # hidden-states correction below: ``alpha_i = alpha_max * sigmoid(w_i)`` - # is a per-in-block-position learnable scalar, initialized to a - # monotonically increasing ramp (weak prefix, strong suffix). Adds only - # ``block_size`` scalar parameters. NOTE: the idea doc recommends the - # hidden-states-injection variant over this logits-bias variant; this is - # provided so both heads can carry an independent alpha curve. - self.markov_pos_adaptive = markov_pos_adaptive - self.markov_alpha_max = markov_alpha_max - self.markov_alpha_start = markov_alpha_start - self.markov_alpha_end = markov_alpha_end - self.markov_smooth_lambda = markov_smooth_lambda self.enable_confidence_head = enable_confidence_head self.confidence_head_with_markov = confidence_head_with_markov - # TreeFlash hidden-states correction (formula (1)): a lightweight SwiGLU - # applied to the drafter's output hidden state, conditioned on the - # previous token's embedding, added back in residual form. - self.enable_hidden_correction = enable_hidden_correction - # Intermediate width of the correction SwiGLU. ``None`` -> lightweight - # default of ``hidden_size`` (kept small since this rides on top of the - # full backbone MLP). - self.hidden_correction_intermediate_size = hidden_correction_intermediate_size - # Position-adaptive transition strength for the hidden-states correction - # (``h'_i = h_i + alpha_i * delta_i``). ``alpha_i = alpha_max * sigmoid(w_i)`` - # is a per-in-block-position learnable scalar, initialized to a - # monotonically increasing ramp (weak prefix correction, strong suffix - # correction) to match the per-position drift-accumulation prior. Adds - # only ``block_size`` scalar parameters, no extra compute at inference. - self.hidden_correction_pos_adaptive = hidden_correction_pos_adaptive - self.hidden_correction_alpha_max = hidden_correction_alpha_max - self.hidden_correction_alpha_start = hidden_correction_alpha_start - self.hidden_correction_alpha_end = hidden_correction_alpha_end - # Optional smoothness regularizer weight lambda for the alpha curve - # (``lambda * sum_i (alpha_i - alpha_{i-1})^2``); 0 disables it. - self.hidden_correction_smooth_lambda = hidden_correction_smooth_lambda - # Block length K (number of in-block positions). Required to size the - # position-adaptive alpha vector; injected by the trainer from its - # ``block_size`` knob when not set explicitly in the config JSON. - self.block_size = block_size - - -class PositionAdaptiveAlpha(nn.Module): - """Per-in-block-position transition strength ``alpha_i`` (length ``block_size``). - - Implements ``alpha_i = alpha_max * sigmoid(w_i)`` with a learnable logit - ``w_i`` initialized so ``alpha_i`` follows a monotonically increasing ramp - from ``alpha_start`` to ``alpha_end`` (weak prefix correction, strong suffix - correction — matching the per-position drift-accumulation prior). Adds only - ``block_size`` scalar parameters and no inference-time compute beyond a - broadcasted multiply. - - Optionally exposes a smoothness regularizer - ``smooth_lambda * sum_i (alpha_i - alpha_{i-1})^2`` via :meth:`smooth_loss`. - """ - - def __init__( - self, - *, - block_size: int, - alpha_max: float = 1.0, - alpha_start: float = 0.1, - alpha_end: float = 1.0, - smooth_lambda: float = 0.0, - ): - super().__init__() - if block_size is None or int(block_size) <= 0: - raise ValueError( - "PositionAdaptiveAlpha requires a positive block_size (the block " - f"length K); got {block_size!r}. It is injected by the trainer " - "from ``dflash_block_size``; set it in the config JSON when " - "building the model outside the trainer." - ) - self.block_size = int(block_size) - self.alpha_max = float(alpha_max) - self.smooth_lambda = float(smooth_lambda) - - # Initialize the ramp in alpha-space, then invert the sigmoid to obtain - # the logit initialization ``w_i``. - if self.block_size == 1: - ramp = torch.full((1,), float(alpha_end)) - else: - ramp = torch.linspace(float(alpha_start), float(alpha_end), self.block_size) - frac = (ramp / self.alpha_max).clamp(1e-4, 1.0 - 1e-4) - w_init = torch.log(frac / (1.0 - frac)) - self.alpha_logit = nn.Parameter(w_init) - - def alpha(self) -> torch.Tensor: - """Return the current ``alpha`` vector of shape ``[block_size]``.""" - return self.alpha_max * torch.sigmoid(self.alpha_logit) - - def smooth_loss(self) -> Optional[torch.Tensor]: - """Smoothness penalty on the alpha curve, or ``None`` when disabled.""" - if self.smooth_lambda <= 0.0 or self.block_size < 2: - return None - a = self.alpha() - diff = a[1:] - a[:-1] - return self.smooth_lambda * (diff * diff).sum() class VanillaMarkov(nn.Module): @@ -160,35 +46,18 @@ def __init__( *, vocab_size: int, markov_rank: int, - pos_adaptive: bool = False, - block_size: Optional[int] = None, - alpha_max: float = 1.0, - alpha_start: float = 0.1, - alpha_end: float = 1.0, - smooth_lambda: float = 0.0, ): super().__init__() self.vocab_size = int(vocab_size) self.markov_rank = int(markov_rank) self.markov_head_type = "vanilla" - assert self.markov_rank > 0, ( - f"VanillaMarkov requires markov_rank > 0, got {self.markov_rank}." - ) + assert ( + self.markov_rank > 0 + ), f"VanillaMarkov requires markov_rank > 0, got {self.markov_rank}." self.markov_w1 = nn.Embedding(self.vocab_size, self.markov_rank) - # TODO: markow_w2 out_features should match "draft_vocab_size" if pruning is used. + # TODO: markov_w2 out_features should match "draft_vocab_size" if pruning is used. self.markov_w2 = nn.Linear(self.markov_rank, self.vocab_size, bias=False) - # Position-adaptive per-in-block-position strength on the logit bias. - self.pos_alpha: Optional[PositionAdaptiveAlpha] = None - if pos_adaptive: - self.pos_alpha = PositionAdaptiveAlpha( - block_size=block_size, - alpha_max=alpha_max, - alpha_start=alpha_start, - alpha_end=alpha_end, - smooth_lambda=smooth_lambda, - ) - def get_prev_embeddings(self, token_ids: torch.Tensor) -> torch.Tensor: return self.markov_w1(token_ids.long()) @@ -204,14 +73,9 @@ def apply_block_logits( *, token_ids: torch.Tensor, ) -> torch.Tensor: - # ``base_logits`` is ``[B, n_blocks, block_size, V]``; the bias is - # broadcast per in-block position when position-adaptive alpha is on. if base_logits.size(2) == 0: return base_logits bias = self.compute_step_bias(token_ids) - if self.pos_alpha is not None: - alpha = self.pos_alpha.alpha().to(bias.dtype) # [block_size] - bias = bias * alpha.view(1, 1, -1, 1) return base_logits + bias @@ -226,109 +90,6 @@ def forward(self, features: torch.Tensor) -> torch.Tensor: return self.proj(features).squeeze(-1) -class HiddenStatesCorrection(nn.Module): - """TreeFlash hidden-states correction (formula (1)). - - Refines the drafter's output hidden state by conditioning on the previous - token, without breaking the ``O(1)`` complexity of the backbone:: - - h'_{t+i} = h_{t+i} + SwiGLU( norm(h_{t+i}) :: norm(e_{t+i-1}) ) - - where - - * ``h_{t+i}`` : the drafter's original output hidden state at position - ``t + i`` (fed to the LM head); - * ``norm(h_{t+i})`` (``h̃_{t+i}`` in the paper) : the RMS-normalized hidden - state; - * ``e_{t+i-1}`` : the RMS-normalized input embedding of the previous token - ``x_{t+i-1}`` (teacher-forced); - * ``::`` : concatenation along the feature dimension; - * ``SwiGLU`` : a Swish-gated linear unit producing a correction with the - same dimension as ``h_{t+i}``. - - The output projection is zero-initialized so the correction starts at 0 and - the model degenerates back to DFlash at initialization (residual form). - """ - - def __init__( - self, - hidden_size: int, - embed_size: int, - intermediate_size: int, - rms_norm_eps: float = 1e-6, - pos_adaptive: bool = False, - block_size: Optional[int] = None, - alpha_max: float = 0.8, - alpha_start: float = 0.1, - alpha_end: float = 0.5, - smooth_lambda: float = 0.0, - ): - super().__init__() - self.hidden_size = int(hidden_size) - self.embed_size = int(embed_size) - self.intermediate_size = int(intermediate_size) - - # Position-adaptive per-in-block-position strength on the residual delta. - self.pos_alpha: Optional[PositionAdaptiveAlpha] = None - if pos_adaptive: - self.pos_alpha = PositionAdaptiveAlpha( - block_size=block_size, - alpha_max=alpha_max, - alpha_start=alpha_start, - alpha_end=alpha_end, - smooth_lambda=smooth_lambda, - ) - - # Normalize each input stream independently before concatenation, so the - # hidden state and the (differently-scaled) token embedding contribute on - # a comparable footing. - self.hidden_norm = DFlashRMSNorm(self.hidden_size, eps=rms_norm_eps) - self.embed_norm = DFlashRMSNorm(self.embed_size, eps=rms_norm_eps) - - in_features = self.hidden_size + self.embed_size - self.gate_proj = nn.Linear(in_features, self.intermediate_size, bias=False) - self.up_proj = nn.Linear(in_features, self.intermediate_size, bias=False) - self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) - - # Residual zero-init: correction = 0 at start -> exactly recovers DFlash. - nn.init.zeros_(self.down_proj.weight) - - def forward( - self, hidden_states: torch.Tensor, prev_token_embeds: torch.Tensor - ) -> torch.Tensor: - """Apply the residual correction. - - Args: - hidden_states: ``[..., hidden_size]`` drafter output hidden states. - prev_token_embeds: ``[..., embed_size]`` input embeddings of the - previous token (same leading shape as ``hidden_states``). - - Returns: - Corrected hidden states, same shape/dtype as ``hidden_states``. - """ - h_norm = self.hidden_norm(hidden_states) - e_norm = self.embed_norm(prev_token_embeds.to(hidden_states.dtype)) - gate_in = torch.cat([h_norm, e_norm], dim=-1) - delta = self.down_proj(F.silu(self.gate_proj(gate_in)) * self.up_proj(gate_in)) - delta = delta.to(hidden_states.dtype) - - # Scale the correction per in-block position: ``h'_i = h_i + alpha_i * delta_i``. - # ``hidden_states`` is ``[B, n_blocks * block_size, hidden]``; fold out the - # block-position axis to broadcast the ``[block_size]`` alpha vector. - if self.pos_alpha is not None: - K = self.pos_alpha.block_size - bsz, n_pos, hid = delta.shape - if n_pos % K != 0: - raise ValueError( - f"HiddenStatesCorrection position count {n_pos} is not a " - f"multiple of block_size {K}; cannot apply position-adaptive alpha." - ) - alpha = self.pos_alpha.alpha().to(delta.dtype) # [block_size] - delta = (delta.view(bsz, -1, K, hid) * alpha.view(1, 1, K, 1)).reshape(bsz, n_pos, hid) - - return hidden_states + delta - - def build_markov_head(config) -> Optional[nn.Module]: markov_rank = int(getattr(config, "markov_rank", 0)) assert markov_rank >= 0, f"markov_rank must be >= 0, got {markov_rank}" @@ -340,12 +101,6 @@ def build_markov_head(config) -> Optional[nn.Module]: return VanillaMarkov( vocab_size=config.vocab_size, markov_rank=markov_rank, - pos_adaptive=bool(getattr(config, "markov_pos_adaptive", False)), - block_size=getattr(config, "block_size", None), - alpha_max=float(getattr(config, "markov_alpha_max", 1.0)), - alpha_start=float(getattr(config, "markov_alpha_start", 0.1)), - alpha_end=float(getattr(config, "markov_alpha_end", 1.0)), - smooth_lambda=float(getattr(config, "markov_smooth_lambda", 0.0)), ) raise NotImplementedError( f"markov_head_type={markov_head_type!r} is not supported yet; only 'vanilla' " @@ -353,34 +108,6 @@ def build_markov_head(config) -> Optional[nn.Module]: ) -def build_hidden_correction(config) -> Optional[nn.Module]: - """Build the TreeFlash hidden-states correction module, or ``None``. - - The previous-token embedding dimension equals the draft ``hidden_size`` - (the drafter reuses the target model's token embedding). The SwiGLU - intermediate width defaults to ``hidden_size`` (lightweight) unless - ``hidden_correction_intermediate_size`` is set on the config. - """ - if not bool(getattr(config, "enable_hidden_correction", False)): - return None - - hidden_size = int(config.hidden_size) - intermediate = getattr(config, "hidden_correction_intermediate_size", None) - intermediate = int(intermediate) if intermediate else hidden_size - return HiddenStatesCorrection( - hidden_size=hidden_size, - embed_size=hidden_size, - intermediate_size=intermediate, - rms_norm_eps=getattr(config, "rms_norm_eps", 1e-6), - pos_adaptive=bool(getattr(config, "hidden_correction_pos_adaptive", False)), - block_size=getattr(config, "block_size", None), - alpha_max=float(getattr(config, "hidden_correction_alpha_max", 0.8)), - alpha_start=float(getattr(config, "hidden_correction_alpha_start", 0.1)), - alpha_end=float(getattr(config, "hidden_correction_alpha_end", 0.5)), - smooth_lambda=float(getattr(config, "hidden_correction_smooth_lambda", 0.0)), - ) - - class DSparkDraftModel(DFlashDraftModel): config_class = DSparkConfig @@ -394,9 +121,6 @@ def __init__(self, config: DSparkConfig): self.markov_head = build_markov_head(config) - # TreeFlash hidden-states correction (formula (1)); ``None`` when disabled. - self.hidden_correction = build_hidden_correction(config) - self.confidence_head: Optional[nn.Module] = None if getattr(config, "enable_confidence_head", False): conf_input_dim = self.hidden_size diff --git a/angelspec/models/draft/llama3_eagle.py b/angelspec/models/draft/llama3_eagle.py index 751f8b8..c573e76 100644 --- a/angelspec/models/draft/llama3_eagle.py +++ b/angelspec/models/draft/llama3_eagle.py @@ -77,12 +77,8 @@ def _import_standard_flash_attn(): import flash_attn as mod from flash_attn import flash_attn_varlen_func as varlen_func from flash_attn.bert_padding import pad_input, unpad_input - from flash_attn.flash_attn_interface import ( - _flash_attn_backward as backward, - ) - from flash_attn.flash_attn_interface import ( - _flash_attn_forward as forward, - ) + from flash_attn.flash_attn_interface import _flash_attn_backward as backward + from flash_attn.flash_attn_interface import _flash_attn_forward as forward from flash_attn.flash_attn_interface import ( _flash_attn_varlen_backward as varlen_backward, ) @@ -1033,12 +1029,12 @@ def _get_block_sparse( _block_sparse_cache[cache_key] = BlockSparseTensorsTorch( mask_block_cnt=cnt.expand(bsz, num_heads, -1).contiguous(), mask_block_idx=idx.expand(bsz, num_heads, -1, -1).contiguous(), - full_block_cnt=f_cnt.expand(bsz, num_heads, -1).contiguous() - if f_cnt is not None - else None, - full_block_idx=f_idx.expand(bsz, num_heads, -1, -1).contiguous() - if f_idx is not None - else None, + full_block_cnt=( + f_cnt.expand(bsz, num_heads, -1).contiguous() if f_cnt is not None else None + ), + full_block_idx=( + f_idx.expand(bsz, num_heads, -1, -1).contiguous() if f_idx is not None else None + ), block_size=block_size, ) return _block_sparse_cache[cache_key] @@ -1493,9 +1489,9 @@ def forward( cache_keys = key_states.unsqueeze(1) cache_values = value_states.unsqueeze(1) - assert attention_mask is not None, ( - "LlamaFlashAttention cached path requires attention_mask" - ) + assert ( + attention_mask is not None + ), "LlamaFlashAttention cached path requires attention_mask" valid_lengths = attention_mask.sum(dim=-1, dtype=torch.long) - lck valid_lengths = valid_lengths.clamp_(0, q_len) diff --git a/angelspec/models/draft/mtp.py b/angelspec/models/draft/mtp.py index 485b3c5..24d07d0 100644 --- a/angelspec/models/draft/mtp.py +++ b/angelspec/models/draft/mtp.py @@ -602,21 +602,21 @@ def __init__(self, config: MTPConfig): self.use_sigmoid = config.moe_router_use_sigmoid H = config.hidden_size - I = config.moe_intermediate_size + inter = config.moe_intermediate_size E = self.num_experts # Fused expert weights [E, in, out] (transpose of nn.Linear.weight [out, in]) # so they feed straight into torch._grouped_mm(mat_a=[T, in], mat_b) → [T, out]. - self.experts_gate_proj = nn.Parameter(torch.empty(E, H, I)) # [E, H, I] - self.experts_up_proj = nn.Parameter(torch.empty(E, H, I)) # [E, H, I] - self.experts_down_proj = nn.Parameter(torch.empty(E, I, H)) # [E, I, H] + self.experts_gate_proj = nn.Parameter(torch.empty(E, H, inter)) # [E, H, I] + self.experts_up_proj = nn.Parameter(torch.empty(E, H, inter)) # [E, H, I] + self.experts_down_proj = nn.Parameter(torch.empty(E, inter, H)) # [E, I, H] self.act_fn = nn.SiLU() # Match nn.Linear's default init (kaiming_uniform_, a=sqrt(5)) with the # correct fan_in: init in the [out, in] layout then transpose into place. with torch.no_grad(): for w, out_dim, in_dim in ( - (self.experts_gate_proj, I, H), - (self.experts_up_proj, I, H), - (self.experts_down_proj, H, I), + (self.experts_gate_proj, inter, H), + (self.experts_up_proj, inter, H), + (self.experts_down_proj, H, inter), ): for e in range(E): tmp = torch.empty(out_dim, in_dim) diff --git a/angelspec/models/draft/treeflash_dspark_dflare.py b/angelspec/models/draft/treeflash_dspark_dflare.py deleted file mode 100644 index 7ae4eec..0000000 --- a/angelspec/models/draft/treeflash_dspark_dflare.py +++ /dev/null @@ -1,55 +0,0 @@ -"""TreeFlash draft model: DFlare backbone + DSpark Markov / confidence heads + -TreeFlash hidden-states correction. - -``TreeflashDSparkDFlareDraftModel`` is a :class:`DFlareDraftModel` (the DFlash -variant with learnable per-layer target fusion) that additionally carries the -DSpark heads — a low-rank Markov logit bias, an optional accept-rate confidence -head, and the TreeFlash hidden-states correction (formula (1), zero-init -residual). It holds only the modules; the DFlare backbone forward is reused -unchanged, and the heads are applied by the :class:`DSparkModel` training wrapper -through its ``_compute_draft_logits`` / ``_compute_extra_loss`` hooks (which read -``draft_model.{markov_head, hidden_correction, confidence_head}`` generically). - -Selected via ``DSparkConfig`` + ``model_arch == "dflare"`` (dispatched in -``AutoEagle3DraftModel.from_config`` / ``dspark_trainer._build_draft_model``). -""" - -from typing import Optional - -import torch.nn as nn - -from angelspec.models.draft.dflare import DFlareDraftModel -from angelspec.models.draft.dspark import ( - AcceptRatePredictor, - DSparkConfig, - build_hidden_correction, - build_markov_head, -) - - -class TreeflashDSparkDFlareDraftModel(DFlareDraftModel): - config_class = DSparkConfig - - def __init__(self, config: DSparkConfig): - super().__init__(config) - - self.markov_rank = int(getattr(config, "markov_rank", 0)) - self.confidence_head_with_markov = bool( - getattr(config, "confidence_head_with_markov", True) - ) - - self.markov_head = build_markov_head(config) - - # TreeFlash hidden-states correction (formula (1)); ``None`` when disabled. - self.hidden_correction = build_hidden_correction(config) - - self.confidence_head: Optional[nn.Module] = None - if getattr(config, "enable_confidence_head", False): - conf_input_dim = self.hidden_size - if self.confidence_head_with_markov: - if self.markov_head is None: - raise ValueError( - "confidence_head_with_markov=True requires a Markov head (markov_rank > 0)." - ) - conf_input_dim += self.markov_rank - self.confidence_head = AcceptRatePredictor(conf_input_dim) diff --git a/angelspec/models/dspark.py b/angelspec/models/dspark.py index cc1a7d1..3894a62 100644 --- a/angelspec/models/dspark.py +++ b/angelspec/models/dspark.py @@ -1,24 +1,3 @@ -"""DSpark training model: DFlash training wrapper + Markov / confidence heads. - -:class:`DSparkModel` reuses the **entire** :class:`DFlashModel` training forward -— anchor sampling, block-causal FlexAttention mask, MASK-token noise, and the -full CE / KL / LK / L1 / D-PACE loss pipeline (with the same label alignment and -loss knobs). The only additions ride on the frozen DFlash subclass hooks so none -of the DFlash loss is re-implemented: - - - ``_compute_draft_logits``: applies the TreeFlash hidden-states correction - (``h' = h + SwiGLU(norm(h) :: norm(e_prev))``) before the LM head and the - Markov logit bias after it, both conditioned on the teacher-forced previous - token. Zero-init keeps these no-ops at start (degenerates to DFlash). - - ``_compute_extra_loss``: adds a confidence-head BCE against the empirical - per-token accept rate ``1 - 0.5 * L1(draft, teacher)`` (+ any - position-adaptive alpha smoothness penalty), surfaced as ``confidence_loss``. - -Total loss = `` + confidence_head_alpha * confidence`` where -```` is the objective selected by the shared knobs -(``loss_objective`` / ``ce_loss_alpha`` / ``l1_loss_alpha`` / ``kl_*`` / ``lk_*``). -""" - import torch import torch.nn.functional as F @@ -39,20 +18,17 @@ def __init__( ce_loss_alpha: float = 0.1, l1_loss_alpha: float = 0.0, kl_loss_weight: float = 0.0, - kl_temperature: float = 1.0, kl_topk: int = 10, - kl_topk_renormalize: bool = True, lk_loss_weight: float = 0.0, lk_loss_type: str = "hybrid", lk_eta: float = 3.0, - lk_temperature: float = 1.0, + e2e_tv_loss_weight: float = 0.0, fp32_lm_head: bool = True, gate_entropy_weight: float = 0.0, confidence_head_alpha: float = 1.0, ): - # Forward the full DFlash loss configuration to the parent so DSpark's - # base loss is identical to DFlash's; ``confidence`` is the only - # DSpark-specific term (added in ``_compute_extra_loss``). + # Forward the full DFlash loss config to the parent; ``confidence`` is + # DSpark's only extra term (added in ``_compute_extra_loss``). super().__init__( draft_model=draft_model, block_size=block_size, @@ -65,18 +41,15 @@ def __init__( ce_loss_alpha=ce_loss_alpha, l1_loss_alpha=l1_loss_alpha, kl_loss_weight=kl_loss_weight, - kl_temperature=kl_temperature, kl_topk=kl_topk, - kl_topk_renormalize=kl_topk_renormalize, lk_loss_weight=lk_loss_weight, lk_loss_type=lk_loss_type, lk_eta=lk_eta, - lk_temperature=lk_temperature, + e2e_tv_loss_weight=e2e_tv_loss_weight, ) self.confidence_head_alpha = float(confidence_head_alpha) - # Handoff buffer for the corrected draft hidden states between the - # ``_compute_draft_logits`` and ``_compute_extra_loss`` hooks within a - # single forward pass. + # Handoff buffer for corrected draft hidden states, between the + # ``_compute_draft_logits`` and ``_compute_extra_loss`` hooks. self._dspark_hidden_4d = None # ------------------------------------------------------------------ @@ -93,18 +66,10 @@ def _compute_draft_logits( """Inject hidden-states correction + Markov bias, then project to logits. ``prev_token_ids`` is ``[B, n_blocks, block_size]`` — the ground-truth - token immediately preceding each draft slot's target (aligned - slot-for-slot with the flattened ``draft_hidden`` layout). + token preceding each draft slot's target (aligned with ``draft_hidden``). """ bsz = draft_hidden.size(0) - # Hidden-states correction (TreeFlash formula (1)), BEFORE the LM head so - # the token distribution is conditioned on the previous token. - if getattr(self.draft_model, "hidden_correction", None) is not None: - prev_embeds = self.draft_model.embed_tokens(prev_token_ids) - prev_embeds = prev_embeds.view(bsz, -1, prev_embeds.size(-1)) - draft_hidden = self.draft_model.hidden_correction(draft_hidden, prev_embeds) - # Cache the corrected hidden states for the confidence head. self._dspark_hidden_4d = draft_hidden.view(bsz, n_blocks, self.block_size, -1) @@ -139,26 +104,12 @@ def _compute_extra_loss( ): """Add the confidence-head BCE against the empirical accept rate. - Uses the SAME objective-weighted validity mask (``flat_weights``) and - weighted-mean reduction as the DFlash CE / distillation terms. Returns - ``(loss, {"confidence_loss": ...})`` so the component is logged via the - shared ``loss_components`` mechanism. + Uses the same objective-weighted mask (``flat_weights``) and weighted-mean + reduction as the DFlash CE / distillation terms. Returns + ``(loss, {"confidence_loss": ...})`` for shared logging. """ confidence_loss = torch.zeros((), device=loss.device, dtype=loss.dtype) - # Position-adaptive alpha smoothness regularizer (independent of the - # confidence head / teacher logits): ``lambda * sum_i (alpha_i - alpha_{i-1})^2`` - # over any head (Markov / hidden-correction) that enables it. - for _head in ( - getattr(self.draft_model, "markov_head", None), - getattr(self.draft_model, "hidden_correction", None), - ): - _pa = getattr(_head, "pos_alpha", None) if _head is not None else None - if _pa is not None: - _reg = _pa.smooth_loss() - if _reg is not None: - loss = loss + _reg.to(loss.dtype) - # Confidence BCE needs the teacher accept-rate target; skip when the head # is off or target last_hidden_states weren't delivered this step. if not self._extra_distill_needed() or teacher_logits_flat is None: diff --git a/angelspec/models/eagle3.py b/angelspec/models/eagle3.py index 1353353..313edc1 100644 --- a/angelspec/models/eagle3.py +++ b/angelspec/models/eagle3.py @@ -127,9 +127,11 @@ def _calculate_loss( use_sum_lazy_loss = self.attention_backend == "usp" if self.gradient_checkpointing and self.training: return torch_checkpoint( - compiled_sum_forward_kl_loss_from_hs - if use_sum_lazy_loss - else compiled_forward_kl_loss_from_hs, + ( + compiled_sum_forward_kl_loss_from_hs + if use_sum_lazy_loss + else compiled_forward_kl_loss_from_hs + ), *args, use_reentrant=False, ) diff --git a/angelspec/models/ops/flex_attention.py b/angelspec/models/ops/flex_attention.py index 1f23c27..aca0b32 100644 --- a/angelspec/models/ops/flex_attention.py +++ b/angelspec/models/ops/flex_attention.py @@ -319,9 +319,9 @@ def build_eagle3_block_mask( Q_BS, KV_BS = _normalize_block_size(BLOCK_SIZE) assert Q_LEN % Q_BS == 0 and KV_LEN % KV_BS == 0 assert Q_BS % KV_BS == 0, f"Q_BS ({Q_BS}) must be a multiple of KV_BS ({KV_BS})" - assert KV_LEN % Q_LEN == 0, ( - f"build_eagle3_block_mask requires KV_LEN to be a multiple of Q_LEN; got Q_LEN={Q_LEN}, KV_LEN={KV_LEN}" - ) + assert ( + KV_LEN % Q_LEN == 0 + ), f"build_eagle3_block_mask requires KV_LEN to be a multiple of Q_LEN; got Q_LEN={Q_LEN}, KV_LEN={KV_LEN}" # Skip the compiled path when nested inside another torch.compile graph. builder = ( diff --git a/angelspec/train_entry.py b/angelspec/train_entry.py index 195db6b..1a55ac0 100644 --- a/angelspec/train_entry.py +++ b/angelspec/train_entry.py @@ -26,38 +26,43 @@ import time os.environ.setdefault("TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS", "ATEN,TRITON") -from collections import namedtuple -from contextlib import contextmanager -from typing import Any, Generator - -import ray -from omegaconf import OmegaConf -from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy - -from angelspec import AutoDraftModelConfig -from angelspec.config.train_config import config_to_flat_args, load_config -from angelspec.config.utils import generate_draft_model_config -from angelspec.controller import ( +# The env var above must be set before importing torch/ray-dependent modules, +# so the following imports intentionally sit below it. +from collections import namedtuple # noqa: E402 +from contextlib import contextmanager # noqa: E402 +from typing import Any, Generator # noqa: E402 + +import ray # noqa: E402 +from omegaconf import OmegaConf # noqa: E402 +from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy # noqa: E402 + +from angelspec import AutoDraftModelConfig # noqa: E402 +from angelspec.config.train_config import config_to_flat_args, load_config # noqa: E402 +from angelspec.config.utils import generate_draft_model_config # noqa: E402 +from angelspec.controller import ( # noqa: E402 AsyncTrainingController, auto_calculate_training_steps, build_mooncake_config, run_training_loop, setup_async_training_with_engines, ) -from angelspec.inference.factory import ( +from angelspec.inference.factory import ( # noqa: E402 prepare_eval_engine, prepare_inference_engines, prepare_score_engine, ) -from angelspec.ray.placement_group import ( +from angelspec.ray.placement_group import ( # noqa: E402 allocate_train_group, create_placement_groups, ) -from angelspec.training.trainer_actor import TrainerActor -from angelspec.transfer.mooncake.utils import launch_mooncake_master -from angelspec.utils.env import get_angelspec_env_vars -from angelspec.utils.logging import init_tracking, logger -from angelspec.utils.usp import validate_dflash_usp_layout, validate_mtp_usp_layout +from angelspec.training.trainer_actor import TrainerActor # noqa: E402 +from angelspec.transfer.mooncake.utils import launch_mooncake_master # noqa: E402 +from angelspec.utils.env import get_angelspec_env_vars # noqa: E402 +from angelspec.utils.logging import init_tracking, logger # noqa: E402 +from angelspec.utils.usp import ( # noqa: E402 + validate_dflash_usp_layout, + validate_mtp_usp_layout, +) _Phase = namedtuple("_Phase", ["name", "duration", "is_async", "blocked"]) diff --git a/angelspec/training/data_fetcher.py b/angelspec/training/data_fetcher.py index 2d8441b..e600371 100644 --- a/angelspec/training/data_fetcher.py +++ b/angelspec/training/data_fetcher.py @@ -36,7 +36,11 @@ from ray.util.queue import Queue as RayQueue from torch.utils.data import DataLoader, IterableDataset -from angelspec.data.utils import deserialize_packed_loss_mask, resolve_loss_mask, unpack_loss_mask +from angelspec.data.utils import ( + deserialize_packed_loss_mask, + resolve_loss_mask, + unpack_loss_mask, +) from angelspec.utils.distributed import ( get_draft_sp_group, get_sp_ring_group, diff --git a/angelspec/training/dflash_trainer.py b/angelspec/training/dflash_trainer.py index 8589504..e608c03 100644 --- a/angelspec/training/dflash_trainer.py +++ b/angelspec/training/dflash_trainer.py @@ -49,7 +49,13 @@ class DFlashTrainer(Trainer): # ``loss_components`` dict; reduced to global means for logging (each key is # reduced only when present, so subclasses just extend this list). Subclasses # (DSpark) add e.g. "confidence_loss". - _extra_loss_component_keys: list[str] = ["ce_loss", "kl_loss", "lk_loss", "l1_loss"] + _extra_loss_component_keys: list[str] = [ + "ce_loss", + "kl_loss", + "lk_loss", + "l1_loss", + "e2e_tv_loss", + ] def __init__(self, args: Namespace): super().__init__(args) @@ -67,22 +73,67 @@ def __init__(self, args: Namespace): self.ce_loss_alpha = getattr(args, "dflash_ce_loss_alpha", 1.0) self.l1_loss_alpha = getattr(args, "dflash_l1_loss_alpha", 0.0) self.kl_loss_weight = float(getattr(args, "dflash_kl_loss_weight", 0.0)) - self.kl_temperature = float(getattr(args, "dflash_kl_temperature", 1.0)) self.kl_topk = int(getattr(args, "dflash_kl_topk", 10)) - self.kl_topk_renormalize = bool(getattr(args, "dflash_kl_topk_renormalize", True)) self.lk_loss_weight = float(getattr(args, "dflash_lk_loss_weight", 0.0)) self.lk_loss_type = str(getattr(args, "dflash_lk_loss_type", "hybrid")) self.lk_eta = float(getattr(args, "dflash_lk_eta", 3.0)) - self.lk_temperature = float(getattr(args, "dflash_lk_temperature", 1.0)) + # Independent e2e multi-step TV loss (added on top; not KL/LK-exclusive). + self.e2e_tv_loss_weight = float(getattr(args, "dflash_e2e_tv_loss_weight", 0.0)) self._lk_enabled = self.lk_loss_weight > 0.0 self._kl_enabled = (not self._lk_enabled) and self.kl_loss_weight > 0.0 - # last_hidden_states (target final norm) is required for KL/LK teacher - # logits; L1 uses raw last_hidden_states directly (no norm). - self._distill_enabled = self._lk_enabled or self._kl_enabled + # last_hidden_states (target final norm) is required for KL/LK/e2e_tv + # teacher logits; L1 uses raw last_hidden_states directly (no norm). + self._distill_enabled = ( + self._lk_enabled or self._kl_enabled or self.e2e_tv_loss_weight > 0.0 + ) # Rolling window of the top-5 candidate-layer set for the gated_sum layer # selection run; drives the topk_jaccard / backbone_size early-stop signals. self._gate_topk_window: deque = deque(maxlen=10) + # ------------------------------------------------------------------ + # Model-build seams (overridable by DFlash-family subclasses) + # ------------------------------------------------------------------ + + def _build_draft_model(self, config): + """Construct the draft model from ``config`` (dispatch by config flags). + + Subclasses (e.g. DFly) override this to build their own draft model. + """ + if getattr(config, "fusion_type", "concat_fc") == "gated_sum": + from angelspec.models.draft.dflash_gated import DFlashGatedDraftModel + + return DFlashGatedDraftModel(config) + elif getattr(config, "model_arch", "dflash") == "dflare": + from angelspec.models.draft.dflare import DFlareDraftModel + + return DFlareDraftModel(config) + return DFlashDraftModel(config) + + def _build_training_wrapper(self, draft_model): + """Wrap ``draft_model`` in the training module (loss / forward plumbing). + + Subclasses (e.g. DFly) override this to swap in a wrapper that injects + their extra behavior through the shared DFlash hooks. + """ + return DFlashModel( + draft_model=draft_model, + block_size=self.block_size, + num_anchors=self.num_anchors, + loss_decay_gamma=self.loss_decay_gamma, + fp32_lm_head=self.fp32_lm_head, + gate_entropy_weight=getattr(self.args, "dflash_gate_entropy_weight", 0.0), + loss_objective=self.loss_objective, + dpace_alpha=self.dpace_alpha, + ce_loss_alpha=self.ce_loss_alpha, + l1_loss_alpha=self.l1_loss_alpha, + kl_loss_weight=self.kl_loss_weight, + kl_topk=self.kl_topk, + lk_loss_weight=self.lk_loss_weight, + lk_loss_type=self.lk_loss_type, + lk_eta=self.lk_eta, + e2e_tv_loss_weight=self.e2e_tv_loss_weight, + ) + def init_model( self, draft_model_config, @@ -131,16 +182,7 @@ def init_model( ) config.target_num_hidden_layers = target_config.num_hidden_layers - if getattr(config, "fusion_type", "concat_fc") == "gated_sum": - from angelspec.models.draft.dflash_gated import DFlashGatedDraftModel - - draft_model = DFlashGatedDraftModel(config) - elif getattr(config, "model_arch", "dflash") == "dflare": - from angelspec.models.draft.dflare import DFlareDraftModel - - draft_model = DFlareDraftModel(config) - else: - draft_model = DFlashDraftModel(config) + draft_model = self._build_draft_model(config) if dist.get_rank() == 0: draft_model.load_embedding( @@ -160,26 +202,7 @@ def init_model( f"{frozen_count:,} frozen (embedding) parameters" ) - dflash_model = DFlashModel( - draft_model=draft_model, - block_size=self.block_size, - num_anchors=self.num_anchors, - loss_decay_gamma=self.loss_decay_gamma, - fp32_lm_head=self.fp32_lm_head, - gate_entropy_weight=getattr(self.args, "dflash_gate_entropy_weight", 0.0), - loss_objective=self.loss_objective, - dpace_alpha=self.dpace_alpha, - ce_loss_alpha=self.ce_loss_alpha, - l1_loss_alpha=self.l1_loss_alpha, - kl_loss_weight=self.kl_loss_weight, - kl_temperature=self.kl_temperature, - kl_topk=self.kl_topk, - kl_topk_renormalize=self.kl_topk_renormalize, - lk_loss_weight=self.lk_loss_weight, - lk_loss_type=self.lk_loss_type, - lk_eta=self.lk_eta, - lk_temperature=self.lk_temperature, - ) + dflash_model = self._build_training_wrapper(draft_model) full_state = dflash_model.state_dict() if dist.get_rank() == 0 else {} @@ -430,7 +453,10 @@ def _compute_opd_loss(self, input_ids: torch.Tensor, opd: dict): """ import ray - from angelspec.models.ops.loss import _OPD_METRIC_KEYS, opd_two_stream_kl_from_hs + from angelspec.models.ops.loss import ( + _OPD_METRIC_KEYS, + opd_two_stream_kl_from_hs, + ) from angelspec.models.ops.tree_layout import build_dflash_opd_batch_layout draft_hidden = opd["draft_hidden"] # [B, n_blocks*bs, D] (grad) diff --git a/angelspec/training/dfly_trainer.py b/angelspec/training/dfly_trainer.py new file mode 100644 index 0000000..587a617 --- /dev/null +++ b/angelspec/training/dfly_trainer.py @@ -0,0 +1,41 @@ +"""DFly trainer — DFlashTrainer with DFly's draft model and training wrapper. + +Reuses the whole DFlash pipeline (FSDP init, optimizer, checkpoint, schedule, +forward/metrics/eval) and overrides only the two model-build seams: builds the +DFly draft model (DFlash FC context + DFlare fusion residual) and wraps it in +``DFlyModel`` for the optional hidden-states correction. Reads the ``dflash_*`` +hyperparameter namespace. +""" + +from angelspec.training.dflash_trainer import DFlashTrainer + + +class DFlyTrainer(DFlashTrainer): + """DFly-specific trainer (DFlash backbone + hidden-states correction).""" + + def _build_draft_model(self, config): + from angelspec.models.draft.dfly import DFlyDraftModel + + return DFlyDraftModel(config) + + def _build_training_wrapper(self, draft_model): + from angelspec.models.dfly import DFlyModel + + return DFlyModel( + draft_model=draft_model, + block_size=self.block_size, + num_anchors=self.num_anchors, + loss_decay_gamma=self.loss_decay_gamma, + fp32_lm_head=self.fp32_lm_head, + gate_entropy_weight=getattr(self.args, "dflash_gate_entropy_weight", 0.0), + loss_objective=self.loss_objective, + dpace_alpha=self.dpace_alpha, + ce_loss_alpha=self.ce_loss_alpha, + l1_loss_alpha=self.l1_loss_alpha, + kl_loss_weight=self.kl_loss_weight, + kl_topk=self.kl_topk, + lk_loss_weight=self.lk_loss_weight, + lk_loss_type=self.lk_loss_type, + lk_eta=self.lk_eta, + e2e_tv_loss_weight=self.e2e_tv_loss_weight, + ) diff --git a/angelspec/training/dspark_trainer.py b/angelspec/training/dspark_trainer.py index 7b4f6ba..e1c76bf 100644 --- a/angelspec/training/dspark_trainer.py +++ b/angelspec/training/dspark_trainer.py @@ -1,12 +1,9 @@ """DSpark trainer — DFlash trainer + Markov / confidence heads. -Subclasses DFlashTrainer, reusing its whole pipeline (FSDP init, optimizer, -checkpoint, LR schedule, target LM head, forward / metrics / eval). Overrides -only ``init_model`` to swap in the DSpark draft model (or TreeFlash when -``model_arch == "dflare"``) and the ``DSparkModel`` wrapper, and extends -``_extra_loss_component_keys`` with ``confidence_loss``. DSparkModel injects its -heads through the shared DFlash hooks, so the loss / metric plumbing is inherited -unchanged. +Subclasses DFlashTrainer, reusing its whole pipeline. Overrides ``init_model`` +to swap in the DSpark draft model and the ``DSparkModel`` wrapper, and extends +``_extra_loss_component_keys`` with ``confidence_loss``. Heads inject through the +shared DFlash hooks, so loss / metric plumbing is inherited unchanged. """ from argparse import Namespace @@ -94,30 +91,8 @@ def init_model( ) config.target_num_hidden_layers = target_config.num_hidden_layers - # Position-adaptive alpha (Markov / hidden-correction) sizes its - # per-position vector from the block length; inject the trainer's - # block_size when the config JSON didn't set it explicitly. - if getattr(config, "block_size", None) is None: - config.block_size = self.block_size - # --- B-seam: swapped model construction --- - # DSpark + model_arch=="dflare" → TreeFlash (DFlare backbone + DSpark - # heads + hidden-states correction); model_arch=="dfly" → DFlareV2 - # (DFlash shared-KV layers + DFlash FC context with a DFlare fusion - # residual + hidden-states correction); otherwise the DFlash-backbone - # DSpark drafter. Mirrors the dispatch in AutoEagle3DraftModel. - if getattr(config, "model_arch", "dflash") == "dflare": - from angelspec.models.draft.treeflash_dspark_dflare import ( - TreeflashDSparkDFlareDraftModel, - ) - - draft_model = TreeflashDSparkDFlareDraftModel(config) - elif getattr(config, "model_arch", "dflash") == "dfly": - from angelspec.models.draft.dfly import DFlyDraftModel - - draft_model = DFlyDraftModel(config) - else: - draft_model = DSparkDraftModel(config) + draft_model = DSparkDraftModel(config) if dist.get_rank() == 0: draft_model.load_embedding( @@ -148,13 +123,11 @@ def init_model( ce_loss_alpha=self.ce_loss_alpha, l1_loss_alpha=self.l1_loss_alpha, kl_loss_weight=self.kl_loss_weight, - kl_temperature=self.kl_temperature, kl_topk=self.kl_topk, - kl_topk_renormalize=self.kl_topk_renormalize, lk_loss_weight=self.lk_loss_weight, lk_loss_type=self.lk_loss_type, lk_eta=self.lk_eta, - lk_temperature=self.lk_temperature, + e2e_tv_loss_weight=self.e2e_tv_loss_weight, fp32_lm_head=self.fp32_lm_head, gate_entropy_weight=getattr(self.args, "dflash_gate_entropy_weight", 0.0), confidence_head_alpha=self.confidence_head_alpha, @@ -234,10 +207,3 @@ def init_model( logger.info(f"[Rank {self.dp_rank}] DSpark model initialized with FSDP2") return 0 - - # DSparkModel now inherits the unified DFlashModel.forward (slot 0 = masked - # anchor, DFlash label alignment), so _forward / eval_forward / _train_step / - # metric aggregation are all inherited from DFlashTrainer unchanged — they - # already thread last_hidden_states + target_norm in, unpack the 6/7-tuple, - # drop slot 0, and reduce loss_components (incl. our confidence_loss via - # _extra_loss_component_keys). diff --git a/angelspec/training/eagle3_trainer.py b/angelspec/training/eagle3_trainer.py index 207beb6..09bb9e4 100644 --- a/angelspec/training/eagle3_trainer.py +++ b/angelspec/training/eagle3_trainer.py @@ -305,9 +305,9 @@ def _forward(self, batch: dict) -> Tuple[List[torch.Tensor], List[torch.Tensor]] target=target, loss_mask=loss_mask, hidden_states=batch["hidden_states"].cuda(), - position_ids=batch.get("position_ids").cuda() - if batch.get("position_ids") is not None - else None, + position_ids=( + batch.get("position_ids").cuda() if batch.get("position_ids") is not None else None + ), ) return plosses, vlosses, acces, acc_counts diff --git a/angelspec/training/mtp_trainer.py b/angelspec/training/mtp_trainer.py index f9eb349..2282a0c 100644 --- a/angelspec/training/mtp_trainer.py +++ b/angelspec/training/mtp_trainer.py @@ -118,7 +118,9 @@ def init_model( mooncake_config=None, ) -> int: if mooncake_config is not None: - from angelspec.transfer.mooncake.utils import check_mooncake_master_available + from angelspec.transfer.mooncake.utils import ( + check_mooncake_master_available, + ) check_mooncake_master_available( mooncake_config.master_server_address, mooncake_config.metadata_server @@ -280,8 +282,9 @@ def init_model( max_seq = getattr(self.args, "max_seq_length", None) if max_seq and sp_size > 1: shard_len = usp_chunk_size(max_seq, sp_size) - hidden = getattr(draft_model_config, "target_hidden_size", None) or getattr( - draft_model_config, "hidden_size" + hidden = ( + getattr(draft_model_config, "target_hidden_size", None) + or draft_model_config.hidden_size ) logger.info( f"[Rank {self.dp_rank}] USP flex warmup: shard_len={shard_len} " @@ -504,9 +507,11 @@ def _forward(self, batch: dict): target_lm_head_weight=self.target_lm_head_weight, loss_mask=loss_mask, hidden_states=draft_input, - position_ids=batch.get("position_ids").cuda() - if batch.get("position_ids") is not None - else None, + position_ids=( + batch.get("position_ids").cuda() + if batch.get("position_ids") is not None + else None + ), ctx_doc_ids=ctx_doc_ids, base_position_ids=base_position_ids, ) diff --git a/angelspec/training/optimizer.py b/angelspec/training/optimizer.py index f40c771..53da238 100644 --- a/angelspec/training/optimizer.py +++ b/angelspec/training/optimizer.py @@ -21,7 +21,10 @@ import torch from angelspec.training.lr_scheduler import LRSchedulerWithWarmup -from angelspec.training.muon_utils import adjust_lr_for_muon, zeropower_via_newtonschulz5 +from angelspec.training.muon_utils import ( + adjust_lr_for_muon, + zeropower_via_newtonschulz5, +) from angelspec.utils.logging import print_on_rank0 diff --git a/angelspec/training/trainer.py b/angelspec/training/trainer.py index 0dd7e47..b240353 100644 --- a/angelspec/training/trainer.py +++ b/angelspec/training/trainer.py @@ -38,7 +38,11 @@ from torch.distributed.device_mesh import init_device_mesh from angelspec.config.mooncake_config import MooncakeConfig -from angelspec.data.utils import DataCollatorWithPadding, DFlashPackingCollator, MTPPackingCollator +from angelspec.data.utils import ( + DataCollatorWithPadding, + DFlashPackingCollator, + MTPPackingCollator, +) from angelspec.training import checkpoint from angelspec.training.data_fetcher import MooncakeDataFetcher, PrefetchedDataFetcher from angelspec.training.fsdp import init_empty_weights @@ -811,7 +815,7 @@ def _maybe_dump(self, batch: dict, step_metrics: dict, step: int, batch_idx: int batch_idx=batch_idx, ) - def _save_dump_data( + def _save_dump_data( # noqa: B027 - optional hook with a default no-op body self, *, batch: dict, diff --git a/angelspec/training/trainer_actor.py b/angelspec/training/trainer_actor.py index 8f1698e..14e0151 100644 --- a/angelspec/training/trainer_actor.py +++ b/angelspec/training/trainer_actor.py @@ -26,6 +26,7 @@ from angelspec import AutoDraftModelConfig from angelspec.models.draft.dflash import DFlashConfig +from angelspec.models.draft.dfly import DFlyConfig from angelspec.models.draft.dspark import DSparkConfig from angelspec.models.draft.mtp import MTPConfig from angelspec.ray.ray_actor import RayActor @@ -80,12 +81,18 @@ def init( draft_model_config = AutoDraftModelConfig.from_file(args.draft_model_config) # Config-based trainer dispatch: DSparkConfig → DSparkTrainer, - # DFlashConfig → DFlashTrainer, MTPConfig → MTPTrainer, else Eagle3. - # DSparkConfig subclasses DFlashConfig, so it MUST be checked first. + # DFlyConfig → DFlyTrainer, DFlashConfig → DFlashTrainer, + # MTPConfig → MTPTrainer, else Eagle3. + # DSparkConfig and DFlyConfig both subclass DFlashConfig, so they MUST be + # checked before the DFlashConfig branch. if isinstance(draft_model_config, DSparkConfig): from angelspec.training.dspark_trainer import DSparkTrainer self._trainer = DSparkTrainer(args) + elif isinstance(draft_model_config, DFlyConfig): + from angelspec.training.dfly_trainer import DFlyTrainer + + self._trainer = DFlyTrainer(args) elif isinstance(draft_model_config, DFlashConfig): from angelspec.training.dflash_trainer import DFlashTrainer diff --git a/angelspec/transfer/mooncake/buffers.py b/angelspec/transfer/mooncake/buffers.py index 12f13ee..09e93b2 100644 --- a/angelspec/transfer/mooncake/buffers.py +++ b/angelspec/transfer/mooncake/buffers.py @@ -217,7 +217,7 @@ def _do_put( def drain(self) -> None: """Wait for every in-flight transfer to finish.""" - for ptr, future in list(self._in_flight.items()): + for _ptr, future in list(self._in_flight.items()): try: future.result() except Exception as exc: diff --git a/angelspec/transfer/mooncake/store.py b/angelspec/transfer/mooncake/store.py index 8df34e1..fde48c5 100644 --- a/angelspec/transfer/mooncake/store.py +++ b/angelspec/transfer/mooncake/store.py @@ -20,7 +20,6 @@ import os import threading -from abc import ABC from typing import Any, Dict, Optional import torch @@ -36,7 +35,7 @@ from angelspec.utils.logging import logger -class MooncakeHiddenStateStore(ABC): +class MooncakeHiddenStateStore: """ Base class for Mooncake Store wrapper to store hidden states from target model. diff --git a/angelspec/utils/profiling.py b/angelspec/utils/profiling.py index 8384d12..f909df9 100644 --- a/angelspec/utils/profiling.py +++ b/angelspec/utils/profiling.py @@ -142,9 +142,9 @@ def stop(self): class _MemrayMemoryProfiler(_BaseMemoryProfiler): def __init__(self, args): super().__init__(args) - assert args.memory_snapshot_num_steps is not None, ( - "In memray, must provide --memory-snapshot-num-steps" - ) + assert ( + args.memory_snapshot_num_steps is not None + ), "In memray, must provide --memory-snapshot-num-steps" def start(self): logger.info("Memray tracker started.") diff --git a/angelspec/utils/usp.py b/angelspec/utils/usp.py index b262f5b..c7b9e28 100644 --- a/angelspec/utils/usp.py +++ b/angelspec/utils/usp.py @@ -125,8 +125,8 @@ def _slice_and_pad(tensor: torch.Tensor, axis: int, pad_value: int = 0): ) attention_mask[:, :valid_len] = 1 - usp_chunk_size = max(local_len - ttt_length, 0) - ring_chunk = usp_chunk_size * sp_ulysses_size + local_chunk_size = max(local_len - ttt_length, 0) + ring_chunk = local_chunk_size * sp_ulysses_size ring_start = ring_rank * ring_chunk position_ids = torch.arange( ring_start, ring_start + ring_chunk, device=input_ids.device, dtype=torch.long diff --git a/configs/vllm_hy3_dfly.yaml b/configs/vllm_hy3_dfly.yaml index d692215..9e28d77 100644 --- a/configs/vllm_hy3_dfly.yaml +++ b/configs/vllm_hy3_dfly.yaml @@ -38,12 +38,16 @@ training: dflash_block_size: 16 dflash_num_target_layers: 5 - dspark_num_anchors: 128 - dspark_num_target_layers: 5 - dspark_loss_decay_gamma: 4.0 - dspark_ce_loss_alpha: 0.1 - dspark_l1_loss_alpha: 0.9 - dspark_confidence_head_alpha: 1.0 + dflash_num_anchors: 128 + + # Loss 策略（两阶段）： + # 冷启动阶段（当前配置）：纯 lk-loss（不用 CE），dpace 位置加权 + # 最终阶段：dflash_lk_loss_weight=0.0，dflash_e2e_tv_loss_weight=1.0 + dflash_loss_objective: dpace # dpace 续接价值加权（作用于 lk 的位置权重） + dflash_dpace_alpha: 0.5 + dflash_lk_loss_weight: 1.0 # =1.0 时 loss=lk，完全不含 CE + dflash_lk_loss_type: hybrid + dflash_e2e_tv_loss_weight: 0.0 inference: inference_engine_type: vllm diff --git a/configs/vllm_qwen3_8b_dfly.yaml b/configs/vllm_qwen3_8b_dfly.yaml index 90c484d..70a7846 100644 --- a/configs/vllm_qwen3_8b_dfly.yaml +++ b/configs/vllm_qwen3_8b_dfly.yaml @@ -2,7 +2,7 @@ model: target_model_path: ${oc.env:TS_SHARE_ROOT,/path/to/share}/model/Qwen_Qwen3-8B trust_remote_code: true draft_model_config: angelspec/config/dfly_qwen3_8b_draft_config.json - lm_head_key: model.embed_tokens.weight + lm_head_key: lm_head.weight dataset: train_data_path: ../examples/data/sample_conversations.jsonl @@ -38,12 +38,16 @@ training: dflash_block_size: 16 dflash_num_target_layers: 5 - dspark_num_anchors: 128 - dspark_num_target_layers: 5 - dspark_loss_decay_gamma: 4.0 - dspark_ce_loss_alpha: 0.1 - dspark_l1_loss_alpha: 0.9 - dspark_confidence_head_alpha: 1.0 + dflash_num_anchors: 128 + + # Loss 策略（两阶段）： + # 冷启动阶段（当前配置）：纯 lk-loss（不用 CE），dpace 位置加权 + # 最终阶段：dflash_lk_loss_weight=0.0，dflash_e2e_tv_loss_weight=1.0 + dflash_loss_objective: dpace # dpace 续接价值加权（作用于 lk 的位置权重） + dflash_dpace_alpha: 0.5 + dflash_lk_loss_weight: 1.0 # =1.0 时 loss=lk，完全不含 CE + dflash_lk_loss_type: hybrid + dflash_e2e_tv_loss_weight: 0.0 inference: inference_engine_type: vllm diff --git a/docs/concepts/dflash.md b/docs/concepts/dflash.md index b833881..9b74662 100644 --- a/docs/concepts/dflash.md +++ b/docs/concepts/dflash.md @@ -17,7 +17,8 @@ pass rather than one per token. sequence, and a mask token plus the anchor token embedding seed each block. - **Loss.** Cross-entropy against the ground-truth tokens, with an exponential positional decay that weights earlier in-block positions more heavily (later positions are harder to predict). - Optional distillation terms (L1, top-K KL) against the target model's logits can be mixed in. + Optional distillation terms against the target model's logits can be mixed in: L1, top-K KL, + or LK. An independent end-to-end multi-step TV term can also be added on top. DFlash is a standalone architecture — it does not inherit the Eagle3 interface, because dual-source KV and block-parallel prediction differ fundamentally from input fusion and diff --git a/docs/concepts/dfly.md b/docs/concepts/dfly.md index 6f12407..6183f82 100644 --- a/docs/concepts/dfly.md +++ b/docs/concepts/dfly.md @@ -23,26 +23,31 @@ hidden-state correction head. ## Configuration -DFly uses `DSparkConfig` with `model_arch: "dfly"`: +DFly uses its own `DFlyConfig` (an extension of `DFlashConfig`) and the +`"Qwen3DFlyModel"` architecture. It is a DFlash-family drafter and does not +depend on DSpark. The hidden-states correction knobs +(`enable_hidden_correction`, `hidden_correction_intermediate_size`) belong to +DFly only: ```json { - "architectures": ["Qwen3DSparkModel"], - "model_type": "qwen3_dspark", - "model_arch": "dfly", - "markov_rank": 0, - "enable_confidence_head": false, + "architectures": ["Qwen3DFlyModel"], + "model_type": "qwen3", "enable_hidden_correction": true } ``` ## Dispatch -The trainer dispatches on `DSparkConfig` + `model_arch == "dfly"`: +The trainer dispatches on the `DFlyConfig` type: - **Model:** `DFlyDraftModel` (in `angelspec/models/draft/dfly.py`) -- **Trainer:** `DSparkTrainer` (shared with DSpark via hooks) -- **Loss:** Inherits the DFlash composable loss (CE + decay/D-PACE + optional KL/LK) +- **Trainer:** `DFlyTrainer` (in `angelspec/training/dfly_trainer.py`) — a thin + subclass of `DFlashTrainer` that swaps in the `DFlyModel` wrapper (in + `angelspec/models/dfly.py`) via the DFlash model-build hooks. Reads the + `dflash_*` hyperparameter namespace. +- **Loss:** Inherits the DFlash composable loss (CE + decay/D-PACE + optional KL/LK, + plus an optional independent end-to-end multi-step TV term) ## Relation to other architectures @@ -50,7 +55,7 @@ The trainer dispatches on `DSparkConfig` + `model_arch == "dfly"`: |---------|--------|--------|------|--------| | Shared KV projection | ✓ | ✗ (separate) | ✓ (base) | ✓ | | Per-layer fusion | ✗ | ✓ | ✓ (residual) | ✗ | -| Hidden-state correction | ✗ | ✗ | ✓ (optional) | ✓ (TreeFlash) | +| Hidden-state correction | ✗ | ✗ | ✓ (TreeFlash) | ✗ | | Markov head | ✗ | ✗ | ✗ | ✓ | | Confidence head | ✗ | ✗ | ✗ | ✓ | diff --git a/docs/concepts/draft_model_family.md b/docs/concepts/draft_model_family.md index ed50e00..96ddd2a 100644 --- a/docs/concepts/draft_model_family.md +++ b/docs/concepts/draft_model_family.md @@ -40,7 +40,8 @@ picks the model. The trainer dispatches on the config type: ``` DFlashConfig → DFlash (model_arch="dflare" → DFlare) -DSparkConfig → DSpark (model_arch="dflare" → DFlare-backbone variant, model_arch="dfly" → DFly) +DSparkConfig → DSpark (model_arch="dflare" → DFlare-backbone variant) +DFlyConfig → DFly (architecture "Qwen3DFlyModel") MTPConfig → MTP LlamaConfig → Eagle3 (Llama family) DeepseekV3Config → Eagle3 (DeepSeek MLA family) diff --git a/docs/concepts/dspark.md b/docs/concepts/dspark.md index ac4ae3a..a5c8e20 100644 --- a/docs/concepts/dspark.md +++ b/docs/concepts/dspark.md @@ -14,10 +14,6 @@ forward. to be accepted by the target model. It is trained with a binary cross-entropy target derived from the agreement between the draft and target (`1 - 0.5 * L1(draft, teacher)`). This signal can be used to prune unlikely branches at serving time. -- **Hidden-states correction (optional).** A residual correction applied to the hidden state - before the `lm_head`, gated by a zero-initialized projection so it is the identity at - initialization and only departs from DFlash as it trains. This is the component shared with the - DFlare-backbone variant of DSpark. ## Loss @@ -31,9 +27,8 @@ loss = + confidence_head_alpha * confidence_loss ## Configuration DSpark is selected by `DSparkConfig` (a superset of `DFlashConfig`). Relevant fields include the -Markov head (`markov_rank`, `markov_head_type`), the confidence head -(`enable_confidence_head`, `confidence_head_with_markov`), and the hidden-states correction -(`enable_hidden_correction`, `hidden_correction_intermediate_size`). Setting `model_arch` to +Markov head (`markov_rank`, `markov_head_type`) and the confidence head +(`enable_confidence_head`, `confidence_head_with_markov`). Setting `model_arch` to `"dflare"` builds DSpark on the DFlare backbone (layer-wise target fusion) instead of DFlash. The loss objective and distillation weights come from the shared `dflash_*` training knobs; see diff --git a/tests/test_dflash.py b/tests/test_dflash.py index 845c810..325480e 100644 --- a/tests/test_dflash.py +++ b/tests/test_dflash.py @@ -14,10 +14,7 @@ import torch -from angelspec.models.dflash import ( - DFlashModel, - _create_dflash_mask_mod, -) +from angelspec.models.dflash import DFlashModel, _create_dflash_mask_mod from angelspec.models.draft.dflash import ( DFlashConfig, DFlashDraftModel, @@ -452,7 +449,7 @@ def test_loss_requires_grad(self): loss.backward() grad_found = False - for name, param in self.model.draft_model.named_parameters(): + for _name, param in self.model.draft_model.named_parameters(): if param.requires_grad and param.grad is not None: if param.grad.abs().sum() > 0: grad_found = True @@ -651,7 +648,7 @@ def test_loss_decreases_over_steps(self): model.train() losses = [] - for step in range(10): + for _step in range(10): optimizer.zero_grad() loss, acc, _, _, _, _ = model( input_ids=input_ids, @@ -1180,7 +1177,7 @@ def test_dflash_yaml_loads(self): try: from angelspec.config.train_config import load_config - except (ImportError, ModuleNotFoundError): + except ImportError: self.skipTest("load_config requires ray (not installed locally)") config_path = os.path.join( diff --git a/tests/test_dfly.py b/tests/test_dfly.py index 47febd0..f03f419 100644 --- a/tests/test_dfly.py +++ b/tests/test_dfly.py @@ -1,19 +1,19 @@ """Tests for the dfly (DFlareV2) draft model. -dfly rides the DSpark path (``DSparkConfig`` + ``model_arch == "dfly"`` → -``DSparkTrainer`` / ``DSparkModel``) so the shared hidden-states correction runs -through the wrapper hook. It is a ``DFlareDraftModel`` whose layers are restored -to DFlash shared-KV layers, with the DFlash FC ``context_proj`` re-added and the -DFlare per-layer fusion applied as a residual. +dfly is a DFlash-family drafter (its own ``DFlyConfig`` / ``"Qwen3DFlyModel"`` → +``DFlyTrainer`` / ``DFlyModel``) so the hidden-states correction runs through the +wrapper hook. It is a ``DFlareDraftModel`` whose layers are restored to DFlash +shared-KV layers, with the DFlash FC ``context_proj`` re-added and the DFlare +per-layer fusion applied as a residual. It has no dependency on DSpark. Covers: -1. auto/dispatch: ``model_arch == "dfly"`` → ``DFlyDraftModel`` (and ``"dflash"`` +1. auto/dispatch: ``DFlyConfig`` → ``DFlyDraftModel`` (and a plain ``DSparkConfig`` still → ``DSparkDraftModel`` — no cross-routing). 2. Structure: DFlash shared-KV layers (no separate ``k_proj_target``), the re-added ``context_proj``, the inherited ``layer_fusion_weights`` / ``context_norm``. 3. hidden_correction: present, zero-init identity, and NO markov / confidence head. 4. ``target_hidden_size != hidden_size`` raises. -5. Tiny forward through the ``DSparkModel`` wrapper: 6-tuple, finite loss, slot-0 +5. Tiny forward through the ``DFlyModel`` wrapper: 6-tuple, finite loss, slot-0 masked, ``loss_components`` keys; and the correction actually runs (perturbing ``down_proj`` changes the loss). 6. state_dict carries the expected backbone/correction keys and none for the @@ -24,15 +24,15 @@ import torch +from angelspec.models.dfly import DFlyModel from angelspec.models.draft.auto import AutoEagle3DraftModel -from angelspec.models.draft.dfly import DFlyDraftModel +from angelspec.models.draft.dfly import DFlyConfig, DFlyDraftModel from angelspec.models.draft.dspark import DSparkConfig, DSparkDraftModel -from angelspec.models.dspark import DSparkModel H, V, BS = 64, 128, 4 -def _config(model_arch="dfly", **kw): +def _base(**kw): base = dict( hidden_size=H, intermediate_size=256, @@ -51,21 +51,24 @@ def _config(model_arch="dfly", **kw): markov_rank=0, enable_confidence_head=False, confidence_head_with_markov=False, - enable_hidden_correction=True, block_size=BS, - model_arch=model_arch, ) base.update(kw) - return DSparkConfig(**base) + return base + + +def _config(**kw): + kw.setdefault("enable_hidden_correction", True) + return DFlyConfig(**_base(**kw)) class TestDflyDispatchAndStructure(unittest.TestCase): def test_auto_dispatch(self): - fly = AutoEagle3DraftModel.from_config(_config(model_arch="dfly")) - ds = AutoEagle3DraftModel.from_config(_config(model_arch="dflash")) + fly = AutoEagle3DraftModel.from_config(_config()) self.assertIsInstance(fly, DFlyDraftModel) + # A plain DSpark config still builds the DSpark drafter (no cross-routing). + ds = AutoEagle3DraftModel.from_config(DSparkConfig(**_base())) self.assertIsInstance(ds, DSparkDraftModel) - # dfly must not be routed to the plain DSpark drafter. self.assertNotIsInstance(ds, DFlyDraftModel) def test_shared_kv_layers(self): @@ -121,13 +124,12 @@ def _build_wrapper(self): torch.manual_seed(0) fly = AutoEagle3DraftModel.from_config(_config()).to(torch.float32) fly.freeze_embedding() - m = DSparkModel( + m = DFlyModel( draft_model=fly, block_size=BS, num_anchors=6, ce_loss_alpha=0.1, l1_loss_alpha=0.9, - confidence_head_alpha=1.0, ) m.eval() return m @@ -151,10 +153,8 @@ def test_forward_six_tuple(self): self.assertTrue(torch.isfinite(loss)) self.assertEqual(lpp.shape[0], BS) self.assertEqual(cpp[0].item(), 0.0) # DFlash convention: slot 0 masked - # confidence_loss key present (=0) even with the head disabled. - self.assertEqual( - set(comps), {"ce_loss", "kl_loss", "lk_loss", "l1_loss", "confidence_loss"} - ) + # DFly rides the plain DFlash loss (no confidence head). + self.assertEqual(set(comps), {"ce_loss", "kl_loss", "lk_loss", "l1_loss"}) def test_correction_actually_runs(self): # With zero-init down_proj the correction is identity; perturbing it must diff --git a/tests/test_mtp.py b/tests/test_mtp.py index 0437cd0..8fa3d66 100644 --- a/tests/test_mtp.py +++ b/tests/test_mtp.py @@ -16,11 +16,7 @@ import torch import torch.nn.functional as F -from angelspec.models.draft.mtp import ( - Hy3MoE, - MTPConfig, - MTPDraftModel, -) +from angelspec.models.draft.mtp import Hy3MoE, MTPConfig, MTPDraftModel from angelspec.models.mtp import MTPModel from angelspec.models.ops.loss import mtp_loss_from_hs @@ -842,7 +838,7 @@ def test_remap_loads_into_draft(self): N = 80 prefix = f"model.layers.{N}." top_level = {"enorm", "hnorm", "eh_proj", "final_layernorm"} - E, H, I = cfg.num_experts, cfg.hidden_size, cfg.moe_intermediate_size + E, H, inter = cfg.num_experts, cfg.hidden_size, cfg.moe_intermediate_size ckpt = {} # non-expert params straight from the model structure for name, p in model.named_parameters(): @@ -858,9 +854,9 @@ def test_remap_loads_into_draft(self): ckpt[prefix + name] = torch.randn_like(p) # per-expert checkpoint weights ([out, in] Linear layout) for e in range(E): - ckpt[f"{prefix}mlp.experts.{e}.gate_proj.weight"] = torch.randn(I, H) - ckpt[f"{prefix}mlp.experts.{e}.up_proj.weight"] = torch.randn(I, H) - ckpt[f"{prefix}mlp.experts.{e}.down_proj.weight"] = torch.randn(H, I) + ckpt[f"{prefix}mlp.experts.{e}.gate_proj.weight"] = torch.randn(inter, H) + ckpt[f"{prefix}mlp.experts.{e}.up_proj.weight"] = torch.randn(inter, H) + ckpt[f"{prefix}mlp.experts.{e}.down_proj.weight"] = torch.randn(H, inter) # Apply the same remap logic the trainer uses (without dist/cuda). import re @@ -899,8 +895,8 @@ def test_remap_loads_into_draft(self): # eh_proj loaded; fused experts loaded with correct [E, in, out] layout. self.assertTrue(torch.equal(model.eh_proj.weight, remapped["eh_proj.weight"])) moe = model.midlayer.mlp - self.assertEqual(tuple(moe.experts_gate_proj.shape), (E, H, I)) - self.assertEqual(tuple(moe.experts_down_proj.shape), (E, I, H)) + self.assertEqual(tuple(moe.experts_gate_proj.shape), (E, H, inter)) + self.assertEqual(tuple(moe.experts_down_proj.shape), (E, inter, H)) # spot-check expert 0 gate slice equals the transposed checkpoint weight self.assertTrue( torch.equal( diff --git a/tests/test_treeflash.py b/tests/test_treeflash.py deleted file mode 100644 index cdffbdb..0000000 --- a/tests/test_treeflash.py +++ /dev/null @@ -1,213 +0,0 @@ -"""Tests for the TreeFlash / DSpark structure overlay. - -Covers the invariants that make the merge safe: -1. HiddenStatesCorrection is zero-initialized → identity at init (degenerates to - DFlash), and non-trivial once the down-proj is perturbed. -2. PositionAdaptiveAlpha: alpha in (0, alpha_max), monotone ramp init, smooth_loss. -3. VanillaMarkov with pos_adaptive=False is byte-identical to the plain bigram - bias (== the pre-merge head), and pos_adaptive=True scales the bias per slot. -4. TreeFlash dispatch + build + tiny forward through the DSparkModel wrapper. -5. Checkpoint-compat guard: the DSpark / TreeFlash state_dict carries the - expected head/correction parameter names. -""" - -import unittest - -import torch - -from angelspec.models.draft.auto import AutoEagle3DraftModel -from angelspec.models.draft.dspark import ( - DSparkConfig, - DSparkDraftModel, - HiddenStatesCorrection, - PositionAdaptiveAlpha, - VanillaMarkov, -) -from angelspec.models.draft.treeflash_dspark_dflare import TreeflashDSparkDFlareDraftModel -from angelspec.models.dspark import DSparkModel - -H, V, BS = 64, 128, 4 - - -def _config(model_arch="dflash", **kw): - base = dict( - hidden_size=H, - intermediate_size=256, - num_hidden_layers=1, - num_attention_heads=4, - num_key_value_heads=2, - vocab_size=V, - rms_norm_eps=1e-6, - max_position_embeddings=512, - rope_theta=10000.0, - num_target_layers=2, - target_hidden_size=H, - target_num_hidden_layers=12, - mask_token_id=V - 1, - markov_rank=16, - enable_confidence_head=True, - confidence_head_with_markov=True, - enable_hidden_correction=True, - block_size=BS, - model_arch=model_arch, - ) - base.update(kw) - return DSparkConfig(**base) - - -class TestHiddenStatesCorrection(unittest.TestCase): - def test_zero_init_is_identity(self): - torch.manual_seed(0) - hc = HiddenStatesCorrection(hidden_size=H, embed_size=H, intermediate_size=H) - self.assertTrue((hc.down_proj.weight == 0).all(), "down_proj must be zero-init") - h = torch.randn(2, 12, H) - e = torch.randn(2, 12, H) - self.assertTrue(torch.allclose(hc(h, e), h), "correction must be identity at init") - - def test_nonzero_after_perturb(self): - torch.manual_seed(0) - hc = HiddenStatesCorrection(hidden_size=H, embed_size=H, intermediate_size=H) - with torch.no_grad(): - hc.down_proj.weight.normal_() - h = torch.randn(2, 12, H) - e = torch.randn(2, 12, H) - self.assertFalse(torch.allclose(hc(h, e), h)) - - def test_pos_adaptive_shape_and_identity(self): - # With zero-init down_proj, identity holds regardless of pos-adaptive alpha. - hc = HiddenStatesCorrection( - hidden_size=H, embed_size=H, intermediate_size=H, pos_adaptive=True, block_size=BS - ) - self.assertIsNotNone(hc.pos_alpha) - h = torch.randn(2, 3 * BS, H) # n_pos multiple of block_size - e = torch.randn(2, 3 * BS, H) - self.assertTrue(torch.allclose(hc(h, e), h)) - - def test_pos_count_not_multiple_raises(self): - hc = HiddenStatesCorrection( - hidden_size=H, embed_size=H, intermediate_size=H, pos_adaptive=True, block_size=BS - ) - with torch.no_grad(): - hc.down_proj.weight.normal_() # make delta nonzero so the reshape runs - with self.assertRaises(ValueError): - hc(torch.randn(2, BS + 1, H), torch.randn(2, BS + 1, H)) - - -class TestPositionAdaptiveAlpha(unittest.TestCase): - def test_alpha_range_and_ramp(self): - pa = PositionAdaptiveAlpha(block_size=BS, alpha_max=0.8, alpha_start=0.1, alpha_end=0.5) - a = pa.alpha() - self.assertEqual(a.shape, (BS,)) - self.assertTrue((a > 0).all() and (a <= 0.8 + 1e-6).all()) - self.assertLess(a[0].item(), a[-1].item()) # monotone ramp at init - - def test_smooth_loss_toggle(self): - self.assertIsNone(PositionAdaptiveAlpha(block_size=BS, smooth_lambda=0.0).smooth_loss()) - reg = PositionAdaptiveAlpha(block_size=BS, smooth_lambda=0.1).smooth_loss() - self.assertIsNotNone(reg) - self.assertGreaterEqual(reg.item(), 0.0) - - def test_requires_positive_block_size(self): - with self.assertRaises(ValueError): - PositionAdaptiveAlpha(block_size=None) - - -class TestVanillaMarkovEquivalence(unittest.TestCase): - def _bias_inputs(self): - torch.manual_seed(0) - base = torch.randn(2, 3, BS, V) - tokens = torch.randint(0, V, (2, 3, BS)) - return base, tokens - - def test_pos_adaptive_off_is_plain_bigram_bias(self): - # pos_adaptive=False must reproduce the plain ``base + bias`` behaviour of - # the pre-merge VanillaMarkov (no per-slot scaling, no extra params). - m = VanillaMarkov(vocab_size=V, markov_rank=16, pos_adaptive=False) - self.assertIsNone(m.pos_alpha) - # no pos_alpha parameters in the state_dict - self.assertFalse(any("pos_alpha" in k for k in m.state_dict())) - base, tokens = self._bias_inputs() - expected = base + m.compute_step_bias(tokens) - self.assertTrue(torch.allclose(m.apply_block_logits(base, token_ids=tokens), expected)) - - def test_pos_adaptive_on_scales_bias(self): - m = VanillaMarkov(vocab_size=V, markov_rank=16, pos_adaptive=True, block_size=BS) - self.assertIsNotNone(m.pos_alpha) - base, tokens = self._bias_inputs() - alpha = m.pos_alpha.alpha() - expected = base + m.compute_step_bias(tokens) * alpha.view(1, 1, -1, 1) - self.assertTrue(torch.allclose(m.apply_block_logits(base, token_ids=tokens), expected)) - - -class TestTreeFlashDispatchAndForward(unittest.TestCase): - def test_auto_dispatch(self): - tf = AutoEagle3DraftModel.from_config(_config(model_arch="dflare")) - ds = AutoEagle3DraftModel.from_config(_config(model_arch="dflash")) - self.assertIsInstance(tf, TreeflashDSparkDFlareDraftModel) - self.assertIsInstance(ds, DSparkDraftModel) - - def test_treeflash_carries_heads(self): - tf = AutoEagle3DraftModel.from_config(_config(model_arch="dflare")) - self.assertIsNotNone(tf.markov_head) - self.assertIsNotNone(tf.hidden_correction) - self.assertIsNotNone(tf.confidence_head) - - def test_treeflash_forward_six_tuple(self): - torch.manual_seed(0) - tf = AutoEagle3DraftModel.from_config(_config(model_arch="dflare")).to(torch.float32) - tf.freeze_embedding() - m = DSparkModel( - draft_model=tf, - block_size=BS, - num_anchors=6, - ce_loss_alpha=0.1, - l1_loss_alpha=0.9, - confidence_head_alpha=1.0, - ) - m.eval() - B, S = 2, 24 - g = torch.Generator().manual_seed(1) - out = m( - input_ids=torch.randint(0, V, (B, S), generator=g), - hidden_states_list=[torch.randn(B, S, H, generator=g) for _ in range(2)], - loss_mask=torch.ones(B, S), - lm_head_weight=torch.randn(V, H, generator=g), - last_hidden_states=torch.randn(B, S, H, generator=g), - ) - self.assertEqual(len(out), 6) - loss, _, lpp, _, cpp, comps = out - self.assertTrue(torch.isfinite(loss)) - self.assertEqual(lpp.shape[0], BS) - self.assertEqual(cpp[0].item(), 0.0) # DFlash convention: slot 0 masked - self.assertEqual( - set(comps), {"ce_loss", "kl_loss", "lk_loss", "l1_loss", "confidence_loss"} - ) - - -class TestCheckpointParamNames(unittest.TestCase): - """Guard the head/correction parameter names so a trained DSpark / - TreeFlash checkpoint keeps loading.""" - - def test_treeflash_state_dict_keys(self): - tf = AutoEagle3DraftModel.from_config(_config(model_arch="dflare")) - keys = set(tf.state_dict()) - for expected in ( - "markov_head.markov_w1.weight", - "markov_head.markov_w2.weight", - "hidden_correction.gate_proj.weight", - "hidden_correction.up_proj.weight", - "hidden_correction.down_proj.weight", - "hidden_correction.hidden_norm.weight", - "hidden_correction.embed_norm.weight", - "confidence_head.proj.weight", - ): - self.assertIn(expected, keys, f"missing checkpoint key: {expected}") - - def test_dspark_no_correction_when_disabled(self): - ds = DSparkDraftModel(_config(model_arch="dflash", enable_hidden_correction=False)) - self.assertIsNone(ds.hidden_correction) - self.assertFalse(any("hidden_correction" in k for k in ds.state_dict())) - - -if __name__ == "__main__": - unittest.main() diff --git a/tools/eval_accept_rate.py b/tools/eval_accept_rate.py index eded91d..bf97456 100755 --- a/tools/eval_accept_rate.py +++ b/tools/eval_accept_rate.py @@ -39,8 +39,8 @@ # Mirror must be set before huggingface_hub is imported anywhere. os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com") -import pandas as pd -from openai import AsyncOpenAI +import pandas as pd # noqa: E402 +from openai import AsyncOpenAI # noqa: E402 HF_ENDPOINT = os.environ.get("HF_ENDPOINT", "https://hf-mirror.com") HF_TOKEN = os.environ.get("HF_TOKEN") or None @@ -51,7 +51,9 @@ # Per-dataset prompt builders (input: one row as a dict) # --------------------------------------------------------------------------- def _arc_challenge(ex: dict) -> str: - choices = "\n".join(f"{l}. {t}" for l, t in zip(ex["choices"]["label"], ex["choices"]["text"])) + choices = "\n".join( + f"{lb}. {t}" for lb, t in zip(ex["choices"]["label"], ex["choices"]["text"]) + ) return ( f"Question: {ex['question']}\n\nChoices:\n{choices}\n\n" "Please select the correct answer and explain your reasoning." @@ -110,13 +112,15 @@ def _mbpp(ex: dict) -> str: def _mmlu(ex: dict) -> str: subject = ex.get("subject", "") lines = [ - f"The following is a multiple choice question about {subject.replace('_', ' ')}." - if subject - else "The following is a multiple choice question.", + ( + f"The following is a multiple choice question about {subject.replace('_', ' ')}." + if subject + else "The following is a multiple choice question." + ), "", f"Question: {ex['question'].strip()}", ] - lines += [f"{l}. {c}" for l, c in zip(_MMLU_LABELS, ex["choices"])] + lines += [f"{lb}. {c}" for lb, c in zip(_MMLU_LABELS, ex["choices"])] lines += [ "", "Please reason step by step, and put your final answer (a single letter A, B, C, or D) " @@ -133,13 +137,15 @@ def _mmlu_pro(ex: dict) -> str: options = ex["options"] labels = _MMLU_PRO_LABELS[: len(options)] lines = [ - f"The following is a multiple choice question about {category}." - if category - else "The following is a multiple choice question.", + ( + f"The following is a multiple choice question about {category}." + if category + else "The following is a multiple choice question." + ), "", f"Question: {ex['question'].strip()}", ] - lines += [f"{l}. {o}" for l, o in zip(labels, options)] + lines += [f"{lb}. {o}" for lb, o in zip(labels, options)] lines += [ "", "Please reason step by step, and put your final answer " diff --git a/tools/generate_data.py b/tools/generate_data.py index 267f7dd..688556a 100644 --- a/tools/generate_data.py +++ b/tools/generate_data.py @@ -6,14 +6,14 @@ 1. Set up one or more SGLang servers for the target model: python3 -m sglang.launch_server \ - --model meta-llama/Llama-3.1-8B-Instruct \ - --mem-fraction-static 0.75 \ - --cuda-graph-max-bs 128 \ - --tp 1 \ - --trust-remote-code \ - --host 0.0.0.0 \ - --port 30000 \ - --dtype bfloat16 + --model meta-llama/Llama-3.1-8B-Instruct \ + --mem-fraction-static 0.75 \ + --cuda-graph-max-bs 128 \ + --tp 1 \ + --trust-remote-code \ + --host 0.0.0.0 \ + --port 30000 \ + --dtype bfloat16 2. Regenerate the dataset: From dd5318746c95a2d37740f82740bdff3a74913be3 Mon Sep 17 00:00:00 2001 From: ali-88123 <1940747290@qq.com> Date: Fri, 31 Jul 2026 10:39:22 +0800 Subject: [PATCH 2/2] delete invalid md --- PR_DESCRIPTION.md | 130 ---------------------------------------------- 1 file changed, 130 deletions(-) delete mode 100644 PR_DESCRIPTION.md diff --git a/PR_DESCRIPTION.md b/PR_DESCRIPTION.md deleted file mode 100644 index d9162d5..0000000 --- a/PR_DESCRIPTION.md +++ /dev/null @@ -1,130 +0,0 @@ -# [Refactor] Promote DFly to a first-class DFlash-family drafter and add an end-to-end multi-step TV loss - -> **Title:** `refactor(drafter): promote DFly to a standalone DFlash-family drafter + add e2e multi-step TV loss` - - -## Summary - -This PR promotes **DFly** from an architecture *variant* -that piggy-backed on the DSpark code path into a standalone, first-class member -of the DFlash drafter family. It gets its own config, model, training wrapper, -and trainer, and no longer depends on DSpark in any way. - -Alongside the refactor, this PR: - -- Removes the dead `treeflash_dspark_dflare` drafter and its test. -- Slims down DSpark by moving the shared TreeFlash hidden-states correction (and - related knobs) out of `dspark.py` and into `dfly.py`, where it now belongs. -- Adds an optional **end-to-end multi-step TV loss** (`γ`-step MTP) to the - DFlash composable loss, replacing the now-unused KL/LK temperature knobs. - -The net effect is a cleaner architecture graph (DFly no longer "rides" DSpark), -less coupling between drafters, and a large reduction in DSpark's surface area -(~+367 / −933 lines overall). - -## Motivation - -Previously, DFly was selected via `DSparkConfig` + `model_arch == "dfly"` and was -dispatched through `DSparkTrainer` / `DSparkModel`. This meant: - -- DFly's behavior was implicit and hard to discover (hidden behind a string flag). -- DSpark carried a lot of machinery (hidden-states correction, position-adaptive - alpha, etc.) that only DFly actually used. -- The `auto` dispatch logic had brittle special-case branches keyed on - `model_arch` string comparison. - -Making DFly a proper config/model/trainer triple removes the cross-routing, -makes the dispatch type-based, and lets each drafter own only what it needs. - -## Changes - -### New — DFly as a first-class drafter - -- **`angelspec/models/dfly.py`** (new): `DFlyModel` training wrapper. Subclasses - `DFlashModel` and overrides `_compute_draft_logits` to apply the optional - TreeFlash hidden-states correction (formula (1)) before the LM head. -- **`angelspec/training/dfly_trainer.py`** (new): `DFlyTrainer`, a thin subclass - of `DFlashTrainer` that only overrides the two model-build seams - (`_build_draft_model` / `_build_training_wrapper`). Reads the `dflash_*` - hyperparameter namespace. -- **`angelspec/models/draft/dfly.py`**: now defines its own `DFlyConfig` - (extends `DFlashConfig`, `model_type = "qwen3"`) and owns the - `HiddenStatesCorrection` module / `build_hidden_correction` helper (moved here - from `dspark.py`). `DFlyDraftModel.config_class` is now `DFlyConfig`. - -### Dispatch / registration - -- **`angelspec/models/draft/auto.py`**: register `DFlyConfig → DFlyDraftModel` - and architecture `"Qwen3DFlyModel" → DFlyConfig`. Removed the - `model_arch == "dfly"` and `model_arch == "dflare"` (TreeFlash) special-case - branches. -- **`angelspec/training/trainer_actor.py`**: add a `DFlyConfig` dispatch branch. - Since both `DSparkConfig` and `DFlyConfig` subclass `DFlashConfig`, they are - checked before the `DFlashConfig` branch. -- **`angelspec/models/__init__.py`** / **`angelspec/models/draft/__init__.py`**: - export `DFlyModel`; drop the `TreeflashDSparkDFlareDraftModel` export. - -### DSpark slim-down - -- **`angelspec/models/draft/dspark.py`**: removed the hidden-states correction, - `PositionAdaptiveAlpha`, position-adaptive Markov knobs, and related - parameters — DSpark now only carries the Markov head and confidence head. -- **`angelspec/models/dspark.py`**: corresponding wrapper cleanup. - -### End-to-end multi-step TV loss - -- **`angelspec/models/dflash.py`**: add `_compute_e2e_tv_loss`, an independent - γ-step MTP TV term added on top of the total loss (not mutually exclusive with - KL/LK), gated on `e2e_tv_loss_weight > 0` and the presence of target - `last_hidden_states`. Emits `e2e_tv_loss` in `loss_components`. - - L_e2e = 1 - (1/γ) * Σ_{j=1..γ} Π_{i=1..j} α_i - -- **`angelspec/config/train_config.py`**: add `dflash_e2e_tv_loss_weight` - (default `0.0`, disabled) and `DatasetConfig.num_proc` (default `64`); remove - the now-unused `dflash_kl_temperature`, `dflash_kl_topk_renormalize`, and - `dflash_lk_temperature`. - -### Removals - -- **`angelspec/models/draft/treeflash_dspark_dflare.py`** (deleted). -- **`tests/test_treeflash.py`** (deleted). -- **`angelspec/config/dflare_dspark_treeflash_qwen3_4b_draft_config.json`** (deleted). - -### Configs - -- **`angelspec/config/dfly_*_draft_config.json`**: switch from - `architectures: ["DSparkDraftModel"]` / `model_type: "dspark"` / - `model_arch: "dfly"` to `architectures: ["Qwen3DFlyModel"]` / - `model_type: "qwen3"`; drop the DSpark-only `markov_rank` / - `enable_confidence_head` / `confidence_head_with_markov` keys. -- **`configs/vllm_qwen3_8b_dfly.yaml`**, **`configs/vllm_hy3_dfly.yaml`**: fix - `lm_head_key` to `lm_head.weight`, migrate `dspark_*` hyperparameters to the - `dflash_*` namespace, and set up the two-stage loss schedule (cold-start - lk-loss → final `e2e_tv_loss`). - -### Docs & tests - -- **`docs/concepts/dfly.md`**: document DFly as an independent `DFlyConfig` / - `Qwen3DFlyModel` / `DFlyTrainer` drafter; update the comparison table. -- **`docs/concepts/dspark.md`**, **`dflash.md`**, **`draft_model_family.md`**: - minor updates reflecting the moved correction module and the new loss term. -- **`tests/test_dfly.py`**: updated to build via `DFlyConfig` and exercise the - `DFlyModel` wrapper; asserts a plain `DSparkConfig` still routes to - `DSparkDraftModel` (no cross-routing). - -## Compatibility / migration notes - -- **Breaking config change:** existing DFly checkpoints/configs using - `architectures: ["DSparkDraftModel"]` + `model_arch: "dfly"` must be migrated - to `architectures: ["Qwen3DFlyModel"]` (see updated `dfly_*_draft_config.json`). -- The removed `dflash_kl_temperature` / `dflash_kl_topk_renormalize` / - `dflash_lk_temperature` training args are no longer accepted. -- `treeflash_dspark_dflare` is gone; any references must be removed. - -## Testing - -- `tests/test_dfly.py` covers auto-dispatch, model structure (shared-KV layers, - re-added `context_proj`, inherited fusion), the zero-init identity of the - hidden correction, a tiny forward through `DFlyModel` (finite loss, correct - `loss_components`, correction actually affects the loss), and state-dict keys.