我在lib/python3.11/site-packages/torch/_C/_VariableFunctions.pyi
文件the function is called here中遇到了PyTorch的function _transformer_encoder_layer_fwd
。
但是我没有找到关于这个函数的任何细节。为什么可以调用这个函数?如何调用?
def _transformer_encoder_layer_fwd(src: Tensor, embed_dim:
_int, num_heads: _int,
qkv_weight: Tensor,
qkv_bias: Tensor,
proj_weight: Tensor,
proj_bias: Tensor,
use_gelu: _bool,
norm_first: _bool,
eps: _float,
norm_weight_1: Tensor,
norm_bias_1: Tensor,
norm_weight_2: Tensor,
norm_bias_2: Tensor,
ffn_weight_1: Tensor,
ffn_bias_1: Tensor,
ffn_weight_2: Tensor,
ffn_bias_2: Tensor,
mask: Optional[Tensor] = None,
mask_type: Optional[_int] = None) -> Tensor:
字符串
...
该函数在torch/_C/_VariableFunctions.pyi
中定义
我试图找到这个函数的任何细节,关于如何调用这个函数。但没有结果。
1条答案
按热度按时间hwamh0ep1#
如“快速路径”部分中所述的here,nn.TransformerEncoderLayer的forward()方法可以使用Flash Attention,这是一种使用融合操作的优化自注意实现。但是,如PyTorch文档中所述,要使用闪光注意,必须满足一系列标准。
从PyTorch的GitHub上的Transformer编码器上的实现来看,这个方法调用很可能是应用Flash Attention的地方。