Đột phá huấn luyện AI: MaxText tự động phục hồi khi lỗi phần cứng chỉ trong vài giây

NGÀY 6 THÁNG 7 NĂM 2026

TLDR (Tóm tắt)

Điều gì sẽ xảy ra khi một máy bị hỏng ngay giữa quá trình huấn luyện đa nút (multi-node)?

Nếu bạn từng huấn luyện các mô hình lớn trên nhiều máy, bạn đã biết câu trả lời: kết nối bị gián đoạn (timeout), mọi worker đều thoát và bạn phải khởi chạy lại toàn bộ công việc từ điểm kiểm tra (checkpoint) gần nhất. Điều này rất tốn kém và đó là cách thức hoạt động của huấn luyện phân tán. Nhưng liệu có cách nào khác không?

Trong bài viết này, chúng ta sẽ khám phá một giải pháp khả thi gọi là huấn luyện đàn hồi (elastic training), sử dụng ngăn xếp JAX AI (MaxText và Pathways) trên Cloud TPU. Chúng ta sẽ huấn luyện một LLM trên nhiều chip TPU trên Google Kubernetes Engine (GKE), cố tình làm cho một worker bị lỗi và quan sát quá trình huấn luyện tự phục hồi tại chỗ mà không cần khởi động lại. Tất cả đều sử dụng cùng một tiến trình, cùng một PID, không cần khởi chạy lại. Trong quá trình chạy thử của chúng tôi, tổng thời gian ngừng hoạt động từ lúc bị ngắt đến bước huấn luyện tiếp theo là dưới hai phút, và phần lớn thời gian đó là chờ Kubernetes lập lịch cho pod thay thế.

Đến cuối bài, bạn sẽ hiểu chính xác ba thành phần nào tạo nên khả năng này, những điểm còn hạn chế của phương pháp này và cách tự mình tái hiện lại toàn bộ quy trình.

Hãy bắt đầu với vấn đề cốt lõi.

Tại sao huấn luyện phân tán lại mong manh đến vậy?

Hãy tưởng tượng bạn đang huấn luyện một mô hình trên nhiều máy (hoặc nút). Trọng số mô hình của bạn được phân mảnh (sharded), vì vậy mỗi máy nắm giữ một phần. Ở mỗi bước huấn luyện, các máy tính toán gradient trên phần dữ liệu của chúng, sau đó một thao tác all-reduce được thực hiện để tất cả trao đổi gradient nhằm đảm bảo mô hình luôn đồng bộ.

Vấn đề nằm ở chỗ: all-reduce cần sự tham gia của tất cả các bên. Nếu một máy biến mất, các máy còn lại sẽ ngồi chờ dữ liệu không bao giờ đến. Cuối cùng, lệnh timeout được kích hoạt, quá trình tập thể thất bại và mọi tiến trình đều thoát. Kết quả là, một máy bị hỏng sẽ kéo theo toàn bộ công việc bị dừng lại.

Giải pháp tiêu chuẩn thường nằm ngoài mã nguồn huấn luyện của bạn. Một bộ lập lịch (Slurm, Kubernetes, Ray, tùy bạn chọn) phát hiện công việc đã chết, phân bổ lại tài nguyên và khởi chạy lại mọi thứ từ đầu. Bạn phải trả giá cho toàn bộ chi phí khởi động lại: lập lịch các pod mới, khởi động các container và tiến trình Python mới, kết nối lại với các bộ tăng tốc và làm nóng bộ tải dữ liệu (dataloader). Và bạn mất mọi bước kể từ checkpoint gần nhất.

Nhưng nếu tiến trình huấn luyện có thể tự bắt lỗi và tiếp tục chạy thì sao? Đó là điều chúng ta sẽ xem xét ngay bây giờ.

MaxText, Pathways và Orbax trên TPU

Nếu bạn mới làm quen với hệ sinh thái TPU, nó có bộ từ vựng riêng. Hãy giới thiệu các thành phần chúng ta sẽ sử dụng và cách chúng kết hợp với nhau.

Đơn vị phần cứng của chúng ta là chip TPU, một bộ tăng tốc được xây dựng xung quanh các đơn vị nhân ma trận. Các chip được nhóm thành một TPU slice, một tập hợp các chip được kết nối với nhau bằng một kết nối chuyên dụng tốc độ cao gọi là ICI (Inter-Chip Interconnect). Trong một slice, các chip được gắn vào các máy ảo (VM) CPU thông thường (trong thiết lập của chúng tôi là bốn host mỗi slice), và chính các host này là nơi chạy các tiến trình worker.

Để lập trình các chip đó, chúng ta sử dụng JAX, một framework mảng và autograd mang phong cách NumPy, sử dụng XLA làm trình biên dịch. Nếu bạn đến từ PyTorch, nó đóng vai trò tương tự. Tuy nhiên, chúng ta sẽ không viết vòng lặp huấn luyện từ đầu. Chúng ta sẽ sử dụng MaxText, một thư viện huấn luyện LLM mã nguồn mở được viết bằng Python thuần và JAX. Bạn chỉ cần cung cấp tên mô hình và cấu hình, nó sẽ cung cấp cho bạn một vòng lặp huấn luyện đã được phân mảnh hoàn toàn.

Hai thành phần nữa sẽ hoàn thiện bức tranh. Pathways là lớp điều phối kết nối tập lệnh Python của chúng ta với tất cả các chip, và đây là chìa khóa cho toàn bộ câu chuyện này vì một lý do mà chúng ta sẽ đề cập ngay sau đây. Và Orbax xử lý việc lưu checkpoint: nó điều phối việc lưu từ bộ điều khiển (controller), trong khi mỗi TPU host ghi phân mảnh trạng thái mô hình của riêng nó trực tiếp lên Cloud Storage song song. Điều đó mang lại cho chúng ta một phương án dự phòng khi có sự cố.

Đây là cách mọi thứ kết hợp với nhau:

Thành phần cần ghi nhớ ở đây là bộ điều khiển đơn (single controller).

Trong hầu hết các trình khởi chạy huấn luyện phân tán, bạn bắt đầu một tiến trình Python trên mỗi nút. Mỗi tiến trình chạy một bản sao giống hệt nhau của tập lệnh và chúng phối hợp với nhau như những thực thể ngang hàng (đây gọi là SPMD, hay Single Program Multiple Data). Với Pathways, chỉ có duy nhất một tiến trình Python chạy trên một máy CPU thông thường, và nó nhìn thấy mọi chip TPU trong cụm như thể chúng là cục bộ. Gọi jax.devices() và bạn sẽ nhận lại tất cả các chip. Bản thân các máy TPU chỉ chạy một tệp nhị phân worker nhẹ để nhận các chương trình đã biên dịch và thực thi chúng.

Tại sao điều này quan trọng đối với việc xử lý lỗi? Bởi vì khi một máy TPU bị hỏng, vẫn còn một tiến trình Python khỏe mạnh đang chạy trên nút CPU có thể can thiệp vào vấn đề đó.

Hãy xem "can thiệp vào vấn đề đó" trông như thế nào.

"Huấn luyện đàn hồi" có nghĩa là gì ở đây?

Hãy làm rõ những gì chúng ta đang xây dựng, vì thuật ngữ "đàn hồi" (elastic) được sử dụng cho rất nhiều thứ.

Về cốt lõi, huấn luyện đàn hồi ở đây có nghĩa là khi phần cứng gặp sự cố, vòng lặp huấn luyện của bạn sẽ nhận được một ngoại lệ (exception) Python thay vì bị chấm dứt tiến trình. Và vì bạn vẫn đang ở trong một tiến trình đang hoạt động, với cấu hình và các thư viện đã được tải sẵn, cùng các TPU slice còn sống vẫn đang chờ lệnh, bạn có nhiều lựa chọn.

Hai ví dụ đơn giản nhưng mạnh mẽ là tạm dừng và tiếp tục (pause and resume) và thay đổi kích thước bản sao (replica resize).

Trong chế độ tạm dừng và tiếp tục, ngoại lệ được bắt lại, bạn đợi slice bị lỗi được thay thế, sau đó tải lại checkpoint khả thi gần nhất và tiếp tục trên toàn bộ mesh. Trong chế độ thay đổi kích thước bản sao, bạn tải lại checkpoint khả thi gần nhất ngay lập tức lên các slice còn sống và quá trình huấn luyện tiếp tục chạy với thông lượng giảm trong khi chờ slice thay thế hoạt động, sau đó mở rộng trở lại kích thước đầy đủ khi nó sẵn sàng.

Cả hai đều có sẵn trong MaxText hiện nay. Bài viết này hướng dẫn về chế độ tạm dừng và tiếp tục, cách đơn giản hơn trong hai cách. Bạn có thể tự viết trình xử lý ngoại lệ đó, nhưng bạn không cần phải làm vậy. Thư viện pathwaysutils cung cấp một decorator gọi là elastic_retry bao bọc toàn bộ hàm huấn luyện, và MaxText đã tích hợp sẵn cho bạn. Khi ngoại lệ lỗi xảy ra, decorator sẽ bắt lấy nó, dọn dẹp mọi trạng thái dang dở, khôi phục checkpoint khả thi gần nhất và gọi lại hàm huấn luyện. Tất cả đều nằm trong cùng một tiến trình.

Cần phải làm rõ tại sao điều này nhanh hơn khởi động lại, vì phục hồi đàn hồi không bỏ qua nhiều bước như bạn nghĩ. Việc gọi lại hàm huấn luyện từ đầu có nghĩa là thiết lập mô hình, bộ tải dữ liệu và khôi phục checkpoint đều chạy lần thứ hai, nhưng bạn cũng sẽ phải trả giá cho tất cả những điều đó khi khởi động lại toàn bộ công việc. Việc lập lịch pod ở đây cũng không miễn phí: worker bị lỗi vẫn cần một pod thay thế được lập lịch trên slice bị ảnh hưởng, và thời gian chờ đó chiếm phần lớn thời gian thực tế. Những gì phục hồi đàn hồi thực sự tiết kiệm cho bạn là mọi thứ xung quanh slice đó. Một lần khởi động lại toàn bộ sẽ phá hủy và lập lịch lại toàn bộ khối lượng công việc, từ pod điều khiển (head), mọi pod worker khỏe mạnh, cho đến tiến trình Python điều khiển mới đi kèm với chúng, trong khi phục hồi đàn hồi để tất cả chúng tiếp tục chạy và chỉ thay thế slice đã chết.

Đáng chú ý là việc biên dịch không nằm trong phần tiết kiệm đó. Pathways duy trì bộ nhớ đệm biên dịch bền vững trong Cloud Storage (mặc định bật), vì vậy việc khởi động lại toàn bộ sẽ tải lại các tệp thực thi XLA đã biên dịch từ bộ nhớ đệm thay vì biên dịch lại từ đầu, và phục hồi đàn hồi cũng phải trả chi phí tương đương khi nó nhập lại hàm huấn luyện trên mesh đã được xây dựng lại. Sự khác biệt giữa hai con đường là việc phá hủy (teardown), không phải việc biên dịch — và trong quá trình chạy của chúng tôi, việc bỏ qua việc phá hủy toàn bộ khối lượng công việc đã tạo ra sự khác biệt giữa hàng trăm giây và vài phút. Thay đổi kích thước bản sao sau đó bổ sung một thứ mà khởi động lại hoàn toàn không thể cung cấp: quá trình huấn luyện vẫn tiếp tục ngay cả khi một số TPU không bao giờ quay trở lại.

Trước khi tiếp tục, một lưu ý nhanh về một cái tên dễ gây nhầm lẫn với những gì chúng ta vừa mô tả. Suspend-resume và elastic pause-resume nghe có vẻ giống nhau nhưng giải quyết các vấn đề khác nhau. Nếu bạn đang chạy trên Spot TPU, các lần tạm dừng theo kế hoạch (planned preemptions) có con đường riêng: tính năng suspend-resume của Pathways lắng nghe thông báo tạm dừng, tự động lưu trạng thái bộ tăng tốc vào Cloud Storage và tiếp tục khi pod được lập lịch lại — không cần mã người dùng. Đó là dành cho các gián đoạn có cảnh báo trước. Huấn luyện đàn hồi, cơ chế mà chúng ta đang tìm hiểu ở đây, là con đường dành cho các lỗi không mong muốn mà không có bất kỳ thông báo nào.

Bây giờ hãy xem xét ba thành phần thực sự phải hợp tác để làm cho điều này hoạt động.

Phục hồi đàn hồi hoạt động như thế nào?

Để phục hồi đàn hồi, ba thành phần độc lập phải hợp tác. Đây là cách chúng được bố trí trên cụm.

Đầu tiên, Pathways phát hiện lỗi. Điều này có thể xuất hiện theo hai cách. Phổ biến nhất, một thao tác đang thực hiện đến worker đã chết bị lỗi và Pathways trả về lỗi DATA_LOSS. Nếu không có gì đang thực hiện, trình quản lý tài nguyên (một container chạy cùng với tập lệnh của chúng ta trên nút CPU) nhận thấy worker đã ngừng gửi tín hiệu heartbeat và trả về DEADLINE_EXCEEDED sau khoảng 10 giây. Dù bằng cách nào, lỗi sẽ xuất hiện trong bước huấn luyện của chúng ta dưới dạng jax.errors.JaxRuntimeError. Lỗi phần cứng đã trở thành một ngoại lệ Python có thể bắt được.

Tiếp theo, decorator elastic_retry bắt lấy ngoại lệ. Decorator này đến từ pathwaysutils; MaxText chỉ đơn giản áp dụng nó xung quanh hàm huấn luyện của mình. Nó bắt ngoại lệ cụ thể đó, ghi lại thông báo "Slice down event detected. Retrying." và chạy đường dẫn phục hồi thay vì để lỗi làm treo tiến trình.

Cuối cùng, Orbax quyết định những gì an toàn để khôi phục. Phần này không chỉ dành riêng cho huấn luyện đàn hồi; đó là cách hoạt động của việc lưu checkpoint Orbax nói chung, và một lần khởi động lại toàn bộ công việc cũng sẽ dựa vào nó theo cách tương tự. Các checkpoint được ghi vào Cloud Storage trong nền trong khi huấn luyện chạy, và một checkpoint chỉ được coi là khả thi khi mọi phân mảnh đã được ghi xong và một tệp đánh dấu commit_success nhỏ được ghi bên cạnh nó. Trong quá trình phục hồi, mã dọn dẹp sẽ kiểm tra thư mục checkpoint mới nhất: nếu không có tệp đánh dấu (nó đang ghi dở khi có sự cố), thư mục đó sẽ bị xóa và chúng ta quay lại checkpoint mới nhất có tệp đánh dấu. Đó là điều đảm bảo chúng ta không bao giờ tải một checkpoint bị lỗi, bất kể chúng ta khởi động lại như thế nào.