fix: share forward pass / loss in DFlash / DSpark

.gitignore

@@ -89,3 +89,5 @@ wandb/log.txt
 
 .claude/
 wandb/
+data/
+profiles/

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()

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,

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"),

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,

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