Распределение обучения

Обучение на нескольких GPU с TRL

Тренеры в TRL запускаются с помощью torchrun — распределённого лаунчера PyTorch. Он запускает один процесс на каждый GPU:

torchrun --nproc_per_node 8 train.py <SCRIPT_ARGS>

CLI trl делает то же самое под капотом, поэтому trl sft ... запускается на всех GPU машины; передайте --nproc_per_node, чтобы использовать меньше.

Под капотом 🤗 Accelerate создаёт по одной модели на каждый GPU. Каждый процесс:

• Обрабатывает собственную партию данных
• Вычисляет потерю и градиенты для этой партии
• Обмен обновлениями градиентов между всеми GPU

Эффективный размер батча рассчитывается как: $$\text{Batch Size} = \text{per\_device\_train\_batch\_size} \times \text{num\_devices} \times \text{gradient\_accumulation\_steps}$$

Чтобы поддерживать постоянный размер батча при масштабировании на несколько GPU, убедитесь, что соответствующим образом обновлены параметры per_device_train_batch_size и gradient_accumulation_steps.

Например, эти конфигурации эквивалентны и должны давать одинаковые результаты:

Использование одной модели на каждый GPU может приводить к высокому потреблению памяти, что не всегда допустимо для больших моделей или GPU с малым объёмом памяти. В таких случаях можно использовать DeepSpeed, который предлагает оптимизации, такие как шардирование модели, Zero Redundancy Optimizer, обучение с смешанной точностью и выгрузка на CPU или NVMe. Подробнее см. в нашем руководстве по DeepSpeed Integration.

Обучение на очень длинных последовательностях описано в отдельном руководстве: Обучение на последовательностях длиннее 1 млн токенов.

Обучение на нескольких узлах

Когда на одном компьютере недостаточно GPU, torchrun позволяет масштабировать обучение на несколько машин (узлов). На каждом узле выполняется одна и та же команда с указанием общего числа узлов, собственного ранга и точки синхронизации (rendezvous) на главном узле.

Вариант 1: Ручной запуск (не HPC)

Вручную выполните следующее на каждом узле:

# Node 0 (main node, IP 10.0.0.1)
torchrun --nnodes 2 --nproc_per_node 8 --node_rank 0 --rdzv_backend c10d --rdzv_endpoint 10.0.0.1:29500 train.py

# Node 1
torchrun --nnodes 2 --nproc_per_node 8 --node_rank 1 --rdzv_backend c10d --rdzv_endpoint 10.0.0.1:29500 train.py

Замените 10.0.0.1 на фактический IP-адрес узла с рангом 0 (основного узла).

Вариант 2: запуск через SLURM (кластеры HPC)

Для кластеров, использующих планировщик заданий SLURM, создайте скрипт задания (например, slurm_job.sh):

#!/bin/bash
#SBATCH --nodes=2
#SBATCH --gpus-per-node=8
#SBATCH --job-name=trl_multi

MAIN_NODE=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
srun torchrun --nnodes 2 --nproc_per_node 8 --node_rank $SLURM_NODEID --rdzv_backend c10d --rdzv_endpoint $MAIN_NODE:29500 train.py

Затем отправьте задание:

sbatch slurm_job.sh

srun запускает по одному torchrun на узел, а каждый из них запускает по одному процессу на GPU.

Ключевые директивы SLURM:

• --nodes=2: запросить 2 вычислительных узла
• --gpus-per-node=8: выделяет 8 GPU на узел (16 всего)
• --job-name: метка для отслеживания в очереди задач

Можно объединить многоузловый режим с DeepSpeed, задав deepspeed в конфигурации обучения. См. руководство по интеграции DeepSpeed.

Дополнительные материалы

• torchrun (элапстичный запуск)
• Документация SLURM Workload Manager — для планирования задач на кластере