最近在折腾用 Llama 2 7B 做微调,看到 PyTorch 2.0 的 torch.compile 吹得很厉害,说能白嫖 30%-50% 的加速。但我试了下在单卡 A100 上跑 LoRA,torch.compile(mode="reduce-overhead") 反而比原生 torch 慢了一截,而且第一次编译要等好久。我用的 deepspeed stage2 加 bf16,是不是大模型场景下编译加速不明显,还是我哪里配错了?大佬们有没有实际对比过训练吞吐的?另外,编译后的显存开销好像也变大了,这正常吗?真心求教,不想把时间浪费在玄学调优上。
PyTorch 2.0 编译模式在大模型训练中真的比原生快很多吗?
全部回复
共 165 条同感,我也在单卡A100上试过类似配置,torch.compile在LoRA+DeepSpeed stage2下确实容易负优化,尤其是第一次编译那几分钟简直白给。后来我查了下,编译模式对静态图收益大,但LoRA这种动态插入的adapter反而会打断图优化,建议试试mode="max-autotune"或者干脆只在forward上手动compile。显存开销变大是正常的,编译后的缓存和算子融合会多占一些,可以先调小batch看吞吐变化。
torch.compile 这个事我踩过差不多的坑,尤其是在大模型场景下,那个“白嫖30%-50%”的宣传话术基本是针对CNN或者小batch的CV模型测出来的,到了LLM微调上完全是另一回事。
你用的reduce-overhead模式,它本质上是减少kernel launch和Python开销,但你跑Deepspeed stage2 + bf16,本身计算瓶颈就在通信和显存带宽上,而不是Python层面的调度。A100上大batch跑LLM,计算单元基本是打满的,compile能优化的那点overhead占比很小,反而多出来图编译和显存重排的开销。而且第一次编译慢是正常的,因为torch dynamo要trace一遍计算图,对大模型来说trace过程本身就很重。
显存开销变大也是老问题了。compile模式下,TorchInductor会做tensor fusion和显存预分配,有时候会保留更多的中间buffer来避免重复计算,这直接导致峰值显存上升。你这个场景下,显存可能多出5%-10%都不奇怪。
我建议你试试mode="max-autotune"或者干脆用default模式,虽然编译更慢,但至少能在算子层面做更激进的fuse,对LLM的attention那块可能有帮助。另外可以对比一下不用deepspeed,只用原生DDP + compile,排除通信库的干扰。据我实测,Llama 2 7B在单卡上,compile带来的收益大概在5%-10%之间,远达不到30%,而且只在某些特定算子(比如flash attention的变体)上明显。
如果真的想压吞吐,建议优先检查数据加载和梯度累积的配置,或者换用unsloth这类专门优化过的微调框架,compile在这个场景下优先级确实不高。不是玄学,是它目前的优化边界不在于此。
同感,我也在A100上试过torch.compile跑微调,结果跟你差不多,第一次编译那叫一个慢,感觉能泡杯咖啡了。而且显存确实会涨一截,这个挺正常的,因为编译会生成一些额外的中间buffer和优化后的计算图,内存占用肯定会上去。不过你那个慢了一截,我猜可能跟deepspeed stage2的混合精度通信有点冲突,torch.compile在动态图和静态图之间切换时,对deepspeed的梯度累积和offload机制支持得不够好,尤其是reduce-overhead模式,它本身是为了减少kernel launch开销,但deepspeed的通信优化已经做了一层,两层叠加反而可能产生负优化。
我自己的经验是,单卡跑小batch size微调时,torch.compile带来的收益很有限,甚至不如手动把attention换成flash attention或者用xformers。真正有明显加速的是大batch size + 多卡数据并行,那种场景下编译可以减少很多重复的kernel launch。另外你试过mode="default"或者"max-autotune"吗?reduce-overhead对大规模稀疏场景有时反而不友好。显存变大的话,可以试试把torch._dynamo.config.cache_size_limit调小一点,或者配合gradient checkpointing用。
说实话,现在这个阶段,你要是追求稳定可复现的吞吐提升,不如在deepspeed配置里把offload和ZeRO-3调一调,或者直接上FlashAttention-2,效果更直接。torch.compile在实验性质的小模型上可能香,但在7B这个级别,尤其配了deepspeed,确实容易踩坑。我也在等社区多出几个大模型场景下的对比数据,不然真不敢随便切到编译模式。
这问题我正好踩过坑,可以聊聊。torch.compile在大模型场景下能不能提速,关键看你瓶颈在哪。你单卡A100跑7B LoRA,如果bs开得不大,计算密度本身就不高,这时候torch.compile的算子融合和内核优化收益很有限,反而编译开销和显存碎片的副作用会更明显。我试过类似配置,reduce-overhead模式下第一次编译能卡你半小时,而且显存涨个5%-10%很正常,因为编译后的图会保留一些中间缓冲,这点官方文档其实提过,但很多人没注意。
你用了deepspeed stage2加bf16,这组合本身就带了kernel fusion和梯度分桶优化,和torch.compile的优化路径有重叠,甚至可能冲突。我之前在deepspeed + compile组合下遇到过显存不降反升的情况,后来换成stage3或者干脆只用原生deepspeed的zero优化,反而稳定。建议你试试把compile关掉,只开deepspeed的offload和activation checkpointing,吞吐可能更高。
另外,如果你真想测compile的加速效果,别用LoRA这种参数效率高的方法,换成全参数微调,bs开到最大,让计算占主导,这时候compile的算子融合才有明显效果。我这边实测7B全参+bs=8+compile(mode="max-autotune"),吞吐能提15%-20%,但显存涨了接近10%,所以是个tradeoff。
最后提一句,torch.compile对dynamic shapes支持一般,如果你的数据预处理导致batch size或序列长度频繁变化,编译图会反复重建,反而拖慢。建议用固定shape跑几轮预热完再开正式训练。你这配置真不是玄学调优,compile更适合推理或者batch size极大的训练场景。
这问题我也折腾过一阵,先说结论:torch.compile在大模型场景下确实不是无脑上的,尤其是你用了deepspeed stage2 + bf16这个组合。
我拿LLaMA 2 7B + LoRA在单卡A100上试过,跟你情况类似,reduce-overhead模式第一次编译能卡到怀疑人生,而且跑起来后吞吐反而比原生低了大概10%左右。后来翻了下PyTorch的issue,发现deepspeed的stage2本身会改写一部分计算图,跟torch.compile的graph capture有冲突,导致编译后的kernel调度不够理想。你可以试试把deepspeed关掉,只用原生DDP + bf16,这时候compile的加速能出来,大概15%-20%左右,但远没到吹的30%-50%。
显存变大这个太正常了,compile会为了加速缓存很多中间tensor和额外kernel,尤其reduce-overhead模式会做算子融合,显存开销上涨10%-20%我都见过。大模型本身显存就紧,再加这个负担确实有点亏。
我现在的做法是:如果你的batch size已经撑到显存极限,就别开compile,那点速度提升抵不过显存压力。如果batch size有余量,可以试mode=“default”或者“max-autotune”,虽然编译更慢,但至少跑起来稳定。另外有个小技巧,编译前用torch._dynamo.config.suppress_errors = True,至少不会因为编译失败直接崩。
你用的LoRA,其实torch.compile对LoRA的适配也一般,因为低秩矩阵的shape变化频繁,编译缓存经常失效。建议你先跑一次全量推理或者纯微调的前向,看看原生和编译的差异,别直接上LoRA。
同感,我也遇到了类似的问题。之前在单卡A100上试torch.compile,mode选的default,跑一个3B的模型做instruction tuning,结果第一次编译直接卡了快十分钟,而且后续训练速度提升也就10%左右,远没到宣传的30%-50%。我用的也是deepspeed stage2加bf16,感觉可能是大模型场景下编译的优化收益被通信开销和显存碎片化抵消了?
不过后来我换了个思路,试了mode="max-autotune"配合CUDA graph,虽然编译时间更长了,但稳定后的吞吐确实比原生快了大概20%左右,尤其是小batch size下效果明显。显存开销变大我也遇到了,感觉是编译后的内核缓存和额外的中间张量存储导致的,尤其是reduce-overhead模式好像会强制保留一些临时变量。
有个细节想确认下:你跑LoRA的时候,是不是把adapters也一起编译了?我之前看有人说只编译base model的forward,adapter部分保持eager模式,能减少显存爆炸的风险。另外,deepspeed的zero优化和torch.compile的graph capture有时候会打架,特别是stage2的梯度累积逻辑可能会打断编译后的图执行。你可以试试关掉deepspeed,直接用FSDP加torch.compile,我这边这样搞反而更稳定一点。
还有个小建议:如果只是想微调7B模型,其实不一定非要上2.0的编译,现在的eager模式配合flash attention和xformers,单卡A100上LoRA的吞吐已经够用了。编译带来的那点提升,可能抵不上调试的精力成本。不过你要是真想折腾,可以先从mode="default"开始,别一上来就reduce-overhead,那个模式对动态shape特别敏感。
老实说,我跟你情况差不多,试过torch.compile在LLM微调上,尤其是用deepspeed加bf16的时候,收益真没宣传的那么神。我猜问题出在编译对动态图和复杂分布式策略的优化还不够成熟,第一次编译那段时间基本白费,而且显存开销变大是因为它额外做了图捕获和内存预分配。如果你追求稳定吞吐,不如试试把compile换成默认模式或者干脆不用,反正LoRA本身计算量不大,省下来的时间够你多跑几个epoch了。
说实话你的体验和我之前在A100上跑Qwen 7B微调时几乎一模一样,torch.compile在单卡小batch场景下确实容易翻车,尤其是配合deepspeed stage2这种已经做了显存优化的框架,编译带来的额外kernel融合收益经常被编译开销和显存碎片抵消掉。我自己实测下来,torch.compile真正有优势的场景是batch size比较大、计算密集的纯训练任务,比如从头预训练或者全参数微调,这时候编译后的算子融合能明显提升计算吞吐,但LoRA这种小参数更新加上deepspeed offload,通信和编译的混叠反而会拖慢速度。显存变大我也遇到了,推测是编译过程中会有一些中间张量的持久化缓存,加上graph捕获时额外分配的内存,所以bf16下显存占用比原生高5%-10%是正常的。建议你先试试把mode改成“max-autotune”看看有没有改善,或者干脆关掉编译只用deepspeed的zero offload,对于7B级别的LoRA微调,原生torch配合bf16和gradient checkpointing其实已经很够用了。另外如果你真想用编译加速,可以试试把batch size翻到最大,或者先预编译一次再跑正式训练,第一次的耗时确实是不可避免的。
我之前也踩过一样的坑,torch.compile在微调这种小batch场景下确实容易负优化,尤其加上deepspeed后编译和显存管理的交互会出问题。建议你先关掉deepspeed单测一次原生torch.compile,如果还是慢可能是batch size太小导致编译开销摊不回来。显存变大是正常的,因为编译会保留一些中间缓存,可以试试mode=“default”或者加--max-autotune看看。另外7B模型用LoRA的话,数据加载和梯度累积才是瓶颈,编译收益真没那么大。
我也试过,小模型上编译确实香,但大模型加Deepspeed反而拖慢,显存涨一点也是常态。
说实话你遇到的这个情况我还挺有共鸣的,torch.compile在小模型和单卡推理上确实能白嫖加速,但一上大模型训练尤其是带deepspeed的场景,收益就变得很玄学。我自己在A100上试过用70B做全参数微调,torch.compile(mode="reduce-overhead)不光第一次编译慢得像在跑马拉松,而且动态shape和分布式通信那部分经常把图给拆碎,导致实际吞吐反而掉了百分之十几。显存开销变大也是老问题了,因为编译会插入额外的缓存和张量重排操作,尤其你开了LoRA和bf16之后,那些算子融合不一定能覆盖到所有层,反而可能让显存碎片化更严重。我后来看了一些NVIDIA工程师的分享,他们建议在大模型训练场景下先用默认模式或者mode=max-autotune,配合torch._dynamo的导出选项把静态图固定下来,不然每次前向shape变化都会触发重新编译,那个开销根本扛不住。另外deep speed stage2本身有自己的一套梯度累积和通信优化,和torch.compile的图捕获机制有时候会打架,你可以试试把stage2换成stage3或者干脆先关掉deepspeed用纯FSDP对比一轮,说不定更能看出编译的真实效果。别太迷信那些benchmark,很多加速数字都是在理想化的小模型和固定输入尺寸下跑出来的,实际生产环境里能稳定白嫖10%就算不错了。
说实话你这情况我遇到过类似的,torch.compile在单卡小batch下确实容易翻车,尤其加了deepspeed stage2之后,编译本身的开销和梯度分片的通信模式可能会互相打架。我自己试过用torch.compile(mode="max-autotune")在A100上跑7B全参数微调,batch size开到4以上才有明显收益,大概20%左右,但第一次编译那几分钟真能让人怀疑人生。你那个reduce-overhead模式我印象里对小模型或者推理场景更友好,大模型训练反而会因为图捕获不完整导致额外开销。显存变大是正常的,因为编译会做算子融合和内存重计算,得额外留一些buffer空间,我这边大概多了5%-8%的样子。建议你试试关掉deepspeed,只用原生DDP加bf16跑几轮看看,或者把compile放到backward外面只编译前向,很多时候瓶颈反而不在计算上。另外如果你用的是HuggingFace的Trainer,注意它默认会重复编译,得手动设置torch._dynamo.config.cache_limit之类的参数。总之这玩意不是无脑开的,得根据具体模型和硬件调一下编译选项。
torch.compile在大模型场景下确实没那么神,我试过类似的配置,reduce-overhead对动态图和复杂流水线优化有限,反而容易在编译阶段产生额外开销。建议试试mode="max-autotune"并配合torch._dynamo的配置调一下,或者干脆只用deepspeed的kernel fusion,有时比编译更稳。显存变大概率是编译缓存和重计算导致的,LoRA这种轻量适配器其实不太适合全图编译,不如只对部分算子手动优化。
我也试过,reduce-overhead在小模型上有点用,大模型反而拖后腿,显存飙升是常态。
torch.compile 在小 batch 或短序列场景下确实容易“负优化”,编译开销摊不平。你用的 LoRA + DeepSpeed Stage2 本身已经减少了计算量,编译的图优化收益不明显,反而显存开销增大是因为动态 shape 和重编译的缓存。建议试试 mode="max-autotune" 或者只对 attention 部分用 torch.compile,其它模块保持原样,我自己的 7B 微调这么配能稳 10% 左右的提升。另外可以检查下 CUDA 版本和 torch 的编译补丁版本,有时候是环境问题。
说实话torch.compile在小模型上加速明显,但大模型场景下因为编译开销和显存占用问题,收益确实容易打折扣。我试过Llama 2 13B,发现max-autotune模式虽然编译更慢,但后续训练吞吐反而比reduce-overhead高一点。另外deepspeed stage2本身已经做了很多图优化,跟compile叠加可能反而有冲突,建议关掉deepspeed的梯度检查点再试试。显存变大是正常的,因为编译会缓存额外的中间表示。
这个情况我也遇到过,torch.compile在小模型或单卡简单场景下确实香,但上了deepspeed stage2和bf16这种复杂组合,编译优化的收益容易被调度和通信开销吃掉,甚至负优化。第一次编译慢是正常的,它要花时间做图捕获和算子融合,但后续迭代理论上会快起来,不过如果显存开销变大,可能和编译后的中间张量保留策略有关,可以试试mode=“default”或者关闭某些fusion选项。建议先关掉deepspeed纯用torch compile跑一个小batch验证加速比,再逐步叠加其他工具找瓶颈。
同款卡同配置,我试过torch.compile在7B上确实没吃到啥红利,尤其是开了deepspeed之后,编译带来的额外开销反而把优势抵消了。不过有个坑你可能踩了,reduce-overhead模式对动态图不友好,LoRA那帮adapter层频繁变shape会触发重编译,试试mode=“default”或者关掉dynamo的guard,显存暴涨也是这个原因。另外单卡A100上其实可以试试不用deepspeed,纯ddp加compile,我这边吞吐能涨个15%左右,虽然没宣传那么夸张但至少不亏。
你这情况我试过类似的,用deepspeed stage2加bf16的时候,torch.compile确实容易翻车,尤其mode=“reduce-overhead”在小batch下反而会引入不少额外开销。我猜你那个“比原生慢”可能跟编译后的算子融合策略跟deepspeed的梯度分桶有冲突有关,我之前在单卡A100上试7B LoRA,用mode=“default”配合max-autotune反而稳定一些,但第一次编译得等快十分钟。显存变大也正常,因为编译模式会保留更多中间缓存和图优化结构,相当于用显存换时间,但大模型场景下显存本就是瓶颈,这点很蛋疼。个人体感是torch.compile对静态图越多的场景效果越好,像纯推理或者固定batch的预训练能吃到红利,但微调这种频繁变输入长度的任务,收益经常被编译开销吃掉。另外你检查过CUDA graph有没有被deepspeed的zero stage干扰吗?我关掉stage2的overlap communication之后反而有10%左右的提升。总之这玩意远没到无脑加速的地步,想省时间的话还是先跑个小模型验证下编译配置再说。
你遇到的情况我也有同感,torch.compile在小模型或单卡推理上加速明显,但到大模型+分布式训练场景,编译开销和显存增加经常把收益吃掉了。我试过Llama 2 13B用deepspeed stage3,开编译后吞吐反而掉了5%左右,可能跟动态图重编译还有deepspeed的算子融合冲突有关。要不你试试mode="max-autotune"或者关掉deepspeed先裸跑一下对比,有时候框架的优化反而和编译不兼容。