Hi, thanks for your awesome PR. In colossalai/trainer/ophooks/__init__.py, I think it is not necessary to expose TestOpHook to the user?

Hi, thanks for your awesome PR. In colossalai/trainer/ophooks/__init__.py, I think it is not necessary to expose TestOpHook to the user?

Sure, I will fix it.

Hey, I polished the code and added a simple GPU memory tracer using the ophooks. It can dump the GPU memory usage curve during #niter iterations to a file. I hope you like the feature.

1. The ophooks are applied on the class engine, therefore I put them under the dir of the engine.
2. I am not sure how to pass the ophooks to the engine constructor. The project seems to recommend initialing the hooks in a config file.

Hi, the ophooks look great to me. As for your second point, we initialize the trainer hooks objects in the python script and pass the hook objects to the trainer instead of defining them in the config file. I am wondering if you can tell us where did you see such usage as it might be deprecated doc and we can update it.

Hi, the ophooks look great to me. As for your second point, we initialize the trainer hooks objects in the python script and pass the hook objects to the trainer instead of defining them in the config file. I am wondering if you can tell us where did you see such usage as it might be deprecated doc and we can update it.

I saw APIs in builder/builder.py

def build_hooks(config, trainer):

Most of the functions in this file are not used in the project.

Noted, these are deprecated after we update some APIs. Some code cleanup is needed.

Meanwhile, I saw some print statements in the op hooks. I would recommend using DistributedLogger instead as the print statement may not be captured by the standard output when logging is used. Using logger also allows the user to see the message in their log txt file.

A general usage of logger will be like this:

from colossalai.logging import get_dist_logger
from colossalai.context import ParallelMode


# this will get the root Python logger
logger = get_dist_logger()

# log on all ranks
logger.info("some message")

# log only on rank 0
logger.info("some messsages", ranks=[0])

# log on rank 0 and rank 1
logger.info("some messsages", ranks=[0, 1])

# log on all data parallel rank 0
logger.info("some messages", ranks=[0], parallel_mode=ParallelMode.DATA)

# save the log
logger.log_to_file('./logs')

You can find the API doc here.

The MR has been polished. Replacing print() with logger.

Awesome, thanks for the contribution :)