megatron.memory

Description

Classes

MemoryBuffer(name, numel, dtype, track_usage)	Contiguous memory buffer.
RingMemBuffer(name, num_buffers, numel, ...)	A ring of memory buffers.

Functions

allocate_mem_buff(name, numel, dtype, ...)	Allocate a memory buffer.
get_mem_buff(name)	Get the memory buffer.