PyTorch Lightning Fabric 多模型与多优化器实战setup 机制、四种组合模式与源码级解析【免费下载链接】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-lightningLightning Fabric 通过统一的fabric.setup()接口让同一套训练代码在单卡、多卡 DDP、FSDP、DeepSpeed 等不同策略下无需修改即可运行。本文聚焦 Fabric 中多模型、多优化器这一高频场景GAN、自编码器、元学习等完整讲解四种模型—优化器组合模式的写法并结合仓库源码剖析setup背后的包装与校验逻辑帮助你写出真正 strategy-agnostic 的训练代码。为什么需要同时管理多个模型与优化器现代深度学习任务中一个训练流程往往不只包含一个模型和一个优化器生成对抗网络GAN生成器Generator与判别器Discriminator是两个独立的模型各自拥有独立的优化器交替更新自编码器Auto-encoder编码器与解码器共享或部分共享参数可能需要按模块施加不同的学习率元学习Meta-learning内循环、外循环各自维护模型与优化器参数更新方式完全不同。在这些场景下模型与优化器的配对关系五花八门。Fabric 给出的核心原则只有一条原文与实现完全一致只要你有一个优化器就应该把模型和优化器一起交给fabric.setup()这样才能让代码真正与所选策略无关strategy-agnostic。这条原则的源码依据在 setup 方法定义fabric.setup()接收一个模型和任意多个优化器统一完成设备搬移、精度转换、策略包装DDP/FSDP/DeepSpeed 等以及_FabricModule/_FabricOptimizer的封装。下面按四种组合模式逐一展开。模式一一个模型 一个优化器最简形态这是最常见的训练形态。模型与优化器必须成对交给setup()import torch from lightning.fabric import Fabric fabric Fabric() # Instantiate model and optimizer model LitModel() optimizer torch.optim.Adam(model.parameters()) # Set up the model and optimizer together model, optimizer fabric.setup(model, optimizer)关键点fabric.setup()返回的model是包装后的_FabricModule返回的optimizer是包装后的_FabricOptimizer必须用返回值覆盖原变量后续训练循环中的optimizer.step()、fabric.backward()才具备跨策略能力根据所选策略不同setup内部会执行不同的包装动作例如 DDPStrategy.setup_module 会把模型包装成torch.nn.parallel.distributed.DistributedDataParallelFSDPStrategy 则包装成FullyShardedDataParallel并要求use_orig_paramsTrue。setup()的完整签名摘自 fabric.pydef setup( self, module: nn.Module, *optimizers: Optimizer, scheduler: Optional[_LRScheduler] None, move_to_device: bool True, _reapply_compile: bool True, ) - Any:各参数含义参数默认值说明module必填要 setup 的torch.nn.Module每次 setup 只能传一个*optimizers空零个或多个优化器按传入顺序返回schedulerNone可选的学习率调度器必须在优化器之后传入move_to_deviceTrue是否自动把模型搬到目标设备设为False时可手动调用fabric.to_device()_reapply_compileTrue若模型此前被torch.compile过策略包装如 DDP/FSDP完成后会用相同设置重新应用编译返回值规则不传优化器时只返回包装后的模型传了优化器且带 scheduler时按(module, *optimizers, scheduler)的顺序返回元组与传入顺序一一对应。模式二一个模型 多个优化器当模型的各部分需要不同的优化器或学习率时例如不同层采用不同的lr可以给同一个模型挂多个优化器# Instantiate model and optimizers model LitModel() optimizer1 torch.optim.SGD(model.layer1.parameters(), lr0.003) optimizer2 torch.optim.SGD(model.layer2.parameters(), lr0.01) # Set up the model and optimizers together model, optimizer1, optimizer2 fabric.setup(model, optimizer1, optimizer2)使用要点多个优化器会按传入顺序原样返回逐一对应到optimizer1、optimizer2训练循环中可分别控制各优化器的更新时机例如按阶段交替step()若还配有调度器写法为model, opt1, opt2, scheduler fabric.setup(model, opt1, opt2, schedulerscheduler)该示例直接来自 setup 的 docstring。模式三多个模型 一个优化器多个模型共享一个优化器时最稳妥的写法是把所有子模型聚合到一个顶层nn.Module下让 Fabric 把它当作一个模型处理class AutoEncoder(torch.nn.Module): def __init__(self): super().__init__() # Group all models under a common nn.Module self.encoder Encoder() self.decoder Decoder()聚合之后所有子模型即可被当作单一模型进行 setup# Instantiate the big model autoencoder AutoEncoder() optimizer ... # Set up the model(s) and optimizer together autoencoder, optimizer fabric.setup(autoencoder, optimizer)这样做的原因可以从setup的实现中得到解释setup的第一个参数始终是一个nn.Module策略包装尤其是 DDP 这类分布式包装也以模块为单位进行把子模型收拢进一个容器模块后optimizer里收集到的encoder与decoder参数仍属于同一个模块树梯度同步、状态广播等分布式逻辑才能正确覆盖全部参数。如果希望各子模型分别参与 setup则应采用下面多模型多优化器的写法而不是强行共用一个优化器。模式四多个模型 多个优化器这是 GAN 类任务的标准结构——两个模型、各自独立的优化器、交替更新。Fabric 支持按模型 优化器配对分多次 setup# Two models generator Generator() discriminator Discriminator() # Two optimizers optimizer_gen torch.optim.SGD(generator.parameters(), lr0.01) optimizer_dis torch.optim.SGD(discriminator.parameters(), lr0.001) # Set up generator generator, optimizer_gen fabric.setup(generator, optimizer_gen) # Set up discriminator discriminator, optimizer_dis fabric.setup(discriminator, optimizer_dis)也可以一次 setup 一个模型加多个优化器甚至多次 setup 组合出任意数量的模型—优化器配对。仓库中提供了完整的 DCGAN 示例examples/fabric/dcgan/train_fabric.py配套说明见 README。该示例的 setup 环节与本文模式四完全对应optimizer_d optim.Adam(discriminator.parameters(), lrlr, betas(beta1, 0.999)) optimizer_g optim.Adam(generator.parameters(), lrlr, betas(beta1, 0.999)) discriminator, optimizer_d fabric.setup(discriminator, optimizer_d) generator, optimizer_g fabric.setup(generator, generator_optimizer)训练循环中判别器与生成器各自执行zero_grad→ 前向 →fabric.backward()→optimizer.step()的独立更新且所有张量real、label、noise都通过fabric.device创建保证设备无关。该示例还展示了fabric.setup_dataloaders()、fabric.print()、fabric.is_global_zero与fabric.barrier()在多进程下的配套用法。可通过sdiff train_torch.py train_fabric.py对比原生 PyTorch 与 Fabric 写法的差异见 README。setup 的底层发生了什么从源码看完整流程理解setup的内部流程有助于判断什么时候该把模型和优化器一起 setup、什么时候必须分开。根据 fabric.py 的实现setup()的执行步骤为合法性校验_validate_setup见 L1200-L1216模型不能是已包装的_FabricModule每个模型只能 setup 一次优化器同理FSDP 策略下若优化器引用了 meta device 上的参数会抛出提示改用分步 setup编译解包若模型已被torch.compile先解包OptimizedModule待策略包装完成后再按原设置重新应用_reapply_compile控制精度转换self._precision.convert_module(module)应用所选的混合精度/低精度配置设备搬移move_to_deviceTrue时把模型及其优化器参数引用搬到目标设备策略包装若传入了优化器调用self._strategy.setup_module_and_optimizers(module, optimizers, scheduler)否则只调用setup_module。默认实现见 strategy.py L150-L161会依次调用策略的setup_module如 DDP 包装与每个优化器的setup_optimizerFabric 包装模型包装为_FabricModule每个优化器包装为_FabricOptimizer并触发on_after_setup回调。校验逻辑中还包含一条重要约束DeepSpeed 与 XLA 策略必须联合 setup。见 fabric.py L1223-L1229if isinstance(self._strategy, (DeepSpeedStrategy, XLAStrategy)): raise RuntimeError( fThe {type(self._strategy).__name__} requires the model and optimizer(s) to be set up jointly through .setup(model, optimizer, ...). )也就是说在这类策略下前文的模式一/二模型与优化器一起 setup是唯一合法路径。进阶需要分步 setup 的场景FSDP 与自定义优化器虽然原则是模型与优化器一起 setup但存在必须分步的情况。Fabric提供了两个配套方法fabric.setup_module(model)只包装模型等价于setup(model)不传优化器fabric.setup_optimizers(*optimizers)只包装优化器且要求至少传入一个优化器。分步写法来自 setup_module 的 docstring# Set up model first (useful for FSDP) model fabric.setup_module(model) # Then create and set up optimizer optimizer torch.optim.Adam(model.parameters()) optimizer fabric.setup_optimizers(optimizer)需要分步的典型场景FSDP 且use_orig_paramsFalse从 FSDPStrategy.setup_module_and_optimizers 的实现可见联合 setup 会强制要求use_orig_paramsTrue若设置为False必须按先setup_module、创建优化器、再setup_optimizer的顺序操作FSDP 下优化器引用 meta device 参数校验逻辑fabric.py L1208-L1216会拒绝这种组合提示先 setup 模型再创建优化器需要自定义优化器包装逻辑setup_optimizer允许策略对优化器做额外处理。例如 FSDP 的 setup_optimizer 会校验优化器是否基于包装后的扁平化参数创建。常见问题与最佳实践小结务必用返回值覆盖变量model, optimizer fabric.setup(model, optimizer)不要只调用不接收返回否则训练仍在未包装的原始对象上进行分布式与精度逻辑不生效。每个对象只 setup 一次_validate_setup会直接对重复 setup 抛出ValueErrorA model should be passed only once to the setup method.。优化器必须在模型 setup 之后创建FSDP 场景需要分步 setup 时先setup_module再从包装后的模型参数创建优化器最后setup_optimizers。DeepSpeed / XLA 只能联合 setup不要尝试用setup_modulesetup_optimizers拆分否则会触发RuntimeError。多模型共享优化器时先聚合把子模型收拢到一个顶层nn.Module如AutoEncoder再整体 setup保证策略包装覆盖全部参数。调度器与优化器一起传scheduler参数只在传了优化器时可用且必须位于优化器之后。四种组合模式的选用可归纳如下表组合推荐写法典型场景1 模型 1 优化器model, opt fabric.setup(model, opt)常规监督训练1 模型 N 优化器model, o1, o2 fabric.setup(model, o1, o2)分层学习率、局部微调N 模型 1 优化器先聚合为顶层nn.Module再 setup自编码器、共享参数模型N 模型 N 优化器按配对多次调用setupGAN、元学习无论哪种组合核心都是把模型与其优化器作为整体交给fabric.setup()让设备、精度、分布式包装的复杂度由 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），仅供参考