[REQUEST] injection_policy for GPTBigCode model

Is your feature request related to a problem? Please describe.
I tried to use injection_policy for from transformers.models.gpt_bigcode.modeling_gpt_bigcode import GPTBigCodeBlock as following:

ds = deepspeed.init_inference(
        model,
        tp={"tp_size": tp_size},
        dtype=torch.half,
        checkpoint=None,
        injection_policy={GPTBigCodeBlock: ("attn.c_proj", "mlp.c_proj")}
    )

Each Block looks like this:

GPTBigCodeBlock(  
        (ln_1): LayerNorm((6144,), eps=1e-05, elementwise_affine=True)  
        (attn): GPTBigCodeAttention(  
          (c_attn): Linear(in_features=6144, out_features=6400, bias=True)  
          (c_proj): Linear(in_features=6144, out_features=6144, bias=True)  
          (attn_dropout): Dropout(p=0.1, inplace=False)  
          (resid_dropout): Dropout(p=0.1, inplace=False)  
        )  
        (ln_2): LayerNorm((6144,), eps=1e-05, elementwise_affine=True)  
        (mlp): GPTBigCodeMLP(  
          (c_fc): Linear(in_features=6144, out_features=24576, bias=True)  
          (c_proj): Linear(in_features=24576, out_features=6144, bias=True)  
          (act): GELUActivation()  
          (dropout): Dropout(p=0.1, inplace=False)  
        )  
      )  

I have 8 GPUS, so the attn.c_proj after splitting should be 6144/8=768, 6400/8=800.
But I encounter this error:
query, key_value = self.c_attn(hidden_states).split((self.embed_dim, 2 * self.kv_dim), dim=2)
File "/opt/conda/envs/ptca/lib/python3.8/site-packages/torch/_tensor.py", line 574, in split
return super(Tensor, self).split_with_sizes(split_size, dim)
RuntimeError: start (768) + length (256) exceeds dimension size (800).

I guess the reason is that the self.kv_dim=128 which is not splitted to 16, thus the length=256 in the error msg.

Describe the solution you'd like
Are there any workaround solutions to this problem since currently I do not how to write a custom injection policy...

Describe alternatives you've considered
Or at least only splitting the MLP layer?

Additional context
None

Start with DeepSpeed's init_inference injection_policy handling and compare it with the referenced transformers.models.gpt_bigcode.modeling_gpt_bigcode.GPTBigCodeBlock structure. Reproduce the 8-GPU configuration and inspect how c_attn dimensions and kv_dim are split. Done means GPTBigCode injection works without the split-size error, or the supported limitation and workaround are documented.