Tại sao chọn DeepSpeed ZeRO‑3 - ZeRO‑3 loại bỏ bộ nhớ trùng lặp bằng cách phân chia optimizer states, gradients và parameters trên tất cả GPU. - Hỗ trợ mô hình lên đến hàng trăm tỷ tham số mà không cần model parallelism phức tạp. - Tích hợp sẵn mixed precision, activation checkpointing và gradient accumulation. Chuẩn bị môi trường 1. Cài đặt driver CUDA 11.8+, cuDNN 8.6 và NCCL 2.18. 2. Tạo virtualenv và cài `torch==2.2`, `deepspeed==0.14`, `transformers==4.40`. 3. Cấu hình SSH không mật khẩu giữa các node, đồng bộ thư mục code qua NFS hoặc rsync. 4. Kiểm tra kết nối NCCL bằng `torchrun --nnodes=2 --nproc_per_node=4 -m torch.distributed.run --standalone test_nccl.py`. Cấu hình Zero (ds_config.json) ```json { "train_batch_size": 32, "gradient_accumulation_steps": 4, "fp16": {"enabled": true}, "zero_optimization": { "stage": 3, "offload_param": {"device": "cpu"}, "offload_optimizer": {"device": "cpu"}, "contiguous_gradients": true, "overlap_comm": true }, "activation_checkpointing": { "partition_activations": true, "cpu_checkpointing": true }, "wall_clock_breakdown": true } ``` Lưu ý: `offload_param/optimizer` về CPU giúp chạy mô hình 70B trên 8×A100 40GB. Quy trình huấn luyện 1. Viết script `train.py` sử dụng `deepspeed.initialize` và `model = AutoModelForCausalLM.from_pretrained(...)` . 2. Khởi chạy với `deepspeed --num_nodes=2 --num_gpus=4 train.py --deepspeed ds_config.json`. 3. Theo dõi memory usage qua `nvidia-smi dmon` và log DeepSpeed (`zero_stage3.log`). 4. Lưu checkpoint định kỳ (`save_interval=500`) và bật `keep_last_n_checkpoints=3`. Mẹo tối ưu thực tế - Bật `overlap_comm` và `contiguous_gradients` để che latency truyền thông. - Sử dụng `gradient_accumulation_steps` lớn khi batch size mỗi GPU nhỏ. - Nếu gặp OOM, tăng `offload_param` sang `nvme` (cần mount NVMe nhanh) hoặc giảm `train_batch_size`. - Áp dụng `activation_checkpointing` cho các layer transformer lớn (attention + feed‑forward). - Chạy benchmark nhỏ (1 epoch) trước khi scale lên toàn bộ dữ liệu để xác nhận throughput. Kết luận DeepSpeed ZeRO‑3 biến việc huấn luyện mô hình Transformer khổng lồ thành quy trình có thể lặp lại trên cụm GPU tiêu dùng. Kết hợp mixed precision, offload CPU/NVMe và activation checkpointing giúp giảm tiêu thụ VRAM tới 70% mà vẫn giữ tốc độ hội tụ. Hãy bắt đầu với cấu hình mẫu trên, điều chỉnh từng tham số theo phần cứng thực tế và theo dõi log để đạt hiệu suất tối ưu.