Fix ckpt convert for qwen35 moe model (#1233)

fix ckpt convert for qwen35 moe model and support uneven pipeline
parallel

Author: wenqingqian

tools/checkpoint/qwen35/README.md

@@ -65,3 +65,19 @@ python convert_qwen35.py --direction meg2hf --meg-path ... --hf-path ... --yaml
 ```bash
 python compare_two_ckpts.py --ref /path/to/ref/release --gen /path/to/gen/release --tp 2
 ```
+
+## WARNING
+
+If torch version < 2.9, maybe need modify (tools/checkpoint/qwen35/qwen35/config.py)
+
+```python
+self.use_linear_proj = cfg.get("vision_patch_embed_linear", True)
+```
+
+to
+
+```python
+self.use_linear_proj = cfg.get("vision_patch_embed_linear", False)
+```
+
+because qwen_vl will use linear instead of conv3d when torch version >= 2.9.

tools/checkpoint/qwen35/convert_qwen35.py

@@ -21,7 +21,9 @@
 """
 
 import argparse
+import shutil
 import sys
+import tempfile
 
 from qwen35.config import Config, detect_model_type
 from qwen35.constants import LN_ADJUSTMENT
@@ -36,11 +38,16 @@ def parse_args():
         choices=["hf2meg", "meg2hf"],
         help="Conversion direction",
     )
-    p.add_argument("--hf-path", required=True, help="Path to HF checkpoint directory")
+    p.add_argument(
+        "--hf-path",
+        default=None,
+        help="For hf2meg: input HF checkpoint directory or ModelScope model ID. "
+        "For meg2hf: output HF checkpoint directory (optional if --ref-path is given).",
+    )
     p.add_argument(
         "--meg-path",
-        required=True,
-        help="For hf2meg: output Megatron checkpoint directory. "
+        default=None,
+        help="For hf2meg: output Megatron checkpoint directory (optional if --ref-path is given). "
         "For meg2hf: input Megatron checkpoint directory.",
     )
     p.add_argument(
@@ -52,12 +59,14 @@ def parse_args():
         "--ref-path",
         default=None,
         help="Reference checkpoint path for validation (Megatron ref for hf2meg, "
-        "HF ref for meg2hf).",
+        "HF ref for meg2hf). When provided, the corresponding output path may be omitted; "
+        "output will be written to a temporary directory, validated, and cleaned up.",
     )
     p.add_argument(
-        "--adjust-embedding",
+        "--ref-skip-value",
         action="store_true",
-        help="Adjust embedding vocab size to match reference checkpoint (hf2meg only)",
+        help="When validating against --ref-path, skip numerical value comparison "
+        "and only compare structure, keys, and shapes.",
     )
     p.add_argument(
         "--adjust-ln",
@@ -71,6 +80,11 @@ def parse_args():
         action="store_true",
         help=argparse.SUPPRESS,
     )
+    p.add_argument(
+        "--adjust-embedding",
+        action="store_true",
+        help="During hf2meg, adjust vocab size to match the reference checkpoint.",
+    )
     p.add_argument(
         "--tp",
         type=int,
@@ -89,11 +103,23 @@ def parse_args():
         default=None,
         help="Override expert model parallel size from YAML (MoE only)",
     )
-    return p.parse_args()
+    return p.parse_args(), p
 
 
 def main():
-    args = parse_args()
+    args, parser = parse_args()
+
+    # Validate required paths based on direction
+    if args.direction == "hf2meg":
+        if not args.hf_path:
+            parser.error("--hf-path is required for hf2meg")
+        if not args.meg_path and not args.ref_path:
+            parser.error("--meg-path is required for hf2meg when --ref-path is not provided")
+    else:
+        if not args.meg_path:
+            parser.error("--meg-path is required for meg2hf")
+        if not args.hf_path and not args.ref_path:
+            parser.error("--hf-path is required for meg2hf when --ref-path is not provided")
 
     # Apply CLI override for layer norm adjustment
     if args.adjust_ln:
@@ -129,11 +155,32 @@ def main():
     print(f"TP={cfg.tp}, PP={cfg.pp}, EP={cfg.ep}")
     print(f"Layers={cfg.num_layers}, hidden={cfg.hidden_size}")
     print(f"LN adjustment: {LN_ADJUSTMENT}")
-
-    if args.direction == "hf2meg":
-        success = converter.run_hf2meg(args.hf_path, args.meg_path, args.ref_path)
-    else:
-        success = converter.run_meg2hf(args.meg_path, args.hf_path, args.ref_path)
+    print(f"Ref skip value: {args.ref_skip_value}")
+
+    temp_dir = None
+    try:
+        if args.direction == "hf2meg":
+            out_dir = args.meg_path
+            if not out_dir:
+                temp_dir = tempfile.mkdtemp(prefix="qwen35_hf2meg_")
+                out_dir = temp_dir
+                print(f"No --meg-path provided; using temporary output: {out_dir}")
+            success = converter.run_hf2meg(
+                args.hf_path, out_dir, args.ref_path, skip_value=args.ref_skip_value
+            )
+        else:
+            out_dir = args.hf_path
+            if not out_dir:
+                temp_dir = tempfile.mkdtemp(prefix="qwen35_meg2hf_")
+                out_dir = temp_dir
+                print(f"No --hf-path provided; using temporary output: {out_dir}")
+            success = converter.run_meg2hf(
+                args.meg_path, out_dir, args.ref_path, skip_value=args.ref_skip_value
+            )
+    finally:
+        if temp_dir:
+            shutil.rmtree(temp_dir, ignore_errors=True)
+            print(f"Cleaned up temporary output directory: {temp_dir}")
 
     sys.exit(0 if success else 1)
 

tools/checkpoint/qwen35/qwen35/config.py

@@ -36,6 +36,10 @@ def __init__(self, yaml_path):
         self.pp = cfg.get("pipeline_model_parallel_size", 1)
         self.ep = cfg.get("expert_model_parallel_size", 1)
 
+        # Uneven PP: optional per-stage layer counts
+        self.decoder_first_pipeline_num_layers = cfg.get("decoder_first_pipeline_num_layers", None)
+        self.decoder_last_pipeline_num_layers = cfg.get("decoder_last_pipeline_num_layers", None)
+
         self.num_layers = _require(cfg, "num_layers")
         self.hidden_size = _require(cfg, "hidden_size")
         self.num_attention_heads = _require(cfg, "num_attention_heads")
@@ -78,6 +82,71 @@ def __init__(self, yaml_path):
     def is_moe(self):
         return self.num_experts is not None and self.num_experts > 0
 
+    @property
+    def pp_layer_counts(self):
+        """Return a list of layer counts per PP rank, supporting uneven splits.
+
+        Uses decoder_first/last_pipeline_num_layers if set; otherwise divides
+        evenly.  For PP=1, returns [num_layers].
+        """
+        if self.pp == 1:
+            return [self.num_layers]
+
+        first = self.decoder_first_pipeline_num_layers
+        last = self.decoder_last_pipeline_num_layers
+
+        if first is None and last is None:
+            # Even split
+            base = self.num_layers // self.pp
+            counts = [base] * self.pp
+            # Distribute remainder to last ranks
+            remainder = self.num_layers - base * self.pp
+            for i in range(remainder):
+                counts[self.pp - 1 - i] += 1
+            return counts
+
+        # Uneven split: derive from first/last constraints
+        if first is not None and last is not None:
+            if self.pp == 2:
+                assert first + last == self.num_layers, (
+                    f"first({first}) + last({last}) != num_layers({self.num_layers})"
+                )
+                return [first, last]
+            # PP > 2: middle ranks share the remainder evenly
+            middle_total = self.num_layers - first - last
+            middle_ranks = self.pp - 2
+            base = middle_total // middle_ranks
+            counts = [first] + [base] * middle_ranks + [last]
+            remainder = middle_total - base * middle_ranks
+            for i in range(remainder):
+                counts[middle_ranks - i] += 1
+            return counts
+
+        if first is not None:
+            # Only first is specified
+            if self.pp == 2:
+                return [first, self.num_layers - first]
+            remaining = self.num_layers - first
+            middle_and_last = self.pp - 1
+            base = remaining // middle_and_last
+            counts = [first] + [base] * middle_and_last
+            remainder = remaining - base * middle_and_last
+            for i in range(remainder):
+                counts[self.pp - 1 - i] += 1
+            return counts
+
+        # Only last is specified
+        if self.pp == 2:
+            return [self.num_layers - last, last]
+        remaining = self.num_layers - last
+        first_and_middle = self.pp - 1
+        base = remaining // first_and_middle
+        counts = [base] * first_and_middle + [last]
+        remainder = remaining - base * first_and_middle
+        for i in range(remainder):
+            counts[first_and_middle - 1 - i] += 1
+        return counts
+
 
 def _require(d, key):
     """Return d[key], raising a clear error if the key is missing."""

tools/checkpoint/qwen35/qwen35/converter.py

@@ -259,14 +259,6 @@ def _split_tp(self, meg_sd):
         vis_head_dim = vis_h // vis_heads
         vis_qg = vis_heads
 
-        num_qg = cfg.num_query_groups
-        kv_ch = cfg.kv_channels
-        heads_per_group = cfg.num_attention_heads // num_qg
-        if cfg.attention_output_gate:
-            total_hpg = 2 * heads_per_group + 2
-        else:
-            total_hpg = heads_per_group + 2
-
         for k, v in meg_sd.items():
             if not isinstance(v, torch.Tensor):
                 for r in range(tp):
@@ -367,10 +359,11 @@ def _split_tp(self, meg_sd):
                 for r in range(tp):
                     shards[r][k] = chunks[r]
             elif "linear_qkv.weight" in k:
-                viewed = v.view(num_qg, total_hpg, kv_ch, cfg.hidden_size)
-                chunks = viewed.chunk(tp, dim=0)
+                # Simple dim-0 chunk / cat for TP sharding of the fused QKV weight.
+                # This is valid for any TP size, including cases where num_query_groups < tp.
+                chunks = v.chunk(tp, dim=0)
                 for r in range(tp):
-                    shards[r][k] = chunks[r].reshape(-1, cfg.hidden_size)
+                    shards[r][k] = chunks[r]
             elif "linear_proj.weight" in k:
                 chunks = v.chunk(tp, dim=1)
                 for r in range(tp):
@@ -483,13 +476,9 @@ def _merge_tp(self, shards):
             elif "out_proj.weight" in k:
                 merged[k] = torch.cat(vals, dim=1)
             elif "linear_qkv.weight" in k:
-                heads_per_group = cfg.num_attention_heads // cfg.num_query_groups
-                if cfg.attention_output_gate:
-                    total_hpg = 2 * heads_per_group + 2
-                else:
-                    total_hpg = heads_per_group + 2
-                viewed = [x.view(-1, total_hpg, cfg.kv_channels, cfg.hidden_size) for x in vals]
-                merged[k] = torch.cat(viewed, dim=0).view(-1, cfg.hidden_size)
+                # TP split is a simple contiguous dim-0 chunk, so merge is
+                # just concatenation along dim 0.
+                merged[k] = torch.cat(vals, dim=0)
             elif "linear_proj.weight" in k:
                 merged[k] = torch.cat(vals, dim=1)
             elif "q_layernorm" in k or "k_layernorm" in k:
@@ -566,7 +555,7 @@ def _adjust_embedding(self, meg_sd, ref_dir):
     # -------------------------------------------------------------------------
     # Public run methods
     # -------------------------------------------------------------------------
-    def run_hf2meg(self, hf_dir, save_dir, ref_dir=None):
+    def run_hf2meg(self, hf_dir, save_dir, ref_dir=None, skip_value=False):
         """Run HF -> Megatron conversion and save release checkpoint."""
         cfg = self.cfg
         hf_dir = ensure_hf_path(hf_dir)
@@ -596,11 +585,11 @@ def run_hf2meg(self, hf_dir, save_dir, ref_dir=None):
         print(f"Saved release checkpoint to {os.path.join(save_dir, 'release')}")
 
         if ref_dir:
-            success = validate_hf2meg_against_ref(shards_dict, cfg, ref_dir)
+            success = validate_hf2meg_against_ref(shards_dict, cfg, ref_dir, skip_value=skip_value)
             return success
         return True
 
-    def run_meg2hf(self, meg_dir, save_dir, ref_dir=None):
+    def run_meg2hf(self, meg_dir, save_dir, ref_dir=None, skip_value=False):
         """Run Megatron -> HF conversion and save safetensors."""
         cfg = self.cfg
         pp_merged = {}
@@ -631,7 +620,7 @@ def run_meg2hf(self, meg_dir, save_dir, ref_dir=None):
         print(f"Saved HF checkpoint to {save_dir}")
 
         if ref_dir:
-            success = validate_meg2hf_against_ref(hf_sd, cfg, ref_dir)
+            success = validate_meg2hf_against_ref(hf_sd, cfg, ref_dir, skip_value=skip_value)
             return success
         return True
 
@@ -655,7 +644,7 @@ class MoEConverter(BaseConverter):
     def __init__(self, cfg: Config, adjust_embedding=False):
         super().__init__(cfg, adjust_embedding=adjust_embedding)
 
-    def run_hf2meg(self, hf_dir, save_dir, ref_dir=None):
+    def run_hf2meg(self, hf_dir, save_dir, ref_dir=None, skip_value=False):
         cfg = self.cfg
         print(f"Loading HF weights from {hf_dir}...")
         hf_sd = load_hf_weights(hf_dir)
@@ -685,14 +674,14 @@ def run_hf2meg(self, hf_dir, save_dir, ref_dir=None):
         save_megatron_release_checkpoint(shards_dict, save_dir, cfg)
         print(f"Saved release checkpoint to {os.path.join(save_dir, 'release')}")
 
-        if ref_dir and cfg.ep == 1:
-            success = validate_hf2meg_against_ref(shards_dict, cfg, ref_dir, use_ep=True)
+        if ref_dir:
+            success = validate_hf2meg_against_ref(
+                shards_dict, cfg, ref_dir, use_ep=True, skip_value=skip_value
+            )
             return success
-        elif ref_dir:
-            print("Validation against reference checkpoint is only supported for EP=1; skipping.")
         return True
 
-    def run_meg2hf(self, meg_dir, save_dir, ref_dir=None):
+    def run_meg2hf(self, meg_dir, save_dir, ref_dir=None, skip_value=False):
         cfg = self.cfg
         pp_merged = {}
         for pp_rank in range(cfg.pp):
@@ -726,6 +715,6 @@ def run_meg2hf(self, meg_dir, save_dir, ref_dir=None):
         print(f"Saved HF checkpoint to {save_dir}")
 
         if ref_dir:
-            success = validate_meg2hf_against_ref(hf_sd, cfg, ref_dir)
+            success = validate_meg2hf_against_ref(hf_sd, cfg, ref_dir, skip_value=skip_value)
             return success
         return True

tools/checkpoint/qwen35/qwen35/io.py

@@ -52,20 +52,18 @@ def find_megatron_shard(meg_dir, tp_rank=0, pp_rank=0, ep_rank=None):
         if os.path.isdir(iter_release):
             candidates.append(iter_release)
 
-    suffixes = []
-    full = f"{tp_rank:02d}"
-    if pp_rank > 0:
-        full = f"{full}_{pp_rank:03d}"
-    if ep_rank is not None and ep_rank > 0:
-        full = f"{full}_{ep_rank:03d}"
-    suffixes.append(full)
-
-    # Fallbacks for older naming conventions
-    if ep_rank is not None and ep_rank > 0 and pp_rank == 0:
-        suffixes.append(f"{tp_rank:02d}_{ep_rank:03d}")
-    if pp_rank > 0 and (ep_rank is None or ep_rank == 0):
+    ep_rank = ep_rank if ep_rank is not None else 0
+    suffixes = [
+        # Standard PP>1 and EP>1 naming: tp_pp_ep
+        f"{tp_rank:02d}_{pp_rank:03d}_{ep_rank:03d}",
+    ]
+    # PP>1 EP=1 naming (only valid when ep_rank==0, otherwise ambiguous)
+    if ep_rank == 0:
         suffixes.append(f"{tp_rank:02d}_{pp_rank:03d}")
-    if pp_rank == 0 and (ep_rank is None or ep_rank == 0):
+    # PP=1 EP>1 naming: tp_ep
+    suffixes.append(f"{tp_rank:02d}_{ep_rank:03d}")
+    # PP=1 EP=1 naming (only valid when both are 0)
+    if pp_rank == 0 and ep_rank == 0:
         suffixes.append(f"{tp_rank:02d}")
 
     for base in candidates:
@@ -81,13 +79,14 @@ def save_megatron_release_checkpoint(shards_dict, save_dir, cfg):
     release_dir = os.path.join(save_dir, "release")
     os.makedirs(release_dir, exist_ok=True)
 
-    use_ep = cfg.ep > 1
     for rank_tuple, shard in shards_dict.items():
-        if use_ep:
+        # Support both (pp, tp) and (pp, tp, ep) tuple layouts regardless of cfg.ep.
+        if len(rank_tuple) == 3:
             pp_rank, tp_rank, ep_rank = rank_tuple
         else:
             pp_rank, tp_rank = rank_tuple
             ep_rank = None
+        use_ep = ep_rank is not None and ep_rank >= 0 and cfg.ep > 1
 
         name = f"mp_rank_{tp_rank:02d}"
         if cfg.pp > 1: