Package: chatopt
RLHF - Colossal-AIImplementation of RLHF (Reinforcement Learning with Human Feedback) powered by Colossal-AI. It supports distributed training and offloading, which can fit extremly large models. More details can be found in theblog.Training process (step 3)Installpipinstall.UsageThe main entrypoint isTrainer. We only support PPO trainer now. We support many training strategies:NaiveStrategy: simplest strategy. Train on single GPU.DDPStrategy: usetorch.nn.parallel.DistributedDataParallel. Train on multi GPUs.ColossalAIStrategy: use Gemini and Zero of ColossalAI. It eliminates model duplication on each GPU and supports offload. It's very useful when training large models on multi GPUs.Simplest usage:fromchatgpt.trainerimportPPOTrainerfromchatgpt.trainer.strategiesimportColossalAIStrategyfromchatgpt.nnimportGPTActor,GPTCritic,RewardModelfromcopyimportdeepcopyfromcolossalai.nn.optimizerimportHybridAdamstrategy=ColossalAIStrategy()withstrategy.model_init_context():# init your model here# load pretrained gpt2actor=GPTActor(pretrained='gpt2')critic=GPTCritic()initial_model=deepcopy(actor).cuda()reward_model=RewardModel(deepcopy(critic.model),deepcopy(critic.value_head)).cuda()actor_optim=HybridAdam(actor.parameters(),lr=5e-6)critic_optim=HybridAdam(critic.parameters(),lr=5e-6)# prepare models and optimizers(actor,actor_optim),(critic,critic_optim),reward_model,initial_model=strategy.prepare((actor,actor_optim),(critic,critic_optim),reward_model,initial_model)# load saved model checkpoint after preparingstrategy.load_model(actor,'actor_checkpoint.pt',strict=False)# load saved optimizer checkpoint after preparingstrategy.load_optimizer(actor_optim,'actor_optim_checkpoint.pt')trainer=PPOTrainer(strategy,actor,critic,reward_model,initial_model,actor_optim,critic_optim,...)trainer.fit(dataset,...)# save model checkpoint after fitting on only rank0strategy.save_model(actor,'actor_checkpoint.pt',only_rank0=True)# save optimizer checkpoint on all ranksstrategy.save_optimizer(actor_optim,'actor_optim_checkpoint.pt',only_rank0=False)For more details, seeexamples/.We also support training reward model with true-world data. Seeexamples/train_reward_model.py.FAQHow to save/load checkpointTo load pretrained model, you can simply use huggingface pretrained models:# load OPT-350m pretrained modelactor=OPTActor(pretrained='facebook/opt-350m')To save model checkpoint:# save model checkpoint on only rank0strategy.save_model(actor,'actor_checkpoint.pt',only_rank0=True)This function must be called afterstrategy.prepare().For DDP strategy, model weights are replicated on all ranks. And for ColossalAI strategy, model weights may be sharded, but all-gather will be applied before returning state dict. You can setonly_rank0=Truefor both of them, which only saves checkpoint on rank0, to save disk space usage. The checkpoint is float32.To save optimizer checkpoint:# save optimizer checkpoint on all ranksstrategy.save_optimizer(actor_optim,'actor_optim_checkpoint.pt',only_rank0=False)For DDP strategy, optimizer states are replicated on all ranks. You can setonly_rank0=True. But for ColossalAI strategy, optimizer states are sharded over all ranks, and no all-gather will be applied. So for ColossalAI strategy, you can only setonly_rank0=False. That is to say, each rank will save a cehckpoint. When loading, each rank should load the corresponding part.Note that different stategy may have different shapes of optimizer checkpoint.To load model checkpoint:# load saved model checkpoint after preparingstrategy.load_model(actor,'actor_checkpoint.pt',strict=False)To load optimizer checkpoint:# load saved optimizer checkpoint after preparingstrategy.load_optimizer(actor_optim,'actor_optim_checkpoint.pt')Todoimplement PPO fine-tuningimplement training reward modelsupport LoRAimplement PPO-ptx fine-tuningintegrate with Raysupport more RL paradigms, like Implicit Language Q-Learning (ILQL)Invitation to open-source contributionReferring to the successful attempts ofBLOOMandStable Diffusion, any and all developers and partners with computing powers, datasets, models are welcome to join and build an ecosystem with Colossal-AI, making efforts towards the era of big AI models from the starting point of replicating ChatGPT!You may contact us or participate in the following ways:Posting anissueor submitting aPRon GitHubJoin the Colossal-AI community onSlack,
andWeChatto share your ideas.Check out and fill in thecooperation proposalSend your proposal to emailcontact@hpcaitech.comThanks so much to all of our amazing contributors!Quick PreviewUp to 7.73 times faster for single server training and 1.42 times faster for single-GPU inferenceUp to 10.3x growth in model capacity on one GPUA mini demo training process requires only 1.62GB of GPU memory (any consumer-grade GPU)Increase the capacity of the fine-tuning model by up to 3.7 times on a single GPUKeep in a sufficiently high running speedCitations@article{Hu2021LoRALA,title={LoRA: Low-Rank Adaptation of Large Language Models},author={Edward J. Hu and Yelong Shen and Phillip Wallis and Zeyuan Allen-Zhu and Yuanzhi Li and Shean Wang and Weizhu Chen},journal={ArXiv},year={2021},volume={abs/2106.09685}}@article{ouyang2022training,title={Training language models to follow instructions with human feedback},author={Ouyang, Long and Wu, Jeff and Jiang, Xu and Almeida, Diogo and Wainwright, Carroll L and Mishkin, Pamela and Zhang, Chong and Agarwal, Sandhini and Slama, Katarina and Ray, Alex and others},journal={arXiv preprint arXiv:2203.02155},year={2022}}