最近在单卡A100上微调一个7B的模型，用HuggingFace的Trainer跑得好好的，但一开torch.compile就各种报错，不是graph break就是显存直接爆掉。看官方文档说能加速30%，可我这边折腾两天了连个完整的epoch都没跑下来。有没有实际在大模型预训练或微调里用起来compile的老哥？想问问你们是怎么处理dynamic shape和自定义算子这些坑的，还是说现在这玩意儿在大模型场景下根本还没法用？
楼主
18天前
PyTorch 2.0的torch.compile在大模型训练上到底靠不靠谱？
请 登录 后发表回复
全部回复
共 3 条
2楼
15天前
7B微调开compile基本就是给自己找罪受，dynamic shape那块坑太多，我后来直接关了。
3楼
13天前
7B微调开compile确实容易踩坑，我这边试下来感觉dynamic shape是最大元凶，batch里seq len一变就疯狂重编译，显存也跟着涨。后来把dynamic设成False、padding到固定长度才勉强跑通，但加速根本没到30%，有时候还不如不开。自定义算子那块更麻烦，rope和flash attention经常直接graph break，得手动加mark_compile_region才行。所以现在大模型场景我基本只在推理或者shape很固定的短序列训练上开，预训练那种变长输入还是先别折腾了。
4楼
13天前
7B模型开compile确实容易炸，我这边dynamic shape直接关掉才勉强跑通，但加速基本没感觉。