/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
colossalai/nn/optimizer/__init__.py
57 строк
1 KB
Hongxin Liu
[misc] fix dist logger (#5782)
05 июн 2024, 10:04
Не верифицирован
05 июн 2024, 10:04
b9d646f
Код
Авторство
О чём код?
from galore_torch import GaLoreAdafactor, GaLoreAdamW from colossalai.logging import get_dist_logger from .came import CAME from .cpu_adam import CPUAdam from .distributed_adafactor import DistributedAdaFactor from .distributed_came import DistributedCAME from .distributed_galore import DistGaloreAwamW from .distributed_lamb import DistributedLamb from .fused_adam import FusedAdam from .fused_lamb import FusedLAMB from .fused_sgd import FusedSGD from .galore import GaLoreAdamW8bit from .hybrid_adam import HybridAdam from .lamb import Lamb from .lars import Lars from .adafactor import Adafactor # noqa __all__ = [ "FusedLAMB", "FusedAdam", "FusedSGD", "Lamb", "Lars", "CPUAdam", "HybridAdam", "DistributedLamb", "DistGaloreAwamW", "GaLoreAdamW", "GaLoreAdafactor", "GaLoreAdamW8bit", "CAME", "DistributedCAME", "Adafactor", "DistributedAdaFactor", ] optim2DistOptim = { GaLoreAdamW8bit: DistGaloreAwamW, Lamb: DistributedLamb, CAME: DistributedCAME, Adafactor: DistributedAdaFactor, } def cast_to_distributed(optim): if optim.__class__ in optim2DistOptim: _logger = get_dist_logger() _logger.info(f"Converting optimizer {optim.__class__.__name__} to its distributed version.", ranks=[0]) if isinstance(optim, GaLoreAdamW8bit): return optim2DistOptim[GaLoreAdamW8bit](optim.param_groups, args=optim.args) return optim2DistOptim[optim.__class__](optim.param_groups) return optim