跳转至

aggregate_logger.py — 这个文件提供了分布式训练环境中的日志工具

模块路径: verl.utils.logger.aggregate_logger

文件概述

这个文件提供了分布式训练环境中的日志工具。在多 GPU/多节点训练中,如果所有进程都打印日志会导致输出混乱,因此需要控制只有特定 rank(通常是 rank 0)才输出日志。

关键代码讲解

1. LocalLogger

class LocalLogger:
    """简单的本地日志器,将度量指标打印到控制台"""
    def __init__(self, print_to_console=True):
        self.print_to_console = print_to_console

    def log(self, data, step):
        if self.print_to_console:
            print(concat_dict_to_str(data, step=step), flush=True)

配合辅助函数格式化输出:

def concat_dict_to_str(dict, step):
    output = [f"step:{step}"]
    for k, v in dict.items():
        if isinstance(v, numbers.Number):
            output.append(f"{k}:{pprint.pformat(v)}")
    return " - ".join(output)

2. DecoratorLoggerBase

class DecoratorLoggerBase:
    """装饰器日志基类,支持 rank 过滤"""
    def __init__(self, role, logger=None, level=logging.DEBUG, rank=0, log_only_rank_0=True):
        self.role = role
        self.rank = rank
        self.log_only_rank_0 = log_only_rank_0
        # 根据是否有 logger 选择输出方式
        self.logging_function = self.log_by_logging if logger else self.log_by_print

    def log_by_print(self, log_str):
        if not self.log_only_rank_0 or self.rank == 0:
            print(f"{self.role} {log_str}", flush=True)

    def log_by_logging(self, log_str):
        if not self.log_only_rank_0 or self.rank == 0:
            self.logger.log(self.level, f"{self.role} {log_str}")

3. 便捷打印函数

def print_rank_0(message):
    """只在 rank 0 打印(或非分布式环境中直接打印)"""
    if torch.distributed.is_initialized():
        if torch.distributed.get_rank() == 0:
            print(message, flush=True)
    else:
        print(message, flush=True)

def print_with_rank(message, rank=0, log_only_rank_0=False):
    """带 rank 信息的打印"""
    if not log_only_rank_0 or rank == 0:
        print(f"[Rank {rank}] {message}", flush=True)

def print_with_rank_and_timer(message, rank=0, log_only_rank_0=False):
    """带 rank 和时间戳的打印"""
    now = datetime.datetime.now()
    message = f"[{now.strftime('%Y-%m-%d %H:%M:%S')}] [Rank {rank}] {message}"
    if not log_only_rank_0 or rank == 0:
        print(message, flush=True)

核心类/函数列表

类/函数名 作用
LocalLogger 简单的本地日志器
DecoratorLoggerBase 装饰器日志基类
print_rank_0 只在 rank 0 打印
print_with_rank 带 rank 信息打印
print_with_rank_and_timer 带 rank 和时间戳打印
log_with_rank 使用 logging 模块记录日志

与其他模块的关系

  • 被 verl 中各组件广泛使用
  • 依赖 torch.distributed 判断分布式环境

小结

分布式训练中日志管理是基础需求。这个模块提供了从简单(print_rank_0)到灵活(DecoratorLoggerBase)的多层次日志工具。所有函数都统一了 flush=True,确保日志及时输出不被缓冲。