tutorials/llm/automodel/multinode.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a45b25c3-08b2-4a7e-b0cd-67293f15c307",
   "metadata": {},
   "source": [
    "# Optimizing Hugging Face Models with Multiple Nodes of GPU with NeMo-Run\n",
    "\n",
    "NeMo 2.0 enables users to perform Supervised Fine-Tuning (SFT) and Parameter-Efficient Fine-Tuning (PEFT) using Hugging Face (HF) Large Language Models (LLMs). It utilizes HF's auto classes to download and load transformer models, and wraps these models as Lightning modules to execute tasks like SFT and PEFT. The goal of this feature is to provide day-0 support for the models available in HF.\n",
    "\n",
    "[AutoModel](https://huggingface.co/docs/transformers/en/model_doc/auto) is the generic model class that is instantiated as one of the model classes from the library when created with the from_pretrained() class method. There are many AutoModel classes in HF, each covering a specific group of transformer model architectures. The AutoModel class primarily loads the base transformer model that converts embeddings to hidden states. For example, a specific AutoModel class like AutoModelForCausalLM includes a causal language modeling head on top of the base model.\n",
    "\n",
    "[NeMo-Run](https://github.com/NVIDIA/NeMo-Run) is a powerful tool designed to streamline the configuration, execution, and management of machine learning experiments across various computing environments.\n",
    "\n",
    "NeMo 2.0 includes wrapper classes for these HF AutoModel classes, making them runnable in NeMo pretraining, SFT, and PEFT workflows by converting them into Lightning modules. Due to the large number of AutoModel classes, NeMo 2.0 currently includes only the widely used auto classes.\n",
    "\n",
    "In this notebook, we will demonstrate a multi-node SFT training example on how to perform SFT with Hugging Face LLMs to make the models more performant on a specific task. We will focus on the models that can be loaded using the HF's `AutoModelForCausalLM` class."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63a50bad-f356-4076-8c5c-66b4481029dc",
   "metadata": {},
   "source": [
    "## Step 0: Install NeMo-Run"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ffd2159",
   "metadata": {},
   "source": [
    "We will be using [NeMo-Run](https://github.com/NVIDIA/NeMo-Run) in this tutorial to launch multi-node job on a SLURM cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28e16913-6a08-4ad8-835e-311fbb5af01d",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install git+https://github.com/NVIDIA/NeMo-Run.git"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67910a2b",
   "metadata": {},
   "source": [
    "### (Optional) Run this tutorial with NeMo container\n",
    "\n",
    "Alternatively, you can use NeMo container\n",
    "\n",
    "#### Launch the NeMo Framework container as follows:\n",
    "\n",
    "Depending on the number of gpus, `--gpus` might need to adjust accordingly:\n",
    "```bash\n",
    "docker run -it -p 8080:8080 -p 8088:8088 --rm --gpus '\\\"device=0,1\\\"' --ipc=host --network host -v $(pwd):/workspace nvcr.io/nvidia/nemo:25.02\\n\",\n",
    "```\n",
    "\n",
    "#### Launch Jupyter Notebook as follows:\n",
    "```bash\n",
    "jupyter notebook --allow-root --ip 0.0.0.0 --port 8088 --no-browser --NotebookApp.token=''\\n\",\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6966670b-2097-47c0-95f2-edaafab0e33f",
   "metadata": {},
   "source": [
    "## Step 1: Setup Huggingface access token\n",
    "\n",
    "Some models have gated access. If you are using one of those models, you will need to obtain access first. Then, set your HF Token by running the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "439a3c6a-8718-4b49-acdb-e7f59db38f59",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# TODO: Replace with your own Hugging Face token\n",
    "HF_TOKEN = '<hf_your_token_here>'\n",
    "os.environ[\"HF_TOKEN\"] = HF_TOKEN"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d652fa70",
   "metadata": {},
   "source": [
    "## Step 2: Import Modules and Prepare the Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "056be19b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "import fiddle as fdl\n",
    "import lightning.pytorch as pl\n",
    "from lightning.pytorch.loggers import WandbLogger\n",
    "from torch.utils.data import DataLoader\n",
    "from typing import Optional\n",
    "\n",
    "import nemo_run as run\n",
    "\n",
    "from nemo import lightning as nl\n",
    "from nemo.collections import llm\n",
    "from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer\n",
    "from nemo.collections.llm.gpt.data.hf_dataset import SquadHFDataModule\n",
    "from nemo.lightning.pytorch.callbacks import JitConfig, JitTransform"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cfe3c7d-9d36-47d2-9107-361025d175a0",
   "metadata": {},
   "source": [
    "We will use the [SQuAD](https://rajpurkar.github.io/SQuAD-explorer/) dataset provided by NeMo 2.0, which is a reading comprehension dataset consisting of questions and answers pairs. The SquadHFDataModule extends the `pl.LightningDataModule`, giving it access to existing data-handling logic. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67e6e4d4-8e0c-4507-b386-22c3d63097c1",
   "metadata": {},
   "source": [
    "## Step 3: Define your SLURM executor with NeMo-Run\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b3bb9ec",
   "metadata": {},
   "source": [
    "It is necessary to setup your SLURM executor function to use NeMo-Run to launch your multinode job. In order to launch a SLURM job, it is necessary to setup your user information with your wanted SLURM cluster. Please change the following part to setup your user information and basic cluster information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "272f6d7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: Change of this part is necessary to continue this tutorial\n",
    "USER = \"\" # Your cluster username\n",
    "HOST = \"\" # Your cluster host address\n",
    "REMOTE_JOB_DIR = \"\" # The path to the directory where the job will be saved\n",
    "ACCOUNT = \"\" # Your cluster account name\n",
    "PARTITION = \"\" # Your cluster partition name\n",
    "TIME = \"04:00:00\" # The time limit for the job, default is \"04:00:00\"\n",
    "CUSTOM_MOUNTS = [] # List of custom mounts, default is None\n",
    "CUSTOM_ENV_VARS = {\"HF_TOKEN\": HF_TOKEN} # Dictionary of custom environment variables, default is None\n",
    "CONTAINER_IMAGE = \"nvcr.io/nvidia/nemo:25.02.rc5\" # The container image to use, default is \"nvcr.io/nvidia/nemo:25.02\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99e77fca",
   "metadata": {},
   "source": [
    "<font color='red'>NOTE:</font> Due to the nature of SLURM clusters, it might be necessary to consult your cluster documentation for SLURM GPU config settings. For example, `gres=\"gpu:8\"` and `gpus_per_node=devices` might be necessery in some cluster but unwanted in another cluster. Please comment out these unwanted GPU configs if necessary. This is an example of how to change the cluster related config settings.\n",
    "\n",
    "```python\n",
    "executor = run.SlurmExecutor(\n",
    "    account=account,\n",
    "    partition=partition,\n",
    "    tunnel=run.SSHTunnel(\n",
    "        user=user,\n",
    "        host=host,\n",
    "        job_dir=remote_job_dir,\n",
    "    ),\n",
    "    nodes=nodes,\n",
    "    ntasks_per_node=devices,\n",
    "    # gpus_per_node=devices,\n",
    "    mem=\"0\",\n",
    "    exclusive=True,\n",
    "    # gres=\"gpu:8\",\n",
    "    packager=run.GitArchivePackager(),\n",
    ")\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a118868-8c6e-44ad-9b2d-3be3994a093b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def slurm_executor(\n",
    "    user: str,\n",
    "    host: str,\n",
    "    remote_job_dir: str,\n",
    "    account: str,\n",
    "    partition: str,\n",
    "    nodes: int,\n",
    "    devices: int,\n",
    "    time: str = \"04:00:00\",\n",
    "    custom_mounts: Optional[list[str]] = None,\n",
    "    custom_env_vars: Optional[dict[str, str]] = None,\n",
    "    container_image: str = \"nvcr.io/nvidia/nemo:25.02\",\n",
    "    retries: int = 0,\n",
    ") -> run.SlurmExecutor:\n",
    "    if not (user and host and remote_job_dir and account and partition and nodes and devices):\n",
    "        raise RuntimeError(\n",
    "            \"Please set user, host, remote_job_dir, account, partition, nodes and devices args for using this \",\n",
    "            \"function.\",\n",
    "        )\n",
    "\n",
    "    mounts = []\n",
    "    if custom_mounts:\n",
    "        mounts.extend(custom_mounts)\n",
    "\n",
    "    env_vars = {\n",
    "        \"TRANSFORMERS_OFFLINE\": \"0\",\n",
    "        \"TORCH_NCCL_AVOID_RECORD_STREAMS\": \"1\",\n",
    "        \"NCCL_NVLS_ENABLE\": \"0\",\n",
    "        \"NVTE_DP_AMAX_REDUCE_INTERVAL\": \"0\",\n",
    "        \"NVTE_ASYNC_AMAX_REDUCTION\": \"1\",\n",
    "    }\n",
    "    if custom_env_vars:\n",
    "        env_vars |= custom_env_vars\n",
    "\n",
    "    # Note: The following part may need to be adjusted,\n",
    "    # Please consult the documentation of the cluster you are using.\n",
    "    # For example, 'gres=\"gpu:8\",' and 'gpus_per_node=devices,' may need to be changed\n",
    "    executor = run.SlurmExecutor(\n",
    "        account=account,\n",
    "        partition=partition,\n",
    "        tunnel=run.SSHTunnel(\n",
    "            user=user,\n",
    "            host=host,\n",
    "            job_dir=remote_job_dir,\n",
    "        ),\n",
    "        nodes=nodes,\n",
    "        ntasks_per_node=devices,\n",
    "        gpus_per_node=devices,\n",
    "        mem=\"0\",\n",
    "        exclusive=True,\n",
    "        gres=\"gpu:8\",\n",
    "        packager=run.GitArchivePackager(),\n",
    "    )\n",
    "\n",
    "    executor.container_image = container_image\n",
    "    executor.container_mounts = mounts\n",
    "    executor.env_vars = env_vars\n",
    "    executor.retries = retries\n",
    "    executor.time = time\n",
    "\n",
    "    return executor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a23943ee-ffa1-497d-a395-3e4767271341",
   "metadata": {},
   "source": [
    "## Step 4: Set Parameters and Start the SFT with a HF Model\n",
    "\n",
    "Now, we will set some of the important variables, including the HF model name, maximum steps, number of GPUs, etc. You can find the details of these parameters below.\n",
    "- `model_name`: Pre-trained HF model or path of a HF model.\n",
    "- `strategy`: Distributed training strategy such as DDP, FSDP, etc. \n",
    "- `devices`: Number of GPUs to be used in the training.\n",
    "- `max_steps`: Number of steps in the training.\n",
    "- `wandb_project`: wandb project.\n",
    "- `use_torch_jit`: Enable torch jit or not.\n",
    "- `ckpt_folder`: Path for the checkpoins.\n",
    "- `num_nodes`: Number of nodes.\n",
    "- `devices`: Number of GPUs per node.\n",
    "- `DATA_PATH`: Path to SQuad dataset in this example.\n",
    "\n",
    "All popular models, including Llama, GPT, Gemma, Mistral, Phi, and Qwen, are supported. After running this workflow, please select another HF model and rerun the notebook with that model. Ensure the chosen model fits within your GPU(s) memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3780a047-febb-4d97-a59a-99d8ee036332",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: In order to use the models like Llama, Gemma, you need to ask for permission on the HF model page and then pass the HF_TOKEN in the next cell.\n",
    "model_name = \"meta-llama/Llama-3.2-1B\"  # HF model name. This can be the path of the downloaded model as well.\n",
    "strategy = \"fsdp2\" # Distributed training strategy such as DDP, FSDP2, etc.\n",
    "max_steps = 100 # Number of steps in the training loop.\n",
    "accelerator = \"gpu\"\n",
    "wandb_project = None # Wandb project name. If None, it will not log to wandb.\n",
    "use_torch_jit = False # torch jit can be enabled.\n",
    "ckpt_folder=\"/opt/checkpoints/automodel_experiments/\" # Path for saving the checkpoint.\n",
    "\n",
    "num_nodes = 2 # Number of nodes to use for training. \n",
    "devices = 8 # Number of GPUs to use for training.\n",
    "\n",
    "DATA_PATH = '/opt/data/squad' # Path to the dataset. This should be a folder containing train and validation files."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cd65e5e-93fa-4ea0-b89d-2f48431b725c",
   "metadata": {},
   "source": [
    "After setting some parameters, we can start the SFT training workflow. Although the SFT workflow with HF models/checkpoints differs slightly from workflows with NeMo models/checkpoints, we still use the same NeMo 2.0 API. The main difference is the model we pass into the `fine-tune` API.\n",
    "\n",
    "<font color='red'>NOTE:</font> For PEFT training workflow, simply change the peft_scheme to 'lora' as shown below\n",
    "\n",
    "```python\n",
    "recipe = llm.hf_auto_model_for_causal_lm.finetune_recipe(\n",
    "    model_name=model_name,\n",
    "    name=exp_name,\n",
    "    num_nodes=num_nodes,\n",
    "    num_gpus_per_node=devices,\n",
    "    peft_scheme='lora', # change to PEFT with LoRA\n",
    "    max_steps=max_steps,\n",
    "    trust_remote_code=True,\n",
    "    attn_implementation='eager',\n",
    ")\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3578630-05b7-4a8c-8b5d-a7d9e847f17b",
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb = None\n",
    "if wandb_project is not None:\n",
    "    model = '_'.join(model_name.split('/')[-2:])\n",
    "    wandb = WandbLogger(\n",
    "        project=wandb_project,\n",
    "        name=f'{model}_dev{devices}_strat_{strategy}',\n",
    "    )\n",
    "\n",
    "callbacks = []\n",
    "if use_torch_jit:\n",
    "    jit_config = JitConfig(use_torch=True, torch_kwargs={'dynamic': False}, use_thunder=False)\n",
    "    callbacks = [JitTransform(jit_config)]\n",
    "\n",
    "callbacks.append(\n",
    "    nl.ModelCheckpoint(\n",
    "        every_n_train_steps=max_steps // 2,\n",
    "        dirpath=ckpt_folder,\n",
    "    )\n",
    ")\n",
    "\n",
    "exp_name = \"HFAutoModelForCausalLM\"\n",
    "\n",
    "# Uses configs from NeMo directly\n",
    "recipe = llm.hf_auto_model_for_causal_lm.finetune_recipe(\n",
    "    model_name=model_name,\n",
    "    name=exp_name,\n",
    "    num_nodes=num_nodes,\n",
    "    num_gpus_per_node=devices,\n",
    "    peft_scheme='none',\n",
    "    max_steps=max_steps,\n",
    "    trust_remote_code=True,\n",
    "    attn_implementation='eager',\n",
    ")\n",
    "\n",
    "recipe.trainer.val_check_interval = 50\n",
    "\n",
    "tokenizer = llm.HFAutoModelForCausalLM.configure_tokenizer(model_name)\n",
    "recipe.data = run.Config(\n",
    "    SquadHFDataModule,\n",
    "    path_or_dataset=DATA_PATH,\n",
    "    split=\"train[:100]\",\n",
    "    pad_token_id=tokenizer.tokenizer.eos_token_id,\n",
    "    tokenizer=run.Config(AutoTokenizer, pretrained_model_name=model_name),\n",
    ")\n",
    "\n",
    "recipe.trainer.strategy = run.Config(\n",
    "    nl.FSDP2Strategy, data_parallel_size=num_nodes * devices, tensor_parallel_size=1\n",
    ")\n",
    "recipe.trainer.plugins = None\n",
    "\n",
    "executor = slurm_executor(\n",
    "    user=USER,\n",
    "    host=HOST,\n",
    "    remote_job_dir=REMOTE_JOB_DIR,\n",
    "    account=ACCOUNT,\n",
    "    partition=PARTITION,\n",
    "    nodes=recipe.trainer.num_nodes,\n",
    "    devices=recipe.trainer.devices,\n",
    "    time=TIME,\n",
    "    custom_mounts=CUSTOM_MOUNTS,\n",
    "    custom_env_vars=CUSTOM_ENV_VARS,\n",
    "    container_image=CONTAINER_IMAGE,\n",
    ")\n",
    "\n",
    "with run.Experiment(f\"{exp_name}\") as exp:\n",
    "    for i in range(1):\n",
    "        exp.add(\n",
    "            recipe,\n",
    "            executor=executor,\n",
    "            name=exp_name,\n",
    "            tail_logs=True if isinstance(executor, run.LocalExecutor) else False,\n",
    "        )\n",
    "\n",
    "    exp.run(sequential=True, detach=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cab7aa3",
   "metadata": {},
   "source": [
    "After launching the multinode job with NeMo-Run, you will be able to check your launched job with `nemo experiment status HFAutoModelForCausalLM_<job_id>`. Also you can your log in both local directory `.nemo_run/experiments/HFAutoModelForCausalLM/HFAutoModelForCausalLM_<job_id>` and remote directory based on the remote directory of your input."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nemo3.10",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}