SpaceX 自研 AI 训练栈：适配 22 万 GPU 集群，大规模训练比 JAX 提速一个数量级以上

【导语：马斯克透露 SpaceX 团队基本完成 V1.0 版本自研 AI 训练栈开发，该系统适配 22 万块 GB300 GPU 集群，相比 JAX 大规模训练提速一个数量级以上，这一成果背后有着复杂的技术考量和工程权衡。】

适配 22 万 GPU 集群的自研训练栈

马斯克在 X 上透露，SpaceX 团队已基本完成 V1.0 版本的自研 AI 训练栈开发。这套系统以 C 语言编写，精确适配 22 万块 GB300 GPU 集群和 800G 网卡配置。GB300 是 NVIDIA 最新的 GPU 架构之一，22 万块的集群规模使算力基础设施达到超级计算机级别。

绕过高级语言抽象开销

团队选择用 C 语言而非 Python 或现有框架来实现训练栈，是为了绕过高级语言带来的抽象开销。在如此大规模的硬件上做训练栈，团队必须在通信带宽、内存访问延迟和计算密度之间做出极其精细的工程权衡。

攻克流水线并行难题

流水线并行是分布式训练中处理大模型的主流方案，其核心是将模型切分为多个阶段，每个设备负责一个阶段的计算，通过重叠通信和计算来隐藏数据传递的延迟。但在 22 万块 GPU 的规模上实现高效的流水线并行，工程难度极高，规模越大，流水线气泡的累积效应越严重，同步开销也越高。SpaceX 团队声称“尽可能接近裸金属”，暗示他们在这一层做了深度的内核级优化。

系统软件层面的优化提速

与 JAX 的比较是关注重点。JAX 是 Google 推出的高性能数值计算框架，广泛应用于 AI 训练大模型场景。SpaceX 认为自研栈相比 JAX 在大规模训练上有“一个数量级以上”的提速，这种差异并非来自算法层面，而是系统软件层面的优化，包括内存布局、计算调度、通信原语和 GPU 利用率等多个环节。

编辑观点：SpaceX 自研 AI 训练栈展现了其在底层技术栈的强大实力，从硬件适配到软件优化，都体现了“全垂直整合”的工程哲学，有望为 AI 训练带来新的突破。