Repository: microsoft/DeepSpeed

Issue:
[BUG] GPT-J inference with kernel inject fails on master (7667988)
**Describe the bug**
Running DeepSpeed inference with kernel inject on the GPT-J model fails on latest master (766798849144fdb51e81204370d4e9411a34cafc) with the following exception:
```
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /home/ray/default/text-generation/inference-test.py:116 in <module>                              │
│                                                                                                  │
│   113 for i in range(iters):                                                                     │
│   114 │   torch.cuda.synchronize()                                                               │
│   115 │   start = time.time()                                                                    │
│ ❱ 116 │   outputs = pipe(inputs,                                                                 │
│   117 │   │   │   num_tokens=args.max_new_tokens,                                                │
│   118 │   │   │   do_sample=(not args.greedy))                                                   │
│   119 │   torch.cuda.synchronize()                                                               │
│                                                                                                  │
│ /home/ray/default/text-generation/utils.py:66 in __call__                                        │
│                                                                                                  │
│    63 │   │   else:                                                                              │
│    64 │   │   │   input_list = inputs                                                            │
│    65 │   │                                                                                      │
│ ❱  66 │   │   outputs = self.generate_outputs(input_list, num_tokens=num_tokens, do_sample=do_   │
│    67 │   │   return outputs                                                                     │
│    68 │                                                                                          │
│    69 │   def _generate_json(self, checkpoint_path=None):                                        │
│                                                                                                  │
│ /home/ray/default/text-generation/utils.py:115 in generate_outputs                               │
│                                                                                                  │
│   112 │   │   │   #       'token_type_ids' kwarg not recognized in Llamma generate function      │
│   113 │   │   │   outputs = self.model.generate(input_tokens.input_ids, **generate_kwargs)       │
│   114 │   │   else:                                                                              │
│ ❱ 115 │   │   │   outputs = self.model.generate(**input_tokens, **generate_kwargs)               │
│   116 │   │   outputs = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)           │
│   117 │   │                                                                                      │
│   118 │   │   return outputs                                                                     │
│                                                                                                  │
│ /home/ray/anaconda3/lib/python3.10/site-packages/deepspeed/inference/engine.py:635 in _generate  │
│                                                                                                  │
│   632 │   │   │   raise NotImplementedError("DeepSpeed does not support `num_beams` > 1, if th   │
│   633 │   │   │   │   │   │   │   │   │     "add your request to: https://github.com/microsoft   │
│   634 │   │                                                                                      │
│ ❱ 635 │   │   return self.module.generate(*inputs, **kwargs)                                     │
│   636                                                                                            │
│                                                                                                  │
│ /home/ray/anaconda3/lib/python3.10/site-packages/torch/utils/_contextlib.py:115 in               │
│ decorate_context                                                                                 │
│                                                                                                  │
│   112 │   @functools.wraps(func)                                                                 │
│   113 │   def decorate_context(*args, **kwargs):                                                 │
│   114 │   │   with ctx_factory():                                                                │
│ ❱ 115 │   │   │   return func(*args, **kwargs)                                                   │
│   116 │                                                                                          │
│   117 │   return decorate_context                                                                │
│   118                                                                                            │
│                                                                                                  │
│ /home/ray/anaconda3/lib/python3.10/site-packages/transformers/generation/utils.py:1565 in        │
│ generate                                                                                         │
│                                                                                                  │
│   1562 │   │   │   )                                                                             │
│   1563 │   │   │                                                                                 │
│   1564 │   │   │   # 13. run sample                                                              │
│ ❱ 1565 │   │   │   return self.sample(                                                           │
│   1566 │   │   │   │   input_ids,                                                                │
│   1567 │   │   │   │   logits_processor=logits_processor,                                        │
│   1568 │   │   │   │   logits_warper=logits_warper,                                              │
│                                                                                                  │
│ /home/ray/anaconda3/lib/python3.10/site-packages/transformers/generation/utils.py:2612 in sample │
│                                                                                                  │
│   2609 │   │   │   model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)  │
│   2610 │   │   │                                                                                 │
│   2611 │   │   │   # forward pass to get next token                                              │
│ ❱ 2612 │   │   │   outputs = self(                                                               │
│   2613 │   │   │   │   **model_inputs,                                                           │
│   2614 │   │   │   │   return_dict=True,                                                         │
│   2615 │   │   │   │   output_attentions=output_attentions,                                      │
│                                                                                                  │
│ /home/ray/anaconda3/lib/python3.10/site-packages/torch/nn/modules/module.py:1501 in _call_impl   │
│                                                                                                  │
│   1498 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1499 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1500 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1501 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1502 │   │   # Do not call functions when jit is used                                          │
│   1503 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1504 │   │   backward_pre_hooks = []                                                           │
│                                                                                                  │
│ /home/ray/anaconda3/lib/python3.10/site-packages/transformers/models/gptj/modeling_gptj.py:853   │
│ in forward                                                                                       │
│                                                                                                  │
│    850 │   │   """                                                                               │
│    851 │   │   return_dict = return_dict if return_dict is not None else self.config.use_return  │
│    852 │   │                                                                                     │
│ ❱  853 │   │   transformer_outputs = self.transformer(                                           │
│    854 │   │   │   input_ids,                                                                    │
│    855 │   │   │   past_key_values=past_key_values,                                              │
│    856 │   │   │   attention_mask=attention_mask,                                                │
│                                                                                                  │
│ /home/ray/anaconda3/lib/python3.10/site-packages/torch/nn/modules/module.py:1501 in _call_impl   │
│                                                                                                  │
│   1498 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1499 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1500 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1501 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1502 │   │   # Do not call functions when jit is used                                          │
│   1503 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1504 │   │   backward_pre_hooks = []                                                           │
│                                                                                                  │
│ /home/ray/anaconda3/lib/python3.10/site-packages/transformers/models/gptj/modeling_gptj.py:688   │
│ in forward                                                                                       │
│                                                                                                  │
│    685 │   │   │   │   │   head_mask[i],                                                         │
│    686 │   │   │   │   )                                                                         │
│    687 │   │   │   else:                                                                         │
│ ❱  688 │   │   │   │   outputs = block(                                                          │
│    689 │   │   │   │   │   hidden_states=hidden_states,                                          │
│    690 │   │   │   │   │   layer_past=layer_past,                                                │
│    691 │   │   │   │   │   attention_mask=attention_mask,                                        │
│                                                                                                  │
│ /home/ray/anaconda3/lib/python3.10/site-packages/torch/nn/modules/module.py:1501 in _call_impl   │
│                                                                                                  │
│   1498 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1499 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1500 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1501 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1502 │   │   # Do not call functions when jit is used                                          │
│   1503 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1504 │   │   backward_pre_hooks = []                                                           │
│                                                                                                  │
│ /home/ray/anaconda3/lib/python3.10/site-packages/deepspeed/model_implementations/transformers/ds │
│ _transformer.py:171 in forward                                                                   │
│                                                                                                  │
│   168 │   │   │                                                                                  │
│   169 │   │   │   presents = (key, value)                                                        │
│   170 │   │   │   self.layer_past = presents if layer_past is None else None                     │
│ ❱ 171 │   │   │   output = self.mlp(attention_output, input, inp_norm, self.attention.attn_ob)   │
│   172 │   │   │                                                                                  │
│   173 │   │   │   if not self.config.pre_layer_norm:                                             │
│   174 │   │   │   │   output = inference_module.layer_norm(output, self.norm_w, self.norm_b, s   │
│                                                                                                  │
│ /home/ray/anaconda3/lib/python3.10/site-packages/torch/nn/modules/module.py:1501 in _call_impl   │
│                                                                                                  │
│   1498 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1499 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1500 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1501 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1502 │   │   # Do not call functions when jit is used                                          │
│   1503 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1504 │   │   backward_pre_hooks = []                                                           │
│                                                                                                  │
│ /home/ray/anaconda3/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/ds_mlp.py:1 │
│ 06 in forward                                                                                    │
│                                                                                                  │
│   103 │   │   │   │   │   │   │   │   │   │     bias=self.inter_b,                               │
│   104 │   │   │   │   │   │   │   │   │   │     weight_out=self.output_w)                        │
│   105 │   │   else:                                                                              │
│ ❱ 106 │   │   │   output, residual_add = self.mlp_gemm_func(input=input,                         │
│   107 │   │   │   │   │   │   │   │   │   │   │   │   │     residual=residual,                   │
│   108 │   │   │   │   │   │   │   │   │   │   │   │   │     weight_interm=self.inter_w,          │
│   109 │   │   │   │   │   │   │   │   │   │   │   │   │     weight_out=self.output_w,            │
│                                                                                                  │
│ /home/ray/anaconda3/lib/python3.10/site-packages/torch/nn/modules/module.py:1501 in _call_impl   │
│                                                                                                  │
│   1498 │   │   if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks   │
│   1499 │   │   │   │   or _global_backward_pre_hooks or _global_backward_hooks                   │
│   1500 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1501 │   │   │   return forward_call(*args, **kwargs)                                          │
│   1502 │   │   # Do not call functions when jit is used                                          │
│   1503 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1504 │   │   backward_pre_hooks = []                                                           │
│                                                                                                  │
│ /home/ray/anaconda3/lib/python3.10/site-packages/deepspeed/ops/transformer/inference/op_binding/ │
│ mlp_gemm.py:68 in forward                                                                        │
│                                                                                                  │
│   65 │   │   │   │   gamma: Optional[torch.Tensor] = None,                                       │
│   66 │   │   │   │   beta: Optional[torch.Tensor] = None):                                       │
│   67 │   │   if self.config.norm_type == NormType.LayerNorm:                                     │
│ ❱ 68 │   │   │   output, residual_add = self.mlp_gemm_func(                                      │
│   69 │   │   │   │   input,                                                                      │
│   70 │   │   │   │   residual,                                                                   │
│   71 │   │   │   │   input_bias,                                                                 │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
TypeError: mlp_gemm_fp16(): incompatible function arguments. The following argument types are supported:
    1. (arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: torch.Tensor, arg4: torch.Tensor, arg5: torch.Tensor, arg6: torch.Tensor, arg7: torch.Tensor, arg8: float, arg9: bool, arg10: bool, arg11: 
torch.Tensor, arg12: torch.Tensor, arg13: bool, arg14: int, arg15: bool) -> List[torch.Tensor]

Invoked with: tensor([[[ 0.1053, -0.0731,  0.0074,  ...,  0.0917,  0.0924, -0.1111],
         [ 0.0710, -0.0090,  0.0092,  ...,  0.0089,  0.0417, -0.1654],
         [-0.0800,  0.1138, -0.1281,  ...,  0.0640,  0.1384, -0.0044],
         ...,
         [-0.0196,  0.0964, -0.0567,  ..., -0.0553,  0.0475,  0.0018],
         [ 0.0242,  0.0211, -0.0890,  ..., -0.3027, -0.0173, -0.2292],
         [ 0.0693,  0.0262, -0.0925,  ..., -0.0804,  0.1093, -0.0806]]],
       device='cuda:0', dtype=torch.float16), tensor([[[ 0.0086,  0.0225, -0.0049,  ..., -0.0125,  0.0019, -0.0077],
         [-0.0086, -0.0046, -0.0257,  ..., -0.0141, -0.0074, -0.0209],
         [ 0.0147,  0.0043,  0.0167,  ..., -0.0083,  0.0008,  0.0096],
         ...,
         [-0.0098,  0.0013,  0.0294,  ..., -0.0185,  0.0043, -0.0280],
         [ 0.0038,  0.0182,  0.0184,  ..., -0.0090, -0.0004, -0.0129],
         [ 0.0175,  0.0097,  0.0035,  ...,  0.0085, -0.0027, -0.0117]]],
       device='cuda:0', dtype=torch.float16), None, Parameter containing:
tensor([[-1.6632e-02,  2.8137e-02,  2.7332e-03,  ..., -2.3514e-02,
          7.4959e-03,  3.9482e-03],
        [-2.0233e-02, -1.3191e-02,  1.2337e-02,  ..., -5.1041e-03,
          4.2686e-03, -4.1428e-03],
        [-1.3817e-02, -7.0870e-05,  1.4000e-03,  ...,  5.4245e-03,
         -1.0519e-03,  1.2566e-02],
        ...,
        [ 4.2076e-03,  2.7115e-02,  7.9651e-03,  ..., -1.0971e-02,
         -1.4496e-02,  1.4030e-02],
        [-1.7128e-03,  6.7635e-03,  9.7351e-03,  ..., -2.5024e-02,
          9.7733e-03,  5.8174e-03],
        [-7.0534e-03,  3.6438e-02, -1.5087e-03,  ...,  7.2908e-04,
          1.9363e-02, -5.4207e-03]], device='cuda:0', dtype=torch.float16), Parameter containing:
tensor([[-0.0027,  0.0021, -0.0100,  ...,  0.0085,  0.0126,  0.0011],
        [ 0.0053,  0.0139,  0.0181,  ...,  0.0169,  0.0010,  0.0126],
        [-0.0036, -0.0032, -0.0045,  ...,  0.0064,  0.0073,  0.0039],
        ...,
        [-0.0052, -0.0079,  0.0102,  ...,  0.0055,  0.0094, -0.0105],
        [-0.0092, -0.0154,  0.0002,  ..., -0.0150, -0.0053, -0.0181],
        [-0.0016, -0.0120,  0.0089,  ...,  0.0137, -0.0035, -0.0088]],
       device='cuda:0', dtype=torch.float16), Parameter containing:
tensor([-0.0587, -0.0428, -0.0135,  ..., -0.0632, -0.0376, -0.0448],
       device='cuda:0', dtype=torch.float16), Parameter containing:
tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.float16), Parameter containing:
tensor([0., 0., 0.,  ..., 0., 0., 0.], device='cuda:0', dtype=torch.float16), 1e-05, True, False, tensor([1.4242e-36]), tensor([2.8699e-42]), False, <ActivationFuncType.GELU: 1>, False
```
Notably, this works as expected with deepspeed==0.9.2.

**To Reproduce**
I've used https://github.com/microsoft/DeepSpeedExamples/tree/master/inference/huggingface/text-generation

`deepspeed --num_gpus 2 inference-test.py --name EleutherAI/gpt-j-6B --batch_size 1 --max_new_tokens 512 --ds_inference --use_kernel`

**Expected behavior**
No exception is raised and model generates succesfully.

**ds_report output**
```
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
async_io ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.0
 [WARNING]  using untested triton version (2.0.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/ray/anaconda3/lib/python3.10/site-packages/torch']
torch version .................... 2.0.1+cu118
deepspeed install path ........... ['/home/ray/anaconda3/lib/python3.10/site-packages/deepspeed']
deepspeed info ................... 0.9.3+76679884, 76679884, master
torch cuda version ............... 11.8
torch hip version ................ None
nvcc version ..................... 11.8
deepspeed wheel compiled w. ...... torch 2.0, cuda 11.8
```

**System info (please complete the following information):**
 - OS: Ubuntu 20.04
 - GPU count and types: g4.12xlarge (4xA10g)
 - DeepSpeed: 0.9.3+76679884
 - Transformers: 4.29.2
 - Python version: Python 3.10.8


Discussion:
Thank you for reporting this bug @Yard1. Unfortunately our unit tests failed to catch this due to a name change of the model and a bug in our tests. Fixed the unit tests bug in #3618 but we're still actively working to resolve the model issues. Thanks!

Resolving patch:
diff --git a/deepspeed/module_inject/containers/base.py b/deepspeed/module_inject/containers/base.py
--- a/deepspeed/module_inject/containers/base.py
+++ b/deepspeed/module_inject/containers/base.py
@@ -249,17 +249,21 @@ def mlp_output_mp(self, mp_replace, reversed_dim=False):
                                                    allocate_tensor=reversed_dim)
 
     def copy_data_to_new_module(self):
-        params = {
-            self.module.mlp.attn_nw: self.attn_nw,
-            self.module.mlp.attn_nb: self.attn_nb,
-            self.module.norm_w: self.input_nw,
-            self.module.norm_b: self.input_nb
-        }
-        for dst, src in params.items():
-            if src is None:
-                dst = src
+        params = {'attn_nw': self.attn_nw, 'attn_nb': self.attn_nb}
+        for key in params:
+            if params[key] is None:
+                setattr(self.module.mlp, key, None)
             else:
-                dst.data.copy_(src.to(get_accelerator().current_device_name()))
+                setattr(self.module.mlp, key,
+                        torch.nn.parameter.Parameter(params[key].to(get_accelerator().current_device_name())))
+
+        params = {'norm_w': self.input_nw, 'norm_b': self.input_nb}
+        for key in params:
+            if params[key] is None:
+                setattr(self.module, key, None)
+            else:
+                setattr(self.module, key,
+                        torch.nn.parameter.Parameter(params[key].to(get_accelerator().current_device_name())))
 
     def transpose(self):
         self.transpose_attention()


Test patch:
diff --git a/tests/unit/inference/test_inference.py b/tests/unit/inference/test_inference.py
--- a/tests/unit/inference/test_inference.py
+++ b/tests/unit/inference/test_inference.py
@@ -49,7 +49,7 @@
     "gpt2",
     "distilgpt2",
     "Norod78/hebrew-bad_wiki-gpt_neo-tiny",
-    #"EleutherAI/gpt-j-6B", # Removed as this is causing OOM errors randomly
+    "EleutherAI/gpt-j-6B",  # bring back this model as we did not catch an error before by merging some changes! TODO: we need to fix the OOM issue later!
     "bigscience/bloom-560m",
 ]
 _opt_models = [