选自medium

作者：KaiyuYue

机器之心编译

编辑：陈

训练大模型时，如何优雅地减少GPU内存消耗？你不妨试试这个TorchShard库，兼具模型并行与数据并行等特点，还具有与PyTorch相同的API设计。

模型并行性能够促进视觉任务的性能。但是目前，还没有一个标准库可以让我们像采用混合精度等其他SOTA技术那样轻松地采用模型并行性。最近，马里兰大学帕克分校计算机科学系的研究者KaiyuYue开源了一个工具TorchShard，这是一个轻量级的引擎，用于将PyTorch张量切片成并行的shard。当模型拥有大量的线性层（例如BERT、GPT）或者很多类（数百万）时，TorchShard可以减少GPU内存并扩展训练规模，它具有与PyTorch相同的API设计。项目


转载请注明地址:http://www.nasaoa.com/nsly/9474.html