diff --git a/.gitignore b/.gitignore index c925588..b03a356 100644 --- a/.gitignore +++ b/.gitignore @@ -89,3 +89,5 @@ wandb/log.txt .claude/ wandb/ +data/ +profiles/ diff --git a/tests/test_dflash.py b/tests/test_dflash.py index e9b673a..a1a21ea 100644 --- a/tests/test_dflash.py +++ b/tests/test_dflash.py @@ -427,7 +427,7 @@ def test_forward_produces_loss_and_acc(self): lm_head_weight = torch.randn(self.V, self.H) with torch.no_grad(): - loss, acc, loss_pp, acc_pp, count_pp = self.model( + loss, acc, loss_pp, acc_pp, count_pp, loss_components = self.model( input_ids=input_ids, hidden_states_list=hidden_states_list, loss_mask=loss_mask, @@ -441,6 +441,7 @@ def test_forward_produces_loss_and_acc(self): self.assertEqual(loss_pp.shape, (self.model.block_size,)) self.assertEqual(acc_pp.shape, (self.model.block_size,)) self.assertEqual(count_pp.shape, (self.model.block_size,)) + self.assertEqual(loss_components, {}) def test_loss_requires_grad(self): """Loss should be differentiable through the draft model.""" @@ -453,7 +454,7 @@ def test_loss_requires_grad(self): lm_head_weight = torch.randn(self.V, self.H) self.model.train() - loss, acc, _, _, _ = self.model( + loss, acc, *_ = self.model( input_ids=input_ids, hidden_states_list=hidden_states_list, loss_mask=loss_mask, @@ -721,7 +722,7 @@ def test_loss_decreases_over_steps(self): losses = [] for step in range(10): optimizer.zero_grad() - loss, acc, _, _, _ = model( + loss, acc, *_ = model( input_ids=input_ids, hidden_states_list=hidden_states_list, loss_mask=loss_mask, @@ -1061,7 +1062,7 @@ def _train_steps( losses, accs = [], [] for _ in range(steps): optimizer.zero_grad() - loss, acc, _, _, _ = model( + loss, acc, *_ = model( input_ids=input_ids, hidden_states_list=hidden_states_list, loss_mask=loss_mask, @@ -1200,7 +1201,7 @@ def test_dflash_uses_ce_loss_not_kl(self): lm_head_weight = torch.randn(V, H) with torch.no_grad(): - loss, acc, _, _, _ = model( + loss, acc, *_ = model( input_ids=input_ids, hidden_states_list=hs_list, loss_mask=loss_mask, @@ -1350,5 +1351,177 @@ def test_min_loss_tokens_validation_passes(self): self.assertGreaterEqual(min_loss, 2 * block_size) +class TestExtraLossComponentAggregation(unittest.TestCase): + class _DummyOptimizer: + def get_learning_rate(self): + return 1e-3 + + def _make_dspark_trainer(self): + from torchspec.training.dspark_trainer import DSparkTrainer + + trainer = object.__new__(DSparkTrainer) + trainer.loss_decay_gamma = 1.0 + trainer.global_step = 0 + trainer.optimizer = self._DummyOptimizer() + return trainer + + @staticmethod + def _steps(loss_key, acc_key, count_key): + base = { + loss_key: torch.tensor([0.0, 1.0, 3.0]), + acc_key: torch.tensor([0.0, 0.5, 1.0]), + count_key: torch.tensor([0.0, 4.0, 4.0]), + } + return [ + { + **base, + "ce_loss": torch.tensor(2.0), + "l1_loss": torch.tensor(0.4), + "confidence_loss": torch.tensor(0.1), + }, + { + **base, + "ce_loss": torch.tensor(4.0), + "l1_loss": torch.tensor(0.6), + "confidence_loss": torch.tensor(0.3), + }, + ] + + def test_dspark_keys_declared(self): + from torchspec.training.dspark_trainer import DSparkTrainer + + self.assertEqual( + DSparkTrainer._extra_loss_component_keys, + ["ce_loss", "l1_loss", "confidence_loss"], + ) + + def test_dflash_declares_no_components(self): + from torchspec.training.dflash_trainer import DFlashTrainer + + self.assertEqual(DFlashTrainer._extra_loss_component_keys, []) + trainer = object.__new__(DFlashTrainer) + # No keys to reduce → no-op even if components are present in the steps. + out = trainer._reduce_loss_components([{"ce_loss": torch.tensor(1.0)}], "train/") + self.assertEqual(out, {}) + + @mock.patch("torchspec.training.dflash_trainer.dist.get_rank", return_value=0) + @mock.patch( + "torchspec.training.dflash_trainer.dist.all_reduce", + side_effect=lambda tensor, op=None: None, + ) + def test_components_in_train_metrics(self, _mock_all_reduce, _mock_get_rank): + trainer = self._make_dspark_trainer() + metrics = trainer._aggregate_metrics( + self._steps("loss_per_position", "acc_per_position", "count_per_position"), + step=1, + grad_norm=torch.tensor(1.0), + ) + self.assertAlmostEqual(metrics["train/ce_loss"], 3.0, places=6) + self.assertAlmostEqual(metrics["train/l1_loss"], 0.5, places=6) + self.assertAlmostEqual(metrics["train/confidence_loss"], 0.2, places=6) + + @mock.patch("torchspec.training.dflash_trainer.dist.get_rank", return_value=0) + @mock.patch( + "torchspec.training.dflash_trainer.dist.all_reduce", + side_effect=lambda tensor, op=None: None, + ) + def test_components_in_eval_metrics(self, _mock_all_reduce, _mock_get_rank): + trainer = self._make_dspark_trainer() + metrics = trainer._aggregate_eval_metrics(self._steps("loss_pp", "acc_pp", "count_pp")) + self.assertAlmostEqual(metrics["eval/ce_loss"], 3.0, places=6) + self.assertAlmostEqual(metrics["eval/l1_loss"], 0.5, places=6) + self.assertAlmostEqual(metrics["eval/confidence_loss"], 0.2, places=6) + + +class TestDFlashL1Loss(unittest.TestCase): + """Opt-in L1 distribution-distillation term (ce_loss_alpha / l1_loss_alpha).""" + + H = 64 + V = 128 + num_target_layers = 2 + + def _model(self, ce_alpha=1.0, l1_alpha=0.0): + config = _make_config( + H=self.H, + intermediate=256, + num_heads=4, + num_kv_heads=2, + V=self.V, + num_target_layers=self.num_target_layers, + target_num_hidden=12, + ) + draft = DFlashDraftModel(config).to(dtype=torch.float32) + draft.freeze_embedding() + return DFlashModel( + draft_model=draft, + block_size=4, + num_anchors=4, + loss_objective="decay", + loss_decay_gamma=7.0, + ce_loss_alpha=ce_alpha, + l1_loss_alpha=l1_alpha, + ) + + def _data(self, B=2, seq_len=64, with_last_hs=True): + input_ids = torch.randint(0, self.V, (B, seq_len)) + hs_list = [torch.randn(B, seq_len, self.H) for _ in range(self.num_target_layers)] + loss_mask = torch.ones(B, seq_len) + lm_head_weight = torch.randn(self.V, self.H) + last_hs = torch.randn(B, seq_len, self.H) if with_last_hs else None + return dict( + input_ids=input_ids, + hidden_states_list=hs_list, + loss_mask=loss_mask, + lm_head_weight=lm_head_weight, + last_hidden_states=last_hs, + ) + + def test_l1_requires_last_hidden_states(self): + torch.manual_seed(0) + model = self._model(l1_alpha=0.5) + kwargs = self._data(with_last_hs=False) + with self.assertRaisesRegex(ValueError, "last_hidden_states"): + model(**kwargs) + + def test_default_is_pure_ce(self): + """ce_alpha=1, l1_alpha=0 (defaults) reproduces the plain-CE loss, and + needs no last_hidden_states.""" + kwargs = self._data(with_last_hs=False) + + def run(ce_alpha, l1_alpha): + torch.manual_seed(11) + with torch.no_grad(): + return self._model(ce_alpha=ce_alpha, l1_alpha=l1_alpha)(**kwargs)[0].item() + + # explicit (1.0, 0.0) == implicit default + self.assertAlmostEqual(run(1.0, 0.0), run(1.0, 0.0), places=6) + + def test_l1_forward_finite_and_grad(self): + torch.manual_seed(0) + model = self._model(ce_alpha=0.1, l1_alpha=0.9) + kwargs = self._data() + model.train() + loss, acc, lpp, *_ = model(**kwargs) + self.assertTrue(torch.isfinite(loss)) + self.assertTrue(torch.all(lpp >= 0)) # per-position metric stays CE + loss.backward() + grad_found = any( + p.requires_grad and p.grad is not None and p.grad.abs().sum() > 0 + for p in model.draft_model.parameters() + ) + self.assertTrue(grad_found, "No gradient flowed to draft model under L1 loss") + + def test_l1_changes_loss_vs_ce_only(self): + """With l1_alpha>0 the loss differs from the ce-only loss (term is active).""" + kwargs = self._data() + + def run(ce_alpha, l1_alpha): + torch.manual_seed(7) # fix anchor sampling so losses are comparable + with torch.no_grad(): + return self._model(ce_alpha=ce_alpha, l1_alpha=l1_alpha)(**kwargs)[0].item() + + self.assertNotAlmostEqual(run(1.0, 0.0), run(1.0, 0.9), places=4) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_dspark.py b/tests/test_dspark.py index 119f1cc..87e3af5 100644 --- a/tests/test_dspark.py +++ b/tests/test_dspark.py @@ -118,7 +118,7 @@ class TestDSparkConfig(unittest.TestCase): def test_subclasses_dflash_and_attrs(self): cfg = _make_dspark_config(markov_rank=32) self.assertIsInstance(cfg, DFlashConfig) # ordering hazard: check DSpark first - self.assertEqual(cfg.model_type, "dspark") + self.assertEqual(cfg.model_type, "qwen3_dspark") self.assertEqual(cfg.markov_rank, 32) self.assertTrue(cfg.enable_confidence_head) @@ -235,8 +235,8 @@ class TestDispatch(unittest.TestCase): def test_json_resolves_to_dspark_config(self): cfg = AutoDraftModelConfig.from_dict( { - "architectures": ["DSparkDraftModel"], - "model_type": "dspark", + "architectures": ["Qwen3DSparkModel"], + "model_type": "qwen3_dspark", "hidden_size": 64, "vocab_size": 128, "num_hidden_layers": 1, diff --git a/tools/convert_to_hf.py b/tools/convert_to_hf.py index b41e84d..02d9bc5 100644 --- a/tools/convert_to_hf.py +++ b/tools/convert_to_hf.py @@ -126,9 +126,13 @@ def set_up_planner( # ── Export fixups ──────────────────────────────────────────────────────────── -# vLLM checkpoints use `layers.0.xxx` for the single decoder layer, -# while our training code uses `midlayer.xxx`. -_WEIGHT_KEY_REMAP = [("midlayer.", "layers.0.")] +# vLLM / SGLang supports a different naming convention for some weight keys. +_WEIGHT_KEY_REMAP = [ + ("midlayer.", "layers.0."), + ("context_proj.", "fc."), + ("context_norm.", "hidden_norm."), + ("final_norm.", "norm."), +] def _remap_weight_keys(tensors: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: @@ -143,6 +147,11 @@ def _remap_weight_keys(tensors: dict[str, torch.Tensor]) -> dict[str, torch.Tens return remapped +_MODEL_TYPE_REMAP = { + "qwen3_dspark": "qwen3", +} + + def _fixup_export_config(raw_config: dict, export_for_vllm: bool = False) -> dict: """Apply model-type-specific fixups to the exported config.""" config = json.loads(json.dumps(raw_config)) @@ -156,6 +165,9 @@ def _fixup_export_config(raw_config: dict, export_for_vllm: bool = False) -> dic if eagle_cfg and key in eagle_cfg: eagle_cfg[key] = [x + 1 for x in eagle_cfg[key]] + if config["model_type"] in _MODEL_TYPE_REMAP: + config["model_type"] = _MODEL_TYPE_REMAP[config["model_type"]] + return config @@ -247,20 +259,17 @@ def _extract_model_weights( return model_state -def _prepare_export_tensors(hf_model, export_for_vllm: bool) -> dict[str, torch.Tensor]: +def _prepare_export_tensors(hf_model) -> dict[str, torch.Tensor]: tensors = hf_model.state_dict() - if export_for_vllm: - logger.info("Exporting vLLM-compatible checkpoint keys") - return _remap_weight_keys(tensors) logger.info("Exporting native checkpoint keys") - return dict(tensors) + return _remap_weight_keys(tensors) def _save_without_vocab_pruning( hf_model, output_dir: str, raw_config: dict, vocab_size: int, export_for_vllm: bool = False ) -> None: version = _get_torchspec_version() - tensors = _prepare_export_tensors(hf_model, export_for_vllm) + tensors = _prepare_export_tensors(hf_model) save_file( tensors, os.path.join(output_dir, "model.safetensors"), diff --git a/torchspec/config/dspark_draft_config.json b/torchspec/config/dspark_draft_config.json index c091754..047d987 100644 --- a/torchspec/config/dspark_draft_config.json +++ b/torchspec/config/dspark_draft_config.json @@ -1,6 +1,6 @@ { - "architectures": ["DSparkDraftModel"], - "model_type": "dspark", + "architectures": ["Qwen3DSparkModel"], + "model_type": "qwen3_dspark", "hidden_size": 4096, "intermediate_size": 12288, "num_hidden_layers": 5, diff --git a/torchspec/config/train_config.py b/torchspec/config/train_config.py index 990554c..ba50bfb 100644 --- a/torchspec/config/train_config.py +++ b/torchspec/config/train_config.py @@ -151,6 +151,8 @@ class TrainingConfig: dflash_dpace_alpha: float = 0.5 dflash_loss_decay_gamma: float = 7.0 dflash_loss_objective: str = "decay" # "decay" or "dpace" + dflash_ce_loss_alpha: float = 1.0 + dflash_l1_loss_alpha: float = 0.0 dflash_num_anchors: int = 512 dflash_num_target_layers: int = 5 diff --git a/torchspec/models/dflash.py b/torchspec/models/dflash.py index da9bad3..b3d8c60 100644 --- a/torchspec/models/dflash.py +++ b/torchspec/models/dflash.py @@ -26,7 +26,7 @@ Matches SpecForge's OnlineDFlashModel (specforge/core/dflash.py). """ -from typing import List, Tuple +from typing import List, Optional, Tuple import torch import torch.nn as nn @@ -109,6 +109,8 @@ def __init__( loss_objective: str = "decay", dpace_alpha: float = 0.5, loss_decay_gamma: float = 7.0, + ce_loss_alpha: float = 1.0, + l1_loss_alpha: float = 0.0, ): super().__init__() loss_objective = loss_objective.lower() @@ -126,6 +128,8 @@ def __init__( self.loss_objective = loss_objective self.dpace_alpha = dpace_alpha self.loss_decay_gamma = loss_decay_gamma + self.ce_loss_alpha = float(ce_loss_alpha) + self.l1_loss_alpha = float(l1_loss_alpha) def _sample_anchor_positions( self, @@ -239,26 +243,25 @@ def _create_noise_embed( return self.draft_model.embed_tokens(noise_ids) - def forward( + def _draft_backbone( self, input_ids: torch.Tensor, hidden_states_list: List[torch.Tensor], loss_mask: torch.Tensor, - lm_head_weight: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Full DFlash training forward pass. + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + """ + Shared DFlash backbone (context features → anchor sampling → noise + embedding → position ids → block-causal mask → draft model forward). - Matches SpecForge's OnlineDFlashModel.forward(). + Both ``DFlashModel.forward`` and ``DSparkModel.forward`` build the draft + hidden states this exact way; only the label/loss tail differs. Keeping + the attention/mask/anchor wiring here gives it a single source of truth. Returns: - loss: scalar training loss (objective-weighted) - accuracy: scalar accuracy (binary mask, no decay) - loss_per_position: [block_size] mean loss at each within-block position - (index 0 is the anchor slot and always 0; indices 1..B-1 are the - predicted tokens at 1..B-1 steps past the anchor) - acc_per_position: [block_size] mean accuracy at each within-block position - count_per_position: [block_size] valid label count at each within-block - position before loss decay is applied + draft_hidden: [B, n_blocks*block_size, D] pre-loss draft hidden states + anchor_positions: [B, n_blocks] sampled anchor positions + block_keep_mask: [B, n_blocks] bool validity of each anchor slot + n_blocks: number of anchor slots (== num_anchors) """ bsz, seq_len = input_ids.shape device = input_ids.device @@ -311,6 +314,39 @@ def forward( noise_embedding=noise_embedding, ) + return draft_hidden, anchor_positions, block_keep_mask, n_blocks + + def forward( + self, + input_ids: torch.Tensor, + hidden_states_list: List[torch.Tensor], + loss_mask: torch.Tensor, + lm_head_weight: torch.Tensor, + last_hidden_states: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict]: + """ + Full DFlash training forward pass. + + Returns: + loss: scalar training loss (objective-weighted) + accuracy: scalar accuracy (binary mask, no decay) + loss_per_position: [block_size] mean loss at each within-block position + (index 0 is the anchor slot and always 0; indices 1..B-1 are the + predicted tokens at 1..B-1 steps past the anchor) + acc_per_position: [block_size] mean accuracy at each within-block position + count_per_position: [block_size] valid label count at each within-block + position before loss decay is applied + loss_components: dict of extra per-component loss scalars for logging + (empty for the base DFlash objective; populated by subclasses). + """ + bsz, seq_len = input_ids.shape + device = input_ids.device + + # 1-6. Shared backbone → draft hidden states + anchor bookkeeping. + draft_hidden, anchor_positions, block_keep_mask, n_blocks = self._draft_backbone( + input_ids, hidden_states_list, loss_mask + ) + # 7. Compute logits via frozen LM head logits = ( self.draft_model.lm_head(draft_hidden) @@ -352,11 +388,31 @@ def forward( # our objective weighting is an addition to the training signal, not the metric. binary_eval_mask = weight_mask.view(-1) - # 9. Cross entropy loss - flat_logits = logits.view(-1, logits.size(-1)) + # 9. Per-token loss: ce_loss_alpha*CE + l1_loss_alpha*L1. + vocab_size = logits.size(-1) + flat_logits = logits.view(-1, vocab_size) flat_targets = target_ids.view(-1) - loss_per_token = F.cross_entropy(flat_logits, flat_targets, reduction="none") - loss_per_token_by_position = loss_per_token.view(bsz, n_blocks, self.block_size) + ce_per_token = F.cross_entropy(flat_logits, flat_targets, reduction="none") + + loss_per_token = self.ce_loss_alpha * ce_per_token + if self.l1_loss_alpha > 0: + if last_hidden_states is None: + raise ValueError( + "DFlash L1 distillation (l1_loss_alpha > 0) requires target " + "last_hidden_states; set inference.store_last_hidden_states=true in the " + "run config." + ) + tgt_idx = (safe_label_indices - 1).clamp(min=0) # [B, n_blocks, block_size] + hdim = last_hidden_states.size(-1) + gather_idx = tgt_idx.reshape(bsz, -1, 1).expand(-1, -1, hdim) + aligned_hidden = torch.gather(last_hidden_states, 1, gather_idx) + target_logits = F.linear(aligned_hidden, lm_head_weight).view(-1, vocab_size) + target_probs = torch.softmax(target_logits.float(), dim=-1) + draft_probs = torch.softmax(flat_logits.float(), dim=-1) + l1_per_token = (draft_probs - target_probs).abs().sum(dim=-1) + loss_per_token = loss_per_token + self.l1_loss_alpha * l1_per_token + + loss_per_token_by_position = ce_per_token.view(bsz, n_blocks, self.block_size) objective_weights = weight_mask if ( @@ -398,11 +454,19 @@ def forward( count_per_position = binary_weights.sum(dim=(0, 1)) count_per_pos = count_per_position.clamp(min=1.0) - loss_per_position = ( - loss_per_token.view(bsz, n_blocks, self.block_size) * binary_weights - ).sum(dim=(0, 1)) / count_per_pos + loss_per_position = (loss_per_token_by_position * binary_weights).sum( + dim=(0, 1) + ) / count_per_pos acc_per_position = (correct.view(bsz, n_blocks, self.block_size).float()).sum( dim=(0, 1) ) / count_per_pos - return loss, accuracy, loss_per_position, acc_per_position, count_per_position + loss_components = {} + return ( + loss, + accuracy, + loss_per_position, + acc_per_position, + count_per_position, + loss_components, + ) diff --git a/torchspec/models/draft/auto.py b/torchspec/models/draft/auto.py index b0dbe05..550b683 100644 --- a/torchspec/models/draft/auto.py +++ b/torchspec/models/draft/auto.py @@ -79,7 +79,7 @@ class AutoDraftModelConfig: "LlamaForCausalLMEagle3": LlamaConfig, "Eagle3DeepseekV2ForCausalLM": DeepseekV3Config, "DFlashDraftModel": DFlashConfig, - "DSparkDraftModel": DSparkConfig, + "Qwen3DSparkModel": DSparkConfig, } @classmethod diff --git a/torchspec/models/draft/dspark.py b/torchspec/models/draft/dspark.py index 429a667..1cf85aa 100644 --- a/torchspec/models/draft/dspark.py +++ b/torchspec/models/draft/dspark.py @@ -48,7 +48,7 @@ class DSparkConfig(DFlashConfig): Configuration for the DSpark draft model. Extends :class:`DFlashConfig`. """ - model_type = "dspark" + model_type = "qwen3_dspark" def __init__( self, @@ -79,6 +79,7 @@ def __init__(self, *, vocab_size: int, markov_rank: int): 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. self.markov_w2 = nn.Linear(self.markov_rank, self.vocab_size, bias=False) def get_prev_embeddings(self, token_ids: torch.Tensor) -> torch.Tensor: diff --git a/torchspec/models/dspark.py b/torchspec/models/dspark.py index 4fd03a9..cffbf16 100644 --- a/torchspec/models/dspark.py +++ b/torchspec/models/dspark.py @@ -45,8 +45,7 @@ import torch.distributed as dist import torch.nn.functional as F -from torchspec.models.dflash import DFlashModel, _create_dflash_mask_mod -from torchspec.models.ops.flex_attention import compile_friendly_create_block_mask +from torchspec.models.dflash import DFlashModel class DSparkModel(DFlashModel): @@ -108,43 +107,9 @@ def forward( bsz, seq_len = input_ids.shape device = input_ids.device - # ---- DFlash backbone (identical to DFlashModel.forward steps 1-7) ---- - context_feature = self.draft_model.extract_context_feature(hidden_states_list) - anchor_positions, block_keep_mask = self._sample_anchor_positions( - seq_len, loss_mask, device - ) - n_blocks = anchor_positions.shape[1] - noise_embedding = self._create_noise_embed(input_ids, anchor_positions, block_keep_mask) - context_position_ids, draft_position_ids = self._create_position_ids( - anchor_positions, seq_len - ) - - draft_len = n_blocks * self.block_size - kv_len = seq_len + draft_len - block_mask = None - if device.type == "cuda": - mask_mod = _create_dflash_mask_mod( - anchor_positions=anchor_positions, - block_keep_mask=block_keep_mask, - ctx_len=seq_len, - block_size=self.block_size, - ) - block_mask = compile_friendly_create_block_mask( - mask_mod=mask_mod, - B=bsz, - H=None, - Q_LEN=draft_len, - KV_LEN=kv_len, - device=device, - ) - - draft_hidden = self.draft_model( - draft_input_ids=None, - context_feature=context_feature, - draft_position_ids=draft_position_ids, - context_position_ids=context_position_ids, - block_mask=block_mask, - noise_embedding=noise_embedding, + # ---- Shared DFlash backbone (steps 1-6) → draft hidden states ---- + draft_hidden, anchor_positions, block_keep_mask, n_blocks = self._draft_backbone( + input_ids, hidden_states_list, loss_mask ) hidden_4d = draft_hidden.view(bsz, n_blocks, self.block_size, -1) diff --git a/torchspec/training/dflash_trainer.py b/torchspec/training/dflash_trainer.py index 195a819..2fcebe0 100644 --- a/torchspec/training/dflash_trainer.py +++ b/torchspec/training/dflash_trainer.py @@ -50,6 +50,7 @@ class DFlashTrainer(Trainer): _draft_config_class = DFlashConfig _anchor_slot_offset = 1 + _extra_loss_component_keys: list[str] = [] def _build_draft_model(self, config): """Instantiate the draft network. Overridden by subclasses.""" @@ -64,6 +65,8 @@ def _build_training_wrapper(self, draft_model): loss_objective=self.loss_objective, dpace_alpha=self.dpace_alpha, loss_decay_gamma=self.loss_decay_gamma, + ce_loss_alpha=self.ce_loss_alpha, + l1_loss_alpha=self.l1_loss_alpha, ) def __init__(self, args: Namespace): @@ -75,6 +78,8 @@ def __init__(self, args: Namespace): self.loss_objective = getattr(args, "dflash_loss_objective", "decay") self.dpace_alpha = getattr(args, "dflash_dpace_alpha", 0.5) self.loss_decay_gamma = getattr(args, "dflash_loss_decay_gamma", 7.0) + self.ce_loss_alpha = getattr(args, "dflash_ce_loss_alpha", 1.0) + self.l1_loss_alpha = getattr(args, "dflash_l1_loss_alpha", 0.0) def init_model( self, @@ -273,7 +278,7 @@ def _split_hidden_states(self, hidden_states: torch.Tensor) -> List[torch.Tensor def _forward( self, batch: dict - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, dict]: device = torch.device("cuda") input_ids = batch["input_ids"].to(device, non_blocking=True) hidden_states = batch["hidden_states"].to(device, non_blocking=True) @@ -283,17 +288,31 @@ def _forward( loss_mask = loss_mask.squeeze(-1) loss_mask = loss_mask.to(device, non_blocking=True) + last_hidden_states = batch.get("last_hidden_states", None) + if last_hidden_states is not None: + last_hidden_states = last_hidden_states.to(device, non_blocking=True) + hidden_states_list = self._split_hidden_states(hidden_states) del hidden_states - loss, accuracy, loss_per_position, acc_per_position, count_per_position = self.model( - input_ids=input_ids, - hidden_states_list=hidden_states_list, - loss_mask=loss_mask, - lm_head_weight=self.target_lm_head_weight, + loss, accuracy, loss_per_position, acc_per_position, count_per_position, loss_components = ( + self.model( + input_ids=input_ids, + hidden_states_list=hidden_states_list, + loss_mask=loss_mask, + lm_head_weight=self.target_lm_head_weight, + last_hidden_states=last_hidden_states, + ) ) - return loss, accuracy, loss_per_position, acc_per_position, count_per_position + return ( + loss, + accuracy, + loss_per_position, + acc_per_position, + count_per_position, + loss_components, + ) def _backward(self, loss: torch.Tensor, accumulation_steps: int = 1) -> torch.Tensor: scaled_loss = loss / accumulation_steps @@ -307,11 +326,14 @@ def _backward(self, loss: torch.Tensor, accumulation_steps: int = 1) -> torch.Te def eval_forward(self, batch: dict) -> dict: """Single forward pass without backward — returns per-position tensors.""" with torch.no_grad(): - _, _, loss_per_position, acc_per_position, count_per_position = self._forward(batch) + _, _, loss_per_position, acc_per_position, count_per_position, loss_components = ( + self._forward(batch) + ) return { "loss_pp": loss_per_position.detach(), "acc_pp": acc_per_position.detach(), "count_pp": count_per_position.detach(), + **loss_components, } def _reduce_position_metrics( @@ -358,6 +380,22 @@ def _compute_scalar_metrics( avg_loss = ((pred_loss_pp * weighted_counts).sum() / safe_weighted_count).item() return avg_loss, avg_acc + def _reduce_loss_components(self, all_step_metrics: list[dict], prefix: str) -> dict: + """ + Reduce extra scalar loss components into ``{prefix}{key}`` global means. + """ + out: dict = {} + for key in self._extra_loss_component_keys: + vals = [m[key] for m in all_step_metrics if key in m] + if not vals: + continue + value = torch.stack([v.float() for v in vals]).mean() + if dist.is_initialized() and dist.get_world_size() > 1: + dist.all_reduce(value, op=dist.ReduceOp.SUM) + value = value / dist.get_world_size() + out[f"{prefix}{key}"] = value.item() + return out + def eval_from_cache(self) -> dict: """Run forward-only eval over all CPU-cached eval samples. @@ -419,6 +457,8 @@ def _aggregate_eval_metrics(self, all_step_metrics: list[dict]) -> dict: metrics[f"eval/ploss_{i}"] = pred_loss_pp[i].item() metrics[f"eval/acc_{i}"] = pred_acc_pp[i].item() + metrics.update(self._reduce_loss_components(all_step_metrics, "eval/")) + if dist.get_rank() == 0: logger.info( f"eval: loss={weighted_avg_loss:.4f}, acc={avg_acc:.4f}, " @@ -445,8 +485,8 @@ def _train_step( evt_bwd_e = torch.cuda.Event(enable_timing=True) evt_fwd_s.record() - loss, accuracy, loss_per_position, acc_per_position, count_per_position = self._forward( - batch + loss, accuracy, loss_per_position, acc_per_position, count_per_position, loss_components = ( + self._forward(batch) ) evt_fwd_e.record() @@ -463,6 +503,7 @@ def _train_step( "total_loss": total_loss.detach(), "_fwd_events": (evt_fwd_s, evt_fwd_e), "_bwd_events": (evt_bwd_s, evt_bwd_e), + **loss_components, } def _aggregate_metrics( @@ -509,6 +550,8 @@ def _aggregate_metrics( metrics[f"train/ploss_{i}"] = pred_loss_pp[i].item() metrics[f"train/acc_{i}"] = pred_acc_pp[i].item() + metrics.update(self._reduce_loss_components(all_step_metrics, "train/")) + # Sub-timing breakdown (forward vs backward) fwd_ms = sum( m["_fwd_events"][0].elapsed_time(m["_fwd_events"][1]) diff --git a/torchspec/training/dspark_trainer.py b/torchspec/training/dspark_trainer.py index d24cc87..766c653 100644 --- a/torchspec/training/dspark_trainer.py +++ b/torchspec/training/dspark_trainer.py @@ -18,19 +18,14 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. -"""DSpark trainer — DFlash trainer + Markov/confidence heads and L1 distillation. +""" +DSpark trainer — DFlash trainer + Markov/confidence heads and L1 distillation. Reuses the entire DFlash training pipeline (FSDP init, optimizer, checkpoint, -metric aggregation, hidden-state capture/transfer) via subclass hooks, and -additionally feeds the target ``last_hidden_states`` into the forward so the -L1 distribution-distillation and confidence-head losses can be computed. +forward, train step, and metric aggregation) via subclass hooks. """ from argparse import Namespace -from typing import Tuple - -import torch -import torch.distributed as dist from torchspec.models.draft.dspark import DSparkConfig, DSparkDraftModel from torchspec.models.dspark import DSparkModel @@ -41,6 +36,7 @@ class DSparkTrainer(DFlashTrainer): """DSpark-specific trainer (DFlash backbone + EAGLE-style heads).""" _draft_config_class = DSparkConfig + _extra_loss_component_keys = ["ce_loss", "l1_loss", "confidence_loss"] def __init__(self, args: Namespace): super().__init__(args) @@ -72,79 +68,3 @@ def _build_training_wrapper(self, draft_model): l1_loss_alpha=self.l1_loss_alpha, confidence_head_alpha=self.confidence_head_alpha, ) - - # ------------------------------------------------------------------ - # Forward — adds target last_hidden_states for L1 / confidence losses - # ------------------------------------------------------------------ - - def _forward( - self, batch: dict - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - device = torch.device("cuda") - input_ids = batch["input_ids"].to(device, non_blocking=True) - hidden_states = batch["hidden_states"].to(device, non_blocking=True) - - loss_mask = batch["loss_mask"] - if loss_mask.dim() == 3: - loss_mask = loss_mask.squeeze(-1) - loss_mask = loss_mask.to(device, non_blocking=True) - - last_hidden_states = batch.get("last_hidden_states", None) - if last_hidden_states is not None: - last_hidden_states = last_hidden_states.to(device, non_blocking=True) - - hidden_states_list = self._split_hidden_states(hidden_states) - del hidden_states - - # DSparkModel.forward returns a 6th element: a dict of per-component loss - # scalars (ce/l1/confidence). Stash it for _train_step to log; return the - # 5-tuple the base trainer expects. - ( - loss, - accuracy, - loss_per_position, - acc_per_position, - count_per_position, - self._last_loss_components, - ) = self.model( - input_ids=input_ids, - hidden_states_list=hidden_states_list, - loss_mask=loss_mask, - lm_head_weight=self.target_lm_head_weight, - last_hidden_states=last_hidden_states, - ) - return loss, accuracy, loss_per_position, acc_per_position, count_per_position - - # ------------------------------------------------------------------ - # Per-component loss logging (ce / l1 / confidence) - # ------------------------------------------------------------------ - - def _train_step( - self, - batch: dict, - accumulation_steps: int, - step: int, - batch_idx: int, - num_batches: int, - ) -> dict: - metrics = super()._train_step(batch, accumulation_steps, step, batch_idx, num_batches) - # Carry the components from the forward that _train_step just ran. - for key, value in getattr(self, "_last_loss_components", {}).items(): - metrics[key] = value - return metrics - - def _aggregate_metrics( - self, all_step_metrics: list[dict], step: int, *, grad_norm: torch.Tensor = None - ) -> dict: - metrics = super()._aggregate_metrics(all_step_metrics, step, grad_norm=grad_norm) - if all_step_metrics: - for key in ("ce_loss", "l1_loss", "confidence_loss"): - vals = [m[key] for m in all_step_metrics if key in m] - if not vals: - continue - value = torch.stack([v.float() for v in vals]).mean() - if dist.is_initialized() and dist.get_world_size() > 1: - dist.all_reduce(value, op=dist.ReduceOp.SUM) - value = value / dist.get_world_size() - metrics[f"train/{key}"] = value.item() - return metrics