diff --git a/miles/backends/megatron_utils/megatron_to_hf/__init__.py b/miles/backends/megatron_utils/megatron_to_hf/__init__.py index 1525c939c0e..25e8b5b2f9c 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/__init__.py +++ b/miles/backends/megatron_utils/megatron_to_hf/__init__.py @@ -9,6 +9,7 @@ from .kimi_vl import convert_kimi_k25_to_hf, convert_kimivl_to_hf from .llama import convert_llama_to_hf from .mimo import convert_mimo_to_hf +from .minimax_m3 import convert_minimax_m3_to_hf from .processors import quantize_params, remove_padding from .qwen2 import convert_qwen2_to_hf from .qwen3_5 import convert_qwen3_5_to_hf @@ -77,6 +78,8 @@ def _convert_to_hf_core(args, model_name, name, param): converted_named_tensors = convert_inkling_to_hf(args, name, param) elif "llama" in model_name: converted_named_tensors = convert_llama_to_hf(args.backend, name, param) + elif "minimax_m3" in model_name: + converted_named_tensors = convert_minimax_m3_to_hf(args.backend, name, param) elif "mimo" in model_name: converted_named_tensors = convert_mimo_to_hf(args.backend, name, param) elif "kimi_k3" in model_name: diff --git a/miles/backends/megatron_utils/megatron_to_hf/minimax_m3.py b/miles/backends/megatron_utils/megatron_to_hf/minimax_m3.py new file mode 100644 index 00000000000..1d7255f20e8 --- /dev/null +++ b/miles/backends/megatron_utils/megatron_to_hf/minimax_m3.py @@ -0,0 +1,71 @@ +import re + +import torch + +_TEXT_MODEL = "language_model.model" + +_DIRECT = { + "module.module.embedding.word_embeddings.weight": f"{_TEXT_MODEL}.embed_tokens.weight", + "module.module.output_layer.weight": "language_model.lm_head.weight", + "module.module.decoder.final_layernorm.weight": f"{_TEXT_MODEL}.norm.weight", +} + +_LAYER_RENAMES = { + "input_layernorm.weight": "input_layernorm.weight", + "pre_mlp_layernorm.weight": "post_attention_layernorm.weight", + "mlp.linear_fc1.layer_norm_weight": "post_attention_layernorm.weight", + "self_attention.linear_proj.weight": "self_attn.o_proj.weight", + "self_attention.q_layernorm.weight": "self_attn.q_norm.weight", + "self_attention.k_layernorm.weight": "self_attn.k_norm.weight", + "self_attention.indexer.index_q_proj.weight": "self_attn.index_q_proj.weight", + "self_attention.indexer.index_k_proj.weight": "self_attn.index_k_proj.weight", + "self_attention.indexer.index_q_norm": "self_attn.index_q_norm.weight", + "self_attention.indexer.index_k_norm": "self_attn.index_k_norm.weight", + "mlp.linear_fc2.weight": "mlp.down_proj.weight", + "mlp.router.weight": "block_sparse_moe.gate.weight", + "mlp.router.expert_bias": "block_sparse_moe.e_score_correction_bias", + "mlp.shared_experts.linear_fc2.weight": "block_sparse_moe.shared_experts.down_proj.weight", +} + + +def convert_minimax_m3_to_hf(args, name, param): + if name in _DIRECT: + return [(_DIRECT[name], param)] + + match = re.match(r"module\.module\.decoder\.layers\.(\d+)\.(.+)", name) + if match is None: + raise ValueError(f"Unknown parameter name: {name}") + layer_idx, rest = match.groups() + prefix = f"{_TEXT_MODEL}.layers.{layer_idx}" + + expert = re.match(r"mlp\.experts\.(linear_fc1|linear_fc2)\.weight(\d+)", rest) + if expert: + kind, expert_idx = expert.groups() + base = f"{prefix}.block_sparse_moe.experts.{expert_idx}" + if kind == "linear_fc1": + gate, up = param.chunk(2, dim=0) + return [(f"{base}.w1.weight", gate), (f"{base}.w3.weight", up)] + return [(f"{base}.w2.weight", param)] + + if rest == "mlp.shared_experts.linear_fc1.weight": + gate, up = param.chunk(2, dim=0) + return [ + (f"{prefix}.block_sparse_moe.shared_experts.gate_proj.weight", gate), + (f"{prefix}.block_sparse_moe.shared_experts.up_proj.weight", up), + ] + if rest == "mlp.linear_fc1.weight": + gate, up = param.chunk(2, dim=0) + return [(f"{prefix}.mlp.gate_proj.weight", gate), (f"{prefix}.mlp.up_proj.weight", up)] + if rest == "self_attention.linear_qkv.weight": + head_dim = args.kv_channels + per_group = args.num_attention_heads // args.num_query_groups + param = param.view(args.num_query_groups, -1, head_dim, args.hidden_size) + q, k, v = torch.split(param, [per_group, 1, 1], dim=1) + return [ + (f"{prefix}.self_attn.q_proj.weight", q.reshape(-1, args.hidden_size)), + (f"{prefix}.self_attn.k_proj.weight", k.reshape(-1, args.hidden_size)), + (f"{prefix}.self_attn.v_proj.weight", v.reshape(-1, args.hidden_size)), + ] + if rest in _LAYER_RENAMES: + return [(f"{prefix}.{_LAYER_RENAMES[rest]}", param)] + raise ValueError(f"Unknown parameter name: {name}") diff --git a/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py b/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py index 2fa7f4ee390..6777d46915f 100644 --- a/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py +++ b/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py @@ -402,7 +402,7 @@ def _check_and_fix_partition(args: Namespace, name: str, partition_stride: int, (GLU/SwiGLU interleaved [gate, up]), so assert partition_stride==2 is removed. But TEGroupedLinear still does not set partition_stride/partition_dim correctly for grouped moe gemm """ - if "linear_fc1.weight" in name and args.backend.swiglu: + if "linear_fc1.weight" in name and (args.backend.swiglu or args.backend.quick_geglu): partition_stride = 2 if partition_dim < 0: partition_dim = 0 diff --git a/miles/kernels/attention/msa/__init__.py b/miles/kernels/attention/msa/__init__.py new file mode 100644 index 00000000000..8f6f0ddece9 --- /dev/null +++ b/miles/kernels/attention/msa/__init__.py @@ -0,0 +1,3 @@ +from miles.kernels.attention.msa.block_max_scores import msa_block_max_scores + +__all__ = ["msa_block_max_scores"] diff --git a/miles/kernels/attention/msa/block_max_scores.py b/miles/kernels/attention/msa/block_max_scores.py new file mode 100644 index 00000000000..fe5dd7d4643 --- /dev/null +++ b/miles/kernels/attention/msa/block_max_scores.py @@ -0,0 +1,83 @@ +"""MSA block scores for training: each query's maximum logit over every causal key block of its sequence.""" + +import torch +import triton +import triton.language as tl +from torch import Tensor + +_BQ = 64 + + +@triton.jit +def _block_max_scores_kernel( + Q, + K, + TILES, + OUT, + stride_qt, + stride_qh, + stride_kt, + stride_oh, + stride_ot, + scale, + D: tl.constexpr, + BQ: tl.constexpr, + BLK: tl.constexpr, +): + tile = tl.program_id(0) + head = tl.program_id(1) + q_start = tl.load(TILES + tile * 3) + seq_start = tl.load(TILES + tile * 3 + 1) + seq_end = tl.load(TILES + tile * 3 + 2) + + offs_q = q_start + tl.arange(0, BQ) + offs_d = tl.arange(0, D) + q_in = offs_q < seq_end + q = tl.load(Q + offs_q[:, None] * stride_qt + head * stride_qh + offs_d[None, :], mask=q_in[:, None], other=0.0) + + # a tile never straddles sequences, so every query shares the key blocks; the last query bounds them + last_block = (tl.minimum(q_start + BQ, seq_end) - 1 - seq_start) // BLK + for block in range(0, last_block + 1): + offs_k = seq_start + block * BLK + tl.arange(0, BLK) + k_in = offs_k < seq_end + k = tl.load(K + offs_k[:, None] * stride_kt + offs_d[None, :], mask=k_in[:, None], other=0.0) + logits = tl.dot(q, tl.trans(k)) + causal = (offs_k[None, :] <= offs_q[:, None]) & k_in[None, :] + best = tl.max(tl.where(causal, logits, float("-inf")), axis=1) * scale + tl.store(OUT + head * stride_oh + offs_q * stride_ot + block, best, mask=q_in) + + +def msa_block_max_scores(q: Tensor, k: Tensor, cu_seqlens: Tensor, block_size: int, scale: float) -> Tensor: + """``q`` [T, H, D], ``k`` [T, D] -> ``[H, T, NB]`` fp32 block scores, ``-inf`` where no key is causal. + + ``NB`` is the block count of the longest sequence; blocks start at each sequence's first token. + """ + T, H, D = q.shape + bounds = cu_seqlens.tolist() + tiles = [ + (q_start, start, end) + for start, end in zip(bounds[:-1], bounds[1:], strict=True) + for q_start in range(start, end, _BQ) + ] + num_blocks = max(-(-(end - start) // block_size) for start, end in zip(bounds[:-1], bounds[1:], strict=True)) + tiles = torch.tensor(tiles, dtype=torch.int32, device=q.device) + out = torch.full((H, T, num_blocks), float("-inf"), dtype=torch.float32, device=q.device) + qc, kc = q.contiguous(), k.contiguous() + _block_max_scores_kernel[(tiles.shape[0], H)]( + qc, + kc, + tiles, + out, + qc.stride(0), + qc.stride(1), + kc.stride(0), + out.stride(0), + out.stride(1), + scale, + D=D, + BQ=_BQ, + BLK=block_size, + num_warps=4, + num_stages=2, + ) + return out diff --git a/miles/kernels/attention/qsa/block_sparse_attention.py b/miles/kernels/attention/qsa/block_sparse_attention.py index 150fd57d45e..86b28ecb9c3 100644 --- a/miles/kernels/attention/qsa/block_sparse_attention.py +++ b/miles/kernels/attention/qsa/block_sparse_attention.py @@ -7,6 +7,9 @@ inclusive key range, sequence start .. query position) removes the extra tokens. """ +import math +from dataclasses import dataclass + import torch import triton import triton.language as tl @@ -246,10 +249,12 @@ def _qsa_bs_dkdv_kernel( stride_ql, stride_ot, stride_oh, + stride_split, T, NB, scale, GROUP: tl.constexpr, + HEADS_PER_SPLIT: tl.constexpr, D: tl.constexpr, BQ: tl.constexpr, BK: tl.constexpr, @@ -257,11 +262,13 @@ def _qsa_bs_dkdv_kernel( ): """dK/dV keyed on the key tile, so nothing needs atomics. - Launched per KV head: the ``GROUP`` query heads sharing a KV head are summed here, - since separate programs would overwrite each other's dK/dV. + Each program sums ``HEADS_PER_SPLIT`` of the ``GROUP`` query heads sharing a KV head into + its own slice of ``DK``/``DV``; the caller adds the slices. """ pid_k = tl.program_id(0) - kv_head = tl.program_id(1) + num_splits: tl.constexpr = GROUP // HEADS_PER_SPLIT + kv_head = tl.program_id(1) // num_splits + split = tl.program_id(1) % num_splits offs_k = pid_k * BK + tl.arange(0, BK) offs_d = tl.arange(0, D) @@ -300,8 +307,8 @@ def _qsa_bs_dkdv_kernel( & q_mask[:, None] ) - for gh in range(0, GROUP): - qh = kv_head * GROUP + gh + for gh in range(0, HEADS_PER_SPLIT): + qh = kv_head * GROUP + split * HEADS_PER_SPLIT + gh q = tl.load( Q + offs_q[:, None] * stride_qt + qh * stride_qh + offs_d[None, :], mask=q_mask[:, None], @@ -327,8 +334,10 @@ def _qsa_bs_dkdv_kernel( dk += tl.dot(tl.trans(ds), q) dv += tl.dot(tl.trans(p.to(do.dtype)), do) - tl.store(DK + offs_k[:, None] * stride_kt + kv_head * stride_kh + offs_d[None, :], dk, mask=k_in[:, None]) - tl.store(DV + offs_k[:, None] * stride_vt + kv_head * stride_vh + offs_d[None, :], dv, mask=k_in[:, None]) + dk_offs = split * stride_split + offs_k[:, None] * stride_kt + kv_head * stride_kh + offs_d[None, :] + dv_offs = split * stride_split + offs_k[:, None] * stride_vt + kv_head * stride_vh + offs_d[None, :] + tl.store(DK + dk_offs, dk, mask=k_in[:, None]) + tl.store(DV + dv_offs, dv, mask=k_in[:, None]) def selection_to_block_bitmap(indices: Tensor, num_tokens: int, block_size: int) -> Tensor: @@ -347,7 +356,7 @@ def selection_to_block_bitmap(indices: Tensor, num_tokens: int, block_size: int) return flags -def _tile_mask(sel: Tensor, block_first: Tensor, block_last: Tensor, bq: int, bk: int) -> Tensor: +def _tile_mask(sel: Tensor, block_first: Tensor, block_last: Tensor, bq: int, bk: int, block_size: int) -> Tensor: """``[NQT, NKT]`` bool: some query in query tile ``i`` selected a block with a token in key tile ``j``. Key tiles are absolute token ranges ``[j * bk, (j + 1) * bk)``, as the kernels walk them. @@ -362,9 +371,11 @@ def _tile_mask(sel: Tensor, block_first: Tensor, block_last: Tensor, bq: int, bk sel = torch.nn.functional.pad(sel, (0, 0, 0, pad_q)) picked = sel.view(nqt, bq, nb).amax(dim=1).to(torch.int32) tile = torch.zeros(nqt, nkt, dtype=torch.int32, device=sel.device) - # block_size <= bk, so a block touches at most the tiles of its first and last token - tile.index_add_(1, block_first.long() // bk, picked) - tile.index_add_(1, block_last.long() // bk, picked) + # a block wider than a tile, or off the tile grid, touches every tile between its first and last token + first_tile = block_first.long() // bk + last_tile = block_last.long() // bk + for off in range(-(-block_size // bk) + 1): + tile.index_add_(1, torch.minimum(first_tile + off, last_tile), picked) return tile > 0 @@ -376,27 +387,48 @@ def _compact_tile_mask(mat: Tensor) -> tuple[Tensor, Tensor]: return order[:, :maxc].contiguous().to(torch.int32), cnt -def build_tile_index(sel: Tensor, block_first: Tensor, block_last: Tensor, bq: int, bk: int) -> tuple[Tensor, Tensor]: +def build_tile_index( + sel: Tensor, block_first: Tensor, block_last: Tensor, bq: int, bk: int, block_size: int +) -> tuple[Tensor, Tensor]: """``[T, NB]`` block flags -> (``klist`` [NQT, maxc] int32, ``kcnt`` [NQT] int32). ``klist[i]`` lists, ascending, the key tiles that at least one query in query-tile ``i`` selected. Per-query exactness still comes from the in-kernel mask; this only decides which tiles are worth visiting. """ - return _compact_tile_mask(_tile_mask(sel, block_first, block_last, bq, bk)) + return _compact_tile_mask(_tile_mask(sel, block_first, block_last, bq, bk, block_size)) -def build_tile_index_pair(sel: Tensor, block_first: Tensor, block_last: Tensor, bq: int, bk: int): +def build_tile_index_pair(sel: Tensor, block_first: Tensor, block_last: Tensor, bq: int, bk: int, block_size: int): """Both directions of the tile map: (klist, kcnt) per query tile, (qlist, qcnt) per key tile. The transposed half is what lets dK/dV be keyed on the key tile and so avoid atomics. """ - tile = _tile_mask(sel, block_first, block_last, bq, bk) + tile = _tile_mask(sel, block_first, block_last, bq, bk, block_size) klist, kcnt = _compact_tile_mask(tile) qlist, qcnt = _compact_tile_mask(tile.t().contiguous()) return klist, kcnt, qlist, qcnt +@dataclass(frozen=True) +class _LaunchConfig: + bq: int + bk: int + num_warps: int + num_stages: int + # programs sharing one KV head's query heads in the dK/dV kernel + dkdv_splits: int = 1 + + +# (forward, backward), measured on GB300 with qwen3.8-next's 4-token QSA blocks and MiniMax-M3's 128-token MSA blocks +_NARROW_BLOCK = (_LaunchConfig(64, 64, 8, 2), _LaunchConfig(128, 64, 8, 1, dkdv_splits=2)) +_WIDE_BLOCK = (_LaunchConfig(64, 128, 4, 1), _LaunchConfig(64, 128, 8, 2, dkdv_splits=4)) + + +def _launch_configs(block_size: int) -> tuple[_LaunchConfig, _LaunchConfig]: + return _WIDE_BLOCK if block_size >= 64 else _NARROW_BLOCK + + class _QSABlockSparseAttn(torch.autograd.Function): @staticmethod def forward(ctx, q, k, v, sel, lo, hi, blk_base, tok_base, block_first, block_last, scale, block_size): @@ -406,11 +438,11 @@ def forward(ctx, q, k, v, sel, lo, hi, blk_base, tok_base, block_first, block_la group = Hq // Hkv qc, kc, vc = q.contiguous(), k.contiguous(), v.contiguous() selc = sel.contiguous() - BQ, BK = 64, 64 - klist, kcnt = build_tile_index(selc, block_first, block_last, BQ, BK) + cfg, _ = _launch_configs(block_size) + klist, kcnt = build_tile_index(selc, block_first, block_last, cfg.bq, cfg.bk, block_size) o = torch.empty(T, Hq, D, device=q.device, dtype=torch.float32) lse = torch.empty(Hq, T, device=q.device, dtype=torch.float32) - grid = (triton.cdiv(T, BQ), Hq) + grid = (triton.cdiv(T, cfg.bq), Hq) _qsa_bs_fwd_kernel[grid]( qc, kc, @@ -439,11 +471,11 @@ def forward(ctx, q, k, v, sel, lo, hi, blk_base, tok_base, block_first, block_la scale, GROUP=group, D=D, - BQ=BQ, - BK=BK, + BQ=cfg.bq, + BK=cfg.bk, BLK=block_size, - num_warps=8, - num_stages=2, + num_warps=cfg.num_warps, + num_stages=cfg.num_stages, ) ctx.save_for_backward(qc, kc, vc, selc, lo, hi, blk_base, tok_base, block_first, block_last, o, lse) ctx.scale = scale @@ -458,12 +490,14 @@ def backward(ctx, grad_out): do = grad_out.contiguous().to(qc.dtype) # delta once, in torch: both backward kernels need it and neither should redo it delta = (do.float() * o.float()).sum(-1).transpose(0, 1).contiguous() + _, cfg = _launch_configs(ctx.block_size) + num_kv_tiles = triton.cdiv(T, cfg.bk) + splits = math.gcd(cfg.dkdv_splits, ctx.group) dq = torch.empty(T, Hq, D, device=qc.device, dtype=torch.float32) - dk = torch.zeros(kc.shape, device=kc.device, dtype=torch.float32) - dv = torch.zeros(vc.shape, device=vc.device, dtype=torch.float32) + dk = torch.empty(splits, *kc.shape, device=kc.device, dtype=torch.float32) + dv = torch.empty(splits, *vc.shape, device=vc.device, dtype=torch.float32) - BQ, BK = 64, 32 - klist, kcnt, qlist, qcnt = build_tile_index_pair(selc, block_first, block_last, BQ, BK) + klist, kcnt, qlist, qcnt = build_tile_index_pair(selc, block_first, block_last, cfg.bq, cfg.bk, ctx.block_size) common = ( qc.stride(0), qc.stride(1), @@ -473,7 +507,7 @@ def backward(ctx, grad_out): vc.stride(1), selc.stride(0), ) - _qsa_bs_dq_kernel[(triton.cdiv(T, BQ), Hq)]( + _qsa_bs_dq_kernel[(triton.cdiv(T, cfg.bq), Hq)]( qc, kc, vc, @@ -498,13 +532,13 @@ def backward(ctx, grad_out): ctx.scale, GROUP=ctx.group, D=D, - BQ=BQ, - BK=BK, + BQ=cfg.bq, + BK=cfg.bk, BLK=ctx.block_size, - num_warps=8, - num_stages=1, + num_warps=cfg.num_warps, + num_stages=cfg.num_stages, ) - _qsa_bs_dkdv_kernel[(triton.cdiv(T, BK), kc.shape[1])]( + _qsa_bs_dkdv_kernel[(num_kv_tiles, kc.shape[1] * splits)]( qc, kc, vc, @@ -524,17 +558,20 @@ def backward(ctx, grad_out): qlist.stride(0), do.stride(0), do.stride(1), + dk.stride(0), T, selc.shape[1], ctx.scale, GROUP=ctx.group, + HEADS_PER_SPLIT=ctx.group // splits, D=D, - BQ=BQ, - BK=BK, + BQ=cfg.bq, + BK=cfg.bk, BLK=ctx.block_size, - num_warps=8, - num_stages=1, + num_warps=cfg.num_warps, + num_stages=cfg.num_stages, ) + dk, dv = dk.sum(dim=0), dv.sum(dim=0) return dq.to(qc.dtype), dk.to(kc.dtype), dv.to(vc.dtype), None, None, None, None, None, None, None, None, None diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 1256b33b6fa..b5fea8bcc6f 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1857,12 +1857,17 @@ def equal(x, y): if (hf_config.model_type == "deepseek_v4" or "deepseekv4" in model_name) and args.context_parallel_size > 1: assert args.allgather_cp, "zigzag CP is not supported for DeepSeek V4." + # MiniMax-M3 names the dense width dense_intermediate_size and the expert width intermediate_size + if hasattr(hf_config, "dense_intermediate_size"): # config-access-exempt: model-specific key + dense_ffn_key, moe_ffn_key = "dense_intermediate_size", "intermediate_size" + else: + dense_ffn_key, moe_ffn_key = "intermediate_size", "moe_intermediate_size" for hf_config_name, megatron_config_name, compare_fn in [ ("hidden_size", "hidden_size", equal), ("num_attention_heads", "num_attention_heads", equal), ("num_hidden_layers", "num_layers", equal), - ("intermediate_size", "ffn_hidden_size", equal), - ("moe_intermediate_size", "moe_ffn_hidden_size", equal), + (dense_ffn_key, "ffn_hidden_size", equal), + (moe_ffn_key, "moe_ffn_hidden_size", equal), ("tie_word_embeddings", "untie_embeddings_and_output_weights", lambda x, y: not x == y), ( "rms_norm_eps", diff --git a/miles/utils/hf_utils/config.py b/miles/utils/hf_utils/config.py index 1638d5496fb..de31db2870c 100644 --- a/miles/utils/hf_utils/config.py +++ b/miles/utils/hf_utils/config.py @@ -15,7 +15,7 @@ from dataclasses import dataclass from pathlib import Path -from transformers import AutoConfig, AutoModelForCausalLM +from transformers import AutoConfig, AutoModelForCausalLM, PretrainedConfig from transformers.models.auto.configuration_auto import CONFIG_MAPPING_NAMES @@ -130,6 +130,9 @@ def load_hf_config( """ register_hf_config_aliases() config = AutoConfig.from_pretrained(checkpoint_path, trust_remote_code=trust_remote_code, **autoconfig_kwargs) + # MiniMax-M3's remote config class keeps the nested text config as a plain dict + if isinstance(getattr(config, "text_config", None), dict): # config-access-exempt: optional key + config.text_config = PretrainedConfig.from_dict(config.text_config) if overrides: for key, value in overrides.items(): diff --git a/miles_plugins/mbridge/__init__.py b/miles_plugins/mbridge/__init__.py index 1c132dde1e9..ebbff3ab3b5 100644 --- a/miles_plugins/mbridge/__init__.py +++ b/miles_plugins/mbridge/__init__.py @@ -11,6 +11,7 @@ from .kimi_k3 import KimiK3Bridge from .kimi_k25 import KimiK25Bridge from .mimo import MimoBridge +from .minimax_m3 import MiniMaxM3Bridge from .qwen3_5 import Qwen3_5Bridge from .qwen3_8_next import Qwen38NextBridge from .qwen3_next import Qwen3NextBridge @@ -23,6 +24,7 @@ "Qwen3_5Bridge", "Qwen38NextBridge", "MimoBridge", + "MiniMaxM3Bridge", "DeepseekV32Bridge", "Glm5NextBridge", "DeepseekV4Bridge", diff --git a/miles_plugins/mbridge/minimax_m3.py b/miles_plugins/mbridge/minimax_m3.py new file mode 100644 index 00000000000..7cbe5ca7886 --- /dev/null +++ b/miles_plugins/mbridge/minimax_m3.py @@ -0,0 +1,111 @@ +"""Bridge for MiniMax-M3 (HF ``MiniMaxM3SparseForConditionalGeneration``, text part only). + +Text weights live under ``language_model.``; the vision tower and projector are never +created on the Megatron side. +""" + +import torch +from transformers import PretrainedConfig + +from mbridge.core import LLMBridge, register_model +from miles_plugins.models.minimax_m3.minimax_m3 import apply_minimax_m3_config + +_LAYER = "language_model.model.layers.{layer_number}" + + +@register_model(["minimax_m3_vl", "minimax_m3"]) +class MiniMaxM3Bridge(LLMBridge): + _CONFIG_MAPPING = {**LLMBridge._CONFIG_MAPPING, "attention_dropout": ("attention_dropout", 0.0)} + + _DIRECT_MAPPING = { + "embedding.word_embeddings.weight": "language_model.model.embed_tokens.weight", + "decoder.final_layernorm.weight": "language_model.model.norm.weight", + "output_layer.weight": "language_model.lm_head.weight", + } + + _ATTENTION_MAPPING = { + "self_attention.linear_proj.weight": [f"{_LAYER}.self_attn.o_proj.weight"], + "self_attention.q_layernorm.weight": [f"{_LAYER}.self_attn.q_norm.weight"], + "self_attention.k_layernorm.weight": [f"{_LAYER}.self_attn.k_norm.weight"], + "self_attention.linear_qkv.weight": [ + f"{_LAYER}.self_attn.q_proj.weight", + f"{_LAYER}.self_attn.k_proj.weight", + f"{_LAYER}.self_attn.v_proj.weight", + ], + "self_attention.indexer.index_q_proj.weight": [f"{_LAYER}.self_attn.index_q_proj.weight"], + "self_attention.indexer.index_k_proj.weight": [f"{_LAYER}.self_attn.index_k_proj.weight"], + "self_attention.indexer.index_q_norm": [f"{_LAYER}.self_attn.index_q_norm.weight"], + "self_attention.indexer.index_k_norm": [f"{_LAYER}.self_attn.index_k_norm.weight"], + } + + # first match wins, so the per-expert and shared-expert keys precede the dense ones + _MLP_MAPPING = { + "mlp.experts.linear_fc1": [ + f"{_LAYER}.block_sparse_moe.experts.{{expert_id}}.w1.weight", + f"{_LAYER}.block_sparse_moe.experts.{{expert_id}}.w3.weight", + ], + "mlp.experts.linear_fc2": [f"{_LAYER}.block_sparse_moe.experts.{{expert_id}}.w2.weight"], + "mlp.shared_experts.linear_fc1.weight": [ + f"{_LAYER}.block_sparse_moe.shared_experts.gate_proj.weight", + f"{_LAYER}.block_sparse_moe.shared_experts.up_proj.weight", + ], + "mlp.shared_experts.linear_fc2.weight": [f"{_LAYER}.block_sparse_moe.shared_experts.down_proj.weight"], + "mlp.router.expert_bias": [f"{_LAYER}.block_sparse_moe.e_score_correction_bias"], + "mlp.router.weight": [f"{_LAYER}.block_sparse_moe.gate.weight"], + "pre_mlp_layernorm.weight": [f"{_LAYER}.post_attention_layernorm.weight"], + "mlp.linear_fc1.layer_norm_weight": [f"{_LAYER}.post_attention_layernorm.weight"], + "mlp.linear_fc1.weight": [f"{_LAYER}.mlp.gate_proj.weight", f"{_LAYER}.mlp.up_proj.weight"], + "mlp.linear_fc2.weight": [f"{_LAYER}.mlp.down_proj.weight"], + } + + _OTHER_MAPPING = { + "input_layernorm.weight": [f"{_LAYER}.input_layernorm.weight"], + } + + def __init__(self, hf_config, *args, **kwargs): + text_config = getattr(hf_config, "text_config", None) or hf_config # config-access-exempt: optional key + # the remote config class keeps the nested text config as a plain dict + if isinstance(text_config, dict): + text_config = PretrainedConfig.from_dict(text_config) + super().__init__(text_config, *args, **kwargs) + + def _weight_name_mapping_mlp(self, name: str) -> list[str]: + layer_number = name.split(".")[2] + for keyword, targets in self._MLP_MAPPING.items(): + if keyword in name: + expert_id = name.split("weight")[-1] if "{expert_id}" in targets[0] else None + return [t.format(layer_number=layer_number, expert_id=expert_id) for t in targets] + raise NotImplementedError(f"Unsupported parameter name: {name}") + + def _weight_to_mcore_format(self, mcore_weights_name: str, hf_weights: list[torch.Tensor]) -> torch.Tensor: + if mcore_weights_name.endswith("mlp.router.expert_bias"): + return hf_weights[0].to(torch.float32) + return super()._weight_to_mcore_format(mcore_weights_name, hf_weights) + + def _build_config(self): + cfg = self.hf_config + config = self._build_base_config( + use_cpu_initialization=False, + num_moe_experts=cfg.num_local_experts, + moe_ffn_hidden_size=cfg.intermediate_size, + ffn_hidden_size=cfg.dense_intermediate_size, + moe_shared_expert_intermediate_size=cfg.shared_intermediate_size * cfg.n_shared_experts, + moe_router_topk=cfg.num_experts_per_tok, + moe_router_score_function="sigmoid", + moe_router_enable_expert_bias=True, + moe_router_bias_update_rate=0.0, + moe_router_topk_scaling_factor=cfg.routed_scaling_factor, + moe_router_load_balancing_type="none", + moe_router_dtype="fp32", + moe_grouped_gemm=True, + moe_layer_freq=list(cfg.moe_layer_freq), + qk_layernorm=True, + layernorm_zero_centered_gamma=True, + add_qkv_bias=False, + add_bias_linear=False, + kv_channels=cfg.head_dim, + rotary_interleaved=False, + ) + apply_minimax_m3_config(config, cfg) + config.activation_func_clamp_value = cfg.swiglu_limit + return config diff --git a/miles_plugins/models/minimax_m3/__init__.py b/miles_plugins/models/minimax_m3/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/miles_plugins/models/minimax_m3/attention.py b/miles_plugins/models/minimax_m3/attention.py new file mode 100644 index 00000000000..8a19cab178c --- /dev/null +++ b/miles_plugins/models/minimax_m3/attention.py @@ -0,0 +1,243 @@ +"""MiniMax-M3 self-attention: GQA restricted to the MSA indexer's top-k key blocks. + +Index head ``h`` scores each causal key block by its max logit and picks, with the +query's own block forced in, the blocks kv head ``h`` attends to. +""" + +import torch +from megatron.core import parallel_state +from megatron.core.extensions.transformer_engine import TELinear +from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region +from megatron.core.transformer.attention import SelfAttention +from megatron.core.transformer.module import MegatronModule +from torch import Tensor + +from miles.kernels.attention.msa import msa_block_max_scores +from miles.kernels.attention.qsa.block_sparse_attention import qsa_block_sparse_attention +from miles_plugins.models.qwen3_8_next.ops.qsa_indexer import PackedBlockLayout, packed_positions + +# topk priority: init blocks, then local blocks, then the best scores; future blocks never +_INIT_BLOCK_SCORE = 1e30 +_LOCAL_BLOCK_SCORE = 1e29 +_INVALID_BLOCK_SCORE = -1e30 + + +def gemma_rmsnorm_fp32(x: Tensor, weight: Tensor, eps: float) -> Tensor: + xf = x.float() + xf = xf * torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return xf * (1.0 + weight.float()) + + +def apply_neox_rope_fp32(x: Tensor, rotary_pos_emb: Tensor, positions: Tensor) -> Tensor: + """Neox rope on the leading ``rotary_dim`` dims of fp32 ``x`` ``[T, heads, head_dim]``.""" + freqs = rotary_pos_emb.flatten(1) + half = freqs.shape[-1] // 2 + freqs = freqs[:, :half].float().contiguous() + cos = freqs.cos()[positions.long()].unsqueeze(-2) + sin = freqs.sin()[positions.long()].unsqueeze(-2) + x1 = x[..., :half] + x2 = x[..., half : 2 * half] + return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin, x[..., 2 * half :]], dim=-1) + + +class MiniMaxM3HeadNorm(torch.nn.Module): + """Per-head gemma RMSNorm that stays in fp32, so rope runs in fp32 and q/k round to bf16 once as in sglang.""" + + def __init__(self, config, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(hidden_size, dtype=config.params_dtype)) + self.weight.sequence_parallel = config.sequence_parallel + self.eps = eps + + def forward(self, x: Tensor) -> Tensor: + return gemma_rmsnorm_fp32(x, self.weight, self.eps) + + +class MiniMaxM3Indexer(MegatronModule): + """Scores key blocks per index head and returns each query's selected block set.""" + + def __init__(self, config): + super().__init__(config) + self.n_heads = config.minimax_index_n_heads + self.head_dim = config.minimax_index_head_dim + self.topk_blocks = config.minimax_index_topk_blocks + self.block_size = config.minimax_index_block_size + self.local_blocks = config.minimax_index_local_blocks + self.init_blocks = config.minimax_index_init_blocks + self.eps = config.layernorm_epsilon + self.index_q_proj = TELinear( + config.hidden_size, + self.n_heads * self.head_dim, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + ) + self.index_k_proj = TELinear( + config.hidden_size, + self.head_dim, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + ) + self.index_q_norm = torch.nn.Parameter(torch.zeros(self.head_dim, dtype=config.params_dtype)) + self.index_k_norm = torch.nn.Parameter(torch.zeros(self.head_dim, dtype=config.params_dtype)) + for p in self.parameters(): + p.requires_grad_(False) + + @torch.no_grad() + def project(self, hidden_states: Tensor, positions: Tensor, rotary_pos_emb: Tensor, heads: list[int]): + """``[T, hidden] -> (q [T, len(heads), D], k [T, D])`` in bf16, normed then roped in fp32.""" + q, _ = self.index_q_proj(hidden_states) + k, _ = self.index_k_proj(hidden_states) + q = q.view(-1, self.n_heads, self.head_dim)[:, heads] + q = gemma_rmsnorm_fp32(q, self.index_q_norm, self.eps) + k = gemma_rmsnorm_fp32(k.view(-1, 1, self.head_dim), self.index_k_norm, self.eps) + q = apply_neox_rope_fp32(q, rotary_pos_emb, positions).to(hidden_states.dtype) + k = apply_neox_rope_fp32(k, rotary_pos_emb, positions).to(hidden_states.dtype) + return q, k.squeeze(1) + + @torch.no_grad() + def select(self, q: Tensor, k: Tensor, cu_seqlens: Tensor, positions: Tensor) -> Tensor: + """Selected blocks as packed start positions ``[H, T, topk_blocks]`` int64, -1 padded.""" + T, H = q.shape[:2] + bs = self.block_size + scores = msa_block_max_scores(q, k, cu_seqlens, bs, q.shape[-1] ** -0.5) + block_ids = torch.arange(scores.shape[-1], device=q.device) + valid_blocks = positions.long() // bs + 1 + valid = block_ids[None, :] < valid_blocks[:, None] + local = valid & (block_ids[None, :] >= (valid_blocks - self.local_blocks).clamp(min=0)[:, None]) + init = valid & (block_ids[None, :] < self.init_blocks) + scores = torch.where(valid, scores, _INVALID_BLOCK_SCORE) + scores = torch.where(local, _LOCAL_BLOCK_SCORE, scores) + scores = torch.where(init, _INIT_BLOCK_SCORE, scores) + kk = min(self.topk_blocks, scores.shape[-1]) + top_score, top_idx = torch.topk(scores, kk, dim=-1) + seq_start = torch.arange(T, device=q.device) - positions.long() + selected = torch.full((H, T, self.topk_blocks), -1, dtype=torch.int64, device=q.device) + selected[..., :kk] = torch.where(top_score > _INVALID_BLOCK_SCORE, seq_start[None, :, None] + top_idx * bs, -1) + return selected + + +def selection_to_flags(selected: Tensor, layout: PackedBlockLayout) -> Tensor: + """``[H, T, K]`` packed block-start positions -> ``[H, T, NB]`` uint8 flags on the packed grid.""" + H, T, _ = selected.shape + flags = torch.zeros(H, T, layout.num_blocks, dtype=torch.uint8, device=selected.device) + keep = selected >= 0 + heads = torch.arange(H, device=selected.device).view(-1, 1, 1).expand_as(selected) + rows = torch.arange(T, device=selected.device).view(1, -1, 1).expand_as(selected) + flags[heads[keep], rows[keep], layout.token_block[selected[keep]]] = 1 + return flags + + +class MiniMaxM3SparseCoreAttention(MegatronModule): + """Core attention restricted, per kv head, to the blocks its index head selected.""" + + def __init__(self, config, owner): + super().__init__(config) + object.__setattr__(self, "_owner", owner) + self.softmax_scale = config.kv_channels**-0.5 + + def forward(self, query: Tensor, key: Tensor, value: Tensor, attention_mask=None, **kwargs): + block_form = self._owner._msa_block_form + assert block_form is not None, "MiniMax-M3 core attention ran without a published block selection" + assert query.dim() == 3, f"MiniMax-M3 sparse attention expects thd inputs, got {tuple(query.shape)}" + flags, lo, hi, blk_base, tok_base, block_first, block_last, block_size = block_form + query, key = query.to(value.dtype), key.to(value.dtype) + T, hq, d = query.shape + hkv = key.shape[1] + group = hq // hkv + outs = [] + for h in range(hkv): + outs.append( + qsa_block_sparse_attention( + query[:, h * group : (h + 1) * group], + key[:, h : h + 1], + value[:, h : h + 1], + flags[h], + lo, + hi, + blk_base, + tok_base, + block_first, + block_last, + self.softmax_scale, + block_size, + ) + ) + return torch.cat(outs, dim=1).reshape(T, hq * d) + + +class MiniMaxM3DenseCoreAttention(MegatronModule): + """Megatron's dense core attention, fed the fp32 q/k that the head norm and rope leave.""" + + def __init__(self, config, core_attention): + super().__init__(config) + self.inner = core_attention + + def forward(self, query: Tensor, key: Tensor, value: Tensor, *args, **kwargs): + return self.inner(query.to(value.dtype), key.to(value.dtype), value, *args, **kwargs) + + +class MiniMaxM3Attention(SelfAttention): + """Megatron GQA self-attention; on sparse layers restricted to the MiniMax (MSA) indexer's blocks.""" + + def __init__(self, config, submodules, layer_number=1, *args, sparse: bool = True, **kwargs): + super().__init__(config, submodules, layer_number, *args, **kwargs) + if config.rotary_interleaved: + raise NotImplementedError("MiniMax-M3 uses neox rope; rotary_interleaved is unsupported") + self.sparse = sparse + self._msa_block_form = None + if not sparse: + self.core_attention = MiniMaxM3DenseCoreAttention(config, self.core_attention) + return + if self.checkpoint_core_attention: + raise NotImplementedError("MiniMax-M3 sparse attention does not support selective core_attn recompute") + self.indexer = MiniMaxM3Indexer(config) + self.core_attention = MiniMaxM3SparseCoreAttention(config, owner=self) + + def forward(self, hidden_states: Tensor, *args, **kwargs): + if not self.sparse: + return super().forward(hidden_states, *args, **kwargs) + if parallel_state.get_context_parallel_world_size() > 1: + raise NotImplementedError("MiniMax-M3 sparse attention does not support context parallelism yet") + packed = kwargs.get("packed_seq_params") + assert packed is not None and packed.cu_seqlens_q is not None, "MiniMax-M3 sparse attention needs thd inputs" + rotary_pos_emb = kwargs.get("rotary_pos_emb") + assert rotary_pos_emb is not None, "the MSA indexer reuses the attention rope table" + if isinstance(rotary_pos_emb, tuple): + rotary_pos_emb = rotary_pos_emb[0] + + with torch.no_grad(): + indexer_states = hidden_states + if self.config.sequence_parallel and parallel_state.get_tensor_model_parallel_world_size() > 1: + indexer_states = gather_from_sequence_parallel_region(hidden_states, tensor_parallel_output_grad=False) + seq = indexer_states.shape[0] + cu = packed.cu_seqlens_q + positions = packed_positions(cu, seq) + tp_rank = parallel_state.get_tensor_model_parallel_rank() + kv_heads_per_rank = self.config.num_query_groups // parallel_state.get_tensor_model_parallel_world_size() + local_heads = list(range(tp_rank * kv_heads_per_rank, (tp_rank + 1) * kv_heads_per_rank)) + q, k = self.indexer.project(indexer_states[:, 0], positions, rotary_pos_emb, local_heads) + selected = self.indexer.select(q, k, cu, positions) + layout = PackedBlockLayout(cu, positions, self.indexer.block_size) + token_start = layout.token_start.to(torch.int32) + self._msa_block_form = ( + selection_to_flags(selected, layout), + token_start, + torch.arange(seq, dtype=torch.int32, device=positions.device), + layout.token_block_start.to(torch.int32), + token_start, + layout.block_first_token.to(torch.int32), + layout.block_last_token.to(torch.int32), + self.indexer.block_size, + ) + try: + return super().forward(hidden_states, *args, **kwargs) + finally: + self._msa_block_form = None diff --git a/miles_plugins/models/minimax_m3/minimax_m3.py b/miles_plugins/models/minimax_m3/minimax_m3.py new file mode 100644 index 00000000000..277fd4fd2b4 --- /dev/null +++ b/miles_plugins/models/minimax_m3/minimax_m3.py @@ -0,0 +1,62 @@ +"""MiniMax-M3 block spec: the GPT decoder with MSA sparse self-attention on the sparse layers. + +The pre-attention norm is kept as its own module (not fused into ``linear_qkv``) because +the indexer reads the normed hidden states too. +""" + +import copy + +from megatron.core.extensions.transformer_engine import TEColumnParallelLinear, TENorm +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_block import get_num_layers_to_build +from megatron.core.transformer.transformer_layer import get_transformer_layer_offset + +from miles.utils.hf_utils.config import load_hf_config +from miles_plugins.models.minimax_m3.attention import MiniMaxM3Attention, MiniMaxM3HeadNorm + + +def apply_minimax_m3_config(config, text_config) -> None: + sparse = text_config.sparse_attention_config + assert sparse["use_sparse_attention"] and sparse["sparse_score_type"] == "max" + config.minimax_sparse_layers = [bool(flag) for flag in sparse["sparse_attention_freq"]] + if not all( + disabled + for is_sparse, disabled in zip(config.minimax_sparse_layers, sparse["sparse_disable_index_value"], strict=True) + if is_sparse + ): + raise NotImplementedError("MiniMax-M3 sparse layers with an index value branch are not supported") + config.minimax_index_n_heads = sparse["sparse_num_index_heads"] + config.minimax_index_head_dim = sparse["sparse_index_dim"] + config.minimax_index_topk_blocks = sparse["sparse_topk_blocks"] + config.minimax_index_block_size = sparse["sparse_block_size"] + config.minimax_index_local_blocks = sparse["sparse_local_block"] + config.minimax_index_init_blocks = sparse["sparse_init_block"] + assert config.minimax_index_n_heads == config.num_query_groups, "index head h selects for kv head h" + # swigluoai: clamp(g, max=7) * sigmoid(1.702 g) * (clamp(u, -7, 7) + 1) + config.glu_linear_offset = 1.0 + + +def get_minimax_m3_spec(args, config, vp_stage=None): + apply_minimax_m3_config(config, load_hf_config(args.hf_checkpoint).text_config) + + kwargs = {"use_transformer_engine": True} + if vp_stage is not None: + kwargs["vp_stage"] = vp_stage + block_spec = get_gpt_decoder_block_spec(config, **kwargs) + + offset = get_transformer_layer_offset(config, vp_stage=vp_stage) + for layer_id in range(get_num_layers_to_build(config, vp_stage=vp_stage)): + layer_spec = copy.deepcopy(block_spec.layer_specs[layer_id]) + attn = layer_spec.submodules.self_attention + attn.submodules.linear_qkv = TEColumnParallelLinear + attn.submodules.q_layernorm = MiniMaxM3HeadNorm + attn.submodules.k_layernorm = MiniMaxM3HeadNorm + layer_spec.submodules.input_layernorm = TENorm + layer_spec.submodules.self_attention = ModuleSpec( + module=MiniMaxM3Attention, + params={**(attn.params or {}), "sparse": config.minimax_sparse_layers[offset + layer_id]}, + submodules=attn.submodules, + ) + block_spec.layer_specs[layer_id] = layer_spec + return block_spec diff --git a/miles_plugins/models/qwen3_8_next/ops/attention.py b/miles_plugins/models/qwen3_8_next/ops/attention.py index 49f57630a8d..32835dfa60a 100644 --- a/miles_plugins/models/qwen3_8_next/ops/attention.py +++ b/miles_plugins/models/qwen3_8_next/ops/attention.py @@ -12,7 +12,7 @@ from torch import Tensor from miles.kernels.attention.qsa import qsa_block_sparse_attention, qsa_sparse_attention -from miles_plugins.models.qwen3_8_next.ops.qsa_indexer import PackedBlockLayout, Qwen38NextQSAIndexer +from miles_plugins.models.qwen3_8_next.ops.qsa_indexer import PackedBlockLayout, Qwen38NextQSAIndexer, packed_positions class Qwen38NextQSACoreAttention(MegatronModule): @@ -85,16 +85,6 @@ def __init__(self, config, submodules, layer_number=1, *args, **kwargs): self._qsa_selection = None self._qsa_block_form = None - @staticmethod - def _packed_positions(cu_seqlens: Tensor, total: int) -> Tensor: - """Positions restarting at 0 for each sequence in a packed batch.""" - idx = torch.arange(total, device=cu_seqlens.device) - starts = cu_seqlens[:-1].long() - seg = torch.zeros(total, dtype=torch.long, device=cu_seqlens.device) - seg[starts[1:]] = 1 - seg = seg.cumsum(0) - return idx - starts[seg] - def forward(self, hidden_states: Tensor, *args, **kwargs): packed = kwargs.get("packed_seq_params") @@ -121,7 +111,7 @@ def forward(self, hidden_states: Tensor, *args, **kwargs): "document." ) self._qsa_cu_seqlens = cu - positions = self._packed_positions(cu, seq) + positions = packed_positions(cu, seq) else: self._qsa_cu_seqlens = None positions = torch.arange(seq, device=hidden_states.device) diff --git a/miles_plugins/models/qwen3_8_next/ops/qsa_indexer.py b/miles_plugins/models/qwen3_8_next/ops/qsa_indexer.py index 050caffc33c..9634a3cf4c6 100644 --- a/miles_plugins/models/qwen3_8_next/ops/qsa_indexer.py +++ b/miles_plugins/models/qwen3_8_next/ops/qsa_indexer.py @@ -65,6 +65,16 @@ def block_causal_mask(query_positions: Tensor, num_blocks: int, compress_ratio: return blocks.unsqueeze(0) < first_invalid.unsqueeze(1) +def packed_positions(cu_seqlens: Tensor, total: int) -> Tensor: + """Positions restarting at 0 for each sequence in a packed batch.""" + idx = torch.arange(total, device=cu_seqlens.device) + starts = cu_seqlens[:-1].long() + seg = torch.zeros(total, dtype=torch.long, device=cu_seqlens.device) + seg[starts[1:]] = 1 + seg = seg.cumsum(0) + return idx - starts[seg] + + class PackedBlockLayout: """Per-sequence compressed-block grid for a packed (thd) batch. diff --git a/scripts/models/minimax-m3-4layer.py b/scripts/models/minimax-m3-4layer.py new file mode 100644 index 00000000000..9ec9183159f --- /dev/null +++ b/scripts/models/minimax-m3-4layer.py @@ -0,0 +1,5 @@ +from model_args_utils import load_sibling_model_args + + +def model_args() -> str: + return load_sibling_model_args(__file__, "minimax-m3", nlayers=4, first_k_dense_replace=2) diff --git a/scripts/models/minimax-m3.py b/scripts/models/minimax-m3.py new file mode 100644 index 00000000000..8f41d465efb --- /dev/null +++ b/scripts/models/minimax-m3.py @@ -0,0 +1,41 @@ +from model_args_utils import moe_layer_freq + + +def model_args(nlayers: int = 60, first_k_dense_replace: int = 3, rotary_base: int = 5000000) -> str: + return ( + "--spec miles_plugins.models.minimax_m3.minimax_m3 get_minimax_m3_spec " + f"--num-layers {nlayers} " + "--hidden-size 6144 " + "--ffn-hidden-size 12288 " + "--num-attention-heads 64 " + "--group-query-attention " + "--num-query-groups 4 " + "--kv-channels 128 " + "--qk-layernorm " + f"--moe-layer-freq {moe_layer_freq(nlayers=nlayers, first_k_dense_replace=first_k_dense_replace)} " + "--num-experts 128 " + "--moe-router-topk 4 " + "--moe-router-score-function sigmoid " + "--moe-router-enable-expert-bias " + "--moe-router-bias-update-rate 0 " + "--moe-router-topk-scaling-factor 2.0 " + "--moe-ffn-hidden-size 3072 " + "--moe-shared-expert-intermediate-size 3072 " + "--moe-router-dtype fp32 " + "--moe-grouped-gemm " + "--moe-permute-fusion " + "--vocab-size 200064 " + f"--rotary-base {rotary_base} " + "--rotary-percent 0.5 " + "--position-embedding-type rope " + "--no-rope-fusion " + "--normalization RMSNorm " + "--norm-epsilon 1e-6 " + "--apply-layernorm-1p " + "--quick-geglu " + "--activation-func-clamp-value 7 " + "--disable-bias-linear " + "--untie-embeddings-and-output-weights " + "--make-vocab-size-divisible-by 64 " + "--enable-experimental " + ) diff --git a/scripts/run_minimax_m3.py b/scripts/run_minimax_m3.py new file mode 100644 index 00000000000..798dc0ee681 --- /dev/null +++ b/scripts/run_minimax_m3.py @@ -0,0 +1,256 @@ +"""MiniMax-M3 DAPO training (bf16 train, bf16 rollout). + +`prepare` downloads the release to `/` and, unless `/` is +that download, converts it with tools/convert_minimax_m3_to_bf16.py. It then converts the model into +the Megatron `/_torch_dist` used as the reference model, unless they +already exist: + MiniMax-M3 the full model (854 GB), 8 nodes x 4 GPUs + MiniMax-M3-4layer `--layers 0,1,3,4`: two dense-attention dense-MLP and two sparse MoE layers, 4 GPUs + +Args: + --model-name: checkpoint directory under `--model-dir`; selects the matching recipe. + --rollout-max-response-len: per-reply generation budget. + --check-weight-update-equal: bitwise weight-sync check; defaults on for the 4-layer model only. + --enable-r3: rollout routing replay for the MoE router. + --train-offload-dir: disk directory for the actor offloaded while the engines generate and the + streamed Adam state. + +The full model expects MILES_SCRIPT_EXTERNAL_RAY=1 with all 8 nodes already joined to the ray cluster. + +python scripts/run_minimax_m3.py --model-name MiniMax-M3 --num-nodes 8 +python scripts/run_minimax_m3.py --model-name MiniMax-M3-4layer --num-nodes 1 +""" + +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + +import typer + +import miles.utils.external_utils.command_utils as U + + +@dataclass(frozen=True) +class _Recipe: + hf_repo: str + megatron_model_type: str + num_gpus: int + # TP / PP / EP of the trainer. + parallel: tuple[int, int, int] + rollout_num_gpus_per_engine: int + # Source decoder layers kept by tools/convert_minimax_m3_to_bf16.py; None keeps all. + layers: str | None = None + # ~13B routed-expert params per GPU at EP8/PP4: fp32 master + Adam fit neither the GPU nor host RAM. + stream_optimizer_state: bool = False + # The weight check snapshots every weight on the host; for the full model that exceeds host RAM. + check_weight_update_equal: bool = True + # A single node runs out of host memory saving the full model, so its conversion is sharded. + convert_args: str = "" + + +_RECIPES = { + "MiniMax-M3": _Recipe( + hf_repo="MiniMaxAI/MiniMax-M3", + megatron_model_type="minimax-m3", + num_gpus=32, + parallel=(4, 4, 8), + rollout_num_gpus_per_engine=8, + stream_optimizer_state=True, + check_weight_update_equal=False, + convert_args="--pipeline-model-parallel-size 4 --expert-model-parallel-size 8", + ), + "MiniMax-M3-4layer": _Recipe( + hf_repo="MiniMaxAI/MiniMax-M3", + megatron_model_type="minimax-m3-4layer", + num_gpus=4, + parallel=(2, 1, 2), + rollout_num_gpus_per_engine=4, + layers="0,1,3,4", + ), +} + + +@dataclass +class ScriptArgs(U.ExecuteTrainConfig): + model_name: Literal["MiniMax-M3", "MiniMax-M3-4layer"] = "MiniMax-M3-4layer" + num_gpus_per_node: int = 4 + run_id: str = U.create_run_id() + model_dir: str = "/root/models" + data_dir: str = "/root/datasets" + megatron_path: str = "/root/Megatron-LM" + train_offload_dir: str = "/root/train_offload" + num_rollout: int = 5 + rollout_max_response_len: int = 4096 + check_weight_update_equal: bool | None = None + enable_r3: bool = True + save: bool = False + extra_args: str = "" + + def __post_init__(self): + if self.check_weight_update_equal is None: + self.check_weight_update_equal = _RECIPES[self.model_name].check_weight_update_equal + + +def prepare(args: ScriptArgs): + backend = args.create_backend() + recipe = _RECIPES[args.model_name] + source = Path(args.model_dir) / recipe.hf_repo.split("/")[1] + target = Path(args.model_dir) / args.model_name + backend.exec_command_cpu(f"mkdir -p {args.model_dir} && hf download {recipe.hf_repo} --local-dir {source}") + if target != source: + layer_args = f"--layers {recipe.layers}" if recipe.layers else "" + backend.exec_command_gpu_once( + f"python {U.repo_base_dir}/tools/convert_minimax_m3_to_bf16.py " + f"--model-dir {source} --save-dir {target} --device cuda {layer_args}", + path_dst=str(target), + ) + backend.convert_checkpoint( + model_name=recipe.megatron_model_type, + megatron_model_type=recipe.megatron_model_type, + num_gpus_per_node=args.num_gpus_per_node, + multinode=args.num_nodes > 1, + num_nodes=args.num_nodes, + extra_args=recipe.convert_args, + dir_dst=args.model_dir, + hf_checkpoint=str(target), + megatron_path=args.megatron_path, + ) + backend.hf_download_dataset("zhuzilin/dapo-math-17k", data_dir=args.data_dir) + + +def execute(args: ScriptArgs): + recipe = _RECIPES[args.model_name] + num_gpus = args.num_nodes * args.num_gpus_per_node + assert num_gpus == recipe.num_gpus, f"{args.model_name} is laid out for {recipe.num_gpus} GPUs, got {num_gpus}" + tp, pp, ep = recipe.parallel + + ckpt_args = ( + f"--hf-checkpoint {args.model_dir}/{args.model_name} " + f"--ref-load {args.model_dir}/{recipe.megatron_model_type}_torch_dist " + ) + if args.save: + save_dir = f"{args.output_dir}/{args.run_id}/checkpoints" + ckpt_args += ( + f"--load {save_dir} --save {save_dir} --save-interval 10 " + "--no-save-optim --no-save-rng --no-load-optim --no-load-rng " + ) + + rollout_args = ( + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type math " + f"--num-rollout {args.num_rollout} " + "--rollout-batch-size 16 " + "--n-samples-per-prompt 16 " + "--rollout-temperature 0.8 " + "--num-steps-per-rollout 1 " + "--balance-data " + f"--prompt-data {args.data_dir}/dapo-math-17k/dapo-math-17k.jsonl " + "--input-key prompt " + f"--rollout-max-response-len {args.rollout_max_response_len} " + ) + + perf_args = ( + f"--tensor-model-parallel-size {tp} " + "--sequence-parallel " + f"--pipeline-model-parallel-size {pp} " + "--context-parallel-size 1 " + f"--expert-model-parallel-size {ep} " + "--expert-tensor-parallel-size 1 " + "--recompute-granularity full " + "--recompute-method uniform " + "--recompute-num-layers 1 " + "--micro-batch-size 1 " + "--max-tokens-per-gpu 8192 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--entropy-coef 0.00 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 3e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + if recipe.stream_optimizer_state: + optimizer_args += "--stream-optimizer-state-to-disk --stream-optimizer-state-moment-dtype bf16 " + + engine_gpus = recipe.rollout_num_gpus_per_engine + sglang_args = ( + f"--rollout-num-gpus-per-engine {engine_gpus} --sglang-tp-size {engine_gpus} --sglang-ep-size {engine_gpus} " + # routing replay needs materialized topk ids + "--sglang-moe-runner-backend triton " + "--sglang-chunked-prefill-size 8192 " + "--sglang-disable-radix-cache " + "--sglang-disable-prefill-cuda-graph " + "--sglang-kv-cache-dtype bfloat16 " + "--sglang-mem-fraction-static 0.7 " + "--router-health-success-threshold 1 " + "--router-health-check-interval-secs 15 " + "--router-health-failure-threshold 40 " + ) + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--attention-softmax-in-fp32 " + "--accumulate-allreduce-grads-in-fp32 " + f"--update-weight-buffer-size {1 * 1024**3} " + f"--actor-num-nodes {args.num_nodes} " + f"--actor-num-gpus-per-node {args.num_gpus_per_node} " + f"--num-gpus-per-node {args.num_gpus_per_node} " + f"--train-memory-margin-bytes {3 * 1024**3} " + "--offload-train-target disk " + f"--offload-train-disk-dir {args.train_offload_dir} " + "--colocate " + "--model-name minimax_m3 " + "--qkv-format thd " + "--rollout-health-check-interval 300 " + "--distributed-timeout-minutes 60 " + "--rollout-health-check-timeout 300 " + ) + if args.check_weight_update_equal: + misc_args += "--check-weight-update-equal --check-weight-update-skip-list vision_tower. multi_modal_projector. patch_merge_mlp. " + if args.enable_r3: + misc_args += "--use-rollout-routing-replay " + + train_args = ( + f"{ckpt_args} {rollout_args} {optimizer_args} {grpo_args} " + f"{U.get_default_wandb_args(__file__, run_id=args.run_id)} " + f"{perf_args} {sglang_args} {misc_args} {args.extra_args} " + ) + + extra_env_vars = { + "SGLANG_HEALTH_CHECK_TIMEOUT": "120", + # inductor's subprocess compile pool deadlocks under the torch_memory_saver LD_PRELOAD + "TORCHINDUCTOR_COMPILE_THREADS": "1", + } + + args.create_backend().execute_train( + train_args=train_args, + config=args, + num_gpus_per_node=args.num_gpus_per_node, + megatron_model_type=recipe.megatron_model_type, + extra_env_vars=extra_env_vars, + megatron_path=args.megatron_path, + ) + + +@U.dataclass_cli +def main(args: ScriptArgs): + prepare(args) + execute(args) + + +if __name__ == "__main__": + typer.run(main) diff --git a/tests/e2e/megatron/model_scripts/test_minimax_m3_4layer_ci.py b/tests/e2e/megatron/model_scripts/test_minimax_m3_4layer_ci.py new file mode 100644 index 00000000000..dd04ef32898 --- /dev/null +++ b/tests/e2e/megatron/model_scripts/test_minimax_m3_4layer_ci.py @@ -0,0 +1,66 @@ +import os + +from scripts.run_minimax_m3 import _RECIPES, ScriptArgs, execute +from tests.ci.ci_register import register_cuda_ci +from tests.ci.metric_history import register_ci_gate + +register_cuda_ci( + est_time=1500, + suite="stage-c-4-gpu-h200", + labels=["megatron", "model-scripts"], + hardware=["hopper", "blackwell"], + num_gpus=4, +) + +register_ci_gate(metric_key="train/grad_norm") +register_ci_gate(metric_key="train/ppo_kl") +register_ci_gate(metric_key="train/train_rollout_logprob_abs_diff") +register_ci_gate(metric_key="train/train_rollout_kl") +# deterministic_random scores each sample by a hash bit, so raw_reward has nothing to regress; track it +register_ci_gate(metric_key="rollout/raw_reward", constraint=None) +register_ci_gate(metric_key="perf/actor_train_mfu") + +_MODEL_ORG = "CharyZeng" + + +def _args() -> ScriptArgs: + return ScriptArgs( + model_name="MiniMax-M3-4layer", + num_nodes=1, + num_gpus_per_node=4, + num_rollout=3, + rollout_max_response_len=512, + extra_args=( + "--ci-test " + "--ci-disable-logprobs-checker " + # the pruned model scores 0 on math, which zeroes every advantage; a pseudo-random + # reward keeps the weights moving so each weight sync carries real deltas + "--rm-type deterministic_random " + "--offload-train-target cpu " + ), + ) + + +def prepare(args: ScriptArgs): + U = args.create_backend() + recipe = _RECIPES[args.model_name] + hf_checkpoint = f"{args.model_dir}/{args.model_name}" + U.exec_command_cpu(f"mkdir -p {args.model_dir} {args.data_dir}") + U.exec_command_cpu(f"hf download {_MODEL_ORG}/{args.model_name} --local-dir {hf_checkpoint}") + U.hf_download_dataset("zhuzilin/dapo-math-17k", data_dir=args.data_dir) + U.convert_checkpoint( + model_name=recipe.megatron_model_type, + megatron_model_type=recipe.megatron_model_type, + num_gpus_per_node=args.num_gpus_per_node, + dir_dst=args.model_dir, + hf_checkpoint=hf_checkpoint, + megatron_path=args.megatron_path, + ) + + +if __name__ == "__main__": + args = _args() + prepare(args) + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) + execute(args) diff --git a/tests/fast/backends/megatron_utils/update_weight/test_backend_partition_options.py b/tests/fast/backends/megatron_utils/update_weight/test_backend_partition_options.py index e3e420224dd..392fb98fa91 100644 --- a/tests/fast/backends/megatron_utils/update_weight/test_backend_partition_options.py +++ b/tests/fast/backends/megatron_utils/update_weight/test_backend_partition_options.py @@ -34,21 +34,22 @@ def test_atomic_export_groups_use_the_backend_query_adapter_rank( ) @pytest.mark.parametrize( - "swiglu,name,stride,dimension,expected", + "swiglu,quick_geglu,name,stride,dimension,expected", [ - (True, "linear_fc1.weight", 1, -1, (2, 0)), - (True, "linear_fc1.weight", 2, 1, (2, 1)), - (False, "linear_fc1.weight", 1, -1, (1, -1)), - (True, "linear_fc2.weight", 1, 0, (1, 1)), + (True, False, "linear_fc1.weight", 1, -1, (2, 0)), + (True, False, "linear_fc1.weight", 2, 1, (2, 1)), + (False, True, "linear_fc1.weight", 1, -1, (2, 0)), + (False, False, "linear_fc1.weight", 1, -1, (1, -1)), + (True, False, "linear_fc2.weight", 1, 0, (1, 1)), ], ) - def test_export_partition_metadata_follows_backend_swiglu_without_rewriting_valid_dimensions( - self, swiglu: bool, name: str, stride: int, dimension: int, expected: tuple[int, int] + def test_export_partition_metadata_follows_the_gated_activation_without_rewriting_valid_dimensions( + self, swiglu: bool, quick_geglu: bool, name: str, stride: int, dimension: int, expected: tuple[int, int] ) -> None: - """SwiGLU fixes missing grouped-FC1 metadata while preserving an explicit partition dimension.""" + """A gated FC1 (SwiGLU or quick-GeGLU) fixes missing grouped-FC1 metadata but keeps an explicit dimension.""" assert ( _check_and_fix_partition( - args=Namespace(backend=Namespace(swiglu=swiglu)), + args=Namespace(backend=Namespace(swiglu=swiglu, quick_geglu=quick_geglu)), name=name, partition_stride=stride, partition_dim=dimension, diff --git a/tests/fast/test_msa_block_max_scores.py b/tests/fast/test_msa_block_max_scores.py new file mode 100644 index 00000000000..91792dd87a8 --- /dev/null +++ b/tests/fast/test_msa_block_max_scores.py @@ -0,0 +1,43 @@ +"""The MSA block scores must equal each query's max causal logit per key block, sequence by sequence.""" + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip("needs CUDA", allow_module_level=True) + +from miles.kernels.attention.msa import msa_block_max_scores # noqa: E402 + +BLOCK = 128 +D = 128 + + +def _reference(q, k, lens, scale): + T, H, _ = q.shape + num_blocks = max(-(-n // BLOCK) for n in lens) + out = torch.full((H, T, num_blocks), float("-inf"), device=q.device) + start = 0 + for n in lens: + logits = torch.einsum("qhd,kd->hqk", q[start : start + n].float(), k[start : start + n].float()) * scale + future = torch.ones(n, n, dtype=torch.bool, device=q.device).triu(1) + logits = logits.masked_fill(future, float("-inf")) + nb = -(-n // BLOCK) + padded = torch.nn.functional.pad(logits, (0, nb * BLOCK - n), value=float("-inf")) + out[:, start : start + n, :nb] = padded.view(H, n, nb, BLOCK).amax(-1) + start += n + return out + + +@pytest.mark.parametrize("lens", [[300], [128, 128], [61, 513, 3], [1, 1, 200], [1000, 129]], ids=str) +def test_block_scores_match_dense_reference(lens): + g = torch.Generator(device="cuda").manual_seed(0) + T, H = sum(lens), 2 + q = torch.randn(T, H, D, device="cuda", generator=g).to(torch.bfloat16) + k = torch.randn(T, D, device="cuda", generator=g).to(torch.bfloat16) + cu = torch.tensor([0, *lens], device="cuda").cumsum(0).to(torch.int32) + got = msa_block_max_scores(q, k, cu, BLOCK, D**-0.5) + want = _reference(q, k, lens, D**-0.5) + assert torch.equal(torch.isinf(got), torch.isinf(want)), "the causal block pattern differs" + finite = torch.isfinite(want) + torch.testing.assert_close(got[finite], want[finite], rtol=1e-4, atol=1e-4) diff --git a/tests/fast/test_qsa_block_sparse_attn.py b/tests/fast/test_qsa_block_sparse_attn.py index a51c3b4a1b3..c25b6e55ccc 100644 --- a/tests/fast/test_qsa_block_sparse_attn.py +++ b/tests/fast/test_qsa_block_sparse_attn.py @@ -149,7 +149,7 @@ def test_tile_index_pair_is_consistent_both_ways(): q, k, v, idx = _single_sequence_case(T=256, budget=32) sel = selection_to_block_bitmap(idx, 256, BLK) block_first = torch.arange(sel.shape[1], dtype=torch.int32, device="cuda") * BLK - klist, kcnt, qlist, qcnt = build_tile_index_pair(sel, block_first, block_first + BLK - 1, 64, 32) + klist, kcnt, qlist, qcnt = build_tile_index_pair(sel, block_first, block_first + BLK - 1, 64, 32, BLK) forward = {(qt, int(klist[qt, i])) for qt in range(kcnt.numel()) for i in range(int(kcnt[qt]))} backward = {(int(qlist[kt, i]), kt) for kt in range(qcnt.numel()) for i in range(int(qcnt[kt]))} diff --git a/tests/fast/test_qsa_tile_index.py b/tests/fast/test_qsa_tile_index.py index 6fa8043f344..e8b192993ad 100644 --- a/tests/fast/test_qsa_tile_index.py +++ b/tests/fast/test_qsa_tile_index.py @@ -23,7 +23,7 @@ ] -def _packed_case(lens, dense: bool, seed: int = 0): +def _packed_case(lens, dense: bool, seed: int = 0, block_size: int = RATIO): """Layout plus a causal, in-sequence block selection; ``dense`` selects every allowed block.""" cu = torch.tensor([0, *lens]).cumsum(0) total = int(cu[-1]) @@ -32,26 +32,26 @@ def _packed_case(lens, dense: bool, seed: int = 0): seg[starts[1:]] = 1 seg = seg.cumsum(0) positions = torch.arange(total) - starts[seg] - layout = PackedBlockLayout(cu, positions, RATIO) + layout = PackedBlockLayout(cu, positions, block_size) g = torch.Generator().manual_seed(seed) sel = torch.zeros(total, layout.num_blocks, dtype=torch.uint8) for t in range(total): own = int(layout.token_block_start[t]) - allowed = torch.arange(own, own + int(positions[t]) // RATIO + 1) + allowed = torch.arange(own, own + int(positions[t]) // block_size + 1) if not dense: allowed = allowed[torch.rand(allowed.numel(), generator=g) < 0.3] - allowed = torch.cat([allowed, torch.tensor([own + int(positions[t]) // RATIO])]) + allowed = torch.cat([allowed, torch.tensor([own + int(positions[t]) // block_size])]) sel[t, allowed] = 1 return layout, positions, sel -def _attended_keys(t, layout, positions, sel): +def _attended_keys(t, layout, positions, sel, block_size): """Keys the kernel mask lets query ``t`` see: block selected, causal, same sequence.""" tok_base = int(layout.token_start[t]) blk_base = int(layout.token_block_start[t]) keys = torch.arange(tok_base, t + 1) - blocks = blk_base + (keys - tok_base) // RATIO + blocks = blk_base + (keys - tok_base) // block_size return keys[sel[t, blocks] != 0] @@ -61,19 +61,20 @@ def _tile_sets(lst, cnt): @pytest.mark.parametrize("lens", PACKS, ids=lambda lens: "-".join(map(str, lens))) @pytest.mark.parametrize("dense", [True, False], ids=["dense", "sparse"]) -def test_every_attended_key_tile_is_listed(lens, dense): - layout, positions, sel = _packed_case(lens, dense) +@pytest.mark.parametrize("block_size", [RATIO, 128], ids=["narrow-blocks", "blocks-wider-than-a-tile"]) +def test_every_attended_key_tile_is_listed(lens, dense, block_size): + layout, positions, sel = _packed_case(lens, dense, block_size=block_size) first, last = layout.block_first_token, layout.block_last_token total = sel.shape[0] - fwd = _tile_sets(*build_tile_index(sel, first, last, BQ, 64)) - klist, kcnt, qlist, qcnt = build_tile_index_pair(sel, first, last, BQ, 32) + fwd = _tile_sets(*build_tile_index(sel, first, last, BQ, 64, block_size)) + klist, kcnt, qlist, qcnt = build_tile_index_pair(sel, first, last, BQ, 32, block_size) bwd_k = _tile_sets(klist, kcnt) bwd_q = _tile_sets(qlist, qcnt) assert len(bwd_q) == -(-total // 32) for t in range(total): - keys = _attended_keys(t, layout, positions, sel) + keys = _attended_keys(t, layout, positions, sel, block_size) missing_fwd = sorted({int(k) for k in keys if int(k) // 64 not in fwd[t // BQ]}) missing_bwd = sorted({int(k) for k in keys if int(k) // 32 not in bwd_k[t // BQ]}) assert not missing_fwd, f"query {t}: forward tile list misses keys {missing_fwd}" diff --git a/tests/snapshots/launch_scripts/py/scripts/run_minimax_m3.py/execute.txt b/tests/snapshots/launch_scripts/py/scripts/run_minimax_m3.py/execute.txt new file mode 100644 index 00000000000..eac77394762 --- /dev/null +++ b/tests/snapshots/launch_scripts/py/scripts/run_minimax_m3.py/execute.txt @@ -0,0 +1,130 @@ +### 0 +pkill -9 sglang; sleep 3; ray stop + --force; pkill -9 ray; pkill -9 miles; sleep 3; pkill -9 ray; pkill -9 miles; pkill -9 redis; true; + +### 1 +export PYTHONUNBUFFERED=1 && ray start + --head + --node-ip-address 127.0.0.1 + --num-gpus 4 + --disable-usage-stats + +### 2 +nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l + +### 3 +export no_proxy=127.0.0.1 && export PYTHONUNBUFFERED=1 && ray job submit + --address="http://127.0.0.1:8265" + --runtime-env-json='{"env_vars": {"PYTHONUNBUFFERED": "1", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "NCCL_NVLS_ENABLE": "0", "no_proxy": "127.0.0.1,127.0.0.1", "MASTER_ADDR": "127.0.0.1", "SGLANG_HEALTH_CHECK_TIMEOUT": "120", "TORCHINDUCTOR_COMPILE_THREADS": "1", "PYTHONPATH": ":/root/Megatron-LM:/frozen/pythonpath"}}' + -- python3 /train.py + --spec miles_plugins.models.minimax_m3.minimax_m3 get_minimax_m3_spec + --num-layers 4 + --hidden-size 6144 + --ffn-hidden-size 12288 + --num-attention-heads 64 + --group-query-attention + --num-query-groups 4 + --kv-channels 128 + --qk-layernorm + --moe-layer-freq '[0,0,1,1]' + --num-experts 128 + --moe-router-topk 4 + --moe-router-score-function sigmoid + --moe-router-enable-expert-bias + --moe-router-bias-update-rate 0 + --moe-router-topk-scaling-factor 2.0 + --moe-ffn-hidden-size 3072 + --moe-shared-expert-intermediate-size 3072 + --moe-router-dtype fp32 + --moe-grouped-gemm + --moe-permute-fusion + --vocab-size 200064 + --rotary-base 5000000 + --rotary-percent 0.5 + --position-embedding-type rope + --no-rope-fusion + --normalization RMSNorm + --norm-epsilon 1e-6 + --apply-layernorm-1p + --quick-geglu + --activation-func-clamp-value 7 + --disable-bias-linear + --untie-embeddings-and-output-weights + --make-vocab-size-divisible-by 64 + --enable-experimental + --hf-checkpoint /root/models/MiniMax-M3-4layer + --ref-load /root/models/minimax-m3-4layer_torch_dist + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type math + --num-rollout 5 + --rollout-batch-size 16 + --n-samples-per-prompt 16 + --rollout-temperature 0.8 + --num-steps-per-rollout 1 + --balance-data + --prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --rollout-max-response-len 4096 + --optimizer adam + --lr 3e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --advantage-estimator grpo + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 + --use-wandb + --wandb-project miles-run_minimax_m3 + --wandb-group 260101-000000-000 + --wandb-key 'frozen-wandb-api-key' + --disable-wandb-random-suffix + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 2 + --expert-tensor-parallel-size 1 + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + --micro-batch-size 1 + --max-tokens-per-gpu 8192 + --rollout-num-gpus-per-engine 4 + --sglang-tp-size 4 + --sglang-ep-size 4 + --sglang-moe-runner-backend triton + --sglang-chunked-prefill-size 8192 + --sglang-disable-radix-cache + --sglang-disable-prefill-cuda-graph + --sglang-kv-cache-dtype bfloat16 + --sglang-mem-fraction-static 0.7 + --router-health-success-threshold 1 + --router-health-check-interval-secs 15 + --router-health-failure-threshold 40 + --attention-dropout 0.0 + --hidden-dropout 0.0 + --attention-softmax-in-fp32 + --accumulate-allreduce-grads-in-fp32 + --update-weight-buffer-size 1073741824 + --actor-num-nodes 1 + --actor-num-gpus-per-node 4 + --num-gpus-per-node 4 + --train-memory-margin-bytes 3221225472 + --offload-train-target disk + --offload-train-disk-dir /root/train_offload + --colocate + --model-name minimax_m3 + --qkv-format thd + --rollout-health-check-interval 300 + --distributed-timeout-minutes 60 + --rollout-health-check-timeout 300 + --check-weight-update-equal + --check-weight-update-skip-list vision_tower. multi_modal_projector. patch_merge_mlp. + --use-rollout-routing-replay + --deploy-component all diff --git a/tests/snapshots/launch_scripts/py/scripts/run_minimax_m3.py/prepare.txt b/tests/snapshots/launch_scripts/py/scripts/run_minimax_m3.py/prepare.txt new file mode 100644 index 00000000000..17908fb4ebc --- /dev/null +++ b/tests/snapshots/launch_scripts/py/scripts/run_minimax_m3.py/prepare.txt @@ -0,0 +1,59 @@ +### 0 +mkdir -p /root/models && hf download MiniMaxAI/MiniMax-M3 + --local-dir /root/models/MiniMax-M3 + +### 1 +rm -rf /root/models/MiniMax-M3-4layer + +### 2 +python /tools/convert_minimax_m3_to_bf16.py + --model-dir /root/models/MiniMax-M3 + --save-dir /root/models/MiniMax-M3-4layer + --device cuda + --layers 0,1,3,4 + +### 3 +PYTHONPATH=:/root/Megatron-LM:/frozen/pythonpath torchrun + --nproc-per-node 4 /tools/convert_hf_to_torch_dist.py + --spec miles_plugins.models.minimax_m3.minimax_m3 get_minimax_m3_spec + --num-layers 4 + --hidden-size 6144 + --ffn-hidden-size 12288 + --num-attention-heads 64 + --group-query-attention + --num-query-groups 4 + --kv-channels 128 + --qk-layernorm + --moe-layer-freq '[0,0,1,1]' + --num-experts 128 + --moe-router-topk 4 + --moe-router-score-function sigmoid + --moe-router-enable-expert-bias + --moe-router-bias-update-rate 0 + --moe-router-topk-scaling-factor 2.0 + --moe-ffn-hidden-size 3072 + --moe-shared-expert-intermediate-size 3072 + --moe-router-dtype fp32 + --moe-grouped-gemm + --moe-permute-fusion + --vocab-size 200064 + --rotary-base 5000000 + --rotary-percent 0.5 + --position-embedding-type rope + --no-rope-fusion + --normalization RMSNorm + --norm-epsilon 1e-6 + --apply-layernorm-1p + --quick-geglu + --activation-func-clamp-value 7 + --disable-bias-linear + --untie-embeddings-and-output-weights + --make-vocab-size-divisible-by 64 + --enable-experimental + --hf-checkpoint /root/models/MiniMax-M3-4layer + --save /root/models/minimax-m3-4layer_torch_dist + +### 4 +hf download + --repo-type dataset zhuzilin/dapo-math-17k + --local-dir /root/datasets/dapo-math-17k diff --git a/tests/snapshots/model_args/minimax-m3-4layer.txt b/tests/snapshots/model_args/minimax-m3-4layer.txt new file mode 100644 index 00000000000..8b0fccca99c --- /dev/null +++ b/tests/snapshots/model_args/minimax-m3-4layer.txt @@ -0,0 +1,60 @@ +"--spec" +"miles_plugins.models.minimax_m3.minimax_m3" +"get_minimax_m3_spec" +"--num-layers" +"4" +"--hidden-size" +"6144" +"--ffn-hidden-size" +"12288" +"--num-attention-heads" +"64" +"--group-query-attention" +"--num-query-groups" +"4" +"--kv-channels" +"128" +"--qk-layernorm" +"--moe-layer-freq" +"[0,0,1,1]" +"--num-experts" +"128" +"--moe-router-topk" +"4" +"--moe-router-score-function" +"sigmoid" +"--moe-router-enable-expert-bias" +"--moe-router-bias-update-rate" +"0" +"--moe-router-topk-scaling-factor" +"2.0" +"--moe-ffn-hidden-size" +"3072" +"--moe-shared-expert-intermediate-size" +"3072" +"--moe-router-dtype" +"fp32" +"--moe-grouped-gemm" +"--moe-permute-fusion" +"--vocab-size" +"200064" +"--rotary-base" +"5000000" +"--rotary-percent" +"0.5" +"--position-embedding-type" +"rope" +"--no-rope-fusion" +"--normalization" +"RMSNorm" +"--norm-epsilon" +"1e-6" +"--apply-layernorm-1p" +"--quick-geglu" +"--activation-func-clamp-value" +"7" +"--disable-bias-linear" +"--untie-embeddings-and-output-weights" +"--make-vocab-size-divisible-by" +"64" +"--enable-experimental" diff --git a/tests/snapshots/model_args/minimax-m3.txt b/tests/snapshots/model_args/minimax-m3.txt new file mode 100644 index 00000000000..d802122817c --- /dev/null +++ b/tests/snapshots/model_args/minimax-m3.txt @@ -0,0 +1,60 @@ +"--spec" +"miles_plugins.models.minimax_m3.minimax_m3" +"get_minimax_m3_spec" +"--num-layers" +"60" +"--hidden-size" +"6144" +"--ffn-hidden-size" +"12288" +"--num-attention-heads" +"64" +"--group-query-attention" +"--num-query-groups" +"4" +"--kv-channels" +"128" +"--qk-layernorm" +"--moe-layer-freq" +"[0,0,0,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]" +"--num-experts" +"128" +"--moe-router-topk" +"4" +"--moe-router-score-function" +"sigmoid" +"--moe-router-enable-expert-bias" +"--moe-router-bias-update-rate" +"0" +"--moe-router-topk-scaling-factor" +"2.0" +"--moe-ffn-hidden-size" +"3072" +"--moe-shared-expert-intermediate-size" +"3072" +"--moe-router-dtype" +"fp32" +"--moe-grouped-gemm" +"--moe-permute-fusion" +"--vocab-size" +"200064" +"--rotary-base" +"5000000" +"--rotary-percent" +"0.5" +"--position-embedding-type" +"rope" +"--no-rope-fusion" +"--normalization" +"RMSNorm" +"--norm-epsilon" +"1e-6" +"--apply-layernorm-1p" +"--quick-geglu" +"--activation-func-clamp-value" +"7" +"--disable-bias-linear" +"--untie-embeddings-and-output-weights" +"--make-vocab-size-divisible-by" +"64" +"--enable-experimental" diff --git a/tools/convert_minimax_m3_to_bf16.py b/tools/convert_minimax_m3_to_bf16.py new file mode 100644 index 00000000000..1075c5c915a --- /dev/null +++ b/tools/convert_minimax_m3_to_bf16.py @@ -0,0 +1,177 @@ +""" +Convert a MiniMax-M3 / MiniMax-M3.1 checkpoint (`MiniMaxM3SparseForConditionalGeneration`) to plain +BF16, optionally keeping a subset of decoder layers for partial-model work. A BF16 source (MiniMax-M3) +is only cut to the kept layers. + +MiniMax-M3.1-Flash stores attention, dense and shared-expert linears in MXFP8 (e4m3 with one ue8m0 scale per +32 inputs) and the routed experts in compressed-tensors NVFP4 (packed e2m1, an e4m3 scale per 16 inputs +and a per-tensor encode scale; w = q * s / g). The output dequantizes both to BF16 and drops +`quantization_config`, so the HF remote code, the Megatron bridge and a BF16 SGLang engine read it as-is. +The vision tower and the `dspark/` draft model are not touched. + +python tools/convert_minimax_m3_to_bf16.py --model-dir --save-dir [--layers 0,1,3,4] +""" + +import copy +import json +import re +import shutil +from argparse import ArgumentParser +from collections import defaultdict +from pathlib import Path + +import torch +from safetensors import safe_open +from safetensors.torch import save_file +from tqdm import tqdm + +E2M1_TABLE = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], + dtype=torch.float32, +) +MXFP8_BLOCK = 32 +NVFP4_BLOCK = 16 +SHARD_BYTES = 5 * 1024**3 +SCALE_SUFFIXES = ("weight_scale_inv", "weight_scale", "weight_global_scale", "input_global_scale") +PER_LAYER_SPARSE_KEYS = ("sparse_attention_freq", "sparse_disable_index_value") + +_LAYER_RE = re.compile(r"^(language_model\.model\.layers\.)(\d+)(\..+)$") + + +def dequant_mxfp8(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + rows, cols = weight.shape + assert scale.shape == (rows, cols // MXFP8_BLOCK), f"{weight.shape=} {scale.shape=}" + blocks = weight.float().view(rows, cols // MXFP8_BLOCK, MXFP8_BLOCK) + return (blocks * torch.exp2(scale.float() - 127.0).unsqueeze(-1)).view(rows, cols).to(torch.bfloat16) + + +def dequant_nvfp4(packed: torch.Tensor, scale: torch.Tensor, global_scale: torch.Tensor) -> torch.Tensor: + """Low nibble is the even element.""" + rows, cols = packed.shape[0], packed.shape[1] * 2 + assert scale.shape == (rows, cols // NVFP4_BLOCK), f"{packed.shape=} {scale.shape=}" + table = E2M1_TABLE.to(packed.device) + values = torch.stack([table[(packed & 0x0F).long()], table[(packed >> 4).long()]], dim=-1).view(rows, cols) + blocks = values.view(rows, cols // NVFP4_BLOCK, NVFP4_BLOCK) * scale.float().unsqueeze(-1) + return (blocks / global_scale.float().reshape(())).view(rows, cols).to(torch.bfloat16) + + +def output_name(name: str, layer_map: dict[int, int] | None) -> str | None: + """Output name of a source tensor, or None when its layer is not kept.""" + match = _LAYER_RE.match(name) + if match is None or layer_map is None: + return name + layer = int(match.group(2)) + if layer not in layer_map: + return None + return f"{match.group(1)}{layer_map[layer]}{match.group(3)}" + + +def convert_config(config: dict, layers: list[int] | None) -> dict: + config = copy.deepcopy(config) + config.pop("quantization_config", None) + if layers is not None: + text = config["text_config"] + text["num_hidden_layers"] = len(layers) + text["moe_layer_freq"] = [text["moe_layer_freq"][layer] for layer in layers] + sparse = text["sparse_attention_config"] + for key in PER_LAYER_SPARSE_KEYS: + sparse[key] = [sparse[key][layer] for layer in layers] + return config + + +class ShardWriter: + def __init__(self, save_dir: Path): + self.save_dir = save_dir + self.buffer: dict[str, torch.Tensor] = {} + self.buffer_bytes = 0 + self.num_shards = 0 + self.weight_map: dict[str, str] = {} + + def add(self, name: str, tensor: torch.Tensor) -> None: + assert name not in self.weight_map and name not in self.buffer, f"duplicate tensor {name}" + self.buffer[name] = tensor.contiguous().cpu() + self.buffer_bytes += tensor.numel() * tensor.element_size() + if self.buffer_bytes >= SHARD_BYTES: + self.flush() + + def flush(self) -> None: + if not self.buffer: + return + self.num_shards += 1 + file_name = f"model-{self.num_shards:05d}.safetensors" + save_file(self.buffer, str(self.save_dir / file_name), metadata={"format": "pt"}) + self.weight_map.update(dict.fromkeys(self.buffer, file_name)) + self.buffer, self.buffer_bytes = {}, 0 + + +def main(model_dir: str, save_dir: str, layers: list[int] | None, device: str): + src, dst = Path(model_dir), Path(save_dir) + dst.mkdir(parents=True, exist_ok=True) + config = json.loads((src / "config.json").read_text()) + assert config["architectures"] == ["MiniMaxM3SparseForConditionalGeneration"], config["architectures"] + num_layers = config["text_config"]["num_hidden_layers"] + assert layers is None or all(0 <= layer < num_layers for layer in layers), layers + layer_map = {old: new for new, old in enumerate(layers)} if layers is not None else None + + weight_map = json.loads((src / "model.safetensors.index.json").read_text())["weight_map"] + names_by_file = defaultdict(list) + for name, file_name in weight_map.items(): + names_by_file[file_name].append(name) + handles = {} + + def load(name: str) -> torch.Tensor: + file_name = weight_map[name] + if file_name not in handles: + handles[file_name] = safe_open(str(src / file_name), framework="pt", device="cpu") + return handles[file_name].get_tensor(name).to(device) + + writer = ShardWriter(dst) + counts = {"mxfp8": 0, "nvfp4": 0, "copied": 0} + for file_name in tqdm(sorted(names_by_file), desc="shards"): + for name in names_by_file[file_name]: + if name.endswith(SCALE_SUFFIXES): + continue + if name.endswith(".weight_packed"): + base = name.removesuffix(".weight_packed") + new_name = output_name(f"{base}.weight", layer_map) + if new_name is None: + continue + tensor = dequant_nvfp4(load(name), load(f"{base}.weight_scale"), load(f"{base}.weight_global_scale")) + counts["nvfp4"] += 1 + else: + new_name = output_name(name, layer_map) + if new_name is None: + continue + scale_name = name.removesuffix(".weight") + ".weight_scale_inv" + if name.endswith(".weight") and scale_name in weight_map: + tensor = dequant_mxfp8(load(name), load(scale_name)) + counts["mxfp8"] += 1 + else: + tensor = load(name) + assert tensor.dtype not in (torch.float8_e4m3fn, torch.uint8), f"unhandled {name} {tensor.dtype}" + counts["copied"] += 1 + writer.add(new_name, tensor) + writer.flush() + + (dst / "model.safetensors.index.json").write_text( + json.dumps({"metadata": {}, "weight_map": writer.weight_map}, indent=2) + ) + (dst / "config.json").write_text(json.dumps(convert_config(config, layers), indent=2, ensure_ascii=False)) + for path in src.iterdir(): + if ( + path.is_file() + and path.suffix != ".safetensors" + and path.name not in ("config.json", "model.safetensors.index.json") + ): + shutil.copy(path, dst / path.name) + print(f"{counts} -> {dst} ({writer.num_shards} shards)") + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument("--model-dir", required=True) + parser.add_argument("--save-dir", required=True) + parser.add_argument("--layers", type=lambda s: [int(x) for x in s.split(",")], default=None) + parser.add_argument("--device", default="cuda") + args = parser.parse_args() + main(args.model_dir, args.save_dir, args.layers, args.device)