def METHOD_NAME(self, device, world_size, use_fsdp):
    self._check_test_params(device, world_size, use_fsdp)
    # Create generate callback
    prompts = ['a', 'bc', 'defg']
    prompt_batch_size = 2
    gen_interval = 2
    generate_cb = Generate(prompts, interval=f'{gen_interval}ba', batch_size=prompt_batch_size, max_new_tokens=5)
    # Create trainer
    train_batches = 6
    trainer = self._create_trainer(device, f'{train_batches}ba', use_fsdp, generate_cb)
    # Mock methods
    state = trainer.state
    model = state.model.module if state.is_model_ddp else state.model
    model.generate = Mock(wraps=model.generate)  # type: ignore
    generate_cb.generate = Mock(wraps=generate_cb.generate)
    trainer.logger.log_table = Mock()
    trainer.fit()
    expected_cb_call_count = math.ceil(train_batches / gen_interval)
    # Assert that the generate callback has been called the correct number of times.
    assert generate_cb.generate.call_count == expected_cb_call_count
    # Assert that model.generate has been called the correct number of times to ensure that prompt batching is correct.
    assert model.generate.call_count == math.ceil(  # type: ignore
        len(prompts) / prompt_batch_size) * expected_cb_call_count
    # Assert that log_table is called on the 0th rank only.
    if dist.get_global_rank() == 0:
        assert trainer.logger.log_table.call_count == expected_cb_call_count
    else:
        trainer.logger.log_table.assert_not_called()