Cấu hình DeepSpeed ZeRO Stage 2 và Stage 3 để tối ưu bộ nhớ GPU
Để đào tạo các mô hình lớn vượt quá dung lượng VRAM của một GPU đơn lẻ, chúng ta cần kích hoạt ZeRO (Zero Redundancy Optimizer). ZeRO Stage 2 phân tán bộ tối ưu hóa, trong khi Stage 3 phân tán thêm cả tham số mô hình và gradient.
Bạn sẽ sửa file cấu hình training script để tích hợp ZeRO. Thay vì chạy script Python gốc, hãy tạo file cấu hình JSON riêng để DeepSpeed đọc, giúp tách biệt logic training và cấu hình phân tán.
Tạo file /opt/deepspeed-config/zero_stage2.json với nội dung sau để tối ưu cho Stage 2 (cân bằng tốt giữa bộ nhớ và tốc độ):
cat > /opt/deepspeed-config/zero_stage2.json
File JSON đã được tạo tại /opt/deepspeed-config/zero_stage2.json. Cấu hình này cho phép DeepSpeed tự động tính toán gradient_accumulation_steps dựa trên tổng batch size và micro batch size.
Để đạt mức tiết kiệm bộ nhớ tối đa cho các mô hình cực lớn (như Llama-3-70B), hãy tạo file /opt/deepspeed-config/zero_stage3.json với cấu hình Stage 3:
cat > /opt/deepspeed-config/zero_stage3.json
File Stage 3 đã sẵn sàng. Lưu ý các tham số offload_optimizer và offload_param đang sử dụng RAM CPU để giảm tải cho GPU, điều này có thể làm chậm tốc độ training nhưng cho phép chạy các mô hình khổng lồ.
Cách kiểm tra: Chạy lệnh nvidia-smi trên worker node trong quá trình training. Bạn sẽ thấy mức sử dụng VRAM giảm đáng kể so với việc chạy không có ZeRO, cho phép tăng batch size mà không bị lỗi OOM (Out Of Memory).
Tinh chỉnh số lượng Worker và Batch Size trong Ray
Ray phân phối các worker trên các node Kubernetes. Việc cân bằng giữa số lượng worker và batch size trên mỗi worker là chìa khóa để tận dụng tối đa băng thông mạng và tính toán GPU.
Bạn cần điều chỉnh số lượng worker trong file Ray Cluster Config (Kubernetes Manifest). Mục tiêu là đảm bảo tổng số GPU được sử dụng khớp với tài nguyên sẵn có và batch size tổng (Global Batch Size) được chia đều.
Sửa file /opt/ray-cluster-config/ray-trainer-cluster.yaml để đặt số worker cụ thể. Giả sử bạn có 4 GPU trên cluster và muốn chạy 4 worker (1 GPU/worker):
cat > /opt/ray-cluster-config/ray-trainer-cluster.yaml
Manifest Kubernetes đã được tạo với 4 worker, mỗi worker cấp phát 1 GPU. Cấu hình này đảm bảo Ray khởi tạo 4 tiến trình training song song.
Trong script Python training của bạn, hãy cấu hình PerWorkerBatchSize để phù hợp với tổng batch size mong muốn. Nếu bạn muốn Global Batch Size = 64 và có 4 worker, thì mỗi worker xử lý 16 mẫu:
import ray
from ray import train
from ray.train import ScalingConfig
# Cấu hình số lượng worker (phải khớp với số pod trong YAML)
scaling_config = ScalingConfig(num_workers=4, use_gpu=True)
# Batch size trên mỗi worker
per_worker_batch_size = 16
# Total batch size = 4 workers * 16 = 64
train_config = train.TrainConfig(
scaling_config=scaling_config,
run_config=train.RunConfig(name="optimization-run"),
)
# Trong hàm training_loop của bạn:
# train.report({"metrics": ...})
Script Python đã được cấu hình để Ray tự động phân phối 16 mẫu cho mỗi worker, tổng cộng 64 mẫu trên toàn cluster. Điều này giúp tận dụng tối đa băng thông GPU và tránh nghẽn cổ chai ở head node.
Cách kiểm tra: Truy cập Ray Dashboard (thường ở port 8265). Vào tab Cluster để xem số lượng worker đang chạy (Active) và tab Tasks để xem thời gian thực hiện mỗi batch. Nếu thời gian đồng đều giữa các worker, cân bằng tải đã thành công.
Sử dụng Kubernetes Resource Quota để quản lý tài nguyên
Trong môi trường Proxmox chia sẻ tài nguyên vật lý, việc một job training chiếm hết tài nguyên có thể gây sập cả cluster. Resource Quota giúp đặt giới hạn cứng (hard limit) cho namespace của training.
Bạn sẽ tạo một ResourceQuota cho namespace ai-training để giới hạn tổng số CPU, RAM và GPU mà các pod trong namespace này có thể tiêu thụ cùng lúc.
Tạo file /opt/k8s-quota/ai-training-quota.yaml để giới hạn tài nguyên:
cat > /opt/k8s-quota/ai-training-quota.yaml
File quota đã được tạo. Nó giới hạn namespace ai-training chỉ được sử dụng tối đa 16 CPU core, 64GB RAM, 4 GPU và 10 pod cùng lúc. Nếu vượt quá, Kubernetes sẽ từ chối khởi tạo pod mới.
Áp dụng quota vào cluster bằng lệnh kubectl:
kubectl apply -f /opt/k8s-quota/ai-training-quota.yaml
Lệnh kubectl apply trả về resourcequota "ai-training-quota" created. Cluster đã bắt đầu thực thi chính sách giới hạn tài nguyên cho namespace này.
Để kiểm tra tình trạng sử dụng tài nguyên hiện tại so với giới hạn, chạy lệnh:
kubectl describe quota ai-training-quota -n ai-training
Kết quả hiển thị bảng Name, Used và Hard. Nếu Used tiệm cận Hard, bạn sẽ thấy cảnh báo hoặc lỗi khi cố tạo thêm pod.
Để verify: Cố gắng tạo thêm một pod với yêu cầu GPU vượt quá giới hạn (ví dụ: 5 GPU khi quota chỉ cho 4). Kubernetes sẽ trả về lỗi exceeded quota, xác nhận cơ chế bảo vệ tài nguyên hoạt động chính xác.
Điều hướng series:
Mục lục: Series: Series: Xây dựng nền tảng AI Training phân tán với DeepSpeed, Ray và Kubernetes trên hạ tầng Proxmox
« Phần 5: Xây dựng và chạy thử nghiệm training mô hình phân tán
Phần 7: Quản lý, giám sát và troubleshooting hệ thống »