Colossal-AI

Colossal-AI is a Python framework for developers and researchers training or serving large AI models on their own GPUs. It provides tools to spread workloads across GPUs and work within available memory.

• Distributed training: Combine data, tensor, pipeline and sequence parallelism through hybrid training, with automatic parallelism and the Zero Redundancy Optimizer (ZeRO) also available.
• Memory management: Gemini manages heterogeneous memory, while cached embeddings let recommendation models train larger embedding tables within a smaller GPU memory budget.
• Distributed inference: Use the framework's parallel components for inference as well as training.
• Model projects: Build on Colossal-LLaMA-2 for domain-specific language models or ColossalChat for a ChatGPT-style system with a reinforcement learning from human feedback (RLHF) pipeline; Open-Sora applies the framework to video generation.

Colossal-AI is open source under Apache 2.0. Single-GPU examples cover GPT-2 and PaLM, including GPT-2 training on an NVIDIA RTX 3080. Multi-GPU benchmarks use Llama-like models on NVIDIA H200 and B200 hardware.

The framework can run on your own hardware. HPC-AI Cloud separately provides hosted GPU environments, while HPC-AI Model APIs provide remote model access with inference running in the cloud.