Chạy Ray trên TPU: Hướng dẫn nền tảng (Phần 1)

20 THÁNG 7, 2026

TL;DR: Nếu bạn đã từng mở rộng quy mô Python với Ray trên GPU, thì giờ đây mã nguồn của bạn có thể chạy trên TPU (Tensor Processing Unit, chip tăng tốc AI của Google) với các API chính thức được hỗ trợ đầy đủ mà bạn đã quen thuộc. Mô hình task-and-actor, JaxTrainer, hay việc triển khai Ray Serve đều có thể hướng trực tiếp đến các TPU được điều phối bởi Google Kubernetes Engine (GKE).

Kể từ phiên bản Ray 2.55, Google Cloud TPU đã trở thành một bộ tăng tốc hạng nhất trong Ray. Điều này có nghĩa là TPU hiện nằm trong các quy trình phát hành của Ray với các hình ảnh (images) được xây dựng sẵn chính thức và được hỗ trợ trên toàn bộ các thư viện cốt lõi, thay vì con đường "thử nghiệm" cũ nơi bạn phải tự xây dựng container và dựa vào sự hỗ trợ từ cộng đồng. Trong loạt bài "Chạy Ray trên TPU" này, bạn sẽ tìm hiểu cách một TPU slice chỉ đơn giản là một bộ tăng tốc khác mà Ray lập lịch (Phần 1), sau đó đi sâu vào từng thư viện (Phần 2).

Ray và TPU: Giới thiệu nhanh

Ray là một framework tính toán phân tán: bạn viết mã Python, và Ray chạy nó trên một cụm dưới dạng các task (hàm không trạng thái) và actor (worker có trạng thái). Đối với Ray, TPU chỉ là một tài nguyên có thể lập lịch khác, giống như cách hoạt động của GPU. Bạn yêu cầu tài nguyên, Ray sẽ đặt công việc của bạn lên đó.

Nhưng có một điều cần lưu ý trước khi chúng ta tiếp tục.

Các chip TPU được kết nối với nhau thành một nhóm cố định gọi là slice: bao gồm nhiều máy chủ (VM) có các chip chia sẻ một liên kết tốc độ cao chuyên dụng gọi là ICI (Inter-Chip Interconnect). Một mô hình đa máy chủ phải nằm gọn trong một slice duy nhất, nếu không các worker sẽ không thể kết nối với nhau và công việc sẽ bị treo.

Nếu bạn tư duy theo hướng GPU, hãy hình dung một slice giống như một hộp chứa nhiều GPU, nơi kết nối tốc độ cao (NVLink) chỉ tồn tại bên trong hộp đó. Nếu bạn chia các worker ra hai hộp mà không có cáp kết nối giữa chúng, thì các thao tác tập thể (collective operations) như các bước all-reduce để đồng bộ hóa gradient sẽ không bao giờ hoàn tất. Quá trình huấn luyện sẽ bị treo. Một TPU slice hoạt động theo cách tương tự: ICI chính là sợi cáp đó, và nó chỉ kết nối được các chip trong cùng một slice.

Đó chính là lý do tại sao "Ray trên TPU" cần những thiết lập đặc biệt. Cần có cơ chế đảm bảo tất cả các worker của bạn đều nằm trên một slice nguyên vẹn. Trên GPU, bạn hầu như không cần bận tâm về điều này; nhưng trên TPU, đây là yếu tố then chốt, và Ray cùng GKE (Google Kubernetes Engine, dịch vụ Kubernetes được quản lý của Google) sẽ xử lý việc đó cho bạn.

Một thuật ngữ nữa bạn sẽ thấy ở khắp mọi nơi là topology: đó là hình dạng của một slice, được viết dưới dạng 4x4 cho một slice 16 chip. Bạn yêu cầu một topology, chứ không phải số lượng chip.

Khi đã hiểu về slicing và topology với TPU, hệ thống Ray hiện tại và quy trình phát triển của bạn vẫn giữ nguyên và nó sẽ chạy trên các TPU slice do GKE cung cấp. Sơ đồ dưới đây mô tả toàn bộ hệ thống trong một hình ảnh: bên trái là mã nguồn bạn viết (các thư viện Ray bạn đã sử dụng); ở giữa là lớp Ray Core thực hiện đặt trước toàn bộ các slice; bên phải là lớp quản lý GKE thực hiện cung cấp phần cứng và gắn nhãn để Ray có thể xác định ranh giới của các slice.

GKE cung cấp một slice và gắn nhãn các máy chủ của nó, Ray Core đọc các nhãn đó để đặt trước toàn bộ slice cùng một lúc, và các lệnh gọi thư viện của bạn nằm ở trên cùng, chỉ cần khai báo topology là đủ. Không cần phải viết mã thủ công để đặt vị trí ở bất cứ đâu. Phần còn lại của bài viết này sẽ đi qua hai lớp dưới cùng là GKE và Ray Core, trong khi Phần 2 sẽ đề cập đến các thư viện AI của Ray.

Cách GKE điều phối Ray trên TPU

Bạn chạy Ray trên TPU thông qua GKE bằng cách sử dụng tiện ích bổ sung Ray Operator.

Shell

Đã sao chép

Cờ lệnh duy nhất đó cài đặt hai thứ quan trọng cho TPU. Thứ nhất là KubeRay, một Kubernetes operator giúp chuyển đổi các tệp YAML của RayCluster, RayService và RayJob thành các cụm Ray đang chạy; đây cũng chính là KubeRay mà bạn sử dụng với GPU. Thứ hai là phần dành riêng cho TPU: Ray TPU webhook, giúp gắn nhãn cho mọi máy chủ TPU với các nhãn như ray.io/tpu-slice-name để Ray có thể nhận biết máy nào được kết nối trong cùng một slice. Nhãn đó chính là sợi dây liên kết mà toàn bộ hệ thống dựa vào.

Từ đó, bạn yêu cầu TPU trong một manifest giống như cách bạn yêu cầu bất kỳ node nào, với nodeSelector cho thế hệ (generation) và topology, cùng với số lượng chip dưới dạng tài nguyên. Một slice đa máy chủ sẽ thêm một trường là numOfHosts.

Shell

Đã sao chép

GKE cung cấp slice, webhook gắn nhãn cho nó, và Ray đọc các nhãn đó. Bạn chỉ cần viết mã Python. Khi tiện ích bổ sung đã sẵn sàng, bạn sẽ thấy pod KubeRay operator đang chạy, và việc áp dụng manifest đó sẽ khởi chạy một head pod cùng với một worker pod cho mỗi máy chủ trong slice. Ví dụ bắt đầu (get-started) trong phần cụm (cluster) sẽ cung cấp tất cả những thứ này bằng Terraform.

Điều thực sự giữ cho các worker của bạn ở cùng nhau là một primitive của Ray Core nằm ngay phía trên lớp này, đó là slice placement group, và đó cũng là nơi phần còn lại của hướng dẫn bắt đầu.

Ray Core trên TPU

Ray Core là lớp cơ sở, là công cụ task-and-actor và bộ lập lịch mà mọi thứ khác đều dựa vào. Hỗ trợ TPU của nó nằm trong API công khai ray.util.tpu, và thực sự chỉ có một hàm cần biết: slice_placement_group. Nó lấy cam kết "giữ các worker của tôi trên một slice nguyên vẹn" từ trước đó và biến nó thành một lệnh gọi duy nhất, đặt trước toàn bộ một slice một cách nguyên tử (tất cả máy chủ hoặc không gì cả) bằng cách khớp với các nhãn của webhook.

Python

Đã sao chép

Cần lưu ý rằng bạn hiếm khi phải tự gọi slice_placement_group. Các thư viện AI của Ray (Data, Train, Serve) sẽ gọi hàm đó thay cho bạn, vì vậy trên thực tế, bạn chỉ cần khai báo topology và chúng sẽ xử lý slice. Bạn chỉ nên sử dụng trực tiếp slice_placement_group khi đang viết một khối lượng công việc phân tán tùy chỉnh không phải là Train, Serve hay Data. Một lưu ý nhỏ: API này là công khai nhưng được đánh dấu là alpha (@PublicAPI(stability="alpha")), vì vậy bạn có thể sử dụng ngay hôm nay nhưng cấu trúc có thể thay đổi giữa các bản phát hành.

Đó là nền tảng. Tiếp theo là các thư viện.

Bây giờ bạn đã có toàn bộ mô hình tư duy: một slice phải được giữ nguyên vẹn, GKE cung cấp và gắn nhãn nó, và Ray Core đặt trước nó như một đơn vị để bạn không bao giờ phải viết mã đặt vị trí thủ công. Mọi thứ bạn xây dựng thực tế đều nằm trên nền tảng đó và tái sử dụng nó.

Trong Phần 2, chúng ta sẽ khám phá cách bạn có thể sử dụng các thư viện AI của Ray trên TPU để phục vụ LLM với vLLM, nạp dữ liệu vào các slice với Ray Data và huấn luyện với JaxTrainer.

Tài nguyên bổ sung

Hiện tại, cảm ơn bạn đã đọc! Và nếu bạn có bất kỳ câu hỏi hoặc phản hồi nào, đừng ngần ngại liên hệ qua mạng xã hội (LinkedIn, X).

Chúc bạn xây dựng thành công!

Trước

Tiếp theo