github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[misc] fix dist logger (#5782)

pull/5787/head
Hongxin Liu 2024-06-05 15:04:22 +08:00 committed by GitHub
parent c46e09715c
commit b9d646fe9e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/nn/optimizer/__init__.py
View File

@ -43,11 +43,11 @@ optim2DistOptim = {
CAME: DistributedCAME,
Adafactor: DistributedAdaFactor,
}
_logger = get_dist_logger()
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):