From 6f3eb3dfa33b522099e75e32717f1c2a86e8281c Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Mon, 22 Jun 2026 22:02:27 +0000 Subject: [PATCH 01/13] megatron-bump: adapt MTP-in-RL to upstream mtp_detach_heads (#6) Megatron dev reimplemented PR #6's MTP-in-RL support natively: - process_mtp_loss derives MTP labels from input_ids when labels is None (RL). - config.mtp_detach_heads detaches output head + MTP embedding gradients. So on the miles side: set config.mtp_detach_heads=True when enable_mtp_training, and stop passing the now-unsupported mtp_kwargs to GPTModel.forward (labels=None + input_ids derivation is equivalent to mtp_labels=batch['tokens']). --- miles/backends/megatron_utils/model.py | 5 +++-- miles/backends/megatron_utils/model_provider.py | 8 ++++++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 44fdaeedcd8..032984f3a04 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -449,8 +449,9 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p "loss_mask": batch["full_loss_masks"], } - if args.enable_mtp_training: - forward_kwargs["mtp_kwargs"] = {"mtp_labels": batch["tokens"]} + # MTP labels: Megatron's process_mtp_loss derives them from input_ids + # (== batch["tokens"]) when labels is None, so no mtp_kwargs is needed. + # MTP head detach is configured via config.mtp_detach_heads (model_provider). if (x := batch["multimodal_train_inputs"]) is not None: forward_kwargs.update(x) diff --git a/miles/backends/megatron_utils/model_provider.py b/miles/backends/megatron_utils/model_provider.py index 03e924637b7..0f2ecf5aed4 100644 --- a/miles/backends/megatron_utils/model_provider.py +++ b/miles/backends/megatron_utils/model_provider.py @@ -156,6 +156,14 @@ def model_provider( assert config is None, "miles builds the config from args, so it expects config to be None" config = core_transformer_config_from_args(args) + # `enable_mtp_training` comes from miles' arg parser; megatron-only arg contexts + # (e.g. the run_megatron debug worker) won't have it, so default to False. + if getattr(args, "enable_mtp_training", False): + # RL MTP training: detach the MTP heads so MTP gradients do not flow into the + # shared output layer / embedding (Megatron implements this via mtp_detach_heads; + # previously miles patched Megatron directly). See bump_docs/01-cherry-pick.md (#6). + config.mtp_detach_heads = True + if args.spec is not None: transformer_layer_spec = import_module(args.spec) # Allow the spec to be a function so that user can use customized Megatron easier. From e5f57d947cdafb388fae9b32c774340faeb32c88 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Tue, 23 Jun 2026 04:37:43 +0000 Subject: [PATCH 02/13] dsv4: adapt miles DeepSeek-V4 to bumped Megatron (dev-native config) Run miles' custom dsv4 attention on the bumped Megatron (miles-main-20260622): - rename plugin/script config to dev-native names: csa_window_size, csa_compress_ratios, csa_compress_rotary_base, o_groups, o_lora_rank, moe_n_hash_layers (mbridge). Keep dsv4_hc_* (mHC precision-branched). Script uses --csa-compress-ratios "[..]" string form; drop --no-activation-func-clamp-shared-expert. - dev API drift: tokenizer _vocab_size_with_padding moved to megatron.core.tokenizers; enable_gloo_process_groups -> use_gloo_process_groups; DeepSeekV4Attention accepts name=. - run_megatron worker: build forward-only model without DDP (wrap_with_ddp=run_backward). Validated against the old-megatron baseline (run-megatron, 4-layer prune): mHC+attention numerically exact, logprob mean within 8e-4. --- miles/backends/megatron_utils/arguments.py | 2 +- miles/backends/megatron_utils/initialize.py | 2 +- miles/backends/megatron_utils/model.py | 4 ++-- .../debug_utils/run_megatron/worker/main.py | 4 +++- miles_plugins/mbridge/deepseekv4.py | 12 ++++++------ miles_plugins/models/deepseek_v4/deepseek_v4.py | 11 ++++++----- .../models/deepseek_v4/ops/compressor.py | 2 +- .../models/deepseek_v4/ops/v4_indexer.py | 2 +- scripts/models/deepseek-v4-flash.sh | 16 +++++++++------- scripts/models/deepseek-v4-pro.sh | 16 +++++++++------- 10 files changed, 39 insertions(+), 32 deletions(-) diff --git a/miles/backends/megatron_utils/arguments.py b/miles/backends/megatron_utils/arguments.py index e64409bbc9d..b70096d2ed2 100644 --- a/miles/backends/megatron_utils/arguments.py +++ b/miles/backends/megatron_utils/arguments.py @@ -1,8 +1,8 @@ import logging import os +from megatron.core.tokenizers.utils.build_tokenizer import vocab_size_with_padding as _vocab_size_with_padding from megatron.training.arguments import parse_args, validate_args -from megatron.training.tokenizer.tokenizer import _vocab_size_with_padding __all__ = ["validate_args", "parse_args", "set_default_megatron_args"] diff --git a/miles/backends/megatron_utils/initialize.py b/miles/backends/megatron_utils/initialize.py index ac9c04b084e..cb8617179c3 100644 --- a/miles/backends/megatron_utils/initialize.py +++ b/miles/backends/megatron_utils/initialize.py @@ -54,7 +54,7 @@ def _initialize_distributed(args, get_embedding_ranks=None, get_position_embeddi order="tp-cp-ep-dp-pp" if not args.use_tp_pp_dp_mapping else "tp-cp-ep-pp-dp", get_embedding_ranks=get_embedding_ranks, get_position_embedding_ranks=get_position_embedding_ranks, - create_gloo_process_groups=args.enable_gloo_process_groups, + create_gloo_process_groups=args.use_gloo_process_groups, ) diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 032984f3a04..4df67e8aadf 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -152,14 +152,14 @@ def setup_model_and_optimizer( optimizer = get_megatron_muon_optimizer( config=config, model_chunks=model, - use_gloo_process_groups=args.enable_gloo_process_groups, + use_gloo_process_groups=args.use_gloo_process_groups, layer_wise_distributed_optimizer="dist" in config.optimizer.lower(), ) else: optimizer = get_megatron_optimizer( config=config, model_chunks=model, - use_gloo_process_groups=args.enable_gloo_process_groups, + use_gloo_process_groups=args.use_gloo_process_groups, ) opt_param_scheduler = get_optimizer_param_scheduler(args, optimizer) return model, optimizer, opt_param_scheduler diff --git a/miles/utils/debug_utils/run_megatron/worker/main.py b/miles/utils/debug_utils/run_megatron/worker/main.py index 359a51c1d60..4c458ff376b 100644 --- a/miles/utils/debug_utils/run_megatron/worker/main.py +++ b/miles/utils/debug_utils/run_megatron/worker/main.py @@ -132,7 +132,9 @@ def _initialize_megatron(args: argparse.Namespace) -> None: def _build_and_load_model(args: argparse.Namespace, script: WorkerScriptArgs) -> list[Any]: model_provider: Callable[..., Any] = get_model_provider_func(args, role=script.role) - model: list[Any] = get_model(model_provider, ModelType.encoder_or_decoder) + # Forward-only runs skip DDP wrapping so the distributed-optimizer grad buffer + # (tens-to-hundreds of GB for large MoE models) is never allocated -> avoids OOM. + model: list[Any] = get_model(model_provider, ModelType.encoder_or_decoder, wrap_with_ddp=script.run_backward) if args.load is not None: load_checkpoint( diff --git a/miles_plugins/mbridge/deepseekv4.py b/miles_plugins/mbridge/deepseekv4.py index fd786732bb7..d607a926c49 100644 --- a/miles_plugins/mbridge/deepseekv4.py +++ b/miles_plugins/mbridge/deepseekv4.py @@ -107,17 +107,17 @@ def _build_config(self): config.dsv4_hc_sinkhorn_iters = getattr(self.hf_config, "hc_sinkhorn_iters", 20) config.dsv4_hc_eps = getattr(self.hf_config, "hc_eps", 1e-6) - config.dsv4_compress_ratios = getattr(self.hf_config, "compress_ratios", None) - config.dsv4_compress_rope_theta = getattr(self.hf_config, "compress_rope_theta", 160000) + config.csa_compress_ratios = getattr(self.hf_config, "compress_ratios", None) + config.csa_compress_rotary_base = getattr(self.hf_config, "compress_rope_theta", 160000) config.dsv4_swiglu_limit = getattr(self.hf_config, "swiglu_limit", 0.0) if config.dsv4_swiglu_limit > 0: config.bias_activation_fusion = False config.activation_func_clamp_value = config.dsv4_swiglu_limit - config.dsv4_o_groups = getattr(self.hf_config, "o_groups", 8) - config.dsv4_o_lora_rank = getattr(self.hf_config, "o_lora_rank", 1024) - config.dsv4_n_hash_layers = getattr(self.hf_config, "n_hash_layers", 3) - config.dsv4_window_size = getattr(self.hf_config, "window_size", 128) + config.o_groups = getattr(self.hf_config, "o_groups", 8) + config.o_lora_rank = getattr(self.hf_config, "o_lora_rank", 1024) + config.moe_n_hash_layers = getattr(self.hf_config, "n_hash_layers", 3) + config.csa_window_size = getattr(self.hf_config, "window_size", 128) return config diff --git a/miles_plugins/models/deepseek_v4/deepseek_v4.py b/miles_plugins/models/deepseek_v4/deepseek_v4.py index b646ebd1cca..fcebd15aee1 100644 --- a/miles_plugins/models/deepseek_v4/deepseek_v4.py +++ b/miles_plugins/models/deepseek_v4/deepseek_v4.py @@ -53,6 +53,7 @@ def __init__( attention_type: str = None, cp_comm_type: str = None, pg_collection=None, + name: str = None, ): _enable_deepseek_v4_tf32() super().__init__(config=config) @@ -74,14 +75,14 @@ def __init__( self.n_heads = config.num_attention_heads self.n_local_heads = self.n_heads // config.tensor_model_parallel_size self.q_lora_rank = config.q_lora_rank - self.o_lora_rank = config.dsv4_o_lora_rank + self.o_lora_rank = config.o_lora_rank self.head_dim = config.kv_lora_rank self.rope_head_dim = config.qk_pos_emb_head_dim self.nope_head_dim = self.head_dim - self.rope_head_dim - self.n_groups = config.dsv4_o_groups + self.n_groups = config.o_groups self.n_local_groups = self.n_groups // config.tensor_model_parallel_size - self.window_size = config.dsv4_window_size - self.compress_ratio = config.dsv4_compress_ratios[layer_id] if config.dsv4_compress_ratios else 0 + self.window_size = config.csa_window_size + self.compress_ratio = config.csa_compress_ratios[layer_id] if config.csa_compress_ratios else 0 self.eps = config.layernorm_epsilon self.use_fp8_qat = config.fp8 is not None @@ -182,7 +183,7 @@ def __init__( else: self.indexer = None - rope_base = config.dsv4_compress_rope_theta if self.compress_ratio else config.rotary_base + rope_base = config.csa_compress_rotary_base if self.compress_ratio else config.rotary_base yarn_disabled = not self.compress_ratio freqs_cis = wrapped_precompute_freqs_cis( config, rope_head_dim=self.rope_head_dim, base=rope_base, yarn_disabled=yarn_disabled diff --git a/miles_plugins/models/deepseek_v4/ops/compressor.py b/miles_plugins/models/deepseek_v4/ops/compressor.py index 55a93f2949d..31806392508 100644 --- a/miles_plugins/models/deepseek_v4/ops/compressor.py +++ b/miles_plugins/models/deepseek_v4/ops/compressor.py @@ -89,7 +89,7 @@ def __init__( self.ape._keep_fp32 = True - base = config.dsv4_compress_rope_theta + base = config.csa_compress_rotary_base assert rope_head_dim == 64 assert base == 160000 freqs_cis = wrapped_precompute_freqs_cis(config, rope_head_dim=rope_head_dim, base=base) diff --git a/miles_plugins/models/deepseek_v4/ops/v4_indexer.py b/miles_plugins/models/deepseek_v4/ops/v4_indexer.py index 45ba7c60301..9e4d7a2cc5a 100644 --- a/miles_plugins/models/deepseek_v4/ops/v4_indexer.py +++ b/miles_plugins/models/deepseek_v4/ops/v4_indexer.py @@ -67,7 +67,7 @@ def __init__(self, config: TransformerConfig, pg_collection=None): cp_group=pg_collection.cp, ) - rope_base = config.dsv4_compress_rope_theta if self.compress_ratio else config.rotary_base + rope_base = config.csa_compress_rotary_base if self.compress_ratio else config.rotary_base freqs_cis = wrapped_precompute_freqs_cis(config, rope_head_dim=self.rope_head_dim, base=rope_base) self.register_buffer("freqs_cis", freqs_cis, persistent=False) diff --git a/scripts/models/deepseek-v4-flash.sh b/scripts/models/deepseek-v4-flash.sh index 3997290de44..328597e0804 100644 --- a/scripts/models/deepseek-v4-flash.sh +++ b/scripts/models/deepseek-v4-flash.sh @@ -11,9 +11,11 @@ printf -v MOE_LAYER_FREQ "[%s]" "$(IFS=', '; echo "${arr[*]}")" if [ ${#COMPRESS_RATIOS[@]} -eq 0 ]; then COMPRESS_RATIOS=(0 0 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 0) fi +# dev's --csa-compress-ratios takes a single Python-list-expression string, e.g. "[0,0,4,128]" +printf -v COMPRESS_RATIOS_STR "[%s]" "$(IFS=','; echo "${COMPRESS_RATIOS[*]}")" ROTARY_SCALING_FACTOR="${ROTARY_SCALING_FACTOR:-16}" -SWIGLU_LIMIT_ARGS=(--activation-func-clamp-value 10 --no-bias-swiglu-fusion --no-activation-func-clamp-shared-expert) +SWIGLU_LIMIT_ARGS=(--activation-func-clamp-value 10 --no-bias-swiglu-fusion) # DeepSeek V4 Flash config MODEL_ARGS=( @@ -66,12 +68,12 @@ MODEL_ARGS=( --experimental-attention-variant dsv4 --dsv4-hc-mult 4 --dsv4-hc-sinkhorn-iters 20 - --dsv4-compress-ratios "${COMPRESS_RATIOS[@]}" - --dsv4-compress-rope-theta 160000 - --dsv4-o-groups 8 - --dsv4-o-lora-rank 1024 - --dsv4-n-hash-layers 3 - --dsv4-window-size 128 + --csa-compress-ratios "${COMPRESS_RATIOS_STR}" + --csa-compress-rotary-base 160000 + --o-groups 8 + --o-lora-rank 1024 + --moe-n-hash-layers 3 + --csa-window-size 128 # DSA Indexer --dsa-indexer-n-heads 64 diff --git a/scripts/models/deepseek-v4-pro.sh b/scripts/models/deepseek-v4-pro.sh index 067d887d967..3fb85c156a1 100644 --- a/scripts/models/deepseek-v4-pro.sh +++ b/scripts/models/deepseek-v4-pro.sh @@ -11,9 +11,11 @@ printf -v MOE_LAYER_FREQ "[%s]" "$(IFS=', '; echo "${arr[*]}")" if [ ${#COMPRESS_RATIOS[@]} -eq 0 ]; then COMPRESS_RATIOS=(128 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 128 4 0) fi +# dev's --csa-compress-ratios takes a single Python-list-expression string, e.g. "[128,128,4,128]" +printf -v COMPRESS_RATIOS_STR "[%s]" "$(IFS=','; echo "${COMPRESS_RATIOS[*]}")" ROTARY_SCALING_FACTOR="${ROTARY_SCALING_FACTOR:-16}" -SWIGLU_LIMIT_ARGS=(--activation-func-clamp-value 10 --no-bias-swiglu-fusion --no-activation-func-clamp-shared-expert) +SWIGLU_LIMIT_ARGS=(--activation-func-clamp-value 10 --no-bias-swiglu-fusion) # DeepSeek V4 Pro config MODEL_ARGS=( @@ -66,12 +68,12 @@ MODEL_ARGS=( --experimental-attention-variant dsv4 --dsv4-hc-mult 4 --dsv4-hc-sinkhorn-iters 20 - --dsv4-compress-ratios "${COMPRESS_RATIOS[@]}" - --dsv4-compress-rope-theta 160000 - --dsv4-o-groups 16 - --dsv4-o-lora-rank 1024 - --dsv4-n-hash-layers 3 - --dsv4-window-size 128 + --csa-compress-ratios "${COMPRESS_RATIOS_STR}" + --csa-compress-rotary-base 160000 + --o-groups 16 + --o-lora-rank 1024 + --moe-n-hash-layers 3 + --csa-window-size 128 # DSA Indexer --dsa-indexer-n-heads 64 From 5903a89c4b4a1d8ca1402ef52f368522e04f5496 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Tue, 23 Jun 2026 04:38:59 +0000 Subject: [PATCH 03/13] docker: bump MEGATRON_BRANCH to miles-main-20260622 --- docker/Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index b259542f558..7ee258be2b0 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -19,7 +19,7 @@ ARG SGLANG_BRANCH=sglang-miles ARG SGLANG_COMMIT="" ARG MEGATRON_REPO=radixark/Megatron-LM -ARG MEGATRON_BRANCH=miles-main +ARG MEGATRON_BRANCH=miles-main-20260622 ARG ENABLE_CUDA_13=1 From c1ac131aa350c006c95785dc4e2feb6352390562 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Tue, 23 Jun 2026 07:06:05 +0000 Subject: [PATCH 04/13] dsv4: convert 4-layer on 1 GPU (PP1) to avoid hash-MoE+PP>1 assertion on new Megatron convert_hf_to_torch_dist auto-forces PP=world_size for >1 GPU; bumped Megatron asserts hash-MoE layers + PP>1 need an explicit pipeline_model_parallel_layout. The 4-layer prune converts fine at PP1 on 1 GPU (validated). Full Flash/Pro use explicit multi-PP convert configs (their PP>1 hash-MoE convert layout is a separate follow-up). --- scripts/run_deepseek_v4.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/scripts/run_deepseek_v4.py b/scripts/run_deepseek_v4.py index 813865376d6..f636555840a 100644 --- a/scripts/run_deepseek_v4.py +++ b/scripts/run_deepseek_v4.py @@ -255,7 +255,11 @@ def _prepare_spmd(args: ScriptArgs): num_gpus_for_convert = actor_num_gpus_per_node if is_4layer: - num_gpus_for_convert = min(num_gpus_for_convert, 4) + # Convert on a single GPU (PP1). convert_hf_to_torch_dist auto-forces PP=world_size when + # >1 GPU, but the bumped Megatron asserts hash-MoE layers + PP>1 require an explicit + # pipeline_model_parallel_layout (which the convert doesn't set). PP1 sidesteps it and is + # plenty for the 4-layer prune. (Full Flash/Pro use explicit multi-PP convert configs.) + num_gpus_for_convert = 1 U.convert_checkpoint( model_name=args.model_name, From b8d0f425bf9801b8009b32899c0962b425b2643c Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Tue, 23 Jun 2026 07:11:31 +0000 Subject: [PATCH 05/13] test: rename enable_gloo_process_groups -> use_gloo_process_groups in lora test mock Matches the dev Megatron arg rename used in model.py (use_gloo_process_groups). --- tests/fast/backends/megatron_utils/test_lora_model_branches.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/fast/backends/megatron_utils/test_lora_model_branches.py b/tests/fast/backends/megatron_utils/test_lora_model_branches.py index 3fb081f6a90..3055ac70e1e 100644 --- a/tests/fast/backends/megatron_utils/test_lora_model_branches.py +++ b/tests/fast/backends/megatron_utils/test_lora_model_branches.py @@ -91,7 +91,7 @@ def _make_args(self, lora_rank=32, role="actor", mode="bridge"): use_checkpoint_opt_param_scheduler=False, override_opt_param_scheduler=False, lr_wsd_decay_style="linear", - enable_gloo_process_groups=False, + use_gloo_process_groups=False, ) @patch(f"{_MODEL_MODULE}.get_optimizer_param_scheduler") From fa36961171e5e3641fdd1ac6d430879f5419c86b Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Tue, 23 Jun 2026 22:55:26 +0000 Subject: [PATCH 06/13] dsv4 step3: indexer-replay (R3) seam on V4Indexer Mirror GLM5: register indexer_replay_manager on V4Indexer (stream_idx=layer_id) and route the indexer top-k through get_topk_fn (flattened to [n_tokens, n_kv]). No-op/transparent when the manager is disabled; DSv4 forward verified unchanged (loss 19.92). Full record/replay-overlap validation needs an RL run with --use-rollout-indexer-replay. --- .../models/deepseek_v4/deepseek_v4.py | 2 +- .../models/deepseek_v4/ops/v4_indexer.py | 20 +++++++++++++++++-- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/miles_plugins/models/deepseek_v4/deepseek_v4.py b/miles_plugins/models/deepseek_v4/deepseek_v4.py index fcebd15aee1..d0898662988 100644 --- a/miles_plugins/models/deepseek_v4/deepseek_v4.py +++ b/miles_plugins/models/deepseek_v4/deepseek_v4.py @@ -171,7 +171,7 @@ def __init__( ) if self.compress_ratio == 4: if os.environ.get("V4_INDEXER_IMPL", "tilelang") == "tilelang": - self.indexer = V4Indexer(config=config, pg_collection=pg_collection) + self.indexer = V4Indexer(config=config, pg_collection=pg_collection, layer_id=layer_id) else: indexer_submodules = DSAIndexerSubmodules( linear_wq_b=TELinear, diff --git a/miles_plugins/models/deepseek_v4/ops/v4_indexer.py b/miles_plugins/models/deepseek_v4/ops/v4_indexer.py index 9e4d7a2cc5a..b7f930ac548 100644 --- a/miles_plugins/models/deepseek_v4/ops/v4_indexer.py +++ b/miles_plugins/models/deepseek_v4/ops/v4_indexer.py @@ -7,6 +7,8 @@ from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.transformer_config import TransformerConfig +from miles.utils.replay_base import indexer_replay_manager + from miles_plugins.models.deepseek_v4.ops.compressor import DeepSeekV4Compressor from miles_plugins.models.deepseek_v4.ops.cp_utils import all_gather_cp, get_freqs_cis_for_cp from miles_plugins.models.deepseek_v4.ops.kernel.tilelang_indexer_fwd import ( @@ -21,7 +23,7 @@ class V4Indexer(MegatronModule): """DSA Indexer for DeepSeek-V4 C4 layers.""" - def __init__(self, config: TransformerConfig, pg_collection=None): + def __init__(self, config: TransformerConfig, pg_collection=None, layer_id: int = 0): super().__init__(config=config) self.hidden_size = config.hidden_size @@ -71,6 +73,11 @@ def __init__(self, config: TransformerConfig, pg_collection=None): freqs_cis = wrapped_precompute_freqs_cis(config, rope_head_dim=self.rope_head_dim, base=rope_base) self.register_buffer("freqs_cis", freqs_cis, persistent=False) + # RL rollout-routing-replay (R3) seam for the sparse-attention indexer topk: lets the miles + # indexer_replay_manager record (on rollout) / replay (on train forward) the top-k KV picks, + # mirroring the MoE routing-replay seam. No-op unless the manager is enabled (record/replay). + indexer_replay_manager.register_to_module(self, "indexer_replay", stream_idx=layer_id) + def forward(self, x: torch.Tensor, qr: torch.Tensor, mask=None, packed_seq_params=None): """Forward pass. @@ -128,7 +135,16 @@ def forward(self, x: torch.Tensor, qr: torch.Tensor, mask=None, packed_seq_param cu_ke = cu_ke[cp_rank * seqlen : (cp_rank + 1) * seqlen] index_scores = batched_indexer_fwd(q, k, weights.float(), cu_ks, cu_ke) + # index_scores: [batch, seqlen, n_kv]; topk over the KV dim. Route through the indexer + # replay manager (flattened to [n_tokens, n_kv], matching the record/replay convention) so + # RL replay can pin the rollout's top-k picks. get_topk_fn is transparent when disabled. topk_count = min(self.index_topk, index_scores.size(-1)) - topk_indices = index_scores.topk(topk_count, dim=-1)[1] + b, s, n_kv = index_scores.shape + + def _original_topk(scores, k): + return scores.topk(k, dim=-1)[1] + + topk_fn = indexer_replay_manager.get_topk_fn(_original_topk, return_probs=False) + topk_indices = topk_fn(index_scores.reshape(b * s, n_kv), topk_count).reshape(b, s, topk_count) return topk_indices From c1064abf16affc0bfa1871380bb6747ab209bb45 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Tue, 23 Jun 2026 23:02:15 +0000 Subject: [PATCH 07/13] dsv4 step4: MILES_DSV4_ATTENTION_BACKEND selector (miles default, native=dsv4_hybrid) get_dsv4_spec picks the dsv4 attention module by env var: 'miles' (default) keeps miles' DeepSeekV4Attention (BSHD + sparse-CP + tilelang, validated); 'native' swaps in dev's dsv4_hybrid sparse attention. Forces miles when TP>1 or CP>1 (dev dsv4_hybrid asserts TP==1, no sparse-CP). Default path verified unchanged (loss 19.92); dev dsv4_hybrid validated via test_dsv4_hybrid_native_parity (unfused passes on H200/SM9; fused needs SM10+). Native end-to-end with a checkpoint additionally needs mbridge dev-name mapping + dsv4_hybrid config coordination (opt-in follow-up; miles is the production path). --- .../models/deepseek_v4/deepseek_v4.py | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/miles_plugins/models/deepseek_v4/deepseek_v4.py b/miles_plugins/models/deepseek_v4/deepseek_v4.py index d0898662988..debbdccf827 100644 --- a/miles_plugins/models/deepseek_v4/deepseek_v4.py +++ b/miles_plugins/models/deepseek_v4/deepseek_v4.py @@ -1,4 +1,5 @@ import copy +import logging import os import einops @@ -24,6 +25,8 @@ from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint +logger = logging.getLogger(__name__) + from miles_plugins.models.deepseek_v4.ops.compressor import DeepSeekV4Compressor from miles_plugins.models.deepseek_v4.ops.cp_utils import ( all_gather_cp, @@ -336,14 +339,40 @@ def _dsv4_attention_module_spec(config, backend=None): ) +def _select_dsv4_attention_backend(config) -> str: + """Pick the DSv4 attention backend for the `dsv4` variant. + + `MILES_DSV4_ATTENTION_BACKEND in {miles, native}` (default `miles`). `native` selects dev's + `dsv4_hybrid` sparse-attention module (flash_mla/cuDNN-DSA or unfused PyTorch); it asserts TP==1 + and has no sparse-CP, so we force `miles` (BSHD + full sparse-CP + tilelang, the validated + production path) whenever TP>1 or CP>1. `miles` stays the default so a `native` limitation never + blocks production. + """ + backend = os.environ.get("MILES_DSV4_ATTENTION_BACKEND", "miles").lower() + if backend not in ("miles", "native"): + raise ValueError(f"MILES_DSV4_ATTENTION_BACKEND must be 'miles' or 'native', got {backend!r}") + if backend == "native" and (config.tensor_model_parallel_size > 1 or config.context_parallel_size > 1): + logger.warning( + "MILES_DSV4_ATTENTION_BACKEND=native is unsupported with TP>1 or CP>1 " + "(dev dsv4_hybrid asserts TP==1 / no sparse-CP); falling back to the miles backend." + ) + backend = "miles" + return backend + + def get_dsv4_spec(args, config, vp_stage): """ Usage: --spec miles_plugins.models.deepseek_v4.deepseek_v4 get_dsv4_spec """ _orig_get_spec = _eav_specs.get_experimental_attention_variant_module_spec + dsv4_backend = _select_dsv4_attention_backend(config) def _patched_get_spec(config, backend=None): if config.experimental_attention_variant == "dsv4": + if dsv4_backend == "native": + # dev's native DSv4-hybrid sparse attention (csa + CSAIndexer). The surrounding + # block (miles mHC, hash-routed MoE) is unchanged; only the attention module swaps. + return _eav_specs.get_dsv4_hybrid_module_spec_for_backend(config=config, backend=backend) return _dsv4_attention_module_spec(config, backend) return _orig_get_spec(config, backend) From 55fbb41c0166dc35a449a387bb423bd8cbfbaf65 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Tue, 23 Jun 2026 23:03:45 +0000 Subject: [PATCH 08/13] test(dumper): retarget _forward_mlp anchors to dev's _forward_mlp_output_with_bias New dev split the MLP into _forward_mlp_output_with_bias + _forward_post_mlp and dropped _sglang_pre_mlp_residual (true-on-policy was deferred). Anchor pre_mlp_residual on the unconditional nvtx_range_push("mlp") (residual set in all branches); keep pre_mlp_layernorm_output and mlp_output_with_bias anchors. Will re-verify after the true-on-policy migration. --- tests/e2e/conftest_dumper.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/e2e/conftest_dumper.py b/tests/e2e/conftest_dumper.py index dbe4322a849..0a63a126500 100644 --- a/tests/e2e/conftest_dumper.py +++ b/tests/e2e/conftest_dumper.py @@ -48,10 +48,10 @@ append: "dumper.dump('layer_input', hidden_states, dims='t[cp:zigzag,sp] 1 h # tp:replicated ep:replicated')" - match: "nvtx_range_pop(suffix=\\"self_attention\\")" append: "dumper.dump('attn_output', attention_output_with_bias[0], dims='t[cp:zigzag,sp] 1 h # tp:replicated ep:replicated')" - - target: megatron.core.transformer.transformer_layer.TransformerLayer._forward_mlp + - target: megatron.core.transformer.transformer_layer.TransformerLayer._forward_mlp_output_with_bias edits: - - match: 'residual = getattr(self, "_sglang_pre_mlp_residual", hidden_states)' - append: "dumper.dump('pre_mlp_residual', residual, dims='t[cp:zigzag,sp] 1 h # tp:replicated ep:replicated')" + - match: 'nvtx_range_push(suffix="mlp")' + prepend: "dumper.dump('pre_mlp_residual', residual, dims='t[cp:zigzag,sp] 1 h # tp:replicated ep:replicated')" - match: "pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states)" append: "dumper.dump('pre_mlp_layernorm_output', pre_mlp_layernorm_output, dims='t[cp:zigzag,sp] 1 h # tp:replicated ep:replicated')" - match: "mlp_output_with_bias = (mlp_output, mlp_output_bias)" From 209c8d1ff171042933d645eefdd028fb059e85bf Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Wed, 24 Jun 2026 00:32:23 +0000 Subject: [PATCH 09/13] test(dumper): retarget BSHD pre-mlp dump to _forward_mlp_output_with_bias The megatron true-on-policy merge relocated the _sglang_pre_mlp_residual extraction out of _forward_mlp into _forward_mlp_output_with_bias, so the BSHD source-patcher anchor 'residual = getattr(self, "_sglang_pre_mlp_residual", hidden_states)' no longer exists. Mirror the THD config: target _forward_mlp_output_with_bias and prepend the pre_mlp_residual dump at the mlp nvtx push (matches the new structure). Fixes test_run_megatron + test_miles_dumper source-patch failures. --- tests/e2e/conftest_dumper.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/e2e/conftest_dumper.py b/tests/e2e/conftest_dumper.py index 0a63a126500..02e4e971364 100644 --- a/tests/e2e/conftest_dumper.py +++ b/tests/e2e/conftest_dumper.py @@ -87,10 +87,10 @@ append: "dumper.dump('layer_input', hidden_states, dims='s[cp:zigzag,sp] 1 h # tp:replicated ep:replicated')" - match: "nvtx_range_pop(suffix=\\"self_attention\\")" append: "dumper.dump('attn_output', attention_output_with_bias[0], dims='s[cp:zigzag,sp] 1 h # tp:replicated ep:replicated')" - - target: megatron.core.transformer.transformer_layer.TransformerLayer._forward_mlp + - target: megatron.core.transformer.transformer_layer.TransformerLayer._forward_mlp_output_with_bias edits: - - match: 'residual = getattr(self, "_sglang_pre_mlp_residual", hidden_states)' - append: "dumper.dump('pre_mlp_residual', residual, dims='s[cp:zigzag,sp] 1 h # tp:replicated ep:replicated')" + - match: 'nvtx_range_push(suffix="mlp")' + prepend: "dumper.dump('pre_mlp_residual', residual, dims='s[cp:zigzag,sp] 1 h # tp:replicated ep:replicated')" - match: "pre_mlp_layernorm_output = self._forward_pre_mlp_layernorm(hidden_states)" append: "dumper.dump('pre_mlp_layernorm_output', pre_mlp_layernorm_output, dims='s[cp:zigzag,sp] 1 h # tp:replicated ep:replicated')" - match: "mlp_output_with_bias = (mlp_output, mlp_output_bias)" From 7201fa5808fd372c6f3923dbd2c79ea61f04ccc8 Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Wed, 24 Jun 2026 01:24:01 +0000 Subject: [PATCH 10/13] fix: accept 'name' kwarg in custom attention __init__ (megatron build_module) New dev's build_module passes name= to every spec-built module (megatron base Attention.__init__ accepts it). miles' custom attention classes overriding __init__ must accept + thread it: DSAMultiLatentAttention/DSAMLASelfAttention (glm5, DeepSeek-V3.2), HuggingfaceAttention base, qwen3_5/qwen3_next Attention. Mirrors the DeepSeekV4Attention name fix. Fixes 'DSAMLASelfAttention.__init__() got an unexpected keyword argument name'. --- miles_plugins/models/glm5/glm5.py | 4 ++++ miles_plugins/models/hf_attention.py | 1 + miles_plugins/models/qwen3_5.py | 2 ++ miles_plugins/models/qwen3_next.py | 2 ++ 4 files changed, 9 insertions(+) diff --git a/miles_plugins/models/glm5/glm5.py b/miles_plugins/models/glm5/glm5.py index d19f54b10a1..774e73861cf 100644 --- a/miles_plugins/models/glm5/glm5.py +++ b/miles_plugins/models/glm5/glm5.py @@ -92,6 +92,7 @@ def __init__( cp_comm_type: str | None = None, model_comm_pgs=None, pg_collection=None, + name: str | None = None, ) -> None: super().__init__( @@ -102,6 +103,7 @@ def __init__( attn_mask_type=attn_mask_type, cp_comm_type=cp_comm_type, pg_collection=pg_collection, + name=name, ) self.query_projection_size = self.config.v_head_dim * self.config.num_attention_heads @@ -333,6 +335,7 @@ def __init__( cp_comm_type: str | None = None, model_comm_pgs=None, pg_collection=None, + name: str | None = None, ): super().__init__( config=config, @@ -344,6 +347,7 @@ def __init__( cp_comm_type=cp_comm_type, model_comm_pgs=model_comm_pgs, pg_collection=pg_collection, + name=name, ) q_down_proj_kwargs = {} if submodules.linear_q_down_proj in [TELinear]: diff --git a/miles_plugins/models/hf_attention.py b/miles_plugins/models/hf_attention.py index 236a698db4a..b152a57a3e3 100644 --- a/miles_plugins/models/hf_attention.py +++ b/miles_plugins/models/hf_attention.py @@ -161,6 +161,7 @@ def __init__( layer_number: int, cp_comm_type: str = "p2p", pg_collection=None, + name: str | None = None, ): super().__init__(config=config) self.args = args diff --git a/miles_plugins/models/qwen3_5.py b/miles_plugins/models/qwen3_5.py index 8c13c17d204..70a107097e4 100644 --- a/miles_plugins/models/qwen3_5.py +++ b/miles_plugins/models/qwen3_5.py @@ -190,6 +190,7 @@ def __init__( layer_number: int, cp_comm_type: str = "p2p", pg_collection=None, + name: str | None = None, ): super().__init__( args, @@ -197,6 +198,7 @@ def __init__( layer_number, cp_comm_type, pg_collection, + name=name, ) # Qwen3.5 is a VLM model with nested text_config self.hf_config = _get_text_config(self.hf_config) diff --git a/miles_plugins/models/qwen3_next.py b/miles_plugins/models/qwen3_next.py index adf1bdb3752..412d1149480 100644 --- a/miles_plugins/models/qwen3_next.py +++ b/miles_plugins/models/qwen3_next.py @@ -201,6 +201,7 @@ def __init__( layer_number: int, cp_comm_type: str = "p2p", pg_collection=None, + name: str | None = None, ): super().__init__( args, @@ -208,6 +209,7 @@ def __init__( layer_number, cp_comm_type, pg_collection, + name=name, ) if Qwen3NextAttention is None: raise ImportError("Please install transformers>=4.35.0 to use Qwen3NextAttention.") From c37c468371e685b8cc2c736ff316a67b30f9798b Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Wed, 24 Jun 2026 02:31:06 +0000 Subject: [PATCH 11/13] fix: stop passing removed moe_use_legacy_grouped_gemm arg to spec builders New dev removed the moe_use_legacy_grouped_gemm argument and the legacy grouped-GEMM concept (grouped_mlp_modules now takes only moe_use_grouped_gemm). miles' model_provider and glm4 spec builders read args.moe_use_legacy_grouped_gemm, which no longer exists -> AttributeError on the Namespace. Drop the kwarg (dev spec fns take **kwargs / the new grouped path is the only one). --- miles/backends/megatron_utils/model_provider.py | 2 -- miles_plugins/models/glm4.py | 1 - 2 files changed, 3 deletions(-) diff --git a/miles/backends/megatron_utils/model_provider.py b/miles/backends/megatron_utils/model_provider.py index 0f2ecf5aed4..52a3adee838 100644 --- a/miles/backends/megatron_utils/model_provider.py +++ b/miles/backends/megatron_utils/model_provider.py @@ -186,7 +186,6 @@ def model_provider( moe_grouped_gemm=args.moe_grouped_gemm, qk_layernorm=args.qk_layernorm, multi_latent_attention=args.multi_latent_attention, - moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm, ) else: transformer_layer_spec = get_gpt_layer_local_spec( @@ -194,7 +193,6 @@ def model_provider( moe_grouped_gemm=args.moe_grouped_gemm, qk_layernorm=args.qk_layernorm, multi_latent_attention=args.multi_latent_attention, - moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm, normalization=args.normalization, use_kitchen=config.use_kitchen, use_true_on_policy_backend=config.true_on_policy_contract is not None, diff --git a/miles_plugins/models/glm4.py b/miles_plugins/models/glm4.py index ba42ea1a6eb..47377143482 100644 --- a/miles_plugins/models/glm4.py +++ b/miles_plugins/models/glm4.py @@ -7,7 +7,6 @@ def get_glm_spec(args, config, vp_stage): moe_grouped_gemm=args.moe_grouped_gemm, qk_layernorm=args.qk_layernorm, multi_latent_attention=args.multi_latent_attention, - moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm, post_self_attn_layernorm=args.post_self_attn_layernorm, post_mlp_layernorm=args.post_mlp_layernorm, ) From 7cfa93eb52b85b410fc6eaae06c5d3ef36d766fb Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Wed, 24 Jun 2026 04:26:33 +0000 Subject: [PATCH 12/13] fix: handle MTP submodule rename transformer_layer -> mtp_model_layer New Megatron renamed MultiTokenPredictionLayerSubmodules.transformer_layer -> mtp_model_layer, so MTP weight names are now mtp.layers.{i}.mtp_model_layer.*. miles' MTP weight converters (mcore<->hf) hardcoded 'transformer_layer' -> 'Unknown MTP component' on rollout weight export (test_mimo_7B_mtp_only_grad). Accept both names across export (mimo/qwen3_next/deepseekv3/glm4moe/qwen3_5 + fp8/ mxfp8/nvfp4 quantizers) and mbridge import (mimo/qwen3_next/qwen3_5/glm4moe_lite/glm4moe). --- .../megatron_utils/megatron_to_hf/deepseekv3.py | 3 ++- .../backends/megatron_utils/megatron_to_hf/glm4moe.py | 3 ++- miles/backends/megatron_utils/megatron_to_hf/mimo.py | 11 +++++++---- .../megatron_to_hf/processors/quantizer_fp8.py | 2 +- .../megatron_to_hf/processors/quantizer_mxfp8.py | 2 +- .../megatron_to_hf/processors/quantizer_nvfp4.py | 2 +- .../backends/megatron_utils/megatron_to_hf/qwen3_5.py | 5 +++-- .../megatron_utils/megatron_to_hf/qwen3_next.py | 7 ++++--- miles_plugins/mbridge/glm4moe.py | 4 ++-- miles_plugins/mbridge/glm4moe_lite.py | 7 +++++-- miles_plugins/mbridge/mimo.py | 11 ++++++----- miles_plugins/mbridge/qwen3_5.py | 5 +++-- miles_plugins/mbridge/qwen3_next.py | 5 +++-- 13 files changed, 40 insertions(+), 27 deletions(-) diff --git a/miles/backends/megatron_utils/megatron_to_hf/deepseekv3.py b/miles/backends/megatron_utils/megatron_to_hf/deepseekv3.py index 22aeeeb3f0e..a7cb1fca812 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/deepseekv3.py +++ b/miles/backends/megatron_utils/megatron_to_hf/deepseekv3.py @@ -149,7 +149,8 @@ def convert_deepseekv3_to_hf(args, name, param): return [(f"model.layers.{layer_idx}.shared_head.norm.weight", param)] else: name = f"module.module.decoder.layers.{layer_idx}.{rest}" - name = name.replace("transformer_layer.", "") + # New Megatron renamed the MTP submodule transformer_layer -> mtp_model_layer. + name = name.replace("transformer_layer.", "").replace("mtp_model_layer.", "") return convert_deepseekv3_to_hf(args, name, param) raise ValueError(f"Unknown parameter name: {name}") diff --git a/miles/backends/megatron_utils/megatron_to_hf/glm4moe.py b/miles/backends/megatron_utils/megatron_to_hf/glm4moe.py index 33a64e6e8f0..fd60770cc6a 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/glm4moe.py +++ b/miles/backends/megatron_utils/megatron_to_hf/glm4moe.py @@ -133,7 +133,8 @@ def convert_glm4moe_to_hf(args, name, param): return [(f"model.layers.{layer_idx}.shared_head.norm.weight", param)] else: name = f"module.module.decoder.layers.{layer_idx}.{rest}" - name = name.replace("transformer_layer.", "") + # New Megatron renamed the MTP submodule transformer_layer -> mtp_model_layer. + name = name.replace("transformer_layer.", "").replace("mtp_model_layer.", "") return convert_glm4moe_to_hf(args, name, param) raise ValueError(f"Unknown parameter name: {name}") diff --git a/miles/backends/megatron_utils/megatron_to_hf/mimo.py b/miles/backends/megatron_utils/megatron_to_hf/mimo.py index 3d9c6c49172..2931968b38a 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/mimo.py +++ b/miles/backends/megatron_utils/megatron_to_hf/mimo.py @@ -52,10 +52,13 @@ def convert_mimo_mtp_param(args, name, param): if component in direct_mappings: return [(direct_mappings[component], param)] - # Handle transformer_layer components - if component.startswith("transformer_layer."): - # Remove "transformer_layer." prefix - transformer_component = component[len("transformer_layer.") :] + # Handle the wrapped transformer-layer components. New Megatron renamed the MTP + # submodule field `transformer_layer` -> `mtp_model_layer`; accept both. + transformer_layer_prefixes = ("transformer_layer.", "mtp_model_layer.") + matched_prefix = next((p for p in transformer_layer_prefixes if component.startswith(p)), None) + if matched_prefix is not None: + # Remove the wrapped-layer prefix + transformer_component = component[len(matched_prefix) :] # Create proxy name for reusing existing Qwen2 conversion functions proxy_name = f"module.module.decoder.layers.{layer_idx}.{transformer_component}" diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py index e3948590022..c15546a8283 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py @@ -29,7 +29,7 @@ def quantize_params_fp8(args, megatron_name, converted_named_params, quantizatio if not match: return converted_named_params layer_idx, rest = match.groups() - rest = rest.replace("transformer_layer.", "") + rest = rest.replace("transformer_layer.", "").replace("mtp_model_layer.", "") else: layer_idx, rest = match.groups() diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_mxfp8.py b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_mxfp8.py index e6e86e4457c..7d8863a5d90 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_mxfp8.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_mxfp8.py @@ -33,7 +33,7 @@ def quantize_params_mxfp8(args, megatron_name, converted_named_params, quantizat if not match: return converted_named_params layer_idx, rest = match.groups() - rest = rest.replace("transformer_layer.", "") + rest = rest.replace("transformer_layer.", "").replace("mtp_model_layer.", "") else: layer_idx, rest = match.groups() diff --git a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py index 6432d5ebed4..f582c4405d4 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py +++ b/miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_nvfp4.py @@ -56,7 +56,7 @@ def quantize_params_nvfp4(args, megatron_name, converted_named_params, quantizat if not match: return converted_named_params layer_idx, rest = match.groups() - rest = rest.replace("transformer_layer.", "") + rest = rest.replace("transformer_layer.", "").replace("mtp_model_layer.", "") else: layer_idx, rest = match.groups() diff --git a/miles/backends/megatron_utils/megatron_to_hf/qwen3_5.py b/miles/backends/megatron_utils/megatron_to_hf/qwen3_5.py index 2aabd86ebac..70d155f76e4 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/qwen3_5.py +++ b/miles/backends/megatron_utils/megatron_to_hf/qwen3_5.py @@ -14,8 +14,9 @@ def _convert_mtp_layer(args, name, param, layer_idx): if "eh_proj.weight" in name: return [("mtp.fc.weight", param)] - if "transformer_layer" in name: - proxy_name = name.replace(f"mtp.layers.{layer_idx}.transformer_layer", f"decoder.layers.{layer_idx}") + _mtp_inner = next((p for p in ("transformer_layer", "mtp_model_layer") if p in name), None) + if _mtp_inner is not None: + proxy_name = name.replace(f"mtp.layers.{layer_idx}.{_mtp_inner}", f"decoder.layers.{layer_idx}") mapped_params = convert_qwen3_5_to_hf(args, proxy_name, param) final_params = [] diff --git a/miles/backends/megatron_utils/megatron_to_hf/qwen3_next.py b/miles/backends/megatron_utils/megatron_to_hf/qwen3_next.py index f248cc0fc31..a2b1a1526bf 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/qwen3_next.py +++ b/miles/backends/megatron_utils/megatron_to_hf/qwen3_next.py @@ -155,9 +155,10 @@ def convert_qwen3_next_to_hf(args, name, param): elif rest == "final_layernorm.weight": return [("mtp.norm.weight", param)] - # transformer_layer components → reuse decoder conversion with mtp prefix - if rest.startswith("transformer_layer."): - transformer_rest = rest[len("transformer_layer.") :] + # transformer_layer components → reuse decoder conversion with mtp prefix. + # New Megatron renamed the MTP submodule transformer_layer -> mtp_model_layer. + if rest.startswith(("transformer_layer.", "mtp_model_layer.")): + transformer_rest = rest.split(".", 1)[1] proxy_name = f"module.module.decoder.layers.{layer_idx}.{transformer_rest}" results = convert_qwen3_next_to_hf(args, proxy_name, param) return [ diff --git a/miles_plugins/mbridge/glm4moe.py b/miles_plugins/mbridge/glm4moe.py index 8bfd65aa172..3b249a3077d 100644 --- a/miles_plugins/mbridge/glm4moe.py +++ b/miles_plugins/mbridge/glm4moe.py @@ -41,14 +41,14 @@ def _weight_name_mapping_mtp(self, name: str, num_layers: int) -> str: elif "mlp" in name: mtp_layer_index = int(re.findall(r"mtp\.layers\.(\d+)\.", name)[0]) name_ = re.sub( - r"^mtp\.layers.\d+.transformer_layer", f"model.layers.{num_layers+mtp_layer_index}", name + r"^mtp\.layers.\d+.(?:transformer_layer|mtp_model_layer)", f"model.layers.{num_layers+mtp_layer_index}", name ) convert_names = self._weight_name_mapping_mlp(name_) break elif "self_attention" in name: mtp_layer_index = int(re.findall(r"mtp\.layers.(\d+)\.", name)[0]) name_ = re.sub( - r"^mtp\.layers.\d+.transformer_layer", f"model.layers.{num_layers+mtp_layer_index}", name + r"^mtp\.layers.\d+.(?:transformer_layer|mtp_model_layer)", f"model.layers.{num_layers+mtp_layer_index}", name ) convert_names = self._weight_name_mapping_attention(name_) break diff --git a/miles_plugins/mbridge/glm4moe_lite.py b/miles_plugins/mbridge/glm4moe_lite.py index c44888cd45a..bbe566d5a3d 100644 --- a/miles_plugins/mbridge/glm4moe_lite.py +++ b/miles_plugins/mbridge/glm4moe_lite.py @@ -130,8 +130,11 @@ def _convert_mtp_param(self, name: str) -> tuple[list[str]]: if name in direct_name_mapping: return [direct_name_mapping[name]] - assert "mtp.layers.0.transformer_layer" in name, "mtp not found" - proxy_name = name.replace("mtp.layers.0.transformer_layer", f"decoder.layers.{mtp_layer_id}") + _mtp_inner = next( + (p for p in ("transformer_layer", "mtp_model_layer") if f"mtp.layers.0.{p}" in name), None + ) + assert _mtp_inner is not None, "mtp not found" + proxy_name = name.replace(f"mtp.layers.0.{_mtp_inner}", f"decoder.layers.{mtp_layer_id}") if "self_attention" in proxy_name or "input_layernorm.weight" in proxy_name: return self._weight_name_mapping_attention(proxy_name) if "mlp" in proxy_name: diff --git a/miles_plugins/mbridge/mimo.py b/miles_plugins/mbridge/mimo.py index a45e7cfdfbb..06799bf8c39 100644 --- a/miles_plugins/mbridge/mimo.py +++ b/miles_plugins/mbridge/mimo.py @@ -76,13 +76,14 @@ def _convert_mtp_param(self, name: str) -> list[str]: if name in direct_name_mapping: return [direct_name_mapping[name]] - # Handle transformer components within MTP - # Check if this is a transformer_layer component - if "transformer_layer" in name: + # Handle transformer components within MTP. New Megatron renamed the MTP submodule + # transformer_layer -> mtp_model_layer; accept both. + _mtp_inner = next((p for p in ("transformer_layer", "mtp_model_layer") if p in name), None) + if _mtp_inner is not None: # Create a proxy name to use with parent class methods - # Convert mtp.layers.{idx}.transformer_layer.* to decoder.layers.{idx}.* + # Convert mtp.layers.{idx}..* to decoder.layers.{idx}.* proxy_name = name.replace( - f"mtp.layers.{mtp_layer_idx}.transformer_layer", + f"mtp.layers.{mtp_layer_idx}.{_mtp_inner}", f"decoder.layers.{mtp_layer_idx}", ) diff --git a/miles_plugins/mbridge/qwen3_5.py b/miles_plugins/mbridge/qwen3_5.py index fe1556ef7d9..7df2daa95ae 100644 --- a/miles_plugins/mbridge/qwen3_5.py +++ b/miles_plugins/mbridge/qwen3_5.py @@ -269,9 +269,10 @@ def _convert_mtp_param(self, name: str) -> list[str]: if name in direct_name_mapping: return [direct_name_mapping[name]] - if "transformer_layer" in name: + _mtp_inner = next((p for p in ("transformer_layer", "mtp_model_layer") if p in name), None) + if _mtp_inner is not None: proxy_name = name.replace( - f"mtp.layers.{mtp_layer_idx}.transformer_layer", + f"mtp.layers.{mtp_layer_idx}.{_mtp_inner}", f"decoder.layers.{mtp_layer_idx}", ) diff --git a/miles_plugins/mbridge/qwen3_next.py b/miles_plugins/mbridge/qwen3_next.py index 953bd29027c..77940752f36 100644 --- a/miles_plugins/mbridge/qwen3_next.py +++ b/miles_plugins/mbridge/qwen3_next.py @@ -130,9 +130,10 @@ def _convert_mtp_param(self, name: str) -> list[str]: if name in direct_mappings: return [direct_mappings[name]] - if "transformer_layer" in name: + _mtp_inner = next((p for p in ("transformer_layer", "mtp_model_layer") if p in name), None) + if _mtp_inner is not None: proxy_name = name.replace( - f"mtp.layers.{mtp_layer_idx}.transformer_layer", + f"mtp.layers.{mtp_layer_idx}.{_mtp_inner}", f"decoder.layers.{mtp_layer_idx}", ) From 75096707f2781352e33492b517a4731df988beff Mon Sep 17 00:00:00 2001 From: Yueming Yuan Date: Wed, 24 Jun 2026 05:16:57 +0000 Subject: [PATCH 13/13] fix: use new Megatron MTP loss tracker API (reduce_loss_in_tracker) New dev rewrote MTPLossLoggingHelper: it stores loss_sums/num_tokens (or loss_values) and only computes tracker['values'] after reduce_loss_in_tracker() (which also does the cross-rank all-reduce). miles read tracker['values'] directly (never populated) -> 'values' not in tracker -> mtp_losses unbound -> UnboundLocalError at extra_metrics. Call reduce_loss_in_tracker(), read values/loss_values, init mtp_losses=None + guard the metric. Fixes test_mimo_7B_mtp_only_grad past weight-export. --- miles/backends/megatron_utils/model.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 4df67e8aadf..57fd13b22e6 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -643,19 +643,23 @@ def train( config.param_sync_func = param_sync_func pre_hook_enabled = True + mtp_losses = None if args.enable_mtp_training: from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper mtp_loss_scale = 1 / num_microbatches[step_id] + # New Megatron tracks MTP loss as loss_sums/num_tokens (or loss_values) rather than a + # pre-divided "values" tensor. reduce_loss_in_tracker() all-reduces across ranks + # (collective: call on all ranks) and computes the per-token loss into tracker["values"]. + MTPLossLoggingHelper.reduce_loss_in_tracker() tracker = MTPLossLoggingHelper.tracker + # here we assume only one mtp layer if "values" in tracker: - values = tracker["values"] - if (x := tracker.get("reduce_group")) is not None: - torch.distributed.all_reduce(values, group=x) - if (x := tracker.get("avg_group")) is not None: - torch.distributed.all_reduce(values, group=x, op=torch.distributed.ReduceOp.AVG) - # here we assume only one mtp layer mtp_losses = (tracker["values"] * mtp_loss_scale).item() + elif "loss_values" in tracker: + mtp_losses = (tracker["loss_values"] * mtp_loss_scale).item() + + if mtp_losses is not None: MTPLossLoggingHelper.clean_loss_in_tracker() # CI check: verify MTP loss is within expected bounds @@ -671,7 +675,7 @@ def train( role_tag = "" if role == "actor" else f"{role}-" extra_metrics = {} - if args.enable_mtp_training: + if args.enable_mtp_training and mtp_losses is not None: extra_metrics["mtp_loss"] = mtp_losses if not disable_optimizer: