From 0831d6556a3683e644e5497e0cdf0c1d81c950de Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Tue, 16 Jun 2026 11:00:52 -0700 Subject: [PATCH 01/19] specdec: per-conversation thinking-mode mix in data synthesis (for MiniMax-M3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a --thinking-modes cycle to server_generate.py so synthetic conversations are generated across a mix of thinking modes — e.g. MiniMax-M3's enabled/disabled/adaptive, passed via chat_template_kwargs — so a DFlash/EAGLE draft trained on the data generalizes across modes. Conversation i uses modes[i % len(modes)] for an even split; the mode is recorded on each output record. Empty (default) sends no thinking_mode, unchanged for models without it. distributed_generate/worker.sh: pass THINKING_MODES through to server_generate.py, and add VLLM_SERVE_EXTRA_ARGS / SGLANG_SERVE_EXTRA_ARGS passthroughs for model-specific serve flags (M3 needs --block-size 128 for MSA sparse attention and --language-model-only for text-only synthesis; KV cache stays bf16 — M3's MSA fused kernel rejects fp8 KV). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- .../distributed_generate/worker.sh | 14 +++++++++-- .../scripts/server_generate.py | 23 +++++++++++++++++++ 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/examples/speculative_decoding/distributed_generate/worker.sh b/examples/speculative_decoding/distributed_generate/worker.sh index 97bf14c014a..01f01bb746b 100644 --- a/examples/speculative_decoding/distributed_generate/worker.sh +++ b/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 diff --git a/examples/speculative_decoding/scripts/server_generate.py b/examples/speculative_decoding/scripts/server_generate.py index 0fb71a0a0a1..46a87d70373 100644 --- a/examples/speculative_decoding/scripts/server_generate.py +++ b/examples/speculative_decoding/scripts/server_generate.py @@ -54,8 +54,20 @@ "--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.", +) args = parser.parse_args() +# 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"): with open(args.data_path) as f: @@ -73,6 +85,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 +125,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 @@ -124,6 +145,8 @@ def generate_data(messages, idx, system_prompt): to_write = {"conversation_id": idx} else: to_write = {"conversation_id": idx, "conversations": 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") From cdc03a6ebb8d6882e34b8dd2333767eb2eb4afdb Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Tue, 16 Jun 2026 11:51:01 -0700 Subject: [PATCH 02/19] specdec(server_generate): accept OAI 'messages' input + OAI-native output While wiring MiniMax-M3 synthesis (prompt set Speculative-Decoding-Dataset-v2 is OAI-format): - Accept both 'conversations' (ShareGPT) and 'messages' (OAI) prompt datasets on input (previously KeyError'd on 'messages'). - Add --output-format {oai,sharegpt} (default oai): emit the OpenAI standard {'messages': [{role, content}, ...]} instead of the legacy {'conversations': [...]}. Pass --output-format sharegpt for the old key. Validated end-to-end on MiniMax-M3-MXFP8 (single-node TP8 H100): the 3-way thinking-mode mix renders correctly (disabled -> direct answer; adaptive -> reasoning captured in content since no reasoning-parser), OAI in/out flows into the vLLM hidden-state dump. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- .../scripts/server_generate.py | 24 +++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/examples/speculative_decoding/scripts/server_generate.py b/examples/speculative_decoding/scripts/server_generate.py index 46a87d70373..6490055edc7 100644 --- a/examples/speculative_decoding/scripts/server_generate.py +++ b/examples/speculative_decoding/scripts/server_generate.py @@ -63,7 +63,17 @@ "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()] @@ -144,7 +154,7 @@ 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: @@ -210,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)): From a02bf8f92b6d9102ebff1e8bee0e3fe6fa686699 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Mon, 22 Jun 2026 09:44:50 -0700 Subject: [PATCH 03/19] specdec(dump): enable vLLM hidden-state dump for MiniMax-M3 (OMNIML-4747) Validated extract_hidden_states on MiniMax-M3-MXFP8 (single-node TP8 H100). Required M3 enablement in compute_hidden_states_vllm.py: - --block-size (M3's MSA sparse attention mandates 128; default None elsewhere). - --enforce-eager: M3's MSA Triton kernel (_gqa_sparse_fwd_kernel) JIT-recompiles per input shape; under cudagraph capture a recompile blows the executor RPC timeout and hangs the engine (sample_tokens timeout). Eager mode + a long VLLM_RPC_TIMEOUT fixes it. - --language-model-only: skip the vision encoder for text-only dumps (M3 is VL). - Read num_hidden_layers from text_config/llm_config for wrapped VL configs (MiniMaxM3VLConfig nests it; previously raised 'no num_hidden_layers attribute'). Output verified: per-conv .pt with input_ids / hidden_states (T,6144) / aux_hidden_states / loss_mask (length-matched) / conversation_id. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- .../compute_hidden_states_vllm.py | 37 ++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py index 77441f8f858..462de6afa65 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py @@ -110,6 +110,25 @@ def parse_args() -> argparse.Namespace: "--trust_remote_code", action="store_true", help="Trust remote code for HF models." ) parser.add_argument("--tp", type=int, default=None, help="Tensor parallel size.") + parser.add_argument( + "--block-size", + type=int, + default=None, + help="KV cache block size. Some models require a specific value — e.g. MiniMax-M3's " + "MSA sparse attention mandates 128. Default (None) lets vLLM choose.", + ) + parser.add_argument( + "--language-model-only", + action="store_true", + help="Skip the vision encoder for text-only dumps (multimodal models, e.g. MiniMax-M3).", + ) + parser.add_argument( + "--enforce-eager", + action="store_true", + help="Disable CUDA graph / torch.compile. Needed for MiniMax-M3: its MSA sparse " + "kernel JIT-recompiles per shape and a recompile can exceed the executor RPC " + "timeout under cudagraph capture, hanging the engine.", + ) parser.add_argument( "--debug-max-num-conversations", type=int, default=None, help="Limit conversations." ) @@ -168,7 +187,12 @@ def keep_conversation(entry): # Resolve the aux-layer indices and append the final-layer output. vLLM saves the # final (un-normed) hidden state when ``num_hidden_layers`` is passed as a layer id. config = AutoConfig.from_pretrained(args.model, trust_remote_code=args.trust_remote_code) - num_hidden_layers = getattr(config, "num_hidden_layers", None) + # Vision-language / wrapped configs (e.g. MiniMax-M3's MiniMaxM3VLConfig) nest the + # text model's layer count under text_config / llm_config rather than at the top level. + text_config = getattr(config, "text_config", None) or getattr(config, "llm_config", None) + num_hidden_layers = getattr(config, "num_hidden_layers", None) or getattr( + text_config, "num_hidden_layers", None + ) if num_hidden_layers is None: raise ValueError(f"model config has no 'num_hidden_layers' attribute: {config}") aux_layer_ids = _resolve_aux_layers_standalone( @@ -244,12 +268,23 @@ def keep_conversation(entry): storage_path.mkdir(parents=True, exist_ok=True) atexit.register(shutil.rmtree, storage_path, ignore_errors=True) + # Model-specific extras (e.g. MiniMax-M3 mandates block_size=128 for MSA sparse + # attention; --language-model-only skips the vision encoder for text-only dumps). + extra_llm_kwargs = {} + if args.block_size is not None: + extra_llm_kwargs["block_size"] = args.block_size + if args.language_model_only: + extra_llm_kwargs["language_model_only"] = True + if args.enforce_eager: + extra_llm_kwargs["enforce_eager"] = True + llm = LLM( model=args.model, tensor_parallel_size=tp, max_model_len=args.max_seq_len, trust_remote_code=args.trust_remote_code, enable_chunked_prefill=False, # required by extract_hidden_states + **extra_llm_kwargs, # With prefix caching on, vLLM serves shared prefixes from cache in block-sized # chunks and the hidden-state connector only emits the freshly-computed suffix, so # the dumped hidden_states come out short by N*block_size vs the full input_ids / From 8961d0358ebe51dce8394b577482150682baa18c Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Mon, 22 Jun 2026 11:04:31 -0700 Subject: [PATCH 04/19] specdec(recipe): add MiniMax-M3 DFlash offline recipe + train template (OMNIML-4747) 2-step offline DFlash recipe for MiniMax-M3 (427B VL-MoE), mirroring MiniMax-M2.7-DFlash: - hf_offline_dflash.yaml: dump (vLLM extract_hidden_states, MXFP8 single-node TP8) + train (FakeBaseModel on bf16). M3-specific: --block-size 128 (MSA), --language-model-only, --enforce-eager + VLLM_RPC_TIMEOUT=1800000 (avoid MSA Triton-kernel JIT RPC-timeout hang), seq-len 8192 end-to-end, mask token 200061 (200054 is a real special token in M3), OVERRIDE_TRANSFORMERS 4.52.4, export-YaRN original_max_position 8192 / factor 24 (tunable; 128 for full 1M). - chat_template_train.jinja: M3 chat template with {% generation %} wrapping the assistant turn (think + content + tool_calls) for answer_only_loss; header + eos sit outside the span, matching the M2.7 convention. Thinking-mode handling preserved verbatim. Validated: generation spans cover exactly the assistant outputs across multi-turn + no-think turns. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- .../chat_template_train.jinja | 256 ++++++++++++++++++ .../MiniMax-M3-DFlash/hf_offline_dflash.yaml | 125 +++++++++ 2 files changed, 381 insertions(+) create mode 100644 tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/chat_template_train.jinja create mode 100644 tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml diff --git a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/chat_template_train.jinja b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/chat_template_train.jinja new file mode 100644 index 00000000000..8d2e37cae3e --- /dev/null +++ b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/chat_template_train.jinja @@ -0,0 +1,256 @@ +{# MiniMax-M3 chat template with {% generation %} tags for answer_only_loss training. + Adapted from https://huggingface.co/MiniMaxAI/MiniMax-M3/blob/main/chat_template.jinja + with {% generation %} / {% endgeneration %} wrapping the assistant turn's output + (think + content + tool_calls), matching the MiniMax-M2.7-DFlash convention: the + ']~b]ai\n' header and the trailing eos sit OUTSIDE the generation span, so the loss + mask covers only what the model produces. Thinking-mode handling is preserved verbatim + so dumps reflect the same enabled/disabled/adaptive mix used during synthesis. +#} +{# ---------- special token variables ---------- #} +{%- set ns_token = ']<]minimax[>[' -%} +{%- set bod_token = ']~!b[' -%} +{%- set bos_token = ']~b]' -%} +{%- set eos_token = '[e~[' -%} +{%- set toolcall_begin_token = ns_token ~ '' -%} +{%- set toolcall_end_token = ns_token ~ '' -%} +{%- set think_begin_token = '' -%} +{%- set think_end_token = '' -%} +{%- set image_token = ']<]image[>[' -%} +{%- set video_token = ']<]video[>[' -%} +{#- Thinking mode: "enabled" / "disabled" / "adaptive" / not defined -#} +{#- Recursive XML renderer for tool_call arguments ======================== -#} +{#- None values are intentionally skipped in mapping iteration so that + `null` (which would round-trip to the literal string "null") + never appears in the rendered tool_call. The convention is: omit the + field entirely. The top-level `_args` loop applies the same rule. + The `val is none` branch below is a safety net only — upstream cleaning + (drop_none_in_tool_arguments) should ensure no None ever reaches here. -#} +{%- macro to_xml(val, ns) -%} +{%- if val is mapping -%} +{%- for k, v in val.items() if v is not none -%} +{{ ns }}<{{ k }}>{{ to_xml(v, ns) }}{{ ns }} +{%- endfor -%} +{%- elif val is iterable and val is not string -%} +{%- for item in val -%} +{{ ns }}{{ to_xml(item, ns) }}{{ ns }} +{%- endfor -%} +{%- elif val is none -%} +{#- Should be unreachable when upstream cleaning is applied. -#} +{%- elif val is boolean -%} +{{ val | tojson }} +{%- else -%} +{{ val }} +{%- endif -%} +{%- endmacro -%} +{#- Tool Rendering Functions ============================================== -#} +{%- macro render_tool_namespace(namespace_name, tool_list) -%} +{%- for tool in tool_list -%} +{{ tool.function | tojson(ensure_ascii=False) }} +{% endfor -%} +{%- endmacro -%} +{%- macro visible_text(content) -%} + {%- if content is string -%} + {{ content }} + {%- elif content is iterable and content is not mapping -%} + {%- for item in content -%} + {%- if item is mapping and item.type == 'text' -%} + {{- item.text }} + {%- elif item is mapping and item.type == 'image' -%} + {{- image_token }} + {%- elif item is mapping and item.type == 'video' -%} + {{- video_token}} + {%- elif item is string -%} + {{- item }} + {%- endif -%} + {%- endfor -%} + {%- elif content is none -%} + {{- '' }} + {%- else -%} + {{- content }} + {%- endif -%} +{%- endmacro -%} +{#- System Message Construction ============================================ -#} +{%- macro build_system_message(system_message) -%} + {%- if system_message and system_message.content -%} + {{- visible_text(system_message.content) }} + {%- else -%} + {{- 'Your model version is MiniMax-M3, developed by MiniMax. Knowledge cutoff: January 2026. Founded in early 2022, MiniMax is a global AI foundation model company committed to advancing the frontiers of AI towards AGI.' }} + {%- endif -%} + + {#- Thinking mode instructions -#} + {{- '\n\n\n' }} + {{- 'You have a thinking capability that allows you to reason step by step before responding. When thinking is enabled, wrap your reasoning in ' ~ think_begin_token ~ think_end_token ~ ' tags before your response. When thinking is disabled, begin your response directly after the ' ~ think_end_token ~ ' prefix. When thinking is adaptive, decide on your own whether to think for the current turn.\n' }} + {%- if thinking_mode is defined -%} + {%- if thinking_mode == "enabled" -%} + {{- 'Current thinking mode: enabled. You MUST think step by step before every response, including after receiving function/tool results.\n' }} + {%- elif thinking_mode == "disabled" -%} + {{- 'Current thinking mode: disabled. Do not output any thinking process.\n' }} + {%- elif thinking_mode == "adaptive" -%} + {{- 'Current thinking mode: adaptive. You are encouraged to think for complex decision-making, multi-step reasoning, or when analyzing function/tool results.\n' }} + {%- endif -%} + {%- else -%} + {{- 'Current thinking mode: adaptive. You are encouraged to think for complex decision-making, multi-step reasoning, or when analyzing function/tool results.\n' }} + {%- endif -%} + {{- '' }} +{%- endmacro -%} +{%- macro build_developer_message(developer_message) -%} + {%- if developer_message and developer_message.content -%} + {{- visible_text(developer_message.content) }} + {%- else -%} + {%- if model_identity is not defined -%} + {%- set model_identity = "You are a helpful assistant." -%} + {%- endif -%} + {{- model_identity }} + {%- endif -%} +{%- endmacro -%} +{#- Main Template Logic ================================================= -#} +{#- Role mapping: root -> system sp (high priority), system/developer -> developer sp (low priority) -#} +{%- set system_message = none -%} +{%- set developer_message = none -%} +{%- set conversation_messages = messages -%} +{%- if messages and messages[0].role == "root" -%} + {%- set system_message = messages[0] -%} + {%- set conversation_messages = messages[1:] -%} + {%- if conversation_messages and conversation_messages[0].role in ["system", "developer"] -%} + {%- set developer_message = conversation_messages[0] -%} + {%- set conversation_messages = conversation_messages[1:] -%} + {%- endif -%} +{%- elif messages and messages[0].role in ["system", "developer"] -%} + {%- set developer_message = messages[0] -%} + {%- set conversation_messages = messages[1:] -%} +{%- endif -%} +{#- Render system sp (higher priority, root role only) -#} +{{- bod_token ~ bos_token ~ 'system' ~ '\n' }} +{{- build_system_message(system_message) }} +{{- eos_token ~ '\n' }} + +{#- Render developer sp (lower priority: system/developer role + tools) -#} +{{- bos_token ~ 'developer' ~ '\n' }} +{{- build_developer_message(developer_message) }} +{%- if tools -%} + {{- '\n\n' ~ '# Tools' ~ '\n' ~ 'You may call one or more tools to assist with the user query.\nHere are the tools available in JSONSchema format:' ~ '\n' }} + {{- '\n' ~ '' ~ '\n' }} + {{- render_tool_namespace("functions", tools) }} + {{- '' ~ '\n\n' }} + {{- 'To call tools, wrap all invocations in a single ' ~ toolcall_begin_token ~ toolcall_end_token ~ ' block. Parameter values containing nested objects or arrays are recursively expanded into XML elements. Example:\n' }} + {{- '\n' ~ toolcall_begin_token ~ '\n' }} + {{- ns_token + '' }} + {{- ns_token + 'value-1' + ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + 'val-a' + ns_token + '' }} + {{- ns_token + 'val-b' + ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + '' }} + {{- ns_token + '\n' }} + {{- ns_token + '' }} + {{- ns_token + 'value-1' + ns_token + '' }} + {{- ns_token + '\n' }} + {{- toolcall_end_token }} +{%- endif -%} +{{- eos_token ~ '\n' }} + +{#- Render messages -#} +{%- set last_tool_call = namespace(name=none) -%} +{%- for message in conversation_messages -%} + {%- if message.role == 'assistant' -%} + {{- bos_token ~ 'ai' ~ '\n' }} + {%- generation -%} + {%- set reasoning_content = '' %} + {%- set content = visible_text(message.content) %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if think_end_token in content %} + {%- set reasoning_content = content.split(think_end_token)[0].strip('\n').split(think_begin_token)[-1].strip('\n') %} + {%- set content = content.split(think_end_token)[-1].strip('\n') %} + {%- endif %} + {%- endif %} + + {%- if reasoning_content -%} + {#- Render thinking for every assistant turn (all-turn visible) -#} + {{- think_begin_token ~ reasoning_content ~ think_end_token }} + {%- else -%} + {#- No thinking rendered → prefix with think_end_token -#} + {{- think_end_token }} + {%- endif -%} + + {%- if content -%} + {{- content }} + {%- endif -%} + {%- if message.tool_calls -%} + {{- toolcall_begin_token ~ '\n' }} + + {%- for tool_call in message.tool_calls -%} + {%- if tool_call.function -%} + {%- set tool_call = tool_call.function -%} + {%- endif -%} +{{- ns_token + '' }} +{%- set _args = tool_call.arguments -%} +{%- for k, v in _args.items() if v is not none %} +{{- ns_token + '<' + k + '>' -}} +{{- to_xml(v, ns_token) -}} +{{- ns_token + '' }} +{%- endfor -%} +{{- ns_token + '' ~ '\n' }} + {%- endfor -%} + + {{- toolcall_end_token }} + {%- if message.tool_calls[-1].function -%} + {%- set last_tool_call.name = message.tool_calls[-1].function.name -%} + {%- else -%} + {%- set last_tool_call.name = message.tool_calls[-1].name -%} + {%- endif -%} + {%- else -%} + {%- set last_tool_call.name = none -%} + {%- endif -%} + {%- endgeneration -%} + {{- eos_token ~ '\n' }} + + {%- elif message.role == 'tool' -%} + {%- if last_tool_call.name is none -%} + {{- raise_exception("Message has tool role, but there was no previous assistant message with a tool call!") }} + {%- endif -%} + {%- if loop.first or (conversation_messages[loop.index0 - 1].role != 'tool') -%} + {{- bos_token ~ 'tool' }} + {%- endif -%} + {{- '\n' }} + {%- if message.content is string -%} + {{- message.content }} + {%- else -%} + {%- for tr in message.content -%} + {%- if tr is mapping and tr.type is defined and tr.type == 'image' -%} + {{- image_token }} + {%- elif tr is mapping and tr.type is defined and tr.type == 'video' -%} + {{- video_token }} + {%- else -%} + {{- tr.output if tr.output is defined else (tr.text if tr.type == 'text' and tr.text is defined else tr) }} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {{- '' }} + {%- if loop.last or (conversation_messages[loop.index0 + 1].role != 'tool') -%} + {{- eos_token ~ '\n' -}} + {%- endif -%} + + {%- elif message.role == 'user' -%} + {{- bos_token ~ 'user' ~ '\n' }} + {{- visible_text(message.content) }} + {{- eos_token ~ '\n' }} + {%- endif -%} +{%- endfor -%} + +{#- Generation prompt -#} +{%- if add_generation_prompt -%} +{{- bos_token ~ 'ai' ~ '\n' }} +{%- if thinking_mode is defined and thinking_mode == "disabled" -%} + {{- think_end_token }} +{%- elif thinking_mode is defined and thinking_mode == "adaptive" -%} + {#- adaptive: no prefix, let model decide -#} +{%- elif thinking_mode is defined and thinking_mode == "enabled" -%} + {#- enabled or not defined: default to think -#} + {{- think_begin_token }} +{%- else -%} + {#- adaptive: no prefix, let model decide -#} +{%- endif -%} +{%- endif -%} diff --git a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml new file mode 100644 index 00000000000..a0e64e7fd91 --- /dev/null +++ b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml @@ -0,0 +1,125 @@ +# DFlash offline speculative decoding training for MiniMax-M3 (427B VL-MoE, 26B active). +# +# 2-step pipeline (mirrors MiniMax-M2.7-DFlash/hf_offline_dflash.yaml). Offline is the +# chosen path for M3 — online FSDP2 training streams the 427B base forward at every step +# and is too slow at scale: +# task_0: Dump base-model hidden states once via vLLM extract_hidden_states. +# task_1: Train the DFlash draft on the dump (FakeBaseModel — loads only lm_head + +# embed_tokens, not the full 427B base). +# +# M3-specific notes (differ from M2.7), all validated 2026-06-22: +# * Dump serves MiniMax-M3-MXFP8 (NVIDIA-published quant) single-node TP8 on H100. M3 +# is not in stable vLLM yet -> image vllm/vllm-openai:minimax-m3. +# * --block-size 128 is MANDATORY for M3's MSA sparse attention. +# * --language-model-only skips the vision encoder (text-only synth/dump). +# * --enforce-eager + VLLM_RPC_TIMEOUT=1800000 are REQUIRED: M3's MSA Triton kernel +# (_gqa_sparse_fwd_kernel) JIT-recompiles per input shape; under cudagraph capture a +# recompile blows the executor RPC timeout (sample_tokens timeout -> EngineDead hang). +# Eager mode + a long RPC timeout avoids the hang. KV cache stays bf16 (M3's MSA fused +# kernel rejects fp8 KV). +# * Training FakeBaseModel reads lm_head + embed_tokens from the bf16 M3 (real weights; +# these tensors are not what MXFP8 quantizes, so dump@MXFP8 / train@bf16 logits stay +# consistent). Per Ye Yu: adhere to published bf16/MXFP8 ckpts, do not self-quantize. +# * Sequence length 8192 (not M2.7's 4096) end-to-end: synth, dump, training — captures +# full reasoning across the enabled/disabled/adaptive mode mix. +# +# Reference: "DFlash: Block Diffusion for Flash Speculative Decoding" (arXiv:2602.06036) +# +# Usage: +# uv run slurm.py --yaml modules/Model-Optimizer/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml --yes + +job_name: MiniMax-M3-DFlash_offline +pipeline: + global_vars: + # bf16 base — used by training's FakeBaseModel (lm_head + embed_tokens) and tokenizer. + hf_model: /hf-local/MiniMaxAI/MiniMax-M3 + # NVIDIA-published MXFP8 quant — used only to serve the dump single-node TP8 on H100. + dump_model: /hf-local/MiniMaxAI/MiniMax-M3-MXFP8 + + # Step 1: Dump base-model hidden states via vLLM extract_hidden_states (TP=8, MXFP8). + task_0: + script: common/eagle3/dump_offline_data_vllm.sh + args: + # Synthetic data from the M3 synth campaign (default.jsonl, 3-way thinking-mode mix), + # cleaned + uploaded. Update the suffix once the cleaned set is published. + - --input-data /hf-local/modelopt/MiniMax-M3-synthetic-data + - --output-dir /scratchspace/dflash_minimax_m3_hidden_states + # Must match the draft model's num_hidden_layers (recipe default: 5). + - --aux-layers dflash + - --answer-only-loss + - --chat-template examples/MiniMax/MiniMax-M3-DFlash/chat_template_train.jinja + - --max-seq-len 8192 + - --tp 8 + # M3 MSA requirements (see header). + - --block-size 128 + - --language-model-only + - --enforce-eager + environment: + - HF_MODEL_CKPT: <> + - TRUST_REMOTE_CODE: "1" + # Survive MSA Triton-kernel JIT recompiles without an executor RPC timeout. + - VLLM_RPC_TIMEOUT: "1800000" + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + container: vllm/vllm-openai:minimax-m3 + + # Step 2: Train DFlash offline on the dumped hidden states. FakeBaseModel avoids loading + # the full 427B — only lm_head + embed_tokens are read from the bf16 checkpoint. + task_1: + script: common/specdec/dflash_online_training.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dflash.yaml + - model.model_name_or_path=<> + - model.trust_remote_code=true + - model.use_fake_base_for_offline=true + - data.mode=offline + - data.offline_data_path=/scratchspace/dflash_minimax_m3_hidden_states + - data.chat_template=examples/MiniMax/MiniMax-M3-DFlash/chat_template_train.jinja + - training.output_dir=/scratchspace/dflash_minimax_m3_offline + - training.num_train_epochs=10 + # bs=1 @ 8192 keeps the activation footprint of M2.7's bs=2 @ 4096 (bs*seqlen equal). + - training.per_device_train_batch_size=1 + - training.learning_rate=1.2e-3 + - training.warmup_steps=100 + - training.training_seq_len=8192 + - training.logging_steps=100 + - training.save_steps=400 + - training.disable_tqdm=true + - training.dp_shard_size=1 + - training.answer_only_loss=true + - training.ddp_timeout=3600 + - training.bf16=false + - dflash.dflash_self_logit_distillation=true + - dflash.dflash_block_size=8 + - dflash.dflash_num_anchors=512 + - dflash.dflash_loss_decay_factor=4.0 + - dflash.dflash_architecture_config.num_hidden_layers=5 + # Mask token id: in M3, 200054 is a real special token, so the first unused reserved + # embedding row is 200061 (M2.7 used 200054). + - dflash.dflash_mask_token_id=200061 + # YaRN rope_scaling injected at EXPORT time only (config.json field; draft weights + # unchanged) -> tunable per export. original_max_position_embeddings = training_seq_len + # (8192). Default factor 24 -> 8192*24 = 196608 served context (matches M2.7's target). + # For M3's full 1M context use factor 128 (8192*128 = 1048576). + - dflash.dflash_export_rope_scaling.type=yarn + - dflash.dflash_export_rope_scaling.factor=24.0 + - dflash.dflash_export_rope_scaling.original_max_position_embeddings=8192 + - dflash.dflash_export_rope_scaling.beta_fast=1.0 + - dflash.dflash_export_rope_scaling.beta_slow=1.0 + - dflash.dflash_export_rope_scaling.mscale=1.0 + - dflash.dflash_export_rope_scaling.mscale_all_dim=1.0 + environment: + - NUM_NODES: "8" + # Offline training uses a lightweight FakeBaseModel, so plain DDP suffices (no + # ACCELERATE_CONFIG / FSDP2 patches). OVERRIDE_TRANSFORMERS pins 4.52.4 for the + # MiniMax-M3 config load. + - OVERRIDE_TRANSFORMERS: "4.52.4" + - MIXED_PRECISION: "no" + slurm_config: + _factory_: "slurm_factory" + nodes: 8 + ntasks_per_node: 1 + gpus_per_node: 8 From 0aa2880fb7cc806650303dec1069a30437e9d830 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Mon, 22 Jun 2026 11:07:25 -0700 Subject: [PATCH 05/19] specdec(recipe): default MiniMax-M3 DFlash export YaRN to full 1M (factor 128) original_max_position_embeddings=8192 (training seq-len) x factor 128 = 1048576 = M3's full 1M context. Export-time tunable (factor 24 -> 196608 for the M2.7-equivalent target). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- .../MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml index a0e64e7fd91..bc848231cc2 100644 --- a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml +++ b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml @@ -102,10 +102,10 @@ pipeline: - dflash.dflash_mask_token_id=200061 # YaRN rope_scaling injected at EXPORT time only (config.json field; draft weights # unchanged) -> tunable per export. original_max_position_embeddings = training_seq_len - # (8192). Default factor 24 -> 8192*24 = 196608 served context (matches M2.7's target). - # For M3's full 1M context use factor 128 (8192*128 = 1048576). + # (8192). factor 128 -> 8192*128 = 1048576 = M3's full 1M context. (Use factor 24 -> + # 196608 to match M2.7's served target instead.) - dflash.dflash_export_rope_scaling.type=yarn - - dflash.dflash_export_rope_scaling.factor=24.0 + - dflash.dflash_export_rope_scaling.factor=128.0 - dflash.dflash_export_rope_scaling.original_max_position_embeddings=8192 - dflash.dflash_export_rope_scaling.beta_fast=1.0 - dflash.dflash_export_rope_scaling.beta_slow=1.0 From b1ebd545df68831d9d7833b357bb7833a676c6ad Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Thu, 25 Jun 2026 11:56:53 -0700 Subject: [PATCH 06/19] specdec(recipe): point MiniMax-M3 DFlash dump at the cleaned synth dataset Synth campaign complete: 2,761,690 clean records / 16,472 shards at /hf-local/modelopt/MiniMax-M3-synthetic-data-clean (empties + finished-markers dropped, provenance joined from source prompts, even 3-way thinking-mode mix). Replaces the pre-campaign placeholder path. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- .../MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml index bc848231cc2..46b9e94515d 100644 --- a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml +++ b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml @@ -40,9 +40,10 @@ pipeline: task_0: script: common/eagle3/dump_offline_data_vllm.sh args: - # Synthetic data from the M3 synth campaign (default.jsonl, 3-way thinking-mode mix), - # cleaned + uploaded. Update the suffix once the cleaned set is published. - - --input-data /hf-local/modelopt/MiniMax-M3-synthetic-data + # Synthetic data from the M3 synth campaign (default.jsonl, even 3-way thinking-mode + # mix), consolidated + cleaned: 2,761,690 records across 16,472 shards (empties + + # finished-markers dropped, provenance joined from the source prompts). + - --input-data /hf-local/modelopt/MiniMax-M3-synthetic-data-clean - --output-dir /scratchspace/dflash_minimax_m3_hidden_states # Must match the draft model's num_hidden_layers (recipe default: 5). - --aux-layers dflash From 34b591e964bb2659eafa8fa2aac2b26fc57fc3d0 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Mon, 29 Jun 2026 09:55:53 -0700 Subject: [PATCH 07/19] specdec(dump): length-bucketing + warm-up for MiniMax-M3 MSA dump (OMNIML-4747) Dry run surfaced that M3's MSA Triton kernel (_gqa_sparse_fwd_kernel) JIT-recompiles per prefill shape (~110s vs ~10s cached). On real 8192-token data hundreds of distinct lengths each recompile, blowing VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS (default 300) -> sample_tokens RPC timeout -> EngineDeadError. (M2.7 uses dense attention + stock vLLM, so it never hits this.) The earlier VLLM_RPC_TIMEOUT 'fix' was a no-op: that env var is unrecognized in this vLLM build and silently ignored. compute_hidden_states_vllm.py: - --length-buckets: right-pad each prompt to the smallest bucket >= its length so the model sees only a handful of prefill shapes. Causal attention leaves the real prefix's hidden states identical to the unpadded forward; pad positions are sliced off on save via real_len. - --warmup: one throwaway prefill per bucket to JIT-compile shapes up front (and populate a persistent TRITON_CACHE_DIR shared across tasks). M3 recipe: add --length-buckets 1024,2048,4096,8192 + --warmup; set VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 (was bogus VLLM_RPC_TIMEOUT) + a persistent TRITON_CACHE_DIR so all parallel dump tasks reuse the warmed kernels. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- .../compute_hidden_states_vllm.py | 59 ++++++++++++++++++- .../MiniMax-M3-DFlash/hf_offline_dflash.yaml | 23 +++++--- 2 files changed, 72 insertions(+), 10 deletions(-) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py index 462de6afa65..eb1203ab80d 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py @@ -129,6 +129,26 @@ def parse_args() -> argparse.Namespace: "kernel JIT-recompiles per shape and a recompile can exceed the executor RPC " "timeout under cudagraph capture, hanging the engine.", ) + parser.add_argument( + "--length-buckets", + type=str, + default="", + help="Comma-separated token-length buckets (e.g. '1024,2048,4096,8192'). When set, each " + "prompt is right-padded to the smallest bucket >= its length, so the model sees only a " + "handful of distinct prefill shapes. This bounds per-shape kernel JIT recompiles for " + "attention that compiles per sequence length (e.g. MiniMax-M3's MSA sparse attention), " + "which otherwise recompiles for hundreds of lengths — each recompile can exceed " + "VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS and kill the engine. Pad positions are dropped on " + "save; causal attention leaves the real prefix's hidden states unchanged. Empty " + "(default) disables bucketing.", + ) + parser.add_argument( + "--warmup", + action="store_true", + help="Before the real dump, run one throwaway prefill per length bucket to JIT-compile " + "the per-shape kernels up front (populating a persistent TRITON_CACHE_DIR shared across " + "tasks), so the timed dump runs on cached kernels. Only meaningful with --length-buckets.", + ) parser.add_argument( "--debug-max-num-conversations", type=int, default=None, help="Limit conversations." ) @@ -214,10 +234,18 @@ def keep_conversation(entry): if args.answer_only_loss: verify_generation_tags(tokenizer.chat_template) + # Length buckets (optional): bound the number of distinct prefill shapes the model sees. + buckets = sorted({int(b) for b in args.length_buckets.split(",") if b.strip()}) + buckets = [b for b in buckets if b <= args.max_seq_len] + pad_token_id = tokenizer.pad_token_id + if buckets: + print(f"Length bucketing enabled: buckets={buckets}, pad_token_id={pad_token_id}") + # Prepare prompts for vLLM prompts = [] conversation_ids = [] loss_masks = [] + real_lens = [] num_skipped_too_long = 0 num_invalid = 0 @@ -243,9 +271,19 @@ def keep_conversation(entry): num_skipped_too_long += 1 continue - prompts.append(TokensPrompt(prompt_token_ids=input_ids.tolist())) + token_ids_list = input_ids.tolist() + # Right-pad to the smallest bucket >= length. Causal attention means the real prefix's + # hidden states are identical to the unpadded forward; pad positions are sliced off on + # save (using real_len). Bounds distinct prefill shapes -> bounded kernel recompiles. + if buckets: + bucket = next((b for b in buckets if b >= num_tokens), buckets[-1]) + if bucket > num_tokens: + token_ids_list = token_ids_list + [pad_token_id] * (bucket - num_tokens) + + prompts.append(TokensPrompt(prompt_token_ids=token_ids_list)) conversation_ids.append(conversation_id) loss_masks.append(loss_mask) + real_lens.append(num_tokens) print( f"Prepared {len(prompts)} prompts ({num_skipped_too_long} skipped too long, {num_invalid} invalid)" @@ -310,13 +348,22 @@ def keep_conversation(entry): ), ) + # Warm up the per-shape kernel JIT (e.g. MiniMax-M3 MSA) once per bucket so the timed dump + # runs on cached kernels and no single step exceeds VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS. With a + # persistent TRITON_CACHE_DIR this also primes the cache for every other parallel task. + if args.warmup and buckets: + print(f"Warming up {len(buckets)} bucket shapes: {buckets}") + warmup_prompts = [TokensPrompt(prompt_token_ids=[pad_token_id] * b) for b in buckets] + llm.generate(warmup_prompts, SamplingParams(max_tokens=1)) + print("Warm-up complete.") + # max_tokens=1: we only need a single forward pass over the prompt tokens. outputs = llm.generate(prompts, SamplingParams(max_tokens=1)) # Save in the same format as compute_hidden_states_hf.py, including loss_mask. num_success = 0 - for conv_id, loss_mask, output in tqdm( - zip(conversation_ids, loss_masks, outputs), total=len(outputs), desc="Saving" + for conv_id, loss_mask, real_len, output in tqdm( + zip(conversation_ids, loss_masks, real_lens, outputs), total=len(outputs), desc="Saving" ): hidden_states_path = output.kv_transfer_params.get("hidden_states_path") if hidden_states_path is None: @@ -329,6 +376,12 @@ def keep_conversation(entry): # extract_layer_ids. Last layer = final output; the rest = aux layers. hidden_states = obj["hidden_states"] + # Drop any right-padding from length-bucketing: keep only the real prefix. Causal + # attention means these positions' states match the unpadded forward. No-op when + # bucketing is disabled (real_len == sequence length). + token_ids = token_ids[:real_len] + hidden_states = hidden_states[:real_len] + output_hidden_states = hidden_states[:, -1, :].cpu() if hidden_states.shape[1] > 1: # Concatenate aux layers along the hidden dim, matching the HF dump format. diff --git a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml index 46b9e94515d..88e55a612bc 100644 --- a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml +++ b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml @@ -12,11 +12,15 @@ # is not in stable vLLM yet -> image vllm/vllm-openai:minimax-m3. # * --block-size 128 is MANDATORY for M3's MSA sparse attention. # * --language-model-only skips the vision encoder (text-only synth/dump). -# * --enforce-eager + VLLM_RPC_TIMEOUT=1800000 are REQUIRED: M3's MSA Triton kernel -# (_gqa_sparse_fwd_kernel) JIT-recompiles per input shape; under cudagraph capture a -# recompile blows the executor RPC timeout (sample_tokens timeout -> EngineDead hang). -# Eager mode + a long RPC timeout avoids the hang. KV cache stays bf16 (M3's MSA fused -# kernel rejects fp8 KV). +# * M3's MSA Triton kernel (_gqa_sparse_fwd_kernel) JIT-recompiles per prefill shape; a +# recompile (~110s) can exceed VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS (default 300) and kill +# the engine (sample_tokens RPC timeout -> EngineDeadError). Three mitigations, all +# required: (1) --enforce-eager (cudagraph capture itself hung); (2) --length-buckets + +# --warmup right-pad prompts to a few fixed lengths so only a handful of shapes compile, +# pre-compiled up front; (3) env VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 to survive those +# compiles + a persistent TRITON_CACHE_DIR so all parallel tasks reuse the warmed kernels. +# KV cache stays bf16 (M3's MSA fused kernel rejects fp8 KV). NB: VLLM_RPC_TIMEOUT is NOT +# a recognized env var in this vLLM build (silently ignored) — use the EXECUTE_MODEL one. # * Training FakeBaseModel reads lm_head + embed_tokens from the bf16 M3 (real weights; # these tensors are not what MXFP8 quantizes, so dump@MXFP8 / train@bf16 logits stay # consistent). Per Ye Yu: adhere to published bf16/MXFP8 ckpts, do not self-quantize. @@ -55,11 +59,16 @@ pipeline: - --block-size 128 - --language-model-only - --enforce-eager + # Bound distinct prefill shapes -> few MSA kernel recompiles, pre-compiled up front. + - --length-buckets 1024,2048,4096,8192 + - --warmup environment: - HF_MODEL_CKPT: <> - TRUST_REMOTE_CODE: "1" - # Survive MSA Triton-kernel JIT recompiles without an executor RPC timeout. - - VLLM_RPC_TIMEOUT: "1800000" + # Survive the (few, bucketed) ~110s MSA kernel recompiles; default is 300s. + - VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1800" + # Persist compiled MSA Triton kernels so all parallel dump tasks reuse them. + - TRITON_CACHE_DIR: /hf-local/modelopt/MiniMax-M3-DFlash-assets/triton_cache slurm_config: _factory_: "slurm_factory" nodes: 1 From a450ca634b9ce112dd387eb24d079fa4683b49e7 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Mon, 29 Jun 2026 10:19:24 -0700 Subject: [PATCH 08/19] specdec(dump): bump max_model_len by 1 for the dummy max_tokens=1 gen A prompt (or a length bucket) exactly equal to max_seq_len + the dump's dummy 1-token generation = max_seq_len+1, which vLLM rejects against max_model_len=max_seq_len ('decoder prompt plus requested output tokens is longer than the maximum model length'). Set max_model_len = max_seq_len + 1. Surfaced by the M3 dry run when --length-buckets padded prompts to exactly the 8192 top bucket. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- .../collect_hidden_states/compute_hidden_states_vllm.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py index eb1203ab80d..57301d853e6 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py @@ -319,7 +319,10 @@ def keep_conversation(entry): llm = LLM( model=args.model, tensor_parallel_size=tp, - max_model_len=args.max_seq_len, + # +1 for the dummy max_tokens=1 generation: the dump only needs the prefill, but vLLM + # validates prompt_len + output_tokens <= max_model_len, and a prompt (or a length + # bucket) can be exactly max_seq_len. Without the +1, max-length prompts are rejected. + max_model_len=args.max_seq_len + 1, trust_remote_code=args.trust_remote_code, enable_chunked_prefill=False, # required by extract_hidden_states **extra_llm_kwargs, From 7a2b096186940aba37f6188efbaee2d812af4db9 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Mon, 29 Jun 2026 11:39:04 -0700 Subject: [PATCH 09/19] specdec(dump): --disable-triton-autotune to fix MiniMax-M3 MoE-kernel TP deadlock (OMNIML-4747) Dry run with bucketing+warmup got further (MSA attention kernel handled, first real prompt fast) but then deadlocked on a SECOND kernel: M3's MoE expert-routing _topk_index_kernel uses a raw @triton.autotune (5 configs). Its ~150s per-shape benchmark runs UNSYNCHRONIZED across the 8 TP ranks; while some ranks benchmark, others block at the hidden-states connector's shared-memory collective -> permanent desync (22min 'no broadcast block') -> idle-GPU reaper kills it. (Serving never hits this: it's decode-dominated with warm kernels and no connector collective. Confirmed via the reaper alert showing 4/8 GPUs idle = the rank split.) --disable-triton-autotune pins every @triton.autotune to a single (last/most-conservative) config so no benchmark runs -> no cross-rank race. Applied in-process (forked workers) and via a sitecustomize injected on PYTHONPATH (spawned workers), since the autotune executes in vLLM's worker subprocesses. Autotune configs are functionally equivalent (perf-only), so pinning is correct. Wired into the M3 recipe alongside the existing bucketing/warmup/timeout/cache fixes. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- .../compute_hidden_states_vllm.py | 58 +++++++++++++++++++ .../MiniMax-M3-DFlash/hf_offline_dflash.yaml | 23 +++++--- 2 files changed, 72 insertions(+), 9 deletions(-) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py index 57301d853e6..37dd69470fa 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py @@ -28,6 +28,8 @@ import atexit import os import shutil +import tempfile +import textwrap from pathlib import Path import torch @@ -46,6 +48,50 @@ "{% if '' in content %}{% set content = content.split('')[-1] %}{% endif %}" ) +# Source for both the in-process patch and the worker sitecustomize: pin every +# @triton.autotune to a single (last/most-conservative) config so the kernel never runs the +# per-shape benchmark. On MiniMax-M3 that benchmark (~150s for the MoE _topk kernel) runs +# unsynchronized across TP ranks and deadlocks the engine<->worker shared-memory collective. +_NO_AUTOTUNE_SRC = textwrap.dedent( + """ + try: + import triton + _orig_autotune = triton.autotune + def _single_config_autotune(*a, **k): + if k.get("configs"): + k["configs"] = [k["configs"][-1]] + elif a and isinstance(a[0], (list, tuple)) and a[0]: + a = ([a[0][-1]], *a[1:]) + return _orig_autotune(*a, **k) + triton.autotune = _single_config_autotune + try: + import triton.runtime.autotuner as _ta + _ta.autotune = _single_config_autotune + except Exception: + pass + except Exception as _e: + import sys + print(f"[no-autotune] patch failed: {_e}", file=sys.stderr) + """ +) + + +def _disable_triton_autotune() -> None: + """Pin @triton.autotune to a single config in this process AND in vLLM's workers. + + Must run before any vLLM/Triton model module is imported (kernels are decorated at import + time). In-process exec covers forked workers (they inherit patched triton from memory); a + sitecustomize on PYTHONPATH covers spawned workers (imported at interpreter startup, before + vLLM). Keeping the last config (smallest BLOCK_SIZE) is always functionally correct — + autotune configs differ only in performance. + """ + exec(_NO_AUTOTUNE_SRC, {}) + d = Path(tempfile.mkdtemp(prefix="no_autotune_")) + (d / "sitecustomize.py").write_text(_NO_AUTOTUNE_SRC) + os.environ["PYTHONPATH"] = f"{d}{os.pathsep}{os.environ.get('PYTHONPATH', '')}" + atexit.register(shutil.rmtree, d, ignore_errors=True) + print(f"[no-autotune] triton.autotune pinned to a single config (sitecustomize at {d})") + def _resolve_aux_layers_standalone( aux_layers: str, num_hidden_layers: int, num_draft: int = 5 @@ -142,6 +188,14 @@ def parse_args() -> argparse.Namespace: "save; causal attention leaves the real prefix's hidden states unchanged. Empty " "(default) disables bucketing.", ) + parser.add_argument( + "--disable-triton-autotune", + action="store_true", + help="Pin every @triton.autotune to a single config (no per-shape benchmark). Required " + "for MiniMax-M3: the MoE _topk_index_kernel autotune (~150s) runs unsynchronized across " + "TP ranks and deadlocks the engine<->worker collective. Applied in-process + via a " + "sitecustomize on PYTHONPATH so vLLM's worker subprocesses inherit it.", + ) parser.add_argument( "--warmup", action="store_true", @@ -166,6 +220,10 @@ def parse_args() -> argparse.Namespace: def main(args: argparse.Namespace) -> None: + # Must happen before importing vLLM (kernels are @triton.autotune-decorated at import). + if args.disable_triton_autotune: + _disable_triton_autotune() + # Import lazily so --help and arg parsing work without vLLM installed. from vllm import LLM, SamplingParams from vllm.config.kv_transfer import KVTransferConfig diff --git a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml index 88e55a612bc..fa49000ae92 100644 --- a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml +++ b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml @@ -12,15 +12,17 @@ # is not in stable vLLM yet -> image vllm/vllm-openai:minimax-m3. # * --block-size 128 is MANDATORY for M3's MSA sparse attention. # * --language-model-only skips the vision encoder (text-only synth/dump). -# * M3's MSA Triton kernel (_gqa_sparse_fwd_kernel) JIT-recompiles per prefill shape; a -# recompile (~110s) can exceed VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS (default 300) and kill -# the engine (sample_tokens RPC timeout -> EngineDeadError). Three mitigations, all -# required: (1) --enforce-eager (cudagraph capture itself hung); (2) --length-buckets + -# --warmup right-pad prompts to a few fixed lengths so only a handful of shapes compile, -# pre-compiled up front; (3) env VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 to survive those -# compiles + a persistent TRITON_CACHE_DIR so all parallel tasks reuse the warmed kernels. -# KV cache stays bf16 (M3's MSA fused kernel rejects fp8 KV). NB: VLLM_RPC_TIMEOUT is NOT -# a recognized env var in this vLLM build (silently ignored) — use the EXECUTE_MODEL one. +# * M3's MSA/MoE Triton kernels JIT-compile + autotune per prefill shape, which kills the +# dump two ways. Four mitigations, all required: (1) --enforce-eager (cudagraph capture +# itself hung); (2) --length-buckets + --warmup right-pad prompts to a few fixed lengths so +# only a handful of attention shapes compile, pre-compiled up front; (3) env +# VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1800 to survive those compiles + a persistent +# TRITON_CACHE_DIR so all parallel tasks reuse the warmed kernels; (4) --disable-triton- +# autotune — the MoE _topk_index_kernel autotune (~150s) runs UNSYNCHRONIZED across the 8 +# TP ranks and deadlocks the engine<->worker shared-memory collective (idle-GPU reaper then +# kills it); pinning autotune to one config removes the race. KV cache stays bf16 (M3's MSA +# fused kernel rejects fp8 KV). NB: VLLM_RPC_TIMEOUT is NOT recognized in this vLLM build +# (silently ignored) — use VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS. # * Training FakeBaseModel reads lm_head + embed_tokens from the bf16 M3 (real weights; # these tensors are not what MXFP8 quantizes, so dump@MXFP8 / train@bf16 logits stay # consistent). Per Ye Yu: adhere to published bf16/MXFP8 ckpts, do not self-quantize. @@ -59,6 +61,9 @@ pipeline: - --block-size 128 - --language-model-only - --enforce-eager + # Pin @triton.autotune to one config: M3's MoE _topk kernel autotune (~150s) races + # across TP ranks and deadlocks the engine<->worker collective (see header). + - --disable-triton-autotune # Bound distinct prefill shapes -> few MSA kernel recompiles, pre-compiled up front. - --length-buckets 1024,2048,4096,8192 - --warmup From ddc8fcc9a4b732dd72a98a5ace408d76a0f7ba04 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Mon, 29 Jun 2026 13:06:06 -0700 Subject: [PATCH 10/19] specdec(dump): --max-num-seqs 1 to eliminate per-shape JIT compile desync (OMNIML-4747) Dry run confirmed --disable-triton-autotune works (0 autotune benchmarks vs the prior 153s ones), but the dump still deadlocked: with autotune gone, plain per-shape kernel COMPILATION (slot_mapping / index_block_score / topk / gqa_sparse) still runs per-rank and desyncs the TP collective at the first real prompt whose batch shape the warm-up didn't cover. Warm-up synchronizes (all ranks compile the same batch together), so the fix is to make the real loop hit only warmed shapes. Bucketing fixes the sequence-length dimension; max_num_seqs=1 fixes the batch dimension, so a batch=1 bucketed warm-up covers every (1, bucket) shape the real loop sees -> no compilation during the timed dump -> no desync. Wired --max-num-seqs 1 into the M3 recipe. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- .../compute_hidden_states_vllm.py | 12 ++++++++++++ .../MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml | 2 ++ 2 files changed, 14 insertions(+) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py index 37dd69470fa..fae981aecc9 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py @@ -188,6 +188,16 @@ def parse_args() -> argparse.Namespace: "save; causal attention leaves the real prefix's hidden states unchanged. Empty " "(default) disables bucketing.", ) + parser.add_argument( + "--max-num-seqs", + type=int, + default=None, + help="vLLM max concurrent sequences. Set 1 for MiniMax-M3: with bucketing the only " + "varying kernel-shape dimension is then sequence length (a handful of buckets), so a " + "batch=1 warm-up covers every shape the real loop hits and NO kernel compiles during " + "the timed dump. Per-shape compilation otherwise desyncs the TP ranks and deadlocks " + "the engine<->worker collective. Default (None) uses vLLM's default.", + ) parser.add_argument( "--disable-triton-autotune", action="store_true", @@ -373,6 +383,8 @@ def keep_conversation(entry): extra_llm_kwargs["language_model_only"] = True if args.enforce_eager: extra_llm_kwargs["enforce_eager"] = True + if args.max_num_seqs is not None: + extra_llm_kwargs["max_num_seqs"] = args.max_num_seqs llm = LLM( model=args.model, diff --git a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml index fa49000ae92..d78a15ca5f1 100644 --- a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml +++ b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml @@ -64,6 +64,8 @@ pipeline: # Pin @triton.autotune to one config: M3's MoE _topk kernel autotune (~150s) races # across TP ranks and deadlocks the engine<->worker collective (see header). - --disable-triton-autotune + # batch=1 so a bucketed warm-up covers every real-loop shape -> no JIT compile in the timed loop. + - --max-num-seqs 1 # Bound distinct prefill shapes -> few MSA kernel recompiles, pre-compiled up front. - --length-buckets 1024,2048,4096,8192 - --warmup From cc91cce0d2f710b956c7caa461babf58e290450a Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Mon, 29 Jun 2026 15:03:02 -0700 Subject: [PATCH 11/19] specdec(dflash): cast offline bf16 hidden states to model dtype in forward (OMNIML-4747) DFlash offline training crashed in the draft forward with 'mat1 and mat2 must have the same dtype: BFloat16 != float': the dumped hidden states are bf16, but with bf16=false the draft (incl. lm_head + input projections) is fp32. The M3 dump saves no logits, so the self-logit-distillation path recomputes them via lm_head(hidden) -> bf16 input x fp32 weight -> error; target_hidden would hit the same downstream. Cast the offline target_hidden and the lm_head input up to the model dtype (lm_head.weight.dtype); bf16 -> fp32 is lossless. Only the offline path is affected (the online path's hidden states already come from the fp32 forward). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- modelopt/torch/speculative/plugins/hf_dflash.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index 9e54ee531ce..84fc5602dab 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -466,12 +466,19 @@ def forward( if self.dflash_offline: assert "base_model_outputs" in kwargs base_outputs = DFlashBaseModelOutput.from_offline_dict(kwargs["base_model_outputs"]) + # Offline hidden states are dumped in bf16; cast to the model dtype (fp32 when + # bf16=false) so matmuls against fp32 weights — the lm_head below and the draft's + # input projections downstream — don't raise 'mat1 and mat2 must have the same + # dtype'. bf16 -> fp32 is lossless. + model_dtype = self._base_model_lm_head.weight.dtype + if base_outputs.target_hidden is not None: + base_outputs.target_hidden = base_outputs.target_hidden.to(model_dtype) if base_outputs.logits is None and self.dflash_self_logit_distillation: # Compute logits from last-layer hidden states for KD loss. # base_model_hidden_states is required on this path — fail fast # with KeyError rather than lm_head(None). out_hiddens = kwargs["base_model_outputs"]["base_model_hidden_states"] - base_outputs.logits = self._base_model_lm_head(out_hiddens) + base_outputs.logits = self._base_model_lm_head(out_hiddens.to(model_dtype)) target_hidden = base_outputs.target_hidden else: # TODO: For co-training the base model, remove no_grad and eval() switch. From 1e32ec16aa3b73793108ad2db20bfef470ab6b01 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Mon, 29 Jun 2026 15:12:19 -0700 Subject: [PATCH 12/19] specdec(fakebase): resolve base dtype from top-level config for VLMs (OMNIML-4747) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root-cause fix for the offline-training 'mat1 and mat2 must have the same dtype: BFloat16 != float' crash, replacing the earlier symptom cast (now reverted). FakeBaseModel.from_source resolved its dtype from getattr(base_cfg, 'dtype', bf16). For a VLM, base_cfg is the nested text config, which often does NOT set torch_dtype — only the top-level config does (MiniMax-M3: top-level torch_dtype=bfloat16, text_config none). base_cfg.dtype then falls back to PyTorch-default fp32, so lm_head/embed_tokens load fp32 while the dumped hidden states are bf16 -> dtype mismatch in the draft forward. (MiniMax-M2.7 is non-VL, so base_cfg is the top-level config carrying torch_dtype=bfloat16 -> bf16 -> never hit this.) Resolve torch_dtype preferring the nested config, then the top-level, then bf16 (mapping a string to torch.dtype), so the fake base loads in the checkpoint dtype and matches the dump. Reverts the hf_dflash.py offline-hidden-states cast, which masked this by upcasting to fp32. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- .../torch/speculative/plugins/hf_dflash.py | 9 +-------- .../speculative/plugins/modeling_fakebase.py | 20 ++++++++++++++++++- 2 files changed, 20 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py index 84fc5602dab..9e54ee531ce 100644 --- a/modelopt/torch/speculative/plugins/hf_dflash.py +++ b/modelopt/torch/speculative/plugins/hf_dflash.py @@ -466,19 +466,12 @@ def forward( if self.dflash_offline: assert "base_model_outputs" in kwargs base_outputs = DFlashBaseModelOutput.from_offline_dict(kwargs["base_model_outputs"]) - # Offline hidden states are dumped in bf16; cast to the model dtype (fp32 when - # bf16=false) so matmuls against fp32 weights — the lm_head below and the draft's - # input projections downstream — don't raise 'mat1 and mat2 must have the same - # dtype'. bf16 -> fp32 is lossless. - model_dtype = self._base_model_lm_head.weight.dtype - if base_outputs.target_hidden is not None: - base_outputs.target_hidden = base_outputs.target_hidden.to(model_dtype) if base_outputs.logits is None and self.dflash_self_logit_distillation: # Compute logits from last-layer hidden states for KD loss. # base_model_hidden_states is required on this path — fail fast # with KeyError rather than lm_head(None). out_hiddens = kwargs["base_model_outputs"]["base_model_hidden_states"] - base_outputs.logits = self._base_model_lm_head(out_hiddens.to(model_dtype)) + base_outputs.logits = self._base_model_lm_head(out_hiddens) target_hidden = base_outputs.target_hidden else: # TODO: For co-training the base model, remove no_grad and eval() switch. diff --git a/modelopt/torch/speculative/plugins/modeling_fakebase.py b/modelopt/torch/speculative/plugins/modeling_fakebase.py index 17150a8690f..51f34fc31fd 100644 --- a/modelopt/torch/speculative/plugins/modeling_fakebase.py +++ b/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), From 5bfc37fe75b33f8b28986b03070184f1829e755c Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Mon, 29 Jun 2026 15:23:47 -0700 Subject: [PATCH 13/19] specdec(recipe): M3 offline training uses transformers 4.57.1 (validated e2e) Validated the full M3 offline path end-to-end (dump -> train -> export): training must use transformers 4.57.1, not M3's 4.52.4. 4.52.4's Trainer imports apex.amp (absent in the training container) and crashes at import; 4.57.1 guards it and loads M3's config/tokenizer fine. With this + the FakeBase VL-dtype fix, training runs and exports the Qwen3 draft with the full-1M YaRN rope_scaling (factor 128) in bf16. The dump step is unaffected (vllm image). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- .../MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml index d78a15ca5f1..b679791c441 100644 --- a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml +++ b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml @@ -131,9 +131,11 @@ pipeline: environment: - NUM_NODES: "8" # Offline training uses a lightweight FakeBaseModel, so plain DDP suffices (no - # ACCELERATE_CONFIG / FSDP2 patches). OVERRIDE_TRANSFORMERS pins 4.52.4 for the - # MiniMax-M3 config load. - - OVERRIDE_TRANSFORMERS: "4.52.4" + # ACCELERATE_CONFIG / FSDP2 patches). transformers 4.57.1 (not M3's 4.52.4): 4.52.4's + # Trainer does `from apex import amp`, which this training container's apex lacks -> + # import crash; 4.57.1 guards it and still loads M3's config/tokenizer fine. (The dump + # step is unaffected — it runs in the vllm:minimax-m3 image with its own transformers.) + - OVERRIDE_TRANSFORMERS: "4.57.1" - MIXED_PRECISION: "no" slurm_config: _factory_: "slurm_factory" From 4f0b72cdbb900f3acb255da4b6783f2ee3bd7921 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Mon, 29 Jun 2026 15:27:25 -0700 Subject: [PATCH 14/19] test(fakebase): regression for VLM base-dtype resolution (OMNIML-4747) Asserts FakeBaseModel.from_source resolves the base dtype from the top-level config when the VLM's nested text_config carries no torch_dtype (the MiniMax-M3 case): lm_head / embed_tokens must load bf16, not fp32. Without the fix they load fp32 and mismatch the bf16 offline hidden states in the DFlash draft forward. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- .../speculative/plugins/test_fakebase.py | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/tests/unit/torch/speculative/plugins/test_fakebase.py b/tests/unit/torch/speculative/plugins/test_fakebase.py index ff8a2c63074..499712b36b2 100644 --- a/tests/unit/torch/speculative/plugins/test_fakebase.py +++ b/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)) From dafbd8bed06888503a05a5107892d9c937d2686b Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Tue, 30 Jun 2026 10:06:24 -0700 Subject: [PATCH 15/19] specdec(dump): chunk generate+save to bound connector staging (OMNIML-4747) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Full-scale dump filled /dev/shm: the script called llm.generate(ALL prompts) then saved afterward, so every conv's hidden states (~190MB) accumulated in the KV connector's staging dir (DFLASH_HS_STAGING_DIR, default /dev/shm RAM, ~1TB/node) -> 'No space left on device' at ~21%. (The fix stack itself held at scale: all 8 tasks hit 100% at ~2.4 it/s, no deadlock.) Process in chunks of --save-chunk-size (default 256): generate a chunk, save + cleanup each output, then the next chunk — so staging holds at most one chunk (~48GB) at a time. Bonus: saves are now incremental, so a task killed at the 4h limit keeps its completed chunks and keep_conversation resume continues from there. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- .../compute_hidden_states_vllm.py | 130 ++++++++++-------- 1 file changed, 74 insertions(+), 56 deletions(-) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py index fae981aecc9..c4a2c9a4717 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py @@ -206,6 +206,15 @@ def parse_args() -> argparse.Namespace: "TP ranks and deadlocks the engine<->worker collective. Applied in-process + via a " "sitecustomize on PYTHONPATH so vLLM's worker subprocesses inherit it.", ) + parser.add_argument( + "--save-chunk-size", + type=int, + default=256, + help="Generate + save + free staging in chunks of this many conversations. The KV " + "connector stages each conv's hidden states (~hundreds of MB) under DFLASH_HS_STAGING_DIR " + "(default /dev/shm, RAM); generating the whole dataset before saving would accumulate it " + "all and fill /dev/shm. Chunking bounds staging to one chunk at a time.", + ) parser.add_argument( "--warmup", action="store_true", @@ -430,63 +439,72 @@ def keep_conversation(entry): llm.generate(warmup_prompts, SamplingParams(max_tokens=1)) print("Warm-up complete.") - # max_tokens=1: we only need a single forward pass over the prompt tokens. - outputs = llm.generate(prompts, SamplingParams(max_tokens=1)) - - # Save in the same format as compute_hidden_states_hf.py, including loss_mask. + # max_tokens=1: single forward pass per prompt. Process in chunks of --save-chunk-size: + # generate the chunk, then save + cleanup each output before the next chunk. The connector + # stages each conv's hidden states (~hundreds of MB) under DFLASH_HS_STAGING_DIR (default + # /dev/shm, RAM); generating the entire dataset before saving would accumulate it all and + # fill /dev/shm (only ~1 TB on a node). Chunking bounds staging to one chunk at a time. + sampling = SamplingParams(max_tokens=1) + chunk = max(1, args.save_chunk_size) num_success = 0 - for conv_id, loss_mask, real_len, output in tqdm( - zip(conversation_ids, loss_masks, real_lens, outputs), total=len(outputs), desc="Saving" - ): - hidden_states_path = output.kv_transfer_params.get("hidden_states_path") - if hidden_states_path is None: - print(f"WARNING: no hidden_states_path for conversation {conv_id}; skipping") - continue - - obj = example_hidden_states_connector.load_hidden_states(hidden_states_path) - token_ids = obj["token_ids"] - # hidden_states: [num_tokens, num_extracted_layers, hidden_size], ordered to match - # extract_layer_ids. Last layer = final output; the rest = aux layers. - hidden_states = obj["hidden_states"] - - # Drop any right-padding from length-bucketing: keep only the real prefix. Causal - # attention means these positions' states match the unpadded forward. No-op when - # bucketing is disabled (real_len == sequence length). - token_ids = token_ids[:real_len] - hidden_states = hidden_states[:real_len] - - output_hidden_states = hidden_states[:, -1, :].cpu() - if hidden_states.shape[1] > 1: - # Concatenate aux layers along the hidden dim, matching the HF dump format. - aux = hidden_states[:, :-1, :].cpu() - aux_hidden_states = aux.reshape(aux.shape[0], -1) - else: - aux_hidden_states = torch.empty(0) - - # loss_mask is sliced to the dumped length below; a shorter loss_mask would slice - # to itself and silently misalign with the hidden states, so guard explicitly. - n_hs = output_hidden_states.shape[0] - if loss_mask.shape[0] < n_hs: - print( - f"WARNING: {conv_id}: loss_mask ({loss_mask.shape[0]}) shorter than hidden " - f"states ({n_hs}); skipping to avoid misalignment" - ) - continue - - output_file = output_dir / f"{conv_id}.pt" - with open(output_file, "wb") as f: - torch.save( - { - "input_ids": token_ids.cpu(), - "hidden_states": output_hidden_states, - "aux_hidden_states": aux_hidden_states, - "loss_mask": loss_mask[: output_hidden_states.shape[0]].cpu(), - "conversation_id": conv_id, - }, - f, - ) - example_hidden_states_connector.cleanup_hidden_states(hidden_states_path) - num_success += 1 + pbar = tqdm(total=len(prompts), desc="Dumping") + for start in range(0, len(prompts), chunk): + sl = slice(start, start + chunk) + outputs = llm.generate(prompts[sl], sampling) + for conv_id, loss_mask, real_len, output in zip( + conversation_ids[sl], loss_masks[sl], real_lens[sl], outputs + ): + pbar.update(1) + hidden_states_path = output.kv_transfer_params.get("hidden_states_path") + if hidden_states_path is None: + print(f"WARNING: no hidden_states_path for conversation {conv_id}; skipping") + continue + + obj = example_hidden_states_connector.load_hidden_states(hidden_states_path) + token_ids = obj["token_ids"] + # hidden_states: [num_tokens, num_extracted_layers, hidden_size], ordered to match + # extract_layer_ids. Last layer = final output; the rest = aux layers. + hidden_states = obj["hidden_states"] + + # Drop any right-padding from length-bucketing: keep only the real prefix. Causal + # attention means these positions' states match the unpadded forward. No-op when + # bucketing is disabled (real_len == sequence length). + token_ids = token_ids[:real_len] + hidden_states = hidden_states[:real_len] + + output_hidden_states = hidden_states[:, -1, :].cpu() + if hidden_states.shape[1] > 1: + # Concatenate aux layers along the hidden dim, matching the HF dump format. + aux = hidden_states[:, :-1, :].cpu() + aux_hidden_states = aux.reshape(aux.shape[0], -1) + else: + aux_hidden_states = torch.empty(0) + + # loss_mask is sliced to the dumped length below; a shorter loss_mask would slice + # to itself and silently misalign with the hidden states, so guard explicitly. + n_hs = output_hidden_states.shape[0] + if loss_mask.shape[0] < n_hs: + print( + f"WARNING: {conv_id}: loss_mask ({loss_mask.shape[0]}) shorter than hidden " + f"states ({n_hs}); skipping to avoid misalignment" + ) + continue + + output_file = output_dir / f"{conv_id}.pt" + with open(output_file, "wb") as f: + torch.save( + { + "input_ids": token_ids.cpu(), + "hidden_states": output_hidden_states, + "aux_hidden_states": aux_hidden_states, + "loss_mask": loss_mask[: output_hidden_states.shape[0]].cpu(), + "conversation_id": conv_id, + }, + f, + ) + example_hidden_states_connector.cleanup_hidden_states(hidden_states_path) + num_success += 1 + pbar.close() print(f"Successfully processed {num_success} out of {len(prompts)} conversations.") From 140947bc6f9998ba15cf3cac79f2788f7399e041 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Tue, 30 Jun 2026 10:44:59 -0700 Subject: [PATCH 16/19] specdec(recipe): leaner M3 offline training (2 nodes x bs=4, not 8 x bs=1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Offline training is lightweight — FakeBaseModel loads only lm_head + embed_tokens (not the 427B base), the draft is a small 5-layer Qwen3, and dflash_num_anchors=512 caps the vocab-logit memory to 512 positions rather than all 8192. So the 8-node x bs=1 config (inherited from M2.7) over-provisions GPUs. Use 2 nodes x 8 GPU x bs=4 = the same effective batch (64) at 1/4 the GPUs; bs is tunable upward given the anchor-bounded logits. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- .../MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml index b679791c441..c40f0831187 100644 --- a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml +++ b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml @@ -97,8 +97,12 @@ pipeline: - data.chat_template=examples/MiniMax/MiniMax-M3-DFlash/chat_template_train.jinja - training.output_dir=/scratchspace/dflash_minimax_m3_offline - training.num_train_epochs=10 - # bs=1 @ 8192 keeps the activation footprint of M2.7's bs=2 @ 4096 (bs*seqlen equal). - - training.per_device_train_batch_size=1 + # Offline training is light (FakeBaseModel = lm_head + embed only; small 5-layer draft; + # dflash_num_anchors=512 caps the vocab-logit memory to 512 positions, not all 8192), so + # we use a higher per-device batch and far fewer nodes than M2.7's 8x bs=1. 2 nodes x 8 + # GPU x bs=4 = effective batch 64 (same as the old 8x8x1), at 1/4 the GPUs. bs is tunable + # upward (anchor-bounded logits leave headroom); drop if a larger draft/seq OOMs. + - training.per_device_train_batch_size=4 - training.learning_rate=1.2e-3 - training.warmup_steps=100 - training.training_seq_len=8192 @@ -129,7 +133,7 @@ pipeline: - dflash.dflash_export_rope_scaling.mscale=1.0 - dflash.dflash_export_rope_scaling.mscale_all_dim=1.0 environment: - - NUM_NODES: "8" + - NUM_NODES: "2" # Offline training uses a lightweight FakeBaseModel, so plain DDP suffices (no # ACCELERATE_CONFIG / FSDP2 patches). transformers 4.57.1 (not M3's 4.52.4): 4.52.4's # Trainer does `from apex import amp`, which this training container's apex lacks -> @@ -139,6 +143,6 @@ pipeline: - MIXED_PRECISION: "no" slurm_config: _factory_: "slurm_factory" - nodes: 8 + nodes: 2 ntasks_per_node: 1 gpus_per_node: 8 From 55e86ca11ff3c535b8f8a649f2bbf3a8a2737d48 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Wed, 1 Jul 2026 19:46:56 -0700 Subject: [PATCH 17/19] specdec(dump): force filter re-eval so resume actually skips done convs (OMNIML-4747) The skip-existing filter (keep_conversation, which checks whether each conv's .pt already exists) was being served from datasets' cache: filter/map results are fingerprinted on the function + dataset, NOT on external disk state. With a persistent HF cache across requeues, the cached 'keep everything' from an early run (few/no .pt) was reused, so resume reported 'Removed 0' and re-dumped/overwrote already-done conversations (observed: 62k .pt rewritten in 2h, total flat). Pass load_from_cache_file=False so the filter re-checks the disk every run. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- .../collect_hidden_states/compute_hidden_states_vllm.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py index c4a2c9a4717..d5197bc7201 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py @@ -275,7 +275,13 @@ def keep_conversation(entry): return not (output_dir / f"{conversation_id}.pt").exists() original_num = len(dataset) - dataset = dataset.filter(keep_conversation) + # load_from_cache_file=False is REQUIRED for resume correctness: keep_conversation depends + # on on-disk .pt state, which is NOT part of datasets' cache fingerprint (only the function + # + dataset are). With a persistent HF cache (e.g. HF_HOME on shared storage across + # requeues), a cached result from an earlier run — when fewer/no .pt existed — is reused, + # so the filter reports "Removed 0" and every resume re-dumps + overwrites already-done + # conversations instead of skipping them. Forcing re-evaluation re-checks the disk each run. + dataset = dataset.filter(keep_conversation, load_from_cache_file=False) print(f"Removed {original_num - len(dataset)} conversations due to existing output files") if args.debug_max_num_conversations is not None: From c6a1caa5f2ac9001fc70f67a20a593b6a8e2dd86 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Thu, 2 Jul 2026 13:48:58 -0700 Subject: [PATCH 18/19] specdec(recipe): M3 offline training bs=2 x grad_accum=2 (bs=4 OOMs) Full-scale training OOM'd at bs=4 on 80GB (73.7GB in use, +6.1GB requested): M3's offline inputs are activation-heavy (6144 hidden x 5 aux layers x 8192 seq) even though the FakeBase + 5-layer draft are small and dflash_num_anchors=512 bounds the logits. bs=2 x grad_accum=2 keeps effective batch 64 (16 ranks / 2 nodes) and fits. Also set PYTORCH_CUDA_ALLOC_CONF= expandable_segments:True to reduce fragmentation. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- .../MiniMax-M3-DFlash/hf_offline_dflash.yaml | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml index c40f0831187..924d0ea57a5 100644 --- a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml +++ b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml @@ -97,12 +97,12 @@ pipeline: - data.chat_template=examples/MiniMax/MiniMax-M3-DFlash/chat_template_train.jinja - training.output_dir=/scratchspace/dflash_minimax_m3_offline - training.num_train_epochs=10 - # Offline training is light (FakeBaseModel = lm_head + embed only; small 5-layer draft; - # dflash_num_anchors=512 caps the vocab-logit memory to 512 positions, not all 8192), so - # we use a higher per-device batch and far fewer nodes than M2.7's 8x bs=1. 2 nodes x 8 - # GPU x bs=4 = effective batch 64 (same as the old 8x8x1), at 1/4 the GPUs. bs is tunable - # upward (anchor-bounded logits leave headroom); drop if a larger draft/seq OOMs. - - training.per_device_train_batch_size=4 + # Offline training uses the lightweight FakeBaseModel (lm_head + embed only) + a small + # 5-layer draft, so far fewer nodes than M2.7's 8x suffice. But bs is still bounded by the + # activation-heavy offline inputs (M3's 6144 hidden x 5 aux layers x 8192 seq): bs=4 OOMs + # on 80 GB. bs=2 x grad_accum=2 = effective batch 64 (16 ranks over 2 nodes), memory-safe. + - training.per_device_train_batch_size=2 + - training.gradient_accumulation_steps=2 - training.learning_rate=1.2e-3 - training.warmup_steps=100 - training.training_seq_len=8192 @@ -134,6 +134,7 @@ pipeline: - dflash.dflash_export_rope_scaling.mscale_all_dim=1.0 environment: - NUM_NODES: "2" + - PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True" # Offline training uses a lightweight FakeBaseModel, so plain DDP suffices (no # ACCELERATE_CONFIG / FSDP2 patches). transformers 4.57.1 (not M3's 4.52.4): 4.52.4's # Trainer does `from apex import amp`, which this training container's apex lacks -> From f41a49c0e16b38ff5512abf5cdfc944cab7ae50e Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Tue, 7 Jul 2026 13:48:34 -0700 Subject: [PATCH 19/19] specdec(recipe): M3 offline training on 1 node (2-node C10d rendezvous unstable) The 2-node training repeatedly crashed on cw-dfw with C10d rendezvous failures (torch.distributed DistNetworkError ~every 20-60min); slurm --requeue doesn't recover it since it's an app-level NCCL/rendezvous crash, not a node failure. Offline training is light enough for a single node anyway (FakeBaseModel + 5-layer draft), which eliminates cross-node rendezvous entirely. 1 node x 8 GPU x bs=2 x grad_accum=4 = effective batch 64 (same as before), stable + hands-off with per-window self-requeue (main.py resumes from checkpoints). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Ye Yu --- .../MiniMax-M3-DFlash/hf_offline_dflash.yaml | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml index 924d0ea57a5..36736e0c741 100644 --- a/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml +++ b/tools/launcher/examples/MiniMax/MiniMax-M3-DFlash/hf_offline_dflash.yaml @@ -98,11 +98,13 @@ pipeline: - training.output_dir=/scratchspace/dflash_minimax_m3_offline - training.num_train_epochs=10 # Offline training uses the lightweight FakeBaseModel (lm_head + embed only) + a small - # 5-layer draft, so far fewer nodes than M2.7's 8x suffice. But bs is still bounded by the - # activation-heavy offline inputs (M3's 6144 hidden x 5 aux layers x 8192 seq): bs=4 OOMs - # on 80 GB. bs=2 x grad_accum=2 = effective batch 64 (16 ranks over 2 nodes), memory-safe. + # 5-layer draft, so a SINGLE node suffices (vs M2.7's 8x). Single-node also avoids the + # cross-node C10d rendezvous that repeatedly crashed the 2-node run on cw-dfw + # (DistNetworkError ~every 20-60min; slurm --requeue doesn't catch app-level NCCL crashes). + # bs is bounded by the activation-heavy offline inputs (6144 hidden x 5 aux layers x 8192 + # seq): bs=4 OOMs on 80 GB. bs=2 x grad_accum=4 on 1 node x 8 GPU = effective batch 64. - training.per_device_train_batch_size=2 - - training.gradient_accumulation_steps=2 + - training.gradient_accumulation_steps=4 - training.learning_rate=1.2e-3 - training.warmup_steps=100 - training.training_seq_len=8192 @@ -133,7 +135,7 @@ pipeline: - dflash.dflash_export_rope_scaling.mscale=1.0 - dflash.dflash_export_rope_scaling.mscale_all_dim=1.0 environment: - - NUM_NODES: "2" + - NUM_NODES: "1" - PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True" # Offline training uses a lightweight FakeBaseModel, so plain DDP suffices (no # ACCELERATE_CONFIG / FSDP2 patches). transformers 4.57.1 (not M3's 4.52.4): 4.52.4's @@ -144,6 +146,6 @@ pipeline: - MIXED_PRECISION: "no" slurm_config: _factory_: "slurm_factory" - nodes: 2 + nodes: 1 ntasks_per_node: 1 gpus_per_node: 8