Lightning Fabric 模型编译加速指南在 Fabric 中正确使用 torch.compile【免费下载链接】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导读本指南以 Lightning Fabric 官方文档《Speed up models by compiling them》为主体系统讲解如何在 Fabric 项目中把 PyTorch 的torch.compile用对、用好包括最基础的编译先行、setup 在后的调用顺序、与ModelParallelStrategy结合进行分布式编译、如何排查 graph break 与重编译导致的训练变慢、以及 CUDA Graphs 与 shape padding 等进阶编译选项。读完本文你将掌握一套可直接复制的torch.compile Fabric 提速方法论并理解fabric.setup()背后自动重放编译_reapply_compile的实现原理避免在 DDP/FSDP 场景下踩坑。一、基础用法在 Fabric 脚本中给模型加一行编译代码在 Fabric 中编译模型非常简单只需在调用fabric.setup()之前给模型套上torch.compileimport torch import lightning as L # Set up Fabric fabric L.Fabric(devices1) # Define the model model ... # Compile the model model torch.compile(model) # fabric.setup() should come after torch.compile() model fabric.setup(model)重要应如上例所示在调用fabric.setup()之前编译模型才能与 Fabric 的各项特性实现最优集成。原因在本文第五节与 DDP/FSDP 结合中会展开Fabric 会在setup()内部自动摘除OptimizedModule包装、完成 DDP/FSDP 包装后再以相同参数重新套回torch.compile。torch.compile()这一行本身并不会立刻做多少事它只是把模型包装成一个已编译模型torch._dynamo.OptimizedModule。真正的优化发生在模型第一次被调用forward()的时刻# 1st execution compiles the model (slow) output model(input) # All future executions will be fast (for inputs of the same size) output model(input) output model(input) ...因此在测量编译模型与普通模型的性能差异时务必把第一次forward()调用排除在计时之外——第一次调用包含了编译时间会让结果严重失真。附带基准测试的完整示例下面的示例测量 TorchVision 的 InceptionV3 在编译前后的加速比import statistics import torch import torchvision.models as models import lightning as L torch.no_grad() def benchmark(model, input, num_iters10): Runs the model on the input several times and returns the median execution time. start torch.cuda.Event(enable_timingTrue) end torch.cuda.Event(enable_timingTrue) times [] for _ in range(num_iters): start.record() model(input) end.record() torch.cuda.synchronize() times.append(start.elapsed_time(end) / 1000) return statistics.median(times) fabric L.Fabric(acceleratorcuda, devices1) model models.inception_v3() input torch.randn(16, 3, 510, 512, devicefabric.device) # Compile! compiled_model torch.compile(model) # Set up the model with Fabric model fabric.setup(model) compiled_model fabric.setup(compiled_model) # warm up the compiled model before we benchmark compiled_model(input) # Run multiple forward passes and time them eager_time benchmark(model, input) compile_time benchmark(compiled_model, input) # Compare the speedup for the compiled execution speedup eager_time / compile_time print(fEager median time: {eager_time:.4f} seconds) print(fCompile median time: {compile_time:.4f} seconds) print(fSpeedup: {speedup:.1f}x)官方文档在 NVIDIA A100 SXM4 40GB、PyTorch 2.2.0、CUDA 12.1 环境下得到的参考输出为Eager median time: 0.0254 seconds Compile median time: 0.0185 seconds Speedup: 1.4x注意两点一是benchmark中通过torch.cuda.Event精确计时且每次迭代都torch.cuda.synchronize()确保测到的是真实的 GPU 执行时间二是在基准之前先做了一次compiled_model(input)的 warm-up这正是为了把首次编译的时间排除在测量之外。二、进阶用法通过 ModelParallelStrategy 的 parallelize_fn 编译torch.compile也可以作为ModelParallelStrategy的parallelize_fn参数的一部分被调用。当torch.compile与torch.distributed.tensorAPI 组合使用时这种方式尤为方便。从源码看ModelParallelStrategy 接收一个parallelize_fn回调该函数以模型和DeviceMesh为输入负责对模型施加各类并行化处理而setup_module内部会在拿到 device mesh 之后调用它见 model_parallel.py。其好处是parallelize在模型被分片shard时才被调用因此torch.compile保证作用在模型分片之上能够捕获分布式算子并一并优化。import lightning as L import torch import torch.nn as nn import torch.nn.functional as F from lightning.pytorch.demos import Transformer from lightning.fabric.strategies.model_parallel import ModelParallelStrategy from torch.distributed._composable.fsdp.fully_shard import fully_shard from torch.distributed.device_mesh import DeviceMesh def parallelize(model: nn.Module, device_mesh: DeviceMesh) - nn.Module: for module in model.modules(): if isinstance(module, (torch.nn.TransformerEncoderLayer, torch.nn.TransformerDecoderLayer)): fully_shard(module, meshdevice_mesh) fully_shard(model, meshdevice_mesh) return torch.compile(model) def train(): L.seed_everything(42) with torch.device(meta): model Transformer( vocab_size50257, nlayers16, nhid4096, ninp1024, nhead32, ) strategy ModelParallelStrategy(data_parallel_size4, tensor_parallel_size1, parallelize_fnparallelize) fabric L.Fabric(precisionbf16-true, strategystrategy) fabric.launch() model fabric.setup(model)示例中的几个关键点模型在with torch.device(meta)上下文中以meta device方式实例化只建结构、不分配内存便于超大模型的分片初始化ModelParallelStrategy(data_parallel_size4, tensor_parallel_size1, ...)指定了数据并行与张量并行的维度划分其中data_parallel_size默认auto取集群节点数tensor_parallel_size默认auto取单节点 GPU 数对每个TransformerEncoderLayer/TransformerDecoderLayer调用fully_shard最后再对整个模型fully_shard这是 FSDP2 风格的逐模块分片parallelize最后返回torch.compile(model)编译发生在分片之后。与 torchao 组合compile(distributed(quantized(model)))当需要叠加其他按类似方式应用的库如 torchao时parallelize_fn让调用顺序一目了然可以轻松得到compile(distributed(quantized(model)))的组合效果import lightning as L import torch import torch.nn as nn import torch.nn.functional as F from lightning.pytorch.demos import Transformer from torch.distributed._composable.fsdp.fully_shard import fully_shard from torch.distributed.device_mesh import DeviceMesh from torchao.float8 import Float8LinearConfig, convert_to_float8_training def parallelize(model: nn.Module, device_mesh: DeviceMesh) - nn.Module: float8_config Float8LinearConfig( pad_inner_dimTrue, ) def module_filter_fn(mod: torch.nn.Module, fqn: str): return fqn ! decoder convert_to_float8_training(model, configfloat8_config, module_filter_fnmodule_filter_fn) for module in model.modules(): if isinstance(module, (torch.nn.TransformerEncoderLayer, torch.nn.TransformerDecoderLayer)): fully_shard(module, meshdevice_mesh) fully_shard(model, meshdevice_mesh) return torch.compile(model) def train(): L.seed_everything(42) with torch.device(meta): model Transformer( vocab_size50257, nlayers16, nhid4096, ninp1024, nhead32, ) strategy ModelParallelStrategy(data_parallel_size4, tensor_parallel_size1, parallelize_fnparallelize) fabric L.Fabric(precisionbf16-true, strategystrategy) fabric.launch() model fabric.setup(model)这里通过module_filter_fn跳过decoder模块通常因为词表大小不满足 float8 的整除要求对模型其余部分做 float8 训练转换然后分片、编译。仓库中提供了可直接运行的完整版本examples/fabric/fp8_distributed_transformer/train.py它在此基础上补充了 WikiText2 数据加载、micro_batch_size1的梯度累积通过fabric.no_backward_sync(model, enabledis_accumulating)关闭累积步的反向同步、Adam 优化器与fabric.setup_optimizers()、以及fabric.is_global_zero控制只在 rank 0 上显示进度条等完整训练循环细节README.md 中还说明了启动方式与 Triton nightly 的安装前提。三、避免 Graph Break让模型被完整编译当torch.compile分析模型forward()中的代码时它会尽量把整段代码编译成图。如果遇到它无法理解的代码区域就会引入所谓的graph break图断裂把代码切成被优化的部分和未被优化的部分。Graph break 并不是致命的——被优化的部分仍然会更快。但如果你想把torch.compile的收益榨干就需要投入精力去重写产生 break 的问题代码段。可以用fullgraphTrue来检查模型是否存在 graph break一旦发生图断裂编译会直接报错# Force an error if there is a graph break in the model model torch.compile(model, fullgraphTrue)需要提醒的是这里产生的报错信息往往相当晦涩你大概率需要借助 PyTorch 官方的 compiler troubleshooting 文档做一番排查才能让模型被完整编译。四、避免重编译输入形状变化是训练突然变慢的元凶如第一节所述模型的实际编译发生在第一次调用forward()时。此时 PyTorch 会检查输入张量并针对该输入的形状、数据类型及其他属性优化出专用代码如果输入形状在所有forward()调用中保持不变PyTorch 会直接复用已生成的编译代码获得最佳加速如果这些属性在后续调用中发生变化PyTorch 会被迫为新形状重新编译模型。若每个 iteration 都触发重编译训练会被显著拖慢。当你的训练突然变慢时很可能是 PyTorch 正在重新编译模型常见触发场景训练代码里包含在不同数据集上的评估步骤或使用了会在训练与验证/测试之间切换的Trainer导致输入形状改变、触发重编译数据集大小不能被 batch size 整除而 dataloader 使用默认的drop_lastFalse——训练循环中的最后一个 batch 会更小从而触发一次重编译。理想情况下应尽量让传入forward()的输入形状保持静态。当确实无法做到时可以让 PyTorch 在编译时把输入形状可能的动态变化纳入考虑# On PyTorch 2.2 model torch.compile(model, dynamicTrue)用dynamicTrue编译出的模型通常比静态形状的编译模型慢一些但可以避免每个 iteration 都付出极端高昂的重编译代价。在 PyTorch 2.2 及更高版本上torch.compile会自动检测动态性通常不再需要手动设置该参数。动态形状对比实验下面的例子演示了输入形状改变导致模型重编译数秒的现象可以通过切换dynamicTrue/False对比计时结果import time import torch import torchvision.models as models import lightning as L fabric L.Fabric(acceleratorcuda, devices1) model models.inception_v3() # dynamicFalse is the default torch._dynamo.config.automatic_dynamic_shapes False compiled_model torch.compile(model) compiled_model fabric.setup(compiled_model) input torch.randn(16, 3, 512, 512, devicefabric.device) t0 time.time() compiled_model(input) torch.cuda.synchronize() print(f1st forward: {time.time() - t0:.2f} seconds.) input torch.randn(8, 3, 512, 512, devicefabric.device) # note the change in shape t0 time.time() compiled_model(input) torch.cuda.synchronize() print(f2nd forward: {time.time() - t0:.2f} seconds.)官方文档在 A100 SXM4 40GB、PyTorch 2.2.0、CUDA 12.1 上的对比数据With automatic_dynamic_shapesTrue: 1st forward: 41.90 seconds. 2nd forward: 89.27 seconds. With automatic_dynamic_shapesFalse: 1st forward: 42.12 seconds. 2nd forward: 47.77 seconds.可见第一次forward因编译耗时约 42 秒第二次改变 batch 大小后若未启用动态形状检测重新编译又耗费约 47 秒。若仍遇到重编译问题可以进一步借助 PyTorch 的 Compile Profiler 深入调查。五、实验性编译选项CUDA Graphs 与 Shape Padding以下可选设置取决于具体模型可能带来额外加速。CUDA Graphs启用 CUDA Graphs 后CUDA 会把所有计算录制进一张图每次调用 forward 和 backward 时直接重放。其前提是模型必须是静态的——输入形状不能变化且模型每次都要执行相同的算子。启用 CUDA Graphs 通常能带来显著加速但有时也会增加模型的内存占用。# Enable CUDA Graphs compiled_model torch.compile(model, modereduce-overhead) # This does the same compiled_model torch.compile(model, options{triton.cudagraphs: True})Shape Padding形状填充模型计算中涉及的张量输入、激活、权重、梯度等的具体形状/尺寸会影响性能。启用 shape padding 后torch.compile可以通过填充把张量扩展到内存对齐更优的尺寸。代价自然是会多消耗一点内存。# Default is False compiled_model torch.compile(model, options{shape_padding: True})torch.compile的完整选项列表以 PyTorch 官方文档为准。六、关于 torch.compile 实践的一点忠告在实际使用中你会发现torch.compile一开始可能并不好用甚至可能对性能起反作用编译可能以晦涩难懂的报错失败调试困难产出明显更慢或内存占用更高的模型并不罕见如果你的模型不属于幸运路径上的那一类就需要在这一阶段投入时间调优编译阶段本身耗时可能长达数分钟。因此官方建议开发阶段不要花太多时间打磨torch.compile而应把它的效果评估放到最后、也就是即将启动长时昂贵的实验时再做。永远把编译模型的耗时与内存占用和原始模型对比七、与 DDP/FSDP 结合fabric.setup 的自动重编译机制如第一节所述官方推荐在fabric.setup()之前编译模型。对于 DDP 和 FSDPfabric.setup()会在内部自动摘除OptimizedModule包装、待模型被 DDP/FSDP 包装完成后以相同的参数重新应用torch.compile从而让编译能够把分布式调用也纳入优化范围。这一机制的源码实现位于 src/lightning/fabric/wrappers.py_capture_compile_kwargswrappers.py在导入 lightning 时就把torch.compile包装起来把用户调用时传入的 kwargs 深拷贝并挂到编译后的模块上compiled_model._compile_kwargs deepcopy(kwargs)_unwrap_compiledwrappers.py负责在setup()时把OptimizedModule外壳去掉、取出_orig_mod和_compile_kwargs。它会在拿不到_compile_kwargs时抛出明确错误Make sure to import lightning before torch.compile is used.——这正是要求先import lightning再调用torch.compile的原因_to_compiledwrappers.py在策略包装如 DDP/FSDP完成后用保存下来的参数重新执行torch.compile(module, **compile_kwargs)。整个流程对应 fabric.py 中setup()的实现先_unwrap_compiled解包 → 交给 strategy 包装 → 再_to_compiled重新编译 → 最后包上_FabricModule。仓库测试 tests/tests_fabric/strategies/test_ddp_integration.py 的test_reapply_compile明确验证了这一行为fabric.setup(compiled_model, _reapply_compileTrue)之后_forward_module是OptimizedModule其_orig_mod是DistributedDataParallel并且断言torch.compile以与用户原始调用相同的参数被再次调用——即编译被重新施加在 DDP 包装后的模块上。如果你在编译 DDP/FSDP 模型时遇到问题可以通过_reapply_compile参数关闭该特性# Choose a distributed strategy like DDP or FSDP fabric L.Fabric(devices2, strategyddp) # Compile the model model torch.compile(model) # Default: fabric.setup() will configure compilation over DDP/FSDP for you model fabric.setup(model, _reapply_compileTrue) # Turn it off if you see issues with DDP/FSDP model fabric.setup(model, _reapply_compileFalse)从 fabric.py 的 API 文档可知_reapply_compileTrue默认表示如果模型之前被torch.compile过则在模型被策略设置完成例如被 DDP、FSDP 包装之后用相同设置重新套回OptimizedModule包装设为False则跳过。该参数同样存在于setup_module()fabric.py适用于 FSDP 那种需要先 setup 模型、再创建优化器的流程。另外值得一提的是FSDP 策略在初始化时会默认设置use_orig_paramsTrue见 fsdp.py这为torch.compile()的配合使用创造了条件。八、深入阅读本文对应的官方文档位于 docs/source-fabric/advanced/compile.rstfabric.setup()/setup_module()的_reapply_compile参数实现见 src/lightning/fabric/fabric.py编译参数捕获与重放机制见 src/lightning/fabric/wrappers.pyModelParallelStrategy的parallelize_fn机制见 src/lightning/fabric/strategies/model_parallel.py完整的 FP8 分布式编译训练示例见 examples/fabric/fp8_distributed_transformer/train.py验证 DDP 下编译重放行为的测试见 tests/tests_fabric/strategies/test_ddp_integration.py。【免费下载链接】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），仅供参考