Unit Scaling是一种模型设计方法,能够让FP16和FP8等低精度数字格式更加易用。我们很高兴地宣布,我们发布了一个PyTorch库来促进Unit Scaling的使用。 
7月,Graphcore(拟未)在ICML上发表了论文《Unit Scaling:开箱即用的低精度训练》[1]。现在,我们发布软件工具,让更多人可以使用这种方法。
支持FP8的硬件的发展大幅提高了用户效率,例如Graphcore® C600 IPU处理器PCIe卡
[2]
。但是,简单地将较高精度的值转换为FP8值往往会导致性能降低。Unit Scaling解决了这一问题,为充分利用FP8硬件进行训练提供了一条简单的路径。

查看库文档[3]
阅读我们的ICML论文[4]
库操作演示
为了向用户展示如何在他们自己的模型中应用Unit Scaling,我们还发布了一个和库配套的notebook。它展示了在FP8下,使用和不使用Unit Scaling的nanoGPT模型训练。
https://ipu.dev/a2FqG0
只需一行代码——model = unit_scale(model)——用户就可以将他们的PyTorch模块转化为单位缩放模型。我们在notebook中通过训练以下模型说明了这一点:
在FP8中直接训练基础模型会导致明显的性能下降。不过,使用Unit Scaling可以恢复全精度。

这种单行转换可应用于任意PyTorch模型,与torch.compile[5]一起使用时,开销可忽略不计。
实施Unit Scaling
单行自动unit_scale()转换是一项实验性功能。我们建议大多数用户通过以下方式手动实施Unit Scaling。

您可以考虑这种导入PyTorch模块/函数的常见方法:
在此设置中,您可以这样实现Unit Scaling,首先添加:
然后将字母nn替换为uu,将F替换为U,让这样这些类和函数变成单位缩放的。例如:

我们的用户指南[6]中还介绍了运用Unit Scaling所需的一些其他注意事项。用户应特别注意正确缩放跳转/残差增量和损失函数。
使用库
Unit Scaling可通过以下方式安装:
尽管我们付出了诸多努力,但unit_scaling毕竟是一个新的库,我们不能保证它完全没有bug或功能已经完全齐备。不过我们将非常愿意为这个库的使用者提供帮助。

您可以通过我们的Slack频道[7]联系我们的开发者社区或提出GitHub问题[8]
[1]https://arxiv.org/abs/2303.11257
[2]https://www.graphcore.cn/c600-pcie%E5%8D%A1/
[3]https://github.com/graphcore-research/unit-scaling
[4]https://arxiv.org/abs/2303.11257
[5]https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html
[6]https://graphcore-research.github.io/unit-scaling/user_guide
[7]https://www.graphcore.ai/join-community
[8]https://github.com/graphcore-research/unit-scaling
获取更多Graphcore资讯,阅读深度技术文章,并与其他创新者们一起交流,请至中国官网graphcore.cn,以及关注Graphcore微信、微博和知乎创新社区。
Graphcore中国官网
Graphcore官方微信
Graphcore微博创新社区
Graphcore知乎创新社区
继续阅读
阅读原文