Build a ChatGPT‑Scale Open‑Source Model with ColossalAI’s End‑to‑End RLHF Pipeline

In recent months, AI applications such as ChatGPT and GPT‑4 have sparked a new wave of large‑model development, prompting both industry giants and researchers to chase the technology.

Colossal‑AI, a leading open‑source AI infrastructure project, released ColossalChat—a practical open‑source implementation that follows the original ChatGPT technical roadmap, using LLaMA as the base model and completing the full RLHF pipeline (supervised fine‑tuning → reward model training → reinforcement learning fine‑tuning).

Key Offerings

Online demo for immediate model interaction without registration.

Complete RLHF training code for 7B and 13B parameter models.

Bilingual (Chinese‑English) dataset of ~104K Q&A pairs, collected from real social‑media queries and expanded via self‑instruct.

4‑bit quantized inference that runs a 7‑billion‑parameter model on a single 4 GB GPU.

Model weights that can be reproduced on a single server with modest compute.

RLHF Stages

Stage 1 – Supervised Fine‑Tuning (SFT) : fine‑tune the base LLaMA model on the bilingual dataset.

Stage 2 – Reward Model (RM) : train a reward model by ranking multiple responses to the same prompt.

Stage 3 – Reinforcement Learning (PPO) : generate experience using SFT, Actor, RM, and Critic models, then update the policy with combined strategy, value, and PTX losses to preserve pretrained knowledge.

Quick Start

Training the SFT model (4‑GPU example):

# Training with a 4‑GPU server
colossalai run --nproc_per_node=4 train_sft.py \
  --pretrain "/path/to/LLaMa-7B/" \
  --model 'llama' \
  --strategy colossalai_zero2 \
  --log_interval 10 \
  --save_path /path/to/Coati-7B \
  --dataset /path/to/data.json \
  --batch_size 4 \
  --accimulation_steps 8 \
  --lr 2e-5

Training the reward model (4‑GPU example):

# Training with a 4‑GPU server
colossalai run --nproc_per_node=4 train_reward_model.py \
  --pretrain "/path/to/LLaMa-7B/" \
  --model 'llama' \
  --strategy colossalai_zero2 \
  --dataset /path/to/datasets

RL training (8‑GPU example):

# Training with an 8‑GPU server
colossalai run --nproc_per_node=8 train_prompts.py prompts.csv \
  --strategy colossalai_zero2 \
  --pretrain "/path/to/Coati-7B" \
  --model 'llama' \
  --pretrain_dataset /path/to/dataset

After obtaining the final weights, low‑bit (4‑bit) GPTQ quantization can reduce inference memory to ~4 GB, enabling deployment on consumer‑grade GPUs (e.g., RTX 3060). Example inference command:

python server.py /path/to/pretrained --quant 4bit --gptq_checkpoint /path/to/coati-7b-4bit-128g.pt --gptq_group_size 128

Performance Optimizations

Colossal‑AI’s ZeRO optimizer with Gemini memory manager removes redundancy, allowing larger models with the same GPU memory, while LoRA low‑rank adaptation reduces fine‑tuning cost by updating only a small matrix.

Collaboration

Developers are encouraged to contribute via GitHub issues or pull requests, join the community Slack/WeChat groups, or contact the team at [email protected].