Low‑Cost Open‑Source Replication of ChatGPT Using Colossal‑AI

On February 14, researchers from UC Berkeley and the National University of Singapore released an open‑source, low‑cost implementation of the ChatGPT training pipeline, attracting widespread interest from the AI community.

ChatGPT’s training consists of three stages: (1) supervised fine‑tuning on a prompt‑response dataset, (2) reward‑model training using human‑ranked responses, and (3) reinforcement learning with proximal policy optimization (PPO) to align the model with human preferences.

These stages demand massive GPU memory (thousands of GB) and at least 64 × 80 GB A100 GPUs, making replication infeasible for most teams.

Colossal‑AI addresses these challenges with multi‑dimensional automatic parallelism, heterogeneous memory management, a large‑scale optimization library, and adaptive task scheduling, enabling efficient large‑model training and inference.

By leveraging ZeRO, Gemini, LoRA, and AutoChunk, Colossal‑AI cuts GPU memory usage roughly in half, allowing the 175‑billion‑parameter model to be trained on 32 GPUs instead of 64, and even provides single‑GPU and 4/8‑GPU variants for smaller models.

Performance tests show up to 7.73× speedup on a single server and 1.42× faster inference per GPU, while supporting up to 80 B‑parameter models on a single consumer‑grade GPU.

Colossal‑AI also offers an out‑of‑the‑box training script; a single line of code selects the system strategy, and the following commands launch the training for various scales:

# Training GPT2‑S on a single GPU torchrun --standalone --nproc_per_node 1 benchmark_gpt_dummy.py --model s --strategy colossalai_gemini_cpu --experience_batch_size 1 --train_batch_size 1 # Training GPT2‑XL on a 4‑GPU machine torchrun --standalone --nproc_per_node 4 benchmark_gpt_dummy.py --model xl --strategy colossalai_zero2 # Training a 175‑billion‑parameter model across 4 × 8‑GPU servers torchrun --nnodes 4 --nproc_per_node 8 \ --rdzv_id=$JOB_ID --rdzv_backend=c10d --rdzv_endpoint=$HOST_NODE_ADDR \ benchmark_gpt_dummy.py --model 175b --strategy colossalai_gemini_cpu --experience_batch_size 1 --train_batch_size 1

Further low‑level optimizations include LoRA (low‑rank adaptation) fine‑tuning, which trains only small A and B matrices, and the Zero+Gemini approach that eliminates memory redundancy and offloads optimizer states to CPU, allowing seamless scaling beyond a single GPU’s memory.

The complete code and detailed documentation are publicly available, aiming to democratize large‑scale LLM training and encourage broader community participation.