Skip to content
This repository has been archived by the owner on Sep 1, 2023. It is now read-only.

Large Scale Distributed Model Training strategy with Colossal AI and Lightning AI

License

Notifications You must be signed in to change notification settings

Lightning-Universe/lightning-ColossalAI

Repository files navigation

Lightning ⚡ Colossal-AI

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

lightning PyPI Status PyPI - Python Version PyPI - Downloads Deploy Docs

General checks CI testing Build Status pre-commit.ci status


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!