[tab] /fast/home/mihir.mahajan/Projects/crowd-pilot/slurm/dev/mihir/berlin/crowd_pilot/nemo/iterating_on_stuff/test_full_finetune_4k_requeue.py (python)
import os

import nemo_run as run
from nemo.collections import llm
from lightning.pytorch.loggers import WandbLogger
from nemo import lightning as nl
from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer


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"
slurm_job_id = os.environ.get("SLURM_JOB_ID", "local")

NAME = f"qwen3_600m_full_4k_tokens_{slurm_job_id}"
# NAME = "qwen3_600m_full_4k_tokens_34940"
DATASET_DIR = "/fast/project/HFMI_SynergyUnit/tab_model/data/nemo_hf_part_jsonl_4k_tokens/" 
LOG_DIR = "/fast/project/HFMI_SynergyUnit/mihir/logs"
BASE_CKPT_PATH = "/fast/project/HFMI_SynergyUnit/tab_model/data/checkpoints/nemo_converted_weights_qwen3-600m/"

MAX_STEPS = 2000
VAL_INTERVAL = 50
NUM_CHECKPOINTS = 10

def recipe_to_dict(recipe, depth=4):
    if depth == 0:
        return dict()
    cfg_dict = dict()
    for k, v in vars(recipe)['__arguments__'].items():
        if isinstance(v, run.Config) or isinstance(v, dict):
            cfg_dict[k] = recipe_to_dict(v, depth-1)
        elif isinstance(v, (int, float, str, bool)):
            cfg_dict[k] = v
        else:
            cfg_dict[k] = str(v)
    return cfg_dict

def get_ckpt_path():
    ckpt_path = f"{LOG_DIR}/{NAME}"
    if not os.path.exists(ckpt_path):
        return BASE_CKPT_PATH
    else:
        return os.path.join(ckpt_path, "checkpoints")

recipe = llm.qwen3_600m.finetune_recipe(
    name=NAME,
    dir=LOG_DIR,
    num_nodes=1,
    num_gpus_per_node=1,
    peft_scheme=None,
    # FIXME(f.srambical): check whether this is supported by ChatDataModule
    packed_sequence=False,
)

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.resume = run.Config(
        nl.AutoResume,
        restore_config=run.Config(nl.RestoreConfig, path=BASE_CKPT_PATH),
        resume_if_exists=True,
        resume_from_directory=os.path.join(LOG_DIR, NAME, "checkpoints"),
    )

    recipe.data = run.Config(
        llm.ChatDataModule,
        dataset_root=DATASET_DIR,
        seq_length=4096,
        global_batch_size=32,
        micro_batch_size=4,
        num_workers=0,
    )

    recipe.trainer.max_steps = MAX_STEPS
    recipe.log.ckpt.every_n_train_steps = MAX_STEPS // NUM_CHECKPOINTS
    recipe.log.ckpt.save_top_k = NUM_CHECKPOINTS
    recipe.trainer.val_check_interval = VAL_INTERVAL
    recipe.log.ckpt.save_weights_only = False
    recipe.log.ckpt.always_save_context = True
    recipe.log.ckpt.save_optim_on_train_end = True

    cfg_dict = recipe_to_dict(recipe)
    recipe.log.wandb = run.Config(
        WandbLogger,
        name=NAME,
        project="crowd-pilot-nemo",
        save_dir=f"{LOG_DIR}/wandb",
        entity="instant-uv",
        config=cfg_dict,
    )

    executor = local_executor_torchrun(nodes=recipe.trainer.num_nodes, devices=recipe.trainer.devices)

    run.run(recipe, executor=executor, name=NAME)

# This condition is necessary for the script to be compatible with Python's multiprocessing module.
if __name__ == "__main__":
    run_finetuning()