AI 棱镜站

主题

PyTorch 深度学习参考资料 ​

官方文档 ​

核心机制 ​

张量系统 ​

  • Tensor: Storage + Stride + Shape 抽象
  • Autograd: 动态计算图, 反向模式自动微分
  • torch.compile: Dynamo + Inductor 两阶段编译
  • CUDA 张量: 设备感知, 异步执行

自动微分 ​

  • Autograd 引擎: 反向传播调度, 梯度累积
  • 自定义 autograd.Function: forward/backward 手动定义
  • torch.func: 函数式变换 (vmap, grad, jvp, vjp)
  • torch.autograd.functional: Hessian, Jacobian 计算

分布式训练 ​

  • torch.distributed: ProcessGroup, Store, Backend
  • DDP: DistributedDataParallel — 梯度 AllReduce
  • FSDP: FullyShardedDataParallel — 参数分片
  • torch.distributed.elastic: 弹性训练, 容错
  • NCCL: GPU 通信后端
  • Gloo: CPU 通信后端

模型编译 ​

  • TorchDynamo: Python 字节码捕获, 图构建
  • TorchInductor: Triton 后端代码生成
  • AOTAutograd: Ahead-of-time autograph 追踪
  • torch.export: 模型导出与部署

关键扩展 ​

  • torchvision: 视觉模型与数据增强
  • torchaudio: 音频 I/O 与变换
  • torchtext: 文本处理
  • torchserve: 模型服务部署
  • PyTorch Lightning: 高级训练框架
  • Hugging Face Accelerate: 分布式训练抽象层

性能优化 ​

  • FlashAttention: Dao et al., 2022 — IO-aware attention
  • Fused Adam: 融合内核优化器
  • AMP: 自动混合精度 (fp16/bf16)
  • Gradient Checkpointing: 以计算换显存
  • torch.profiler: 性能分析工具
  • CUDA Graphs: 减少内核启动开销

关键论文 ​

  • "PyTorch: An Imperative Style, High-Performance Deep Learning Library" (Paszke et al., 2019)
  • "FlashAttention: Fast and Memory-Efficient Exact Attention" (Dao et al., 2022)
  • "TorchDynamo: A Python JIT Compiler for Dynamic Shapes" (2023)
最近更新