谷歌狂喜：JAX性能超越Pytorch、TensorFlow！或成GPU推理训练最快选择

新智元报道

编辑：编辑部

【新智元导读】JAX在最近的基准测试中的性能已经不声不响地超过了Pytorch和TensorFlow，也许未来会有更多的大模型诞生在这个平台上。谷歌在背后的默默付出终于得到了回报。

谷歌力推的JAX在最近的基准测试中性能已经超过Pytorch和TensorFlow，7项指标排名第一。

而且测试并不是在JAX性能表现最好的TPU上完成的。

虽然现在在开发者中，Pytorch依然比Tensorflow更受欢迎。

但未来，也许有更多的大模型会基于JAX平台进行训练和运行。

模型

最近，Keras团队为三个后端（TensorFlow、JAX、PyTorch）与原生PyTorch实现以及搭配TensorFlow的Keras 2进行了基准测试。

首先，他们为生成式和非生成 式人工智 能任务选择了一组主流的计算机视觉和自然语言处理模型：

对于模型的Keras版本，其采用了KerasCV和KerasNLP中已有的实现进行构建。而对于原生的PyTorch版本，则选择了网络上最流行的几个选项：

- 来自HuggingFace Transformers的BERT、Gemma、Mistral

- 来自HuggingFace Diffusers的StableDiffusion

- 来自Meta的SegmentAnything

他们将这组模型称作「Native PyTorch」，以便与使用PyTorch后端的Keras 3版本进行区分。

他们对所有基准测试都使用了合成数据，并在所有LLM训练和推理中使用了bfloat16精度，同时在所有LLM训练中使用了LoRA（微调）。

根据PyTorch团队的建议，他们在原生PyTorch实现中使用了torch.compile(model, mode="reduce-overhead")（由于不兼容，Gemma和Mistral训练除外）。

为了衡量开箱即用的性能，他们使用高级API（例如HuggingFace的Trainer()、标准PyTorch训练循环和Keras model.fit()），并尽可能减少配置。

硬件配置

所有基准测试均使用Google Cloud Compute Engine进行，配置为：一块拥有40GB显存的NVIDIA A100 GPU、12个虚拟CPU和85GB的主机内存。

基准测试结果

表2显示了基准测试结果（以步/毫秒为单位）。每步都涉及对单个数据批次进行训练或预测。

结果是100步的平均值，但排除了第一个步，因为第一步包括了模型创建和编译，这会额外花费时间。

为了确保比较的公平性，对于相同的模型和任务（不论是训练还是推理）都使用相同的批大小。

然而，对于不同的模型和任务，由于它们的规模和架构有所不同，可根据需要调整数据批大小，从而避免因过大而导致内存溢出，或是批过小而导致GPU使用不足。

过小的批大小也会使PyTorch看起来较慢，因为会增加Python的开销。

对于大型语言模型（Gemma和Mistral），测试时也使用了相同的批处理大小，因为它们是相同类型的模型，具有类似数量的参数（7B）。

考虑到用户对单批文本生成的需求，也对批大小为1的文本生成情况进行了基准测试。

关键发现

发现1

不存在「最优」后端。

Keras的三种后端各展所长，重要的是，就性能而言，并没有哪一个后端能够始终胜出。

选择哪个后端最快，往往取决于模型的架构。

这一点突出了选择不同框架以追求最佳性能的重要性。Keras 3可以帮助轻松切换后端，以便为模型找到最合适的选择。

发现2

Keras 3的性能普遍超过PyTorch的标准实现。

相对于原生PyTorch，Keras 3在吞吐量（步/毫秒）上有明显的提升。

特别是，在10个测试任务中，有5个的速度提升超过了50%。其中，最高更是达到了290%。

如果是100%，意味着Keras 3的速度是PyTorch的2倍；如果是0%，则表示两者性能相当

发现3

Keras 3提供一流的「开箱即用」性能。

也就是，所有参与测试的Keras模型都未进行过任何优化。相比之下，使用原生PyTorch实现时，通常需要用户自行进行更多性能优化。

除了上面分享的数据，测试中还注意到在HuggingFace Diffusers的StableDiffusion推理功能上，从版本0.25.0升级到0.3.0时，性能提升超过了100%。

同样，在HuggingFace Transformers中，Gemma从4.38.1版本升级至4.38.2版本也显著提高了性能。

这些性能的提升凸显了HuggingFace在性能优化方面的专注和努力。

对于一些手动优化较少的模型，如SegmentAnything，则使用了研究作者提供的实现。在这种情况下，与Keras相比，性能差距比大多数其他模型更大。

这表明，Keras能够提供卓越的开箱即用性能，用户无需深入了解所有优化技巧即可享受到快速的模型运行速度。

发现4

Keras 3的表现始终优于Keras 2。

例如，SegmentAnything的推理速度提升了惊人的380%，StableDiffusion的训练处理速度提升了150%以上，BERT的训练处理速度也提升了100%以上。

这主要是因为Keras 2在某些情况下直接使用了更多的TensorFlow融合操作，而这可能对于XLA的编译并不是最佳选择。

值得注意的是，即使仅升级到Keras 3并继续使用TensorFlow后端，也能显著提升性能。

结论

框架的性能在很大程度上取决于具体使用的模型。

Keras 3能够帮助为任务选择最快的框架，这种选择几乎总能超越Keras 2和PyTorch实现。

更为重要的是，Keras 3模型无需进行复杂的底层优化，即可提供卓越的开箱即用性能。

参考资料：

https://keras.io/getting_started/benchmarks/