model = FSDP(
      model,
      device_id=torch.cuda.current_device(),
  )

In this case FSDP will simply wrap the whole model in a single FSDP unit.
Running on an NVIDIA A100-SXM4–40GB with 8 GPUs, we are able to reach 2.3 TFlops and 95% GPU memory utilization with a batch size of 14.
However, since T5 is a transformer model, we are better served to leverage the transformer wrapper for this model. 
To use that, we need to isolate the layer class for the transformer, and then pass it in to create our transformer wrapper. 
from transformers.models.t5.modeling_t5 import T5Block

And now we can create our Transformer wrapper: 
transformer_auto_wrapper_policy = functools.partial(
        transformer_auto_wrap_policy,
        transformer_layer_cls={
            T5Block,  # < ---- Your Transformer layer class
        },
    )
