Repository: pytorch
Issue: Semi-Structured Sparsity Causing Memory Overflow and Unsupported Operations in LLaVA Model Inference on A100 GPU
Labels: module: sparse,triaged

## Issue description

I am following the method of this blog [Accelerating Generative AI with PyTorch: Segment Anything, Fast](https://pytorch.org/blog/accelerating-generative-ai/#sparse-semi-structured-24-sparsity) to accelerate [LLaVA](https://llava-vl.github.io/)  inference.

When I experimented with the Semi-structured sparsity solution, I encountered a problem.

First, when I apply semi-structured sparse to the entire model, the memory explodes (cuda out of memory). But my A100 has 40960 MiB of video memory, which logically shouldn't be the case.

LLaVA is composed of Llama + CLIP, and the model structure is as follows.

```
LlavaLlamaForCausalLM(
   (model): LlavaLlamaModel(
     (embed_tokens): Embedding(32000, 5120, padding_idx=0)
     (layers): ModuleList(
       (0-39): 40 x LlamaDecoderLayer(
         (self_attn): LlamaAttention(
           (q_proj): Linear(in_features=5120, out_features=5120, bias=False)
           (k_proj): Linear(in_features=5120, out_features=5120, bias=False)
           (v_proj): Linear(in_features=5120, out_features=5120, bias=False)
           (o_proj): Linear(in_features=5120, out_features=5120, bias=False)
           (rotary_emb): LlamaRotaryEmbedding()
         )
         (mlp): LlamaMLP(
           (gate_proj): Linear(in_features=5120, out_features=13824, bias=False)
           (up_proj): Linear(in_features=5120, out_features=13824, bias=False)
           (down_proj): Linear(in_features=13824, out_features=5120, bias=False)
           (act_fn): SiLUActivation()
         )
         (input_layernorm): LlamaRMSNorm()
         (post_attention_layernorm): LlamaRMSNorm()
       )
     )
     (norm): LlamaRMSNorm()
     (vision_tower): CLIPVisionTower(
       (vision_tower): CLIPVisionModel(
         (vision_model): CLIPVisionTransformer(
           (embeddings): CLIPVisionEmbeddings(
             (patch_embedding): Conv2d(3, 1024, kernel_size=(14, 14), stride=(14, 14), bias=False)
             (position_embedding): Embedding(577, 1024)
           )
           (pre_layrnorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
           (encoder): CLIPEncoder(
             (layers): ModuleList(
               (0-23): 24 x CLIPEncoderLayer(
                 (self_attn): CLIPAttention(
                   (k_proj): Linear(in_features=1024, out_features=1024, bias=True)
                   (v_proj): Linear(in_features=1024, out_features=1024, bias=True)
                   (q_proj): Linear(in_features=1024, out_features=1024, bias=True)
                   (out_proj): Linear(in_features=1024, out_features=1024, bias=True)
                 )
                 (layer_norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
                 (mlp): CLIPMLP(
                   (activation_fn): QuickGELUActivation()
                   (fc1): Linear(in_features=1024, out_features=4096, bias=True)
                   (fc2): Linear(in_features=4096, out_features=1024, bias=True)
                 )
                 (layer_norm2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
               )
             )
           )
            (post_layernorm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
         )
       )
     )
     (mm_projector): Sequential(
       (0): Linear(in_features=1024, out_features=5120, bias=True)
       (1): GELU(approximate='none')
       (2): Linear(in_features=5120, out_features=5120, bias=True)
     )
   )
   (lm_head): Linear(in_features=5120, out_features=32000, bias=False)
)
```

So I performed semi-structured sparse on vision_tower alone, but encountered another error.

![image](https://github.com/pytorch/pytorch/assets/45281765/c0ee5a83-601c-4358-92ca-8b68cecc0429)

The cause of this error has not been found for the time being, so I chose to only perform semi-structured sparse on the first layer of Llama, that is, apply_sparse(model.model.layers[0]). It did not burst the video memory, but the following error was reported.

![image](https://github.com/pytorch/pytorch/assets/45281765/16b3058d-2e3c-4bd3-9644-7aa7d32ee51b)

In the Llama model structure of LLaVA, the Linear layer of the attention layer does not contain bias. It is speculated that the above error may be caused by the lack of bias.

In addition, in the PyTorch documentation, semi-structured sparse only supports the following operations, and bias is also required in Linear.

![image](https://github.com/pytorch/pytorch/assets/45281765/1ee2232c-5a8b-45b4-ac1c-0e52df8627b9)

https://pytorch.org/docs/stable/sparse.html#sparse-semi-structured-tensors

In LLaVA, except Vision Tower, only the Linear of mm_projector contains bias. Sure enough, `apply_sparse(models.model.mm_projector)` succeeded, but the inference speed did not increase.

Performing semi-structured sparse on mm_projector, memory usage: 28194 MiB, time-consuming: 149.9 ms (average of 200 pictures).
Without semi-structured sparse mm_projector, memory usage: 28034 MiB, time consumption: 150.7 ms (average of 200 pictures).

my question is:
1. Why does the video memory explode? How is semi-structured sparse converted?
2. Is it possible to make the self-attention layer in Llama support semi-structured sparse?
3. Is it possible to support semi-structured sparse in Vision Tower?

## System Info

![image](https://github.com/pytorch/pytorch/assets/45281765/2b085557-a23a-4773-9aca-92bc807b37d7)
![image](https://github.com/pytorch/pytorch/assets/45281765/195fbfb6-f2ef-45b1-8437-2ffcd02dafd1)



cc @alexsamardzic @nikitaved @pearu @cpuhrsch @amjames @bhosmer @jcaip