AWS phát hành Ray Train DLC để đơn giản hóa việc huấn luyện AI phân tán trên EKS

Amazon Web Services (AWS) vừa ra mắt Ray Train Deep Learning Container (DLC), một hình ảnh (image) được cấu hình sẵn nhằm tối ưu hóa quy trình huấn luyện máy học phân tán trên Amazon Elastic Kubernetes Service (EKS). Được công bố vào tháng 10 năm 2026, container này tích hợp sẵn Ray, PyTorch và các thư viện mạng cần thiết cho tính toán hiệu năng cao, cho phép kỹ sư triển khai các tác vụ huấn luyện đa nút (multi-node) mà không phải tự quản lý chuỗi phụ thuộc phức tạp.

Chuyện gì đã xảy ra

Nút thắt chính khi mở rộng quy mô huấn luyện PyTorch trên các cụm thường nằm ở hạ tầng container hơn là bản thân mã mô hình. Các kỹ sư thường dành nhiều thời gian để đồng bộ hóa runtime CUDA, các bản build PyTorch, NVIDIA Collective Communications Library (NCCL) và các thư viện người dùng của Elastic Fabric Adapter (EFA). Sự không khớp giữa các thành phần này có thể gây ra suy giảm hiệu năng ngầm định, trong đó NCCL chuyển sang sử dụng socket TCP chậm thay vì kết nối EFA tốc độ cao. Để giải quyết vấn đề này, AWS đã xây dựng, kiểm thử và công bố một hình ảnh thống nhất trên Amazon Elastic Container Registry (ECR) Public.

DLC mới này do AWS xây dựng và bảo trì, bao gồm việc vá lỗi stack và xác thực thông qua các bài kiểm tra huấn luyện KubeRay với EFA đa nút. Phiên bản phát hành bao gồm Ray 2.58.0, PyTorch 2.13.0, CUDA 13.0.2 và NCCL 2.29.7, cùng với plugin AWS OFI NCCL giúp kết nối NCCL với EFA. Bằng cách cung cấp dưới dạng hình ảnh công khai, AWS loại bỏ nhu cầu các đội ngũ phải tự xây dựng và bảo trì container riêng cho mỗi lần cập nhật các thư viện nền tảng này.

Hướng dẫn kỹ thuật đi kèm minh họa cách sử dụng DLC này để tinh chỉnh (fine-tune) mô hình Qwen2.5-1.5B bằng PyTorch Fully Sharded Data Parallel (FSDP) trên tám GPU NVIDIA L4. Thiết lập này sử dụng hai instance g6.12xlarge trong một Availability Zone duy nhất, được quản lý bởi operator KubeRay. Cấu hình này cho phép các thao tác thu thập dữ liệu liên nút (inter-node collectives) chạy ngay lập tức qua EFA, vì DLC đã chứa sẵn stack mạng cần thiết.

Cách thức hoạt động

Ray Train DLC hoạt động như một hình ảnh duy nhất, độc lập với vai trò, phục vụ cả pod head và worker của Ray. Trên Kubernetes, operator KubeRay sử dụng manifest RayCluster để tạo các pod cần thiết, chỉ định worker kết nối tới node head để phối hợp. Vì hình ảnh chứa tất cả driver và thư viện cần thiết, người dùng không cần xây dựng image tùy chỉnh. Thay vào đó, họ cấu hình cụm thông qua biến môi trường và yêu cầu tài nguyên trong manifest Kubernetes.

Việc xử lý mạng được tự động hóa thông qua các biến môi trường đã được cấu hình sẵn. DLC mặc định sử dụng EFA cho giao tiếp qua cài đặt FI_PROVIDER=efa và loại trừ các giao diện loopback khỏi cấu hình NCCL. Nếu một tác vụ yêu cầu các gói Python bổ sung không có trong base image, tính năng runtime_env của Ray sẽ cài đặt chúng tại thời điểm gửi job, giữ nguyên base image. Sự tách biệt này đảm bảo container ổn định trong khi vẫn linh hoạt cho các nhu cầu thử nghiệm cụ thể.

Bên trong container, Ray Train quản lý việc hình thành nhóm tiến trình phân tán. Kỹ sư viết vòng lặp huấn luyện dưới dạng hàm Python tiêu chuẩn, và Ray sẽ thực thi nó trên tất cả GPU khả dụng. Đối với khối lượng công việc FSDP, Ray chia nhỏ trọng số mô hình, gradient và trạng thái optimizer trên các GPU, giúp giảm yêu cầu bộ nhớ trên mỗi thiết bị. DLC đảm bảo rằng lưu lượng giao tiếp nặng nề sinh ra từ các thao tác gather và reduce-scatter của FSDP truyền qua giao diện băng thông cao EFA thay vì socket mạng tiêu chuẩn.

Chi tiết quan trọng

• Image Tag: Mẫu ví dụ sử dụng public.ecr.aws/deep-learning-containers/ray:train-ml-cuda-v1.1, bao gồm Ray 2.58.0 và PyTorch 2.13.0.
• Hỗ trợ phần cứng: Kiến trúc tham chiếu sử dụng instance g6.12xlarge với GPU NVIDIA L4 và giao diện EFA, triển khai trong một Availability Zone duy nhất.
• Stack mạng: Bao gồm CUDA 13.0.2, NCCL 2.29.7 và plugin AWS OFI NCCL, đã được kiểm thử để đảm bảo sử dụng EFA thay vì fallback về TCP.
• Phương pháp cấu hình: Hình dạng cụm và cài đặt truyền tải được xác định qua biến môi trường và yêu cầu tài nguyên Kubernetes, không phải bằng cách rebuild image.
• Ước tính chi phí: Thiết lập mẫu tốn khoảng $9.50 mỗi giờ, bao gồm hai instance GPU, control plane EKS và NAT Gateway.
• Tương thích Driver: Entrypoint của container tự động bật các thư viện tương thích ngược (forward-compatibility) của CUDA nếu driver NVIDIA trên host cũ hơn phiên bản CUDA trong image.

Tại sao điều này quan trọng

Đối với kỹ sư phần mềm và đội ngũ ML, bản phát hành này giảm đáng kể gánh nặng vận hành của việc huấn luyện phân tán. Quản lý các phiên bản tương thích của CUDA, NCCL và driver EFA là một quy trình mong manh, thường dẫn đến các vấn đề hiệu năng khó nhận biết. Bằng cách cung cấp một hình ảnh công khai đã được kiểm thử, AWS cho phép các đội ngũ tập trung vào phát triển mô hình thay vì gỡ lỗi hạ tầng. Khả năng ghim (pin) các phiên bản cụ thể hoặc sử dụng tag di động cho các bản cập nhật bảo mật mang lại sự linh hoạt cho cả nghiên cứu tái lập được lẫn sự ổn định trong sản xuất.

Việc tích hợp với KubeRay và EKS cũng đơn giản hóa quản lý cụm. Kỹ sư có thể xác định toàn bộ hạ tầng huấn luyện trong một manifest duy nhất, chỉ định số lượng GPU và giao diện EFA như các tài nguyên Kubernetes tiêu chuẩn. Cách tiếp cận này tận dụng các công cụ Kubernetes hiện có và tránh nhu cầu về script tùy chỉnh để quản lý vòng đời cụm Ray. Kết quả là một lộ trình đáng tin cậy và có khả năng mở rộng tốt hơn để chạy các job tinh chỉnh mô hình ngôn ngữ lớn trên hạ tầng AWS.

Bạn có thể làm gì

• Xem xét repository mẫu: Truy cập mã nguồn và manifest được cung cấp để hiểu cấu trúc của một job Ray Train trên EKS.
• Kiểm tra hạn mức (quota): Đảm bảo tài khoản AWS của bạn có đủ dung lượng On-Demand cho các instance dòng G, cụ thể là 96 vCPU cho hai node g6.12xlarge.
• Cài đặt công cụ cần thiết: Thiết lập AWS CLI v2, eksctl, kubectl, Helm và envsubst để quản lý cụm và gửi job.
• Thử nghiệm DLC cục bộ: Pull public image từ ECR để xác minh tính tương thích với môi trường phát triển cục bộ trước khi triển khai lên cụm.
• Cấu hình biến môi trường: Sử dụng scripts/env.sh được cung cấp để thiết lập các tham số cụm như region, loại instance và image tag.
• Giám sát hiệu năng mạng: Bật NCCL_DEBUG=INFO trong manifest của bạn để xác nhận rằng EFA đang được sử dụng cho giao tiếp liên nút trong quá trình huấn luyện.