Elosztott betanítás a Ray Train használatával

Important

Ez a funkció nyilvános előzetes verzióban van.

Ez a példa elosztott adat-párhuzamos finomhangolást futtat a Ray TrainTorchTrainer 8 H100 GPU-jával egyetlen csomóponton. Egy bootstrap-szkript elindít egy Ray-fürtöt a csomóponton, majd a Ray Train-illesztő GPU-nként egy feldolgozót indít el, a modellt DDP-ben burkolja, és automatikusan szétszedi az adathalmazt a feldolgozók között.

Egy nyilvános modellt (Qwen2.5-3B) finomhangol, így az változtatás nélkül, Hugging Face-token nélkül is futtatható.

A számítási feladat a következőket végzi el:

Feltölti a helyi projektet a következővel code_source: snapshot: .
Elindít egy Ray-fejet mind a 8 GPU-val, majd futtatja a Ray Train illesztőprogramot.
A ray.train.torch.prepare_model és a prepare_data_loader használatával kezeli a DDP-beágyazást, az eszközök hozzárendelését és az elosztott mintavételezést.
Metrikákat naplóz az MLflow-ba.

Prerequisites

A parancssori air felület telepítve és hitelesítve. Lásd : Az AI-futtatókörnyezeti parancssori felület telepítése.

Projektelrendezés

Hozzon létre egy könyvtárat a következő fájlokkal.

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

1. lépés: A számítási feladat YAML-jének írása

train.yaml egyetlen GPU_8xH100 csomópontot kér. A függőségek soron belül, a environment alatt vannak deklarálva (a futtatókörnyezettel együtt: version), és a command elindít egy Ray-fürtöt a csomóponton, majd futtatja az illesztőt, így a munkaterhelésnek nincs szüksége külön függőségfájlra vagy indítószkriptre:

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

Az inline command elindít egy Ray headet a csomóponton található összes GPU-val, lefuttatja az illesztőprogramot a(z) python train_ray.py használatával, majd leállítja a fürtöt. Emellett tartalmaz egy feldolgozó ágat is, amely a fejhez csatlakozik, így ugyanaz a parancs működik, ha a feladatot több csomópontra skálázza.

2. lépés: A Ray Train-illesztőprogram meghatározása

train_ray.py meghatároz egy train_func elemet, amely minden feldolgozón fut, valamint egy main elemet, amely a TorchTrainer elemet úgy konfigurálja, hogy a fürt összes GPU-ját használja. prepare_model DDP-be csomagolja a modellt, és áthelyezi a worker GPU-jára. prepare_data_loader elosztott mintavevőt ad hozzá:

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

A teljes szkript a teljes betanítási szkriptben jelenik meg a lap végén.

3. lépés: A futtatás beküldése

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

4. lépés: A futtatás vizsgálata

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

A Ray-fej és az illesztőprogram is a 0. csomóponton fut, ezért egyetlen csomópontról naplózza a streamet.

Ahol az eredmények megjelennek

A ray.train.report használatával jelentett és az MLflow-val naplózott metrikák a experiment_name alatt megnevezett MLflow-kísérletben jelennek meg, amely a munkaterület MLflow felületén tekinthető meg.

Teljes betanítási szkript

A teljes train_ray.py másolás-beillesztéshez:

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

További erőforrások

Visszajelzés

Hasznosnak találta ezt az oldalt?

Last updated on 2026-07-24