原贴链接

大家好，我们正在为非NVIDIA GPU构建AI堆栈。我和我的联合创始人过去5年在Google和Meta的ML基础设施团队工作，我们正在利用这些经验为TPU、TRN和AMD GPU等芯片组构建LLM调优和服务堆栈。

我们从Google TPU开始，并为其构建了一个类似runpod的UI。为什么？在大规模云上进行AI训练的开发工作流程是破碎的。你只需要一个安装了PyTorch/JAX的加速器VM，连接到存储以加载数据和写入训练日志。但大型云服务使其变得不必要的复杂。

我们的UI层位于app.felafax.ai。你可以启动任意大小的TPU VM，从8个芯片到1024个芯片。我们还提供了常见用例作为模板，如LLaMA 3.1和Gemma微调。该pod预装了依赖项，并提供了一个笔记本供你运行微调。

在TPU上进行LLaMA 3.1微调比我们最初想象的要复杂得多！我们首先尝试了PyTorch XLA路线。虽然这看起来像是直接的选择（LLaMA 3是基于PyTorch的，HuggingFace库也是基于PyTorch的），但事实并非如此。PyTorch与XLA的集成在LazyTensors上很笨拙。存在很大的裂缝——Bitsandbytes在XLA上不起作用，甚至在许多情况下HuggingFace库会抛出奇怪的错误。

在与PyTorch斗争之后，我们将LLaMA 3.1转换为JAX。这在TPU上运行得更好，但我们必须构建许多支持库——LoRA、量化库（如bnb）等。我们才刚刚开始构建这些库，并将其视为绿色空间！

那么，我们为什么要这样做？NVIDIA的垄断不会持续，对行业也不利。还有其他芯片组，如TPU，它们便宜得多但没有人使用。关于TPU v5p的有趣事实：它带有8个芯片，每个芯片有96GB VRAM。它的性能相当于四个NVIDIA H100，但便宜5倍。

我们的请求：请访问我们的平台app.felafax.ai，体验在最新一代Google TPU上的微调。我们提供50美元的信用额度（我们仍然是一家小型初创公司:P）。你可以直接运行LLaMA 3.1微调。

请告诉我们你的想法或任何问题！

讨论总结

本次讨论主要聚焦于非 NVIDIA GPU 的 AI 基础设施建设，特别是 Google TPU 和 AMD GPU 的应用。讨论中涉及了多个关键话题，包括 TPU 和 AMD GPU 的性能与可用性、JAX 和 PyTorch XLA 的技术挑战、模型微调的资源需求、以及教程和硬件普及的需求。总体而言，讨论呈现出对非 NVIDIA GPU 技术前景的积极态度，同时也指出了当前面临的诸多挑战和改进空间。

主要观点

  1. 👍 TPU 性能优越但难以购买
    • 支持理由：TPU 在性能上具有优势且成本较低，但主要由 Google 控制，普通用户难以购买。
    • 反对声音：TPU 可以通过云服务租赁使用。
  2. 🔥 AMD GPU 更适合普通开发者
    • 正方观点：AMD GPU 的可用性更高，更适合普通开发者使用。
    • 反方观点：TPU 性能更优越，但购买和使用门槛较高。
  3. 💡 JAX 在 TPU 上的应用表现良好
    • 解释：JAX 在 TPU 上的应用表现优于 PyTorch XLA，但需要构建许多支持库。
  4. 👀 PyTorch XLA 使用体验不佳
    • 解释：PyTorch XLA 缺乏良好的 eager 模式，模型分片 API 不够清晰，许多库在使用时会抛出奇怪的错误。
  5. 🚀 期待在 AMD GPU 上进行 LLaMA 3.1 微调
    • 解释：JAX 和 XLA 在 AMD GPU 上的应用表现良好，期待未来在 AMD GPU 上进行 LLaMA 3.1 微调。

金句与有趣评论

  1. “😂 Except TPUs are all kept by google. I can’t really buy used ones like the nvidia cards.”
    • 亮点：指出了 TPU 难以购买的现状，与 NVIDIA GPU 形成鲜明对比。
  2. “🤔 I haven’t even tried the app yet, but one interesting feature would be auto calculating requirements and maybe even speed.”
    • 亮点：提出了自动计算模型微调资源需求和速度的功能设想，具有创新性。
  3. “👀 Honestly PyTorch XLA feels like a duck taped version of XLA to get Pytorch working on TPUs.”
    • 亮点：形象地描述了 PyTorch XLA 的使用体验不佳，引起了广泛共鸣。

情感分析

讨论的总体情感倾向较为积极，多数评论者对非 NVIDIA GPU 的 AI 基础设施建设表示支持，并期待未来的技术进步。主要分歧点在于 TPU 和 AMD GPU 的可用性和性能比较，以及 PyTorch XLA 的使用体验。这些分歧可能源于不同用户的技术背景和使用需求。

趋势与预测

  • 新兴话题：自动计算模型微调资源需求和速度的功能设想，可能引发后续讨论。
  • 潜在影响：随着非 NVIDIA GPU 技术的进步和普及，AI 基础设施建设的多样性和成本效益将进一步提升，可能对整个 AI 行业产生深远影响。

详细内容：

标题：打破 NVIDIA 垄断，探索非 NVIDIA GPU 的 AI 基础设施

最近，Reddit 上一则关于为非 NVIDIA GPU 构建 AI 栈的帖子引起了广泛关注。该帖获得了众多点赞和大量评论。帖子的作者称，他和联合创始人凭借在谷歌和 Meta 的 ML infra 团队积累的五年经验，正在为诸如 TPU、TRN 和 AMD GPU 等芯片组打造一个 LLM 调优及服务栈。他们以谷歌 TPUs 为起点，开发了类似 runpod 的 UI，并指出在大云平台上进行 AI 训练的开发工作流程存在诸多问题，让其变得过于复杂。

他们的 UI 层位于 app.felafax.ai，用户能在此启动各种规模的 TPU VM。他们还提供了常见用例模板，如 LLaMA 3.1 和 Gemma 微调。然而，让 LLaMA 3.1 在 TPU 上实现微调的过程比想象中复杂得多。

有人指出，TPUs 都被谷歌掌控，无法像 NVIDIA 卡那样购买二手的。也有人表示，如果能让 TPUs 更易用是很不错的，但更期待对 AMD 卡的支持。还有人提到，TPUs 虽不能购买但可在云端租用。

有人认为，如果能有自动计算需求甚至速度的有趣功能就好了。有人对 XLA 与 PyTorch 的结合提出疑问，有人详细解释称 PyTorch XLA 存在诸多痛点，如没有良好的即时模式、模型分片 API 不够清晰、很多库会报错以及参考示例少等。

有人支持该项目及背后的动机，但指出真正的问题在于配套的库，在 BNB 等库能在非 NVIDIA 硬件上运行之前，大多数开发工作难以真正摆脱 CUDA。

有人对 TPUs 感兴趣并注册了账号，但觉得界面更适合熟悉 TPUs 的人，希望能有完整的热门模型训练教程，以及在个人 GPU 上运行训练模型而不降低性能的教程。

这一讨论的核心问题在于如何打破 NVIDIA 的垄断，让非 NVIDIA GPU 更好地服务于 AI 开发，以及如何解决在这一过程中遇到的技术难题和推广使用的问题。