# 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](hpcaitech__ColossalAI__UserGuide.md).

---

## Table of Contents

1. [Executive Summary](#1-executive-summary)
2. [System Architecture](#2-system-architecture)
3. [Component Inventory](#3-component-inventory)
4. [System Requirements](#4-system-requirements)
5. [Plugin Configuration Reference](#5-plugin-configuration-reference)
6. [Inference Engine Configuration](#6-inference-engine-configuration)
7. [Environment Variables](#7-environment-variables)
8. [Deployment](#8-deployment)
9. [Operations](#9-operations)
10. [Performance Tuning](#10-performance-tuning)
11. [Security Analysis](#11-security-analysis)
12. [Repository Structure & Internals](#12-repository-structure--internals)
13. [Module Reference](#13-module-reference)
14. [Code Patterns & Conventions](#14-code-patterns--conventions)
15. [Testing Strategy](#15-testing-strategy)
16. [Known Issues & Potential Improvements](#16-known-issues--potential-improvements)
17. [Contributing](#17-contributing)
18. [Quick Reference](#18-quick-reference)

---

## 1. Executive Summary

| Aspect | Detail |
|---|---|
| What it is | A PyTorch-based distributed training & inference framework. |
| Distribution unit | One OS process per GPU, coordinated by `torch.distributed` (NCCL backend). |
| Central abstraction | `Booster` + a `Plugin` that encapsulates a parallelism strategy. |
| Parallelism supported | Data, 1D/2D/2.5D/3D tensor, pipeline (incl. zero-bubble), sequence, expert (MoE), and ZeRO (stages 1–3 / Gemini). |
| Memory management | ZeRO sharding + chunk-based heterogeneous offload to CPU and NVMe (Gemini). |
| Acceleration | Fused CUDA/C++ kernels (built lazily or with `BUILD_EXT=1`), flash attention, FP8, CUDA graphs. |
| Inference | Dedicated engine with continuous batching, paged KV-cache, speculative decoding. |
| Applications | RLHF chat (ColossalChat), domain LLM (Colossal-LLaMA), MoE, evaluation, RAG Q&A. |
| Platform | Linux + NVIDIA GPUs (CC ≥ 7.0). Windows blocked. |
| License | Apache 2.0. |

**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

| Component | Package | Responsibility |
|---|---|---|
| Launch / init | `colossalai/initialize.py` | `launch`, `launch_from_torch/slurm/openmpi`; sets up `torch.distributed`. |
| Booster | `colossalai/booster` | High-level training facade; precision, device, checkpoint orchestration. |
| Plugins | `colossalai/booster/plugin` | Parallelism strategies (the main config surface). |
| Shardformer | `colossalai/shardformer` | Tensor-parallel/sequence-parallel sharding of HF transformer models via policies. |
| Pipeline | `colossalai/pipeline` | Pipeline schedules (1F1B, interleaved, zero-bubble), stage manager, P2P comm. |
| ZeRO | `colossalai/zero` | `low_level` (ZeRO-1/2) and `gemini` (ZeRO-3 + chunk-based offload). |
| MoE | `colossalai/moe` | Mixture-of-Experts layers and expert parallelism. |
| Auto-parallel | `colossalai/auto_parallel` | Automatic parallelization solver (tensor-shard strategies). |
| Optimizers | `colossalai/nn/optimizer` | CPU/GPU/fused/distributed optimizers. |
| Checkpoint IO | `colossalai/checkpoint_io` | Sharded, safetensors, async, distributed-tensor-aware checkpointing. |
| Cluster | `colossalai/cluster` | `ProcessGroupMesh`, `DistCoordinator`. |
| Device | `colossalai/device` | Device mesh and topology. |
| Accelerator | `colossalai/accelerator` | Hardware abstraction (CUDA/CPU; NPU/XPU stubs). |
| Quantization | `colossalai/quantization` | FP8 hooks, bitsandbytes integration. |
| Lazy init | `colossalai/lazy` | Meta-device lazy initialization to build huge models without OOM. |
| Inference | `colossalai/inference` | Engine, config, request handler, KV-cache, batch bucket. |
| Extensions | `extensions/` | Native CUDA/C++ kernels + Python bindings. |
| CLI | `colossalai/cli` | `colossalai run` (launcher) + `colossalai check`. |

> [!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

| Requirement | Value | Source |
|---|---|---|
| OS | Linux only (WSL2 for Windows) | `setup.py` raises `RuntimeError` on Windows |
| Python | ≥ 3.7 (`python_requires=">=3.6"` in setup, README states ≥ 3.7) | `setup.py`, README |
| PyTorch | ≥ 2.2 (runtime pin `torch>=2.2.0,<=2.5.1`) | `requirements/requirements.txt` |
| CUDA | ≥ 11.0 | README |
| GPU compute capability | ≥ 7.0 (V100, RTX 20-series and newer) | README |
| Build tools (for kernels) | `ninja`, matching CUDA toolkit, C++ compiler | `extensions/`, Dockerfile |

### 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`

| Argument | Default | Purpose |
|---|---|---|
| `broadcast_buffers` | `True` | Broadcast module buffers at forward start. |
| `bucket_cap_mb` | `25` | DDP gradient bucket size (MB). |
| `find_unused_parameters` | `False` | Detect params not used in forward. |
| `check_reduction` | `False` | Debug: verify reduction occurred. |
| `gradient_as_bucket_view` | `False` | Use bucket views to save memory. |
| `static_graph` | `False` | Enable DDP static-graph optimization. |
| `fp8_communication` | `False` | Compress gradient comm to FP8. |

### 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)

| Argument | Default | Purpose |
|---|---|---|
| `stage` | `1` | `1` shards optimizer state; `2` also shards gradients. |
| `precision` | `'fp16'` | `fp16` / `bf16` / `fp32`. |
| `initial_scale` | `2**32` | Initial loss-scale for fp16. |
| `min_scale` / `max_scale` | `1` / `2**32` | Loss-scale bounds. |
| `growth_factor` / `backoff_factor` | `2` / `0.5` | Loss-scale adjustment factors. |
| `growth_interval` | `1000` | Steps between scale growth attempts. |
| `hysteresis` | `2` | Overflow tolerance before backing off. |
| `max_norm` | `0.0` | Gradient clipping max-norm (`0` = off). |
| `norm_type` | `2.0` | Norm order for clipping. |
| `reduce_bucket_size_in_m` | `12` | Gradient reduce bucket (millions of elements). |
| `communication_dtype` | `None` | Dtype for gradient comm. |
| `overlap_communication` | `True` | Overlap comm with compute. |
| `overlap_allgather` | `False` | Overlap parameter all-gather. |
| `cpu_offload` | `False` | Offload optimizer state to CPU. |
| `master_weights` | `True` | Keep fp32 master weights. |
| `cast_inputs` | `True` | Auto-cast inputs to compute dtype. |
| `fp8_communication` / `use_fp8` | `False` / `False` | FP8 comm / FP8 compute. |
| `extra_dp_size` | `1` | Extra data-parallel replication factor. |

### 5.4 `GeminiPlugin` (ZeRO-3 + heterogeneous memory)

| Argument | Default | Purpose |
|---|---|---|
| `chunk_config_dict` | `None` | Manual chunk configuration. |
| `chunk_init_device` | `None` | Device for chunk initialization. |
| `placement_policy` | `'static'` | `static` or `auto` placement of chunks across GPU/CPU. |
| `enable_gradient_accumulation` | `False` | Allow grad accumulation under Gemini. |
| `max_prefetch` | `0` | Number of chunks to prefetch. |
| `shard_param_frac` | `1.0` | Fraction of params to shard. |
| `offload_optim_frac` | `0.0` | Fraction of optimizer state offloaded to CPU. |
| `offload_param_frac` | `0.0` | Fraction of params offloaded to CPU (with `static` policy). |
| `warmup_non_model_data_ratio` | `0.8` | Memory budget for non-model data during warmup. |
| `steady_cuda_cap_ratio` | `0.9` | Steady-state GPU memory cap. |
| `precision` | `'fp16'` | `fp16` or `bf16` (only these two supported). |
| `master_weights` | `True` | Keep fp32 master weights. |
| `pin_memory` | `False` | Pin CPU memory for faster transfer. |
| `force_outputs_fp32` | `False` | Cast outputs to fp32. |
| `strict_ddp_mode` | `False` | Enforce DDP-style replication. |
| `search_range_m` / `min_chunk_size_m` | `32` / `32` | Chunk-size search range / minimum (millions). |
| `hidden_dim` | `None` | Model hidden dim hint for chunk sizing. |
| `gpu_margin_mem_ratio` | `0.0` | Reserved GPU memory margin. |
| `initial_scale` … `max_scale` | `2**16` … `2**32` | Loss-scale schedule (same fields as above). |
| `max_norm` / `norm_type` | `0.0` / `2.0` | Gradient clipping. |
| `tp_size` | `1` | Tensor-parallel degree combined with Gemini. |
| `extra_dp_size` | `1` | Extra DP factor. |
| `enable_all_optimization` | `False` | Enable Shardformer optimizations. |
| `enable_fused_normalization` | `False` | Fused LayerNorm. |
| `enable_flash_attention` | `False` | Flash attention. |
| `enable_sequence_parallelism` | `False` | Sequence parallelism. |
| `enable_jit_fused` | `False` | JIT-fused ops. |
| `enable_async_reduce` | `True` | Asynchronous reduction. |
| `use_fp8` / `fp8_communication` | `False` / `False` | FP8 compute / comm. |
| `verbose` | `False` | Verbose logging. |

### 5.5 `HybridParallelPlugin` (TP + PP + SP + ZeRO)

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

| Argument | Default | Purpose |
|---|---|---|
| `tp_size` | (required) | Tensor-parallel degree. |
| `pp_size` | (required) | Pipeline-parallel degree. |
| `sp_size` | `None` | Sequence-parallel degree. |
| `precision` | `'fp16'` | Compute precision. |
| `zero_stage` | `0` | ZeRO stage layered on top (`0`–`2`). |
| `enable_all_optimization` | `False` | Turn on fused norm + flash attn + JIT together. |
| `enable_fused_normalization` | `False` | Fused LayerNorm. |
| `enable_flash_attention` | `False` | Flash attention. |
| `enable_jit_fused` | `False` | JIT fusion. |
| `enable_sequence_parallelism` | `False` | Sequence parallelism. |
| `sequence_parallelism_mode` | `None` | One of `split_gather`, `ring`, `all_to_all`, `ring_attn`. |
| `parallel_output` | `True` | Keep output sharded across TP ranks. |
| `num_microbatches` / `microbatch_size` | `None` / `None` | Pipeline micro-batching (set one when `pp_size>1`). |
| `pp_style` | `'1f1b'` | Pipeline schedule: `1f1b`, interleaved, or zero-bubble. |
| `num_model_chunks` | `1` | Chunks per stage for interleaved schedule. |
| `scheduler_nodes` | `None` | Custom zero-bubble schedule nodes. |
| `num_layers_per_stage` | `None` | Manual layer-to-stage assignment. |
| `gradient_checkpoint_config` | `None` | Activation checkpointing config. |
| `zero_bucket_size_in_m` | `12` | ZeRO comm bucket (millions). |
| `cpu_offload` | `False` | Offload to CPU. |
| `overlap_communication` | `True` | Overlap comm/compute. |
| `overlap_p2p` | `True` | Overlap pipeline P2P. |
| `overlap_allgather` | `False` | Overlap all-gather. |
| `make_vocab_size_divisible_by` | `64` | Pad vocab for TP divisibility. |
| `dp_outside` | `True` | Place DP as the outermost mesh axis. |
| `custom_policy` | `None` | Custom Shardformer policy. |
| `fp8_communication` / `use_fp8` | `False` / `False` | FP8. |
| `inner_ring_size` | `None` | Ring size for ring-attention SP. |
| Loss-scale + DDP args | (as above) | `initial_scale`, `broadcast_buffers`, `ddp_bucket_cap_mb=25`, `find_unused_parameters`, `static_graph`, etc. |

> [!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:

| Field | Default | Purpose |
|---|---|---|
| `max_batch_size` | `8` | Max concurrent requests. |
| `max_input_len` / `max_output_len` | `256` / `256` | Prompt / generation length caps. |
| `dtype` | `torch.float16` | Compute dtype (`fp16`/`bf16`/`fp32`). |
| `kv_cache_dtype` | `None` | KV-cache dtype (e.g. for FP8 cache). |
| `prompt_template` | `None` | Chat prompt template name. |
| `do_sample` | `False` | Sampling vs greedy. |
| `beam_width` | `1` | Beam search width. |
| `prefill_ratio` | `1.2` | Buffer ratio for prefill scheduling. |
| `pad_input` | `False` | Pad inputs to equal length. |
| `early_stopping` | `False` | Stop beams early. |
| `top_k` / `top_p` / `temperature` | `50` / `1.0` / `1.0` | Sampling controls. |
| `no_repeat_ngram_size` | `0` | Block repeated n-grams. |
| `repetition_penalty` | `1.0` | Penalize repeats. |
| `forced_eos_token_id` | `None` | Force EOS token. |
| `ignore_eos` | `False` | Continue past EOS. |
| `use_spec_dec` | `False` | Speculative decoding. |
| `max_n_spec_tokens` | `5` | Speculative tokens per step. |
| `glimpse_large_kv` | `False` | Large-KV glimpse optimization. |
| `block_size` | `16` | Paged KV-cache block size. |
| `tp_size` / `pp_size` | `1` / `1` | Parallelism for inference. |
| `micro_batch_size` / `micro_batch_buffer_size` | `1` / `None` | Pipeline micro-batching. |
| `use_cuda_kernel` | `False` | Use fused CUDA kernels. |
| `high_precision` | `False` | Higher-precision math paths. |
| `use_cuda_graph` | `False` | CUDA-graph capture to cut launch overhead. |
| `max_context_len_to_capture` | `512` | Max context for graph capture. |
| `enable_streamingllm` | `False` | StreamingLLM attention. |
| `start_token_size` / `generated_token_size` | `4` / `512` | StreamingLLM window sizes. |
| `patched_parallelism_size` | `1` | Patched-parallelism degree (diffusion). |

`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:

| Variable | Read at | Effect |
|---|---|---|
| `BUILD_EXT` | `setup.py` (install) | `1` = compile CUDA/C++ kernels at install time. |
| `FORCE_CUDA` | `extensions/cuda_extension.py` | `1` = force CUDA extension build even without a visible GPU. |
| `TORCH_CUDA_ARCH_LIST` | `extensions/utils.py` | Target GPU architectures for kernel compilation. |
| `CUDA_DEVICE_MAX_CONNECTIONS` | set to `1` by `colossalai/initialize.py` | Ensures deterministic kernel-launch ordering so communication overlaps correctly. |
| `RANK`, `LOCAL_RANK`, `WORLD_SIZE`, `MASTER_ADDR`, `MASTER_PORT` | `launch_from_torch` | Standard torch.distributed coordinates. |
| `SLURM_PROCID`, `SLURM_NPROCS` | `launch_from_slurm` | SLURM rank/world size. |
| `OMPI_COMM_WORLD_*` | `launch_from_openmpi` | OpenMPI rank/local-rank/world size. |
| `INFER_NTK_ALPHA` | legacy smoothquant llama | NTK scaling alpha (legacy inference). |

> [!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`

```bash
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](https://hub.docker.com/r/hpcaitech/colossalai) 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.

| Task | How |
|---|---|
| Verify environment | `colossalai check -i` |
| Start a job | `colossalai run --nproc_per_node N [--hostfile ...] train.py` |
| Monitor | Standard tools (`nvidia-smi`, `dmesg`, your scheduler's logs). The dist logger prints rank-0 messages. |
| Logging | `colossalai.logging.get_dist_logger()` — rank-aware; use `ranks=[0]` to log once. |
| Checkpoint | `booster.save_model(..., shard=True, use_async=True)` for large models. |
| Resume | `booster.load_model/optimizer/lr_scheduler`. |
| Profiling | ColossalChat ships `profiling.sh`/`profile_grpo.py`; pipeline profilers live in `colossalai/fx/profiler` and `legacy/utils/profiler`. |

### 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:

```python
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](#16-known-issues--potential-improvements)). 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

| Module | Key public symbols |
|---|---|
| `colossalai` | `launch`, `launch_from_torch`, `launch_from_slurm`, `launch_from_openmpi`, `__version__` |
| `colossalai.booster` | `Booster`, `Plugin`, `Accelerator` |
| `colossalai.booster.plugin` | `TorchDDPPlugin`, `TorchFSDPPlugin`, `LowLevelZeroPlugin`, `GeminiPlugin`, `HybridParallelPlugin`, `MoeHybridParallelPlugin` |
| `colossalai.shardformer` | `ShardFormer`, `ShardConfig`, `ModelSharder`, `GradientCheckpointConfig`, `PipelineGradientCheckpointConfig` |
| `colossalai.zero` | `GeminiDDP`, `GeminiOptimizer`, `GeminiAdamOptimizer`, `LowLevelZeroOptimizer`, `zero_model_wrapper`, `zero_optim_wrapper`, `get_static_torch_model` |
| `colossalai.nn.optimizer` | `HybridAdam`, `CPUAdam`, `FusedAdam`, `FusedLAMB`, `FusedSGD`, `Lamb`, `Lars`, `Adafactor`, `CAME`, `GaLoreAdamW`, `GaLoreAdamW8bit`, `GaLoreAdafactor`, `DistributedLamb`, `DistributedCAME`, `DistributedAdaFactor`, `DistGaloreAwamW`, `cast_to_distributed` |
| `colossalai.inference` | `InferenceEngine`, `InferenceConfig`, `RequestHandler` |
| `colossalai.cluster` | `ProcessGroupMesh`, `DistCoordinator` |
| `colossalai.checkpoint_io` | `GeneralCheckpointIO`, `HybridParallelCheckpointIO`, `MoECheckpointIO`, `CheckpointIO`, `CheckpointIndexFile` |
| `colossalai.lazy` | `LazyInitContext` |
| `colossalai.logging` | `get_dist_logger` |

### The `Booster` surface

| Method | Signature (abbreviated) |
|---|---|
| `boost` | `(model, optimizer=None, criterion=None, dataloader=None, lr_scheduler=None)` |
| `backward` | `(loss, optimizer)` |
| `execute_pipeline` | `(data_iter, model, criterion, optimizer=None, return_loss=True, return_outputs=False)` |
| `no_sync` | `(model=None, optimizer=None)` (context manager) |
| `enable_lora` | `(model, pretrained_dir=None, lora_config=None, bnb_quantization_config=None, quantize=False)` |
| `save_model` | `(model, checkpoint, shard=False, gather_dtensor=True, prefix=None, size_per_shard=1024, use_safetensors=False, use_async=False)` |
| `load_model` | `(model, checkpoint, strict=True, low_cpu_mem_mode=True, num_threads=1)` |
| `save_optimizer` / `load_optimizer` | sharded-aware optimizer checkpointing |
| `save_lr_scheduler` / `load_lr_scheduler` | scheduler checkpointing |
| `save_lora_as_pretrained` | `(model, checkpoint, use_safetensors=False)` |

---

## 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.).

| Aspect | Detail |
|---|---|
| Runner | `pytest`. |
| Distributed tests | Spawn multi-process groups; require GPUs (NCCL). |
| Model zoo | `tests/kit/model_zoo` provides standard models for sharding/parallelism tests. |
| CI | GitHub Actions (`.github/workflows`); see the repo's workflow README. |

> [!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:

| Area | Symptom |
|---|---|
| `LowLevelZeroPlugin` + various optimizers (Adagrad, FusedLAMB, Lamb, GaLoreAdamW8bit) | Failures with small `initial_scale` and/or CPU offload (`KeyError: 'sum'`, runtime errors). |
| `LowLevelZeroPlugin` + `use_fp8=True` | `shape mismatch error: shape [32, 512] is invalid for input of size 512`. |
| `GeminiPlugin` + CPUAdam | `AssertionError: div_scale should remain default`. |
| `GeminiPlugin` + `offload_optim_frac=1.0` + small `min_chunk_size_m` | `RuntimeError: expected input to be on cuda during training`. |
| `HybridParallelPlugin` + `enable_all_optimization=True` | `TypeError in ColoAttention.attention()`. |
| `HybridParallelPlugin` + `use_fp8=True` | matmul dimension-divisibility runtime error. |
| Sequence parallelism in HybridPlugin | "cannot split tensor evenly". |
| `PipelineStageManager` | `AttributeError: object has no attribute 'stage_indices'`. |

> [!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 `RuntimeError`s 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

```bash
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

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

### Minimal training scaffold

```python
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

| Situation | Plugin |
|---|---|
| Model fits on one GPU | `TorchDDPPlugin` |
| Optimizer memory tight | `LowLevelZeroPlugin(stage=2, cpu_offload=True)` |
| Model too big for GPU | `GeminiPlugin(placement_policy="auto", offload_optim_frac=...)` |
| Largest transformers | `HybridParallelPlugin(tp_size, pp_size)` |
| Mixture of Experts | `MoeHybridParallelPlugin(tp_size, pp_size, ep_size)` |
| Native FSDP | `TorchFSDPPlugin` |

### Checkpoint

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

---

_Generated June 2026 from hpcaitech/ColossalAI at 4f9953b. AI-generated documentation. The library is checked by sampling (three repositories so far); this guide was not individually reviewed. Verify against the project before relying on it. Unofficial; not affiliated with or reviewed by the project._