Add Deepspeed Transformer kernel, but encounter error  IndexError: tuple index out of range
Author: dancingpipi
Pull request: no

I have compare the Bert-base pretraining performance between [Nvidia Tensorflow ](https://github.com/NVIDIA/DeepLearningExamples/tree/master/TensorFlow/LanguageModeling/BERT) and huggingface + Deepspeed.
My environment:
cuda: 11.2
torch: 1.10
deepspeed: 0.5.8
transformers: 4.13.0.dev0
machine: single node with 8 A100

The result show that huggingface are 20% slower than TF while XLA is enable. 
I know that huggingface does not support deepspeed's transformer kernel, so I manually replaced it, the code is as follows:
```
import time
import code
import logging
import sys
import os
from typing import Optional
import glob

import datasets
from dataclasses import dataclass, field
import transformers
from transformers.integrations import TensorBoardCallback
from datasets import concatenate_datasets

from deepspeed.ops.transformer import (
    DeepSpeedTransformerConfig,
    DeepSpeedTransformerLayer
)

logger = logging.getLogger(__name__)

datasets.set_caching_enabled(False)

@dataclass
class CustomArguments:
  model_name: Optional[str] = field(default="bert-base-uncased")
  model_path: Optional[str] = field(default=None)
  train_data_dir: Optional[str] = field(default=None)
  validation_file: Optional[str] = field(default=None)
  max_seq_length: Optional[int] = field(default=128)
  preprocessing_num_worker: Optional[int] = field(default=os.cpu_count())
  use_pretrain_collator: Optional[bool] = field(default=True)

  line_by_line: Optional[bool] = field(default=True)
  vocab_path: Optional[str] = field(default=None)
  model_conf_path: Optional[str] = field(default=None)

def gen_ds_bert_config(training_args, config):
    bert_config = DeepSpeedTransformerConfig(
        batch_size=384,
        hidden_size=config.hidden_size,
        intermediate_size=config.intermediate_size,
        heads=config.num_attention_heads,
        attn_dropout_ratio=config.attention_probs_dropout_prob,
        hidden_dropout_ratio=config.hidden_dropout_prob,
        num_hidden_layers=config.num_hidden_layers,
        initializer_range=0.02,
        layer_norm_eps=1e-8,
        local_rank=training_args.local_rank,
        fp16=training_args.fp16,
        pre_layer_norm=False,
        training=True
    )
   return bert_config

def inject_ds_enc_layer(model, training_args, config):
    for i in range(config.num_hidden_layers):
        bert_config = gen_ds_bert_config(training_args, config)
        model.bert.encoder.layer[i] = DeepSpeedTransformerLayer(bert_config)


def main():
  parser = transformers.HfArgumentParser((CustomArguments, transformers.TrainingArguments))
  custom_args, training_args = parser.parse_args_into_dataclasses()

  logging.basicConfig(
      format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
      datefmt="%m/%d/%Y %H:%M:%S",
      handlers=[logging.StreamHandler(sys.stdout)],
  )

  log_level = training_args.get_process_log_level()
  logger.setLevel(log_level)
  datasets.utils.logging.set_verbosity(log_level)
  transformers.utils.logging.set_verbosity(log_level)
  transformers.utils.logging.enable_default_handler()
  transformers.utils.logging.enable_explicit_format()

  logger.warning(
      f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
      + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
  )
  # Set the verbosity to info of the Transformers logger (on main process only):
  logger.info(f"Training/evaluation parameters {training_args}")

  tokenizer = transformers.BertTokenizer(
                      vocab_file=custom_args.vocab_path,
                      do_lower_case=False,
                      max_length=custom_args.max_seq_length)

  model_config = transformers.BertConfig.from_pretrained(
                                          custom_args.model_conf_path)
  if custom_args.model_path is not None:
    print(f"training continue from: f{custom_args.model_path}!")
    model = transformers.BertForPreTraining.from_pretrained(custom_args.model_path)
  elif custom_args.model_conf_path:
    print(f"training from scratch!")
    model = transformers.BertForPreTraining(config=model_config)
  else:
    raise ValueError("no model config or model path provied, please check!")

  model.resize_token_embeddings(len(tokenizer))
  inject_ds_enc_layer(model, training_args, model_config)

  if custom_args.train_data_dir is None:
    raise ValueError("train_data_dir must be specified!")

  print("start loading data!")
  start = time.time()
  data_files = glob.glob(custom_args.train_data_dir.rstrip("/") + "/*")
  data = [datasets.load_from_disk(data_file) for data_file in data_files]
  train_dataset = concatenate_datasets(data)
  train_dataset = train_dataset.shuffle()
  end = time.time()
  print(f"loading data cost: {end - start} s")


  data_collator = None
  if custom_args.use_pretrain_collator:
    data_collator = transformers.DataCollatorForLanguageModeling(
      tokenizer=tokenizer,
      mlm_probability=0.15,
      mlm=True,
      # pad_to_multiple_of=8
    )

  code.interact(local=locals())
  trainer = transformers.Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=None,
    tokenizer=tokenizer,
    data_collator=data_collator,
    callbacks=[TensorBoardCallback]
  )

  trainer.train()

if __name__ == "__main__":
  main()
```

But the error occur: 
```
  File "/usr/local/anaconda3/lib/python3.7/site-packages/deepspeed/ops/transformer/transformer.py", line 607, in forward
    self.config)
  File "/usr/local/anaconda3/lib/python3.7/site-packages/deepspeed/ops/transformer/transformer.py", line 179, in forward
    if inp_size[1] % 16 != 0:
IndexError: tuple index out of range
...
RuntimeError: CUDA error: an illegal memory access was encountered
```
It seem the input dim is not correctly.
Did I do something wrong, hope to help me point it out
 

Comments:
By the way, does Huggingface support similar XLA or kernel fusion capabilities?

Ok, problem solved, it's because the origin huggingface bert will return a tuple and choose the first element while deepspeed return the element

you can enable return_tuple in DeepSpeedTransformerConfig