$ show hpcaitech/ColossalAI
ColossalAI is a comprehensive distributed training and inference framework for large-scale deep learning models, providing multiple parallelism strategies (data, tensor, pipeline, sequence, mixture-of-experts), ZeRO memory optimization, automatic parallelization via FX tracing, and high-performance custom CUDA/Triton kernels. The architecture is library-based with a plugin-style booster for orchestrating training and inference, a shardformer for model-level parallelism policies, and a legacy subsystem for older APIs. It includes application-level tools: ColossalChat for RLHF (PPO, DPO, GRPO), ColossalQA for retrieval-augmented question answering, ColossalEval for LLM evaluation with GPT judges, and Colossal-LLaMA for domain-specific LLM training. The codebase is actively maintained with recent commits, but has notable technical drift due to outdated dependency pins in some examples and a large legacy namespace that coexists with newer implementations.
The LLM identified 3 files worth highlighting for reuse:
Abstract base class for training plugins defining the lifecycle methods (configure, prepare_dataloader, no_sync, enable_lora). This pattern allows clean separation of parallelism strategies (TorchDDP, Gemini, HybridParallel, etc.) and is reusable in any PyTorch training system. The contract is minimal and well-defined.
ShardFormer class that applies a policy-based model transformation for tensor and pipeline parallelism. It uses ModulePolicyDescription to replace layers with distributed counterparts (e.g., linear, attention, embedding). This pattern is reusable for any model that needs to be sharded across devices, with a clean separation of policy and execution.
Decorator-based registry for meta-profiling functions that estimate compute and memory costs of PyTorch operations. Each operation registers a handler that produces a TrainCycleItem (memory and compute for fwd/bwd). This pattern is reusable for building cost models for new op sets and is cleanly separated from the auto-parallel solver.