Vyladění LLM s více uzly pomocí FSDP

Important

Tato funkce je ve verzi Public Preview.

V tomto příkladu se provádí supervizované doladění (SFT) modelu Llama-3.1-8B na 16 GPU H100 rozložených do 2 uzlů pomocí torchrun a PyTorch Fully Sharded Data Parallel (FSDP). FSDP rozděluje parametry modelu, gradienty a stavy optimalizátoru napříč všemi 16 ranky, takže se model s 8 miliardami parametrů i jeho stav optimalizátoru bez problémů vejdou do paměti GPU.

Úloha provádí následující akce:

• Nahraje místní projekt pomocí code_source: snapshot.
• Spustí jeden proces na GPU pomocí torchrun, přičemž používá proměnné prostředí pro rendezvous, které AI Runtime nastavuje na každém uzlu.
• Načte model s omezeným přístupem z Hugging Face pomocí tajného klíče v Databricks.
• Zaznamená metriky do MLflow a zapíše konsolidovaný kontrolní bod do svazku katalogu Unity.

Prerequisites

• Rozhraní příkazového air řádku se nainstalovalo a ověřilo. Viz Instalace rozhraní příkazového řádku AI Runtime.
• Svazek v katalogu Unity, do kterého můžete zapisovat data pro kontrolní bod výstupu.
• Přístup k chráněnému modelu na Hugging Face a přístupový token uložený jako tajný kód Databricks (viz níže).

Získání přístupu k modelu na Hugging Face

Llama-3.1-8B je model s omezeným přístupem, takže musíte požádat o přístup a poskytnout token k jeho stažení:

1. Otevřete stránku modelu na adrese meta-llama/Llama-3.1-8B a přijměte licenci, abyste mohli požádat o přístup. Počkejte, až se udělí přístup.
2. Vytvořte přístupový token Hugging Face s oprávněním ke čtení .

Uložit token jako tajný klíč Databricks

Úloha načte token z tajného kódu Databricks místo jeho pevného kódování. Vytvořte secret scope a přidejte svůj token:

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

train.yaml odkazuje na ni jako my_scope/hf_token. Nahraďte obor a klíč vlastním oborem.

Rozložení projektu

Vytvořte adresář s následujícími soubory.

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

Krok 1: Zápis YAML úlohy

train.yaml požaduje 16 GPU ve formě dvou uzlů GPU_8xH100, připojí token Hugging Face jako secret a předá hyperparametry skriptu prostřednictvím bloku parameters. Závislosti jsou deklarovány přímo pod environment (s bitovou kopií klienta version). Balíček torch je součástí základní image AI Runtime, takže jsou uvedeny pouze doplňky:

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

Modul AI Runtime běží command jednou pro každý uzel a na každém uzlu nastavuje proměnné prostředí pro rendezvous (NUM_NODES, NODE_RANK, LOCAL_WORLD_SIZE, MASTER_ADDR a MASTER_PORT). torchrun je přečte a spustí tak samostatný proces pro každé GPU, takže příkaz na řádku představuje celý spouštěč. Není potřeba žádný samostatný spouštěcí skript.

Krok 2: Zápis trénovacího skriptu FSDP

train.py inicializuje skupinu procesů, zabalí každý blok transformátoru v FSDP, trénuje na tokenizované datové sadě instrukcí a uloží konsolidovaný kontrolní bod z pořadí 0. Klíčové části:

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

Proces s hodností 0 shromáždí celý slovník stavu (přesunutý do paměti CPU) a zapíše jej do svazku 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)

Úplný skript je uveden v úplném trénovacím skriptu na konci této stránky.

Krok 3: Odeslání spuštění

Ověřte konfiguraci a pak odešlete a sledujte protokoly:

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

Krok 4: Kontrola spuštění

Distribuované běhy probíhají napříč více uzly. Slouží --node ke čtení protokolů z určitého uzlu:

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

Kam se výsledky uloží

• Metriky a parametry: Zaznamenáno do experimentu MLflow pojmenovaného v experiment_name. Zobrazte je v uživatelském rozhraní MLflow v pracovním prostoru.
• Kontrolní bod po doladění: Zapsáno do svazku v katalogu Unity v parameters.output_dir.

Úplný trénovací skript

Kompletní train.py pro kopírování a vložení:

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

Další zdroje informací

• Distribuované trénování pomocí Ray Trainu
• Referenční informace k YAML pro úlohy