Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 16 additions & 2 deletions src/megatron/bridge/models/conversion/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -309,9 +309,23 @@ def get_causal_lm_class_name_via_auto_map(
def conform_config_to_reference(
hf_config_dict: dict[str, object], reference_config: dict[str, object]
) -> dict[str, object]:
"""Return a projected hf_config_dict onto the reference key set, imputing missing keys with reference values."""
"""Return hf_config_dict projected onto reference keys, imputing missing values from reference.

Nested config dictionaries are projected recursively so export-time config
reconstruction does not lose sub-config fields that were not emitted by the
Megatron-side config.
"""
reference_config_keys = set(reference_config.keys())
filtered_config_dict = {key: value for (key, value) in hf_config_dict.items() if key in reference_config_keys}
filtered_config_dict = {}
for key, value in hf_config_dict.items():
if key not in reference_config_keys:
continue

reference_value = reference_config[key]
if isinstance(value, dict) and isinstance(reference_value, dict):
value = conform_config_to_reference(value, reference_value)
filtered_config_dict[key] = value

for key, value in reference_config.items():
if key not in filtered_config_dict:
filtered_config_dict[key] = value
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -354,6 +354,12 @@ def __init__(
self.text_config = text_config
self.audio_token_id = audio_token_id

def to_cfg_dict(self) -> dict[str, object]:
"""Return an instantiable config dictionary for Megatron Bridge run_config.yaml."""
config = self.to_dict()
config["_target_"] = f"{self.__class__.__module__}.{self.__class__.__qualname__}"
return config


class Qwen3ASRConfig(PretrainedConfig):
"""
Expand Down
Loading
Loading