google / jetstream-pytorch

PyTorch/XLA integration with JetStream (https://github.com/google/JetStream) for LLM inference"
Apache License 2.0
33 stars 14 forks source link

V5e8 ray #159

Closed FanhaiLu1 closed 1 month ago

FanhaiLu1 commented 1 month ago

This PR add ray support for v5e-8 vms.

Thanks Richard for adding 8 chips TPU support in ray. But there is an issue to get get_current_pod_worker_count for v5e-8, this RP add parameter to let engineer specify worker count.