PyTorch Lightning - 深度学习训练框架

技能概述

PyTorch Lightning 是一个用于组织 PyTorch 代码的深度学习框架，通过消除样板代码和实现训练自动化，让你专注于模型构建而非工程实现，同时支持从单机到大规模分布式训练的无缝扩展。

适用场景

1. 构建可扩展的深度学习项目

当你需要将 PyTorch 代码组织成清晰、可维护的结构时，PyTorch Lightning 提供 LightningModule 和 LightningDataModule 来分离模型逻辑、数据处理和训练流程，特别适合中大型深度学习项目的专业开发。

2. 多GPU和分布式训练

需要在不同硬件配置下训练模型时，Lightning 统一了多GPU、多节点、TPU 训练的配置方式。通过简单的策略选择（DDP、FSDP、DeepSpeed），即可从笔记本开发无缝扩展到集群环境，无需修改核心代码。

3. 实验跟踪和生产部署

在需要系统化记录训练指标、管理模型检查点、集成实验追踪平台（W&B、MLflow、TensorBoard）时，Lightning 提供标准化的日志记录和回调机制，确保实验可重复、可对比、可部署。

核心功能

1. LightningModule - 模型代码组织

将 PyTorch 模型的训练、验证、测试和预测逻辑结构化到统一模块中。通过定义六个核心方法（training_step、validation_step、test_step、predict_step、configure_optimizers 和初始化），消除手动编写训练循环的繁琐代码。设备管理、梯度操作、混合精度训练全部由框架自动处理，让代码更简洁、更易调试。

2. Trainer - 训练流程自动化

一个 Trainer 对象替代数百行训练循环代码，自动处理设备选择、分布式策略、梯度累积、早停、检查点保存、进度显示等通用训练任务。支持从快速开发测试（fast_dev_run=True）到生产级大规模训练的完整工作流，配置选项丰富但使用简单。

3. 分布式训练策略支持

内置主流分布式训练后端：DDP 适合中小型模型（<500M 参数），FSDP 优化大型模型（500M+ 参数）的内存效率，DeepSpeed 提供最前沿的分布式优化技术。策略切换仅需修改一行配置，无需重构代码，支持从单卡到数千 GPU 的平滑扩展。

常见问题

PyTorch Lightning 和原生 PyTorch 有什么区别？

PyTorch Lightning 建立在 PyTorch 之上，本质上仍然是 PyTorch，但通过结构化封装消除了大量样板代码。原生 PyTorch 需要手动编写训练循环、设备管理、指标记录等，而 Lightning 将这些标准化到 LightningModule 和 Trainer 中。你保留对模型的完全控制，但获得自动化设备处理、分布式支持、实验追踪等工程能力。迁移成本很低，核心算法逻辑几乎不需要修改。

如何选择 DDP、FSDP 和 DeepSpeed 三种训练策略？

根据模型大小和硬件资源选择：DDP（分布式数据并行）适合中小型模型如 ResNet、小型 Transformer，实现简单、稳定性高；FSDP（完全分片数据并行）适合 500M 参数以上的大型模型，通过模型分片显著降低显存占用，是 Lightning 推荐的大规模训练方案；DeepSpeed 提供更细粒度的优化（ZeRO 优化器、梯度压缩等），适合需要极致性能优化或特殊需求的场景。可以从 DDP 开始，根据显存和性能需求逐步升级。

使用 PyTorch Lightning 需要什么环境配置？

基础要求是 Python 3.10+ 和 PyTorch（支持 CPU 或 GPU）。GPU 训练需要 CUDA 兼容的 PyTorch 版本。安装方式简单：uv pip install lightning 即可获得核心功能。如需使用特定日志工具（W&B、MLflow）或高级分布式策略，可额外安装对应包。Lightning 2.6+ 是当前推荐版本，文档齐全，社区活跃。开发时推荐使用 fast_dev_run=True 快速验证代码正确性，生产环境则可配置确定性训练（deterministic=True）确保结果可重复。