`bridge.models.gpt.gpt_builder`#

Module Contents#

Classes#

`GPTModelConfig`	Bridge GPT config with transitional output-logit dtype support.
`GPTModelBuilder`	Bridge GPT builder that prevents silent fallback on older MCore.

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(…)’

class bridge.models.gpt.gpt_builder.GPTModelConfig#

Bases: megatron.training.models.gpt.GPTModelConfig

Bridge GPT config with transitional output-logit dtype support.

builder: ClassVar[str]#: ‘megatron.bridge.models.gpt.gpt_builder.GPTModelBuilder’

logit_dtype: torch.dtype | None#: None

class bridge.models.gpt.gpt_builder.GPTModelBuilder#

Bases: megatron.training.models.gpt.GPTModelBuilder

Bridge GPT builder that prevents silent fallback on older MCore.

build_model( pg_collection: megatron.core.process_groups_config.ProcessGroupCollection, pre_process: bool | None = None, post_process: bool | None = None, vp_stage: int | None = None, ) → megatron.core.models.gpt.GPTModel#: Build a GPT stage after validating MCore logit-dtype support.

bridge.models.gpt.gpt_builder.transformer_engine_layer_spec( config: bridge.models.gpt.gpt_builder.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: bridge.models.gpt.gpt_builder.GPTModelConfig, ) → megatron.core.transformer.ModuleSpec#: Layer specification for quantization with ModelOpt.

bridge.models.gpt.gpt_builder.default_layer_spec( config: bridge.models.gpt.gpt_builder.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…

bridge.models.gpt.gpt_builder#