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

.. automodule:: core.tensor_parallel.cross_entropy
   :members:
   :undoc-members:
   :show-inheritance:

tensor_parallel.data module

.. automodule:: core.tensor_parallel.data
   :members:
   :undoc-members:
   :show-inheritance:

tensor_parallel.layers module

.. automodule:: core.tensor_parallel.layers
   :members:
   :undoc-members:
   :show-inheritance:

tensor_parallel.mappings module

.. automodule:: core.tensor_parallel.mappings
   :members:
   :undoc-members:
   :show-inheritance:

tensor_parallel.random module

.. automodule:: core.tensor_parallel.random
   :members:
   :undoc-members:
   :show-inheritance:

tensor_parallel.utils module

.. automodule:: core.tensor_parallel.utils
   :members:
   :undoc-members:
   :show-inheritance:

Module contents

.. automodule:: core.tensor_parallel
   :members:
   :undoc-members:
   :show-inheritance: