Upload configs/finetune_teleyaml.py with huggingface_hub

configs/finetune_teleyaml.py

@@ -0,0 +1,119 @@
+#!/usr/bin/env python3
+"""
+Custom fine-tuning script for TeleYAML models.
+Wraps the standard Nemotron-3-Nano finetune with custom LoRA parameters.
+
+Usage:
+    torchrun --nproc-per-node=2 /scripts/nemo-configs/finetune_teleyaml.py \
+        --lora-dim 64 --lora-alpha 128 --lora-dropout 0.05 \
+        --config-file /scripts/nemo-configs/teleyaml-v3.yaml
+"""
+
+import argparse
+import logging
+import os
+import sys
+
+import torch
+from omegaconf import OmegaConf
+
+from megatron.bridge.peft.lora import LoRA
+from megatron.bridge.recipes.nemotronh.nemotron_3_nano import (
+    nemotron_3_nano_finetune_config as finetune_config,
+)
+from megatron.bridge.training.finetune import finetune
+from megatron.bridge.training.gpt_step import forward_step
+from megatron.bridge.training.utils.omegaconf_utils import (
+    apply_overrides,
+    create_omegaconf_dict_config,
+    parse_hydra_overrides,
+)
+
+# Import custom processor directly
+from megatron.bridge.data.hf_processors.teleyaml import process_teleyaml_example
+
+logger = logging.getLogger(__name__)
+
+# Target modules for Nemotron-3-Nano (Mamba + MLP layers)
+MAMBA_TARGET_MODULES = [
+    "linear_qkv",
+    "linear_proj", 
+    "linear_fc1",
+    "linear_fc2",
+    "in_proj",
+    "out_proj",
+]
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="TeleYAML Fine-tuning with Custom LoRA")
+    
+    # LoRA parameters
+    parser.add_argument("--lora-dim", type=int, default=32, help="LoRA rank dimension (default: 32)")
+    parser.add_argument("--lora-alpha", type=int, default=32, help="LoRA alpha scaling (default: 32)")
+    parser.add_argument("--lora-dropout", type=float, default=0.0, help="LoRA dropout rate (default: 0.0)")
+    
+    # Standard args from original script
+    parser.add_argument("--config-file", type=str, help="Path to YAML config file")
+    parser.add_argument("--packed-sequence", action="store_true", help="Use sequence packing")
+    parser.add_argument("--seq-length", type=int, default=2048, help="Sequence length")
+    
+    args, cli_overrides = parser.parse_known_args()
+    return args, cli_overrides
+
+
+def main():
+    args, cli_overrides = parse_args()
+    
+    # Build custom LoRA config with our parameters
+    lora_config = LoRA(
+        target_modules=MAMBA_TARGET_MODULES,
+        dim=args.lora_dim,
+        alpha=args.lora_alpha,
+        dropout=args.lora_dropout,
+    )
+    
+    print(f"LoRA Config: dim={args.lora_dim}, alpha={args.lora_alpha}, dropout={args.lora_dropout}")
+    
+    # Get base config, passing our custom LoRA object instead of "lora" string
+    cfg = finetune_config(
+        seq_length=args.seq_length,
+        peft=lora_config,  # Pass the LoRA object, not "lora" string
+        packed_sequence=args.packed_sequence,
+    )
+    cfg.model.seq_length = args.seq_length
+    
+    # Convert to OmegaConf for merging
+    merged_omega_conf, excluded_fields = create_omegaconf_dict_config(cfg)
+    
+    # Load YAML config if provided
+    if args.config_file:
+        if not os.path.exists(args.config_file):
+            print(f"ERROR: Config file not found: {args.config_file}")
+            sys.exit(1)
+        yaml_overrides = OmegaConf.load(args.config_file)
+        merged_omega_conf = OmegaConf.merge(merged_omega_conf, yaml_overrides)
+        print(f"Loaded config from: {args.config_file}")
+    
+    # Apply CLI overrides
+    if cli_overrides:
+        merged_omega_conf = parse_hydra_overrides(merged_omega_conf, cli_overrides)
+    
+    # Apply merged config back to ConfigContainer
+    final_overrides = OmegaConf.to_container(merged_omega_conf, resolve=True)
+    apply_overrides(cfg, final_overrides, excluded_fields)
+    
+    # CRITICAL: Set the processor function directly (bypasses Hydra _target_ issue)
+    cfg.dataset.process_example_fn = process_teleyaml_example
+    print(f"Using custom processor: {process_teleyaml_example.__name__}")
+    
+    # Start training
+    print("Starting fine-tuning...")
+    finetune(config=cfg, forward_step_func=forward_step)
+    
+    if torch.distributed.is_initialized():
+        torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()