fix qwen35 ckpt convert (#1231)

fix qwen35 ckpt convert and add loss_mask

Author: wenqingqian

flagscale/models/megatron/qwen35/qwen35_model.py

@@ -191,6 +191,7 @@ def forward(
         video_input_mask: torch.Tensor = None,
         attention_mask: torch.Tensor = None,
         labels: torch.Tensor = None,
+        loss_mask: torch.Tensor | None = None,
         inference_params: InferenceParams = None,
         packed_seq_params: PackedSeqParams = None,
         extra_block_kwargs: dict = None,
@@ -274,6 +275,7 @@ def forward(
             attention_mask=attention_mask,
             decoder_input=combined_embeddings,
             labels=labels,
+            loss_mask=loss_mask,
             inference_params=inference_params,
             packed_seq_params=packed_seq_params,
             visual_pos_masks=visual_pos_masks,

flagscale/train/megatron/train_qwen35.py

@@ -372,6 +372,7 @@ def forward_step(data_iterator, model: Qwen35Model):
             video_input_mask=video_input_mask,
             attention_mask=attention_mask,
             labels=labels,
+            loss_mask=loss_mask,
         )
 
     return output_tensor, partial(loss_func, loss_mask, model=model)

tools/checkpoint/qwen35/README.md

@@ -0,0 +1,67 @@
+# Qwen3.5 Checkpoint Converter
+
+This directory provides a unified entry point for bidirectional Qwen3.5 checkpoint conversion between HuggingFace and Megatron.
+
+## Main Files
+
+| File/Directory | Description |
+| --- | --- |
+| [convert_qwen35.py](convert_qwen35.py) | Unified conversion entry point, supports `hf2meg` and `meg2hf` |
+| [run.sh](run.sh) | Shell wrapper for `convert_qwen35.py` |
+| [qwen35/](qwen35/) | Shared conversion library (Config, GDN/Attention/Vision/MTP/MLP, TP/PP/EP sharding, IO, validation) |
+
+## Usage
+
+### HF → Megatron
+
+```bash
+./run.sh hf2meg \
+    --yaml /workspace/FlagScale/examples/qwen35/conf/train/4b_nv_baseline.yaml \
+    --hf-path Qwen/Qwen3.5-4B \
+    --meg-path /path/to/output \
+    [--ref-path /path/to/ref/megatron]
+```
+
+For `hf2meg`, `--hf-path` can be either a local HF checkpoint directory or a ModelScope model ID (e.g. `Qwen/Qwen3.5-4B`). When the path is not found locally, the converter automatically downloads it from ModelScope.
+
+### Megatron → HF
+
+```bash
+./run.sh meg2hf \
+    --yaml /workspace/FlagScale/examples/qwen35/conf/train/4b_nv_baseline.yaml \
+    --meg-path /path/to/megatron/checkpoint \
+    --hf-path /path/to/output \
+    [--ref-path /path/to/ref/hf]
+```
+
+### Direct Python Invocation
+
+```bash
+python convert_qwen35.py --direction hf2meg --hf-path ... --meg-path ... --yaml ...
+python convert_qwen35.py --direction meg2hf --meg-path ... --hf-path ... --yaml ...
+```
+
+## Parameters
+
+- `--direction {hf2meg,meg2hf}`: **required**, conversion direction
+- `--hf-path PATH|ID`: **required**, HF checkpoint directory (input for hf2meg, output for meg2hf). For `hf2meg`, a ModelScope model ID such as `Qwen/Qwen3.5-4B` is also accepted; the checkpoint will be downloaded automatically if it is not found locally.
+- `--meg-path PATH`: **required**, Megatron checkpoint directory (output for hf2meg, input for meg2hf)
+- `--yaml PATH`: **required**, training config yaml (provides TP/PP/EP and model structure)
+- `--ref-path PATH`: optional, reference checkpoint path for post-conversion comparison
+- `--tp N` / `--pp N` / `--ep N`: optional, override parallel sizes from yaml
+- `--adjust-embedding`: during hf2meg, adjust vocab size to match the reference checkpoint
+- `--adjust-ln`: enable legacy layer norm adjustment (off by default; usually not needed for Qwen3.5)
+
+## Output Format
+
+- **Megatron**: `{save_dir}/release/mp_rank_*/model_optim_rng.pt`, plus `latest_checkpointed_iteration.txt`
+- **HF**: `{save_dir}/model.safetensors`
+
+## Validation
+
+- If `--ref-path` is provided, the converter automatically compares keys and shapes with the reference checkpoint
+- Use `compare_two_ckpts.py` for stricter per-tensor value comparison:
+
+```bash
+python compare_two_ckpts.py --ref /path/to/ref/release --gen /path/to/gen/release --tp 2
+```

tools/checkpoint/qwen35/convert_qwen35.py

@@ -0,0 +1,142 @@
+#!/usr/bin/env python3
+"""Unified Qwen3.5 HF <-> Megatron checkpoint converter.
+
+Supports both dense and MoE models. Model type is auto-detected from the
+input checkpoint / YAML; users do not need to specify it.
+
+Examples:
+    # HF -> Megatron
+    python convert_qwen35.py --direction hf2meg \
+        --hf-path /path/to/hf \
+        --meg-path /path/to/output \
+        --yaml /path/to/train.yaml \
+        [--ref-path /path/to/ref]
+
+    # Megatron -> HF
+    python convert_qwen35.py --direction meg2hf \
+        --meg-path /path/to/meg/checkpoint \
+        --hf-path /path/to/output \
+        --yaml /path/to/train.yaml \
+        [--ref-path /path/to/hf/ref]
+"""
+
+import argparse
+import sys
+
+from qwen35.config import Config, detect_model_type
+from qwen35.constants import LN_ADJUSTMENT
+from qwen35.converter import DenseConverter, MoEConverter
+
+
+def parse_args():
+    p = argparse.ArgumentParser(description="Unified Qwen3.5 checkpoint converter")
+    p.add_argument(
+        "--direction",
+        required=True,
+        choices=["hf2meg", "meg2hf"],
+        help="Conversion direction",
+    )
+    p.add_argument("--hf-path", required=True, help="Path to HF checkpoint directory")
+    p.add_argument(
+        "--meg-path",
+        required=True,
+        help="For hf2meg: output Megatron checkpoint directory. "
+        "For meg2hf: input Megatron checkpoint directory.",
+    )
+    p.add_argument(
+        "--yaml",
+        required=True,
+        help="Path to training YAML config (provides TP/PP/EP and model shapes)",
+    )
+    p.add_argument(
+        "--ref-path",
+        default=None,
+        help="Reference checkpoint path for validation (Megatron ref for hf2meg, "
+        "HF ref for meg2hf).",
+    )
+    p.add_argument(
+        "--adjust-embedding",
+        action="store_true",
+        help="Adjust embedding vocab size to match reference checkpoint (hf2meg only)",
+    )
+    p.add_argument(
+        "--adjust-ln",
+        action="store_true",
+        help="Enable legacy layer norm adjustment (add/subtract 1.0). "
+        "Only use this if the model stores raw gamma values instead of "
+        "zero-centered weights (default: disabled for Qwen3.5).",
+    )
+    p.add_argument(
+        "--no-adjust-ln",
+        action="store_true",
+        help=argparse.SUPPRESS,
+    )
+    p.add_argument(
+        "--tp",
+        type=int,
+        default=None,
+        help="Override tensor model parallel size from YAML",
+    )
+    p.add_argument(
+        "--pp",
+        type=int,
+        default=None,
+        help="Override pipeline model parallel size from YAML",
+    )
+    p.add_argument(
+        "--ep",
+        type=int,
+        default=None,
+        help="Override expert model parallel size from YAML (MoE only)",
+    )
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    # Apply CLI override for layer norm adjustment
+    if args.adjust_ln:
+        import qwen35.constants
+
+        qwen35.constants.LN_ADJUSTMENT = True
+    if args.no_adjust_ln:
+        import qwen35.constants
+
+        qwen35.constants.LN_ADJUSTMENT = False
+
+    cfg = Config(args.yaml)
+    if args.tp is not None:
+        cfg.tp = args.tp
+    if args.pp is not None:
+        cfg.pp = args.pp
+    if args.ep is not None:
+        cfg.ep = args.ep
+
+    # Auto-detect model type from whichever input is available
+    hf_input = args.hf_path if args.direction == "hf2meg" else None
+    meg_input = args.meg_path if args.direction == "meg2hf" else None
+    model_type = detect_model_type(
+        hf_dir=hf_input,
+        meg_dir=meg_input,
+        yaml_path=args.yaml,
+    )
+    converter_cls = MoEConverter if model_type == "moe" else DenseConverter
+    converter = converter_cls(cfg, adjust_embedding=args.adjust_embedding)
+
+    print(f"Direction: {args.direction}")
+    print(f"Model type: {model_type}")
+    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)
+
+    sys.exit(0 if success else 1)
+
+
+if __name__ == "__main__":
+    main()