enable autoTP for MPT

[sywangyi]

No description provided.

[sywangyi]

@yao-matrix @delock please help review

[sywangyi]

verified on model https://huggingface.co/mosaicml/mpt-7b https://huggingface.co/mosaicml/mpt-30b

[delock]

Hi @tjruwase is @molly-smith owner of AutoTP related changes? There are contribution to enable AutoTP for different LLM model, should we invite Molly to review?

[delock]

@sywangyi can you also add mpt into supported models in docs?
https://github.com/microsoft/DeepSpeed/blob/master/docs/_tutorials/automatic-tensor-parallelism.md

[sywangyi]

@sywangyi can you also add mpt into supported models in docs? https://github.com/microsoft/DeepSpeed/blob/master/docs/_tutorials/automatic-tensor-parallelism.md

done

[tjruwase]

Hi @tjruwase is @molly-smith owner of AutoTP related changes? There are contribution to enable AutoTP for different LLM model, should we invite Molly to review?

@delock, sure I will ask for Molly for help. Thanks!

[tjruwase]

Hi @tjruwase is @molly-smith owner of AutoTP related changes? There are contribution to enable AutoTP for different LLM model, should we invite Molly to review?

@delock, sure I will ask for Molly for help. Thanks!

@delock, July 4th is this week and so there might be some delay with reviews as team members take some time off. Apologies for the inconvenience.

[delock]

Hi @tjruwase is @molly-smith owner of AutoTP related changes? There are contribution to enable AutoTP for different LLM model, should we invite Molly to review?

@delock, sure I will ask for Molly for help. Thanks!

@delock, July 4th is this week and so there might be some delay with reviews as team members take some time off. Apologies for the inconvenience.

Got it, thanks for reminding!

[molly-smith]

@tjruwase @sywangyi @delock AutoTP creates a dictionary of which model layers require all reduce and is implemented in deepspeed/module_inject/auto_tp.py. This part is actually already functional for MPT. However, I assume the lower level implementation of tensor parallelism is not functional for MPT. I am not the code owner for this part of TP. I think @RezaYazdaniAminabadi is the right person to review these changes.

[tjruwase]

@molly-smith, thanks for the clarification.

@RezaYazdaniAminabadi, can you please help review this?

[minionKP]

Hi @sywangyi, I was trying out deepspeed for MPT30B as well and saw you had added the changes.

I just tried out your changes by building deepspeed from source by cloning the repository with changes and doing pip install .

# tried with both replace_with_kernel_inject as True and False
ds_engine = deepspeed.init_inference(model, mp_size=4, dtype=torch.float16, replace_with_kernel_inject=True)
model = ds_engine.module
# evaluate model

I am getting the following error when replace_with_kernel_inject is False

│ a00f1a5307c5ca947e72e89693c/attention.py:41 in scaled_multihead_dot_product_attention            │
│                                                                                                  │
│    38 │   │   attn_bias = attn_bias[:, :, _s_q:, _s_k:]                                          │
│    39 │   │   if attn_bias.size(-1) != 1 and attn_bias.size(-1) != s_k or (attn_bias.size(-2)    │
│    40 │   │   │   raise RuntimeError(f'attn_bias (shape: {attn_bias.shape}) is expected to bro   │
│ ❱  41 │   │   attn_weight = attn_weight + attn_bias                                              │
│    42 │   min_val = torch.finfo(q.dtype).min                                                     │
│    43 │   if key_padding_mask is not None:                                                       │
│    44 │   │   if attn_bias is not None:                                                          │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
RuntimeError: The size of tensor a (16) must match the size of tensor b (64) at non-singleton dimension 1

and w/o it I am getting OOM with 4xA100 40GB machines.

Anything I should be doing differently?

[sywangyi]

Hi @sywangyi, I was trying out deepspeed for MPT30B as well and saw you had added the changes.

I just tried out your changes by building deepspeed from source by cloning the repository with changes and doing pip install .

# tried with both replace_with_kernel_inject as True and False
ds_engine = deepspeed.init_inference(model, mp_size=4, dtype=torch.float16, replace_with_kernel_inject=True)
model = ds_engine.module
# evaluate model

I am getting the following error when replace_with_kernel_inject is False

│ a00f1a5307c5ca947e72e89693c/attention.py:41 in scaled_multihead_dot_product_attention            │
│                                                                                                  │
│    38 │   │   attn_bias = attn_bias[:, :, _s_q:, _s_k:]                                          │
│    39 │   │   if attn_bias.size(-1) != 1 and attn_bias.size(-1) != s_k or (attn_bias.size(-2)    │
│    40 │   │   │   raise RuntimeError(f'attn_bias (shape: {attn_bias.shape}) is expected to bro   │
│ ❱  41 │   │   attn_weight = attn_weight + attn_bias                                              │
│    42 │   min_val = torch.finfo(q.dtype).min                                                     │
│    43 │   if key_padding_mask is not None:                                                       │
│    44 │   │   if attn_bias is not None:                                                          │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
RuntimeError: The size of tensor a (16) must match the size of tensor b (64) at non-singleton dimension 1

and w/o it I am getting OOM with 4xA100 40GB machines.

Anything I should be doing differently?

I could run mpt30B 4TP successfully in 4 RTX 8000 using AutoTP, seem the attn_bias is incorrect in your side, could you have a check if https://github.com/microsoft/DeepSpeed/pull/3861/files#diff-27be33d45da8a29a59628046c212ecbdb630e85a9bd987a98431272f1472a3fbR63 is correctly called.

[minionKP]

Hi @sywangyi, I was trying out deepspeed for MPT30B as well and saw you had added the changes.
I just tried out your changes by building deepspeed from source by cloning the repository with changes and doing pip install .

# tried with both replace_with_kernel_inject as True and False
ds_engine = deepspeed.init_inference(model, mp_size=4, dtype=torch.float16, replace_with_kernel_inject=True)
model = ds_engine.module
# evaluate model

I am getting the following error when replace_with_kernel_inject is False

│ a00f1a5307c5ca947e72e89693c/attention.py:41 in scaled_multihead_dot_product_attention            │
│                                                                                                  │
│    38 │   │   attn_bias = attn_bias[:, :, _s_q:, _s_k:]                                          │
│    39 │   │   if attn_bias.size(-1) != 1 and attn_bias.size(-1) != s_k or (attn_bias.size(-2)    │
│    40 │   │   │   raise RuntimeError(f'attn_bias (shape: {attn_bias.shape}) is expected to bro   │
│ ❱  41 │   │   attn_weight = attn_weight + attn_bias                                              │
│    42 │   min_val = torch.finfo(q.dtype).min                                                     │
│    43 │   if key_padding_mask is not None:                                                       │
│    44 │   │   if attn_bias is not None:                                                          │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
RuntimeError: The size of tensor a (16) must match the size of tensor b (64) at non-singleton dimension 1

and w/o it I am getting OOM with 4xA100 40GB machines.
Anything I should be doing differently?

I could run mpt30B 4TP successfully in 4 RTX 8000 using AutoTP, seem the attn_bias is incorrect in your side, could you have a check if https://github.com/microsoft/DeepSpeed/pull/3861/files#diff-27be33d45da8a29a59628046c212ecbdb630e85a9bd987a98431272f1472a3fbR63 is correctly called.

I did a mistake in my setup with your changes. Can confirm it now works. Thanks

[sywangyi]

@mrwyattii I have moved the model specific ops to a specific file. could you review the change?