core.models.T5.t5_spec

Module Contents

Functions

_te	The Transformer Engine provider, built on first use.
encoder_model_with_transformer_engine_default_spec	T5 encoder TE spec (uses Transformer Engine components).
decoder_model_with_transformer_engine_default_spec	T5 decoder TE spec (uses Transformer Engine components).
encoder_model_with_local_spec	T5 encoder local spec (uses Megatron-Core components).
decoder_model_with_local_spec	T5 decoder local spec (uses Megatron-Core components).
get_t5_encoder_with_transformer_engine_block_spec	T5 encoder block spec for Transformer Engine
get_t5_decoder_with_transformer_engine_block_spec	T5 decoder block spec for Transformer Engine
get_t5_encoder_with_local_block_spec	T5 encoder block spec for local (uses Megatron-Core components)
get_t5_decoder_with_local_block_spec	T5 decoder block spec for local (uses Megatron-Core components)

Data

_local

API

core.models.T5.t5_spec._te()

The Transformer Engine provider, built on first use.

Deferred so that importing this module does not require Transformer Engine – only building a Transformer Engine spec does. That is what the HAVE_TE guard used to buy.

core.models.T5.t5_spec._local

‘get_backend(…)’

core.models.T5.t5_spec.encoder_model_with_transformer_engine_default_spec() → megatron.core.transformer.spec_utils.ModuleSpec

T5 encoder TE spec (uses Transformer Engine components).

core.models.T5.t5_spec.decoder_model_with_transformer_engine_default_spec() → megatron.core.transformer.spec_utils.ModuleSpec

T5 decoder TE spec (uses Transformer Engine components).

core.models.T5.t5_spec.encoder_model_with_local_spec() → megatron.core.transformer.spec_utils.ModuleSpec

T5 encoder local spec (uses Megatron-Core components).

core.models.T5.t5_spec.decoder_model_with_local_spec() → megatron.core.transformer.spec_utils.ModuleSpec

T5 decoder local spec (uses Megatron-Core components).

core.models.T5.t5_spec.get_t5_encoder_with_transformer_engine_block_spec(

num_layers: int,

) → megatron.core.transformer.transformer_block.TransformerBlockSubmodules

T5 encoder block spec for Transformer Engine

Parameters:

config (TransformerConfig) – config, containing number of layers for encoder

core.models.T5.t5_spec.get_t5_decoder_with_transformer_engine_block_spec(

num_layers: int,

) → megatron.core.transformer.transformer_block.TransformerBlockSubmodules

T5 decoder block spec for Transformer Engine

Parameters:

config (TransformerConfig) – config, containing number of layers for decoder

core.models.T5.t5_spec.get_t5_encoder_with_local_block_spec(

num_layers: int,

) → megatron.core.transformer.transformer_block.TransformerBlockSubmodules

T5 encoder block spec for local (uses Megatron-Core components)

Parameters:

num_layers (int) – number of encoder layers

core.models.T5.t5_spec.get_t5_decoder_with_local_block_spec(

num_layers: int,

) → megatron.core.transformer.transformer_block.TransformerBlockSubmodules

T5 decoder block spec for local (uses Megatron-Core components)

Parameters:

num_layers (int) – number of decoder layers