memory.py — 内存缓冲区管理

文件路径: verl/utils/megatron/memory.py

文件概述

提供 MemoryBuffer 类，用于预分配连续的 GPU 内存作为缓冲区，供 Megatron 的梯度归约等操作使用。

核心类

class MemoryBuffer:
    def __init__(self, numel, numel_padded, dtype):
        self.data = torch.zeros(numel_padded, dtype=dtype, device=get_device_id(), requires_grad=False)

    def zero(self):
        """重置缓冲区为零"""
        self.data.zero_()

    def get(self, shape, start_index):
        """从缓冲区中获取指定形状的视图"""
        buffer_tensor = self.data[start_index:end_index]
        return buffer_tensor.view(shape)

预分配大块内存并按需切分视图，避免频繁的小块内存分配。

与其他模块的关系

• 依赖 device.py 获取设备 ID
• 被 Megatron 的梯度缓冲区管理使用

小结

简单的内存池实现，优化 GPU 内存分配性能。