def _assert_saved_model_is_equal(fabric, model, checkpoint_path):
    """Convert the saved checkpoint to a single file with the model weights consolidated to easily verify the full
    weights in float32 precision."""
    from deepspeed.utils.zero_to_fp32 import convert_zero_checkpoint_to_fp32_state_dict

    assert isinstance(fabric.strategy, DeepSpeedStrategy)

    # carry out the check only on rank 0
    if fabric.is_global_zero:
        if fabric.strategy.config["zero_optimization"]["stage"] in (2, 3):
            single_ckpt_path = checkpoint_path / "single_model.pt"
            # the tag is hardcoded in DeepSpeedStrategy
            convert_zero_checkpoint_to_fp32_state_dict(checkpoint_path, single_ckpt_path, tag="checkpoint")
            state_dict = torch.load(single_ckpt_path, weights_only=False)
        else:
            # 'checkpoint' is the tag, hardcoded in DeepSpeedStrategy
            single_ckpt_path = checkpoint_path / "checkpoint" / "mp_rank_00_model_states.pt"
            state_dict = torch.load(single_ckpt_path, weights_only=False)["module"]

        model = model.cpu()

        # assert model parameters are identical after loading
        for orig_param, saved_model_param in zip(model.parameters(), state_dict.values()):
            # perform the equality check in the same precision
            saved_model_param = saved_model_param.cpu().to(orig_param.dtype)
            assert torch.equal(orig_param, saved_model_param)

    fabric.barrier()
Convert the saved checkpoint to a single file with the model weights consolidated to easily verify the full
    weights in float32 precision.
tests/tests_fabric/strategies/test_deepspeed_integration.py