PyTorch

进入训练数据通路:理解 WebDataset、Mosaic Streaming 与自定义 IterableDataset 的设计取舍;调优 DataLoader 的 num_workers、prefetch_factor、pin_memory、persistent_workers 与 shared memory,定位 GPU starvation。

进入 Transformer 算子层加速:理解 FlashAttention v2/v3、PyTorch SDPA、xFormers 与 Apex Fused Kernels 的适用边界,动手替换 attention,并用 TFLOPS / MFU 判断优化是否真的生效。

学习训练性能分析的最小闭环:计算 MFU / HFU,用 Nsight Systems 抓一段训练 step,通过 NVTX、CUDA kernel、NCCL timeline 识别 compute、communication 与 pipeline bubble。

进入 PyTorch 2.x 编译器栈: TorchDynamo 如何从 Python 代码捕获 FX 图, guards 如何决定复用/重编译/fallback, Inductor 如何融合算子并生成 Triton/C++ kernel;最后对一个小模型做 eager vs compile benchmark。

拆开混合精度训练:FP32 / TF32 / FP16 / BF16 / FP8 的数值范围、精度与溢出风险;理解 torch.amp.autocast 的算子选择、GradScaler 的动态 loss scaling,并跑通一个可对比的 AMP benchmark。

拆开 PyTorch 的 caching allocator:Segment / Block 两级结构、small/large pool 分裂、stream 池隔离;掌握 memory_summary / memory_snapshot 工具链,学会 OOM 排查的标准流程,并理解 expandable_segments 如何根治碎片化。

拆开 PyTorch 算子实现层:c10 是底座,ATen 是算子库,native_functions.yaml 是真理之源。从代码生成管线到 TORCH_LIBRARY 注册机制,跑通一个自定义 CUDA 算子。

用 ~200 行 Python 手写一个 mini-autograd:理解动态计算图的构建、链式法则的层层应用、拓扑排序的反向遍历。在最小可运行实现里看清 PyTorch backward() 的真面目。

打开 PyTorch 黑盒:Tensor 与 Storage 的分离、Dispatcher 的多重派发机制、Autograd Engine 的工作原理。跟踪一行 a + b 从 Python 一路调到 CUDA kernel 的完整路径。

系统学习 GPU profiling 工具链:nsys 看系统级 timeline、ncu 做 kernel 深度分析、py-spy 抓 Python 热点、torch.profiler 串起训练全链路,建立先测量再优化的工程习惯。