Awesome

Lightning ⚡ Colossal-AI

Efficient Large-Scale Distributed Training with Colossal-AI and Lightning AI

Installation

pip install -U lightning-colossalai

Usage

Simply set the strategy argument in the Trainer:

import lightning as L

trainer = L.Trainer(strategy="colossalai", precision="16-mixed", devices=...)

For more fine-grained tuning of Colossal-AI's parameters, pass the strategy object to the Trainer:

import lightning as L
from lightning_colossalai import ColossalAIStrategy

strategy = ColossalAIStrategy(...)
trainer = L.Trainer(strategy=strategy, precision="16-mixed", devices=...)

Find all configuration options in the docs!