`bridge.models.bagel.checkpoint`#

Strict Bridge adapter from a native BAGEL checkpoint to MCore.

Native BAGEL stores separate Q/K/V and gate/up projections. MCore stores those projections in group-interleaved QKV and concatenated GLU tensors, respectively. This module owns that representation change and verifies it by applying the inverse mapping to every fused destination tensor after it is copied.

The adapter intentionally supports only an unsharded TP=1, PP=1, non-VP model. Loading a full native tensor into a partial MCore model would otherwise appear to work until a missing or incorrectly sharded parameter is used.

Module Contents#

Classes#

`NativeCheckpointLoadReport`	Counts proving complete source consumption and target verification.
`_TensorSource`	Track exact, single-use consumption of a safetensors file.
`_LLMConfig`	BAGEL language configuration fields required by the remapper.

Functions#

`_fuse_qkv`	Convert native [Q, K, V] tensors to MCore’s group-interleaved layout.
`_unfuse_qkv`	Invert :func:`_fuse_qkv` without relying on the source tensors.
`_fuse_glu`	Convert native gate/up projections to MCore’s [gate; up] FC1.
`_unfuse_glu`	Invert :func:`_fuse_glu` into native gate and up projections.
`_reference_on_destination`
`_assert_equal_to_source`
`_get_high_precision_init_val`
`_clear_high_precision_init_val`
`_preserve_fp32_main_init`	Expose the FP32 source through Megatron-FSDP’s main-weight init hook.
`_copy_parameter`
`_copy_linear`
`_require_bias_free`
`_pre_mlp_norm_weight`
`_copy_attention_branch`
`_copy_mlp_branch`
`_validate_unsharded_model`
`_language_keys`
`_metadata_int`
`_validate_metadata`
`_initialize_language`
`_initialize_vision`
`_initialize_auxiliary`
`_assert_exact_target_coverage`	Require every semantic MCore parameter and position buffer exactly once.
`initialize_bagel_from_native_checkpoint`	Initialize an unsharded MCore BAGEL model from native BAGEL safetensors.

Data#

`_FORMAT_VERSION`
`_HIGH_PRECISION_VALUE_ATTR`
`_ACTIVE_TARGETS`
`_AUXILIARY_KEYS`
`_LANGUAGE_GLOBAL_KEYS`
`_LANGUAGE_LAYER_SUFFIXES`

API#

bridge.models.bagel.checkpoint._FORMAT_VERSION#: 1

bridge.models.bagel.checkpoint._HIGH_PRECISION_VALUE_ATTR#: ‘_bagel_native_high_precision_init_val’

bridge.models.bagel.checkpoint._ACTIVE_TARGETS: contextvars.ContextVar[dict[int, str] | None]#: ‘ContextVar(…)’

bridge.models.bagel.checkpoint._AUXILIARY_KEYS#: ‘frozenset(…)’

bridge.models.bagel.checkpoint._LANGUAGE_GLOBAL_KEYS#: ‘frozenset(…)’

bridge.models.bagel.checkpoint._LANGUAGE_LAYER_SUFFIXES#: (‘input_layernorm.weight’, ‘input_layernorm_moe_gen.weight’, ‘post_attention_layernorm.weight’, ‘pos…

class bridge.models.bagel.checkpoint.NativeCheckpointLoadReport#

Counts proving complete source consumption and target verification.

source_tensors_consumed: int#: None

target_tensors_verified: int#: None

fp32_main_tensors_preserved: int#: None

class bridge.models.bagel.checkpoint._TensorSource(checkpoint)#

Track exact, single-use consumption of a safetensors file.

Initialization

take(key: str) → torch.Tensor#

require_exact_keys(expected: set[str]) → None#

assert_all_consumed() → None#

property consumed_count: int#

class bridge.models.bagel.checkpoint._LLMConfig#

Bases: typing.Protocol

BAGEL language configuration fields required by the remapper.

num_hidden_layers: int#: None

hidden_size: int#: None

num_attention_heads: int#: None

num_key_value_heads: int#: None

head_dim: int | None#: None

bridge.models.bagel.checkpoint._fuse_qkv( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, *, num_attention_heads: int, num_query_groups: int, head_dim: int, ) → torch.Tensor#: Convert native [Q, K, V] tensors to MCore’s group-interleaved layout.

bridge.models.bagel.checkpoint._unfuse_qkv( fused: torch.Tensor, *, num_attention_heads: int, num_query_groups: int, head_dim: int, ) → tuple[torch.Tensor, torch.Tensor, torch.Tensor]#: Invert :func:_fuse_qkv without relying on the source tensors.

bridge.models.bagel.checkpoint._fuse_glu(gate: torch.Tensor, up: torch.Tensor) → torch.Tensor#: Convert native gate/up projections to MCore’s [gate; up] FC1.

bridge.models.bagel.checkpoint._unfuse_glu( fused: torch.Tensor, *, ffn_hidden_size: int, ) → tuple[torch.Tensor, torch.Tensor]#: Invert :func:_fuse_glu into native gate and up projections.

bridge.models.bagel.checkpoint._reference_on_destination( source: torch.Tensor, destination: torch.Tensor, ) → torch.Tensor#

bridge.models.bagel.checkpoint._assert_equal_to_source( destination: torch.Tensor, source: torch.Tensor, name: str, ) → None#

bridge.models.bagel.checkpoint._get_high_precision_init_val( parameter: torch.nn.Parameter, ) → torch.Tensor#

bridge.models.bagel.checkpoint._clear_high_precision_init_val(parameter: torch.nn.Parameter) → None#

bridge.models.bagel.checkpoint._preserve_fp32_main_init( destination: torch.Tensor, source: torch.Tensor, name: str, ) → None#: Expose the FP32 source through Megatron-FSDP’s main-weight init hook.

bridge.models.bagel.checkpoint._copy_parameter( destination: torch.Tensor, source: torch.Tensor, name: str, ) → None#

bridge.models.bagel.checkpoint._copy_linear( reader: bridge.models.bagel.checkpoint._TensorSource, module, prefix: str, ) → None#

bridge.models.bagel.checkpoint._require_bias_free(module, name: str) → None#

bridge.models.bagel.checkpoint._pre_mlp_norm_weight(layer, branch_suffix: str) → torch.Tensor#

bridge.models.bagel.checkpoint._copy_attention_branch( reader: bridge.models.bagel.checkpoint._TensorSource, mcore_attention, source_layer_prefix: str, *, source_suffix: str, destination_suffix: str, num_attention_heads: int, num_query_groups: int, head_dim: int, ) → None#

bridge.models.bagel.checkpoint._copy_mlp_branch( reader: bridge.models.bagel.checkpoint._TensorSource, mcore_mlp, source_prefix: str, ) → None#

bridge.models.bagel.checkpoint._validate_unsharded_model(language_model) → None#

bridge.models.bagel.checkpoint._language_keys(num_layers: int) → set[str]#

bridge.models.bagel.checkpoint._metadata_int(metadata: Mapping[str, str], key: str) → int#

bridge.models.bagel.checkpoint._validate_metadata( metadata: Mapping[str, str] | None, *, expected_model_seed: int, expected_world_size: int, ) → None#

bridge.models.bagel.checkpoint._initialize_language( reader: bridge.models.bagel.checkpoint._TensorSource, language_model, llm_config: bridge.models.bagel.checkpoint._LLMConfig, ) → None#

bridge.models.bagel.checkpoint._initialize_vision( reader: bridge.models.bagel.checkpoint._TensorSource, mimo_model, ) → None#

bridge.models.bagel.checkpoint._initialize_auxiliary( reader: bridge.models.bagel.checkpoint._TensorSource, mimo_model, ) → None#

bridge.models.bagel.checkpoint._assert_exact_target_coverage( mimo_model: torch.nn.Module, initialized_targets: Mapping[int, str], *, require_fp32_main_values: bool, ) → None#: Require every semantic MCore parameter and position buffer exactly once.

bridge.models.bagel.checkpoint.initialize_bagel_from_native_checkpoint( mimo_model: torch.nn.Module, checkpoint_path: str, *, expected_model_seed: int, expected_world_size: int, validate_metadata: bool = True, llm_config: bridge.models.bagel.checkpoint._LLMConfig, ) → bridge.models.bagel.checkpoint.NativeCheckpointLoadReport#

Initialize an unsharded MCore BAGEL model from native BAGEL safetensors.

The checkpoint must contain the complete native Bagel.state_dict(). Exported initialization checkpoints also validate format_version, model_seed, and world_size metadata. Official BAGEL EMA releases omit that metadata and may set validate_metadata=False. Every tensor must map exactly once; missing and additional tensors are rejected before any parameter is modified.

bridge.models.bagel.checkpoint#