Hệ thống Tunix của Google đã tháo gỡ điểm nghẽn vốn ngăn cản việc sử dụng TPU hiệu quả cho học tăng cường dựa trên tác nhân (agentic reinforcement learning - RL) quy mô lớn. Bằng cách tách biệt công việc tạo dữ liệu tương tác khỏi công việc cập nhật chính sách (policy), Tunix đẩy hiệu suất sử dụng TPU từ mức chỉ vài phần trăm lên mức gần như tối đa, giúp cắt giảm đáng kể sự lãng phí tài nguyên tính toán.

Điểm nghẽn trong agentic RL

Agentic RL khác với việc huấn luyện mô hình ngôn ngữ theo kiểu "token tiếp theo" quen thuộc. Một tác nhân (agent) phải gửi các lệnh gọi API, thực thi mã hoặc thực hiện các bước trong một môi trường mô phỏng, sau đó phản ứng với kết quả. Do đó, vòng lặp huấn luyện là đồng bộ: mô hình đưa ra một hành động, môi trường chạy, kết quả trả về, và chỉ sau đó mô hình mới nhận được bản cập nhật gradient. Khi một bước thực hiện trong môi trường mất vài giây, phần cứng TPU đắt đỏ sẽ rơi vào trạng thái nhàn rỗi, và hiệu suất sử dụng được báo cáo có thể xuống dưới 10%. Sự kém hiệu quả này trực tiếp dẫn đến hóa đơn điện toán đám mây cao hơn và các chu kỳ nghiên cứu chậm hơn.

Kiến trúc tách rời của Tunix

Tunix giải quyết vấn đề bằng cách tách hai giai đoạn—tạo quỹ đạo (trajectory generation) và tối ưu hóa chính sách (policy optimization)—vào các nhóm phần cứng riêng biệt.

  • Các actor bất đồng bộ (Asynchronous actors) chạy trên các CPU hoặc GPU giá rẻ. Mỗi actor liên tục tương tác với môi trường được chỉ định, ghi lại các hành động và quan sát, sau đó truyền các quỹ đạo (trajectories) thu được vào một kho lưu trữ chung.
  • Các bộ học liên tục (Continuous learners) chiếm dụng các TPU Pod chuyên dụng. Bộ học sẽ lấy các batch dữ liệu từ bộ đệm trung tâm và thực hiện cập nhật gradient mà không cần chờ đợi bất kỳ actor đơn lẻ nào hoàn thành một lượt rollout.
  • Bộ đệm thông lượng cao (High-throughput buffer) nằm ở giữa, đóng vai trò là khu vực lưu trữ tạm thời cho các quỹ đạo. Vì bộ học có thể đọc nhanh bằng tốc độ bộ đệm cung cấp dữ liệu, TPU không bao giờ bị đình trệ.

Hiệu quả cuối cùng là một đường ống huấn luyện nơi các TPU luôn bận rộn gần như mọi lúc, đẩy hiệu suất sử dụng lên gần mức 100%.

Các rào cản kỹ thuật và cách Tunix vượt qua

Các episode có độ dài thay đổi và việc biên dịch lại XLA

Trình biên dịch XLA của JAX tối ưu hóa cho các hình dạng tensor cố định. Tuy nhiên, các tác vụ agentic tạo ra các chuỗi có độ dài khác nhau, điều này thường sẽ kích hoạt việc biên dịch lại rất tốn kém. Tunix đóng gói các chuỗi ngắn hơn lại với nhau và nhóm các episode có độ dài tương tự vào các "thùng" (buckets), giữ cho các hình dạng (shapes) ổn định đủ lâu để XLA có thể tái sử dụng các nhân (kernels) đã biên dịch. Kết quả là thông lượng ổn định mà không gặp phải chi phí biên dịch dư thừa vốn có thể làm tê liệt hiệu suất.

Mở rộng các mô hình khổng lồ trên nhiều chip TPU

Việc huấn luyện các tác nhân với hơn 70 tỷ tham số đòi hỏi phải phân tán trọng số và dữ liệu trên nhiều nút TPU. Tunix sử dụng nguyên mẫu ShardMap của JAX để phân mảnh (shard) cả tham số mô hình và các kích hoạt (activations), cho phép bộ học giữ toàn bộ mô hình trong bộ nhớ trong khi vẫn nạp dữ liệu cho nó ở tốc độ cao. Chiến lược phân mảnh này giúp việc huấn luyện các mô hình vốn trước đây nằm ngoài khả năng của một TPU pod đơn lẻ trở nên khả thi.

Gradient bị cũ (stale) từ các đường ống tách rời

Khi các actor chạy nhanh hơn bộ học, dữ liệu chúng cung cấp có thể trở nên "cũ" (stale) so với chính sách hiện tại. Tunix giảm thiểu sự sai lệch này bằng hai cơ chế: lấy mẫu tầm quan trọng (importance-sampling) để tái trọng số các mẫu cũ nhằm phản ánh mức độ liên quan của chúng, và một ngưỡng độ cũ (staleness threshold) có thể cấu hình để loại bỏ các quỹ đạo vượt quá độ tuổi thiết lập sẵn. Cùng với nhau, chúng giữ cho việc học ổn định ngay cả khi đường ống chạy bất đồng bộ.

Những điều người triển khai cần lưu ý

  • Kiểm tra độ trễ (Latency audit) – Lợi ích của việc tách rời phụ thuộc vào thời gian phản hồi của môi trường. Các nhóm nên đo lường độ trễ đầu cuối (end-to-end) và đảm bảo rằng các nhóm actor được định cỡ đủ lớn để giữ cho bộ đệm luôn đầy.
  • Thiết kế nhóm worker (Worker pool design) – CPU hoặc GPU giá rẻ có thể chứa nhiều actor, nhưng việc cấp phát quá mức (oversubscribing) có thể gây ra tranh chấp về mạng hoặc lưu trữ. Một nhóm cân bằng, khớp với tốc độ nạp của bộ đệm là điều thiết yếu.
  • Tính bền vững của bộ đệm (Buffer robustness) – Kho lưu trữ trung tâm phải xử lý được tốc độ ghi và đọc cao mà không trở thành một điểm nghẽn mới. Việc chọn một hệ thống lưu trữ có độ trễ đuôi (tail latency) thấp và băng thông đủ lớn là một phần không thể thương lượng của kiến trúc.

Nhược điểm tiềm ẩn

Kiến trúc tách rời làm tăng thêm nhiều thành phần phức tạp: các đội ngũ phần cứng riêng biệt, một bộ đệm duy trì liên tục và logic điều phối để thực thi các giới hạn về độ cũ của dữ liệu.

Kết luận

Tunix cho thấy chi phí chủ yếu trong agentic RL không nằm ở bản thân mô hình mà ở thời gian nhàn rỗi do các vòng lặp tương tác đồng bộ gây ra. Bằng cách chuyển công việc rollout sang phần cứng giá rẻ và cung cấp dữ liệu cho một TPU pod đang học liên tục từ một bộ đệm thông lượng cao, Google đã biến vấn đề hiệu suất sử dụng dưới 10% thành một quy trình làm việc đạt gần mức công suất tối đa.