[Paddle] Add nn layer

[Wong4j]

Add the following nn layers (FP16/BF16):

• FusedScaleMaskSoftmax - softmax.py
• DotProductAttention, MultiHeadAttention - attention.py
• TransformerLayer - transformer.py

[Wong4j]

Could you please help me review the code? @timmoon10
cc @jeng1220@Tom-Zheng

[zlsh80826]

/te-ci

[Wong4j]

/te-ci

[mingxu1067]

/te-ci

[mingxu1067]

/te-ci

[mingxu1067]

/te-ci

[jeng1220]

LGTM.
@timmoon10 ,
All UT passed. Could you please review and merge this PR if everything looks good?

[timmoon10]

Overall LGTM.

[timmoon10]

The PyTorch attention module has logic to automatically choose the backend based on what configurations are supported by Flash Attention or cuDNN. Instead of making the transformer_engine backend specific to cuDNN, how about we add a fused_attention backend and make the transformer_engine backend choose between paddle and fused_attention? If Flash Attention is ever ported to Paddle, we could also add a flash_attention backend. In general, I think it would be a good idea to move away from the strict dichotomy between transformer_engine and paddle backends.

[jeng1220]

I am not sure what you are saying.

This is TE/PyTorch:

classDotProductAttention(...):
if (condition):
# calling cuDNN fused attentionelse:
# calling flash_attention from hazyresearch

This is what we made for TE/PaddlePaddle for now:

classDotProductAttention(...):
if (condition):
# transformer_engine backed, calling cuDNN fused attentionelse:
# paddle backend, calling PaddlePaddle Primitives 

And what do you mean?

As far as I know, Baidu refers to flash_attention of hazyresearch and is developing its own version, but I am not sure the current status. Once complete, it will be in PaddlePaddle repo instead of hazyresearch repo.

[timmoon10]

The PyTorch behavior is more like:

classDotProductAttention(...):
if (condition1):
# calling flash_attention from hazyresearchelif (condition2):
# calling cuDNN fused attentionelse:
# calling PyTorch primitives

This decision is made during the forward pass, not at construction time, and it can change depending on the sequence lengths. Also, it's possible for the attention to call PyTorch primitives while other Transformer layer components use TE kernels.

[jeng1220]

So, there are 2 things.

1. add 3rd option - calling flash_attention op (Baidu's version).
   I need to sync with Baidu to know their status. This will take some time. Also, Paddle primitives should be automatically converted to flash by the Paddle deep learning compiler if the compiler is ready, so we might not need the 3rd option. I think cuDNN is still our top priority.
2. move the logic of kernel selecting from constructor into forward.
   It is easy to do because PaddlePaddle supports dynamic graph mode. Although it has maximum flexibility, it will block CUDA Graph integration if kernel can be dynamically changed at runtime, won't it?

I think both 1. and 2. need more time to have better understanding. Considering schedule pressure, I hope this PR can be merged first, then 1. and 2. can be follow up in the future.

If I misunderstand anything, please correct me.
Thanks

[timmoon10]

Yea, I think we can defer until later. Most of the backend logic in the PyTorch attention module is related to Flash Attention:

TransformerEngine/transformer_engine/pytorch/attention.py

Lines 1086 to 1113 in e4f9e76

	use_flash_attention=self.use_flash_attention
	use_fused_attention=self.use_fused_attention
	if (query_layer.dtypenotin [torch.bfloat16, torch.float16]
	orkey_layer.dtypenotin [torch.bfloat16, torch.float16]
	orvalue_layer.dtypenotin [torch.bfloat16, torch.float16]
	):
	use_flash_attention=False
	ifkey_layer.shape[-1] >64:
	ifself.device_compute_capabilityin (8.6, 8.7):
	use_flash_attention=False
	elifnot_flash_attn_2_availableandself.device_compute_capability==8.9:
	use_flash_attention=False
	ifnot_flash_attn_2_availableandself.num_gqa_groups!=self.num_attention_heads:
	use_flash_attention=False
	ifself.attn_mask_type=="padding"andattention_maskisnotNone:
	use_flash_attention=False
	use_fused_attention=False
	ifcore_attention_bias_type!="no_bias"orcore_attention_biasisnotNone:
	use_flash_attention=False
	ifis_in_onnx_export_mode():
	use_flash_attention=False
	use_fused_attention=False

It seems that this logic changes rapidly, as the cuDNN team adds kernels or fixes bugs.

[ksivaman]

/te-ci