PyTorch/XLA 文档

PyTorch/XLA 是一个 Python 包，它使用 XLA 深度学习编译器连接 PyTorch 深度学习框架和 Cloud TPU。

了解 PyTorch/XLA

• PyTorch/XLA 概览
• 在 XLA 设备上运行 PyTorch
• PyTorch/XLA API
• 动态 shape
• Eager Mode + Compile API
• PJRT Runtime
• 故障排除

了解加速器

• 了解 TPU
• 了解 GPU

PyTorch/XLA 特性

• 通过 Pallas 实现自定义 Kernel
• Torch Export 到 StableHLO
• 通过 Triton 实现自定义 GPU Kernel
• 使用 scan 和 scan_layers 的指南

提高 PyTorch/XLA 工作负载性能

• 自动混合精度
• PyTorch/XLA SPMD 用户指南
• PyTorch/XLA SPMD 高级主题
• 分布式检查点
• 在 GPU 上运行 SPMD
• 如何进行 DistributedDataParallel(DDP)
• TorchDynamo 在 PyTorch XLA 中的集成
• 使用 while_loop 优化内存利用率
• Fully Sharded Data Parallel 在 PyTorch XLA 中
• 使用 SPMD 实现 Fully Sharded Data Parallel
• 用于 XLA 的量化操作 (实验性功能)
• PyTorch/XLA 中重复编译的来源

贡献 PyTorch/XLA

• 配置开发环境
• Codegen 迁移指南
• OP Lowering 指南
• 自定义硬件插件
• PyTorch/XLA 中的 Bazel