重要
這項功能目前處於 公開預覽版。
此範例使用 以及 PyTorch torchrun，在分散於 2 個節點上的 16 顆 H100 GPU 上對 Llama-3.1-8B 進行監督式微調（SFT）。
FSDP 將模型參數、梯度與優化器狀態分片涵蓋所有 16 個等級，使 8B 參數模型及其優化器狀態能舒適地放入 GPU 記憶體中。
工作量會做出以下效果：
- 使用
code_source: snapshot上傳本地專案。 - 使用 AI 執行時在每個節點上設定的同步環境變數，以
torchrun為每張 GPU 啟動一個處理程序。 - 使用 Databricks 密鑰讀取 Hugging Face 上的受限存取模型。
- 將度量記錄到 MLflow，並將整合後的檢查點寫入 Unity Catalog 磁碟區。
先決條件
-
airCLI 已安裝並完成驗證。 請參閱 安裝 AI 執行階段 CLI。 - 可供寫入輸出檢查點的 Unity Catalog 磁碟區。
- 對 Hugging Face 上受限制模型的存取權限，以及儲存在 Databricks Secret 中的存取權杖（見下文）。
在 Hugging Face 上取得該模型的存取權
Llama-3.1-8B 是門控式模型，因此您必須申請存取權並提供令牌才能下載：
- 開啟模型頁面 meta-llama/Llama-3.1-8B ，並接受授權以申請存取權限。 等到存取權限被核准後再說。
- 建立一個擁有閱讀權限的 Hugging Face 存取權杖。
將該令牌儲存為 Databricks 秘密
工作負載會從 Databricks 密鑰 讀取權杖，而不是將其硬式編碼。 建立一個秘密範圍並加入你的代幣：
databricks secrets create-scope my_scope
databricks secrets put-secret my_scope hf_token
train.yaml 將其引用為 my_scope/hf_token。 把瞄準鏡和鑰匙換成你自己的。
專案版面配置
建立一個包含以下檔案的目錄。
multinode_llm_sft/
├── train.yaml # air workload config (inline dependencies + torchrun launcher)
└── train.py # FSDP fine-tuning script
步驟 1：撰寫 YAML 工作負載
train.yaml 請求 16 顆 GPU 作為兩個 GPU_8xH100 節點，將 Hugging Face 標記掛載為秘密，並透過區 parameters 塊將超參數傳遞給腳本。 相依關係會在 environment 下方以內嵌方式宣告（搭配用戶端映像 version）。
torch 套件已隨附於 AI 執行階段的基底映像中，因此這裡只列出額外項目：
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 會在每個節點上執行一次 command，並在每個節點上設定 rendezvous 環境變數（NUM_NODES、NODE_RANK、LOCAL_WORLD_SIZE、MASTER_ADDR 和 MASTER_PORT）。
torchrun 讀取指令來啟動每個 GPU 的一個程序，所以內嵌指令就是整個啟動器。
不需要另外的啟動器腳本。
步驟二：撰寫FSDP訓練腳本
train.py 初始化程序群組，將每個變壓器區塊包裹在 FSDP，訓練於分詞化指令資料集，並將合併檢查點從第 0 階儲存下來。 關鍵部分：
# 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 會彙整完整的狀態字典（已卸載到 CPU），並將其寫入 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)
完整腳本列在本頁末尾的 完整訓練腳本 中。
步驟三：提交跑量
驗證設定，然後提交並監控日誌：
air run -f train.yaml --dry-run
air run -f train.yaml --watch
步驟 4：檢查執行結果
分散式執行跨越多個節點。 用 --node 來讀取特定節點的日誌：
air get run <run-id>
air logs <run-id> --node 0
air logs <run-id> --node 1
結果落在哪裡
-
指標與參數：已記錄至名為
experiment_name的 MLflow 實驗。 在工作區的 MLflow UI 中檢視它們。 -
微調後的檢查點：寫入
parameters.output_dir中的 Unity Catalog 磁碟區。
完整訓練腳本
完整的 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()