core.models.T5.t5_spec

Module Contents

Functions

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)

API

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