DFlash for MiniMax-M3 (WIP): synthesis thinking-mode mix

examples/speculative_decoding/distributed_generate/worker.sh

@@ -20,10 +20,15 @@ BACKEND="$2"
 JOBS_PER_NODE="$3"
 SYSTEM_PROMPT="$4"
 
+# Optional model-specific serve flags via env, appended to the serve command. E.g. for
+# MiniMax-M3: VLLM_SERVE_EXTRA_ARGS="--block-size 128 --language-model-only" (--block-size
+# 128 is mandatory for M3's MSA sparse attention; --language-model-only skips the vision
+# encoder for text-only synthesis; KV cache stays bf16 — M3's MSA fused kernel rejects
+# fp8 KV).
 if [ "$BACKEND" == "vllm" ]; then
-    vllm serve /model/  --tensor-parallel-size 8 --served-model-name model --port 8000 --host 0.0.0.0 --trust-remote-code &
+    vllm serve /model/  --tensor-parallel-size 8 --served-model-name model --port 8000 --host 0.0.0.0 --trust-remote-code ${VLLM_SERVE_EXTRA_ARGS:-} &
 else
-    python3 -m sglang.launch_server --model-path /model --served-model-name model --tp 8 --port 8000 --host 0.0.0.0 --trust-remote-code &
+    python3 -m sglang.launch_server --model-path /model --served-model-name model --tp 8 --port 8000 --host 0.0.0.0 --trust-remote-code ${SGLANG_SERVE_EXTRA_ARGS:-} &
 fi
 # Wait for server to start up by polling the health endpoint
 echo "Waiting for server to start..."
@@ -59,6 +64,11 @@ if [ "$mpi_rank" -eq 0 ]; then
         if [ -n "$SYSTEM_PROMPT" ]; then
             cmd+=" --system_prompt $SYSTEM_PROMPT"
         fi
+        # Optional: cycle thinking modes for a mixed dataset (e.g. MiniMax-M3
+        # THINKING_MODES="enabled,disabled,adaptive").
+        if [ -n "${THINKING_MODES:-}" ]; then
+            cmd+=" --thinking-modes $THINKING_MODES"
+        fi
         echo "Running command: $cmd"
         eval $cmd
     done

examples/speculative_decoding/scripts/server_generate.py

@@ -54,7 +54,29 @@
     "--log_empty_conversations", action="store_true", help="Log empty conversations"
 )
 parser.add_argument("--system_prompt", nargs="+", type=str, default="", help="System prompt")
+parser.add_argument(
+    "--thinking-modes",
+    type=str,
+    default="",
+    help="Comma-separated thinking modes to cycle through per conversation, passed to the "
+    "server via chat_template_kwargs (e.g. 'enabled,disabled,adaptive' for MiniMax-M3). "
+    "Conversation i uses modes[i %% len(modes)], giving an even mix across the dataset. "
+    "Empty (default) sends no thinking_mode, preserving behavior for models without it.",
+)
+parser.add_argument(
+    "--output-format",
+    type=str,
+    default="oai",
+    choices=["oai", "sharegpt"],
+    help="Output chat format: 'oai' writes the OpenAI standard ({'messages': [{role, "
+    "content}, ...]}); 'sharegpt' writes the legacy {'conversations': [...]} key. Both "
+    "use role/content message dicts.",
+)
 args = parser.parse_args()
+MESSAGES_KEY = "messages" if args.output_format == "oai" else "conversations"
+
+# Parse the thinking-mode cycle; empty -> no thinking_mode injected.
+THINKING_MODES = [m.strip() for m in args.thinking_modes.split(",") if m.strip()]
 
 
 if args.data_path.endswith("jsonl"):
@@ -73,6 +95,14 @@ def generate_data(messages, idx, system_prompt):
     try:
         model_name = args.model
 
+        # Cycle thinking modes per conversation for an even mix across the dataset (e.g.
+        # MiniMax-M3 enabled/disabled/adaptive). Passed via chat_template_kwargs; empty
+        # list -> not sent.
+        thinking_mode = THINKING_MODES[idx % len(THINKING_MODES)] if THINKING_MODES else None
+        extra_body = (
+            {"chat_template_kwargs": {"thinking_mode": thinking_mode}} if thinking_mode else {}
+        )
+
         if args.chat:
             output_messages = []
 
@@ -105,6 +135,7 @@ def generate_data(messages, idx, system_prompt):
                         messages=output_messages,
                         max_tokens=args.max_tokens,
                         temperature=args.temperature,
+                        extra_body=extra_body,
                     )
                     if response.choices[0].finish_reason == "length":
                         break
@@ -123,7 +154,9 @@ def generate_data(messages, idx, system_prompt):
                     return
                 to_write = {"conversation_id": idx}
             else:
-                to_write = {"conversation_id": idx, "conversations": output_messages}
+                to_write = {"conversation_id": idx, MESSAGES_KEY: output_messages}
+                if thinking_mode:
+                    to_write["thinking_mode"] = thinking_mode
             with open(args.output_path, "a") as f:
                 # write in share gpt format
                 f.write(json.dumps(to_write) + "\n")
@@ -187,7 +220,17 @@ def generate_data(messages, idx, system_prompt):
     for idx, sample in enumerate(data):
         if idx in finished_ids:
             continue
-        future = executor.submit(generate_data, sample["conversations"], idx, system_prompt)
+        # Accept both ShareGPT ("conversations") and OAI-chat ("messages") prompt datasets
+        # (e.g. Speculative-Decoding-Dataset-v2 uses "messages"). generate_data already
+        # handles the from/value and role/content message shapes.
+        sample_messages = sample.get("conversations")
+        if sample_messages is None:
+            sample_messages = sample.get("messages")
+        if sample_messages is None:
+            raise KeyError(
+                f"sample {idx} has neither 'conversations' nor 'messages'; keys: {list(sample)}"
+            )
+        future = executor.submit(generate_data, sample_messages, idx, system_prompt)
         futures.append(future)
 
     for future in tqdm.tqdm(concurrent.futures.as_completed(futures), total=len(futures)):

modelopt/torch/speculative/plugins/modeling_fakebase.py

@@ -161,13 +161,31 @@ def from_source(cls, source: str, trust_remote_code: bool = False) -> "FakeBaseM
             ),
             orig_config,
         )
+
+        # Resolve the base model's dtype so lm_head / embed_tokens load in the checkpoint's
+        # dtype (e.g. bf16), matching the dumped hidden states in offline training. For VLMs
+        # base_cfg is the nested text config, which often does NOT set torch_dtype — only the
+        # top-level config does (e.g. MiniMax-M3: top-level torch_dtype=bfloat16, text_config
+        # has none). Reading base_cfg.dtype there falls back to PyTorch-default fp32, so the
+        # head loads fp32 while the dump is bf16 -> 'mat1 and mat2 must have the same dtype' in
+        # the draft forward. Prefer the nested torch_dtype, then the top-level, then bf16; map
+        # a string (e.g. "bfloat16") to the torch.dtype.
+        def _resolve_dtype(*cfgs):
+            for cfg in cfgs:
+                dt = getattr(cfg, "torch_dtype", None)
+                if isinstance(dt, str):
+                    dt = getattr(torch, dt, None)
+                if isinstance(dt, torch.dtype):
+                    return dt
+            return torch.bfloat16
+
         # Extract necessary info for spec training from base config
         config = FakeBaseConfig(
             num_hidden_layers=getattr(base_cfg, "num_hidden_layers", None),
             hidden_size=getattr(base_cfg, "hidden_size", None),
             vocab_size=getattr(base_cfg, "vocab_size", None),
             max_position_embeddings=getattr(base_cfg, "max_position_embeddings", None),
-            dtype=getattr(base_cfg, "dtype", torch.bfloat16),
+            dtype=_resolve_dtype(base_cfg, orig_config),
             tie_word_embeddings=getattr(base_cfg, "tie_word_embeddings", False),
             num_attention_heads=getattr(base_cfg, "num_attention_heads", None),
             num_key_value_heads=getattr(base_cfg, "num_key_value_heads", None),

tests/unit/torch/speculative/plugins/test_fakebase.py

@@ -65,6 +65,39 @@ def test_fakebase_local_happy_path(fake_checkpoint):
     assert model.embed_tokens.weight.shape == torch.Size([_VOCAB_SIZE, _HIDDEN_SIZE])
 
 
+def test_fakebase_vlm_dtype_from_top_level_config(tmp_path, monkeypatch):
+    """Regression (OMNIML-4747): for a VLM the language-model sub-config often carries no
+    torch_dtype — only the top-level config sets it (e.g. MiniMax-M3: top-level
+    torch_dtype=bfloat16, text_config none). FakeBaseModel must resolve the base dtype from the
+    top-level config rather than falling back to fp32; an fp32 head mismatches the bf16 offline
+    hidden states in the draft forward ('mat1 and mat2 must have the same dtype')."""
+    text_cfg = transformers.PretrainedConfig()
+    text_cfg.hidden_size = _HIDDEN_SIZE
+    text_cfg.vocab_size = _VOCAB_SIZE
+    text_cfg.num_hidden_layers = 2
+    text_cfg.max_position_embeddings = 128
+    text_cfg.tie_word_embeddings = False
+    text_cfg.torch_dtype = None  # the nested config carries no dtype (the M3 case)
+
+    cfg = transformers.PretrainedConfig()
+    cfg.model_type = "minimax_m3_vl"
+    cfg.text_config = text_cfg  # marks this as a VLM -> base_cfg becomes text_cfg
+    cfg.torch_dtype = "bfloat16"  # only the top-level config sets the dtype
+    monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", lambda *a, **kw: cfg)
+
+    safetensors.torch.save_file(
+        {
+            "lm_head.weight": torch.zeros(_VOCAB_SIZE, _HIDDEN_SIZE),
+            "embed_tokens.weight": torch.zeros(_VOCAB_SIZE, _HIDDEN_SIZE),
+        },
+        tmp_path / "model.safetensors",
+    )
+
+    model = FakeBaseModel.from_source(str(tmp_path))
+    assert model.lm_head.weight.dtype == torch.bfloat16
+    assert model.embed_tokens.weight.dtype == torch.bfloat16
+
+
 def test_fakebase_missing_index_raises(tmp_path, fake_config):
     with pytest.raises(FileNotFoundError, match="safetensors"):
         FakeBaseModel.from_source(str(tmp_path))