diff --git a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py index f0901955..16ca5c5a 100644 --- a/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py +++ b/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py @@ -154,7 +154,9 @@ def _parse_transformer_config(transformer_config: dict) -> dict: if dist.is_master(): run_config_name = f"{checkpoint_name}/modelopt_run_config.yaml" - config_dict = _parse_transformer_config(copy.deepcopy(model[0].config.__dict__)) + # We avoid deepcopy here since some attributes in Megatron-Bridge config cannot be + # deepcopy. + config_dict = _parse_transformer_config(model[0].config.__dict__) config_dict["nvidia_modelopt_version"] = modelopt.__version__ with open(run_config_name, "w") as f: yaml.dump(config_dict, f, default_flow_style=False)