Ottimizzazione LLM a più nodi con FSDP

Questo esempio esegue l'ottimizzazione con supervisione (SFT) di Llama-3.1-8B in 16 GPU H100 distribuite su 2 nodi usando torchrun e PyTorch Fully Sharded Data Parallel (FSDP). FSDP suddivide tra tutti i 16 processi i parametri del modello, i gradienti e gli stati dell'ottimizzatore, in modo che il modello da 8 miliardi di parametri e il relativo stato dell'ottimizzatore possano rientrare comodamente nella memoria della GPU.

Il carico di lavoro esegue le operazioni seguenti:

• Carica il progetto locale con code_source: snapshot.
• Avvia un processo per GPU con torchrun, usando le variabili di ambiente di rendezvous impostate dal runtime di intelligenza artificiale in ogni nodo.
• Legge un modello ad accesso limitato da Hugging Face tramite un segreto Databricks.
• Registra le metriche in MLflow e scrive il checkpoint consolidato in un volume di Unity Catalog.

Prerequisites

• La air CLI è installata e autenticata. Consulta Installazione dell'AI Runtime CLI.
• Un volume del catalogo Unity in cui è possibile scrivere per il checkpoint di output.
• Accesso al modello ad accesso limitato su Hugging Face, oltre a un token di accesso memorizzato come segreto di Databricks (vedere di seguito).

Ottenere l'accesso al modello in Hugging Face

Llama-3.1-8B è un modello gestito, quindi è necessario richiedere l'accesso e fornire un token per scaricarlo:

1. Aprire la pagina del modello in meta-llama/Llama-3.1-8B e accettare la licenza per richiedere l'accesso. Attendere fino a quando non viene concesso l'accesso.
2. Creare un token di accesso Hugging Face con autorizzazione di lettura .

Memorizzare il token come segreto di Databricks

Il carico di lavoro legge il token da un segreto di Databricks invece di codificarlo direttamente nel codice. Creare un ambito segreto e aggiungere il token:

databricks secrets create-scope my_scope
databricks secrets put-secret my_scope hf_token

train.yaml fa riferimento a esso come my_scope/hf_token. Sostituisci l'ambito e la chiave con i valori corrispondenti ai tuoi.

Layout del progetto

Creare una directory con i file seguenti.

multinode_llm_sft/
├── train.yaml          # air workload config (inline dependencies + torchrun launcher)
└── train.py            # FSDP fine-tuning script

Passaggio 1: Scrivere il carico di lavoro YAML

train.yaml richiede 16 GPU sotto forma di due nodi GPU_8xH100, monta il token di Hugging Face come segreto e passa gli iperparametri allo script tramite il blocco parameters. Le dipendenze vengono dichiarate in linea sotto environment (con l'immagine client version). Il torch pacchetto viene fornito nell'immagine di base del runtime di intelligenza artificiale, quindi sono elencati solo i componenti aggiuntivi:

experiment_name: air-multinode-llama-sft

environment:
  version: '4'
  dependencies:
    - transformers>=4.45
    - datasets>=3.0
    - huggingface_hub>=0.34
    - accelerate>=0.34
    # The base image ships fsspec 2023.5.0, which is too old for modern
    # huggingface_hub and breaks dataset/model downloads. Pin a newer fsspec.
    - fsspec>=2024.6.1

# 16 GPUs across 2 nodes (GPU_8xH100 = 8 H100 per node).
compute:
  num_accelerators: 16
  accelerator_type: GPU_8xH100

code_source:
  type: snapshot
  snapshot:
    root_path: .

command: |
  cd $CODE_SOURCE_PATH
  # air sets NUM_NODES, NODE_RANK, LOCAL_WORLD_SIZE, MASTER_ADDR, and MASTER_PORT on each node.
  torchrun \
    --nnodes="$NUM_NODES" \
    --node_rank="$NODE_RANK" \
    --nproc_per_node="${LOCAL_WORLD_SIZE:-8}" \
    --master_addr="$MASTER_ADDR" \
    --master_port="$MASTER_PORT" \
    train.py

# Pin NCCL control-plane traffic to eth0 so cross-node rendezvous works.
env_variables:
  NCCL_SOCKET_IFNAME: eth0
  HF_HOME: /tmp/hf

# Gated model download needs a Hugging Face token. Replace with your own
# Databricks secret in the form "scope/key".
secrets:
  HF_TOKEN: 'my_scope/hf_token'

max_retries: 1
timeout_minutes: 120

# Surfaced to train.py via HYPERPARAMETERS_PATH.
parameters:
  model_name: meta-llama/Llama-3.1-8B
  dataset_name: tatsu-lab/alpaca
  max_seq_len: 1024
  per_device_batch_size: 4
  gradient_accumulation_steps: 2
  learning_rate: 0.00002
  max_steps: 100
  output_dir: /Volumes/main/default/air_checkpoints/llama31-8b-sft

Il runtime AI viene eseguito command una volta per nodo e imposta le variabili di ambiente di rendezvous (NUM_NODES, NODE_RANK, LOCAL_WORLD_SIZE, MASTER_ADDR e MASTER_PORT) su ogni nodo. torchrun li usa per avviare un processo per ogni GPU, quindi il comando in linea costituisce l'intero comando di avvio. Non è necessario alcun script di avvio separato.

Passaggio 2: Scrivere lo script di training FSDP

train.py inizializza il gruppo di processi, esegue il wrapping di ogni blocco di trasformatori in FSDP, esegue il training su un set di dati di istruzioni con token e salva un checkpoint consolidato dal rango 0. Le parti chiave:

# Shard each transformer block independently so no single GPU holds the full model.
auto_wrap_policy = functools.partial(
    transformer_auto_wrap_policy, transformer_layer_cls={LlamaDecoderLayer}
)
model = FSDP(
    model,
    auto_wrap_policy=auto_wrap_policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    mixed_precision=MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
        buffer_dtype=torch.bfloat16,
    ),
    device_id=local_rank,
    use_orig_params=True,
)

Il rank 0 raccoglie lo state dict completo (scaricato sulla CPU) e lo scrive nel volume di Unity Catalog:

save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, save_policy):
    cpu_state = model.state_dict()
if rank == 0:
    model.module.save_pretrained(output_dir, state_dict=cpu_state)
    tokenizer.save_pretrained(output_dir)

Lo script completo è elencato in Script di training completo alla fine di questa pagina.

Passaggio 3: Invia il run

Convalidare la configurazione, quindi inviare e controllare i log:

air run -f train.yaml --dry-run
air run -f train.yaml --watch

Passaggio 4: Controllare l'esecuzione

Le esecuzioni distribuite si estendono su più nodi. Usare --node per leggere i log da un nodo specifico:

air get run <run-id>
air logs <run-id> --node 0
air logs <run-id> --node 1

Dove atterrare i risultati

• Metriche e parametri: registrati nell'esperimento MLflow denominato in experiment_name. Visualizzarli nell'interfaccia utente MLflow dell'area di lavoro.
• Checkpoint sottoposto a fine-tuning: scritto nel volume di Unity Catalog in parameters.output_dir.

Script di training completo

Il file completo train.py per la copia-incolla:

#!/usr/bin/env python3
"""Multi-node FSDP supervised fine-tuning of Llama-3.1-8B.

Launched via ``torchrun`` from the workload YAML ``command`` across 2 nodes x 8 H100 (16 ranks). Each rank
owns one GPU. The model is sharded with PyTorch FSDP (full shard + bf16), trained on
an instruction dataset, and the consolidated checkpoint is written to a Unity Catalog
Volume by rank 0. Metrics are logged to MLflow.

Hyperparameters are read from the YAML block passed by ``air`` via HYPERPARAMETERS_PATH.
"""

import functools
import os

import mlflow
import torch
import torch.distributed as dist
import yaml
from datasets import load_dataset
from torch.distributed.fsdp import FullStateDictConfig, FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision, ShardingStrategy, StateDictType
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from torch.utils.data import DataLoader, DistributedSampler
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.models.llama.modeling_llama import LlamaDecoderLayer


def load_params() -> dict:
    """Read the hyperparameters block that `air` materializes from the YAML `parameters:`."""
    path = os.environ.get("HYPERPARAMETERS_PATH")
    if path and os.path.exists(path):
        with open(path) as f:
            return yaml.safe_load(f) or {}
    return {}


def build_dataset(tokenizer, dataset_name: str, max_seq_len: int):
    """Tokenize an instruction dataset into fixed-length causal-LM examples."""
    raw = load_dataset(dataset_name, split="train")

    def format_example(row):
        instruction = row.get("instruction", "")
        context = row.get("input", "")
        response = row.get("output", "")
        prompt = f"### Instruction:\n{instruction}\n\n"
        if context:
            prompt += f"### Input:\n{context}\n\n"
        text = f"{prompt}### Response:\n{response}{tokenizer.eos_token}"
        out = tokenizer(text, truncation=True, max_length=max_seq_len, padding="max_length")
        out["labels"] = out["input_ids"].copy()
        return out

    cols = raw.column_names
    tokenized = raw.map(format_example, remove_columns=cols)
    # Emit torch tensors so the default DataLoader collate stacks them into [B, L] batches.
    tokenized.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])
    return tokenized


def main():
    rank = int(os.environ["RANK"])
    local_rank = int(os.environ["LOCAL_RANK"])
    world_size = int(os.environ["WORLD_SIZE"])

    dist.init_process_group(backend="nccl")
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")

    p = load_params()
    model_name = p.get("model_name", "meta-llama/Llama-3.1-8B")
    dataset_name = p.get("dataset_name", "tatsu-lab/alpaca")
    max_seq_len = int(p.get("max_seq_len", 1024))
    batch_size = int(p.get("per_device_batch_size", 4))
    grad_accum = int(p.get("gradient_accumulation_steps", 2))
    lr = float(p.get("learning_rate", 2e-5))
    max_steps = int(p.get("max_steps", 100))
    output_dir = p.get("output_dir", "/tmp/llama-sft")

    if rank == 0:
        print(f"World size={world_size} | model={model_name} | dataset={dataset_name}", flush=True)

    # --- Model & data --------------------------------------------------------
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
    model.config.use_cache = False  # incompatible with gradient checkpointing / FSDP training
    model.gradient_checkpointing_enable()

    # Shard each transformer block independently so no single GPU holds the full model.
    auto_wrap_policy = functools.partial(transformer_auto_wrap_policy, transformer_layer_cls={LlamaDecoderLayer})
    model = FSDP(
        model,
        auto_wrap_policy=auto_wrap_policy,
        sharding_strategy=ShardingStrategy.FULL_SHARD,
        mixed_precision=MixedPrecision(
            param_dtype=torch.bfloat16,
            reduce_dtype=torch.bfloat16,
            buffer_dtype=torch.bfloat16,
        ),
        device_id=local_rank,
        use_orig_params=True,
    )

    dataset = build_dataset(tokenizer, dataset_name, max_seq_len)
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
    loader = DataLoader(dataset, batch_size=batch_size, sampler=sampler, drop_last=True)

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)

    # --- MLflow (rank 0 only) ------------------------------------------------
    # AI Runtime injects MLFLOW_RUN_ID and configures the databricks tracking URI on
    # the node, so logging works without DATABRICKS_HOST/TOKEN. Gate on MLFLOW_RUN_ID
    # so the script also runs cleanly off-platform (e.g. locally) where it is unset.
    use_mlflow = rank == 0 and bool(os.environ.get("MLFLOW_RUN_ID"))
    if use_mlflow:
        mlflow.start_run(run_id=os.environ.get("MLFLOW_RUN_ID"))
        mlflow.log_params({"model_name": model_name, "lr": lr, "batch_size": batch_size, "world_size": world_size})

    # --- Training loop -------------------------------------------------------
    model.train()
    sampler.set_epoch(0)
    step = 0
    optimizer.zero_grad()
    for micro_step, batch in enumerate(loader):
        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels = batch["labels"].to(device)

        out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
        (out.loss / grad_accum).backward()

        if (micro_step + 1) % grad_accum == 0:
            model.clip_grad_norm_(1.0)
            optimizer.step()
            optimizer.zero_grad()
            step += 1
            if rank == 0:
                print(f"step={step}/{max_steps} loss={out.loss.item():.4f}", flush=True)
                if use_mlflow:
                    mlflow.log_metric("train_loss", out.loss.item(), step=step)
            if step >= max_steps:
                break

    # --- Save consolidated checkpoint to the UC Volume (rank 0) --------------
    save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
    with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, save_policy):
        cpu_state = model.state_dict()
    if rank == 0:
        os.makedirs(output_dir, exist_ok=True)
        model.module.save_pretrained(output_dir, state_dict=cpu_state)
        tokenizer.save_pretrained(output_dir)
        print(f"Saved checkpoint to {output_dir}", flush=True)
        if use_mlflow:
            mlflow.end_run()

    dist.barrier()
    dist.destroy_process_group()


if __name__ == "__main__":
    main()

Risorse aggiuntive

• Training distribuito con Ray Train
• Informazioni di riferimento su YAML del carico di lavoro