PyTorch 2.15 sửa lỗi "một GPU hỏng" khi huấn luyện AI

PyTorch 2.15 đã chính thức ra mắt với một thay đổi đáng chú ý: backend cho thư viện truyền thông tập hợp NCCL của Nvidia được viết lại hoàn toàn. Mục tiêu là để một GPU hỏng không còn khiến toàn bộ tác vụ huấn luyện phải chạy lại từ đầu.

Đây là bước tiến nữa trong việc đưa PyTorch từ công cụ dành cho nhà nghiên cứu trở thành trụ cột của hoạt động AI trong môi trường sản xuất.

Vì sao phải viết lại backend NCCL

NCCL (Nvidia Collective Communications Library, đọc là "Nickel") cung cấp cho PyTorch bộ nguyên thủy truyền thông giữa các GPU, cho phép chúng trao đổi dữ liệu trong cùng một node hoặc giữa nhiều node trong cluster. Đây là thành phần thiết yếu khi chạy suy luận trên các mạng nơ-ron lớn và khi xây dựng, vận hành mô hình ngôn ngữ lớn (LLM).

Theo Tristan Rice, kỹ sư phần mềm của Meta chuyên về tính toán phân tán và là nhà phát triển cốt lõi của PyTorch, giao diện cũ với thư viện mạng này đã "chìm trong nợ kỹ thuật" và buộc phải viết lại. Ông chia sẻ điều này trong buổi livestream trên YouTube hôm thứ Năm.

Backend cũ, ProcessGroupNCCL, có thể sinh ra hàng chục communicator cho mỗi process group, dẫn tới tình trạng deadlock (khóa chết) chập chờn và hiệu năng khó lường.

Quy tắc mới và khả năng phục hồi lỗi

Phiên bản mới áp dụng một quy tắc nghiêm ngặt: mỗi process group chỉ có duy nhất một communicator trong giai đoạn khởi tạo. Mặc định, điều này loại bỏ nhiễu lẫn giữa các communicator và mọi phức tạp kéo theo. Nếu chương trình cần song song hóa, nó phải khai báo thêm process group.

Thư viện mới cũng xử lý tốt hơn các GPU chậm (straggler) nhờ khả năng cấu hình lại dễ dàng hơn. Khi một GPU gặp sự cố hoặc chạy chậm, một bộ điều phối bên thứ ba có thể loại bỏ GPU đó và thậm chí gán worker thay thế.

Ngoài ra còn có thêm "flight recorder" (hộp đen ghi nhật ký) giúp quản trị viên hiểu rõ hơn chuyện gì đã xảy ra. Trước đây, khi một cluster GPU bị treo — có thể do timeout hoặc deadlock — người dùng chỉ nhận được thông báo chung chung. Flight recorder mới ghi lại toàn bộ thao tác dưới dạng chuỗi ảnh chụp (snapshot) để có thể kiểm tra sau khi lỗi xảy ra.

Rice cho biết backend NCCL đã được dọn dẹp, vẫn mang tên ProcessGroupNCCL, sẽ cho nhà phát triển một "trang giấy trắng", giúp việc thêm tính năng mới dễ dàng hơn. ProcessGroupNCCL phiên bản cũ vẫn nằm trong gói 2.15 để có thể quay lui, nhưng nếu backend mới hoạt động đúng như kỳ vọng, phiên bản cũ sẽ bị loại bỏ trong bản 2.16 dự kiến ra mắt đầu năm sau.

Các cập nhật đáng chú ý khác

PyTorch 2.15 còn có một số thay đổi đáng chú ý:

• Python transport API mới cho phép truyền dữ liệu một chiều bất đồng bộ, lý tưởng cho các tác vụ huấn luyện cần đồng bộ trọng số giữa trainer và generator.
• Cải thiện xử lý song song trên thiết bị của Apple.
• Hỗ trợ GPU Blackwell thế hệ tiếp theo của Nvidia.

Tổng cộng, bản cập nhật gồm 3.117 commit từ 583 người đóng góp.

"Mọi phòng lab ML lớn (OpenAI, DeepMind, Meta, Nvidia) đều dùng PyTorch. Nó vận hành từ mô hình diffusion, LLM cho tới hệ thống gợi ý và robot. Việc được sử dụng rộng rãi như vậy không phải ngẫu nhiên", Will Stewart, đồng sáng lập nhà cung cấp phần mềm AI Northflank, nhận định trong một bài blog.

Góc nhìn

Với các kỹ sư và doanh nghiệp đang vận hành cluster GPU để huấn luyện mô hình, thay đổi này chạm đúng vào nỗi đau lớn nhất: chi phí của một lần hỏng hóc. Khi một job huấn luyện kéo dài hàng giờ hoặc hàng ngày bị treo vì một GPU lỗi, việc phải khởi động lại từ đầu không chỉ tốn thời gian mà còn tốn tiền thuê máy. Flight recorder và khả năng loại bỏ GPU chậm giúp việc gỡ lỗi có bằng chứng cụ thể thay vì đoán mò.

Việc siết quy tắc một communicator cho mỗi process group cũng là lời nhắc rằng các dự án đang phụ thuộc vào cách cấu hình NCCL cũ nên kiểm tra lại trước khi nâng cấp. Backend cũ vẫn có mặt trong 2.15 để quay lui, nhưng sẽ biến mất ở 2.16 — đây là thời điểm hợp lý để rà soát mã và chuyển sang cách khai báo process group mới.