Colossal-AI Framework, in depth: the Booster and its plugins, Gemini chunk memory, HybridParallelPlugin with Shardformer, and sizing a 13B run

The Booster and its plugins

The plugins currently exported by colossalai.booster.plugin are below. TorchFSDPPlugin is only exported on PyTorch 1.12 or newer.

A minimal training program

Every Colossal-AI program follows the same five steps: initialize the distributed environment, build the plugin and booster, create the model and optimizer, call boost, and train through the booster. The launch function now takes only a backend, a seed and a verbosity flag. Older tutorials pass a config argument, which current releases no longer accept.

import colossalai
import torch
from colossalai.booster import Booster
from colossalai.booster.plugin import GeminiPlugin
from colossalai.nn.optimizer import HybridAdam
from transformers import AutoModelForCausalLM

colossalai.launch_from_torch(seed=42)          # reads RANK, WORLD_SIZE from torchrun

plugin = GeminiPlugin(precision="bf16", placement_policy="static",
                      offload_optim_frac=0.0)
booster = Booster(plugin=plugin)

model = AutoModelForCausalLM.from_pretrained("path/to/model")
optimizer = HybridAdam(model.parameters(), lr=2e-5)   # runs on GPU or CPU state
criterion = lambda outputs, inputs: outputs.loss

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

for batch in dataloader:
    batch = {k: v.cuda() for k, v in batch.items()}
    outputs = model(**batch)
    loss = criterion(outputs, batch)
    booster.backward(loss, optimizer)                # not loss.backward()
    optimizer.step()
    optimizer.zero_grad()

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

Launch it with torchrun --nproc_per_node 8 train.py, adding the usual --nnodes and rendezvous flags for multi-node runs. Two details matter. First, call booster.backward rather than loss.backward(), because the plugin may need to scale the loss or hook gradient reduction. Second, use HybridAdam with Gemini or offloading. It keeps optimizer math wherever each piece of state lives, CPU or GPU, and a stock PyTorch Adam does not handle that split.

Gemini: chunks and placement

Gemini is Colossal-AI's version of fully sharded training, and its distinguishing feature is chunks. Instead of handling each parameter tensor separately, Gemini packs consecutive parameters, in initialization order, into fixed-size contiguous chunks. Gathering, reducing and moving data then happen one chunk at a time. That turns many small collectives and copies into fewer large ones, which use bandwidth better, and it reduces memory fragmentation from temporary communication buffers. A short search picks the chunk size, guided by hidden_dim and min_chunk_size_m.

Placement is the other half. With placement_policy="static", you set what lives where through three fractions: shard_param_frac (how much of the parameter state is sharded, default 1.0), offload_optim_frac and offload_param_frac (how much optimizer state and parameter state goes to CPU memory, both default 0.0). With "auto", Gemini measures memory during warm-up steps and moves chunks between CPU and GPU to stay within budget. Static placement is predictable and easy to reason about. Auto placement fits larger models on smaller GPUs, but throughput then depends on PCIe traffic you did not plan.

The docs describe Gemini as designed for clusters with high cross-node bandwidth and medium to small scale. Because every step all-gathers parameters, a large job on a slow inter-node network spends its time waiting. That is where the hybrid plugin fits better.

HybridParallelPlugin and Shardformer

HybridParallelPlugin combines tensor parallelism (tp_size), pipeline parallelism (pp_size), optional sequence parallelism (sp_size, with sequence_parallelism_mode set to split_gather, ring or all_to_all) and ZeRO stage 0, 1 or 2 across the remaining data-parallel ranks. It does not re-implement your model. Shardformer holds a policy per supported Hugging Face architecture that says how to split each linear layer, embedding and attention block, and which layers go to which pipeline stage. The docs state that only models Shardformer supports work with this plugin, so check the list before you plan a run around it.

from colossalai.booster.plugin import HybridParallelPlugin
from colossalai.lazy import LazyInitContext

plugin = HybridParallelPlugin(tp_size=8, pp_size=2, zero_stage=1,
                              microbatch_size=1, precision="bf16")
booster = Booster(plugin=plugin)

with LazyInitContext(default_device=torch.cuda.current_device()):
    model = AutoModelForCausalLM.from_config(config)   # no full copy materialized
optimizer = HybridAdam(model.parameters(), lr=1e-4)
model, optimizer, criterion, dataloader, _ = booster.boost(
    model, optimizer, criterion, dataloader)

it = iter(dataloader)
for step in range(num_steps):
    out = booster.execute_pipeline(it, model, criterion, optimizer,
                                   return_loss=True)
    optimizer.step()
    optimizer.zero_grad()
    if out["loss"] is not None:                          # only the last stage has it
        log(step, out["loss"].item())

With pipeline parallelism you no longer call the model yourself. execute_pipeline takes a data iterator, splits each batch into microbatches, runs the forward and backward schedule across stages, and returns a dict. Only the last stage computes a loss, so guard any logging. LazyInitContext delays allocating weights until the plugin knows how to shard them. The docs recommend it when 4N >= M for an N-billion-parameter model on M GB GPUs, because otherwise every rank briefly holds a full copy. For pipeline schedules in general, see pipeline parallelism in depth.

Worked example: a 13B model on 16 GPUs

Size a 13-billion-parameter model on two nodes of eight 80 GB GPUs, 16 GPUs in all, with bf16 weights and Adam. Mixed-precision Adam needs about 16 bytes per parameter: 2 for bf16 weights, 2 for gradients, and 12 for the fp32 master weights and two moments. For 13B that is 26 GB of weights, 26 GB of gradients and 156 GB of optimizer state, 208 GB in total, before activations.

Every option except DDP fits, so the decision rests on communication. LowLevelZero stage 2 keeps full weights on each GPU and only adds a reduce-scatter, which suits a small job with short sequences. Gemini leaves the most room for activations, so it suits long sequences or bigger batches, but it all-gathers weights across nodes on every step. The hybrid layout keeps the heavy tensor-parallel traffic on NVLink inside each node and sends only data-parallel gradient traffic between nodes. It is usually fastest on two or more nodes, provided the model is one Shardformer supports. A good process is to benchmark the top two candidates for 50 steps each and compare tokens per second per GPU.

Picking a plugin as a procedure

The worked example generalizes into a short procedure. Write it down as code so the choice is repeatable and reviewable, rather than something one engineer remembers. The thresholds below are starting points. Replace them with what your own benchmarks show.

def pick_plugin(params_b, gpu_gb, gpus_per_node, nodes, hf_supported, fast_internode):
    state_gb = 16 * params_b                 # bf16 weights + grads + fp32 Adam
    world = gpus_per_node * nodes
    headroom = 0.6 * gpu_gb                  # leave ~40% for activations
    if state_gb <= headroom:
        return "TorchDDPPlugin"
    weights_gb = 2 * params_b
    if weights_gb + 14 * params_b / world <= headroom:
        return "LowLevelZeroPlugin(stage=2)"
    if nodes > 1 and hf_supported and not fast_internode:
        return "HybridParallelPlugin(tp_size=gpus_per_node, zero_stage=1)"
    if state_gb / world <= headroom:
        return "GeminiPlugin(placement_policy='static')"
    return "GeminiPlugin(placement_policy='auto')  # or add GPUs"

The order of the tests reflects cost. Each step down adds communication or host traffic, so the procedure takes the cheapest layout that fits. The fast_internode flag captures the docs' warning that Gemini assumes good cross-node bandwidth: on a slow network, the hybrid plugin keeps the heavy traffic inside the node. For the 13B example, with 80 GB GPUs, two nodes and a supported model, the procedure returns LowLevelZero stage 2, because 26 GB of weights plus 11.4 GB of sharded state is under the 48 GB headroom. That matches the table. Treat the result as the first candidate to benchmark, not the final answer.

Checkpoints

Save and load through the booster so each plugin can handle its own sharding. booster.save_model(model, path, shard=True) writes sharded files with an index, and size_per_shard (in MB, default 1024) controls file size. use_safetensors=True writes safetensors files. Optimizer and learning-rate scheduler state have their own save and load calls on the booster. Load into a model that has already been boosted with the same plugin, so the shards land where the plugin expects. To change layout, for example to resume a hybrid run under Gemini, save a consolidated model and restart the optimizer state, or test the reshard path on a small model first. Not every plugin supports every checkpoint format, so check the plugin's docs before you rely on sharded saves.

Failure modes

• Calling loss.backward(). Mixed-precision scaling and gradient hooks are skipped. Training may run but diverge or never reduce gradients.
• Stock optimizers with offload. Gemini offload expects HybridAdam. A plain torch.optim.Adam cannot update optimizer state that lives on the CPU.
• Unsupported architectures in the hybrid plugin. A model with no Shardformer policy fails at boost time, or a custom layer stays unsplit and replicated.
• ZeRO 2 with pipelines or accumulation. The docs warn that ZeRO 2 does not support local gradient accumulation and is a poor match for pipeline parallelism. Use ZeRO 1 there.
• Logging the loss on every rank under pipelines. Non-final stages get None and the script crashes on .item().
• Auto placement surprises. placement_policy="auto" fits the model, but PCIe traffic halves throughput. Profile before committing to a long run.
• Old tutorials. Passing config= to the launch function, or using the pre-Booster engine API, fails on current releases.

Trade-offs

Colossal-AI's strength is that one API covers DDP, ZeRO, chunked offloading and 3D parallelism, with ready-made policies for popular Hugging Face models. That makes it easy to compare strategies on your own hardware. The costs are a smaller community than PyTorch FSDP or DeepSpeed, an API that has changed between releases, and a hybrid plugin limited to the architectures Shardformer covers. If you live in PyTorch and want the most conservative path, FSDP2 is the default. If you need maximum throughput at thousands of GPUs, Megatron-based stacks have more production mileage. Colossal-AI fits best for fine-tuning and mid-scale pre-training of supported models, where trying several layouts quickly matters more than the last few percent.

What to do next

1. Install the release that matches your PyTorch and CUDA, and run one of the repository's examples on a single node before touching your own code.
2. Do the 16-bytes-per-parameter sizing for your model and cross off plugins that cannot fit.
3. Check that your model architecture has a Shardformer policy before planning a hybrid run.
4. Benchmark the two best candidate plugins for 50 steps each and compare tokens per second per GPU and peak memory.
5. Use LazyInitContext for any model where 4N is at least M.
6. Save a sharded checkpoint, then restore and resume it, before starting a long run.