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

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

Các trainer trong TRL sử dụng 🤗 Accelerate để cho phép huấn luyện phân tán trên nhiều GPU hoặc node. Để làm điều này, trước tiên hãy tạo một tệp cấu hình 🤗 Accelerate bằng cách chạy

accelerate config

và trả lời các câu hỏi theo cấu hình đa GPU / đa node của bạn. Sau đó, bạn có thể khởi động huấn luyện phân tán bằng cách chạy:

accelerate launch train.py

Chúng tôi cũng cung cấp các tệp cấu hình trong thư mục examples có thể được dùng làm mẫu. Để sử dụng các mẫu này, chỉ cần truyền đường dẫn đến tệp cấu hình khi khởi chạy một tác vụ, ví dụ:

accelerate launch --config_file examples/accelerate_configs/multi_gpu.yaml train.py <SCRIPT_ARGS>

Điều này tự động phân phối khối lượng công việc trên tất cả các GPU khả dụng.

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, TRL có thể mở rộng huấn luyện trên nhiều máy (node) bằng cách sử dụng 🤗 Accelerate.

Cấu hình Accelerate

Tạo tệp cấu hình accelerate (ví dụ: multi_node.yaml) cho huấn luyện đa nút. Các trường chính:

compute_environment: LOCAL_MACHINE
distributed_type: MULTI_GPU
num_machines: 2
machine_rank: 0  # 0 for main node, 1 for second node
main_process_ip: 10.0.0.1  # IP of rank 0 node
main_process_port: 29500
num_processes: 16  # total processes across nodes
mixed_precision: bf16
use_cpu: false
same_network: true

Điều chỉnh num_processes cho khớp với tổng số GPU trên tất cả các node.

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

Ra mắt

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)
accelerate launch --config_file multi_node.yaml --machine_rank 0 train.py

# Node 1
accelerate launch --config_file multi_node.yaml --machine_rank 1 train.py

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

srun accelerate launch --config_file multi_node.yaml train.py

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

sbatch slurm_job.sh

SLURM tự động phân phối quá trình huấn luyện trên tất cả các node và GPU được yêu cầu, và srun cấu hình các biến môi trường cần thiết cho liên lạc đa node.

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 distributed_type: DEEPSPEED và thêm khối deepspeed_config. Xem hướng dẫn tích hợp DeepSpeed.

Đọc thêm

• Accelerate: Kịch bản khởi chạy
• Accelerate: Thư viện ví dụ
• Tài liệu SLURM Workload Manager - Để lên lịch công việc cụm