def _assert_checkpoint_equivalence(
    tmp_path: pathlib.Path,
    expected_normal_checkpoints: int,
    expected_hf_checkpoints: int,
    trainer: Trainer,
    batches_per_epoch: int,
    precision: torch.dtype,
    model: str,
    tokenizer: PreTrainedTokenizerBase,
    original_model: ComposerModel,
    fsdp_state_dict_type: Optional[str] = None,
    peft_config: Optional[dict] = None,
):
    """Asserts the equivalence of checkpoints.

    Asserts equivalence of checkpoints between the original mpt model and the converted hf model.

    Args:
        tmp_path (str): The path to the temporary directory where the checkpoints are saved.
        expected_normal_checkpoints (int): The expected number of normal checkpoints.
        expected_hf_checkpoints (int): The expected number of HuggingFace checkpoints.
        trainer (Trainer): The trainer object used for training the model.
        batches_per_epoch (int): The number of batches per epoch.
        precision (torch.dtype): The precision of the model.
        model (str): The type of model ('mpt', 'neo', or 'llama2').
        tokenizer (PreTrainedTokenizerBase): The model tokenizer.
        original_model (ComposerModel): The original model object.
        fsdp_state_dict_type (Optional[str], optional): The type of FSDP state dict. Defaults to None.
        peft_config (Optional[dict], optional): The PEFT configuration. Defaults to None.
    """
    loaded_model = None
    loaded_tokenizer = None
    # Only rank zero is saving the huggingface checkpoints, so only check
    # for equivalence on rank zero
    if dist.get_global_rank() == 0:
        normal_checkpoints = [
            name for name in os.listdir(os.path.join(tmp_path, 'checkpoints'))
            if name != 'huggingface'
        ]
        huggingface_checkpoints = list(
            os.listdir(os.path.join(tmp_path, 'checkpoints', 'huggingface')),
        )

        checkpoint_files = os.listdir(
            os.path.join(
                tmp_path,
                'checkpoints',
                'huggingface',
                huggingface_checkpoints[-1],
            ),
        )
        if peft_config is not None:
            assert 'adapter_config.json' in checkpoint_files
            assert 'adapter_model.safetensors' in checkpoint_files

        assert len(normal_checkpoints) == expected_normal_checkpoints
        assert len(huggingface_checkpoints) == expected_hf_checkpoints

        # Patch flash_attn package to be empty to simulate loading the model in
        # an environment without flash attention installed
        with patch.dict('sys.modules', {'flash_attn': None}):
            if peft_config is not None:
                composer_model = trainer.state.model.module if trainer.state.is_model_ddp else trainer.state.model
                composer_model.model.base_model.save_pretrained( # type: ignore
                    tmp_path / 'base-model',
                )

            checkpoint_path = os.path.join(
                tmp_path,
                'checkpoints',
                'huggingface',
                f'ba{batches_per_epoch}',
            )

            if peft_config is not None:
                with open(
                    os.path.join(checkpoint_path, 'adapter_config.json'),
                ) as _f:
                    adapter_config = json.load(_f)

                adapter_config['base_model_name_or_path'] = str(
                    tmp_path / 'base-model',
                )

                with open(
                    os.path.join(checkpoint_path, 'adapter_config.json'),
                    'w',
                ) as _f:
                    json.dump(adapter_config, _f)

            # Load the last huggingface checkpoint
            loaded_model = transformers.AutoModelForCausalLM.from_pretrained(
                checkpoint_path,
                trust_remote_code=True,
                torch_dtype=precision,
            )

        # Check that the loaded model has the correct precision, and then set it back
        # to the original for the equivalence check
        if peft_config is None:
            assert loaded_model.config.torch_dtype == precision
            loaded_model.config.torch_dtype = original_model.model.config.torch_dtype  # type: ignore

        if model == 'mpt':
            # Check that we have correctly set these attributes, and then set them back
            # to the original for the equivalence check
            assert loaded_model.config.attn_config['attn_impl'] == 'torch'
            assert loaded_model.config.init_device == 'cpu'
            loaded_model.config.attn_config[
                'attn_impl'
            ] = original_model.model.config.attn_config[  # type: ignore
                'attn_impl']
            loaded_model.config.init_device = original_model.model.config.init_device  # type: ignore

        loaded_tokenizer = transformers.AutoTokenizer.from_pretrained(
            os.path.join(
                tmp_path,
                'checkpoints',
                'huggingface',
                f'ba{batches_per_epoch}',
            ),
            trust_remote_code=True,
        )

        check_hf_model_equivalence(
            trainer.state.model.model.to(precision) if fsdp_state_dict_type  # type: ignore
            is not None else trainer.state.model.module.model.to(precision), # type: ignore
            loaded_model,
            just_lora=peft_config is not None,
        )
        check_hf_tokenizer_equivalence(tokenizer, loaded_tokenizer)
Asserts the equivalence of checkpoints.

    Asserts equivalence of checkpoints between the original mpt model and the converted hf model.

    Args:
        tmp_path (str): The path to the temporary directory where the checkpoints are saved.
        expected_normal_checkpoints (int): The expected number of normal checkpoints.
        expected_hf_checkpoints (int): The expected number of HuggingFace checkpoints.
        trainer (Trainer): The trainer object used for training the model.
        batches_per_epoch (int): The number of batches per epoch.
        precision (torch.dtype): The precision of the model.
        model (str): The type of model ('mpt', 'neo', or 'llama2').
        tokenizer (PreTrainedTokenizerBase): The model tokenizer.
        original_model (ComposerModel): The original model object.
        fsdp_state_dict_type (Optional[str], optional): The type of FSDP state dict. Defaults to None.
        peft_config (Optional[dict], optional): The PEFT configuration. Defaults to None.
    
tests/a_scripts/inference/test_convert_composer_to_hf.py