Pelatihan terdistribusi dengan Ray Train

Important

Fitur ini ada di Pratinjau Umum.

Contoh ini menjalankan fine-tuning data-paralel terdistribusi dengan Ray Train milik TorchTrainer pada 8 GPU H100 dalam satu node. Skrip bootstrap menjalankan kluster Ray pada node, lalu driver Ray Train meluncurkan satu worker untuk setiap GPU, membungkus model dengan DDP, dan membagi dataset menjadi shard di seluruh worker secara otomatis.

Ini melakukan fine-tuning pada model publik (Qwen2.5-3B), sehingga model tersebut dapat langsung digunakan tanpa token Hugging Face.

Beban kerja ini melakukan hal berikut:

  • Mengunggah proyek lokal dengan code_source: snapshot.
  • Memulai kepala Ray dengan semua 8 GPU, kemudian menjalankan driver Ray Train.
  • Menggunakan ray.train.torch.prepare_model dan prepare_data_loader untuk menangani pembungkusan DDP, penempatan perangkat, dan pengambilan sampel terdistribusi.
  • Mencatat metrik ke MLflow.

Prasyarat

Tata letak proyek

Buat direktori dengan file berikut.

ray_train_distributed/
├── train.yaml          # air workload config (inline dependencies + Ray bootstrap)
└── train_ray.py        # Ray Train TorchTrainer driver + per-worker training

Langkah 1: Tulis YAML beban kerja

train.yaml meminta satu GPU_8xH100 node. Dependensi dinyatakan sebaris di bawah environment (dengan runtime version), dan command memulai kluster Ray pada node kemudian menjalankan driver, sehingga beban kerja tidak memerlukan file dependensi terpisah atau skrip peluncur:

experiment_name: air-ray-train-distributed

environment:
  version: 'databricks_ai_v5'
  dependencies:
    # databricks_ai_v5 (the databricks-ai runtime) already ships ray, transformers, datasets, and
    # huggingface_hub, so they no longer need to be listed here. It does ship fsspec
    # 2023.5.0, which is too old for modern huggingface_hub and breaks dataset/model
    # downloads, so pin a newer fsspec.
    - fsspec>=2024.6.1

# 8 H100 on a single node. Ray Train launches one worker per GPU.
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100

code_source:
  type: snapshot
  snapshot:
    root_path: .

command: |
  cd $CODE_SOURCE_PATH
  RAY_HEAD_PORT=6379
  GPUS_PER_NODE=${LOCAL_WORLD_SIZE:-8}
  if [ "${NODE_RANK:-0}" = "0" ]; then
    echo "NODE_RANK=0: starting Ray head with $GPUS_PER_NODE GPU(s)..."
    ray start --head --port=$RAY_HEAD_PORT --num-gpus="$GPUS_PER_NODE" --dashboard-host=0.0.0.0
    python train_ray.py
    ray stop
  else
    echo "NODE_RANK=$NODE_RANK: connecting to Ray head at $MASTER_ADDR:$RAY_HEAD_PORT..."
    for i in $(seq 1 12); do
      if ray start --address="$MASTER_ADDR:$RAY_HEAD_PORT" --num-gpus="$GPUS_PER_NODE" --block 2>/dev/null; then
        break
      fi
      echo "Attempt $i failed, retrying in 5s..."
      sleep 5
    done
  fi

max_retries: 0
timeout_minutes: 90
env_variables:
  NCCL_SOCKET_IFNAME: eth0

Sebaris command memulai kepala Ray dengan semua GPU pada simpul, menjalankan driver dengan python train_ray.py, lalu menghentikan kluster. Ini juga termasuk cabang pekerja yang bergabung dengan kepala, sehingga perintah yang sama terus berfungsi jika Anda menskalakan pekerjaan ke beberapa simpul.

Langkah 2: Tentukan driver Ray Train

train_ray.py mendefinisikan train_func yang berjalan di setiap worker dan main yang mengonfigurasi TorchTrainer untuk menggunakan semua GPU di cluster. prepare_model membungkus model dalam DDP dan memindahkannya ke GPU pekerja. prepare_data_loader menambahkan sampler terdistribusi:

def train_func(config: dict):
    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.bfloat16)
    model.config.use_cache = False
    model = prepare_model(model)              # DDP wrap + device placement

    loader = DataLoader(dataset, batch_size=config["batch_size"], shuffle=True, drop_last=True)
    loader = prepare_data_loader(loader)      # distributed sampler + GPU transfer
    optimizer = torch.optim.AdamW(model.parameters(), lr=config["lr"])
    ...
    ray.train.report({"loss": out.loss.item(), "step": step})


def main():
    ray.init(address="auto")
    total_gpus = int(ray.cluster_resources().get("GPU", 0))
    trainer = TorchTrainer(
        train_func,
        train_loop_config={"lr": 2e-5, "batch_size": 4, "max_steps": 100},
        scaling_config=ScalingConfig(num_workers=total_gpus, use_gpu=True),
    )
    trainer.fit()

Skrip lengkap tercantum dalam Skrip pelatihan lengkap di akhir halaman ini.

Langkah 3: Kirim eksekusi

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

Langkah 4: Periksa proses eksekusi

air get run <run-id>
air logs <run-id>

Kepala Ray dan driver keduanya berjalan pada node 0, sehingga log mengalir dari satu simpul.

Tempat hasil ditampilkan

Metrik yang dilaporkan dengan ray.train.report dan dicatat dengan MLflow muncul di eksperimen MLflow bernama di experiment_name, dapat dilihat di ruang kerja MLflow UI.

Skrip pelatihan penuh

train_ray.py lengkap untuk salin-tempel:

#!/usr/bin/env python3
"""Distributed data-parallel fine-tuning with Ray Train on a single 8x H100 node.

The workload `command` starts a Ray head with 8 GPUs and runs this script. Ray Train's
TorchTrainer launches one worker per GPU (8 total), wraps the model in DDP, shards
the dataset across workers, and aggregates metrics. Each worker runs `train_func`.

Uses a public model (no Hugging Face token required) so the example runs as-is.
"""

import os

import mlflow
import ray
import ray.train
import torch
from datasets import load_dataset
from ray.train import RunConfig, ScalingConfig
from ray.train.torch import TorchTrainer, prepare_data_loader, prepare_model
from torch.utils.data import DataLoader
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_NAME = "Qwen/Qwen2.5-3B"
DATASET_NAME = "tatsu-lab/alpaca"
MAX_SEQ_LEN = 1024


def build_dataset(tokenizer):
    raw = load_dataset(DATASET_NAME, split="train[:8000]")

    def format_example(row):
        prompt = f"### Instruction:\n{row['instruction']}\n\n"
        if row.get("input"):
            prompt += f"### Input:\n{row['input']}\n\n"
        text = f"{prompt}### Response:\n{row['output']}{tokenizer.eos_token}"
        out = tokenizer(text, truncation=True, max_length=MAX_SEQ_LEN, padding="max_length")
        out["labels"] = out["input_ids"].copy()
        return out

    return raw.map(format_example, remove_columns=raw.column_names)


def train_func(config: dict):
    """Runs on every Ray Train worker (one per GPU)."""
    rank = ray.train.get_context().get_world_rank()

    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
    # prepare_model moves the model to this worker's GPU and wraps it in DDP.
    model = prepare_model(model)

    dataset = build_dataset(tokenizer).with_format("torch")
    loader = DataLoader(dataset, batch_size=config["batch_size"], shuffle=True, drop_last=True)
    # prepare_data_loader injects a DistributedSampler and moves batches to the GPU.
    loader = prepare_data_loader(loader)

    optimizer = torch.optim.AdamW(model.parameters(), lr=config["lr"])

    # 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": MODEL_NAME, "lr": config["lr"], "batch_size": config["batch_size"]})

    model.train()
    step = 0
    for batch in loader:
        out = model(
            input_ids=batch["input_ids"],
            attention_mask=batch["attention_mask"],
            labels=batch["labels"],
        )
        out.loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        optimizer.zero_grad()
        step += 1

        ray.train.report({"loss": out.loss.item(), "step": step})
        if use_mlflow:
            mlflow.log_metric("train_loss", out.loss.item(), step=step)
        if step >= config["max_steps"]:
            break

    if use_mlflow:
        mlflow.end_run()


def main():
    ray.init(address="auto")
    total_gpus = int(ray.cluster_resources().get("GPU", 0))
    print(f"Ray cluster ready: {total_gpus} GPU(s)", flush=True)

    trainer = TorchTrainer(
        train_func,
        train_loop_config={"lr": 2e-5, "batch_size": 4, "max_steps": 100},
        scaling_config=ScalingConfig(num_workers=total_gpus, use_gpu=True),
        run_config=RunConfig(storage_path="/tmp/ray_results", name="qwen-sft"),
    )
    result = trainer.fit()
    print(f"Training finished. Final metrics: {result.metrics}", flush=True)

    ray.shutdown()


if __name__ == "__main__":
    main()

Sumber daya tambahan

  • Last updated on