Trenowanie rozproszone za pomocą programu Ray Train

Important

Ta funkcja jest dostępna w wersji beta. Administratorzy obszaru roboczego mogą kontrolować dostęp do tej funkcji ze strony Podglądy . Zobacz Zarządzanie wersjami zapoznawczami usługi Azure Databricks.

W tym przykładzie uruchamiane jest rozproszone dostrajanie równoległe względem danych z użyciem Ray TrainTorchTrainer na 8 procesorach graficznych H100 w obrębie jednego węzła. Skrypt startowy uruchamia klaster Ray na węźle, a następnie sterownik Ray Train uruchamia po jednym procesie roboczym na każdy procesor GPU, opakowuje model w DDP i automatycznie dzieli zbiór danych pomiędzy procesy robocze.

Dostraja publicznie dostępny model (Qwen2.5-3B), dzięki czemu działa od razu bez tokenu Hugging Face.

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

  • Przesyła lokalny projekt za pomocą polecenia code_source: snapshot.
  • Uruchamia głowicę Ray ze wszystkimi 8 procesorami GPU, a następnie uruchamia sterownik Ray Train.
  • Używa ray.train.torch.prepare_model i prepare_data_loader do obsługi opakowywania w DDP, rozmieszczania na urządzeniach i próbkowania rozproszonego.
  • Rejestruje metryki w usłudze MLflow.

Prerequisites

Układ projektu

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

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

Krok 1. Zapisywanie obciążenia YAML

train.yaml żąda jednego GPU_8xH100 węzła. Zależności są deklarowane bezpośrednio pod environment (wraz ze środowiskiem uruchomieniowym version), a command uruchamia klaster Ray na węźle, po czym uruchamia driver, więc zadanie nie wymaga oddzielnego pliku z zależnościami ani skryptu uruchamiającego:

experiment_name: air-ray-train-distributed

environment:
  version: 'AI5'
  dependencies:
    # AI5 (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

Dyrektywa inline command uruchamia głowicę Ray ze wszystkimi GPU w węźle, uruchamia sterownik za pomocą python train_ray.py, a następnie zatrzymuje klaster. Zawiera również odgałęzienie robocze, które dołącza do węzła głównego, dzięki czemu to samo polecenie nadal działa po skalowaniu zadania do wielu węzłów.

Krok 2. Definiowanie sterownika Ray Train

train_ray.py definiuje train_func, który działa na każdym węźle roboczym, oraz main, który konfiguruje TorchTrainer do używania wszystkich procesorów GPU w klastrze. prepare_model opakowuje model za pomocą DDP i przenosi go na GPU procesu roboczego. prepare_data_loader dodaje rozproszony próbnik:

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()

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

Krok 3: Prześlij uruchomienie

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

Krok 4. Sprawdzanie przebiegu

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

Węzeł główny Ray i sterownik działają na węźle 0, więc logi są przesyłane strumieniowo z jednego węzła.

Gdzie wyniki lądują

Metryki raportowane za pomocą ray.train.report i rejestrowane w MLflow pojawiają się w eksperymencie MLflow o nazwie podanej w experiment_name, widocznym w interfejsie użytkownika MLflow obszaru roboczego.

Pełny skrypt szkoleniowy

Kompletny train_ray.py do kopiowania i wklejania:

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

Dodatkowe zasoby

  • Last updated on