Colossal-AI

Overview

Colossal-AI is an open source system for training and running very large neural networks across many GPUs. It brings several kinds of parallelism and several memory-saving tricks under one interface, so that you can choose a strategy by changing a plugin rather than rewriting your training loop. It covers data parallelism, tensor parallelism, pipeline parallelism, sequence parallelism and sharded optimizers, and it can move data between GPU memory, CPU memory and disk when the GPU runs short.

The main reason to look at Colossal-AI is simplicity of choice. Large-model training has many moving parts, and different libraries cover different parts. A unified system tries to give you one entry point. This page explains what that means in practice, how the main ideas work with everyday analogies, how it compares with Megatron-LM, DeepSpeed ZeRO and FSDP, and how the same scaling methods help models that learn from raw tables.

What it is

Colossal-AI is a Python library built on PyTorch. It wraps a model, an optimizer and a data loader, and it applies a chosen parallel and memory strategy to them. The project came from a company and research group that focus on high performance computing for AI, and it is developed in the open.

At the center is a high-level object, often called a booster, which takes your model, optimizer, loss and data loader and returns versions of them that are ready for distributed training. You pick a plugin, which is a named strategy, and the booster does the rest.

Typical plugins cover these strategies:

• A plain data-parallel plugin built on the standard PyTorch approach.
• A sharded data-parallel plugin built on the PyTorch FSDP approach.
• A low-level sharded optimizer plugin in the style of ZeRO.
• A heterogeneous memory plugin, often called Gemini, which manages tensors across GPU and CPU memory.
• A hybrid parallel plugin that mixes tensor, pipeline, sequence and data parallelism.
• A mixture-of-experts plugin for models with expert layers.

What is included

• Parallel layers and model wrappers for transformers.
• A chunk-based memory manager that decides where each tensor lives.
• Schedules for pipeline parallelism.
• Fused kernels and optimizers.
• Tools for checkpointing in sharded form and for converting checkpoints.
• Example recipes for training and fine-tuning large language models and for inference serving.

What it is not

It is not a model hub and it is not a hosted service. It is not a replacement for a cluster scheduler. It does not remove the need to understand memory and communication. It reduces the amount of glue code you write.

Where it came from

The project grew out of the practical pain of training large models with several tools at once. Around 2021, researchers described a unified system that supported more forms of tensor parallelism than the usual one-dimensional split. The paper presented the idea that different partitioning schemes, such as splitting a matrix in two dimensions or three dimensions, can lower communication cost on some network layouts.

Early versions focused on tensor parallelism in several dimensions, pipeline parallelism and a sharded optimizer. A later focus was heterogeneous training, in which the system uses CPU memory and even disk as extra room, so that a single machine with a modest GPU can train a model that would otherwise not fit. The memory manager for this idea was called Gemini.

The project also developed example applications, including a chat model training recipe with supervised fine-tuning and preference training, and later work on serving and on diffusion models. The interface changed over time from a configuration-file style to a more PyTorch-like booster and plugin style, which fits how most users already write training loops.

You should treat details such as exact plugin names and arguments as version dependent. Always read the documentation for the release you install.

How it works

Training a very large model is like moving a huge household. You have too many boxes for one truck. You can use several trucks that each carry a copy of the essentials, which is data parallelism. You can split the furniture across trucks, which is model parallelism. And you can rent a storage unit for things that you do not need right now, which is offloading to CPU memory or disk. Colossal-AI helps you combine all three.

Step 1: wrap the pieces with a booster

You create a plugin, hand it to a booster, and call a boost function with your model, optimizer, criterion and data loader. The booster returns wrapped objects. Your training loop then calls the booster's backward and the optimizer step, instead of calling them directly. This single point of control is how the system can switch strategy without changing the rest of your code.

Step 2: choose the parallel layout

With the hybrid plugin, you set sizes for tensor, pipeline and data parallelism, as well as whether to use sequence parallelism and which optimizer sharding stage to apply. The sizes multiply to the number of GPUs.

world_size = tp_size * pp_size * dp_size

If you leave tensor and pipeline sizes at one, you get a data-parallel run with optional sharding. This is a good place to start.

Step 3: tensor parallelism in more than one dimension

One-dimensional tensor parallelism splits a weight matrix along one axis, as in tensor parallelism. It needs a full activation on every device and an all-reduce after each block. Higher-dimensional schemes arrange the devices in a grid and split both the weights and the activations over the grid. They send smaller messages to more neighbors. The gain depends on the network and the number of devices, and the schemes are most useful at larger scales. In many practical cases, the one-dimensional scheme is the best default, and you should test the others only if communication is your bottleneck.

Step 4: pipeline parallelism

Layers are divided into stages. Micro-batches flow through the stages in a schedule so that devices work at the same time. The booster builds the stage assignment and the schedule for you, and the loop you write calls a single function that runs one pipeline step. As with other pipeline systems, the idle time at the ends of the batch is the bubble, and more micro-batches shrink it. See pipeline parallelism.

Step 5: sharded optimizer and gradients

The sharded plugins split the optimizer state, and at higher stages the gradients and the parameters, across data-parallel workers. This is the idea behind DeepSpeed ZeRO. The cost of sharding the parameters is extra gather traffic, because a layer must collect its full weights before it runs. Lower stages trade less saving for less communication.

Step 6: heterogeneous memory with Gemini

GPU memory is fast and small. CPU memory is slower and larger. Disk is slowest and largest. The Gemini manager groups tensors into chunks and moves chunks between GPU and CPU as needed. It tracks how tensors are used during the first steps and then plans which chunks to keep on the GPU and which to push out. This is like a librarian who keeps the most-used books on the desk and the others on shelves nearby.

The approach lets you train larger models on limited hardware, with a slowdown from the data movement. The slowdown is smaller when the CPU link is fast and when the chunks that move are not needed right away.

Step 7: activation checkpointing and fused kernels

Activation checkpointing drops some intermediate values and recomputes them in the backward pass, as in gradient checkpointing. Fused kernels combine small operations, which saves memory traffic. The system includes a fused optimizer and can use fast attention such as FlashAttention where supported.

Step 8: mixed precision and low precision

Training runs in mixed precision. The booster can take a precision argument such as bfloat16 or float16, and it keeps higher precision copies of weights where needed. Some recipes also use low-bit formats for communication or storage, though you should test quality when you do.

Step 9: checkpointing

Sharded training produces sharded checkpoints. The library offers saving and loading that is aware of the layout, and it offers tools to load a checkpoint back under a different layout. Always test this with a small model before a long run.

Step 10: inference

The same project has tools for serving large models with tensor parallel inference. For production serving, compare with specialized engines such as vLLM and check which features you need.

Key concepts and terms

A worked example

This example shows the shape of a training script with the booster API, followed by a small memory calculation. Names are placeholders and the details change by version, so treat the code as a sketch to compare with the documentation.

A training loop sketch

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

# Start the distributed environment. Launched with torchrun on every node.
colossalai.launch_from_torch()

model = build_model()               # your PyTorch module
optimizer = HybridAdam(model.parameters(), lr=1e-4)
criterion = torch.nn.CrossEntropyLoss()
dataloader = build_dataloader()     # your data loader

plugin = HybridParallelPlugin(
    tp_size=2,                      # tensor parallel size
    pp_size=1,                      # pipeline parallel size
    precision="bf16",
    zero_stage=1,                   # shard optimizer state
    enable_fused_normalization=True,
    enable_flash_attention=True,
)
booster = Booster(plugin=plugin)

model, optimizer, criterion, dataloader, _ = booster.boost(
    model, optimizer, criterion=criterion, dataloader=dataloader
)

for step, batch in enumerate(dataloader):
    outputs = model(**batch)
    loss = criterion(outputs.logits, batch["labels"])
    booster.backward(loss, optimizer)
    optimizer.step()
    optimizer.zero_grad()
    if step % 50 == 0:
        print("step", step, "loss", loss.item())

booster.save_model(model, "/ckpt/example-run", shard=True)

If you use pipeline parallelism, the inner loop changes: you call a function on the booster that runs the full pipeline step over a data iterator, and it returns the loss on the last stage.

Launching

torchrun --nproc_per_node 8 --nnodes 2 --node_rank $NODE_RANK \
  --master_addr node0.example.internal --master_port 29500 \
  train.py --plugin hybrid --tp 2 --batch-size 8

Switching strategy

To compare strategies, change only the plugin.

A hand calculation for offloading

Suppose a 7 billion parameter model trains with Adam in mixed precision. A common estimate is 16 bytes per parameter for weights, gradients and optimizer state, so about 112 GB. A single 24 GB GPU cannot hold that. With a heterogeneous plugin, the system might keep the 16-bit weights and some gradients on the GPU, which is about 14 GB for weights, and push the 12-byte optimizer part to CPU memory, which is about 84 GB. The CPU needs plenty of RAM, and each step pays the cost of moving data across the bus and running the optimizer on the CPU. The job runs, but slower than it would on a larger GPU. The point of the example is the budget method: sum up the parts, and decide which part can live in the slower memory.

Worked numeric example: choosing sizes for sixteen GPUs

Suppose you have two servers with eight GPUs each, and a model with 32 layers and 32 attention heads.

1. Try tensor size 4 and pipeline size 1. The layer and head counts divide evenly. Data-parallel size is 16 / 4 = 4.
2. Check memory. If weights per GPU are too large, raise tensor size to 8, which still fits inside one server. Data-parallel size is now 2.
3. If activations are the problem, add checkpointing before adding pipeline stages, because checkpointing costs only extra compute.
4. If you still need more room, use pipeline size 2 with one stage per server. Data-parallel size is now 1 with tensor size 8. Add micro-batches so the bubble stays small: with 2 stages and 16 micro-batches, the bubble is about 1/16 of a step.
5. Compare tokens per second for each setup, and pick the best one that also trains stably.

This step-by-step approach is easy to follow and to explain to a team. It also shows that no single setting is best for every case.

Typical failure messages and what they mean

Strengths and limitations

Strengths

• One interface for many strategies. Switching plugins is cheaper than rewriting code.
• Strong memory options. Sharding plus heterogeneous memory lets small clusters train large models.
• Multiple tensor parallel schemes. You can test higher-dimensional splits when the network favors them.
• A path from single GPU to cluster. The same loop works at many sizes.
• Example recipes for large language model training and fine-tuning reduce the time to a first run.
• Open source, with active development.

Limitations

• Fast-moving API. Plugin names and arguments change between versions, and older guides may not run.
• Smaller community than some alternatives, which means fewer answered questions for rare problems.
• Performance depends on tuning. Defaults are not always the best for your model and network.
• Offloading has a speed cost. It is a way to fit, not a way to go fast.
• Some features are tied to specific model families or require model wrappers.
• Debugging distributed issues is still hard, as with any library.
• Compatibility with new hardware and new number formats depends on release timing.

How it compares with alternatives

In practice, many teams test two or three of these on a small version of their job. The best one depends on model size, hardware, network and team skills.

Role in tabular foundation models

A tabular foundation model learns from raw tables. Each table has rows and columns, mixed types such as numbers, categories, dates and text, missing cells, and sometimes columns with no useful names. A model like TabPFN reads labeled rows as context and predicts for new rows without updating its weights. Larger models of this kind use attention over cells, rows and columns. Colossal-AI does not define such a model. It supplies ways to train and serve one when it becomes too large or too memory hungry for a single device.

Fitting bigger table models

Scaling a table model means more layers, a wider hidden size, or more experts. The same memory problems appear as for language models: weights, gradients and optimizer state. Sharded optimizers cut these across data-parallel workers. Tensor parallelism splits attention heads and feed-forward matrices. Heterogeneous memory helps a research group with a small cluster prototype a larger model before they move to a bigger system.

Long row contexts

In-context learning on tables benefits from many labeled rows. The sequence length grows with rows times columns. Sequence parallelism and fast attention reduce the memory for these long inputs. If attention over the full table is too costly, a design can use separate attention over rows and over columns, which keeps the cost lower. The parallel layers handle either design.

Synthetic and real pre-training data

Pre-training on generated tables yields an almost unlimited stream of data. A booster-based loop pairs well with an on-the-fly generator, because the data loader can create tables with random structure, random column types and random missing patterns. The key point for distributed training is to give each data-parallel worker a different random seed so workers do not see the same tables, and to bucket tables by size so that micro-batches have similar cost.

Missing values, imputation and anomaly detection

Masking cells is the natural training objective for imputation. The model sees a table with hidden cells and learns to fill them. This works with any of the parallel strategies, since masks are part of the data. The same trained model can score how surprising each cell is, which gives an anomaly signal. Such tasks need no special distributed code, but they do need careful evaluation on held-out tables.

Continual learning

When new table families arrive, a team may continue training. The sharded checkpoint tools let you resume under a different layout, which helps when the available cluster changes between runs. To limit forgetting, mix replayed old data with new data and track scores on earlier tasks.

Semantic and sector classification

When column names and sample values carry meaning, a model can classify a column's meaning or a table's industry sector. A text encoder plus a table encoder can be trained in the same system. Hybrid parallelism handles both blocks in one job. For the research and architecture context behind these ideas, see /model/research, /model/architecture and /model/benchmarks.

Serving

After training, the model may need to answer many small requests, each carrying a table. Serving needs batching by table size, careful handling of missing values and stable numeric output. Check whether the serving tools you choose support your architecture, or export the model to a general engine such as ONNX Runtime if the architecture allows.

Design thinking for a table model training job

Tabular training jobs differ from text jobs in a few ways that affect the choice of strategy.

• Example size varies a lot. A table with five columns and fifty rows is far cheaper than a table with hundreds of columns and thousands of rows. Data parallel workers that receive tables of different cost will finish steps at different times, so the fastest wait for the slowest. Bucketing by cost fixes most of this.
• Attention is often applied along two axes. The cost is then the sum of a column-wise part and a row-wise part, not one large square. Tensor parallelism splits both parts in the same way, by heads and by feed-forward width.
• The vocabulary is small or absent. Cell values are embedded through numeric and category encoders rather than a large token table, so the embedding layer is cheap and does not need vocabulary splitting.
• Targets are often a single value per row, so the output head is small and the loss is cheap. Most memory goes into activations, which makes recomputation and sequence-aware splitting the most useful tools.

Example plan for a mid-size table model

A plausible plan for a model of a few billion parameters on one server of eight GPUs:

1. Start with plain data parallelism and a sharded optimizer. Measure memory.
2. If activations do not fit for the longest tables, turn on activation checkpointing.
3. If weights still do not fit, set tensor parallel size to two or four.
4. Keep the global batch size fixed while you change layouts, so loss curves stay comparable.
5. Add masked-cell reconstruction and label prediction tasks to the evaluation loop, and run them every few thousand steps.

This plan uses the minimum amount of splitting needed, which keeps communication low and debugging simple.

Questions to ask before you adopt it

• Does the model really need more than one GPU?
• Does the team have time to learn a new API, or would a simpler wrapper meet the need?
• Which features matter most: memory fit, speed, or ease of switching?
• Will the cluster network support tensor parallelism, or only data parallelism?
• How will checkpoints be converted for serving?

Clear answers here prevent most failed adoptions.

Practical guidance and best practices

Start small and simple

Begin with the plain data-parallel plugin on a few GPUs. Confirm that loss falls and throughput is stable. Then add one feature at a time: sharding, then tensor parallelism, then pipeline parallelism, then offloading.

Pick the strategy from the bottleneck

• If memory is the limit and the network is slow, use sharding and offloading before tensor parallelism.
• If the network is fast inside servers, use tensor parallelism inside servers.
• If compute is the limit, check kernel fusion and precision before adding more parallelism.

Keep tensor parallel groups inside a server

Tensor parallelism needs high bandwidth. Place groups on devices that share fast links. Use pipeline or data parallelism across servers.

Use micro-batches wisely

With pipeline parallelism, increase the number of micro-batches until the bubble is small, but keep each micro-batch large enough to use the GPU well.

Use mixed precision with care

Use bfloat16 when your hardware supports it. If you use float16, keep loss scaling on. Log gradient norms to catch instability.

Save sharded checkpoints and test restore

Save in the sharded form for speed, and test loading on a small run. If you need a single file for sharing, run a conversion step, and then compare outputs.

Pin versions

Record the library version, the PyTorch version, the CUDA version and the driver version for each run. The API changes enough that a version mismatch is a common source of confusing errors.

Profile before tuning

Use a profiler to see where time goes. Common findings: the data loader is slow, a collective dominates the step, or many small kernels waste time.

Measure with the metric you care about

Tokens per second per GPU is a good comparison measure across layouts. Also measure time to a fixed loss, since a layout that changes the effective batch size may change how fast the model learns.

Keep a baseline

Keep a simple baseline, such as plain data parallelism on a small model, so you can check that a complex layout gives the same loss curve at the start.

Document the layout

Write down the sizes of tensor, pipeline and data parallelism, the micro-batch size, the global batch size and the precision for each run. Store them with the checkpoint.

Installing and checking the environment

Install the library into a clean environment that already has a matching PyTorch and CUDA build. Some features compile kernels at install or first use, so the compiler and CUDA toolkit must match the runtime. After install, run a tiny script that launches two processes, performs an all-reduce and prints the result. If that works, the network and the launcher are healthy. Only then run the real model.

python -c "import torch; print(torch.__version__, torch.cuda.device_count())"
torchrun --nproc_per_node 2 check_allreduce.py

A decision table for a first strategy

Estimating communication cost

A rough way to think about it: each form of parallelism has a typical message size and frequency.

• Data parallelism sends gradients once per step. The message size equals the number of parameters held by a replica.
• Tensor parallelism sends activations several times per layer. The messages are smaller but very frequent.
• Pipeline parallelism sends activations once per stage boundary per micro-batch. The messages are small and few.
• Sharded parameters send weights before each layer runs, and then drop them.

From this you can see why tensor parallelism wants the fastest links, and why pipeline parallelism suits slow links between servers.

Working with data loaders

Large runs are often limited by data. Prepare data in advance, store it in a format that supports fast random reads, and use several loader workers per process. Pin memory for faster transfer to the GPU. For table data, store each table with a small header that records its shape and column types, so the loader can bucket by size without reading the whole file.

Logging and metrics

Log at least these values on a fixed interval: loss, learning rate, gradient norm, tokens per second, memory use, and step time. Log from one rank to avoid clutter, but also record a per-rank step time summary every so often to spot stragglers. Keep logs next to checkpoints so each checkpoint can be traced back to its run.

Reliability habits

• Run a short smoke test every time you change the layout.
• Save a checkpoint just before any risky change, such as an upgrade.
• Keep a record of which data shards were consumed, so a resumed run does not repeat data.
• Use a fixed seed for model initialization so you can reproduce a run when you hunt a bug.
• Set timeouts on collective calls so that failures show up as errors instead of silent hangs.

Fine-tuning versus pre-training

Most users fine-tune rather than pre-train. For fine-tuning, combine the sharded plugin with parameter-efficient methods such as LoRA to cut memory further, and consider bitsandbytes for low-precision loading when the model is large. For pre-training, favor the hybrid plugin and long-run reliability features. The two jobs have different goals: fine-tuning cares about quick turnaround on small data, pre-training cares about steady throughput for days or weeks.

Testing a plugin change safely

When you move from one plugin to another, compare the first hundred steps of loss with the same seed and the same global batch size. Small differences are normal because of floating point order, but a large gap points to a bug in the setup, such as a wrong gradient scaling or a broken data split. Keep the comparison plots in your run notes.

A note on reading benchmarks

Published speed comparisons between training libraries depend heavily on model size, hardware, network, precision and batch size. A claim that one tool is faster usually holds only for the setup tested. When you read such claims, check the setup, then repeat the test with your own model and a short run. Do not carry numbers from one cluster to another.

Common mistakes

• Starting with the most complex plugin. Add features one at a time.
• Using offloading and expecting speed. It saves memory at the cost of time.
• Putting tensor parallel groups across slow links.
• Forgetting that the layer count and head count must divide evenly by the parallel sizes.
• Comparing runs with different global batch sizes and calling the results a speed test.
• Reusing code from an old release without checking API changes.
• Not testing checkpoint conversion or restore before a long job.
• Ignoring CPU memory. Heavy offloading needs plenty of RAM and fast pinned transfers.
• Leaving random seeds identical across data-parallel workers, so they see the same data.
• Treating a hang as a bug in the model when a node or network problem is the cause.
• Using a higher-dimensional tensor split without measuring. It is not always faster.
• Skipping evaluation on the target task and trusting only training loss.

Frequently asked questions

What is Colossal-AI used for?

Colossal-AI is used to train, fine-tune and serve large neural networks across many GPUs or on limited hardware. It offers data, tensor, pipeline and sequence parallelism, sharded optimizers and memory offloading through one booster interface, so users can choose a strategy by selecting a plugin.

How is Colossal-AI different from DeepSpeed?

Both offer sharded optimizers and offloading, so they overlap. DeepSpeed is configured mainly through a config file and is known for its ZeRO stages. Colossal-AI uses a booster and plugin interface and adds several tensor parallel schemes and a chunk-based memory manager. The best choice depends on your model, your hardware and which features you need.

What is Gemini in Colossal-AI?

Gemini is the heterogeneous memory manager. It groups tensors into chunks and moves them between GPU memory, CPU memory and, in some setups, disk, based on how soon each chunk is needed. This lets a single machine train a model that is larger than its GPU memory, though training becomes slower.

Do I need many GPUs to use it?

No. You can run it on a single GPU for testing, and the heterogeneous memory option is aimed at small setups. The parallel features show their value when you have several GPUs, and tensor and pipeline parallelism matter most when the model does not fit on one device.

Does Colossal-AI work with Hugging Face models?

It provides policies and wrappers that support many common transformer models from the Transformers ecosystem, so that you can load a model and apply a plugin. Support depends on the model family and the release, so check the list of supported models for the version you use.

What is the difference between tensor parallelism and ZeRO?

Tensor parallelism splits the math of each layer across devices, so each device computes a slice of every layer. ZeRO keeps the whole layer computation on each device but shards the stored state, such as optimizer data, across workers. Tensor parallelism needs fast links, while ZeRO trades memory for extra communication of weights at higher stages.

Can I use it for models that read tables?

Yes, if the model is a PyTorch network. The library does not care what the inputs are. A table model with attention over rows and columns can use the same plugins. You still need a data loader that handles variable table shapes, masks for missing cells and a sensible batching scheme.

Is it good for inference?

It includes tools for serving large models across several GPUs. For high-throughput production serving, compare it with specialized engines, check which model architectures they support, and test latency and memory with your own traffic pattern.

How do I know which plugin to pick?

Start from the question of what limits you. If the model fits, use plain data parallelism. If memory is tight, add a sharded optimizer. If the model is still too large, add tensor or pipeline parallelism. If you only have a small machine, try the heterogeneous plugin. Measure throughput and memory for each choice on a short run.

How should I handle version changes?

Pin the library version for each project, read the release notes before upgrading, and run a short regression test that checks the loss curve and the checkpoint loading. Treat examples from older blog posts as guides to the ideas, not as code to paste.

Glossary

• Booster: the object that applies a plugin to a model, optimizer and data loader.
• Chunk: a group of tensors stored and moved as one unit by the memory manager.
• Data parallelism: copies of the model that see different data and share gradients.
• Gemini: the heterogeneous memory manager of the project.
• Heterogeneous training: training that uses GPU, CPU and disk memory together.
• Hybrid parallelism: a mix of tensor, pipeline and data parallelism.
• Micro-batch: a small slice of a batch used in pipeline schedules.
• Offloading: moving tensors out of GPU memory to a slower store.
• Pipeline bubble: idle time at the start and end of a pipeline step.
• Plugin: a strategy object that tells the booster how to parallelize.
• Sharding: splitting stored state across workers.
• Tensor parallelism: splitting layer weights across devices.
• ZeRO: a family of techniques that shard optimizer state, gradients and parameters.

Summary

Colossal-AI is a unified open source system for training and serving large models. A booster applies a plugin that selects data, tensor, pipeline or sequence parallelism, sharded optimizers or heterogeneous memory, so you can change strategy without rewriting your loop. Its strengths are the single interface and the memory options. Its weaknesses are a fast-moving API and the usual difficulty of distributed debugging.

For tabular foundation models, it provides the scaling layer: sharding and tensor splits for bigger table models, sequence tools for long row contexts, masking-friendly training for imputation and anomaly detection, resumable sharded checkpoints for continual learning, and a path to run all of this on small hardware. Compare it with Megatron-LM, DeepSpeed ZeRO and FSDP, pair it with FlashAttention and mixed precision, and read more about the model side at /model/research, /model/architecture and /model/benchmarks.

Related

It belongs to Training tooling in my stack, under LLMs, retrieval, and evaluation. Others in the same group: