bridge.models.gpt.gpt_builder

Module Contents

Functions

transformer_engine_layer_spec	Create a Transformer Engine layer specification based on the provided config.
transformer_engine_full_layer_spec	Create a full Transformer Engine layer specification with autocast support.
local_layer_spec	Create a local layer specification without Transformer Engine.
modelopt_transformer_layer_spec	Layer specification for quantization with ModelOpt.
default_layer_spec	Determine the most appropriate layer specification based on availability.

Data

logger
__all__

API

bridge.models.gpt.gpt_builder.logger

‘getLogger(…)’

bridge.models.gpt.gpt_builder.transformer_engine_layer_spec(

config: megatron.training.models.gpt.GPTModelConfig,

) → megatron.core.transformer.ModuleSpec

Create a Transformer Engine layer specification based on the provided config.

bridge.models.gpt.gpt_builder.transformer_engine_full_layer_spec(

config: megatron.bridge.models.transformer_config.TransformerConfig,

) → megatron.core.transformer.ModuleSpec

Create a full Transformer Engine layer specification with autocast support.

Parameters:

config – GPT configuration object

Returns:

Module specification for full TE layers

Return type:

ModuleSpec

bridge.models.gpt.gpt_builder.local_layer_spec(

config: megatron.bridge.models.transformer_config.TransformerConfig,

) → megatron.core.transformer.ModuleSpec

Create a local layer specification without Transformer Engine.

Parameters:

config – GPT configuration object

Returns:

Module specification for local implementation layers

Return type:

ModuleSpec

bridge.models.gpt.gpt_builder.modelopt_transformer_layer_spec(

config: megatron.training.models.gpt.GPTModelConfig,

) → megatron.core.transformer.ModuleSpec

Layer specification for quantization with ModelOpt.

bridge.models.gpt.gpt_builder.default_layer_spec(

config: megatron.training.models.gpt.GPTModelConfig,

) → megatron.core.transformer.ModuleSpec

Determine the most appropriate layer specification based on availability.

bridge.models.gpt.gpt_builder.__all__

[‘GPTModelConfig’, ‘GPTModelBuilder’, ‘mtp_block_spec’, ‘transformer_engine_layer_spec’, ‘transforme…