Guide for fine-tuning Llama/CodeLlama models
from modal import Stub, Image, Volume, Secret

N_GPUS = 4
GPU_MEM = 80
BASE_MODELS = {
    "base7": "meta-llama/Llama-2-7b-hf",
    "chat7": "meta-llama/Llama-2-7b-chat-hf",
    "chat13": "meta-llama/Llama-2-13b-chat-hf",
    "code7": "codellama/CodeLlama-7b-hf",
    "code34": "codellama/CodeLlama-34b-hf",
    "instruct7": "codellama/CodeLlama-7b-Instruct-hf",
    "instruct13": "codellama/CodeLlama-13b-Instruct-hf",
    "instruct34": "codellama/CodeLlama-34b-Instruct-hf",
    # Training 70B requires experimental flag fsdp_peft_cpu_offload_for_save.
    "chat70": "meta-llama/Llama-2-70b-chat-hf",
    "base70": "meta-llama/Llama-2-70b-hf",
}

image = (
    Image.micromamba()
    .micromamba_install(
        "cudatoolkit=11.8",
        "cudnn=8.1.0",
        "cuda-nvcc",
        channels=["conda-forge", "nvidia"],
    )
    .apt_install("git")
    .pip_install(
        "llama-recipes @ git+https://github.com/modal-labs/llama-recipes.git@6636910761b70ada964409960129c5a4e9c2c049",
        extra_index_url="https://download.pytorch.org/whl/nightly/cu118",
        pre=True,
    )
    .pip_install("huggingface_hub==0.17.1", "hf-transfer==0.1.3", "scipy")
    .env(dict(HUGGINGFACE_HUB_CACHE="/pretrained", HF_HUB_ENABLE_HF_TRANSFER="1"))
)

stub = Stub("llama-finetuning", image=image, secrets=[Secret.from_name("huggingface")])

# Download pre-trained models into this volume.
stub.pretrained_volume = Volume.persisted("example-pretrained-vol")

# Save trained models into this volume.
stub.results_volume = Volume.persisted("example-results-vol")

VOLUME_CONFIG = {
    "/pretrained": stub.pretrained_volume,
    "/results": stub.results_volume,
}

common.py
modal-labs/llama-finetuning
import datasets

from llama_recipes.datasets.utils import Concatenator


def get_custom_dataset(dataset_config, tokenizer, split):
    full_dataset = datasets.load_dataset(
        "json", data_files="./modal_docs.jsonl", split="train"
    )

    # Since the dataset has no train/test split, we create one and select it
    dataset = full_dataset.train_test_split(
        train_size=0.95,
        test_size=0.05,
        seed=42,
    )["train" if split == dataset_config.train_split else "test"]

    dataset = dataset.map(
        lambda x: tokenizer(x["text"]), remove_columns=list(dataset.features)
    )

    dataset = dataset.map(Concatenator(), batched=True)

    return dataset

datasets/local_dataset.py
modal-labs/llama-finetuning
import datasets

from llama_recipes.datasets.utils import Concatenator


B_INST, E_INST = "[INST] ", " [/INST]"
B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"


def format_text(row, tokenizer):
    text = (
        B_INST
        + B_SYS
        + "You are an advanced SQL assistant that uses this SQL table schema to generate a SQL query which answers the user question.\n"
        + row["context"]
        + E_SYS
        + row["question"]
        + E_INST
        + "\n[SQL]\n"
        + row["answer"]
        + "\n[/SQL]"
        + "</s>"
    )

    return tokenizer(text)


def get_custom_dataset(dataset_config, tokenizer, split):
    full_dataset = datasets.load_dataset("b-mc2/sql-create-context", split="train")

    # Since the dataset has no train/test split, we create one and select it
    dataset = full_dataset.train_test_split(
        train_size=10000,
        test_size=200,
        seed=42,
    )["train" if split == dataset_config.train_split else "test"]

    dataset = dataset.map(
        lambda x: format_text(x, tokenizer), remove_columns=list(dataset.features)
    )

    dataset = dataset.map(Concatenator(), batched=True, batch_size=None)

    return dataset

datasets/sql_dataset.py
modal-labs/llama-finetuning
from modal import Image, gpu, method

import subprocess
import os

from common import stub, BASE_MODELS, VOLUME_CONFIG

tgi_image = (
    Image.from_registry("ghcr.io/huggingface/text-generation-inference:1.0.3")
    .dockerfile_commands("ENTRYPOINT []")
    .pip_install("text-generation", "transformers>=4.33.0")
    .env(dict(HUGGINGFACE_HUB_CACHE="/pretrained"))
)


@stub.function(image=tgi_image, volumes=VOLUME_CONFIG, timeout=60 * 20)
def merge(run_id: str, commit: bool = False):
    from text_generation_server.utils.peft import download_and_unload_peft

    os.mkdir(f"/results/{run_id}/merged")
    subprocess.call(f"cp /results/{run_id}/*.* /results/{run_id}/merged", shell=True)

    print(f"Merging weights for fine-tuned {run_id=}.")
    download_and_unload_peft(f"/results/{run_id}/merged", None, False)

    if commit:
        print("Committing merged model permanently (can take a few minutes).")
        stub.results_volume.commit()


@stub.cls(
    image=tgi_image,
    gpu=gpu.A100(count=1, memory=40),
    allow_concurrent_inputs=100,
    volumes=VOLUME_CONFIG,
)
class Model:
    def __init__(self, base: str = "", run_id: str = ""):
        from text_generation import AsyncClient
        import socket
        import time

        model = f"/results/{run_id}/merged" if run_id else BASE_MODELS[base]

        if run_id and not os.path.isdir(model):
            merge.local(run_id)  # local = run in the same container

        print(f"Loading {model} into GPU ... ")
        launch_cmd = ["text-generation-launcher", "--model-id", model, "--port", "8000"]
        self.launcher = subprocess.Popen(launch_cmd, stdout=subprocess.DEVNULL)

        self.client = None
        while not self.client and self.launcher.returncode is None:
            try:
                socket.create_connection(("127.0.0.1", 8000), timeout=1).close()
                self.client = AsyncClient("http://127.0.0.1:8000", timeout=60)
            except (socket.timeout, ConnectionRefusedError):
                time.sleep(1.0)

        assert self.launcher.returncode is None

    def __exit__(self, _exc_type, _exc_value, _traceback):
        self.launcher.terminate()

    @method()
    async def generate(self, prompt: str):
        result = await self.client.generate(prompt, max_new_tokens=512)

        return result.generated_text


@stub.local_entrypoint()
def main(prompt: str, base: str, run_id: str = "", batch: int = 1):
    print(f"Running completion for prompt:\n{prompt}")

    print("=" * 20 + "Generating without adapter" + "=" * 20)
    for output in Model(base).generate.map([prompt] * batch):
        print(output)

    if run_id:
        print("=" * 20 + "Generating with adapter" + "=" * 20)
        for output in Model(base, run_id).generate.map([prompt] * batch):
            print(output)

inference.py
modal-labs/llama-finetuning
from modal import gpu, Mount

from common import stub, N_GPUS, GPU_MEM, BASE_MODELS, VOLUME_CONFIG


@stub.function(
    volumes=VOLUME_CONFIG,
    memory=1024 * 100,
    timeout=3600 * 4,
)
def download(model_name: str):
    from huggingface_hub import snapshot_download
    from transformers.utils import move_cache

    try:
        snapshot_download(model_name, local_files_only=True)
        print(f"Volume contains {model_name}.")
    except FileNotFoundError:
        print(f"Downloading {model_name} (no progress bar) ...")
        snapshot_download(model_name)
        move_cache()

        print("Committing /pretrained directory (no progress bar) ...")
        stub.pretrained_volume.commit()


def library_entrypoint(config):
    from llama_recipes.finetuning import main

    main(**config)


@stub.function(
    volumes=VOLUME_CONFIG,
    mounts=[
        Mount.from_local_dir("./datasets", remote_path="/root"),
    ],
    gpu=gpu.A100(count=N_GPUS, memory=GPU_MEM),
    timeout=3600 * 12,
)
def train(train_kwargs):
    from torch.distributed.run import elastic_launch, parse_args, config_from_args

    torch_args = parse_args(["--nnodes", "1", "--nproc_per_node", str(N_GPUS), ""])
    print(f"{torch_args=}\n{train_kwargs=}")

    elastic_launch(
        config=config_from_args(torch_args)[0],
        entrypoint=library_entrypoint,
    )(train_kwargs)

    print("Committing results volume (no progress bar) ...")
    stub.results_volume.commit()


@stub.local_entrypoint()  # Runs locally to kick off remote training job.
def main(
    dataset: str,
    base: str = "chat7",
    run_id: str = "",
    num_epochs: int = 10,
    batch_size: int = 16,
):
    print(f"Welcome to Modal Llama fine-tuning.")

    model_name = BASE_MODELS[base]
    print(f"Syncing base model {model_name} to volume.")
    download.remote(model_name)

    if not run_id:
        import secrets

        run_id = f"{base}-{secrets.token_hex(3)}"
    elif not run_id.startswith(base):
        run_id = f"{base}-{run_id}"

    print(f"Beginning run {run_id=}.")
    train.remote(
        {
            "model_name": BASE_MODELS[base],
            "output_dir": f"/results/{run_id}",
            "batch_size_training": batch_size,
            "lr": 3e-4,
            "num_epochs": num_epochs,
            "val_batch_size": 1,
            # --- Dataset options ---
            "dataset": "custom_dataset",
            "custom_dataset.file": dataset,
            # --- FSDP options ---
            "enable_fsdp": True,
            "low_cpu_fsdp": True,  # Optimization for FSDP model loading (RAM won't scale with num GPUs)
            "fsdp_config.use_fast_kernels": True,  # Only works when FSDP is on
            "fsdp_config.fsdp_activation_checkpointing": True,  # Activation checkpointing for fsdp
            "pure_bf16": True,
            # --- Required for 70B ---
            "fsdp_config.fsdp_cpu_offload": True,
            "fsdp_peft_cpu_offload_for_save": True,  # Experimental
            # --- PEFT options ---
            "use_peft": True,
            "peft_method": "lora",
            "lora_config.r": 8,
            "lora_config.lora_alpha": 16,
        }
    )

    print(f"Training completed {run_id=}.")
    print(
        f"Test: `modal run inference.py --base {base} --run-id {run_id} --prompt '...'`."
    )

train.py
modal-labs/llama-finetuning
from modal import Mount, gpu

from common import stub, BASE_MODELS


@stub.function(
    volumes={
        "/pretrained": stub.pretrained_volume,
        "/results": stub.results_volume,
    },
    mounts=[
        Mount.from_local_dir("./datasets", remote_path="/root"),
    ],
)
def dataset(base: str = "chat7", dataset: str = "local_dataset.py"):
    from llama_recipes.utils.dataset_utils import get_custom_dataset
    from llama_recipes.configs.datasets import custom_dataset
    from llama_recipes.utils.config_utils import update_config

    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(BASE_MODELS[base])
    tokenizer.add_special_tokens({"pad_token": "<PAD>"})

    config = custom_dataset()
    update_config(config, file=dataset)

    BLOCK = "=" * 20

    for split in [config.train_split, config.test_split]:
        dataset = get_custom_dataset(config, tokenizer, split)
        print(f"{split}: {len(dataset)} sequences")

        sample = tokenizer.decode(dataset[0]["input_ids"])[:500]
        print(f"{BLOCK} Sample {BLOCK}\n{sample} ...")
        print(f"{BLOCK} Tokens {BLOCK}\n{dataset[0]['input_ids'][:25]} ...\n")

validate_dataset.py
modal-labs/llama-finetuning