Add Ministral3 DeepSpec reproduction experiments

config/dspark/dspark_ministral3_3b.py

@@ -0,0 +1,65 @@
+import os
+
+from deepspec.trainer import Ministral3DSparkTrainer
+
+BASE_TB_DIR = "/mnt/vast/runs/andy/deepspec_tensorboard"
+BASE_CKPT_DIR = "/mnt/vast/runs/andy/deepspec_checkpoints"
+project_name = "deepspec"
+exp_name = "dspark_block7_ministral3_3b"
+seed = 42
+
+model = dict(
+    target_model_name_or_path="mistralai/Ministral-3-3B-Instruct-2512",
+    block_size=7,
+    num_draft_layers=5,
+    target_layer_ids=[1, 7, 13, 19, 24],
+    mask_token_id=11,
+    num_anchors=512,
+    markov_rank=256,
+    markov_head_type="vanilla",
+    confidence_head_alpha=1.0,
+    confidence_head_with_markov=True,
+    loss_decay_gamma=4.0,
+    ce_loss_alpha=0.1,
+    l1_loss_alpha=0.9,
+)
+
+train = dict(
+    trainer_cls=Ministral3DSparkTrainer,
+    lr=6.0e-4,
+    warmup_ratio=0.04,
+    weight_decay=0.0,
+    precision="bf16",
+    local_batch_size=1,
+    global_batch_size=512,
+    num_train_epochs=10,
+    max_train_steps=None,
+    max_grad_norm=1.0,
+    sharding_strategy="no_shard",
+    torch_compile=False,
+)
+
+logging = dict(
+    logging_steps=10,
+    checkpointing_steps=3000,
+    save_only_checkpointing_steps=None,
+    keep_last_checkpoints=None,
+)
+
+data = dict(
+    target_cache_path=None,
+    chat_template="ministral3",
+    max_length=4096,
+    num_workers=4,
+)
+
+
+def finalize_cfg(cfg):
+    logging_cfg = dict(cfg["logging"])
+    project_name = str(cfg["project_name"])
+    exp_name = str(cfg["exp_name"])
+    logging_cfg["checkpoint_dir"] = os.path.join(BASE_CKPT_DIR, project_name, exp_name)
+    logging_cfg["tensorboard_dir"] = os.path.join(BASE_TB_DIR, project_name, exp_name)
+    cfg["logging"] = logging_cfg
+
+    return cfg

config/eagle3/eagle3_ministral3_3b.py

@@ -0,0 +1,58 @@
+import os
+
+from deepspec.trainer import Ministral3Eagle3Trainer
+
+BASE_TB_DIR = os.path.expanduser("~/tensorboard")
+BASE_CKPT_DIR = os.path.expanduser("~/checkpoints")
+project_name = "deepspec"
+exp_name = "eagle3_ttt7_ministral3_3b"
+seed = 0
+
+model = dict(
+    target_model_name_or_path="mistralai/Ministral-3-3B-Instruct-2512",
+    target_layer_ids=[1, 7, 13, 19, 24],
+    ttt_length=7,
+    step_loss_decay=0.8,
+    loss_type="soft_ce",
+    draft_num_hidden_layers=1,
+)
+
+train = dict(
+    trainer_cls=Ministral3Eagle3Trainer,
+    lr=6.0e-4,
+    warmup_ratio=0.04,
+    weight_decay=0.0,
+    precision="bf16",
+    local_batch_size=1,
+    global_batch_size=512,
+    num_train_epochs=10,
+    max_train_steps=None,
+    max_grad_norm=1.0,
+    sharding_strategy="no_shard",
+    torch_compile=False,
+)
+
+logging = dict(
+    logging_steps=10,
+    checkpointing_steps=3000,
+    save_only_checkpointing_steps=None,
+    keep_last_checkpoints=None,
+)
+
+data = dict(
+    target_cache_path=None,
+    chat_template="ministral3",
+    max_length=4096,
+    num_workers=4,
+)
+
+
+def finalize_cfg(cfg):
+    logging_cfg = dict(cfg["logging"])
+    project_name = str(cfg["project_name"])
+    exp_name = str(cfg["exp_name"])
+    logging_cfg["checkpoint_dir"] = os.path.join(BASE_CKPT_DIR, project_name, exp_name)
+    logging_cfg["tensorboard_dir"] = os.path.join(BASE_TB_DIR, project_name, exp_name)
+    cfg["logging"] = logging_cfg
+
+    return cfg

config/eagle3/eagle3_ministral3_3b_e2e_tv.py

@@ -0,0 +1,20 @@
+import os
+
+from config.eagle3.eagle3_ministral3_3b import *  # noqa: F403
+
+BASE_TB_DIR = "/mnt/vast/runs/andy/deepspec_tensorboard"
+BASE_CKPT_DIR = "/mnt/vast/runs/andy/deepspec_checkpoints"
+exp_name = "eagle3_ttt7_ministral3_3b_e2e_tv"
+model = dict(model)  # noqa: F405
+model["loss_type"] = "e2e_tv"
+
+
+def finalize_cfg(cfg):
+    logging_cfg = dict(cfg["logging"])
+    project_name = str(cfg["project_name"])
+    exp_name = str(cfg["exp_name"])
+    logging_cfg["checkpoint_dir"] = os.path.join(BASE_CKPT_DIR, project_name, exp_name)
+    logging_cfg["tensorboard_dir"] = os.path.join(BASE_TB_DIR, project_name, exp_name)
+    cfg["logging"] = logging_cfg
+
+    return cfg

config/eagle3/eagle3_ministral3_3b_e2e_tv_from_kl.py

@@ -0,0 +1,26 @@
+import os
+
+from config.eagle3.eagle3_ministral3_3b import *  # noqa: F403
+
+BASE_TB_DIR = "/mnt/vast/runs/andy/deepspec_tensorboard"
+BASE_CKPT_DIR = "/mnt/vast/runs/andy/deepspec_checkpoints"
+KL_DRAFT_CHECKPOINT = (
+    "/mnt/vast/home/andy/checkpoints/deepspec/"
+    "eagle3_ttt7_ministral3_3b_10k_3000steps/step_latest"
+)
+
+exp_name = "eagle3_ttt7_ministral3_3b_e2e_tv_from_kl"
+model = dict(model)  # noqa: F405
+model["loss_type"] = "e2e_tv"
+model["init_draft_model_name_or_path"] = KL_DRAFT_CHECKPOINT
+
+
+def finalize_cfg(cfg):
+    logging_cfg = dict(cfg["logging"])
+    project_name = str(cfg["project_name"])
+    exp_name = str(cfg["exp_name"])
+    logging_cfg["checkpoint_dir"] = os.path.join(BASE_CKPT_DIR, project_name, exp_name)
+    logging_cfg["tensorboard_dir"] = os.path.join(BASE_TB_DIR, project_name, exp_name)
+    cfg["logging"] = logging_cfg
+
+    return cfg

config/eagle3/eagle3_ministral3_3b_e2e_tv_from_kl_lr1e4.py

@@ -0,0 +1,5 @@
+from config.eagle3.eagle3_ministral3_3b_e2e_tv_from_kl import *  # noqa: F403
+
+exp_name = "eagle3_ttt7_ministral3_3b_e2e_tv_from_kl_lr1e4"
+train = dict(train)  # noqa: F405
+train["lr"] = 1.0e-4

config/eagle3/eagle3_ministral3_3b_e2e_tv_from_kl_lr1e4_8gpu.py

@@ -0,0 +1,3 @@
+from config.eagle3.eagle3_ministral3_3b_e2e_tv_from_kl_lr1e4 import *  # noqa: F403
+
+exp_name = "eagle3_ttt7_ministral3_3b_e2e_tv_from_kl_lr1e4_8gpu"

config/eagle3/eagle3_ministral3_3b_e2e_tv_ttt5.py

@@ -0,0 +1,21 @@
+import os
+
+from config.eagle3.eagle3_ministral3_3b import *  # noqa: F403
+
+BASE_TB_DIR = "/mnt/vast/runs/andy/deepspec_tensorboard"
+BASE_CKPT_DIR = "/mnt/vast/runs/andy/deepspec_checkpoints"
+exp_name = "eagle3_ttt5_ministral3_3b_e2e_tv"
+model = dict(model)  # noqa: F405
+model["loss_type"] = "e2e_tv"
+model["ttt_length"] = 5
+
+
+def finalize_cfg(cfg):
+    logging_cfg = dict(cfg["logging"])
+    project_name = str(cfg["project_name"])
+    exp_name = str(cfg["exp_name"])
+    logging_cfg["checkpoint_dir"] = os.path.join(BASE_CKPT_DIR, project_name, exp_name)
+    logging_cfg["tensorboard_dir"] = os.path.join(BASE_TB_DIR, project_name, exp_name)
+    cfg["logging"] = logging_cfg
+
+    return cfg

deepspec/data/parser.py

@@ -50,6 +50,16 @@ def get(self, name):
     ),
 )
 
+TEMPLATE_REGISTRY.register(
+    "ministral3",
+    ChatTemplate(
+        assistant_header="[/INST]",
+        user_header="[INST]",
+        system_prompt=None,
+        end_of_turn_token="</s>",
+    ),
+)
+
 
 class GeneralParser:
     def __init__(self, tokenizer, chat_template):

deepspec/eval/__init__.py

@@ -1,12 +1,22 @@
 from .base_evaluator import BaseEvaluator, DraftProposal, VerificationResult
-from .dspark import Gemma4DSparkEvaluator, Qwen3DSparkEvaluator
-from .eagle3 import Gemma4Eagle3Evaluator, Qwen3Eagle3Evaluator
+from .dspark import (
+    Gemma4DSparkEvaluator,
+    Ministral3DSparkEvaluator,
+    Qwen3DSparkEvaluator,
+)
+from .eagle3 import (
+    Gemma4Eagle3Evaluator,
+    Ministral3Eagle3Evaluator,
+    Qwen3Eagle3Evaluator,
+)
 
 __all__ = [
     "BaseEvaluator",
     "DraftProposal",
     "Gemma4Eagle3Evaluator",
     "Gemma4DSparkEvaluator",
+    "Ministral3DSparkEvaluator",
+    "Ministral3Eagle3Evaluator",
     "Qwen3Eagle3Evaluator",
     "Qwen3DSparkEvaluator",
     "VerificationResult",

deepspec/eval/base_evaluator.py

@@ -702,12 +702,38 @@ def record_dataset_metrics(
         )
         self.metrics_rows.append(metrics_row)
         self.print_dataset_result(metrics_row)
+        self.write_output_json(complete=False)
         return metrics_row
 
+    def write_output_json(self, *, complete: bool) -> None:
+        if (
+            self.args.output_json is None
+            or dist.get_rank() != 0
+            or not self.metrics_rows
+        ):
+            return
+        output_path = Path(self.args.output_json)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        tmp_path = output_path.with_suffix(output_path.suffix + ".tmp")
+        with tmp_path.open("w", encoding="utf-8") as handle:
+            json.dump(
+                {
+                    "target_model": self.args.target_name_or_path,
+                    "draft_model": self.args.draft_name_or_path,
+                    "step": self.args.step,
+                    "complete": complete,
+                    "rows": self.metrics_rows,
+                },
+                handle,
+                indent=2,
+            )
+        tmp_path.replace(output_path)
+
     def report_results(self) -> None:
         if dist.get_rank() == 0 and self.metrics_rows:
             if self.args.tensorboard_dir is not None:
                 self.log_tensorboard()
+            self.write_output_json(complete=True)
         self.print_results()
 
     def evaluate(self) -> None:

deepspec/eval/dspark/__init__.py

@@ -1,4 +1,8 @@
-from .evaluator import Gemma4DSparkEvaluator, Qwen3DSparkEvaluator
+from .evaluator import (
+    Gemma4DSparkEvaluator,
+    Ministral3DSparkEvaluator,
+    Qwen3DSparkEvaluator,
+)
 from .draft_ops import (
     DSparkDraftProposal,
     build_dspark_proposal,
@@ -8,6 +12,7 @@
 
 __all__ = [
     "Gemma4DSparkEvaluator",
+    "Ministral3DSparkEvaluator",
     "Qwen3DSparkEvaluator",
     "DSparkDraftProposal",
     "build_dspark_proposal",

deepspec/eval/dspark/draft_ops.py

@@ -8,10 +8,11 @@
 from deepspec.eval.base_evaluator import DraftProposal
 from deepspec.utils.sampling import logits_to_probs
 from deepspec.modeling.dspark.gemma4 import Gemma4DSparkModel
+from deepspec.modeling.dspark.ministral3 import Ministral3DSparkModel
 from deepspec.modeling.dspark.qwen3 import Qwen3DSparkModel
 
 
-DSparkModel = Qwen3DSparkModel | Gemma4DSparkModel
+DSparkModel = Qwen3DSparkModel | Gemma4DSparkModel | Ministral3DSparkModel
 
 
 @dataclass

deepspec/eval/dspark/evaluator.py

@@ -4,7 +4,7 @@
 from types import SimpleNamespace
 
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache
+from transformers import AutoTokenizer, DynamicCache
 
 from deepspec.eval.base_evaluator import (
     BaseEvaluator,
@@ -21,7 +21,9 @@
 )
 from deepspec.modeling.dspark.common import extract_context_feature
 from deepspec.modeling.dspark.gemma4 import Gemma4DSparkModel
+from deepspec.modeling.dspark.ministral3 import Ministral3DSparkModel
 from deepspec.modeling.dspark.qwen3 import Qwen3DSparkModel
+from deepspec.modeling.target_utils import load_target_causal_lm, load_target_tokenizer
 from deepspec.utils import jsonable
 
 
@@ -66,7 +68,7 @@ def _build_confidence_head_recorder(self) -> ConfidenceHeadRecorder | None:
         )
 
     def build_models(self) -> tuple[object, Qwen3DSparkModel, AutoTokenizer]:
-        target_model = AutoModelForCausalLM.from_pretrained(
+        target_model = load_target_causal_lm(
             self.args.target_name_or_path,
             dtype=torch.bfloat16,
             attn_implementation=self.EVAL_ATTN_IMPLEMENTATION,
@@ -79,7 +81,7 @@ def build_models(self) -> tuple[object, Qwen3DSparkModel, AutoTokenizer]:
         ).to(self.device).eval()
         assert_no_final_target_layer(target_model, draft_model.target_layer_ids)
         assert 0.0 <= float(self.args.confidence_threshold) <= 1.0
-        tokenizer = AutoTokenizer.from_pretrained(self.args.target_name_or_path)
+        tokenizer = load_target_tokenizer(self.args.target_name_or_path)
         return target_model, draft_model, tokenizer
 
     def _init_context(
@@ -223,3 +225,7 @@ def print_results(self) -> None:
 
 class Gemma4DSparkEvaluator(Qwen3DSparkEvaluator):
     draft_model_cls = Gemma4DSparkModel
+
+
+class Ministral3DSparkEvaluator(Qwen3DSparkEvaluator):
+    draft_model_cls = Ministral3DSparkModel

deepspec/eval/eagle3/__init__.py

@@ -1,3 +1,11 @@
-from .evaluator import Gemma4Eagle3Evaluator, Qwen3Eagle3Evaluator
+from .evaluator import (
+    Gemma4Eagle3Evaluator,
+    Ministral3Eagle3Evaluator,
+    Qwen3Eagle3Evaluator,
+)
 
-__all__ = ["Gemma4Eagle3Evaluator", "Qwen3Eagle3Evaluator"]
+__all__ = [
+    "Gemma4Eagle3Evaluator",
+    "Ministral3Eagle3Evaluator",
+    "Qwen3Eagle3Evaluator",
+]