memory_utils.py — 显存管理与可视化

文件路径: verl/utils/memory_utils.py

文件概述

提供 GPU 显存使用的监控、快照和可视化工具。用于调试 OOM（显存不足）问题和优化显存使用。

核心功能

1. 显存可视化

def enable_memory_visualize(trace_alloc_max_entries=100000, stack_depth=32):
    """开启 CUDA 内存分配追踪，用于后续生成内存快照"""
    torch.cuda.memory._record_memory_history(
        max_entries=trace_alloc_max_entries,
        stacks=stack_depth
    )

2. 内存快照

class MemorySnapshotSampler:
    def dump_memory_snapshot(self, out_dir, tag, sub_dir=None):
        """导出 CUDA 内存快照文件，可用 PyTorch Memory Viz 工具查看"""
        snapshot = torch.cuda.memory._snapshot()
        # 保存为 pickle 文件

3. 显存日志

def log_gpu_memory_usage(head, logger=None):
    """打印当前 GPU 显存使用情况"""
    # 输出: memory allocated, memory reserved, device memory used/total

与其他模块的关系

• 被 profiler/profile.py 中的 TorchMemoryProfiler 使用
• 被各个 worker 在关键节点调用以监控显存

小结

memory_utils.py 是显存调试的好帮手，特别是在优化大模型训练时定位显存瓶颈。