📖 AI 原生基础设施:从云原生到 AI 原生的架构演进指南

PatrickStar

通过分块动态内存管理和异构卸载,实现大规模预训练模型并行训练的高效框架。

Tencent · Since 2021-04-02
正在加载评分...

详细介绍

PatrickStar 是一个基于 PyTorch 的并行训练框架,面向大规模预训练模型(PTM)。它通过分块(chunk-based)的动态内存管理与异构训练策略,将非计算中所需的数据高效地卸载到 CPU,从而在更少 GPU 的条件下训练更大的模型,降低 OOM 风险并提升可用性。项目由腾讯 NLP/WeChat AI 团队维护,旨在将大模型训练能力普惠化。

主要特性

  • 分块动态内存调度:按计算窗口管理激活值与参数,减少 GPU 内存占用。
  • 异构卸载:将大部分非即时计算数据搬到 CPU,支持混合 CPU/GPU 内存使用。
  • 高效通信与扩展性:为多卡与多节点场景优化集体通信,兼顾性能与可扩展性。
  • 与 PyTorch 兼容:使用类似 DeepSpeed 的配置格式,降低迁移成本。

使用场景

适用于预训练和大规模微调场景,尤其在硬件受限的集群上想训练数十亿至百亿级参数模型时表现突出。也适合性能对比基准、深度学习框架研究与高校/科研单位的教学与实验环境。

技术特点

PatrickStar 采用 chunk-based 内存管理与动态调度算法,在运行时只保留当前计算所需分块,其他分块可异步迁移;同时优化集体通信以提高多卡效率。项目提供基准与示例脚本,能在 V100/A100 等常见 GPU 集群上训练数十亿参数模型,许可证为 BSD-3-Clause。

PatrickStar
评分详情
🏋️ 模型训练 ⚡ 优化 🏗️ 框架 📁 项目 🖥️ 机器学习平台