megatron.memory.MemoryBuffer#

class megatron.memory.MemoryBuffer(name, numel, dtype, track_usage)#

Bases: object

Contiguous memory buffer. Allocate a contiguous memory of type dtype and size numel. It is used to reduce memory fragmentation.

Usage: After the allocation, the _start index is set tot the first: index of the memory. A memory chunk starting from _start index can be allocated for an input tensor, with the elements of the tensor being coppied. The buffer can be reused by resetting the _start index.

add(tensor)#: Allocate a chunk of memory from the buffer to tensor and copy the values.

get_data()#: Return the data currently in use.

is_in_use()#: Whether the current buffer hold on to any memory.

numel_in_use()#: Return number of elements in use.

print_average_usage()#: Print memory usage average over time. We would like this value to be as high as possible.

reset()#: Reset the buffer start index to the beginning of the buffer.