Tunix: Thư viện của Google giúp tăng tốc huấn luyện AI tác tử (Agentic RL) trên TPU

21 THÁNG 7, 2026

Trọng tâm của việc căn chỉnh (alignment) LLM đã nhanh chóng chuyển dịch từ các chatbot tĩnh sang các quy trình làm việc có tính tác nhân (agentic workflows) năng động. Các mô hình ngày nay không chỉ đơn thuần là trò chuyện—chúng thực thi các suy luận đa bước, gọi các API bên ngoài và tương tác với các môi trường phức tạp.

Việc huấn luyện các tác nhân suy luận gặp phải những thách thức và nút thắt đặc biệt. Sự phát triển gần đây của quá trình huấn luyện RL (học tăng cường) cho tác nhân đã chuyển dịch từ căn chỉnh đơn lượt sang ra quyết định đa lượt với các tương tác môi trường phức tạp và việc sử dụng công cụ. Sự thay đổi này đặt ra những thách thức mới về mặt cơ sở hạ tầng đối với hiệu suất và hiệu quả triển khai; khi một tác nhân tạm dừng để thực thi mã, truy vấn cơ sở dữ liệu hoặc chờ đợi kết quả tìm kiếm trên web, hiệu suất sử dụng bộ tăng tốc AI đắt tiền sẽ giảm mạnh do các TPU phải nhàn rỗi chờ đợi các bước của môi trường.

Tunix—thư viện hậu huấn luyện của Google—giải quyết triệt để nút thắt này trong bản phát hành mới nhất, giới thiệu một khung làm việc hiệu quả, có khả năng tùy biến để huấn luyện các tác nhân LLM ở quy mô lớn. Tunix giữ cho các bộ tăng tốc được sử dụng tối đa trên hai phương diện:

Ngoài việc điều phối, RL cho tác nhân đòi hỏi khả năng quan sát chuyên biệt. Trong khi các trình phân tích hiệu năng tiêu chuẩn như XProf cung cấp các dấu vết (traces) chi tiết ở cấp độ toán tử, thì chi phí vận hành cao của chúng giới hạn việc chỉ có thể thu thập dữ liệu trong thời gian ngắn và rời rạc. Tunix giới thiệu cơ chế đo lường liên tục, nhẹ, được xây dựng trực tiếp xung quanh các chỉ số RL đặc thù của miền. Bằng cách tương quan các chỉ số vòng lặp cấp cao này với dòng thời gian của TPU, các nhà phát triển có được cái nhìn toàn diện về hiệu quả thực thi để nhanh chóng phát hiện và giải quyết các nút thắt của hệ thống.

Cuối cùng, Tunix được xây dựng để tối đa hóa thông lượng TPU, giữ cho các môi trường có tính mô-đun và làm cho hiệu quả huấn luyện đa lượt trở nên hoàn toàn minh bạch. Dưới đây là cách thức hoạt động của nó.

1. Triển khai không đồng bộ & tách biệt: Thời gian nhàn rỗi gần bằng 0, thông lượng tối đa

Đạt được thông lượng phần cứng đỉnh cao đồng nghĩa với việc giữ cho các TPU luôn bận rộn. Tunix thực hiện điều này bằng cách kết hợp các triển khai không đồng bộ để loại bỏ các khoảng trống thực thi và các tác nhân chậm chạp, cùng với một đường ống (pipeline) tách biệt liên tục truyền dữ liệu đến trình huấn luyện.

Triển khai không đồng bộ

Trong RL cho tác nhân, việc tạo quỹ đạo (rollout) là giai đoạn tốn thời gian nhất. Tuy nhiên, kiến trúc triển khai đồng bộ truyền thống tạo ra hai vấn đề lớn, như được mô tả trong hình dưới đây.

Tunix giải quyết vấn đề này bằng Công cụ thu thập quỹ đạo không đồng bộ (Asynchronous Trajectory Collector Engine).

Kiến trúc này chồng lấp hoàn toàn quá trình suy luận mô hình, thực thi công cụ và tính toán phần thưởng, giúp duy trì hiệu suất sử dụng phần cứng ở mức cao.

Đường ống triển khai & huấn luyện tách biệt

Trong khi các triển khai không đồng bộ giải quyết các nút thắt về tạo quỹ đạo, một thách thức quan trọng khác đối với hiệu quả phần cứng trong quy trình RL đầu-cuối là kết nối các quỹ đạo năng động, có độ dài thay đổi và có khả năng kéo dài với một vòng lặp huấn luyện đồng bộ nghiêm ngặt. Một cách tiếp cận đơn giản dựa vào điểm đồng bộ hóa buộc bộ tăng tốc phải chờ đợi cho đến khi toàn bộ lô quỹ đạo hoàn tất trước khi bắt đầu bước huấn luyện, khiến TPU huấn luyện bị bỏ đói.

Tunix loại bỏ nút thắt này bằng cách tách biệt việc triển khai và huấn luyện thành một đường ống sản xuất-tiêu thụ liên tục (được minh họa trong sơ đồ dưới đây):

Ngay khi một nhóm quỹ đạo hoàn tất, nó sẽ được xử lý hậu kỳ, chấm điểm và truyền trực tiếp vào trình huấn luyện. Đường ống này đảm bảo trình huấn luyện đồng bộ luôn được cung cấp dữ liệu liên tục, tối đa hóa thông lượng đầu-cuối.

2. Trừu tượng hóa tác nhân và môi trường có khả năng tùy biến – Môi trường OSS cắm-và-chạy

Một điểm gây khó khăn lớn trong các khung làm việc RL là sự gắn kết cứng nhắc của thuật toán với vòng lặp môi trường. Việc sửa đổi cơ sở mã để hỗ trợ một tiêu chuẩn phần mềm nguồn mở (OSS) mới như SWE-bench, WebArena hoặc một công cụ trò chơi tùy chỉnh thường đòi hỏi phải viết lại rất nhiều.

Tunix giải quyết vấn đề này bằng kiến trúc tách biệt, có khả năng tùy biến. Bằng cách cung cấp một ranh giới API rõ ràng, Tunix tự động hóa việc gọi bước và quản lý vòng đời để bạn có thể tập trung hoàn toàn vào logic tương tác cốt lõi.

Tại sao điều này quan trọng: Bạn có thể tích hợp bất kỳ môi trường RL nguồn mở nào trong vài phút. Vì logic của tác nhân và môi trường hoàn toàn tách biệt với quy trình huấn luyện, việc thay thế một trình xác minh toán học đơn lượt bằng một thiết bị đầu cuối bash tương tác không đòi hỏi bất kỳ sửa đổi nào đối với mã huấn luyện của bạn. Để chứng minh sức mạnh của thiết kế có khả năng tùy biến này, chúng tôi sẽ giới thiệu một vài ví dụ về việc các tác nhân, mô hình hoặc môi trường mới có thể được sử dụng dễ dàng như thế nào. Bạn có thể tìm thấy các ví dụ chi tiết hơn về Tác nhân/Môi trường tùy chỉnh trong các công thức (recipes) của chúng tôi.

Ví dụ 1: Tác nhân dựng sẵn so với Tác nhân tùy chỉnh

Tunix cung cấp các lớp tích hợp sẵn như ModelAgent và ToolAgent hoạt động ngay lập tức thông qua cấu hình.

Python

Đã sao chép

Ngoài ra, bạn có thể tự xây dựng Tác nhân tùy chỉnh của riêng mình và thêm logic cụ thể về cách xử lý phản hồi của mô hình. Tunix sẽ tự động kết nối tác nhân này vào quy trình huấn luyện đầu-cuối. Ví dụ: SWEAgent, FrozenLakeAgent

Python

Đã sao chép

Ví dụ 2: Đưa vào các môi trường tùy chỉnh

Tương tự như Tác nhân, Tunix cung cấp một số môi trường dựng sẵn bao gồm TaskEnvironment, ToolEnvironment. Ngoài ra, bạn cũng có thể đưa vào môi trường tùy chỉnh của riêng mình bằng cách chỉ cần triển khai một vài API chính, bao gồm bất kỳ môi trường nguồn mở nào như ví dụ về Gymnasium dưới đây.

Python

Đã sao chép

3. Loại bỏ "hộp đen": Phân tích hiệu năng nhẹ chuyên biệt cho RL

Khi chạy huấn luyện tác nhân không đồng bộ ở quy mô lớn, việc ghi nhật ký truyền thống không còn hiệu quả. Bạn cần khả năng quan sát chi tiết nhưng đặc thù cho từng miền để xác định các vấn đề về hiệu suất: Nút thắt nằm ở giai đoạn tạo dữ liệu? Việc gọi công cụ mất quá nhiều thời gian? Hay trình tải dữ liệu quá chậm?

Các trình phân tích hiệu năng tiêu chuẩn như XProf cung cấp các dấu vết chi tiết ở cấp độ toán tử để hiểu hiệu suất vi mô như thực thi nhân và mô hình. Tuy nhiên, việc thu thập các dấu vết kéo dài với các công cụ này thường rất tốn kém và việc xác định các nút thắt cấp vĩ mô giữa những dữ liệu nhiễu cấp thấp vẫn rất khó khăn. Đối với các quy trình phức tạp của RL cho tác nhân, các nhà phát triển cần một cái nhìn vĩ mô, nhẹ nhàng, được xây dựng trên các chỉ số đặc thù của miền, ánh xạ trực tiếp đến các giai đoạn RL.

Tunix mang đến bức tranh toàn cảnh này bằng cách theo dõi cẩn thận một tập hợp tối thiểu các chỉ số quan trọng đặc thù cho RL, đại diện cho cả đường ống toàn cục (cách các giai đoạn triển khai, huấn luyện và đồng bộ trọng số tương tác) và các bước phụ quan trọng (mỗi lần gọi mô hình, tương tác môi trường, v.v.). Vì chúng nhẹ, các chỉ số này chạy liên tục trong suốt toàn bộ công việc huấn luyện. Người dùng có thể nhanh chóng xác định nơi quy trình đang bị đình trệ trên toàn hệ thống, sau đó triển khai một công cụ như XProf để gỡ lỗi mục tiêu và chuyên sâu hơn.

Hình trên minh họa một dấu vết Perfetto được thu thập từ một công việc huấn luyện tác nhân đa lượt, trình bày chi tiết dòng thời gian thực thi theo giai đoạn trên các luồng CPU và thiết bị TPU. Như đã chứng minh qua dấu vết, hiệu suất sử dụng thiết bị TPU cao hơn nhiều so với các luồng CPU, vốn có thời gian nhàn rỗi chủ yếu do độ trễ thực thi môi trường.