ColossalAI

Use this skill when the task names ColossalAI, colossalai, colossalai run, Booster, GeminiPlugin, HybridParallelPlugin, ShardFormer, Colossal-Inference, ColossalChat, Colossal-LLaMA, ColossalEval, ColossalQA, or ColossalMoE.

ColossalAI is a Linux-oriented distributed PyTorch and large-model system. Most meaningful training and inference workflows require CUDA-capable PyTorch, torchrun or ColossalAI launch setup, and model/data assets chosen by the user.

Quick Start

1. Verify the package and backend before giving GPU or distributed advice:

   python scripts/check_colossalai_environment.py --check-cli
   

2. Route to the focused sub-skill for detailed commands and recovery steps.
3. Keep hardware, model weights, datasets, credentials, and optional extensions explicit. Do not promise a CUDA, Apex, flash-attn, TensorNVMe, vLLM, or application package path until that dependency is installed and checked.
4. Prefer synthetic or helper-generated command checks before running long training, benchmark, model-download, or service workflows.

Route By Task

• Install, check, and launch: use sub-skills/installation-and-launch/SKILL.md for installation variants, colossalai check -i, colossalai run, hostfiles, torchrun, SLURM/OpenMPI launch APIs, and port/NCCL startup failures.
• Booster training: use sub-skills/booster-training/SKILL.md for Booster, plugin selection, training loop structure, dataloaders, checkpointing, LoRA, mixed precision, ZeRO, Gemini, and memory-aware training.
• Parallelism and sharding: use sub-skills/parallelism-and-sharding/SKILL.md for tensor/pipeline/sequence parallelism, ShardFormer, topology sizing, auto-parallel, activation checkpointing, NVMe offload, and process-group reasoning.
• Inference and serving: use sub-skills/inference-and-serving/SKILL.md for InferenceConfig, InferenceEngine, LLM generation, speculative decoding, Stable Diffusion 3 patched parallelism, serving/client benchmarks, and inference errors.
• First-party applications: use sub-skills/application-recipes/SKILL.md for ColossalChat/Coati, Colossal-LLaMA, ColossalEval, ColossalQA, and ColossalMoE environment isolation, command anatomy, and limitations.

Core Import and CLI Signals

python - <<'PY'
import colossalai, torch
print('colossalai', colossalai.__version__)
print('torch', torch.__version__, 'cuda', torch.version.cuda, 'available', torch.cuda.is_available())
PY
colossalai --help
colossalai check -i

Use colossalai check -i to distinguish PyTorch CUDA compatibility, system CUDA visibility, and AOT extension build information. System CUDA version: N/A usually means CUDA_HOME or nvcc is not visible; PyTorch CUDA wheels can still run CUDA kernels when the driver supports them.

Cross-Cutting References and Helpers

• references/installation.md summarizes package, Python, PyTorch, CUDA, and optional extension constraints that affect all workflows.
• references/troubleshooting.md covers install/import/backend/CLI/distributed failures that are not owned by one sub-skill.
• references/testing-and-verification.md lists safe checks, native candidates, and expensive examples to avoid by default.
• references/repo-provenance.md records the source revision, package version, and evidence paths used to build this skill.
• scripts/check_colossalai_environment.py performs a safe import, version, CUDA, and optional CLI-help diagnostic without downloading models or starting services.

Operating Rules

• Initialize distributed state before constructing most Booster plugins. For ordinary scripts, use torchrun or colossalai run; for manual setup, call colossalai.launch(...) with rank/world size/host/port.
• Treat app packages as separate environments unless proven compatible. The core ColossalAI package uses modern PyTorch; ColossalQA and some application stacks pin older or conflicting dependencies.
• Do not run benchmark-scale examples, large model generation, multi-node jobs, or service listeners unless the user has supplied assets, hardware, and approval.
• Do not route generic PyTorch, DeepSpeed, Lightning, Transformers, vLLM, or LangChain tasks here unless ColossalAI owns the launcher, plugin, sharding, inference, or application workflow in the user request.