Why is the first dimension sequence_length as opposed to batch_size in the input tensor

[vgoklani]

Why does the library use sequence_length as the first dimension of the input-tensor as opposed to the batch_size?

Is this just a choice of convention from RNNs or is the difference performance related?

From the example code:

bmm1 = torch.bmm(query.transpose(0, 1), key.transpose(0, 1).transpose(1, 2)) / self.norm_factor

https://github.com/NVIDIA/TransformerEngine/blob/main/docs/examples/quickstart_utils.py#L93

I see two successive transpose(0, 1) operations?

Thanks!

[ksivaman]

Having either batch size or sequence length as the first dimension of the input tensor can be beneficial for performance depending on the attention implementation and some other details. Our choice to stick with sequence length as the first dimension was so that we can add support for sequence parallelism. That being said, we will soon add a batch_first argument in our APIs that would optionally support input tensors with batch_size as their first dimension.

[vgoklani]

that makes sense, thank you!

[bryangopal]

Having either batch size or sequence length as the first dimension of the input tensor can be beneficial for performance depending on the attention implementation and some other details. Our choice to stick with sequence length as the first dimension was so that we can add support for sequence parallelism. That being said, we will soon add a batch_first argument in our APIs that would optionally support input tensors with batch_size as their first dimension.

any updates on the batch_first addition?

[timmoon10]

The attention module has logic to handle multiple formats (e.g. SBHD, BSHD). See:

TransformerEngine/transformer_engine/pytorch/attention.py

Line 1821 in 666539f

qkv_format: str = "sbhd",

However, we haven't exposed this in the Transformer layer yet. Pinging @cyanguwa.

[bryangopal]

@cyanguwa just following up, would appreciate an update!