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,确保日志及时输出不被缓冲。