FSDP ile çok düğümlü LLM ince ayarı

Important

Bu özellik Genel Önizleme aşamasındadır.

Bu örnek, ve PyTorch torchrun kullanarak 2 düğüme yayılmış 16 H100 GPU üzerinde Llama-3.1-8B modelinin denetimli ince ayarını (SFT) çalıştırır. FSDP, model parametrelerini, gradyanları ve optimizer durumlarını 16 rankın tümüne dağıtarak 8 milyar parametreli modelin ve optimizer durumunun GPU belleğine rahatça sığmasını sağlar.

İş yükü aşağıdakileri yapar:

• Yerel projeyi code_source: snapshot ile karşıya yükler.
• AI Runtime’ın her düğümde ayarladığı randevu ortam değişkenlerini kullanarak, torchrun ile GPU başına bir işlem başlatır.
• Databricks gizli anahtarını kullanarak Hugging Face'ten erişim kısıtlamalı bir modeli okur.
• Ölçümleri MLflow'a kaydeder ve birleştirilmiş denetim noktasını bir Unity Kataloğu birimine yazar.

Prerequisites

• air CLI yüklendi ve kimliği doğrulandı. Bkz. AI Runtime CLI'yı yükleme.
• Çıktı denetim noktası için yazma iznine sahip olduğunuz bir Unity Catalog birimi.
• Hugging Face'teki erişim kısıtlamalı modele erişim ve Databricks gizli anahtarı olarak depolanan bir erişim belirteci (aşağıya bakın).

Hugging Face'teki modele erişin

Llama-3.1-8B erişim kısıtlamalı bir modeldir, bu nedenle onu indirmek için erişim talebinde bulunmanız ve bir erişim belirteci sağlamanız gerekir:

1. Meta-llama/Llama-3.1-8B konumundaki model sayfasını açın ve erişim isteğinde bulunma lisansını kabul edin. Erişim verilene kadar bekleyin.
2. Hugging Face erişim belirteci için okuma izni oluşturun.

Belirteci Databricks gizli anahtarı olarak depolayın

İş yükü, belirteci sabit kodlamak yerine onu bir Databricks secret'ından okur. Gizli bir kapsam oluşturun ve token'ınızı ekleyin:

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

train.yaml, bundan my_scope/hf_token olarak bahseder. Kapsamı ve anahtarı kendinizinkilerle değiştirin.

Proje düzeni

Aşağıdaki dosyaları içeren bir dizin oluşturun.

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

1. Adım: İş yükü YAML'sini yazma

train.yaml, 16 GPU’yu iki GPU_8xH100 düğüm olarak talep eder, Hugging Face belirtecini bir secret olarak bağlar ve hiperparametreleri parameters bloğu aracılığıyla betiğe iletir. Bağımlılıklar, environment altında satır içinde bildirilir (istemci görüntüsü version ile birlikte). torch paketi AI Runtime temel görüntüsüyle birlikte gelir, bu nedenle burada yalnızca ekler listelenmiştir:

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 Çalışma Zamanı düğüm başına bir kez çalışır command ve her düğümde randevu ortam değişkenlerini (NUM_NODES, NODE_RANK, LOCAL_WORLD_SIZE, MASTER_ADDRve MASTER_PORT) ayarlar. torchrun GPU başına bir işlem başlatmak için bunları okur; dolayısıyla satır içi komutun tamamı başlatıcıdır. Ayrı başlatıcı betiği gerekmez.

2. Adım: FSDP eğitim betiğini yazma

train.py işlem grubunu başlatır, her transformatör bloğunu FSDP ile sarar, tokenleştirilmiş bir komut veri kümesi üzerinde eğitir ve rank 0'dan konsolide edilmiş bir denetim noktası kaydeder. Anahtar parçalar:

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

Rank 0, tam durum sözlüğünü toplar (CPU’ya aktarılmıştır) ve bunu Unity Catalog volume’üne yazar:

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)

Betiğin tamamı, bu sayfanın sonundaki Tam eğitim betiği bölümünde listelenmiştir.

3. Adım: Çalıştırmayı gönderme

Yapılandırmayı doğrulayın, ardından gönderin ve günlükleri izleyin:

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

4. Adım: Çalıştırmayı inceleyin

Dağıtılmış çalıştırmalar birden çok düğümü kapsar. Belirli bir düğümden günlükleri okumak için --node kullanın:

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

Sonuçların geldiği yer

• Metrikler ve parametreler: Adı experiment_name içinde belirtilen MLflow deneyine kaydedilir. Bunları çalışma alanı MLflow kullanıcı arabiriminde görüntüleyin.
• İnce ayarlanmış kontrol noktası: parameters.output_dir içindeki Unity Catalog birimine yazılır.

Tam eğitim betiği

Kopyala-yapıştır için eksiksiz train.py:

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

Ek kaynaklar

• Ray Train ile dağıtılmış eğitim
• İş yükü YAML başvurusu