Hugging Face ❤️ TorchTPU:原生支持TPU
huggingface ·
Hugging Face 推出了 TorchTPU,这是一个原生的 PyTorch 后端,允许开发者直接在 Google 的张量处理单元(TPU)上运行模型,而无需使用 JAX 或复杂的变通方案。此次更新通过将 TPU 视为普通设备来简化工作流程,使其能够与 Transformers 和 Diffusers 等流行库无缝集成。用户现在可以使用简单的“device_map='tpu'”参数从 Hugging Face Hub 加载模型,并支持推理和训练的 eager 模式及编译模式。该生态系统还通过 Kernels Hub 提供了优化的 Pallas 内核,显著提升了 MoE 解码等特定操作的性能。测试结果显示,96.5% 的 Transformers 测试套件在 TPU 上通过,相关改进正通过 GitHub Issue 持续报告。这一发展消除了此前的软件障碍,使更广泛的 PyTorch 社区能够访问高性能 AI 加速能力。
huggingface