Распределенное обучение с помощью Ray Train

Important

Эта функция доступна в общедоступной предварительной версии.

В этом примере выполняется распределенное дообучение с параллелизмом данных с использованием Ray TrainTorchTrainer на 8 графических процессорах H100 в рамках одного узла. Скрипт начальной загрузки запускает кластер Ray на узле, затем драйвер Ray Train запускает одну рабочую роль на GPU, упаковывает модель в DDP и сегментирует набор данных между рабочими ролей автоматически.

Она дообучает общедоступную модель (Qwen2.5-3B), поэтому её можно запускать как есть без токена Hugging Face.

Рабочая нагрузка выполняет следующие действия:

Загружает локальный проект с помощью code_source: snapshot.
Запускает головной узел Ray со всеми 8 графическими процессорами, затем запускает драйвер Ray Train.
Использует ray.train.torch.prepare_model и prepare_data_loader обрабатывает оболочку DDP, размещение устройств и распределенную выборку.
Записывает метрики в MLflow.

Необходимые условия

Интерфейс командной строки air установлен и прошел аутентификацию. См. раздел "Установка интерфейса командной строки среды выполнения ИИ".

Структура проекта

Создайте каталог со следующими файлами.

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

Шаг 1. Напишите YAML-файл рабочей нагрузки

train.yaml запрашивает один GPU_8xH100 узел. Зависимости объявляются непосредственно в environment (с использованием среды выполнения version), а command запускает на узле кластер Ray, а затем драйвер, поэтому рабочей нагрузке не требуется отдельный файл зависимостей или скрипт запуска:

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

Встроенная команда command запускает головной узел Ray со всеми GPU на узле, затем запускает драйвер с помощью python train_ray.py, а после этого останавливает кластер. Он также включает ветку worker, которая подключается к head-узлу, поэтому та же команда продолжит работать при масштабировании задания на несколько узлов.

Шаг 2. Определите драйвер Ray Train

train_ray.py определяет train_func, который запускается на каждом рабочем узле, и main, который настраивает TorchTrainer на использование всех GPU в кластере. prepare_model оборачивает модель в DDP и перемещает её на GPU рабочего процесса. prepare_data_loader добавляет распределённый семплер:

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

Полный скрипт указан в полнофункционалном скрипте обучения в конце этой страницы.

Шаг 3: Отправьте запуск

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

Шаг 4. Проверка выполнения

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

Головной узел Ray и драйвер оба работают на узле 0, поэтому журналы передаются с одного узла.

Куда попадают результаты

Метрики, передаваемые с помощью ray.train.report и регистрируемые с помощью MLflow, отображаются в эксперименте MLflow, имя которого указано в experiment_name; его можно просмотреть в интерфейсе MLflow рабочей области.

Полный сценарий обучения

Полный текст train_ray.py для копирования и вставки:

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

Дополнительные ресурсы

Обратная связь

Были ли сведения на этой странице полезными?

Last updated on 2026-07-24