megatron.core.tensor_parallel.utils

Description

Classes

VocabUtility()

Split the vocabulary into world_size chunks and return the first and last index of the vocabulary belonging to the rank partition: Note that indices in [fist, last)

Functions

gather_split_1d_tensor(tensor)	Opposite of split_tensor_into_1d_equal_chunks.
split_tensor_along_last_dim(tensor, ...[, ...])	Split a tensor along its last dimension.
split_tensor_into_1d_equal_chunks(tensor[, ...])	Break a tensor into equal 1D chunks across tensor parallel ranks.