企业官网建设流程全解析

PyTorch 代码迁移到 Lightning Fabric 完全指南：五步将原生训练脚本扩展到任意规模

【免费下载链接】pytorch-lightningPretrain, finetune ANY AI model of ANY size on 1 or 10,000+ GPUs with zero code changes.项目地址: https://gitcode.com/gh_mirrors/py/pytorch-lightning

导读

本文基于 PyTorch Lightning 仓库中 convert.rst 文档，系统讲解如何把一段普通的 PyTorch 训练脚本改造为使用 Lightning Fabric 的版本。Fabric 是 PyTorch Lightning 提供的轻量级加速层，它不引入 Trainer 或 LightningModule 这类重型抽象，只对原生 PyTorch 的训练循环做最小侵入式包装。读完本文，你将掌握五步迁移法的每个细节、底层Fabric类的工作原理、可选的进阶改造点，以及如何用一条命令行在单卡、多卡乃至多节点上运行同一份脚本。

为什么选择 Fabric：不改训练循环，只改样板代码

在开始迁移之前，先明确 Fabric 的设计定位：它不接管你的训练循环。for循环、优化器调度、评估逻辑全部由你自己编写，Fabric 只负责三件“脏活”：

• 设备管理：自动把模型和数据放到正确的设备（CPU/GPU/TPU）上；
• 分布式初始化：自动拉起多进程、配置进程组，无需手动调用torch.distributed系列 API；
• 精度与策略切换：通过一个标志即可在单卡、DDP、FSDP、DeepSpeed 等策略之间切换。

从 fabric.py 类文档 可以看到，Fabric 的核心特性正是“自动放置模型与数据到设备”“自动支持混合精度与双精度”“无缝切换硬件与分布式策略”“自动派生进程”以及“多节点支持”。

迁移的成本极低——官方文档给出的是五步改动，且全部是“删除样板代码”和“替换方法调用”，不涉及模型结构或训练逻辑的重写。

五步迁移法：核心改动一览

官方文档把迁移过程浓缩为以下五步：

Step 1：在训练代码的开头创建Fabric对象：

from lightning.fabric import Fabric fabric = Fabric()

Step 2：如果打算使用多个设备（例如多 GPU），调用fabric.launch()：

fabric.launch()

注意：launch()的作用是初始化分布式执行所需的所有进程。从源码看，launch()内部调用_wrap_and_launch，再通过策略的 launcher 来派生子进程（见 fabric.py）。如果你是通过fabric run ...命令行启动脚本，则必须删除代码里的launch()，否则会抛出RuntimeError。

Step 3：对每个「模型 + 优化器」组合调用fabric.setup()，对所有 DataLoader 调用fabric.setup_dataloaders()：

model, optimizer = fabric.setup(model, optimizer) dataloader = fabric.setup_dataloaders(dataloader)

Step 4：删除所有.to和.cuda调用，设备迁移交给 Fabric：

- model.to(device) - batch.to(device)

Step 5：用fabric.backward(loss)替换loss.backward()：

- loss.backward() + fabric.backward(loss)

以上五步就是让脚本具备 Fabric 加速能力所需的全部代码改动。

完整前后对比：一个 Transformer 语言模型示例

官方文档用一段基于lightning.pytorch.demos中的WikiText2数据集和Transformer模型的代码，展示了迁移前后的完整差异：

import torch from lightning.pytorch.demos import WikiText2, Transformer + import lightning as L - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + fabric = L.Fabric(accelerator="cuda", devices=8, strategy="ddp") + fabric.launch() dataset = WikiText2() dataloader = torch.utils.data.DataLoader(dataset) model = Transformer(vocab_size=dataset.vocab_size) optimizer = torch.optim.SGD(model.parameters(), lr=0.1) - model = model.to(device) + model, optimizer = fabric.setup(model, optimizer) + dataloader = fabric.setup_dataloaders(dataloader) model.train() for epoch in range(20): for batch in dataloader: input, target = batch - input, target = input.to(device), target.to(device) optimizer.zero_grad() output = model(input, target) loss = torch.nn.functional.nll_loss(output, target.view(-1)) - loss.backward() + fabric.backward(loss) optimizer.step()

对比可以清楚看到迁移的收益：手写device判断、.to(device)、.backward()这些与训练逻辑无关的代码全部消失，取而代之的是三个等价的 Fabric 调用。原来只能跑单卡 GPU（或回退 CPU）的脚本，现在通过accelerator="cuda", devices=8, strategy="ddp"即可在 8 张 GPU 上做数据并行训练。

迁移完成后，可以直接从终端运行：

python path/to/your/script.py

深入源码：迁移背后发生了什么

Fabric 构造参数

Fabric的构造函数（见 fabric.py）接受以下关键参数：

构造时Fabric会通过_Connector（见 connector.py）完成策略选择、加速器初始化、精度解析等配置工作。值得注意的是，如果检测到脚本是通过 CLI 启动的，构造阶段就会直接执行setup_environment()初始化进程环境（见 fabric.py）。

setup：模型与优化器的统一包装

fabric.setup()（见 fabric.py）的内部流程是：

1. 先做参数校验，并处理torch.compile包装器（若模型被编译过，会先解包、设置完成后再重新应用）；
2. 通过精度插件转换模型（如precision.convert_module）；
3. 若move_to_device=True（默认），把模型移动到目标设备；
4. 交给策略的setup_module_and_optimizers()或setup_module()完成包装——在 DDP 策略下模型会被包成DistributedDataParallel，在 FSDP 下则被分片；
5. 最后把模型包装成_FabricModule、优化器包装成_FabricOptimizer（见 wrappers.py），返回给用户。

值得注意的细节：

• setup支持一次传入多个优化器和可选的调度器：model, opt1, opt2, scheduler = fabric.setup(model, opt1, opt2, scheduler=scheduler)；
• 对于 FSDP 这类需要“先建模型、后建优化器”的策略，可以拆分为fabric.setup_module(model)和fabric.setup_optimizers(optimizer)两步（见 fabric.py）；
• DeepSpeed 与 XLA 策略要求模型与优化器联合设置，setup_optimizers对这两种策略会直接抛错。

setup_dataloaders：分布式采样的自动化

fabric.setup_dataloaders()（见 fabric.py）做了三件事：

1. 自动替换/包装分布式采样器：当策略需要分布式采样器（如 DDP）时，会用DistributedSampler重新实例化 DataLoader（use_distributed_sampler=True，默认开启）；
2. 注入 worker 初始化函数：保证数据加载 worker 进程的随机种子与全局种子一致；
3. 包装为_FabricDataLoader：让每次迭代产出的 batch 自动被移动到正确设备（move_to_device=True时）。

因此迁移后训练循环里不需要再写input, target = input.to(device), target.to(device)。

backward：精度与策略感知的反向传播

fabric.backward()（见 fabric.py）并非简单地转发给loss.backward()，而是把反向传播交给当前策略和精度插件处理：

• 在混合精度（AMP）下，它会进入精度插件的backward逻辑，配合GradScaler工作；
• 在 DeepSpeed 策略下，它把损失传给DeepSpeedEngine执行反向，并且当设置了多个模型时要求显式传入model参数；
• 普通场景下等价于loss.backward()，但在分布式策略下由统一入口保证一致性。

命令行启动：fabric run与python两种方式

官方文档展示的迁移结果是直接运行python path/to/your/script.py——此时加速器、设备数、策略等由代码中的Fabric(...)参数决定。

仓库同时提供了更灵活的 CLI 启动方式fabric run。其实现位于 cli.py，支持以下核心选项：

fabric run train_fabric.py fabric run train_fabric.py --accelerator=gpu fabric run train_fabric.py --accelerator=gpu --devices=4 fabric run train_fabric.py --strategy=deepspeed --devices=8 --num-nodes=2

两种方式的取舍：代码内传参适合固定配置的脚本；fabric run把运行参数与代码解耦，并且会在构造Fabric时自动初始化进程环境（这正是前文提到的“CLI 模式下不可再调用launch()”的原因）。

实战对照：MNIST 图像分类器的迁移前后

仓库的 image_classifier 示例 提供了同一模型的两份实现，非常适合对照学习：

• train_torch.py：纯 PyTorch 版本，手动判断 CUDA/MPS 设备，手动.to(device)，仅支持单 GPU；
• train_fabric.py：Fabric 版本，仅改动十余行。

用仓库 README 中推荐的方式可以直接查看两文件的逐行差异：

sdiff train_torch.py train_fabric.py

Fabric 版本中值得注意的几处进阶用法：

fabric = Fabric() # 用 rank_zero_first 让 0 号进程先下载数据，其余进程等待 with fabric.rank_zero_first(local=False): train_dataset = MNIST(DATASETS_PATH, download=fabric.is_global_zero, train=True, transform=transform) # 数据加载器一次设置两个 train_loader, test_loader = fabric.setup_dataloaders(train_loader, test_loader) # 模型不再调用 .to(device) model = Net() model, optimizer = fabric.setup(model, optimizer) # 用 fabric.backward 代替 loss.backward() fabric.backward(loss) # 分布式评估：跨进程聚合 loss test_loss = fabric.all_gather(test_loss).sum() / len(test_loader.dataset) # 分布式保存：确保由正确的进程写文件 fabric.save(path="mnist_cnn.pt", state=model.state_dict())

启动方式也随 CLI 而变得灵活：

# CPU fabric run train_fabric.py # GPU（CUDA 或 M1 Mac） fabric run train_fabric.py --accelerator=gpu # 多 GPU fabric run train_fabric.py --accelerator=gpu --devices=4

可选进阶改造：四类锦上添花的替换

完成五步迁移后，官方文档还给出了四类可选的代码升级建议，它们能让脚本在分布式场景下更健壮、更省显存：

1. 用 Fabric 的保存/加载方法替换torch.save/torch.load

torch.save/torch.load在分布式环境下由谁写、写什么、如何合并分片等问题都需要手工处理。Fabric 的save/load（见 fabric.py）把“哪个进程保存、如何保存”的决定权交给策略：例如 DDP 策略只在 0 号进程写文件，FSDP 策略则从每个 rank 保存分片。同时支持通过filter参数在保存时过滤指定参数（如不保存 bias）。

state = {"model": model, "optimizer": optimizer, "epoch": epoch} fabric.save("checkpoint.pth", state)

详细说明见 checkpoint 指南。

2. 用 Fabric 的集合通信方法替换torch.distributed原语

fabric.barrier()、fabric.broadcast()、fabric.all_gather()、fabric.all_reduce()等封装了torch.distributed的 barrier、broadcast、all_gather、all_reduce 操作（见 fabric.py），好处是：

• 单设备策略下自动退化为 no-op，脚本无需为“是否分布式”写分支；
• 支持传入嵌套的 tensor 集合（dict/list/tuple）并递归处理；
• all_reduce默认执行 mean 归约，也可指定sum。

详细说明见 分布式通信指南。

3. 梯度累积时使用no_backward_sync()

如果你在训练循环里实现了梯度累积，多卡场景下默认每次backward都会触发梯度同步通信，白白增加开销。fabric.no_backward_sync(model)上下文管理器可以在累积期间跳过梯度同步：

for batch_idx, batch in enumerate(dataloader): with fabric.no_backward_sync(model, enabled=(batch_idx % 8 != 0)): output = model(batch) loss = criterion(output, target) fabric.backward(loss) if batch_idx % 8 == 0: optimizer.step() optimizer.zero_grad()

其实现位于 fabric.py，对单设备策略自动退化为空操作。详细说明见 梯度累积指南。

4. 在init_module()上下文中初始化模型

大模型在 CPU 上完整初始化会带来不必要的峰值内存占用。fabric.init_module()（见 fabric.py）让模型参数直接在目标设备上以正确的数据类型创建，从而降低峰值显存；加载大模型检查点时可设置empty_init=True用未初始化内存创建参数。详细说明见 模型初始化指南。

总结与下一步

五步迁移法的本质，是把原生 PyTorch 脚本中与训练逻辑无关的设备样板代码（.to(device)、device判断、分布式原语、保存/加载细节）全部委托给 Fabric 的统一入口，从而换来“改一个标志位即可切换硬件与策略”的能力。迁移后的脚本既能python直接运行，也能用fabric run从命令行注入加速器、设备数与策略参数。

如果希望继续深入，推荐按以下顺序阅读仓库中的相关文档：

• Fabric 示例总览：覆盖计算机视觉、NLP、强化学习等领域的完整示例；
• 加速器指南：了解各硬件上的运行细节；
• 自建 Trainer 指南：基于 Fabric 搭建完全定制化的训练器。

【免费下载链接】pytorch-lightningPretrain, finetune ANY AI model of ANY size on 1 or 10,000+ GPUs with zero code changes.项目地址: https://gitcode.com/gh_mirrors/py/pytorch-lightning