Lightning âš¡ Colossal-AI
Efficient Large-Scale Distributed Training with Colossal-AI and Lightning AI
Installation
pip install -U lightning-colossalaiUsage
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!
![dependabot[bot]](https://avatars.githubusercontent.com/u/49699333?size=120)
