↔ [Converter] Add operations to accelerate Transformer encoder

[narendasan]

torch.nn.TransformerEncoder

• Function Schema:

  • aten::chunk(Tensor(a -> *) self, int chunks, int dim=0) -> Tensor(a)[]
  • aten::_native_multi_head_attention(Tensor query, Tensor key, Tensor value, int embed_dim, int num_head, Tensor qkv_weight, Tensor qkv_bias, Tensor proj_weight, Tensor proj_bias, Tensor? mask=None, bool need_weights=True, bool average_attn_weights=True, int? mask_type=None) -> (Tensor, Tensor)
  • aten::add.t(t[] a, t[] b) -> t[]
  • aten::all.bool(bool[] self) -> bool
  • prim::requires_grad(Tensor a) -> bool
  • aten::any.bool(bool[] self) -> bool
  • aten::_transformer_encoder_layer_fwd(Tensor src, int embed_dim, int num_heads, Tensor qkv_weight, Tensor qkv_bias, Tensor proj_weight, Tensor proj_bias, bool use_gelu, bool norm_first, float eps, Tensor 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=None, int? mask_type=None) -> Tensor
  • aten::len.str(str s) -> int
  • aten::str(t elem) -> str
  • prim::device(Tensor a) -> Device
  • aten::is_grad_enabled() -> bool
  • aten::find(str self, str substr, int start=0, int end=-1) -> int
  • prim::is_cuda(Tensor a) -> bool

• Original PyTorch API:

  • https://pytorch.org/docs/stable/generated/torch.nn.TransformerEncoder.html

• Relevant TensorRT Documentation:

Alternatives

Additional context

[Christina-Young-NVIDIA]

This issue needs to be subdivided into multiple tasks (too large for single sprint).

[apbose]

FX implementation update:
Passes with PR 1751 and PR 1708, model-

class Sample(torch.nn.Module):
    def __init__(self):
        super(Sample, self).__init__()
        self.encoder_layer = torch.nn.TransformerEncoderLayer(d_model=512, nhead=8)
        self.transformer_encoder = torch.nn.TransformerEncoder(self.encoder_layer, num_layers=6)

    def forward(self, x):
        out = self.transformer_encoder(x)
        return out

Some observations-

• The aot_module_simplified at present introduces operations which are then split by the splitter in the lowering pass due to which the model with just the single layer works. However the no of splits is too huge due to which the compilation is slow. Also it prints out the aten currently supported and the unsupported operations
  Prints-

Supported node types in the model:
torch.ops.aten.add.Tensor: ((torch.float32, torch.float32), {})
torch.ops.aten.add.Tensor: ((torch.float32,), {})
torch.ops.aten.sub.Tensor: ((torch.float32, torch.float32), {})
torch.ops.aten.mul.Tensor: ((torch.float32, torch.float32), {})
torch.ops.aten.relu.default: ((torch.float32,), {})

Unsupported node types in the model:
torch.ops.aten.permute.default: ((torch.float32,), {})
torch.ops.aten.reshape: ((torch.float32,), {})
torch.ops.aten.addmm.default: ((torch.float32, torch.float32, torch.float32), {})
torch.ops.aten.unsqueeze.default: ((torch.float32,), {})
torch.ops.aten.squeeze.dim: ((torch.float32,), {})
torch.ops.aten.select.int: ((torch.float32,), {})
torch.ops.aten._scaled_dot_product_efficient_attention.default: ((torch.float32, torch.float32, torch.float32), {})
torch_tensorrt.fx.passes.lower_basic_pass_aten.aten_operator_getitem: ((), {})
torch.ops.aten.var_mean.correction: ((torch.float32,), {})
torch.ops.aten.rsqrt.default: ((torch.float32,), {})

• Just using backend= fx2trt shows the same error as before
  Unsupported node types in the model:
  torch.nn.modules.transformer.TransformerEncoder: ((torch.float32,), {}) with just one split from the splitter. The aot_module_simplified decomposes it into operations, due to which it goes till the point where it reports unsupported ops

[apbose]

Added issues #1790, #1724, #1791, #1792, #1793, #1794

[apbose]

#1790 - No prior implementation in acc, will take more time
#1724 - Trivial operation since implementation in acc exists
#1791- Trivial operation since implementation in acc exists
#1792- Implementation exists, but tests might take time
#1793- No prior implementation in acc, will take more time
#1794 - Trivial operation since implementation in acc exists