Tối ưu hóa huấn luyện PyTorch với khả năng chịu lỗi trên Databricks AI Runtime

Ở quy mô lớn, hiệu suất huấn luyện của bạn được quyết định bởi một chỉ số duy nhất: "goodput" (hiệu suất thực tế), tức là tỷ lệ thời gian GPU của bạn dành cho việc tính toán hiệu quả thay vì chờ đợi hoặc khôi phục sau lỗi. Vì lỗi GPU là điều khó tránh khỏi ở quy mô lớn, khả năng tự động khôi phục nhanh chóng sau sự cố là cách duy nhất để duy trì goodput cao và quản lý tổng chi phí GPU của bạn.

Hai hệ thống con quyết định sự thành bại của quá trình khôi phục đó, nhưng cả hai thường bị xem nhẹ: đường ống dữ liệu (data pipeline) cung cấp dữ liệu cho các bộ tăng tốc (accelerators) và cơ chế checkpointing (lưu điểm kiểm tra) giúp ghi lại trạng thái để công việc có thể tiếp tục. Nếu làm sai một trong hai, mỗi sự cố sẽ khiến bạn lãng phí thời gian chờ của GPU nhiều hơn mức cần thiết. Ngay cả khi không có sự cố, một đường ống dữ liệu không theo kịp các bộ tăng tốc sẽ âm thầm làm "đói" GPU và bào mòn goodput chắc chắn như khi xảy ra lỗi hệ thống. Chúng tôi sẽ đi sâu vào các cơ chế và sự đánh đổi của cả hai, cũng như cách mỗi yếu tố định hình goodput và tổng chi phí GPU của bạn. Hãy xem hướng dẫn về hiệu suất và khả năng phục hồi (Training performance and resiliency guide) để biết các ví dụ và mã nguồn.

Đối với khía cạnh hạ tầng của cùng vấn đề này, về cách một cụm máy chủ phát hiện và cô lập các GPU không ổn định trước khi chúng làm gián đoạn công việc, hãy xem bài viết đi kèm: Cách chúng tôi duy trì độ tin cậy của GPU trên Databricks AI.

Tại sao lỗi là điều khó tránh khỏi ở quy mô lớn

Khi số lượng GPU trong một tác vụ tăng lên, xác suất để tác vụ đó hoàn thành mà không bị gián đoạn sẽ giảm nhanh chóng. Một mô hình ước tính nhanh từ bài viết đi kèm của Databricks giả định mỗi GPU có tỷ lệ lỗi hàng năm khoảng 1%. Dựa trên giả định đó, bài viết lưu ý rằng "một tác vụ chạy trên 256 GPU trong 30 ngày có khoảng 19% khả năng gặp lỗi. Với 1.024 GPU, con số đó tăng lên 57%." và đây chỉ mới là các vấn đề ở cấp độ hạ tầng.

Để đưa ước tính đó vào thực tế, siêu máy tính delta với 608 GPU H100 đã gặp lỗi mỗi 1,9 giờ, nghĩa là đối với một tác vụ 32 GPU, thời gian trung bình giữa các lần lỗi (MTBF) sẽ là 36 giờ. Bài học chính ở đây là tác vụ huấn luyện của bạn có khả năng sẽ thất bại tại một thời điểm nào đó, và việc đưa ra các quyết định đúng đắn có thể giúp mô hình của bạn có khả năng phục hồi tốt hơn và giảm tổng thời gian lãng phí khi sự cố xảy ra.

Tác động 1: Định dạng checkpoint quyết định tần suất bạn có thể lưu

Checkpointing là nơi quyết định sự thành bại của khả năng phục hồi, và cơ chế bạn chọn có ảnh hưởng trực tiếp đến tần suất bạn có thể lưu. Đây là đòn bẩy lớn nhất đối với goodput của bạn: nếu bạn checkpoint mỗi ngày một lần, thì khi xảy ra lỗi, bạn sẽ phải chạy lại trung bình 12 giờ công việc trùng lặp để đưa hệ thống trở lại trạng thái trước khi lỗi xảy ra.

Nút thắt cổ chai của torch.save nguyên khối (monolithic)

Checkpoint đầu tiên mà hầu hết các nhóm viết là một lệnh torch.save đơn giản trên rank 0. Tùy thuộc vào cách mô hình của bạn được huấn luyện, có khả năng xảy ra hai vấn đề:

Hành vi chặn (blocking) này khiến GPU của bạn nhàn rỗi, làm giảm goodput. Nhưng có một cách để giảm thời gian GPU dành cho việc checkpoint: API checkpoint phân tán (distributed checkpoint API) của Torch.

Checkpoint phân tán (DCP): mỗi rank ghi shard (phân đoạn) riêng của mình

Checkpoint phân tán của PyTorch đảo ngược thiết kế truyền thống. Mỗi rank ghi shard riêng biệt của nó song song với nhau, cùng với một tệp.metadata nhỏ mô tả cách các shard này hợp thành các tensor hoàn chỉnh.

Thời gian lưu giảm dần theo tỷ lệ xấp xỉ 1/N với số lượng rank và vì tệp.metadata ghi lại bố cục toàn cục, cùng một checkpoint đó có thể được tải lại trên một số lượng GPU khác nhau. DCP lập kế hoạch lại các byte mà mỗi rank mới cần, vì vậy việc khôi phục trên một cụm có dung lượng giảm sau khi mất node vẫn hoạt động bình thường.

DCP đáng giá ngay cả đối với các tác vụ song song dữ liệu (data-parallel) thông thường

Một giả định phổ biến là DCP chỉ dành cho các mô hình được phân đoạn (sharded), rằng một tác vụ song song dữ liệu (DDP), nơi mỗi rank giữ một bản sao giống hệt nhau của các trọng số, sẽ không nhận được lợi ích gì. Không phải vậy, DCP phân đoạn trạng thái mô hình và ghi nó song song trên mỗi worker ngay cả đối với các tác vụ huấn luyện DDP.

Đây cũng là API tương tự mà bạn sẽ cần vào ngày bạn chuyển sang FSDP hoặc song song tensor, vì vậy việc áp dụng sớm có nghĩa là bạn không bao giờ phải viết lại mã phục hồi vào thời điểm tồi tệ nhất.

Lưu không đồng bộ (asynchronous) giúp tần suất lưu gần như miễn phí

Ngay cả với việc ghi song song, một lệnh lưu đồng bộ vẫn chặn quá trình huấn luyện cho đến khi các byte được ghi bền vững vào bộ nhớ; đối với một checkpoint lớn gửi đến ổ đĩa từ xa, điều này tốn hàng chục giây thời gian nhàn rỗi của bộ tăng tốc. async_save chia nhỏ thao tác: sao chép nhanh vào bộ đệm tạm (staging buffer), sau đó tải lên nền trong khi quá trình huấn luyện vẫn tiếp tục.

Vòng lặp huấn luyện chỉ tốn thời gian cho việc sao chép tạm, không phải việc tải lên. Một checkpoint trước đây tốn hàng chục giây thời gian nhàn rỗi giờ đây gần như không tốn gì, đây chính xác là điều giúp việc checkpoint thường xuyên trong phần tiếp theo trở nên khả thi.

Trên AI Runtime, UCVolumeWriter và UCVolumeReader triển khai DCP dựa trên các UC volume, thực hiện I/O tạm thông qua NVMe cục bộ và chỉ đánh dấu checkpoint hoàn tất khi dữ liệu của nó đã được ghi hoàn toàn. Xem hướng dẫn về hiệu suất và khả năng phục hồi để biết chi tiết đầy đủ và ví dụ về mã nguồn.

Những điều trên không bao gồm thời gian lưu trữ mạng cho torch.save.

Tác động 2: Tần suất checkpoint quyết định chi phí khôi phục của bạn

Đây là nơi các yếu tố kết hợp với nhau. Khi một tác vụ thất bại, nó mất tất cả dữ liệu kể từ checkpoint hợp lệ cuối cùng và phải tính toán lại. Vì vậy, công việc lãng phí dự kiến cho mỗi lần lỗi bằng khoảng một nửa khoảng thời gian giữa các lần checkpoint, và việc lưu không đồng bộ giá rẻ cho phép bạn rút ngắn khoảng thời gian đó.

Giảm khoảng thời gian này xuống 10 lần sẽ giảm thời gian dự kiến để khôi phục xuống 10 lần. Hãy nhớ lại con số ~8,6 lần gián đoạn mỗi ngày của Llama 3: với tỷ lệ lỗi đó, việc checkpoint mỗi 2 giờ có nghĩa là bạn dự kiến lãng phí 8,6 giờ mỗi ngày để huấn luyện lại, đạt goodput 64%. Nếu checkpoint mỗi 30 phút, bạn chỉ mất 2,15 giờ, đạt goodput 91%.

Việc khôi phục cũng phải tự động. Khi khởi động lại, tác vụ sẽ tìm checkpoint gần nhất đã ghi xong, bỏ qua bất kỳ checkpoint nào bị dở dang do sự cố, và tiếp tục từ đó mà không cần sự can thiệp của con người. DCP làm cho điều này trở nên đáng tin cậy: tệp.metadata chỉ được ghi sau khi tất cả các shard đã hoàn tất, vì vậy sự hiện diện của nó là một dấu hiệu đáng tin cậy để chọn "checkpoint này đã hoàn tất".

Tác động 3: Tải dữ liệu (dataloading) quyết định liệu GPU của bạn có bao giờ nhàn rỗi hay không

Một tác vụ huấn luyện tiến hành theo tốc độ của đầu vào chậm nhất. Khi các bộ tăng tốc chờ đợi batch tiếp theo, goodput của bạn bị giảm do GPU đơn giản là đang nhàn rỗi. Cách duy nhất để khắc phục vấn đề này là đảm bảo đường ống đầu vào của bạn chồng lấp việc chuẩn bị dữ liệu cho bước tiếp theo với việc tính toán ở bước hiện tại như trong hình dưới đây:

Chúng tôi thường thấy các khách hàng chuyển sang chồng lấp việc tải dữ liệu với tính toán đã giảm được 20–50% thời gian thực tế (wall-clock time).

Chi phí đọc trực tiếp từ bộ nhớ từ xa

Trên một nền tảng được quản lý, dữ liệu huấn luyện nằm trong bộ nhớ đối tượng từ xa (remote object storage). Trên AI Runtime, các Unity Catalog (UC) volume được hiển thị dưới dạng các mount mạng.

Việc đọc tệp trực tiếp từ mount đó mỗi khi truy cập sẽ gắn thời gian thực hiện bước của bạn với độ trễ mạng và tải lại cùng các tệp đó mỗi epoch. Giải pháp là một dataloader sao chép từng tệp vào bộ nhớ cục bộ nhanh khi truy cập lần đầu, phục vụ các lần đọc tiếp theo từ bộ nhớ đệm cục bộ đó và tìm nạp các tệp sắp tới song song trong khi GPU tính toán.

Với AI Runtime, UCVolumeDataset và DataLoader thực hiện chính xác điều này (xem hướng dẫn để biết ví dụ về mã nguồn). UCVolumeDataset truyền phát tệp từ một UC volume, lưu vào bộ nhớ đệm từng tệp vào NVMe cục bộ khi truy cập lần đầu, và phân chia tệp giữa các rank và worker để mỗi bộ tăng tốc nhận được một phần dữ liệu riêng biệt, không chồng lấp. DataLoader của chúng tôi là một lớp con thay thế cho PyTorch DataLoader với các mặc định được tinh chỉnh cho đường dẫn này, vì vậy các tệp được tìm nạp và lưu vào bộ nhớ đệm đồng thời trong khi GPU tính toán thay vì từng tệp một trên luồng huấn luyện.

Ví dụ: huấn luyện một mô hình hình ảnh từ các tệp UC