feat: enhance convert_lora_to_hf script to support exporting LoRA adapters in HuggingFace PEFT format

Signed-off-by: ruit <ruit@nvidia.com>

Author: RayenTian

examples/converters/convert_lora_to_hf.py

@@ -1,10 +1,20 @@
-"""Merge a Megatron LoRA adapter checkpoint with its base model and export to HuggingFace format.
+"""Export a Megatron LoRA adapter checkpoint to HuggingFace format.
 
-This is helpful when one wants to train the model using Megatron with LoRA adapter and then convert it to HuggingFace format
-for inference and evaluation.
+This script supports two workflows:
+
+1. Merge the base model and LoRA adapter, then export a standard HuggingFace model.
+2. Export only the LoRA adapter to a HuggingFace PEFT-compatible directory without merging.
 
 Usage (requires mcore extra):
 
+    # Export adapter only (recommended when you want PEFT format)
+    uv run --extra mcore python examples/converters/convert_lora_to_hf.py \
+        --adapter-only \
+        --adapter-ckpt results/dpo_glm5/step_5/policy/weights/iter_0000000 \
+        --hf-model-name zai-org/GLM-5 \
+        --hf-ckpt-path ./hf_lora_adapter
+
+    # Merge base model + adapter and export a full HF checkpoint
     uv run --extra mcore python examples/converters/convert_lora_to_hf.py \
         --base-ckpt ~/.cache/huggingface/nemo_rl/zai-org/GLM-5/iter_0000000 \
         --adapter-ckpt results/dpo_glm5/step_5/policy/weights/iter_0000000 \
@@ -29,13 +39,13 @@
 
 def parse_args():
     parser = argparse.ArgumentParser(
-        description="Merge Megatron LoRA adapter with base model and export to HF"
+        description="Export Megatron LoRA checkpoint to HuggingFace format"
     )
     parser.add_argument(
         "--base-ckpt",
         type=str,
-        required=True,
-        help="Path to base model Megatron checkpoint (iter_XXXXXXX directory)",
+        default=None,
+        help="Path to base model Megatron checkpoint (iter_XXXXXXX directory). Required unless --adapter-only is set.",
     )
     parser.add_argument(
         "--adapter-ckpt",
@@ -53,9 +63,47 @@ def parse_args():
         "--hf-ckpt-path",
         type=str,
         required=True,
-        help="Output path for merged HF checkpoint",
+        help="Output path for the exported HF checkpoint or adapter directory",
     )
-    return parser.parse_args()
+    parser.add_argument(
+        "--adapter-only",
+        action="store_true",
+        help="Export only the LoRA adapter in HuggingFace PEFT format without merging into the base model.",
+    )
+    args = parser.parse_args()
+    if not args.adapter_only and not args.base_ckpt:
+        parser.error("--base-ckpt is required unless --adapter-only is set")
+    return args
+
+
+def export_lora_adapter_to_hf(
+    adapter_ckpt: str,
+    hf_model_name: str,
+    hf_ckpt_path: str,
+) -> str:
+    """Export a Megatron LoRA checkpoint to HuggingFace PEFT adapter format.
+
+    Args:
+        adapter_ckpt: Path to the LoRA adapter Megatron checkpoint (iter_XXXXXXX directory).
+        hf_model_name: HuggingFace model identifier for the base model.
+        hf_ckpt_path: Output directory for the HuggingFace PEFT adapter files.
+
+    Returns:
+        The *hf_ckpt_path* that was written to.
+
+    Raises:
+        FileExistsError: If *hf_ckpt_path* already exists.
+    """
+    if os.path.exists(hf_ckpt_path):
+        raise FileExistsError(f"Output path already exists: {hf_ckpt_path}")
+
+    from megatron.bridge import AutoBridge
+
+    bridge = AutoBridge.from_hf_pretrained(hf_model_name, trust_remote_code=True)
+    logger.info("Exporting LoRA adapter in HuggingFace PEFT format...")
+    bridge.export_adapter_ckpt(adapter_ckpt, hf_ckpt_path)
+    logger.info(f"Done! HF adapter saved to: {hf_ckpt_path}")
+    return hf_ckpt_path
 
 
 def merge_lora_to_hf(
@@ -86,13 +134,16 @@ def merge_lora_to_hf(
     from megatron.bridge import AutoBridge
     from megatron.bridge.peft.lora import LoRA
     from megatron.bridge.training.checkpointing import (
+        _generate_model_state_dict,
         _load_model_weights_from_checkpoint,
+        apply_peft_adapter_filter_to_state_dict,
     )
     from megatron.bridge.training.model_load_save import (
         load_model_config,
         megatron_cpu_init_context,
         temporary_distributed_context,
     )
+    from megatron.core import dist_checkpointing
 
     bridge = AutoBridge.from_hf_pretrained(hf_model_name, trust_remote_code=True)
 
@@ -140,9 +191,10 @@ def merge_lora_to_hf(
             lora_B_init_method=peft_section.get("lora_B_init_method", "zero"),
             a2a_experimental=peft_section.get("a2a_experimental", False),
         )
-        model_cfg.peft = peft
 
-        logger.info("Building model with LoRA wrappers on CPU...")
+        logger.info(
+            "Building base model on CPU (LoRA wrappers applied after base weights are loaded)..."
+        )
         if hasattr(model_cfg, "finalize"):
             model_cfg.finalize()
         with megatron_cpu_init_context(model_cfg):
@@ -159,8 +211,25 @@ def merge_lora_to_hf(
         _load_model_weights_from_checkpoint(base_ckpt, megatron_model, strict=False)
         gc.collect()
 
+        logger.info("Applying LoRA wrappers to model...")
+        megatron_model = peft(megatron_model, training=False)
+        gc.collect()
+
         logger.info(f"Loading LoRA adapter from {adapter_ckpt}...")
-        _load_model_weights_from_checkpoint(adapter_ckpt, megatron_model, strict=False)
+        adapter_sharded_state_dict = _generate_model_state_dict(megatron_model, {})
+        adapter_sharded_state_dict = apply_peft_adapter_filter_to_state_dict(
+            adapter_sharded_state_dict, peft
+        )
+        loaded_adapter_state_dict = dist_checkpointing.load(
+            adapter_sharded_state_dict, adapter_ckpt
+        )
+        model_key = (
+            "model"
+            if "model" in loaded_adapter_state_dict
+            else next(k for k in loaded_adapter_state_dict if k.startswith("model"))
+        )
+        for m in megatron_model:
+            m.load_state_dict(loaded_adapter_state_dict[model_key], strict=False)
         gc.collect()
 
         logger.info("Saving merged model in HuggingFace format...")
@@ -183,12 +252,19 @@ def merge_lora_to_hf(
 
 def main():
     args = parse_args()
-    merge_lora_to_hf(
-        base_ckpt=args.base_ckpt,
-        adapter_ckpt=args.adapter_ckpt,
-        hf_model_name=args.hf_model_name,
-        hf_ckpt_path=args.hf_ckpt_path,
-    )
+    if args.adapter_only:
+        export_lora_adapter_to_hf(
+            adapter_ckpt=args.adapter_ckpt,
+            hf_model_name=args.hf_model_name,
+            hf_ckpt_path=args.hf_ckpt_path,
+        )
+    else:
+        merge_lora_to_hf(
+            base_ckpt=args.base_ckpt,
+            adapter_ckpt=args.adapter_ckpt,
+            hf_model_name=args.hf_model_name,
+            hf_ckpt_path=args.hf_ckpt_path,
+        )
 
 
 if __name__ == "__main__":