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/miles_plugins/models/hf_attention.py b/miles_plugins/models/hf_attention.py index c353ae7b29d..7abe09b0eed 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,57 @@ 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 + + +# 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: + 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 +82,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 +106,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..354c216a651 100644 --- a/tools/convert_hf_to_torch_dist.py +++ b/tools/convert_hf_to_torch_dist.py @@ -18,6 +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 +from miles_plugins.models.hf_attention import _load_hf_config def patch_weight_to_mcore_format_preserve_fp32(): @@ -100,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 @@ -130,7 +132,11 @@ 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(hf_model_path)) # Patch to preserve FP32 precision for _keep_fp32 params patch_weight_to_mcore_format_preserve_fp32()