megatron.model.distributed.MemoryBuffer

class megatron.model.distributed.MemoryBuffer(numel, numel_padded, dtype)

Bases: object

get(shape, start_index)

Return a tensor with the input shape as a view into the 1-D data starting at start_index.

zero()

Reset the buffer to zero.