PyTorch

一句话定位

PyTorch 是提供张量计算、自动求导、神经网络组件及设备执行能力的机器学习框架。

它不负责什么

  • 不直接替代集群编排系统。
  • 不等同于面向在线请求的 LLM 推理调度器。
  • 不自动决定具体项目的分布式与服务架构。

在 LLM 系统中的因果位置

1
2
3
4
5
Python model code
→ Tensor / autograd graph
→ compiler or eager dispatcher
→ CUDA/CPU kernel
→ device memory 与同步

训练时 PyTorch 负责 forward、backward 和 optimizer 接口，分布式训练再叠加 DDP/FSDP/DeepSpeed；推理时它是算子执行底座，不等于 vLLM/SGLang 这样的请求级调度器。性能问题要先区分 Python/graph 开销、kernel 算术吞吐、内存带宽和 collective 通信。

最小验证路径

  • 固定 dtype、设备、batch/sequence，比较 eager 与 torch.compile 后的数值和吞吐。
  • 用 profiler 看 kernel、同步点、显存峰值，不用单个 GPU utilization 数字下结论。
  • 分布式时记录 rank、通信后端、all-reduce 时间和数据 sampler，避免把框架启动问题误判为模型问题。

相关节点：Training Memory Accounting、Triton Language。

PyTorch 性能问题要沿 dispatch 链拆开

一个 torch 调用不是一个固定成本的黑盒。Python 侧先构造 tensor 和 graph，dispatcher 再按 device、dtype、layout 和 backend 选择 kernel；某些组合会落到 fused/FlashAttention 路径，另一些会落到通用实现。torch.compile 还会增加 graph capture、guard 和首次编译成本。因而“启用 compile 后变慢”可能只是把编译时间算进了首个请求，也可能是动态 shape 导致反复 recompilation，不能只看平均 wall time。

训练和推理的显存账本也不同。训练要保存反向所需 activation、gradient 和 optimizer state；推理通常关闭 autograd，但长上下文服务会把显存转移到 KV Cache。线上推理应明确使用 inference_mode/等价路径，并核对是否有意外的 requires_grad、CPU↔GPU 拷贝或同步点；否则会把训练期的内存行为误带进 serving。

从 profiler 读出因果而不是截图

1
2
3
4
5
端到端 wall time
→ Python/queue/tokenize
→ kernel duration 与 launch gap
→ HBM 读写、同步、collective
→ peak memory 与 allocator 状态

如果 kernel 时间短但 launch gap 长，先看 Python 循环、stream 同步和小算子融合；如果 Tensor Core 不满而 HBM 接近上限，优先查 layout、dtype 和 memory-bound；多卡时把 NCCL 等待单独计时。改用 Triton 或 CUDA 只能作用于已定位的热点 op，不能解决请求级 batching、KV 分配或网络路由。

最小报告应保存 PyTorch/CUDA/driver 版本、硬件、dtype、shape、warmup、编译时间、峰值显存、kernel trace 和数值误差。没有这些上下文，“PyTorch 很慢”不是可复现的问题描述。