[tab] slurm/dev/mihir/berlin/crowd_pilot/nemo/fine_tune.py (python)
from dataclasses import dataclass
import tyro
from nemo.collections import llm
import nemo_run as run
import os

os.environ["NEMORUN_HOME"] = "/fast/project/HFMI_SynergyUnit/tab_model/franz/nemorun_home"
os.environ["NEMO_CACHE_DIR"] = "/fast/project/HFMI_SynergyUnit/tab_model/franz/nemo_cache"

@dataclass
class Args:
    # training parameters
    name: str = "train_crowd_code_dataset"
    model_name: str = "qwen3_600m"
    seq_length: int = 4096
    global_batch_size: int = 128
    micro_batch_size: int = 1
    num_nodes: int = 1
    num_gpus_per_node: int = 1
    peft_scheme: str = "none" # "lora", "none"
    packed_sequence: bool = False
    # directories 
    ckpt_dir: str = ""
    data_dir: str = ""
    restore_ckpt_dir: str = ""
    # distributed training hyperparameters
    tensor_model_parallel_size: int = 1
    expert_model_parallel_size: int = 1
    recompute_granularity: str = "full"
    recompute_method: str = "uniform"
    recompute_num_layers: int = 1
    ckpt_async_save: bool = False



get_model = {
    "qwen3_600m": llm.qwen3_600m,
    "qwen3_30b_a3b": llm.qwen3_30b_a3b,
}

def get_recipe(args: Args):
    model = get_model[args.model_name]
    recipe = model.finetune_recipe(
            name=args.name,
            dir=args.ckpt_dir,
            num_nodes=args.num_nodes,
            num_gpus_per_node=args.num_gpus_per_node,
            peft_scheme=args.peft_scheme,
            packed_sequence=args.packed_sequence)
    return recipe

def local_executor_torchrun(nodes: int, devices: int) -> run.LocalExecutor:
    # Env vars for jobs are configured here
    env_vars = {
        "TORCH_NCCL_AVOID_RECORD_STREAMS": "1",
        "NCCL_NVLS_ENABLE": "0",
        "NVTE_DP_AMAX_REDUCE_INTERVAL": "0",
        "NVTE_ASYNC_AMAX_REDUCTION": "1",
    }

    executor = run.LocalExecutor(ntasks_per_node=devices, launcher="torchrun", env_vars=env_vars)

    return executor

def run_finetuning(recipe, args: Args):
    recipe.resume.restore_config.path = "/fast/project/HFMI_SynergyUnit/tab_model/data/checkpoints/nemo_converted_weights_qwen3-coder-30b-a3b-instruct/"
    recipe.data = run.Config(
        llm.ChatDataModule,
        dataset_root=args.data_dir,
        seq_length=args.seq_length,
        global_batch_size=args.global_batch_size,
        micro_batch_size=args.micro_batch_size,
        num_workers=0,
    )

    recipe.trainer.strategy.tensor_model_parallel_size = args.tensor_model_parallel_size
    recipe.trainer.strategy.expert_model_parallel_size = args.expert_model_parallel_size
    recipe.model.config.recompute_granularity = args.recompute_granularity
    recipe.model.config.recompute_method = args.recompute_method
    recipe.model.config.recompute_num_layers = args.recompute_num_layers
    recipe.trainer.strategy.ckpt_async_save = args.ckpt_async_save

    executor = local_executor_torchrun(recipe.trainer.num_nodes, recipe.trainer.devices)
    run.run(recipe, executor=executor, name=args.name)

def main():
    args = tyro.cli(Args)
    print("Parsed hyperparameters:")
    for k, v in vars(args).items():
        print(f"{k}: {v}")

    num_nodes = os.environ["SLURM_NNODES"]
    num_gpus_per_node = os.environ["SLURM_GPUS_ON_NODE"]

    assert args.num_nodes == num_nodes, f"num_nodes mismatch: {args.num_nodes} != {num_nodes}"
    assert args.num_gpus_per_node == num_gpus_per_node, f"num_gpus_per_node mismatch: {args.num_gpus_per_node} != {num_gpus_per_node}"

    recipe = get_recipe(args)
    print("Recipe:")
    for k, v in vars(recipe).items():
        print(f"{k}: {v}")
        
    print("Starting finetuning...")
    run_finetuning(recipe, args)
    print("Finetuning completed successfully.")
    
if __name__ == "__main__":
    main()