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

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

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

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

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