最近在调一个nlp分类模型,数据量不大(大概2万条),单卡3090训练。试了torch.compile,用默认模式跑,第一轮确实慢(编译开销),但后面几个epoch速度只提升了10%左右,和官方宣传的“训练提速30%-50%”差很多。而且用动态shape(比如变长序列padding到不同长度)直接报错,回退到eager模式。想问问大家:是不是小模型小数据量根本没必要上compile?还是我哪里设置不对?有一说一,模型里用了transformers的BertForSequenceClassification,加了个自定义分类头。有没有类似场景的老哥分享下实际收益?现在纠结要不要为了部署时的推理加速去折腾这个。
PyTorch 2.0的compile到底值不值得用?静态图加速不明显还报错
全部回复
共 107 条2万条数据确实太小了,compile那点编译开销摊到几个epoch根本回不了本,官方30%+基本是大模型长训练才跑得出来。动态shape报错是老毛病了,Bert变长输入要么pad到固定长度,要么上dynamic=True试试,但不一定稳。你这种情况我建议训练阶段就别折腾了,推理部署时再compile收益反而更明显。
2万条数据配3090,瓶颈大概率在数据加载和Python开销上,compile能压的那点算子时间占比本来就小,10%挺正常。动态shape报错是老毛病了,得开dynamic=True或者手动mark_dynamic,不然它默认按第一次的shape编译,一变就炸。Bert这种本身算子就挺规整的,收益确实不如那些自定义op多的模型明显,我之前有个类似规模的分类任务也是差不多10%出头。推理部署的话倒是可以试试,但建议先拿torch.export走一遍看能不能过,别到上线才发现fallback了。
小数据量真没必要硬上compile,编译开销都摊不平。动态shape报错是硬伤,等2.1修吧。
小数据量确实没必要硬上compile,编译开销都回不来。变长序列报错很正常,我后来直接关掉动态shape才勉强跑通。
我这边也是3090,数据和模型规模跟你差不多,compile默认模式确实就10%上下,宣传那个数字得是ResNet那种大吞吐场景才跑得出来。动态shape报错可以试试mode="reduce-overhead"或者设dynamic=True,但小模型编译开销摊不回来,收益很有限。你这个量级老实说没必要硬上,部署推理倒是可以单独compile一下试试,训练阶段先别折腾了。
小数据量确实没啥必要,编译开销都摊不回来,动态shape报错也是老毛病了,推理部署再上也不迟。
2万条数据配3090,瓶颈大概率在数据加载和python开销上,编译收益自然被稀释了。动态shape那块试试dynamic=True或者mark_dynamic,不然变长序列确实容易炸。我这边类似规模的小模型,compile后也就快个15%左右,推理端倒是收益更明显些。你要是主要为了部署推理,不如直接上onnx或者TensorRT,比折腾compile省心。