如何用PyTorch Lightning跑HuggingFace Transformer(TPU)

  1. TPU本身非常擅长处理transformer类的结构,而且可以白嫖(Colab)。
  2. HuggingFace Transformer本身是Transformer结构中最重要的库。但是他们使用的是自己的trainer,API经常改变。导致如果你要加一些自己的训练trick就比较困难。

所以目标是用PyTorch Lightning跑Hugging Face的Transformer。并且要求是在TPU上。代码在这里。别忘了把runtime改成TPU。

注意目前来说,TPU训练还是完全没有达到很好的效果的。相比之下,最大问题在于整个PyTorch生态圈不支持tf step_per_execution的这个。换句话说,是一点点喂给tpu而不是一次。另外,Lightning当中也没有控制prefetch的问题。后者好办,前者不太好解决。

有时间我会在做一个TF2版本和Jax/Flax版本。

猜你喜欢

转载自blog.csdn.net/weixin_42812353/article/details/111878548