README.md

---

language:
- en
tags:
- pytorch
- causal-lm
- muddformer
license: mit
---

In comparison with Pythia-2.8B, MUDDPythia-2.8B is a pretrained language model on the Pile with 300B tokens, which uses a simple yet effective method to address the limitations of residual connections and enhance cross-layer information flow in Transformers. Please see downstrem evaluations and more details in the paper[(MUDDFormer: Breaking Residual Bottlenecks in Transformers via Multiway Dynamic Dense Connections)](https://arxiv.org). In addition, we open-source Jax training code on [(Github)](https://github.com/Caiyun-AI/MUDDFormer/).

We recommend <strong>compiled version</strong> of MUDDPythia with *torch.compile* for inference acceleration. Please refer to Generation section for compile implementation.

# Usage

## Env
 
```

pip install transformers==4.35.0 torch==2.5.1

```

## Generation 

```

import time

from transformers import AutoTokenizer, AutoModelForCausalLM

import torch



import os

os.environ['TOKENIZERS_PARALLELISM'] = 'false'



device = torch.device('cuda:0')

dtype = torch.bfloat16

MAX_BATCH_SIZE = 1

MAX_SEQ_LENGTH = 2048

NUM_TOKENS_TO_GENERATE = 10

COMPILE = True

OPTIMIZED_COMPPILE = False 



if OPTIMIZED_COMPPILE: 

    import torch._dynamo.config

    import torch._inductor.config

    torch._dynamo.config.cache_size_limit = 64

    torch._inductor.config.coordinate_descent_tuning = True

    torch._inductor.config.triton.unique_kernel_names = True

    torch._inductor.config.fx_graph_cache = True



tokenizer = AutoTokenizer.from_pretrained("Caiyun-AI/MUDDPythia-2.8B")

model = AutoModelForCausalLM.from_pretrained("Caiyun-AI/MUDDPythia-2.8B", trust_remote_code=True)



_ = model.to(device=device,dtype=dtype)

with torch.device(device):

    model.setup_caches(max_batch_size=MAX_BATCH_SIZE, max_seq_length=MAX_SEQ_LENGTH,dtype=dtype)



def decode_one_token(model, cur_token, input_pos):

    logits = model(cur_token, input_pos=input_pos, return_tensor=True)

    new_token = torch.argmax(logits[:, -1], dim=-1)[:,None]

    return new_token



prompt = "Beijing is the capital of China. London is the capital of"

input_ids = tokenizer.encode(prompt, return_tensors='pt')



compiled_decode_one_token = torch.compile(decode_one_token,mode="reduce-overhead", fullgraph=True) if COMPILE else None



print('Start generating tokens, but it will take a few minutes to compile at the first time.')

for i in range(10):

    t0 = time.time()

    with torch.no_grad():

        generated_ids = model.generate(input_ids.to(device),num_tokens_to_generate=NUM_TOKENS_TO_GENERATE, compiled_decode_one_token=compiled_decode_one_token)

        text = tokenizer.decode(generated_ids[0])

        if i ==0:

            print(f'Generated text: {text}')

    t1 = time.time()

    print(f'Time consumed at iteration {i}: {t1-t0}s')

```