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