Hugging Face 接入 TorchTPU,Transformers 与 Diffusers 原生支持 TPU
Hugging Face 宣布接入 TorchTPU 这一 PyTorch 原生 TPU 后端,Transformers、Diffusers 和 Kernels 可直接在 TPU 上运行,不再需要 Optimum TPU 或 torch_xla。
加载模型时传入 device_map="tpu" 即可,Diffusers 只需 pipe.to("tpu"),编译模式再加一次 compile 调用;超出单芯片显存的模型可通过 TensorParallelConfig 在加载时分片到多芯片。
关注点
Hugging Face 生态原生支持 TPU 后,PyTorch 用户不必再迁到 torch_xla 或 JAX,现有模型代码改一个设备参数即可运行。