verl/utils/memory_buffer.py

• MemoryBuffer · class · L24-L48 — class MemoryBuffer
• init · method · L30-L34 — def init(self, numel: int, numel_padded: int, dtype: torch.dtype)
• zero · method · L36-L38 — def zero(self)
• get · method · L40-L48 — def get(self, shape, start_index)
• calc_padded_numel · function · L51-L55 — def calc_padded_numel(shape: torch.Size, dtype: torch.dtype)
• get_weight_buffer_meta_from_module · function · L58-L65 — def get_weight_buffer_meta_from_module(module: nn.Module) -> Dict[str, Dict]
• build_memory_buffer · function · L68-L94 — def build_memory_buffer(weight_buffer_meta: Dict[str, Dict]) -> Dict[torch.dtype, MemoryBuffer]
• build_memory_reference_from_module · function · L97-L110 — def build_memory_reference_from_module(module: torch.nn.Module, memory_buffers: Dict[torch.dtype, MemoryBuffer], maintain_weight=True)
• build_memory_reference · function · L113-L137 — def build_memory_reference(weight_buffer_meta: Dict[str, Dict], memory_buffers: Dict[torch.dtype, MemoryBuffer])
• MemoryBufferModuleWrapper · class · L140-L157 — class MemoryBufferModuleWrapper
• init · method · L146-L151 — def init(self, module: nn.Module)
• get_memory_buffers · method · L153-L154 — def get_memory_buffers(self)
• get_weight_buffer_meta · method · L156-L157 — def get_weight_buffer_meta(self)
• MegatronMemoryBufferForRollout · class · L160-L214 — class MegatronMemoryBufferForRollout(object)
• init · method · L175-L179 — def init(self, transform_memory_param_fn)
• initialize_weight_buffer · method · L181-L197 — def initialize_weight_buffer(self, weight_buffer_meta_pp: List[Dict[str, Dict]])
• build_memory_reference · method · L199-L202 — def build_memory_reference(self)
• named_parameters · method · L205-L206 — def named_parameters(self)
• weight_buffers · method · L209-L210 — def weight_buffers(self)
• memory_buffers · method · L213-L214 — def memory_buffers(self)