PyTorch/XLA 文档

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

了解 Pytorch/XLA

• Pytorch/XLA 概述
• 在 XLA 设备上使用 PyTorch
• PyTorch/XLA API
• 动态形状
• Eager 模式 + Compile API
• PJRT 运行时
• 故障排除

了解加速器

• 了解 TPU
• 了解 GPU

PyTorch/XLA 功能

• 通过 Pallas 自定义内核
• Torch 导出到 StableHLO
• 通过 Triton 自定义 GPU 内核
• 使用 scan 和 scan_layers 的指南

提高 Pytorch/XLA 工作负载性能

• 自动混合精度
• PyTorch/XLA SPMD 用户指南
• PyTorch/XLA SPMD 高级主题
• 分布式检查点
• 在 GPU 上运行 SPMD
• 如何进行 DistributedDataParallel(DDP)
• TorchDynamo 集成到 PyTorch XLA 中
• 使用 while_loop 优化内存利用率
• PyTorch XLA 中的完全分片数据并行
• 使用 SPMD 的完全分片数据并行
• XLA 的量化操作（实验性功能）
• Pytorch/XLA 中重新编译的来源

为 Pytorch/XLA 做贡献

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