transformers pytorch mistral text-generation conversational license:apache-2.0 autotrain_compatible endpoints_compatible text-generation-inference region:us
---
license: apache-2.0
---

### To Run:

```python
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "nm-testing/zephyr-50sparse-24"
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_id)
inputs = tokenizer("Hello my name is", return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=20)
print(tokenizer.batch_decode(outputs)[0])
```

### Utility to check if the model is 2:4:
- Note: this samples 10 rows randomly from each weight of the model and looks at every group of 4 in that row

```python
import torch
import random

@torch.no_grad()
def check_24(model, num_rows_to_sample=10, _verbose=False):
    BLOCK_SIZE = 4
    MAX_NON_ZEROS = 2
    
    for idx, layer in enumerate(model.model.layers):
        print(f"----- Checking layer idx {idx}")
        weights = [
            layer.self_attn.q_proj,
            layer.self_attn.k_proj,
            layer.self_attn.v_proj,
            layer.self_attn.o_proj,
            layer.mlp.gate_proj,
            layer.mlp.up_proj,
            layer.mlp.down_proj,
        ]

        for weigth in weights:
            num_rows, num_cols = weight.shape
            sampled_row_idxs = random.choices(range(num_rows), k=num_rows_to_sample)
            if _verbose:
                print(f"Sampled row idxs = {sampled_row_idxs}")

            for i in sampled_row_idxs:
                for j in range(0,num_cols - BLOCK_SIZE,BLOCK_SIZE):
                    block = weight[i, j:j+BLOCK_SIZE]
                    num_nonzero = torch.count_nonzero(block)
                    if num_nonzero > MAX_NON_ZEROS:
                        print(f"found a non 2:4 segment at block {i,j}")
```