Affinement LLM multi-nœud avec FSDP

info

Aperçu

Cette fonctionnalité est en aperçu public.

Cet exemple exécute l'affinement supervisé (SFT) de Llama-3.1-8B. sur 16 GPU H100 répartis sur 2 nœuds à l'aide de torchrun et de PyTorch Fully Sharded Data Parallel (FSDP). FSDP partitionne les paramètres du modèle, les gradients et les états de l'optimiseur sur les 16 rangs afin que le modèle de 8B de paramètres et son état d'optimiseur tiennent confortablement dans la mémoire GPU.

La charge de travail réalise les étapes suivantes :

Importe le projet local avec code_source: snapshot.
Lance un processus par GPU avec torchrun, en utilisant les variables d'environnement de rendez-vous que l'AI Runtime définit sur chaque nœud.
Lit un modèle à accès restreint de Hugging Face à l'aide d'un secret Databricks.
Logs metrics to MLflow et écrit le point de contrôle consolidé dans un volume Unity Catalog.

Prérequis

Le CLI air est installé et authentifié. Voir Installer le CLI AI Runtime.
Un volume Unity Catalog sur lequel vous pouvez écrire pour le point de contrôle de sortie.
Accès au modèle à accès restreint sur Hugging Face, plus un jeton d'accès stocké en tant que secret Databricks (voir ci-dessous).

Accéder au modèle sur Hugging Face

Llama-3.1-8B est un modèle à accès restreint, vous devez donc demander l'accès et fournir un jeton pour le download :

Ouvrez la page du modèle à l’adresse meta-llama/Llama-3.1-8B et acceptez la licence pour demander l’accès. Attendez que l’accès soit accordé.
Créez un jeton d'accès Hugging Face avec une autorisation de lecture .

Stocker le jeton comme un secret Databricks.

La charge de travail lit le jeton à partir d’un secret Databricks au lieu de le coder en dur. Créez un Secret Scope et ajoutez votre jeton :

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

train.yaml le référence en tant que my_scope/hf_token. Remplacez la portée et la clé par les vôtres.

Project Layout

Créez un répertoire avec les fichiers suivants.

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

Étape 1 : Rédiger le YAML de la charge de travail

train.yaml demande 16 GPU comme deux nœuds GPU_8xH100, monte le jeton Hugging Face en tant que secret et transmet les hyperparamètres au script via le bloc parameters. Les dépendances sont déclarées en ligne sous environment (avec l'image client version). Le package torch est inclus dans l'image de base d'AI Runtime, de sorte que seuls les éléments supplémentaires sont listés :

YAML
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

AI Runtime exécute command une fois par nœud et définit les variables d'environnement de rendez-vous (NUM_NODES, NODE_RANK, LOCAL_WORLD_SIZE, MASTER_ADDR et MASTER_PORT) sur chaque nœud. torchrun les lit pour lancer un processus par GPU, de sorte que la commande en ligne est le lanceur entier. Aucun script de lancement séparé n'est nécessaire.

Étape 2 : Écrire le script d’entraînement FSDP

train.py initialise le groupe de processus, encapsule chaque bloc de transformeur dans FSDP, s'entraîne sur un dataset d'instructions tokenisées et enregistre un point de contrôle consolidé à partir du rang 0. Les éléments clés :

Python
# 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,
)

Le rang 0 collecte l'intégralité du dictionnaire d'état (déchargé vers le CPU) et l'écrit dans le volume Unity Catalog :

Python
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)

Le script complet est répertorié dans Script d'entraînement complet à la fin de cette page.

Étape 3 : Soumettre l'exécution

Validez la configuration, puis soumettez et surveillez les Logs :

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

Étape 4 : Inspecter l'exécution

Les exécutions distribuées s'étendent sur plusieurs nœuds. Utilisez --node pour lire les logs d’un nœud spécifique :

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

Où les résultats aboutissent

**Métriques et paramètres :** Enregistrés dans l'Experimentation MLflow nommée experiment_name dans. Visualisez-les dans l'interface utilisateur MLflow du Workspace.
Point de contrôle affiné : Écrit dans le volume Unity Catalog de parameters.output_dir.

Script de formation complet

Le train.py complet pour le copier-coller :

Python
#!/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()

Prérequis​

Accéder au modèle sur Hugging Face​

Stocker le jeton comme un secret Databricks.​

Project Layout​

Étape 1 : Rédiger le YAML de la charge de travail​

Étape 2 : Écrire le script d’entraînement FSDP​

Étape 3 : Soumettre l'exécution​

Étape 4 : Inspecter l'exécution​

Où les résultats aboutissent​

Script de formation complet​

Ressources supplémentaires​