Phân phối huấn luyện

Đào tạo đa GPU với TRL

Các trainer trong TRL được khởi chạy bằng torchrun, trình khởi chạy phân tán của PyTorch. Nó khởi động một tiến trình cho mỗi GPU:

torchrun --nproc_per_node 8 train.py <SCRIPT_ARGS>

CLI trl thực hiện cùng một điều ở bên trong, vì vậy trl sft ... chạy trên mọi GPU của máy; truyền --nproc_per_node để sử dụng ít hơn.

Bên trong, 🤗 Accelerate tạo một mô hình cho mỗi GPU. Mỗi tiến trình:

• Xử lý lô dữ liệu của riêng nó
• Tính toán loss và gradients cho batch đó
• Chia sẻ các bản cập nhật gradient trên tất cả các GPU

Kích thước lô hiệu dụng được tính như sau: $$\text{Batch Size} = \text{per\_device\_train\_batch\_size} \times \text{num\_devices} \times \text{gradient\_accumulation\_steps}$$

Để duy trì kích thước batch nhất quán khi mở rộng lên nhiều GPU, hãy đảm bảo cập nhật per_device_train_batch_size và gradient_accumulation_steps cho phù hợp.

Ví dụ, các cấu hình sau tương đương và nên cho ra kết quả giống nhau:

Việc gán một mô hình cho mỗi GPU có thể dẫn đến mức sử dụng bộ nhớ cao, điều này có thể không khả thi đối với các mô hình lớn hoặc GPU có bộ nhớ thấp. Trong những trường hợp như vậy, bạn có thể tận dụng DeepSpeed, cung cấp các tối ưu hóa như chia nhỏ mô hình (model sharding), Zero Redundancy Optimizer, huấn luyện độ chính xác hỗn hợp và offload sang CPU hoặc NVMe. Xem hướng dẫn DeepSpeed Integration của chúng tôi để biết thêm chi tiết.

Huấn luyện trên các chuỗi rất dài có hướng dẫn riêng: Huấn luyện vượt quá 1 triệu Token.

Đào tạo đa nút

Khi một máy không có đủ GPU, torchrun có thể mở rộng việc huấn luyện trên nhiều máy (node). Mỗi node chạy cùng một lệnh với tổng số node, rank của chính nó và một điểm rendezvous trên node chính.

Tùy chọn 1: Khởi động thủ công (Không HPC)

Chạy các lệnh sau trên từng nút một cách thủ công:

# 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

Thay 10.0.0.1 bằng địa chỉ IP thực tế của nút rank 0 (chính).

Tùy chọn 2: Khởi động SLURM (Các cụm HPC)

Đối với các cụm sử dụng trình lên lịch công việc SLURM, hãy tạo một kịch bản công việc (ví dụ: 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

Sau đó, gửi tác vụ:

sbatch slurm_job.sh

srun khởi động một torchrun cho mỗi node, và mỗi torchrun khởi động một tiến trình cho mỗi GPU.

Các chỉ thị SLURM chính:

• --nodes=2: Yêu cầu 2 nút tính toán
• --gpus-per-node=8: Phân bổ 8 GPU cho mỗi node (tổng cộng 16)
• --job-name: Nhãn để theo dõi trong hàng đợi công việc

Bạn có thể kết hợp multi-node với DeepSpeed bằng cách đặt deepspeed trong cấu hình huấn luyện. Xem hướng dẫn tích hợp DeepSpeed.

Đọc thêm

• torchrun (Khởi động co giãn)
• Tài liệu SLURM Workload Manager - Để lên lịch công việc cụm