heterogeneous-distributed-training-framework/docs/source/api-guide/tensor_parallel.rst

tensor_parallel package

This package contains an implementation for tensor parallelism in transformer models (see Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism and Reducing Activation Recomputation in Large Transformer Models for details).

Submodules

tensor_parallel.cross_entropy module

core.tensor_parallel.cross_entropy

tensor_parallel.data module

core.tensor_parallel.data

tensor_parallel.layers module

core.tensor_parallel.layers

tensor_parallel.mappings module

core.tensor_parallel.mappings

tensor_parallel.random module

core.tensor_parallel.random

tensor_parallel.utils module

core.tensor_parallel.utils

Module contents

core.tensor_parallel