`bridge.models.transformer_config`#

Bridge wrapper classes for Megatron Core transformer configurations.

These classes provide deferred post-initialization to support the Bridge configuration override system while maintaining compatibility with Megatron Core’s post_init behavior.

Module Contents#

Classes#

`TransformerConfig`	Megatron Core TransformerConfig with deferred post-init.
`MLATransformerConfig`	Megatron Core MLATransformerConfig with deferred post-init.
`HeterogeneousTransformerConfig`	Megatron Core HeterogeneousTransformerConfig with deferred post-init.

Functions#

`_safe_asdict`	Shallow asdict variant that preserves handles like process groups.
`_resolve_string_fields`	Resolve string-valued fields to their runtime types.
`_normalize_attention_backend`	Resolve an unset recipe backend to MCore’s supported automatic selection.
`_set_moe_expert_tensor_parallel_default`	Default expert tensor parallelism to one when expert parallelism is enabled.
`_enable_safe_hybridep_dispatch`	Enable uneven-input padding for an eager HybridEP THD recipe.
`_patch_yarn_concentration_factor`	Patch MCore _yarn_get_concentration_factor_from_config for None handling.

Data#

_HYBRIDEP_PADDING_FIELDS

API#

bridge.models.transformer_config._safe_asdict(obj, skip_keys: set[str]) → dict#

Shallow asdict variant that preserves handles like process groups.

dataclasses.asdict performs a deep copy of every leaf value, which breaks objects that should remain shared references (e.g., ProcessGroupCollection). This helper mirrors the structure of asdict but returns leaf objects as-is so they are not deep-copied.

bridge.models.transformer_config._resolve_string_fields( config: megatron.core.transformer.transformer_config.TransformerConfig, ) → None#

Resolve string-valued fields to their runtime types.

Handles activation_func (e.g. model.activation_func=silu) and dtype fields params_dtype / pipeline_dtype (e.g. model.params_dtype=bf16) when they arrive as strings via CLI overrides.

bridge.models.transformer_config._normalize_attention_backend( config: megatron.core.transformer.transformer_config.TransformerConfig, ) → None#

Resolve an unset recipe backend to MCore’s supported automatic selection.

Explicit backend choices and Transformer Engine environment variables are left unchanged. MCore validates that those process-wide settings agree.

bridge.models.transformer_config._HYBRIDEP_PADDING_FIELDS#: (‘moe_hybridep_pad_uneven_dispatch_inputs’, ‘moe_hybridep_pad_variable_tokens’)

bridge.models.transformer_config._set_moe_expert_tensor_parallel_default( config: megatron.core.transformer.transformer_config.TransformerConfig, ) → None#

Default expert tensor parallelism to one when expert parallelism is enabled.

Megatron Core resolves an unset expert tensor parallel size to the tensor model parallel size. Bridge MoE configurations instead scale expert layers with expert parallelism by default, while preserving explicit expert tensor parallel values.

bridge.models.transformer_config._enable_safe_hybridep_dispatch( config: megatron.core.transformer.transformer_config.TransformerConfig, *, uses_thd: bool, ) → None#

Enable uneven-input padding for an eager HybridEP THD recipe.

The combined training recipe owns the tensor layout, so it requests this path only when its dataset produces THD packed-sequence metadata. CUDA-graph configs retain their explicit setting because the padding path’s host scalar synchronization is not capture-safe; those configs require equal inputs.

bridge.models.transformer_config._patch_yarn_concentration_factor() → None#

Patch MCore _yarn_get_concentration_factor_from_config for None handling.

TransformerConfig declares yarn_rotary_scaling_factor as float | None = None, but MCore uses hasattr(), which returns True for dataclass fields set to None. This causes a crash for non-YaRN models. Use getattr + is not None instead.

TODO: Remove once upstream MCore merges the fix.

class bridge.models.transformer_config.TransformerConfig#

Bases: megatron.core.transformer.transformer_config.TransformerConfig

Megatron Core TransformerConfig with deferred post-init.

This class inherits from Megatron Core’s TransformerConfig but defers the execution of post_init() until finalize() is explicitly called. This allows for field modifications after construction but before computed fields are calculated.

Usage: # Create config with deferred post-init config = TransformerConfig(num_layers=32, hidden_size=4096)

# Modify fields as needed
config.seq_length = 8192
config.tensor_model_parallel_size = 2

# Finalize to compute derived fields
config.finalize()

_NO_COPY_KEYS#: None

tensor_parallel_num_weight_shards: int | None#: None

expert_tensor_parallel_num_weight_shards: int | None#: None

gtp_weight_remat_size: int#: ‘field(…)’

expert_gtp_weight_remat_size: int#: ‘field(…)’

yarn_rotary_scaling_factor: float | None#: None

yarn_original_max_position_embeddings: int | None#: None

yarn_beta_fast: float | None#: None

yarn_beta_slow: float | None#: None

yarn_mscale: float | None#: None

yarn_mscale_all_dim: float | None#: None

yarn_correction_range_round_to_int: bool | None#: None

__post_init__() → None#

Skip MCore post_init during initial construction.

The original post_init logic is deferred until finalize() is called. This allows for field modifications after construction without invalidating computed fields.

finalize() → None#

Execute the deferred MCore post-init logic.

This method calls the original Megatron Core TransformerConfig.post_init() to compute derived fields based on the current field values. It can be called multiple times safely.

_finalize_gtp_weight_shards() → None#: Derive rematerialization sizes from optional logical weight-shard counts.

__deepcopy__(memo)#

Custom deepcopy to preserve process group handles when cloning configs.

Certain attributes (_pg_collection, etc.) should remain shared references rather than being wiped or re-created during deepcopy. TODO: This is a temporary hack. Once providers stop embedding the Transformer config and instead hold the MCore config as an attribute, we can remove this.

asdict() → dict#: Return a dict view without deep-copying shared handles (e.g., process groups).

class bridge.models.transformer_config.MLATransformerConfig#

Bases: bridge.models.transformer_config.TransformerConfig, megatron.core.transformer.transformer_config.MLATransformerConfig

Megatron Core MLATransformerConfig with deferred post-init.

This class inherits from Megatron Core’s MLATransformerConfig but defers the execution of post_init() until finalize() is explicitly called. This allows for field modifications after construction but before computed fields are calculated.

Usage: # Create config with deferred post-init config = MLATransformerConfig(num_layers=32, hidden_size=4096)

# Modify fields as needed
config.q_lora_rank = 1536
config.kv_lora_rank = 512

# Finalize to compute derived fields
config.finalize()

multi_latent_attention: bool#: True

__post_init__() → None#

Skip MCore post_init during initial construction.

The original post_init logic is deferred until finalize() is called. This allows for field modifications after construction without invalidating computed fields.

finalize() → None#

Execute the deferred MCore post-init logic.

This method calls the original Megatron Core MLATransformerConfig.post_init() to compute derived fields based on the current field values. It can be called multiple times safely.

class bridge.models.transformer_config.HeterogeneousTransformerConfig#

Bases: bridge.models.transformer_config.TransformerConfig, megatron.core.transformer.heterogeneous.heterogeneous_config.HeterogeneousTransformerConfig

Megatron Core HeterogeneousTransformerConfig with deferred post-init.

This class inherits from both our lazy TransformerConfig and Megatron Core’s HeterogeneousTransformerConfig. The MRO ensures that our lazy post-init behavior is preserved while maintaining all heterogeneous functionality.

CRITICAL: The inheritance order is important for MRO:

TransformerConfig (our lazy version) comes first
MCoreHeterogeneousTransformerConfig comes second

Usage: # Create config with deferred post-init config = HeterogeneousTransformerConfig( num_layers=32, hidden_size=4096, heterogeneous_layers_config_encoded_json=json_string )

# Modify fields as needed
config.seq_length = 8192
config.tensor_model_parallel_size = 2

# Finalize to compute derived fields and parse heterogeneous config
config.finalize()

__post_init__() → None#

Skip MCore post_init during initial construction.

The original post_init logic is deferred until finalize() is called. This allows for field modifications after construction without invalidating computed fields.

finalize() → None#

Execute the deferred MCore post-init logic.

This method calls the original Megatron Core HeterogeneousTransformerConfig.post_init() to compute derived fields and parse the heterogeneous block configurations. It can be called multiple times safely.

get_config_for_layer( layer_number: int, ) → megatron.core.transformer.transformer_config.TransformerConfig#: Return a layer-specific TransformerConfig without deep-copying process groups.

bridge.models.transformer_config#

Module Contents#

Classes#

Functions#

Data#

API#

`bridge.models.transformer_config`#