bridge.models.conversion.utils

Module Contents

Functions

mcore_to_hf_window_size	Convert an MCore inclusive attention window to the Hugging Face token count.
unwrap_model	Unwrap a model (or list of models) to the underlying module. Extends megatron.core.utils.unwrap_model with awareness of MegatronFSDP.
weights_verification_table	Returns a table comparing weights between a Hugging Face model and a Megatron-LM model.
get_module_and_param_from_name	Get parameter from specific VP stage, ensuring that parameter attributes are preserved. Supports both absolute and relative parameter names.
remove_non_pickleables	Remove non-pickleable objects from a configuration object recursively.
extract_sort_key	Extract sorting key based on layer and expert numbers.
moe_experts_stored_packed	Whether a checkpoint stores routed MoE experts fused rather than per-expert.
get_causal_lm_class_name_via_auto_map	Return CausalLM class name via config.auto_map if available; otherwise None.
conform_config_to_reference	Return hf_config_dict projected onto reference keys, filling missing values from reference_config.
persistent_buffers	Return an iterator over persistent module buffers, yielding both the name of the buffer as well as the buffer itself.
is_modelopt_dynamic_module	Check if a module is a modelopt dynamic module.

API

bridge.models.conversion.utils.mcore_to_hf_window_size(

window_size: int | list[int] | tuple[int, int] | None,

) → int | None

Convert an MCore inclusive attention window to the Hugging Face token count.

MCore represents a causal sliding window as (left, right), where the current token is included in the Hugging Face window size. Checkpoint YAML reloads tuples as lists, so both sequence types are accepted. Providers that already store the Hugging Face scalar representation pass through unchanged.

bridge.models.conversion.utils.unwrap_model(model, module_instances=None)

Unwrap a model (or list of models) to the underlying module. Extends megatron.core.utils.unwrap_model with awareness of MegatronFSDP.

bridge.models.conversion.utils.weights_verification_table(bridge, megatron_model) → rich.table.Table

Returns a table comparing weights between a Hugging Face model and a Megatron-LM model.

Parameters:

• bridge (AutoBridge) – The bridge object containing model information.

• megatron_model – The Megatron-LM model instance.

Returns:

A rich Table object with the comparison.

Return type:

Table

bridge.models.conversion.utils.get_module_and_param_from_name(

models: megatron.core.transformer.module.MegatronModule | List[megatron.core.transformer.module.MegatronModule],

param_name: str,

vp_stage: Optional[int] = None,

) → Tuple[torch.nn.Module, torch.Tensor] | Tuple[torch.nn.Module, torch.Tensor, Tuple]

Get parameter from specific VP stage, ensuring that parameter attributes are preserved. Supports both absolute and relative parameter names.

Parameters:

• models – List of Megatron model instances or a submodule

• param_name – Dot-separated parameter name (can be absolute or relative to models)

• vp_stage – Virtual pipeline stage index (None for single stage)

Returns:

Tuple of (module, parameter) where module owns the parameter

Raises:

ValueError – If vp_stage is out of range or parameter doesn’t exist

.. rubric:: Examples

Basic usage with full model:

module, param = get_module_and_param_from_name( … models=full_model, … param_name=”transformer.layers.0.attention.query.weight” … )

Usage with model list and VP stage:

module, param = get_module_and_param_from_name( … models=[model1, model2, model3], … param_name=”layers.0.mlp.dense.bias”, … vp_stage=1 … )

Usage with submodule and relative path:

linear_module = model.transformer.layers[0].mlp.dense module, param = get_module_and_param_from_name( … models=linear_module, … param_name=”weight” … )

Usage with submodule and absolute path (automatic suffix matching):

linear_module = model.transformer.layers[0].mlp.dense module, param = get_module_and_param_from_name( … models=linear_module, … param_name=”transformer.layers.0.mlp.dense.weight” … )

Automatically matches “weight” suffix and returns the parameter

Edge case with partial path matching:

attention_module = model.transformer.layers[0].attention module, param = get_module_and_param_from_name( … models=attention_module, … param_name=”layers.0.attention.query.weight” … )

Matches “query.weight” suffix within the attention module

bridge.models.conversion.utils.remove_non_pickleables(

obj,

max_depth: int = 3,

current_depth: int = 0,

)

Remove non-pickleable objects from a configuration object recursively.

This utility function identifies and removes objects that cannot be pickled for inter-process communication, including functions, bound methods, partial functions, and other problematic callables.

Parameters:

• obj – The object to clean

• max_depth – Maximum recursion depth (default: 3)

• current_depth – Current recursion depth (internal use)

Returns:

The cleaned object with non-pickleables removed

bridge.models.conversion.utils.extract_sort_key(param_name: str)

Extract sorting key based on layer and expert numbers.

bridge.models.conversion.utils.moe_experts_stored_packed(

hf_pretrained,

layers_prefix: str,

default: bool = False,

) → bool

Whether a checkpoint stores routed MoE experts fused rather than per-expert.

Fused layout keys look like {layers_prefix}<L>.mlp.experts.gate_up_proj / .down_proj (a single stacked tensor per projection), while the per-expert layout (what save_pretrained writes) uses {layers_prefix}<L>.mlp.experts.<i>.gate_proj|up_proj|down_proj.weight. Which one a checkpoint uses depends on the transformers version, so the bridge selects mappings accordingly.

Parameters:

• hf_pretrained – The loaded HF model wrapper (its state.source provides the checkpoint keys).

• layers_prefix – HF prefix up to and including layers. (e.g. "model.layers." for an LLM, "model.language_model.layers." for a VLM).

• default – Value returned when the checkpoint keys are unavailable (e.g. building mappings without a loaded checkpoint, as in unit tests) or no routed-expert keys are found.

Returns:

True if experts are stored fused, False if per-expert, default if undetermined.

Raises:

ValueError – if the checkpoint mixes fused and per-expert layouts.

bridge.models.conversion.utils.get_causal_lm_class_name_via_auto_map(

config: transformers.configuration_utils.PretrainedConfig,

) → str | None

Return CausalLM class name via config.auto_map if available; otherwise None.

If auto_map[“AutoModelForCausalLM”] is present in the config, returns the class name string extracted from the mapping value by splitting on ‘.’ and taking the last segment. Returns None if auto_map is not set.

bridge.models.conversion.utils.conform_config_to_reference(

hf_config_dict: dict[str, object],

reference_config: dict[str, object],

) → dict[str, object]

Return hf_config_dict projected onto reference keys, filling missing values from reference_config.

bridge.models.conversion.utils.persistent_buffers(

model: torch.nn.Module,

) → Iterable[Tuple[str, torch.Tensor]]

Return an iterator over persistent module buffers, yielding both the name of the buffer as well as the buffer itself.

bridge.models.conversion.utils.is_modelopt_dynamic_module(module)

Check if a module is a modelopt dynamic module.