def transform_config(
        self,
        original_config: PretrainedConfig,
    ) -> PretrainedConfig:
        """Transform the model config before saving.

        Args:
            original_config (Any): The original model config.

        Returns:
            The transformed model config.
        """
        copied_config = copy.deepcopy(original_config)
        if copied_config.model_type == 'mpt':
            copied_config.attn_config['attn_impl'] = 'torch'
            copied_config.init_device = 'cpu'
            if 'moe_world_size' in getattr(copied_config, 'ffn_config', {}):
                copied_config.ffn_config['moe_world_size'] = 1
        return copied_config
Transform the model config before saving.

        Args:
            original_config (Any): The original model config.

        Returns:
            The transformed model config.
        
llmfoundry/callbacks/hf_checkpointer.py