Colossal-AI Admin & Developer Guide

This guide is for the engineers, ML platform operators, and architects who deploy, tune, and extend Colossal-AI. It documents the system architecture, the complete configuration surface of every parallelism plugin and the inference engine, deployment and operations, security considerations, and the internals needed to understand or contribute to the codebase.

Colossal-AI (colossalai, version 0.5.0 at the documented commit) is a unified deep-learning system for large-scale parallel training and inference. Its job is to take a model expressed in ordinary PyTorch and distribute it across many GPUs and nodes using data, tensor, pipeline, sequence, and expert parallelism — composed freely — plus ZeRO-style memory sharding and heterogeneous (CPU/NVMe) offloading.

For installation and day-to-day usage, see the companion User Guide.

Table of Contents

1. Executive Summary
2. System Architecture
3. Component Inventory
4. System Requirements
5. Plugin Configuration Reference
6. Inference Engine Configuration
7. Environment Variables
8. Deployment
9. Operations
10. Performance Tuning
11. Security Analysis
12. Repository Structure & Internals
13. Module Reference
14. Code Patterns & Conventions
15. Testing Strategy
16. Known Issues & Potential Improvements
17. Contributing
18. Quick Reference

1. Executive Summary

Operational character: Colossal-AI is a library and toolkit, not a long-running service. There is no server, database, or network listener in the core (the only HTTP surfaces are demo/web-chatbot servers shipped under applications/). Operating it means provisioning GPU clusters, configuring the distributed launcher, and tuning plugin parameters — not running a daemon.

2. System Architecture

Layered view

+-----------------------------------------------------------------------+
|  Applications  (applications/)                                          |
|  ColossalChat (RLHF) | Colossal-LLaMA | ColossalMoE | Eval | QA (RAG)  |
+-----------------------------------------------------------------------+
|  High-level training API                                               |
|  Booster  --->  Plugin (TorchDDP/FSDP, LowLevelZero, Gemini,           |
|                          HybridParallel, MoeHybridParallel)            |
|  Inference: InferenceEngine + InferenceConfig + RequestHandler          |
+-----------------------------------------------------------------------+
|  Parallelism & memory primitives                                       |
|  shardformer (TP for transformers) | pipeline (schedules, stage mgr)   |
|  zero (low_level, gemini/chunk)     | moe | auto_parallel | tensor      |
|  nn/optimizer (HybridAdam, CPUAdam, fused, GaLore, ...)                 |
|  checkpoint_io (sharded/safetensors) | amp | quantization | lazy        |
+-----------------------------------------------------------------------+
|  Cluster & device layer                                                |
|  cluster (ProcessGroupMesh, DistCoordinator) | device | accelerator    |
|  context | logging | initialize (launch_*)                             |
+-----------------------------------------------------------------------+
|  Native acceleration                                                   |
|  extensions/  (C++/CUDA kernels: fused adam, layernorm, softmax,       |
|               flash attention, moe) built via PyTorch BuildExtension   |
+-----------------------------------------------------------------------+
|  PyTorch  ->  torch.distributed (NCCL)  ->  NVIDIA GPUs / CUDA          |
+-----------------------------------------------------------------------+

Runtime flow of a training job

colossalai run --nproc_per_node N train.py
        |
        v  (spawns N processes per node, sets RANK/LOCAL_RANK/WORLD_SIZE/MASTER_ADDR/PORT)
train.py: colossalai.launch_from_torch()
        |        -> dist.init_process_group(backend="nccl", init_method="tcp://host:port")
        |        -> set CUDA device, seed, CUDA_DEVICE_MAX_CONNECTIONS=1
        v
plugin = SomePlugin(...)           # builds a ProcessGroupMesh defining DP/TP/PP/EP axes
booster = Booster(plugin=plugin)
model,opt,... = booster.boost(...) # plugin.configure(): shard params, wrap optimizer,
        |                          #   apply Shardformer policies, set up checkpoint IO
        v
training loop: booster.backward()  (or booster.execute_pipeline() when pp_size>1)
        |
        v
booster.save_model(..., shard=True)  # sharded/safetensors checkpoint

The key structural decision is the ProcessGroupMesh (colossalai/cluster): a multi-dimensional grid of process groups. Each plugin lays out the world's GPUs along named axes (data-parallel, tensor-parallel, pipeline-parallel, ZeRO, expert-parallel) and the rest of the system communicates along those axes.

3. Component Inventory

Note

colossalai/legacy holds the previous-generation API (config-file-driven Engine/Trainer, the 1D/2D/2.5D/3D layer implementations, old inference/serving). New code should use the Booster API; legacy is retained for backward compatibility.

4. System Requirements

Core runtime dependencies (requirements/requirements.txt)

torch>=2.2.0,<=2.5.1, numpy, tqdm, psutil, packaging, rich, click, fabric, contexttimer, ninja, safetensors, einops, pydantic, ray, sentencepiece, protobuf, transformers==4.51.3, peft>=0.7.1,<=0.13.2, bitsandbytes>=0.39.0, rpyc==6.0.0, fastapi, uvicorn==0.29.0, galore_torch, diffusers==0.29.0.

Important

transformers is pinned exactly (==4.51.3). Shardformer's model policies target specific transformer internals, so upgrading transformers independently can break tensor-parallel sharding. Treat this pin as load-bearing.

5. Plugin Configuration Reference

Plugins are the primary configuration surface — most "tuning" in Colossal-AI is choosing a plugin and setting its constructor arguments. Below is the complete, source-derived argument list for each plugin, with defaults.

Note

These arguments are read directly from the plugin __init__ signatures in colossalai/booster/plugin/. There is no YAML/JSON config file for the Booster API; configuration is Python keyword arguments.

5.1 TorchDDPPlugin

5.2 TorchFSDPPlugin

Wraps PyTorch-native Fully Sharded Data Parallel (requires torch >= 1.12). Accepts FSDP's configuration objects — cpu_offload (CPUOffload), mixed_precision (MixedPrecision), sharding_strategy (ShardingStrategy), backward_prefetch (BackwardPrefetch), and the full-state-dict config — passed through to torch.distributed.fsdp.FullyShardedDataParallel.

Caution

Importing TorchFSDPPlugin on torch < 1.12 raises RuntimeError: FSDP is not supported while torch version under 1.12.0.

5.3 LowLevelZeroPlugin (ZeRO-1 / ZeRO-2)

5.4 GeminiPlugin (ZeRO-3 + heterogeneous memory)

5.5 HybridParallelPlugin (TP + PP + SP + ZeRO)

Required: tp_size (int), pp_size (int). Selected arguments:

Important

tp_size × pp_size × dp_degree must equal the total world size, where dp_degree is inferred. Mis-sized meshes fail at plugin construction. make_vocab_size_divisible_by matters: an embedding vocab not divisible by tp_size will be padded — verify your tokenizer vocab against this.

5.6 MoeHybridParallelPlugin

Extends HybridParallelPlugin with expert parallelism. Required: tp_size, pp_size, ep_size (expert-parallel degree). Adds moe_dp_outside (default True). All other arguments mirror HybridParallelPlugin. Note overlap_communication defaults to False here (vs True for HybridParallel).

6. Inference Engine Configuration

InferenceConfig (a dataclass in colossalai/inference/config.py) is consumed by InferenceEngine. Complete field list:

InferenceEngine(model_or_path, tokenizer=None, inference_config=None, verbose=False, model_policy=None) accepts an nn.Module, a path string, or a DiffusionPipeline. The main method is engine.generate(request_ids=None, prompts=None).

7. Environment Variables

Colossal-AI is configured mostly through Python arguments, not environment variables. The few variables that exist:

Note

colossalai/initialize.py unconditionally sets CUDA_DEVICE_MAX_CONNECTIONS=1 at import time. This is intentional (it makes computation/communication overlap deterministic, per NVIDIA Megatron-LM guidance) and should not be overridden.

8. Deployment

Docker (recommended for clusters)

The provided docker/Dockerfile builds a GPU image:

• Base: hpcaitech/cuda-conda:12.1
• Python 3.10, PyTorch 2.1.2 + CUDA 12.1 (conda)
• NVIDIA Apex compiled with C++/CUDA extensions
• RDMA/InfiniBand support packages installed
• Passwordless SSH preconfigured (for multi-node colossalai run)
• Colossal-AI installed with BUILD_EXT=1

cd ColossalAI
docker build -t colossalai ./docker            # requires NVIDIA Container Runtime as default
docker run -ti --gpus all --rm --ipc=host colossalai bash

Important

--ipc=host is required so PyTorch dataloader workers and NCCL can use sufficient shared memory. Omitting it causes mysterious dataloader/NCCL hangs. The Dockerfile builds kernels, so the build host must expose a GPU (NVIDIA Container Runtime as default docker build runtime).

Prebuilt images are published on DockerHub per release.

Bare-metal cluster

1. Provision Linux nodes with matching CUDA/driver versions and GPUs (CC ≥ 7.0).
2. Install PyTorch ≥ 2.2, then BUILD_EXT=1 pip install colossalai.
3. Ensure passwordless SSH between nodes for multi-node launch.
4. Create a hostfile and launch with colossalai run --hostfile ....

Multi-node launch internals

colossalai run (in colossalai/cli/launcher/) parses the hostfile into a HostInfoList, applies --include/--exclude filters, then uses a MultiNodeRunner to SSH into each host and start torchrun-style processes with the computed rank assignments. The master address/port and --extra_launch_args are forwarded to the underlying torch launcher.

Tip

For InfiniBand/RDMA clusters, base your image on the provided Dockerfile (it installs libibverbs, librdmacm, perftest, etc.) so NCCL can use the high-speed fabric.

9. Operations

Because Colossal-AI is a training library, "operations" means running and monitoring jobs rather than a service.

Logging

get_dist_logger() returns a process-aware logger. Always pass ranks=[0] for messages you want once per job rather than once per GPU:

from colossalai.logging import get_dist_logger
logger = get_dist_logger()
logger.info("Epoch complete", ranks=[0])

10. Performance Tuning

A practical escalation ladder, from cheapest to most invasive:

1. Precision — use bf16 where supported; consider use_fp8=True for comm/compute on Hopper+.
2. ZeRO — LowLevelZeroPlugin(stage=2); add cpu_offload=True if optimizer state is the bottleneck.
3. Gemini — for models exceeding GPU memory, tune offload_optim_frac / offload_param_frac and placement_policy="auto".
4. Communication overlap — keep overlap_communication=True, overlap_p2p=True, overlap_allgather=True; tune bucket sizes (reduce_bucket_size_in_m, ddp_bucket_cap_mb, zero_bucket_size_in_m).
5. Shardformer optimizations — enable_all_optimization=True (fused norm, flash attention, JIT) on transformer models.
6. Tensor + pipeline parallelism — HybridParallelPlugin(tp_size, pp_size); choose pp_style (1f1b vs interleaved vs zero-bubble) and tune num_microbatches.
7. Activation checkpointing — pass a gradient_checkpoint_config to trade compute for memory.
8. Lazy init — use colossalai.lazy.LazyInitContext to construct huge models on the meta device, avoiding host-memory OOM before sharding.

Tip

The benchmark in the README shows the value of composing strategies: for a 7B model on 8×B200, zero1(dp2)+tp2+pp4 reached ~25.8 samp/s vs ~17.1 samp/s for zero2(dp8) on H200. The winning configuration depends on model size, sequence length, and interconnect — measure on your hardware.

Caution

Pipeline parallelism (pp_size>1) changes the training loop contract: you must use booster.execute_pipeline(...). Setting num_microbatches too low underfills the pipeline (idle "bubble"); too high inflates memory. Start near num_microbatches ≈ 4 × pp_size.

11. Security Analysis

Colossal-AI is a compute library, so its threat surface differs from a network service. Key considerations:

Trust model

• No authentication or network listener in the core. The core library does not open sockets except via torch.distributed (NCCL) on the cluster's trusted interconnect. Treat the training network as a trusted, isolated environment — NCCL traffic is unauthenticated and unencrypted.
• Multi-node launch uses SSH. colossalai run SSHes into hosts from the hostfile. Secure these with key-based auth and a restricted network. The Docker image enables passwordless SSH for intra-cluster convenience — do not expose those containers to untrusted networks.

Code / model-loading risks

• Checkpoint loading executes deserialization. Loading PyTorch .pt/pickle checkpoints can execute arbitrary code. Prefer use_safetensors=True (the safetensors format is non-executable) for any checkpoint of uncertain provenance.
• Remote model downloads. Applications pull models/datasets from Hugging Face Hub and ModelScope. Pin and verify model sources; untrusted custom-code models (trust_remote_code) run arbitrary Python.
• use_cuda_kernel/runtime build compiles native code with your toolchain — ensure the source tree is the trusted upstream.

Application web surfaces

The only HTTP endpoints in the tree are demo servers under applications/ (ColossalChat web chatbot, ColossalQA webui — FastAPI/uvicorn) and a single POST /check_correctness code-verifier used in RL reward computation (applications/ColossalChat/start_code_verifier.py).

Warning

The code-verifier (/check_correctness) and the RL "code reward" path execute model-generated code to score it. This is sandbox-sensitive: run it only in an isolated, untrusted-code-safe environment (container with no secrets, no network egress, resource limits). Never expose these demo servers to the public internet without authentication and sandboxing — they are examples, not hardened services.

FP8 / experimental paths

Several FP8 and plugin-combination paths have open correctness bugs (see Known Issues). Validate numerics on a small run before committing a large job to an experimental flag.

12. Repository Structure & Internals

ColossalAI/
├── colossalai/            # core library
│   ├── booster/           # Booster + plugins (the main API)
│   ├── shardformer/       # transformer TP/SP sharding via policies
│   ├── pipeline/          # pipeline schedules, stage manager, P2P
│   ├── zero/              # low_level (ZeRO 1/2), gemini (ZeRO-3 + chunks)
│   ├── moe/               # mixture-of-experts
│   ├── auto_parallel/     # automatic parallelization solver
│   ├── tensor/            # ColoTensor, d_tensor, moe_tensor, padded_tensor
│   ├── nn/                # optimizers, losses, lr schedulers
│   ├── checkpoint_io/     # sharded/safetensors/async checkpoint IO
│   ├── inference/         # inference engine, config, KV-cache
│   ├── cluster/           # ProcessGroupMesh, DistCoordinator
│   ├── device/            # device mesh / topology
│   ├── accelerator/       # hardware abstraction
│   ├── quantization/      # FP8 hooks, bitsandbytes
│   ├── lazy/              # meta-device lazy init
│   ├── amp/               # mixed precision
│   ├── fx/                # graph tracing/profiling
│   ├── cli/               # colossalai run / check
│   ├── initialize.py      # launch_* entry points
│   └── legacy/            # previous-generation API (config-driven Engine/Trainer)
├── extensions/            # native CUDA/C++ kernels + pybind
├── applications/          # ColossalChat, Colossal-LLaMA, ColossalMoE, Eval, QA
├── examples/              # per-model training/inference examples
├── tests/                 # pytest suite (mirrors colossalai/ layout)
├── docker/                # Dockerfile
├── docs/                  # source docs (en + zh-Hans)
├── requirements/          # requirements.txt
└── setup.py               # build, BUILD_EXT, console_scripts

How sharding actually happens (Shardformer)

enable_all_optimization=True (or a custom_policy) drives Shardformer (colossalai/shardformer). A Policy maps a Hugging Face model's submodules to sharded replacements (column/row-parallel linear, fused/sequence-parallel attention, fused normalization). ShardConfig records the TP group, sequence-parallel mode, and which optimizations are on. This is why transformers is version-pinned: policies reference specific module class names and forward signatures.

Pipeline schedules

colossalai/pipeline/schedule implements OneForwardOneBackwardSchedule (1F1B), InterleavedSchedule, and ZeroBubbleVPipeScheduler. PipelineStageManager tracks which layers live on each stage. booster.execute_pipeline delegates to the plugin, which runs the selected schedule over micro-batches.

ZeRO and Gemini chunks

colossalai/zero/low_level implements ZeRO-1/2 by sharding optimizer state (and optionally gradients) across the DP group. colossalai/zero/gemini implements ZeRO-3 with a chunk abstraction: parameters are grouped into fixed-size chunks that the placement_policy migrates between GPU, CPU, and NVMe based on a live memory tracer (gemini/memory_tracer). HybridAdam/CPUAdam exist precisely to run optimizer math wherever the chunk currently lives.

13. Module Reference

The Booster surface

14. Code Patterns & Conventions

• Plugin = strategy object. All plugins subclass Plugin (ABC) and implement configure, supported_devices, supported_precisions, control_precision/device/checkpoint_io, support_no_sync/lora, no_sync, enable_lora, prepare_dataloader, get_checkpoint_io. Pipeline plugins also implement execute_pipeline (PipelinePluginBase); DP plugins inherit dataloader creation from DPPluginBase.
• Wrappers over inheritance. Boosted models/optimizers are ModelWrapper/OptimizerWrapper instances (colossalai/interface) with .unwrap() to recover the underlying object — important when interoperating with raw PyTorch/HF code.
• Rank-aware logging. Use get_dist_logger(...).info(msg, ranks=[0]) to avoid N-way log spam.
• Capability negotiation. Booster.__init__ asks the plugin control_device(), control_precision(), control_checkpoint_io() to decide who owns each concern — this is why setting device=/mixed_precision= on the Booster is ignored when the plugin controls them (a warning is logged).
• Native extensions are optional. Kernels build lazily; pure-Python fallbacks exist for most paths. BUILD_EXT=1 forces ahead-of-time compilation.
• Process-group mesh. Parallelism axes are coordinates in a ProcessGroupMesh; this keeps DP/TP/PP/EP group creation consistent across plugins.

15. Testing Strategy

The tests/ tree mirrors colossalai/ (e.g. tests/test_booster, test_shardformer, test_zero, test_pipeline, test_checkpoint_io, test_moe, test_infer, test_optimizer, test_fp8, test_lora). Tests are GPU- and distributed-heavy: many spawn multiple processes via the helpers in colossalai/testing and tests/kit (which includes a model_zoo of reference models across transformers, torchvision, timm, diffusers, etc.).

Note

Most of tests/ cannot run on a CPU-only or single-GPU box meaningfully — they validate cross-rank correctness. Reproducing CI locally requires a multi-GPU node.

16. Known Issues & Potential Improvements

Open issues observed in the community (GitHub)

These reflect real, currently-open bug reports — useful as a "tread carefully" map of experimental combinations:

Important

The common thread is FP8 and aggressive offload combinations are still maturing. For production runs, validate your exact (plugin, optimizer, precision, offload) tuple on a short job before scaling.

Suggested improvements (from code reading)

• Extension build UX. Kernel build failures surface as terse RuntimeErrors at first use. A pre-flight colossalai check that compiles and validates each extension would catch toolchain issues earlier.
• Config validation. Because configuration is free-form keyword args, an invalid (tp_size, pp_size, world_size) combination or unsupported precision for a plugin is only caught at runtime. A validation layer that fails fast with actionable messages would reduce wasted cluster time.
• Legacy surface area. colossalai/legacy is large and overlaps the modern API. Clear deprecation timelines would reduce confusion for new users who find legacy tutorials.
• Windows/CPU developer story. The all-or-nothing Linux+GPU requirement makes contribution onboarding hard; a CPU-mock test tier would help.

17. Contributing

• License: Apache 2.0.
• Contribution flow: fork → branch → PR against main, following CONTRIBUTING.md. Pre-commit hooks are configured (pre-commit is a dev dependency; formatting via black/autoflake, isort).
• CI runs on GitHub Actions; tests require GPUs for the distributed paths.
• For substantial proposals, the project also accepts proposals via [email removed] and discussion on the forum/Slack.

To cite the system, use the ICPP '23 BibTeX entry in README.md ("Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training").

18. Quick Reference

Install

pip install colossalai                 # no kernel build
BUILD_EXT=1 pip install colossalai     # build CUDA kernels now
pip install colossalai-nightly         # weekly main build
colossalai check -i                    # verify environment

Launch

colossalai run --nproc_per_node 8 train.py
colossalai run --hostfile hosts --master_addr h0 --nproc_per_node 8 train.py

Minimal training scaffold

import colossalai
from colossalai.booster import Booster
from colossalai.booster.plugin import HybridParallelPlugin
from colossalai.nn.optimizer import HybridAdam

colossalai.launch_from_torch()
plugin = HybridParallelPlugin(tp_size=2, pp_size=2, zero_stage=1, precision="bf16",
                              num_microbatches=8, enable_all_optimization=True)
booster = Booster(plugin=plugin)
model, optimizer, *_ = booster.boost(model, HybridAdam(model.parameters(), lr=1e-4))

Plugin selection cheat-sheet

Checkpoint

booster.save_model(model, "ckpt/model", shard=True, use_safetensors=True)
booster.load_model(model, "ckpt/model")