3. 深度学习框架选择

说实话，选框架这件事，我见过太多人纠结了。

有人一上来就问「哪个框架最好」，也有人因为网上吵得凶就反复横跳。我个人习惯是：先搞清楚你要做什么，再选工具。别让工具定义你的问题。

3.1 PyTorch vs TensorFlow vs JAX

这三个框架，我都有过实战经验。简单说说我的感受。

PyTorch

PyTorch 现在基本是学术界和工业界的默认选择。为什么？因为它的调试体验太好了。你想想看，写 Python 代码的时候，你可以随时 print 一个张量的形状，随时打断点。这在 TensorFlow 1.x 时代简直是奢望。

我在项目中遇到过最典型的场景：一个研究员写了个新的注意力机制，用 PyTorch 半天就调通了。换成 TensorFlow 的话，光是把动态图转静态图就能卡两天。

但 PyTorch 也有短板。生产部署时，TorchScript 的兼容性偶尔会让人头疼。嗯，这里要注意：如果你要做大规模生产部署，最好提前测试 TorchScript 的导出流程。

TensorFlow

TensorFlow 2.x 之后，体验好了很多。但说实话，它的 API 设计还是有点「历史包袱」。Keras 虽然好用，但一旦你要做底层定制，就会碰到各种奇怪的限制。

我曾经在一个推荐系统项目里，需要自定义一个非常规的梯度更新逻辑。在 PyTorch 里就是重写 backward() 的事，在 TensorFlow 里折腾了两天才搞定。

不过 TensorFlow 的生态确实强。TF Serving、TF Lite、TF.js，一套下来很完整。如果你的团队有成熟的 DevOps 支持，TensorFlow 的生产部署体验其实比 PyTorch 更顺。

JAX

JAX 是个异类。它不像一个传统框架，更像是一个「可微分的 NumPy」。我第一次用 JAX 的时候，感觉像是回到了写纯 NumPy 的日子，但突然多了自动微分和 JIT 编译。

JAX 的核心优势在于函数式编程和 XLA 编译。你写的每个函数都是纯函数，没有副作用。这意味着 JAX 可以自动做很多优化，比如算子融合、内存复用。

但 JAX 的学习曲线比较陡。我建议至少要有 3 个月以上的 PyTorch 经验，再考虑上手 JAX。否则你会被它的「函数式思维」搞疯。

3.2 框架性能对比基准

很多人喜欢看 benchmark 选框架。但我要泼一盆冷水：benchmark 只能参考，不能迷信。

为什么？因为 benchmark 测试的场景往往很理想化。你的模型结构、数据加载、硬件配置，跟 benchmark 环境可能完全不同。

我整理了一份常见的性能对比维度：

维度	PyTorch	TensorFlow	JAX
单卡训练速度	中等	中等	快（XLA 编译）
多卡扩展性	好（DDP 成熟）	好（MirroredStrategy）	优秀（pmap）
内存效率	中等	中等	高（函数式无状态）
调试体验	优秀	一般	较差（编译后难调试）
部署便捷性	一般	优秀	一般

我曾经做过一个对比实验：同样的 ResNet-50，在 8 张 A100 上训练。PyTorch 的 DDP 大概能跑到 95% 的线性加速，TensorFlow 的 MirroredStrategy 大概 92%，JAX 的 pmap 能到 98%。

但注意，这个差距在更复杂的模型（比如 类）上会缩小。因为通信开销占比变大了，框架本身的优化空间就小了。

3.3 框架与硬件兼容性

选框架的时候，一定要考虑你的硬件环境。我见过太多人买了 A100，结果发现框架不支持某些算子，只能跑在 CPU 上。

目前主流硬件兼容情况：

• NVIDIA GPU： 三个框架都支持得很好。但要注意 CUDA 版本和 cuDNN 版本的匹配。我建议用 NVIDIA 官方推荐的组合，别自己瞎试。
• AMD GPU： PyTorch 有 ROCm 支持，TensorFlow 也有，但 JAX 的支持还比较有限。如果你用 AMD 卡，建议优先选 PyTorch。
• Google TPU： JAX 是亲儿子，TensorFlow 也支持，PyTorch 需要额外配置。如果你要用 TPU，JAX 是最省心的选择。
• Apple Silicon： PyTorch 的 MPS 后端已经比较成熟了，TensorFlow 的 Metal 支持也还行。JAX 在 Mac 上主要靠 CPU。

3.4 分布式训练框架概览

单卡训练已经不够用了。现在随便一个 LLM 都是几十亿参数，不搞分布式根本跑不动。

这里我重点讲三个分布式框架：Horovod、DeepSpeed、Megatron。

Horovod

Horovod 是 Uber 开源的，主打「简单易用」。它的核心思想是：你只需要在单卡代码里加几行，就能变成分布式。

import horovod.torch as hvd
hvd.init()
torch.cuda.set_device(hvd.local_rank())
optimizer = hvd.DistributedOptimizer(optimizer)
hvd.broadcast_parameters(model.state_dict(), root_rank=0)

嗯，就这么几行。但 Horovod 的问题在于：它只做数据并行，不支持模型并行。如果你的模型大到一张卡放不下，Horovod 就帮不了你了。

DeepSpeed

DeepSpeed 是微软开源的，它的 ZeRO 优化器是真正的杀手锏。ZeRO 把模型状态（优化器状态、梯度、参数）分散到多张卡上，让单卡能训练更大的模型。

我曾经用 DeepSpeed 在 8 张 V100 上训练了一个 13B 的模型。如果没有 ZeRO，这个模型至少需要 32 张卡。说白了，DeepSpeed 就是帮你省钱。

# DeepSpeed 配置示例
deepspeed_config = {
    "train_batch_size": 32,
    "gradient_accumulation_steps": 4,
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu"
        }
    }
}

Megatron

Megatron 是 NVIDIA 开源的，主打「模型并行」。它把 Transformer 的每一层切分到多张卡上，让超大模型训练成为可能。

Megatron 的核心是张量并行和流水线并行。张量并行把一层内的矩阵乘法拆开，流水线并行把不同层放到不同卡上。两者结合，可以训练千亿参数级别的模型。

但 Megatron 的代码比较底层，上手难度高。我建议至少要有 3 个月以上的分布式训练经验，再碰 Megatron。

最后说一句：框架只是工具，别在选工具上花太多时间。我见过有人为了「哪个框架更好」吵了一个月，结果模型还没开始写。先动手，遇到问题再换，这才是务实的态度。