gpu-tpu-pedia

GPU 和 TPU 学习与实践知识库

---

项目结构

TPU

1. HunyuanVideo-1.5 TPU 推理 ⭐

在 Google Cloud TPU v6e-8 上运行 HunyuanVideo-1.5 视频生成模型。

关键特性：

• 支持 TPU v6e-8（8 chips, 256GB HBM）
• JAX + torchax 实现 TPU 原生推理
• Splash Attention 优化
• 张量并行权重分片 + DeepCache 加速

性能数据（121帧 720p, 50步）：

文档：

• 项目 README
• TPU 推理（原生版本）
• TPU 推理（Diffusers 版本）
• GPU→TPU 迁移指南
• GPU 推理参考

---

2. CogVideoX TPU 加速

在 TPU 上运行 CogVideoX 视频生成模型，JAX/Flax 原生实现。

关键特性：

• Splash Attention TPU 优化
• Flax VAE 解码器（解决 OOM）
• FSDP/Tensor Parallel 模型分片
• BFloat16 全流程优化

性能：

• 第一次运行（含编译）：~45 秒
• 后续运行：~18 秒
• 加速比：2.44x

文档：

• CogVideoX README

---

3. CogVideoX VAE PyTorch→JAX 迁移 📚

完整的 PyTorch 到 JAX/Flax 迁移方法论，基于 2,013 行 VAE 代码的实战经验。

核心价值：

• 1,150+ 行迁移圣经文档
• 17 个单元测试全部通过
• 数值精度 MAE < 0.6
• JIT 加速 112x

关键教训：

1. 数据格式陷阱：Channel-First vs Channel-Last
2. GroupNorm 必须在 channel-first 计算才能匹配 PyTorch
3. JIT 是性能关键：不仅快 100x+，还能解决 OOM
4. 时序模型特殊性：CausalConv 不能时间分片

文档：

• 项目 README
• PyTorch→JAX 迁移圣经
• 范式转换指南

---

4. PyTorch→JAX 入门教程

从零开始学习如何在 TPU 上运行 HuggingFace 模型。

教程内容：

1. 在 JAX 中运行 HuggingFace 模型
2. 分布式运行 HuggingFace 模型
3. 进阶：使用 torchax
4. 完整示例代码

---

5. TPU 图像处理示例

展示 GPU→TPU 图像处理迁移，包含 crop/resize/blur 等操作。

关键发现：

• TPU tracing 首次较慢（0.58秒），后续快 20x（0.028秒）
• float32 精度与 GPU 完全一致
• bfloat16 精度下降 218 倍，需谨慎使用

文档：

• 图像处理 README

---

GPU

1. DeepEP on GKE B200

在 Google Kubernetes Engine (GKE) 上部署 DeepSeek 的 DeepEP 框架。

关键特性：

• NVIDIA B200 GPU（8x/node）
• RDMA 网络配置
• 节点内/节点间测试

技术栈：

• DOCA OFED v3.0.0
• NVIDIA Driver 575
• CUDA Toolkit 12.9
• NVSHMEM 3.2.5
• PyTorch (CUDA 12.9)

文档：

• DeepEP README

---

2. HunyuanVideo-1.5 GPU 推理

在 NVIDIA H100 8卡上运行 HunyuanVideo-1.5 视频生成。

关键特性：

• Flash Attention 2/3、SageAttention、Sparse Attention
• Sequence Parallelism 多卡并行
• DeepCache 加速（1.83x）

Attention 性能对比：

文档：

• GPU 推理 README

---

贡献

欢迎提交 Issues 和 Pull Requests！

许可

本项目采用开源许可。详见各子项目的许可声明。