Repository: microsoft/DeepSpeed

Issue:
[BUG]TypeError: allocate_workspace_fp16(): incompatible function arguments.
**Describe the bug**
When I run DiffusionPipeline, 
`Time to load transformer_inference op: 23.22636890411377 seconds
[2023-04-13 14:24:52,241] [INFO] [logging.py:96:log_dist] [Rank -1] DeepSpeed-Attention config: {'layer_id': 0, 'hidden_size': 320, 'intermediate_size': 1280, 'heads': 8, 'num_hidden_layers': -1, 'fp16': True, 'pre_layer_norm': True, 'local_rank': -1, 'stochastic_mode': False, 'epsilon': 1e-12, 'mp_size': 1, 'q_int8': False, 'scale_attention': True, 'triangular_masking': False, 'local_attention': False, 'window_size': 256, 'rotary_dim': -1, 'rotate_half': False, 'rotate_every_two': True, 'return_tuple': True, 'mlp_after_attn': True, 'mlp_act_func_type': <ActivationFuncType.GELU: 1>, 'specialized_mode': False, 'training_mp_size': 1, 'bigscience_bloom': False, 'max_out_tokens': 4096, 'min_out_tokens': 1, 'scale_attn_by_inverse_layer_idx': False, 'enable_qkv_quantization': False, 'use_mup': False, 'return_single_tuple': False, 'set_empty_params': False, 'transposed_mode': False}
Installed CUDA version 11.3 does not match the version torch was compiled with 11.7 but since the APIs are compatible, accepting this combination
Using /home/yinian.lw/.cache/torch_extensions/py39_cu117 as PyTorch extensions root...
No modifications detected for re-loaded extension module transformer_inference, skipping build step...
Loading extension module transformer_inference...
Time to load transformer_inference op: 0.0314486026763916 seconds
Installed CUDA version 11.3 does not match the version torch was compiled with 11.7 but since the APIs are compatible, accepting this combination
Using /home/yinian.lw/.cache/torch_extensions/py39_cu117 as PyTorch extensions root...
Creating extension directory /home/yinian.lw/.cache/torch_extensions/py39_cu117/spatial_inference...
Detected CUDA files, patching ldflags
Emitting ninja build file /home/yinian.lw/.cache/torch_extensions/py39_cu117/spatial_inference/build.ninja...
Building extension module spatial_inference...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
[1/3] /home/yinian.lw/cuda-11.3/bin/nvcc  -DTORCH_EXTENSION_NAME=spatial_inference -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -I/home/yinian.lw/miniconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/spatial/includes -I/home/yinian.lw/miniconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/includes -isystem /home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/include -isystem /home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/include/torch/csrc/api/include -isystem /home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/include/TH -isystem /home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/include/THC -isystem /home/yinian.lw/cuda-11.3/include -isystem /home/yinian.lw/miniconda3/include/python3.9 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_86,code=compute_86 -gencode=arch=compute_86,code=sm_86 --compiler-options '-fPIC' -O3 --use_fast_math -std=c++14 -U__CUDA_NO_HALF_OPERATORS__ -U__CUDA_NO_HALF_CONVERSIONS__ -U__CUDA_NO_HALF2_OPERATORS__ -gencode=arch=compute_86,code=sm_86 -gencode=arch=compute_86,code=compute_86 -c /home/yinian.lw/miniconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/spatial/csrc/opt_bias_add.cu -o opt_bias_add.cuda.o 
[2/3] c++ -MMD -MF pt_binding.o.d -DTORCH_EXTENSION_NAME=spatial_inference -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -I/home/yinian.lw/miniconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/spatial/includes -I/home/yinian.lw/miniconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/includes -isystem /home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/include -isystem /home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/include/torch/csrc/api/include -isystem /home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/include/TH -isystem /home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/include/THC -isystem /home/yinian.lw/cuda-11.3/include -isystem /home/yinian.lw/miniconda3/include/python3.9 -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++14 -O3 -std=c++14 -g -Wno-reorder -c /home/yinian.lw/miniconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/spatial/csrc/pt_binding.cpp -o pt_binding.o 
[3/3] c++ opt_bias_add.cuda.o pt_binding.o -shared -L/home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/lib -lc10 -lc10_cuda -ltorch_cpu -ltorch_cuda_cu -ltorch_cuda_cpp -ltorch -ltorch_python -L/home/yinian.lw/cuda-11.3/lib64 -lcudart -o spatial_inference.so
Loading extension module spatial_inference...
Time to load spatial_inference op: 18.57782793045044 seconds
**** found and replaced unet w. <class 'deepspeed.model_implementations.diffusers.unet.DSUNet'>
DeepSpeed Inference Engine initialized
Traceback (most recent call last):
  File "/home/yinian.lw/AIGC/script/Hongbao520Img2Img/text2img.py", line 45, in <module>
    ds_results = measure_latency(pipe,prompt)
  File "/home/yinian.lw/AIGC/script/Hongbao520Img2Img/text2img.py", line 15, in measure_latency
    _ =  pipe(prompt)
  File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py", line 667, in __call__
    noise_pred = self.unet(
  File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/deepspeed/model_implementations/diffusers/unet.py", line 40, in forward
    self._create_cuda_graph(*inputs, **kwargs)
  File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/deepspeed/model_implementations/diffusers/unet.py", line 52, in _create_cuda_graph
    ret = self._forward(*inputs, **kwargs)
  File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/deepspeed/model_implementations/diffusers/unet.py", line 73, in _forward
    return self.unet(sample, timestamp, encoder_hidden_states, return_dict)
  File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/diffusers/models/unet_2d_condition.py", line 582, in forward
    sample, res_samples = downsample_block(
  File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/diffusers/models/unet_2d_blocks.py", line 837, in forward
    hidden_states = attn(
  File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/diffusers/models/transformer_2d.py", line 265, in forward
    hidden_states = block(
  File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/deepspeed/ops/transformer/inference/diffusers_transformer_block.py", line 91, in forward
    out_attn_1 = self.attn_1(out_norm_1)
  File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/yinian.lw/miniconda3/lib/python3.9/site-packages/deepspeed/ops/transformer/inference/diffusers_attention.py", line 184, in forward
    self.allocate_workspace(self.config.hidden_size, self.config.heads,
TypeError: allocate_workspace_fp16(): incompatible function arguments. The following argument types are supported:
    1. (arg0: int, arg1: int, arg2: int, arg3: int, arg4: int, arg5: int, arg6: bool, arg7: int, arg8: int, arg9: int) -> None

Invoked with: 320, 8, 4096, 2, 32, 1, False, 0, 4096`


**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

 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-devel package with yum
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [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/yinian.lw/miniconda3/lib/python3.9/site-packages/torch']
torch version .................... 1.13.1+cu117
deepspeed install path ........... ['/home/yinian.lw/miniconda3/lib/python3.9/site-packages/deepspeed']
deepspeed info ................... 0.9.0+970d827f, 970d827f, master
torch cuda version ............... 11.7
torch hip version ................ None
nvcc version ..................... 11.3
deepspeed wheel compiled w. ...... torch 1.13, cuda 11.7
`

**System info (please complete the following information):**
 - GPU type:A10
 - Python version: 3.9
 - diffusers version: 0.14.0
 - DeepSpeed version: 0.9.0+970d827f




Discussion:
Hi @yinian-lw. This bug was caused by a recent change. I'll work on a solution. In the meantime, if you install this version of DeepSpeed, you should be able to run Stable Diffusion:
`pip install git+https://github.com/microsoft/deepspeed@6eca037`

Resolving patch:
diff --git a/deepspeed/ops/transformer/inference/diffusers_attention.py b/deepspeed/ops/transformer/inference/diffusers_attention.py
--- a/deepspeed/ops/transformer/inference/diffusers_attention.py
+++ b/deepspeed/ops/transformer/inference/diffusers_attention.py
@@ -59,7 +59,7 @@ def selfAttention_fp(input, context, input_mask):
             scale = (1 / norm_factor) * (1 / norm_factor)
             if do_flash_attn and context == None:
                 qkv_out = linear_func(input, attn_qkvw, attn_qkvb if attn_qkvb is not None else attn_qkvw, attn_qkvb
-                                      is not None, do_flash_attn, config.heads)
+                                      is not None, do_flash_attn, config.heads, False)
 
                 context_layer = triton_flash_attn_kernel(qkv_out[0], qkv_out[1], qkv_out[2], scale,
                                                          input.shape[-2] % 128 == 0)
@@ -82,7 +82,7 @@ def selfAttention_fp(input, context, input_mask):
                 attention_scores = (torch.matmul(query, key.transpose(-1, -2)) * scale).softmax(dim=-1)
                 context_layer = _transpose_for_context(torch.matmul(attention_scores, value))
 
-            output = linear_func(context_layer, attn_ow, attn_ob, do_out_bias, False, config.heads)
+            output = linear_func(context_layer, attn_ow, attn_ob, do_out_bias, False, config.heads, False)
             return output
 
         output = selfAttention_fp(input, context, input_mask)
@@ -184,7 +184,7 @@ def forward(self, input, context=None, input_mask=None):
             self.allocate_workspace(self.config.hidden_size, self.config.heads,
                                     input.size()[1],
                                     input.size()[0], DeepSpeedDiffusersAttention.layer_id, self.config.mp_size, False,
-                                    0, self.config.max_out_tokens)
+                                    0, self.config.max_out_tokens, self.config.min_out_tokens)
         output = DeepSpeedDiffusersAttentionFunction.apply(input, context, input_mask, self.config, self.attn_qkvw,
                                                            self.attn_qw, self.attn_kw, self.attn_vw, self.attn_qkvb,
                                                            self.num_attention_heads_per_partition, self.norm_factor,


Test patch: