optimum-tpu 文档
Jetstream Pytorch 和 PyTorch XLA 之间的区别
加入 Hugging Face 社区
并获得增强的文档体验
开始使用
Jetstream Pytorch 和 PyTorch XLA 之间的区别
本指南向 optimum-tpu 用户解释 Jetstream Pytorch 和 PyTorch XLA 之间的区别，因为它们是 TGI 中可用的两种后端。
JetStream PyTorch 是一个基于 PyTorch XLA 构建的高性能推理引擎。它在 TPU 上运行大型语言模型 (LLM) 时针对吞吐量和内存效率进行了优化。
| 特性 | Jetstream Pytorch | PyTorch XLA |
|---|---|---|
| 训练 | ❌ | ✅ |
| 服务 | ✅ | ✅ |
| 性能 | 更高的服务性能 | 标准性能 |
| 灵活性 | 仅限于服务 | 完整的 PyTorch 生态系统 |
| 用例 | 生产推理 | 开发和训练 |
| 集成 | 针对部署优化 | 标准 PyTorch 工作流程 |
注意：默认情况下，optimum-tpu 使用 PyTorch XLA 进行训练，使用 Jetstream Pytorch 进行服务。
您可以配置 optimum-tpu 以使用任一版本进行 TGI 服务。您可以通过在 docker 运行参数中设置 -e JETSTREAM_PT_DISABLE=1 来在 TGI 中使用 Pytorch XLA 后端。
您可以找到更多关于以下内容的信息
- PyTorch XLA：https://pytorch.ac.cn/xla/ 和 https://github.com/pytorch/xla
- Jetstream Pytorch：https://github.com/AI-Hypercomputer/jetstream-pytorch