十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

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

PyTorch 代码迁移到 Lightning Fabric 完全指南:五步将原生训练脚本扩展到任意规模 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(acceleratorcuda, devices8, strategyddp) fabric.launch() dataset WikiText2() dataloader torch.utils.data.DataLoader(dataset) model Transformer(vocab_sizedataset.vocab_size) optimizer torch.optim.SGD(model.parameters(), lr0.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的脚本现在通过acceleratorcuda, devices8, strategyddp即可在 8 张 GPU 上做数据并行训练。迁移完成后可以直接从终端运行python path/to/your/script.py深入源码迁移背后发生了什么Fabric 构造参数Fabric的构造函数见 fabric.py接受以下关键参数参数取值默认值说明acceleratorcpu、cuda、mps、gpu、tpu、autoauto运行硬件gpu会自动在 CUDA 与 MPS 之间选择strategydp、ddp、ddp_spawn、deepspeed、fsdp、autoauto多设备运行策略devicesint、list、str、autoauto每节点设备数量或指定 GPU 编号num_nodesint1分布式训练的节点数precision64、32、16-mixed、bf16-mixed等None精度设置None时按设备取默认plugins单个或列表None自定义插件callbacks单个或列表None回调对象loggers单个或列表None日志记录器构造时Fabric会通过_Connector见 connector.py完成策略选择、加速器初始化、精度解析等配置工作。值得注意的是如果检测到脚本是通过 CLI 启动的构造阶段就会直接执行setup_environment()初始化进程环境见 fabric.py。setup模型与优化器的统一包装fabric.setup()见 fabric.py的内部流程是先做参数校验并处理torch.compile包装器若模型被编译过会先解包、设置完成后再重新应用通过精度插件转换模型如precision.convert_module若move_to_deviceTrue默认把模型移动到目标设备交给策略的setup_module_and_optimizers()或setup_module()完成包装——在 DDP 策略下模型会被包成DistributedDataParallel在 FSDP 下则被分片最后把模型包装成_FabricModule、优化器包装成_FabricOptimizer见 wrappers.py返回给用户。值得注意的细节setup支持一次传入多个优化器和可选的调度器model, opt1, opt2, scheduler fabric.setup(model, opt1, opt2, schedulerscheduler)对于 FSDP 这类需要“先建模型、后建优化器”的策略可以拆分为fabric.setup_module(model)和fabric.setup_optimizers(optimizer)两步见 fabric.pyDeepSpeed 与 XLA 策略要求模型与优化器联合设置setup_optimizers对这两种策略会直接抛错。setup_dataloaders分布式采样的自动化fabric.setup_dataloaders()见 fabric.py做了三件事自动替换/包装分布式采样器当策略需要分布式采样器如 DDP时会用DistributedSampler重新实例化 DataLoaderuse_distributed_samplerTrue默认开启注入 worker 初始化函数保证数据加载 worker 进程的随机种子与全局种子一致包装为_FabricDataLoader让每次迭代产出的 batch 自动被移动到正确设备move_to_deviceTrue时。因此迁移后训练循环里不需要再写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 --acceleratorgpu fabric run train_fabric.py --acceleratorgpu --devices4 fabric run train_fabric.py --strategydeepspeed --devices8 --num-nodes2选项说明--accelerator硬件加速器可选 CPU/CUDA/GPU/TPU 等--strategy多设备策略自动排除需要 spawn/fork 或需额外配置的策略见 cli.py--devices每节点设备数默认1--num-nodes/--num_nodes节点数默认1--node-rank/--node_rank当前节点在集群中的序号两种方式的取舍代码内传参适合固定配置的脚本fabric run把运行参数与代码解耦并且会在构造Fabric时自动初始化进程环境这正是前文提到的“CLI 模式下不可再调用launch()”的原因。实战对照MNIST 图像分类器的迁移前后仓库的 image_classifier 示例 提供了同一模型的两份实现非常适合对照学习train_torch.py纯 PyTorch 版本手动判断 CUDA/MPS 设备手动.to(device)仅支持单 GPUtrain_fabric.pyFabric 版本仅改动十余行。用仓库 README 中推荐的方式可以直接查看两文件的逐行差异sdiff train_torch.py train_fabric.pyFabric 版本中值得注意的几处进阶用法fabric Fabric() # 用 rank_zero_first 让 0 号进程先下载数据其余进程等待 with fabric.rank_zero_first(localFalse): train_dataset MNIST(DATASETS_PATH, downloadfabric.is_global_zero, trainTrue, transformtransform) # 数据加载器一次设置两个 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(pathmnist_cnn.pt, statemodel.state_dict())启动方式也随 CLI 而变得灵活# CPU fabric run train_fabric.py # GPUCUDA 或 M1 Mac fabric run train_fabric.py --acceleratorgpu # 多 GPU fabric run train_fabric.py --acceleratorgpu --devices4可选进阶改造四类锦上添花的替换完成五步迁移后官方文档还给出了四类可选的代码升级建议它们能让脚本在分布式场景下更健壮、更省显存1. 用 Fabric 的保存/加载方法替换torch.save/torch.loadtorch.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_initTrue用未初始化内存创建参数。详细说明见 模型初始化指南。总结与下一步五步迁移法的本质是把原生 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创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表