Dostrajanie modułu LLM z wieloma węzłami za pomocą protokołu FSDP

Ważna

Ta funkcja jest dostępna w publicznej wersji testowej.

Ten przykład uruchamia nadzorowane dostrajanie (SFT) modelu Llama-3.1-8B na 16 kartach GPU H100 rozmieszczonych w 2 węzłach z użyciem torchrun i PyTorch Fully Sharded Data Parallel (FSDP). FSDP dzieli parametry modelu, gradienty i stany optymalizatora na fragmenty między wszystkie 16 rang, dzięki czemu model o 8 miliardach parametrów i jego stan optymalizatora bez problemu mieszczą się w pamięci GPU.

Obciążenie robocze realizuje następujące działania:

• Przesyła lokalny projekt za pomocą polecenia code_source: snapshot.
• Uruchamia jeden proces na każdy procesor GPU za pomocą torchrun, używając zmiennych środowiskowych rendezvous, które AI Runtime ustawia na każdym węźle.
• Odczytuje model o ograniczonym dostępie z Hugging Face za pomocą klucza tajnego Databricks.
• Rejestruje metryki w MLflow i zapisuje skonsolidowany punkt kontrolny w wolumenie Unity Catalog.

Wymagania wstępne

• Interfejs wiersza poleceń air został zainstalowany i uwierzytelniony. Zobacz Instalowanie interfejsu wiersza polecenia środowiska uruchomieniowego AI.
• Wolumin Unity Catalog, do którego można zapisywać dane wyjściowego punktu kontrolnego.
• Dostęp do modelu z ograniczonym dostępem w serwisie Hugging Face oraz token dostępu przechowywany jako sekret Databricks (patrz poniżej).

Uzyskaj dostęp do modelu na Hugging Face

Llama-3.1-8B to model z bramką, więc musisz zażądać dostępu i podać token, aby go pobrać:

1. Otwórz stronę modelu w meta-llama/Llama-3.1-8B i zaakceptuj licencję, aby zażądać dostępu. Poczekaj na udzielenie dostępu.
2. Utwórz token dostępu Hugging Face z uprawnieniem do odczytu.

Przechowaj token jako wpis tajny w usłudze Databricks

Zadanie robocze odczytuje token z wpisu tajnego w Databricks, zamiast wpisywać go na stałe w kodzie. Utwórz tajny zakres i dodaj swój token:

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

train.yaml odwołuje się do niego jako my_scope/hf_token. Zastąp zakres i klucz własnymi.

Układ projektu

Utwórz katalog z następującymi plikami.

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

Krok 1. Zapisywanie obciążenia YAML

train.yaml wymaga 16 procesorów GPU w postaci dwóch węzłów GPU_8xH100, montuje token Hugging Face jako sekret i przekazuje hiperparametry do skryptu za pośrednictwem bloku parameters. Zależności są deklarowane bezpośrednio pod environment (z obrazem klienta version). Pakiet torch jest zawarty w obrazie bazowym środowiska AI Runtime, dlatego wymieniono tylko pakiety dodatkowe:

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

Środowisko uruchomieniowe sztucznej inteligencji jest uruchamiane command raz na węzeł i ustawia zmienne środowiskowe rendezvous (NUM_NODES, NODE_RANK, LOCAL_WORLD_SIZEMASTER_ADDR, i MASTER_PORT) w każdym węźle. torchrun odczytuje je, aby uruchomić po jednym procesie na każdy układ GPU, więc polecenie podane bezpośrednio stanowi cały program uruchamiający. Nie jest wymagany żaden oddzielny skrypt uruchamiania.

Krok 2. Pisanie skryptu trenowania FSDP

train.py Inicjuje grupę procesów, opakowuje każdy blok transformatora w fsDP, trenuje na tokenizowanym zestawie danych instrukcji i zapisuje skonsolidowany punkt kontrolny z rangi 0. Kluczowe elementy:

# 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 0 zbiera pełny słownik stanu (przeniesiony do pamięci CPU) i zapisuje go w woluminie 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)

Pełny skrypt znajduje się w sekcji Pełny skrypt szkoleniowy na końcu tej strony.

Krok 3: Prześlij uruchomienie

Zweryfikuj konfigurację, a następnie prześlij i obejrzyj dzienniki:

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

Krok 4. Sprawdzanie przebiegu

Rozproszone uruchomienia obejmują wiele węzłów. Użyj --node, aby odczytać dzienniki z określonego węzła:

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

Gdzie wyniki lądują

• Metryki i parametry: rejestrowane w eksperymencie MLflow o nazwie podanej w experiment_name. Wyświetl je w interfejsie użytkownika MLflow w przestrzeni roboczej.
• Punkt kontrolny po dostrojeniu: Zapisywany w woluminie Unity Catalog w parameters.output_dir.

Pełny skrypt szkoleniowy

Kompletny train.py do kopiowania i wklejania:

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

Dodatkowe zasoby

• Trenowanie rozproszone za pomocą programu Ray Train
• Informacje referencyjne YAML dla obciążenia roboczego