`bridge.models.bagel.provider`#

Megatron Core model provider for BAGEL-7B-MoT.

Module Contents#

Classes#

BagelModelProvider

Build the Qwen2-MoT, packed SigLIP, and diffusion BAGEL model.

Functions#

`_complete_mot_layer_sharded_state_dict`	Include MoT layer norms omitted by the current upstream sharded-state helper.
`gelu_pytorch_tanh`	Apply the connector activation used by official BAGEL.

Data#

logger

API#

bridge.models.bagel.provider.logger#: ‘getLogger(…)’

bridge.models.bagel.provider._complete_mot_layer_sharded_state_dict( layer: torch.nn.Module, prefix: str = '', sharded_offsets: tuple = (), metadata: dict[str, Any] | None = None, ) → dict[str, Any]#: Include MoT layer norms omitted by the current upstream sharded-state helper.

bridge.models.bagel.provider.gelu_pytorch_tanh(value: torch.Tensor) → torch.Tensor#: Apply the connector activation used by official BAGEL.

class bridge.models.bagel.provider.BagelModelProvider#

Bases: megatron.bridge.models.gpt_provider.GPTModelProvider

Build the Qwen2-MoT, packed SigLIP, and diffusion BAGEL model.

num_layers: int#: 28

hidden_size: int#: 3584

num_attention_heads: int#: 28

num_query_groups: int#: 4

ffn_hidden_size: int#: 18944

seq_length: int#: 36864

max_position_embeddings: int#: 32768

vocab_size: int#: 152064

activation_func: Any#: None

gated_linear_unit: bool#: True

normalization: str#: ‘RMSNorm’

rms_norm_eps: float#: 1e-06

layernorm_epsilon: float#: 1e-06

position_embedding_type: str#: ‘rope’

rotary_base: float#: 1000000.0

attention_dropout: float#: 0.0

hidden_dropout: float#: 0.0

add_bias_linear: bool#: False

add_qkv_bias: bool#: True

qk_layernorm: bool#: True

bias_activation_fusion: bool#: True

bias_dropout_fusion: bool#: True

apply_rope_fusion: bool#: True

persist_layer_norm: bool#: True

share_embeddings_and_output_weights: bool#: False

bf16: bool#: True

params_dtype: torch.dtype#: None

pipeline_dtype: torch.dtype#: None

use_flex_attention: bool#: True

moe_token_dispatcher_type: str#: ‘alltoall’

bagel_repo: str | None#: None

model_path: str | None#: None

official_config_values: dict[str, Any] | None#: None

vision_model_path: str | None#: None

vae_path: str | None#: None

latent_patch_size: int#: 2

max_latent_size: int#: 64

max_num_patch_per_side: int#: 70

timestep_shift: float#: 1.0

ce_weight: float#: 1.0

mse_weight: float#: 1.0

ce_loss_reweighting: bool#: False

recompute_vit: bool#: False

freeze_vit: bool#: False

freeze_llm: bool#: False

native_model_checkpoint: str | None#: None

native_model_seed: int | None#: None

native_world_size: int | None#: None

validate_native_checkpoint_metadata: bool#: True

reference_training_seed: int | None#: None

reference_training_world_size: int | None#: None

reset_reference_training_rng: bool#: False

finalize() → None#: Validate the first supported BAGEL training topology.

_get_num_floating_point_operations_with_runtime_stats( *, batch_size: int, seqlen_sum: int | None, seqlen_squared_sum: int | None, cross_seqlen_sum: int | None = None, cross_seqlen_product_sum: int | None = None, ) → float#: Return the official BAGEL Qwen2 training FLOPs estimate.

_official_config() → Any#: Load the official local BAGEL configuration without model weights.

provide( pre_process=None, post_process=None, vp_stage=None, ) → torch.nn.Module#: Instantiate PR #3635’s MCore BAGEL model.

bridge.models.bagel.provider#