Skip to content

Minimax m3 weight free export - #1317

Open
amarquic wants to merge 3 commits into
quic:mainfrom
amarquic:minimax-m3-weight-free-export
Open

Minimax m3 weight free export#1317
amarquic wants to merge 3 commits into
quic:mainfrom
amarquic:minimax-m3-weight-free-export

Conversation

@amarquic

@amarquic amarquic commented Sep 9, 2026

Copy link
Copy Markdown
Contributor

Fixes weight-free (meta-device) torch.export/torch.onnx.export for MiniMax-M3's language decoder via QEFFAutoModelForImageTextToText. Each fix was
root-caused from a real export/compile failure surfaced by iteratively running examples/text_generation/minimax_m3_decode_only.py against the real
MiniMax-M3 checkpoint (and, for the final dtype-leak issue, a standalone repro harness that exercises the full 60-layer config directly —
weight-free export only needs the config at trace time, not real weights).

  • checkpoint_key_resolver.py: resolve block_sparse_moe.e_score_correction_bias — the checkpoint stores this buffer one level shallower (directly on
    the MoE block) than the live module tree (nested under mlp.gate/block_sparse_moe.gate).
  • modeling_minimax_m3_vl.py:
    • Avoid duplicate nn.Module registration of the shared MoEWeights instance (object.setattr instead of a plain attribute assignment), which
      otherwise made torch.export emit duplicate gate/up/down initializers under two FQNs — only one of which promote_initializers_and_build_spec would
      promote, leaving the other as an un-promoted meta tensor that fails at ONNX save().
    • Avoid device-bound lifted scalar constants in two shape[1] == torch.tensor(1, device=...) comparisons — the meta-device input made these
      dataless lifted constants that fail ONNX serialization.
    • Replace _qeff_minimax_clamp's torch.maximum/torch.minimum against a torch.tensor(..., device=hidden_states.device) constant with
      torch.clamp(min=, max=), which accepts plain Python scalars directly.
    • Cast hidden_states to lm_head.weight.dtype consistently at all four lm_head call sites (two were missing the cast the other two already had),
      fixing a float16/float32 Equal/matmul dtype mismatch the QAIC compiler rejected.
    • Cast the sigmoid router's top_k_weights back to the router's native (float16) dtype before returning. routing_weights is intentionally upcast to
      float32 for numerically-stable sigmoid scoring (matching upstream HF), but the derived top_k_weights was never cast back down, so moe_decode_bmm's
      down * topk_weights silently promoted the MoE block's output to float32. This was invisible with --num-layers 4 (the MoE layer was the last layer
      traced, and the lm_head cast mopped it up), but fatal on the real 60-layer model as soon as a later float16-weighted layer consumed the float32
      activation.
  • blocked_attention_forwards.py: in blocked_kv_attention_forward_headpar_offline (the function MiniMax's blocking_mode="kv_headpar" dispatches to),
    compare start_index as a plain Python int rather than wrapping it in torch.tensor(start_index, device=query.device). The wrapped form traces to a
    device-bound lifted tensor placeholder that has no data under meta-device tracing; the plain-int comparison traces to a scalar-comparison op with
    the constant baked in as an attribute instead. Scoped to only this function, not the ~9 other occurrences of the same pattern elsewhere in the file
    that back other (non-MiniMax) blocking modes.

kdulla and others added 3 commits September 3, 2026 18:45
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
Signed-off-by: Kushal Dulla <kdulla@qti.qualcomm.com>
…oder

Adds the checkpoint-key resolution, duplicate-submodule, lifted-constant,
and dtype-promotion fixes needed to trace and export MiniMax-M3's decoder
under torch.export/torch.onnx.export with weight_free=True (meta-device
parameters, no real weights loaded at trace time):

- checkpoint_key_resolver: resolve MiniMax's flattened
  block_sparse_moe.e_score_correction_bias checkpoint key (buffer sits one
  level shallower in the checkpoint than in the live module tree).
- modeling_minimax_m3_vl: avoid duplicate nn.Module registration of
  MoEWeights (object.__setattr__ instead of plain assignment); avoid
  device-bound lifted scalar constants in shape/threshold comparisons;
  cast lm_head input to lm_head.weight.dtype consistently across all
  forward paths; use torch.clamp instead of torch.maximum/minimum against
  a device-bound constant; cast the sigmoid-router's top_k_weights back
  to the router's native dtype so a float32 MoE output doesn't leak into
  later float16-weighted layers (only visible once more than one MoE
  layer runs in sequence).
- blocked_attention_forwards: compare start_index as a plain Python int
  (not a device-bound tensor) in blocked_kv_attention_forward_headpar_offline,
  scoped to the function MiniMax's kv_headpar blocking mode dispatches to.

Validated by iteratively running
examples/text_generation/minimax_m3_decode_only.py against the real
MiniMax-M3 checkpoint and a local reproduction harness exercising the
full 60-layer config without requiring real weights (weight-free export
only needs the config at trace time).

Signed-off-by: amarshar <amarshar@qti.qualcomm.com>
@amarquic
amarquic force-pushed the minimax-m3-weight-free-export branch from faac0f3 to 33a08dd Compare September 11, 2026 07:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants