[fix] resolve SP/CP gradient issue in fla and qwen35 ckpt conversion

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)

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

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])

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":

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()