• 一個開源完整的基于 PyTorch 的 ChatGPT 等效實現(xiàn)流程，涵蓋所有 3 個階段，可以幫助你構建基于預訓練模型的 ChatGPT 式服務。

• 提供了一個迷你演示訓練過程供用戶試玩，它只需要 1.62GB 的 GPU 顯存，并且可能在單個消費級 GPU 上實現(xiàn)，單GPU模型容量最多提升10.3 倍。

• 與原始 PyTorch 相比，單機訓練過程最高可提升7.73 倍，單 GPU 推理速度提升 1.42 倍，僅需一行代碼即可調用。

• 在微調任務上，同樣僅需一行代碼，就可以在保持足夠高的運行速度的情況下，最多提升單GPU的微調模型容量3.7 倍。

• 提供多個版本的單 GPU 規(guī)模、單節(jié)點多 GPU 規(guī)模和原始 1750 億參數(shù)規(guī)模。還支持從 Hugging Face 導入 OPT、GPT-3、BLOOM 和許多其他預訓練的大型模型到你的訓練過程中。

from chatgpt.nn import GPTActor, GPTCritic, RewardModel
from chatgpt.trainer import PPOTrainer
from chatgpt.trainer.strategies import ColossalAIStrategy

strategy = ColossalAIStrategy(stage=3, placement_policy='cuda')

with strategy.model_init_context():
    actor = GPTActor().cuda()
    critic = GPTCritic().cuda()
    initial_model = deepcopy(actor).cuda()
    reward_model = RewardModel(deepcopy(critic.model)).cuda()

trainer = PPOTrainer(strategy, actor, critic, reward_model, initial_model,...)
trainer.fit(prompts)