This parameter
| layer_number: int, default = `None` |
| layer number of the current `DotProductAttention` when multiple such modules |
| are concatenated, for instance in consecutive transformer blocks. |
should indicate that the number must be 1-index based. An assertion could be added to enforce it.
Otherwise, if the user accidentally passes 0, then https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py#L206 will raise a division by 0 error after this multiplication: https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py#L198
This parameter
TransformerEngine/transformer_engine/pytorch/attention.py
Lines 719 to 721 in 574f1b4
should indicate that the number must be 1-index based. An assertion could be added to enforce it.
Otherwise, if the user accidentally passes
0, then https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py#L206 will raise a division by 0 error after this multiplication: https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py#L198