Efficient Large-Scale Distributed Training with Colossal-AI and Lightning AI
pip install -U lightning-colossalai
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!