    def test_default_distributed_with_sharded_state(self):
        """
        Simulates using the Checkpointer during distributed training with a sharded model.
        """
        world_size = 2
        default_num_to_keep = 2
        num_epochs = 5
        target = [(e, 0) for e in range(num_epochs - default_num_to_keep, num_epochs)]

        checkpointers = [Checkpointer(serialization_dir=self.TEST_DIR) for _ in range(world_size)]
        for i, checkpointer in enumerate(checkpointers):
            checkpointer._rank = i
            checkpointer.state_is_sharded = True

        for epochs_completed in range(num_epochs):
            for batches_completed in [0, 5, 10]:
                for i, checkpointer in enumerate(checkpointers):
                    state = {
                        "epochs_completed": epochs_completed,
                        "batches_in_epoch_completed": batches_completed,
                        "rank": i,
                    }
                    checkpointer.maybe_save_checkpoint(
                        FakeTrainer(model_state=state, training_state=state),
                        epochs_completed,
                        batches_completed,
                    )

        for i, checkpointer in enumerate(checkpointers):
            checkpoint = checkpointer.load_checkpoint()
            assert checkpoint is not None
            model_state, training_state = checkpoint
            assert model_state["rank"] == i
            assert training_state["rank"] == i

            models, training = self.retrieve_and_delete_saved(shard=i)
            assert models == training == target
Simulates using the Checkpointer during distributed training with a sharded model.
tests/training/checkpointer_test.py