Lightning Fabric 完全指南：零样板代码扩展 PyTorch 训练与自建 Trainer
发布时间：2026/9/20 14:46:05 作者：尧图编辑部 阅读量：1,286
Lightning Fabric 完全指南零样板代码扩展 PyTorch 训练与自建 Trainer【免费下载链接】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导读本文围绕当前仓库中 src/lightning_fabric/README.md 展开系统讲解 Lightning Fabric 的核心设计、最小代码改动迁移方法、硬件与分布式策略配置以及基于 Fabric 原语自建 Trainer 的完整实战。读完本文你将掌握如何在不重构 PyTorch 代码的前提下把单机 CPU 脚本无缝扩展到多 GPU、多节点乃至 TPU并能用fabric.setup/fabric.backward/fabric.save等原语搭建属于自己的训练循环。Fabric 是什么为专家级控制而生的轻量扩展层Fabric 是 Lightning 生态中面向 PyTorch 原教旨主义者的加速层其定位与 Lightning Trainer 形成互补Fabric保留你对训练循环的完全控制不强制任何 epoch 结构只提供一套工具箱原语把设备摆放、分布式初始化、精度切换等样板代码接走Lightning Trainer把 checkpoint、日志、早停等几十项工程能力全部封装适合希望开箱即用的场景。README 开篇即给出定位Run on any device at any scale with expert-level control over PyTorch training loop and scaling strategy. You can even write your own Trainer.在任意设备、任意规模上运行同时对 PyTorch 训练循环与扩展策略保持专家级控制甚至可以自写 Trainer。在 src/lightning/fabric/fabric.py 的Fabric类文档中核心能力被归纳为自动设备摆放、自动混合精度与双精度支持、CPU/GPU/TPU 与分布式策略间无缝切换、自动进程派生无需启动工具以及多节点支持。从源码结构看Fabric 由acceleratorsCPU/CUDA/MPS/XLA、strategiesDP/DDP/DeepSpeed/FSDP/ModelParallel 等、pluginsprecision、checkpoint_io、cluster_environment与wrappers_FabricModule、_FabricOptimizer、_FabricDataLoader四层组件构成全部集中在 src/lightning/fabric 目录下Fabric类本身只负责编排这些组件。最小迁移把一段 PyTorch 训练脚本改成 Fabric可复制README 用一张改动对照表直观展示了迁移所需的全部改动。下面给出可直接运行的完整版本原文档中 diff 与完整代码的合并import lightning as L import torch; import torchvision as tv dataset tv.datasets.CIFAR10(data, downloadTrue, trainTrue, transformtv.transforms.ToTensor()) fabric L.Fabric() fabric.launch() model tv.models.resnet18() optimizer torch.optim.SGD(model.parameters(), lr0.001) model, optimizer fabric.setup(model, optimizer) dataloader torch.utils.data.DataLoader(dataset, batch_size8) dataloader fabric.setup_dataloaders(dataloader) model.train() num_epochs 10 for epoch in range(num_epochs): for batch in dataloader: inputs, labels batch optimizer.zero_grad() outputs model(inputs) loss torch.nn.functional.cross_entropy(outputs, labels) fabric.backward(loss) optimizer.step()对照原始 PyTorch 写法全部改动只有四处用标注 import lightning as L fabric L.Fabric() fabric.launch() model, optimizer fabric.setup(model, optimizer) - device cuda if torch.cuda.is_available() else cpu - model.to(device) dataloader fabric.setup_dataloaders(dataloader) - inputs, labels inputs.to(device), labels.to(device) - loss.backward() fabric.backward(loss)其中fabric.setup()内部会按顺序执行精度转换self._precision.convert_module、设备迁移_move_model_to_device、策略封装strategy.setup_module_and_optimizers并最终把模型包成_FabricModule、优化器包成_FabricOptimizer见 fabric.pyfabric.backward()则把loss.backward()转发给策略从而让 AMP 梯度缩放、DeepSpeed 引擎等底层逻辑自动生效fabric.py。零改动切换硬件与规模accelerator / devices / num_nodesREADME 强调换硬件不需要改代码。全部通过Fabric构造参数完成# 使用当前可用硬件自动探测 fabric Fabric() # 使用 GPUCUDA 或 Apple Silicon 的 MPS fabric Fabric(acceleratorgpu) # 单机 8 卡 fabric Fabric(acceleratorgpu, devices8) # 32 节点 × 8 卡 256 个 GPU fabric Fabric(acceleratorgpu, devices8, num_nodes32) # 使用 TPU fabric Fabric(acceleratortpu)各参数在 Fabric.init中的语义源码 docstring 依据acceleratorcpu/cuda/mps/gpu/tpu/auto默认auto自动探测当前可用硬件_Connector._choose_auto_accelerator负责选择devices整数设备数量、列表或字符串指定具体设备编号、auto按节点生效默认autonum_nodes参与训练的机器节点数默认1strategy跨设备运行策略dp/ddp/ddp_spawn/deepspeed/fsdp/auto默认autoprecision64双精度、32全精度、16-mixedFP16 自动混合精度、bf16-mixedBF16 AMP传None时按设备选择默认值。这些参数的实际解析发生在 _Connector 中它先校验各 flag 之间的冲突再实例化 Accelerator 与 Strategy。值得注意的是connector 还会读取LT_ACCELERATOR、LT_DEVICES等环境变量由 CLI 设置这就是CLI 参数覆盖代码参数的实现机制。开箱即用的分布式策略与混合精度README 展示了策略与精度的一行切换# 先进分布式训练技术 fabric Fabric(strategyddp) fabric Fabric(strategydeepspeed) fabric Fabric(strategyfsdp) # 切换精度 fabric Fabric(precision16-mixed) fabric Fabric(precision64)在源码层面每种策略都是一个独立类DataParallelStrategydp、DDPStrategyddp含_DDP_FORK_ALIASES、DeepSpeedStrategydeepspeed内置 ZeRO 阶段、CPU/NVMe offload 等大量配置项、FSDPStrategyfsdp支持sharding_strategy、cpu_offload、activation_checkpointing、state_dict_type等详见 src/lightning/fabric/strategies。策略通过STRATEGY_REGISTRY注册注册表实现在 src/lightning/fabric/strategies/registry.py因此strategy参数天然可扩展。精度则由 src/lightning/fabric/plugins/precision 下的插件实现MixedPrecision16-mixed / bf16-mixed、HalfPrecision16-true / bf16-true、DoublePrecision64、BitsandbytesPrecisionNF4/FP4/INT8 量化与TransformerEnginePrecision等。16-mixed这类 AMP 精度通过torch.amp.GradScaler自动完成梯度缩放见 amp.py无需手动scaler.scale/unscale。注意DeepSpeed、XLA 策略要求模型与优化器联合 setupfabric.setup(model, optimizer)不支持单独setup_optimizers见 fabric.py 的说明使用torch.compile时setup会自动剥除并重套OptimizedModule包装_reapply_compile参数默认True。设备样板代码全接管告别.to(device)/.cuda()README 用一个 diff 强调这一点# no more of this! - model.to(device) - batch.to(device)背后的实现是三层隐式搬移fabric.setup(model, optimizer)在move_to_deviceTrue时调用_move_model_to_device把模型放到self.device策略的根设备上fabric.setup_dataloaders(dataloader)把 dataloader 包装为_FabricDataLoader其__iter__在产出 batch 时自动调用move_data_to_devicewrappers.py分布式训练时setup_dataloaders还会自动替换DistributedSampler并注入worker_init_fn保证多进程数据切分与随机种子一致fabric.py。如果你希望手动控制设备把move_to_deviceFalse传入setup/setup_dataloaders再显式调用fabric.to_device(obj)即可。用 Fabric 原语自建 TrainerBYOTFabric 的所有特性都是**按需开启opt-in**的因此你可以用它拼装出自己的 Trainer。README 给出了最小骨架import lightning as L class MyCustomTrainer: def __init__(self, acceleratorauto, strategyauto, devicesauto, precision32-true): self.fabric L.Fabric(acceleratoraccelerator, strategystrategy, devicesdevices, precisionprecision) def fit(self, model, optimizer, dataloader, max_epochs): self.fabric.launch() model, optimizer self.fabric.setup(model, optimizer) dataloader self.fabric.setup_dataloaders(dataloader) model.train() for epoch in range(max_epochs): for batch in dataloader: input, target batch optimizer.zero_grad() output model(input) loss loss_fn(output, target) self.fabric.backward(loss) optimizer.step()仓库中提供了比骨架完整得多的可运行实现位于 examples/fabric/build_your_own_trainertrainer.py 中的MyCustomTrainer在构造参数上对齐了Fabric的全部关键项accelerator、strategy、devices、precision、plugins、callbacks、loggers并额外实现了max_epochs/max_steps/grad_accum_steps梯度累积/limit_train_batches/limit_val_batches/validation_frequency/checkpoint_dir等训练控制项训练循环通过fabric.call(on_train_batch_start)等回调钩子实现扩展点支持与 LightningModule 的 hook 互相调用README 注释中强调为 Lightning Trainer 编写的回调不一定适用于自建 Trainercheckpoint 通过self.fabric.save(path, state)落盘、self.fabric.load(path, state)恢复支持断点续训ckpt_pathrun.py 给出 MNIST 用法定义带training_step/configure_optimizers/validation_step的LightningModule然后trainer.fit(model, train_loader, val_loader)即可。运行方式README 依据cd examples/fabric/build_your_own_trainer python run.py依赖安装build_your_own_trainer/README.mdpip install lightning2.0 torchmetrics0.11 torchvision0.14 torch1.13 tqdm需要注意BYOT 示例目前不支持 FSDP 策略源码中显式raise NotImplementedError且仅支持单一优化器。安装与命令行启动安装先创建虚拟环境推荐 Python 3.8 及以上版本然后pip install -U lightninglightning是统一包名安装后import lightning as L即可同时获得 Fabric 与 Lightning Trainer 两套 API。用 CLI 启动分布式训练普通脚本已包含fabric.setup等调用可以直接用 Fabric CLI 启动lightning run model path/to/train.pyCLI 的实现位于 src/lightning/fabric/cli.py底层会调用torchrun派生进程_torchrun_launch并把参数写入LT_ACCELERATOR、LT_STRATEGY、LT_DEVICES、LT_NUM_NODES、LT_PRECISION等环境变量供_Connector读取覆盖代码中的默认值。因此通过 CLI 启动时不要在代码里再调用fabric.launch()否则会抛出RuntimeErrorfabric.pyCLI 常用选项--acceleratorcpu/gpu/cuda/mps/tpu/auto、--strategy从策略注册表中过滤 spawn/fork/notebook/xla/tpu/offload 后的可用项、--devices每节点设备数或设备列表、--num-nodes、--node-rank、--main-address、--main-port默认 29400、--precision脚本自身的参数追加在命令末尾透传CLI 还提供lightning run model ... consolidate子命令用于把 FSDP 分片 checkpoint 合并成单个可用torch.load()加载的文件支持 S3/GCS 等远程文件系统路径。在 Notebook / 交互环境中启动非 CLI 场景如 Jupyter Notebook直接调用fabric.launch()即可完成进程初始化对于 spawn/fork 类策略如ddp_spawn、XLAlaunch()需要传入一个接收fabric作为首参的函数fabric.pydef train_function(fabric): model, optimizer fabric.setup(model, optimizer) # ... training code ... fabric Fabric(acceleratortpu, devices8) fabric.launch(train_function)FAQ如何选择 Fabric 与 Lightning Trainer什么时候用 Fabric最小代码改动想用多 GPU 或 DeepSpeed 等高级策略但不想重构现有 PyTorch 代码只想尽快把规模提上去最大控制权希望自己掌控训练/推理逻辑直至单次 optimizer 调用不被迫套用 Trainer 的标准 epoch 结构——迭代式训练、元学习、交叉验证等灵活算法都可以直接实现README 原话With great power comes a great responsibility.最大灵活性所有特性 opt-in用原语工具箱自建 Trainer。什么时候用 Lightning Trainer想要 checkpoint、日志、早停等开箱即用的工程化能力少踩坑想要现成的良好默认值作为起点希望代码模块化、可读、易分享。能否与 LightningModule / Callback 混用可以。Fabric 完全兼容LightningModule与 Lightning CallbackREADME 明确说明选择如何组织代码、复用已有模型与回调都随你。源码中fabric.setup()检测到模型带_fabric属性即 LightningModule时会自动把模块注册进回调列表fabric.py并在setup完成后触发on_after_setup钩子。结语Fabric 的价值在于把控制权还给你把样板代码收走四行改动即可完成 PyTorch 脚本的迁移一行参数即可切换硬件、策略与精度而setup/backward/save/load等原语足以支撑你写出自己的 Trainer。如果追求少操心可以无缝回到 Lightning 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），仅供参考
尧图网站开发 