def test_mpt_generate_multi_gpu(
    attn_impl: str,
    use_alibi: bool,
    tie_word_embeddings: bool,
    build_tiny_mpt: Callable[..., ComposerMPTCausalLM],
    mpt_tokenizer: PreTrainedTokenizerBase,
):
    """Tests mpt generation with mutiple gpus.

    and generations of different lengths.
    """
    device = get_device('gpu')

    model = build_tiny_mpt(
        tie_word_embeddings=tie_word_embeddings,
        attn_config={
            'attn_impl': attn_impl,
            'attn_uses_sequence_id': False,
            'alibi': use_alibi,
        },
    )
    model = device.module_to_device(model)

    model.eval()

    model.model = FSDP(model.model)

    with get_precision_context('amp_bf16'):
        _ = model.generate(
            device.tensor_to_device(
                mpt_tokenizer('hello',
                              return_tensors='pt')['input_ids'],  # type: ignore
            ),
            max_new_tokens=3,
            eos_token_id=EOS_TOKEN_ID,
            use_cache=True,
            synced_gpus=True,
        )
Tests mpt generation with mutiple gpus.

    and generations of different lengths.
    
tests/models/test_mpt_gen.py