diff --git a/colossalai/logging/logging.py b/colossalai/logging/logging.py index 089308188..1f0ea6c1b 100644 --- a/colossalai/logging/logging.py +++ b/colossalai/logging/logging.py @@ -5,11 +5,12 @@ import colossalai import logging from pathlib import Path from typing import Union +from rich.logging import RichHandler from colossalai.context.parallel_mode import ParallelMode _FORMAT = 'colossalai - %(name)s - %(asctime)s %(levelname)s: %(message)s' -logging.basicConfig(level=logging.INFO, format=_FORMAT) +logging.basicConfig(level=logging.INFO, format=_FORMAT, handlers=[RichHandler()]) class DistributedLogger: diff --git a/requirements/requirements-test.txt b/requirements/requirements-test.txt index 69b82ff84..3fa37c0f2 100644 --- a/requirements/requirements-test.txt +++ b/requirements/requirements-test.txt @@ -1,3 +1,3 @@ pytest rpyc -matplotlib \ No newline at end of file +matplotlib diff --git a/requirements/requirements-zero.txt b/requirements/requirements-zero.txt index 816211e72..e69de29bb 100644 --- a/requirements/requirements-zero.txt +++ b/requirements/requirements-zero.txt @@ -1 +0,0 @@ -deepspeed \ No newline at end of file diff --git a/requirements/requirements.txt b/requirements/requirements.txt index 469b4d323..a05697b37 100644 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -6,3 +6,4 @@ psutil tensorboard packaging pre-commit +rich