From d1d486ad09e8e355c74b53b83497a2afdc46c19d Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Wed, 1 Apr 2026 03:26:53 +0000 Subject: [PATCH 1/4] Fix Qwen3.5/Qwen3-next support for unsupported HF model types - Add fallback config loading when transformers doesn't recognize the model_type (reads config.json directly and constructs a namespace) - Fix all-gather backward pass for duplicated computation (use custom autograd function that returns local gradient instead of reduce-scatter) - Add tensor_parallel_output_grad=False for sequence parallel gather - Compute layer_types from full_attention_interval when config doesn't expose it directly - Patch tied lm_head in checkpoint converter when lm_head.weight is absent from HF weights - Proactively clear memory when free GPU memory is low before distributed calls Made-with: Cursor --- miles/utils/reloadable_process_group.py | 5 +- miles_plugins/models/hf_attention.py | 68 ++++++++++++++++++++-- miles_plugins/models/qwen3_5.py | 16 ++++-- miles_plugins/models/qwen3_next.py | 11 +++- tools/convert_hf_to_torch_dist.py | 75 ++++++++++++++++++++++++- 5 files changed, 160 insertions(+), 15 deletions(-) diff --git a/miles/utils/reloadable_process_group.py b/miles/utils/reloadable_process_group.py index af4878b4e71..55d9731e84b 100644 --- a/miles/utils/reloadable_process_group.py +++ b/miles/utils/reloadable_process_group.py @@ -5,7 +5,7 @@ import torch import torch.distributed as dist -from miles.utils.memory_utils import print_memory +from miles.utils.memory_utils import available_memory, clear_memory, print_memory logger = logging.getLogger(__name__) @@ -275,6 +275,9 @@ def reload_process_groups(): @contextmanager def _wrap_low_level_call(): try: + mem_info = available_memory() + if mem_info["free_GB"] < 3: + clear_memory() yield except Exception as e: mem_info = print_memory("after torch distributed error") diff --git a/miles_plugins/models/hf_attention.py b/miles_plugins/models/hf_attention.py index c353ae7b29d..d3858456fbe 100644 --- a/miles_plugins/models/hf_attention.py +++ b/miles_plugins/models/hf_attention.py @@ -1,3 +1,5 @@ +import json +import os from abc import ABC, abstractmethod import torch @@ -6,7 +8,56 @@ from megatron.core.inference.contexts import BaseInferenceContext from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.transformer.module import MegatronModule -from transformers import AutoConfig + + +def _load_hf_config(checkpoint_path): + """Load HF config with fallback for unsupported model types.""" + try: + from transformers import AutoConfig + + return AutoConfig.from_pretrained(checkpoint_path, trust_remote_code=True) + except (ValueError, KeyError): + config_path = os.path.join(checkpoint_path, "config.json") + with open(config_path) as f: + config_dict = json.load(f) + + _DTYPE_MAP = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32} + + def _fix_dtype(d): + if "torch_dtype" in d: + d["torch_dtype"] = _DTYPE_MAP.get(d["torch_dtype"], d["torch_dtype"]) + if "dtype" in d: + d["dtype"] = _DTYPE_MAP.get(d["dtype"], d["dtype"]) + + _fix_dtype(config_dict) + ns = type("HFConfig", (), config_dict)() + if "text_config" in config_dict: + _fix_dtype(config_dict["text_config"]) + ns.text_config = type("TextConfig", (), config_dict["text_config"])() + return ns + + +class _AllGatherForDuplicatedComputation(torch.autograd.Function): + """All-gather whose backward just returns the local gradient slice (no reduce). + + Use this instead of ``dist.nn.all_gather`` when the computation after the + gather is *duplicated* across ranks (same weights, same full input -> + identical gradients). The default ``all_gather`` backward performs a + reduce-scatter, which incorrectly sums ``world_size`` identical copies of + the gradient. + """ + + @staticmethod + def forward(ctx, x, group): + ctx.group = group + ctx.rank = dist.get_rank(group=group) + out = [torch.empty_like(x) for _ in range(dist.get_world_size(group=group))] + dist.all_gather(out, x.contiguous(), group=group) + return tuple(out) + + @staticmethod + def backward(ctx, *grads): + return grads[ctx.rank], None class HuggingfaceAttention(MegatronModule, ABC): @@ -30,7 +81,7 @@ def __init__( # Note that megatron layer_number starts at 1 self.layer_number = layer_number self.hf_layer_idx = layer_number - 1 - self.hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True) + self.hf_config = _load_hf_config(args.hf_checkpoint) # hardcode to fa2 at the moment. self.hf_config._attn_implementation = "flash_attention_2" @@ -54,15 +105,22 @@ def forward( cu_seqlens = packed_seq_params.cu_seqlens_q if self.args.sequence_parallel: + # tensor_parallel_output_grad=False: the linear attention after this + # gather is NOT TP-sharded (duplicated on all ranks), so the backward + # should split (not reduce-scatter) to avoid inflating gradients by TP. hidden_states = tensor_parallel.gather_from_sequence_parallel_region( - hidden_states, group=mpu.get_tensor_model_parallel_group() + hidden_states, + tensor_parallel_output_grad=False, + group=mpu.get_tensor_model_parallel_group(), ) if mpu.get_context_parallel_world_size() > 1: cp_size = mpu.get_context_parallel_world_size() - hidden_states_list = dist.nn.all_gather( + # Use custom all-gather whose backward returns local gradient + # instead of reduce-scatter, since the computation is duplicated. + hidden_states_list = _AllGatherForDuplicatedComputation.apply( hidden_states, - group=mpu.get_context_parallel_group(), + mpu.get_context_parallel_group(), ) # TODO: preprocess this for each batch to prevent tolist in the training step diff --git a/miles_plugins/models/qwen3_5.py b/miles_plugins/models/qwen3_5.py index b8455fe9e0f..a796c8c49c5 100644 --- a/miles_plugins/models/qwen3_5.py +++ b/miles_plugins/models/qwen3_5.py @@ -7,7 +7,6 @@ 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 transformers import AutoConfig from transformers.activations import ACT2FN try: @@ -16,11 +15,7 @@ except ImportError: pass -from .hf_attention import HuggingfaceAttention - - -def _load_hf_config(hf_checkpoint): - return AutoConfig.from_pretrained(hf_checkpoint, trust_remote_code=True) +from .hf_attention import HuggingfaceAttention, _load_hf_config def _get_text_config(hf_config): @@ -132,6 +127,7 @@ def forward( initial_state=None, output_final_state=False, use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, ) z_shape_og = z.shape @@ -209,6 +205,14 @@ def get_qwen3_5_spec(args, config, vp_stage): hf_config = _load_hf_config(args.hf_checkpoint) text_config = _get_text_config(hf_config) + # Compute layer_types if the config class doesn't expose it + if not hasattr(text_config, "layer_types"): + interval = getattr(text_config, "full_attention_interval", 4) + n = text_config.num_hidden_layers + text_config.layer_types = [ + "full_attention" if (i + 1) % interval == 0 else "linear_attention" for i in range(n) + ] + for layer_id in range(num_layers_to_build): if text_config.layer_types[layer_id + offset] == "linear_attention": layer_specs = copy.deepcopy(transformer_layer_spec.layer_specs[layer_id]) diff --git a/miles_plugins/models/qwen3_next.py b/miles_plugins/models/qwen3_next.py index 482dcfe92f4..92e39ff318d 100644 --- a/miles_plugins/models/qwen3_next.py +++ b/miles_plugins/models/qwen3_next.py @@ -7,9 +7,10 @@ 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 transformers import AutoConfig from transformers.activations import ACT2FN +from .hf_attention import _load_hf_config + try: from fla.modules import FusedRMSNormGated, ShortConvolution from fla.ops.gated_delta_rule import chunk_gated_delta_rule @@ -214,7 +215,13 @@ def get_qwen3_next_spec(args, config, vp_stage): num_layers_to_build = get_num_layers_to_build(config, vp_stage=vp_stage) offset = get_transformer_layer_offset(config, vp_stage=vp_stage) - hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True) + hf_config = _load_hf_config(args.hf_checkpoint) + + # Compute layer_types if the config class doesn't expose it + if not hasattr(hf_config, "layer_types"): + interval = getattr(hf_config, "full_attention_interval", 4) + n = hf_config.num_hidden_layers + hf_config.layer_types = ["full_attention" if (i + 1) % interval == 0 else "linear_attention" for i in range(n)] for layer_id in range(num_layers_to_build): if hf_config.layer_types[layer_id + offset] == "linear_attention": diff --git a/tools/convert_hf_to_torch_dist.py b/tools/convert_hf_to_torch_dist.py index 9ff33e61307..0e0064f57ec 100644 --- a/tools/convert_hf_to_torch_dist.py +++ b/tools/convert_hf_to_torch_dist.py @@ -1,4 +1,5 @@ import gc +import json import os import shutil from functools import wraps @@ -20,6 +21,73 @@ from miles.utils.memory_utils import print_memory +def _get_hf_weight_names(checkpoint_path: str) -> set[str]: + index_path = os.path.join(checkpoint_path, "model.safetensors.index.json") + if os.path.exists(index_path): + with open(index_path) as f: + index = json.load(f) + return set(index.get("weight_map", {}).keys()) + + single_file = os.path.join(checkpoint_path, "model.safetensors") + if os.path.exists(single_file): + from safetensors import safe_open + + with safe_open(single_file, framework="pt") as f: + return set(f.keys()) + + return set() + + +def _load_hf_config_with_fallback(checkpoint_path: str): + """Load HF config with fallback for model types unknown to transformers.""" + try: + from transformers import AutoConfig + + return AutoConfig.from_pretrained(checkpoint_path, trust_remote_code=True) + except (ValueError, KeyError): + config_path = os.path.join(checkpoint_path, "config.json") + with open(config_path) as f: + config_dict = json.load(f) + + dtype_map = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32} + + def fix_dtype(d): + if "torch_dtype" in d: + d["torch_dtype"] = dtype_map.get(d["torch_dtype"], d["torch_dtype"]) + if "dtype" in d: + d["dtype"] = dtype_map.get(d["dtype"], d["dtype"]) + + fix_dtype(config_dict) + ns = type("HFConfig", (), config_dict)() + if "text_config" in config_dict: + fix_dtype(config_dict["text_config"]) + ns.text_config = type("TextConfig", (), config_dict["text_config"])() + return ns + + +def _patch_bridge_for_tied_lm_head(bridge, hf_model_path: str): + """If lm_head is absent in HF weights, fallback output_layer to embeddings.""" + direct_mapping = getattr(bridge, "_DIRECT_MAPPING", None) + if not isinstance(direct_mapping, dict): + return + + output_key = "output_layer.weight" + if direct_mapping.get(output_key) != "lm_head.weight": + return + + hf_weight_names = _get_hf_weight_names(hf_model_path) + if "lm_head.weight" in hf_weight_names: + return + + embedding_key = direct_mapping.get("embedding.word_embeddings.weight") + if embedding_key and embedding_key in hf_weight_names: + direct_mapping[output_key] = embedding_key + print( + f"[Patch] lm_head.weight not found in {hf_model_path}; " + f"fallback map output_layer.weight -> {embedding_key}" + ) + + def patch_weight_to_mcore_format_preserve_fp32(): original_method = Bridge._weight_to_mcore_format @@ -130,7 +198,12 @@ def main(): # Load model hf_model_path = args.hf_checkpoint - bridge = AutoBridge.from_pretrained(hf_model_path, trust_remote_code=True) + try: + bridge = AutoBridge.from_pretrained(hf_model_path, trust_remote_code=True) + except (ValueError, KeyError): + # Fallback for configs with model_type unknown to installed transformers. + bridge = AutoBridge.from_config(_load_hf_config_with_fallback(hf_model_path)) + _patch_bridge_for_tied_lm_head(bridge, hf_model_path) # Patch to preserve FP32 precision for _keep_fp32 params patch_weight_to_mcore_format_preserve_fp32() From a7d3e015cdf3451af834711aa775bdecda5043e9 Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Wed, 1 Apr 2026 18:44:29 +0000 Subject: [PATCH 2/4] fix impl --- miles_plugins/mbridge/qwen3_5.py | 17 +++++++++++++ tools/convert_hf_to_torch_dist.py | 41 ------------------------------- 2 files changed, 17 insertions(+), 41 deletions(-) diff --git a/miles_plugins/mbridge/qwen3_5.py b/miles_plugins/mbridge/qwen3_5.py index d0d0c5fd3b2..ee629d009f2 100644 --- a/miles_plugins/mbridge/qwen3_5.py +++ b/miles_plugins/mbridge/qwen3_5.py @@ -117,6 +117,23 @@ def _get_text_config(self): return self.hf_config.text_config return self.hf_config + def _is_tied_word_embeddings(self) -> bool: + tie_word_embeddings = getattr(self.hf_config, "tie_word_embeddings", None) + if tie_word_embeddings is None and hasattr(self.hf_config, "text_config"): + tie_word_embeddings = getattr(self.hf_config.text_config, "tie_word_embeddings", None) + return bool(tie_word_embeddings) + + def _adjust_mapping_for_shared_weights(self): + self._DIRECT_MAPPING = dict(self._DIRECT_MAPPING) + if self._is_tied_word_embeddings(): + embed_key = self._DIRECT_MAPPING["embedding.word_embeddings.weight"] + self._DIRECT_MAPPING["output_layer.weight"] = embed_key + + def _get_hf_shared_weight_keys(self) -> list[str]: + if self._is_tied_word_embeddings(): + return [self._DIRECT_MAPPING["embedding.word_embeddings.weight"]] + return [] + def _supports_transformer_config_kwarg(self, kwarg_name: str) -> bool: """Check whether the current TransformerConfig accepts a given kwarg.""" transformer_config_class = getattr(self, "TransformerConfigClass", None) diff --git a/tools/convert_hf_to_torch_dist.py b/tools/convert_hf_to_torch_dist.py index 0e0064f57ec..07475256954 100644 --- a/tools/convert_hf_to_torch_dist.py +++ b/tools/convert_hf_to_torch_dist.py @@ -21,23 +21,6 @@ from miles.utils.memory_utils import print_memory -def _get_hf_weight_names(checkpoint_path: str) -> set[str]: - index_path = os.path.join(checkpoint_path, "model.safetensors.index.json") - if os.path.exists(index_path): - with open(index_path) as f: - index = json.load(f) - return set(index.get("weight_map", {}).keys()) - - single_file = os.path.join(checkpoint_path, "model.safetensors") - if os.path.exists(single_file): - from safetensors import safe_open - - with safe_open(single_file, framework="pt") as f: - return set(f.keys()) - - return set() - - def _load_hf_config_with_fallback(checkpoint_path: str): """Load HF config with fallback for model types unknown to transformers.""" try: @@ -65,29 +48,6 @@ def fix_dtype(d): return ns -def _patch_bridge_for_tied_lm_head(bridge, hf_model_path: str): - """If lm_head is absent in HF weights, fallback output_layer to embeddings.""" - direct_mapping = getattr(bridge, "_DIRECT_MAPPING", None) - if not isinstance(direct_mapping, dict): - return - - output_key = "output_layer.weight" - if direct_mapping.get(output_key) != "lm_head.weight": - return - - hf_weight_names = _get_hf_weight_names(hf_model_path) - if "lm_head.weight" in hf_weight_names: - return - - embedding_key = direct_mapping.get("embedding.word_embeddings.weight") - if embedding_key and embedding_key in hf_weight_names: - direct_mapping[output_key] = embedding_key - print( - f"[Patch] lm_head.weight not found in {hf_model_path}; " - f"fallback map output_layer.weight -> {embedding_key}" - ) - - def patch_weight_to_mcore_format_preserve_fp32(): original_method = Bridge._weight_to_mcore_format @@ -203,7 +163,6 @@ def main(): except (ValueError, KeyError): # Fallback for configs with model_type unknown to installed transformers. bridge = AutoBridge.from_config(_load_hf_config_with_fallback(hf_model_path)) - _patch_bridge_for_tied_lm_head(bridge, hf_model_path) # Patch to preserve FP32 precision for _keep_fp32 params patch_weight_to_mcore_format_preserve_fp32() From e34fbabe38d497efa7851ddb8da7fb3c653d3ac3 Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Wed, 1 Apr 2026 20:31:52 +0000 Subject: [PATCH 3/4] fix(qwen35): revert rpg low-mem clear and reuse hf config loader --- miles/utils/reloadable_process_group.py | 5 +--- tools/convert_hf_to_torch_dist.py | 32 +++---------------------- 2 files changed, 4 insertions(+), 33 deletions(-) diff --git a/miles/utils/reloadable_process_group.py b/miles/utils/reloadable_process_group.py index 55d9731e84b..af4878b4e71 100644 --- a/miles/utils/reloadable_process_group.py +++ b/miles/utils/reloadable_process_group.py @@ -5,7 +5,7 @@ import torch import torch.distributed as dist -from miles.utils.memory_utils import available_memory, clear_memory, print_memory +from miles.utils.memory_utils import print_memory logger = logging.getLogger(__name__) @@ -275,9 +275,6 @@ def reload_process_groups(): @contextmanager def _wrap_low_level_call(): try: - mem_info = available_memory() - if mem_info["free_GB"] < 3: - clear_memory() yield except Exception as e: mem_info = print_memory("after torch distributed error") diff --git a/tools/convert_hf_to_torch_dist.py b/tools/convert_hf_to_torch_dist.py index 07475256954..354c216a651 100644 --- a/tools/convert_hf_to_torch_dist.py +++ b/tools/convert_hf_to_torch_dist.py @@ -1,5 +1,4 @@ import gc -import json import os import shutil from functools import wraps @@ -19,33 +18,7 @@ from miles.backends.megatron_utils.model_provider import get_model_provider_func from miles.utils.logging_utils import configure_logger from miles.utils.memory_utils import print_memory - - -def _load_hf_config_with_fallback(checkpoint_path: str): - """Load HF config with fallback for model types unknown to transformers.""" - try: - from transformers import AutoConfig - - return AutoConfig.from_pretrained(checkpoint_path, trust_remote_code=True) - except (ValueError, KeyError): - config_path = os.path.join(checkpoint_path, "config.json") - with open(config_path) as f: - config_dict = json.load(f) - - dtype_map = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32} - - def fix_dtype(d): - if "torch_dtype" in d: - d["torch_dtype"] = dtype_map.get(d["torch_dtype"], d["torch_dtype"]) - if "dtype" in d: - d["dtype"] = dtype_map.get(d["dtype"], d["dtype"]) - - fix_dtype(config_dict) - ns = type("HFConfig", (), config_dict)() - if "text_config" in config_dict: - fix_dtype(config_dict["text_config"]) - ns.text_config = type("TextConfig", (), config_dict["text_config"])() - return ns +from miles_plugins.models.hf_attention import _load_hf_config def patch_weight_to_mcore_format_preserve_fp32(): @@ -128,6 +101,7 @@ def ceildiv(a, b): def main(): if torch.version.hip: import megatron.core.dist_checkpointing.strategies.filesystem_async as filesystem_async_module + from miles.utils.rocm_checkpoint_writer import ROCmFileSystemWriterAsync filesystem_async_module.FileSystemWriterAsync = ROCmFileSystemWriterAsync @@ -162,7 +136,7 @@ def main(): bridge = AutoBridge.from_pretrained(hf_model_path, trust_remote_code=True) except (ValueError, KeyError): # Fallback for configs with model_type unknown to installed transformers. - bridge = AutoBridge.from_config(_load_hf_config_with_fallback(hf_model_path)) + bridge = AutoBridge.from_config(_load_hf_config(hf_model_path)) # Patch to preserve FP32 precision for _keep_fp32 params patch_weight_to_mcore_format_preserve_fp32() From 181d4481fd2b2363c90058763e25bc1884193735 Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Wed, 1 Apr 2026 20:45:06 +0000 Subject: [PATCH 4/4] chore(hf_attention): add note on common hf config fallback path --- miles_plugins/models/hf_attention.py | 1 + 1 file changed, 1 insertion(+) diff --git a/miles_plugins/models/hf_attention.py b/miles_plugins/models/hf_attention.py index d3858456fbe..7abe09b0eed 100644 --- a/miles_plugins/models/hf_attention.py +++ b/miles_plugins/models/hf_attention.py @@ -10,6 +10,7 @@ from megatron.core.transformer.module import MegatronModule +# Common fallback path for HF config loading; may be migrated elsewhere later. def _load_hf_config(checkpoint_path): """Load HF config with fallback for unsupported model types.""" try: