Chạy Ray trên TPU (Phần 2): Tối ưu hóa thư viện AI cho hạ tầng Google

24 THÁNG 7, 2026

Tóm tắt: Phần 2 trong tổng số 2 phần. Phần 1 đã đề cập đến ý tưởng phần cứng duy nhất bạn cần và hai lớp bên dưới (GKE và Ray Core). Phần này giới thiệu các thư viện mà bạn thực sự sử dụng để xây dựng, bao gồm Ray Serve, Ray Data và Ray Train.

Tóm tắt lại

Tóm tắt nhanh nếu đây là lần đầu bạn đọc bài viết này. Việc chạy Ray trên TPU có một lưu ý quan trọng: các chip TPU được kết nối thành các nhóm cố định gọi là slice (các máy chủ ảo host VM chia sẻ một liên kết tốc độ cao gọi là ICI), và một mô hình đa máy chủ (multi-host) phải nằm gọn trong một slice hoàn chỉnh, 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.

Google Kubernetes Engine (GKE) với tiện ích bổ sung Ray Operator sẽ cấp phát các slice và gắn nhãn cho các host của chúng, đồng thời một thành phần cơ bản của Ray Core là slice_placement_group sẽ đặt trước toàn bộ một slice cùng lúc. Bạn chỉ cần khai báo cấu trúc liên kết (hình dạng slice, ví dụ: 4x4 cho 16 chip) và các thư viện bên dưới sẽ tự động xử lý việc sắp xếp cho bạn.

Với việc Ray Core xử lý việc sắp xếp bên dưới, tất cả các thư viện đều tuân theo cùng một mô hình: khai báo cấu trúc liên kết và để Ray Core đặt trước slice. Điểm khác biệt giữa các thư viện chỉ nằm ở nơi bạn khai báo. Chúng ta sẽ đi theo thứ tự mà hầu hết các nhóm áp dụng, bắt đầu với Ray Serve.

Ray Serve trên TPU

Serving (phục vụ mô hình) là nơi hầu hết các nhóm bắt đầu. Một mô hình cần nhiều GPU để chạy có thể hoạt động trên một host TPU duy nhất, và TPU thường là lựa chọn sẵn có và tiết kiệm chi phí hơn cho việc suy luận (inference). Ray Serve cung cấp cho bạn các tính năng thông thường như tự động mở rộng (autoscaling), cân bằng tải và kết hợp đa mô hình; trên TPU, nó phục vụ các LLM thông qua vLLM, một engine có lưu lượng xử lý cao.

Rất tiếc, trình duyệt của bạn không hỗ trợ phát video này.

Trường hợp khó khăn là khi một mô hình quá lớn so với một host (ví dụ: một tensor được phân mảnh theo kiểu tensor-parallel trên 16 chip). Đó là lúc Serve giải quyết vấn đề bằng một trường bổ sung duy nhất: topology.

Văn bản thuần

Đã sao chép

Trường đó rất đáng để tìm hiểu, vì nếu thiết lập sai, đây sẽ là lỗi điển hình của TPU đa máy chủ. Khi đã thiết lập topology, backend TPU của Serve sẽ bỏ qua nhóm sắp xếp (placement group) thông thường ban đầu và chuyển quyền cho bản sao (replica), từ đó tạo ra một slice placement group khi khởi động. Sự trì hoãn đó giúp các worker của mô hình tensor-parallel nằm trên cùng một lưới ICI chia sẻ. Nếu bỏ qua nó, Serve sẽ quay lại sử dụng các gói (bundle) theo từng chip; trên một mô hình đa máy chủ, các gói này có thể bị phân tán trên hai slice, và vì không có ICI giữa các slice, các worker sẽ không bao giờ hoàn thành bước tính toán tập thể đầu tiên. Bạn sẽ không gặp lỗi crash, mà thay vào đó là một deployment bị kẹt ở trạng thái DEPLOYING mãi mãi trong khi bạn lãng phí hàng giờ TPU để tìm kiếm một lỗi thực chất chỉ là thiếu một dòng YAML. Vì vậy, hãy nhớ rằng trường topology tạo nên sự khác biệt.

Trong thực tế, bạn triển khai một RayService (được khuyến nghị dùng thay vì RayCluster thô cho môi trường production) trên một image vLLM TPU đã được công bố, đợi nó đạt trạng thái Running và curl vào endpoint. Các hướng dẫn chính thức của GKE đã bao gồm Llama 3 8B và Mistral 7B trên v5e, Llama 3.1 70B trên v6e và Stable Diffusion. Bước serve trong ví dụ bắt đầu nhanh (get-started) sẽ hướng dẫn toàn bộ quá trình triển khai từ đầu đến cuối.

Ray Data trên TPU: cung cấp dữ liệu cho các bộ tăng tốc với iter_jax_batches

Một bộ tăng tốc nhanh chỉ hữu ích khi dữ liệu được cung cấp liên tục, và TPU đủ nhanh để khiến một bộ tải dữ liệu thông thường trở thành nút thắt cổ chai. Đó là vấn đề mà iter_jax_batches giải quyết. Nó cung cấp cho bạn các batch đã là mảng JAX và đã được phân mảnh trên thiết bị (device-sharded), vì vậy một pipeline đầu vào huấn luyện hoặc một công việc suy luận batch lớn sẽ lấy dữ liệu trực tiếp từ pipeline của Ray Data mà không bị chậm trễ do quá trình sao chép NumPy-sang-JAX ở phía host.

Python

Đã sao chép

API iter_jax_batches thực hiện phân mảnh thiết bị cho bạn và xử lý batch cuối cùng không đủ kích thước (batch không chia hết cho kích thước batch của bạn) bằng cách cho phép chọn rõ ràng giữa drop (bỏ), pad (đệm) hoặc raise (báo lỗi), thay vì gặp lỗi về hình dạng (shape error) sau ba giờ chạy.

Bạn có thể sử dụng nó làm đầu vào cho công việc JaxTrainer, và nó cũng hữu ích không kém khi chạy độc lập cho việc suy luận batch ngoại tuyến trên một tập dữ liệu lớn trên slice TPU. Nó mới được thêm vào Ray gần đây, và bước dữ liệu trong ví dụ bắt đầu nhanh sử dụng nó để chuẩn bị tập dữ liệu và suy luận batch.

Ray Train trên TPU: huấn luyện phân tán với JaxTrainer

Huấn luyện từng là phần khó hiểu nhất của Ray trên TPU do vấn đề về cấu trúc liên kết và việc phải tính toán hình dạng slice trong mã nguồn. JaxTrainer giải quyết vấn đề đó. Nó mang vòng lặp huấn luyện của Ray Train (checkpointing, khả năng chịu lỗi, mở rộng đa slice) đến với JAX – thư viện mảng và tự động phân biệt của Google, cũng là framework gốc cho TPU. Bạn chỉ cần cung cấp một hàm huấn luyện và hình dạng slice, Ray sẽ khởi chạy một worker trên mỗi host, kết nối chúng thành một lưới duy nhất và chạy hàm của bạn trên từng worker.

Python

Đã sao chép

Có hai điều trong đoạn mã này bạn cần lưu ý để tiết kiệm thời gian gỡ lỗi. Lệnh import jax nằm bên trong train_loop_per_worker, không phải ở đầu file, vì mỗi worker khởi tạo JAX trong ngữ cảnh TPU riêng của nó; nếu import ở phạm vi module, bạn sẽ gặp các lỗi khởi tạo thiết bị khó hiểu trước khi bước đầu tiên bắt đầu. Và topology="4x4" là toàn bộ khai báo sắp xếp, dòng mã từng là một khối mã phối hợp viết tay. Khi đặt cạnh JaxTrainer cho GPU hoặc TorchTrainer, sự khác biệt thực sự duy nhất là use_tpu=True và một cấu trúc liên kết thay vì số lượng GPU.

Phần còn lại nó sẽ tự chạy. Điều này là do Ray Train quản lý vòng lặp, bạn có được tính năng checkpointing và khởi động lại khi có lỗi, giúp các tác vụ TPU dài trên các tài nguyên có thể bị thu hồi (preemptible) thực sự hoàn thành, và cấu trúc liên kết có thể mở rộng sang đa slice (Ray tự xử lý việc phối hợp giữa các slice) khi một slice là không đủ. Bước huấn luyện trong ví dụ bắt đầu nhanh là một quy trình JaxTrainer DPO hoàn chỉnh.

Hai phần bổ sung cuối cùng: Docker image cho TPU và các chỉ số dashboard

Là một phần của hỗ trợ bộ tăng tốc hạng nhất, Ray hiện xuất bản các image chính thức rayproject/ray:*-tpu với stack JAX/TPU (jax[tpu], flax, optax, orbax-checkpoint) và các công cụ profiling đã được cài đặt sẵn, vì vậy bạn không cần phải tự tay xây dựng môi trường TPU. Bạn chỉ cần sử dụng image có gắn thẻ -tpu làm nền tảng.

Và để giám sát, Ray Dashboard – giao diện web tích hợp sẵn của Ray để theo dõi trạng thái cụm và công việc – hiện hiển thị mức sử dụng TPU và bộ nhớ bên cạnh CPU và GPU trong tab Cluster, với ray.util.tpu.init_jax_profiler cho phép gắn trình profiler JAX trên mỗi worker vào dashboard.

Tóm tắt

Trong hướng dẫn dành cho nhà phát triển về Ray trên TPU này, chúng ta đã đi qua toàn bộ hành trình từ cách Ray chạy trên TPU đến việc thực hiện các khối lượng công việc AI.

Phần 1 đã chỉ ra rằng việc chạy Ray trên TPU có một lưu ý quan trọng: giữ một mô hình đa máy chủ trên một slice nguyên vẹn, và GKE (thông qua tiện ích Ray Operator) cùng Ray Core (thông qua slice_placement_group) sẽ xử lý điều đó cho bạn. Phần này đã đặt các thư viện AI lên trên: Ray Serve lập lịch nhóm (gang-schedule) một mô hình đa máy chủ vào một slice với một trường accelerator_config.topology duy nhất, Ray Data cung cấp các batch JAX-native cho slice thông qua iter_jax_batches, và JaxTrainer chạy một vòng lặp huấn luyện phân tán từ một ScalingConfig. Vẫn là Ray mà bạn đã sử dụng trên GPU, nay đã có mặt trên TPU.

Tiếp theo là gì

Và nhiều thứ khác đang đến. Nhóm Ray trên Google Cloud đang mở rộng hỗ trợ TPU từ đây: tích hợp sâu hơn giữa Ray Data và Ray LLM trên TPU, SkyRL trên TPU đa máy chủ cho học tăng cường (reinforcement learning) và hậu huấn luyện, cùng hỗ trợ dynamic super/sub-slice đều nằm trong lộ trình. Đối với bước tiếp theo của bạn, lời khuyên của tôi là: hãy clone ví dụ bắt đầu nhanh, dựng cụm, sau đó chạy serve, data hoặc train. Hoặc chỉ cần bật --enable-ray-operator trên một cụm và chạy một tác vụ Ray trên một slice nhỏ để xem nó hoạt động. Bạn không cần phải trở thành chuyên gia TPU mới có thể sử dụng nó, hãy cứ thử xem sao. Hiện tại, cảm ơn bạn đã đọc! Nếu bạn có thêm câu hỏi hoặc phản hồi, đừ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!

Tài nguyên bổ sung

Bạn mới tham gia? Phần 1 giải thích về các slice, GKE và Ray Core, nền tảng mà mọi thứ ở trên được xây dựng dựa vào.

Trước đó