Views
No views yet
xlstm-jax framework.xlstm, which now uses the mlstm_kernels package for triton kernels (tested on python 3.11):1pip install xlstm
2pip install accelerate
3pip install 'transformers @ git+https://github.com/huggingface/transformers.git@main'pip install 'triton @ git+https://github.com/triton-lang/triton.git@main'1from transformers import AutoModelForCausalLM, AutoTokenizer
2import torch
3
4xlstm = AutoModelForCausalLM.from_pretrained("NX-AI/xLSTM-7b", device_map="auto")
5
6# this is a fork of EleutherAI/gpt-neox-20b
7tokenizer = AutoTokenizer.from_pretrained("NX-AI/xLSTM-7b")
8
9tokens = tokenizer("Explain quantum computing in simple terms.", return_tensors='pt')['input_ids'].to(device="cuda")
10
11# Get the BOS token ID from the tokenizer
12bos_id = tokenizer.bos_token_id
13
14# Prepend BOS
15bos_tensor = torch.tensor([[bos_id]], device=tokens.device, dtype=tokens.dtype)
16tokens_with_bos = torch.cat([bos_tensor, tokens], dim=1)
17
18out = xlstm.generate(tokens_with_bos, max_new_tokens=20)
19
20print(tokenizer.decode(out[0]))1from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
2import torch
3
4xlstm_config = AutoConfig.from_pretrained("NX-AI/xLSTM-7b")
5xlstm_config.step_kernel = "native"
6xlstm_config.chunkwise_kernel = "chunkwise--native_autograd"
7xlstm_config.sequence_kernel = "native_sequence__native"
8
9xlstm = AutoModelForCausalLM.from_pretrained("NX-AI/xLSTM-7b",
10 config=xlstm_config, device_map="auto")
11
12# Load the tokenizer
13tokenizer = AutoTokenizer.from_pretrained("NX-AI/xLSTM-7b")
14
15# Your prompt
16prompt = "Explain quantum computing in simple terms."
17
18# Tokenize and send to the same device as the model
19inputs = tokenizer(prompt, return_tensors="pt")['input_ids'].to(xlstm.device)
20
21# Get the BOS token ID from the tokenizer
22bos_id = tokenizer.bos_token_id
23
24# Prepend BOS
25bos_tensor = torch.tensor([[bos_id]], device=xlstm.device, dtype=inputs.dtype)
26tokens_with_bos = torch.cat([bos_tensor, inputs], dim=1)
27
28# Generate
29outputs = xlstm.generate(
30 tokens_with_bos,
31 max_new_tokens=200, # adjust for output length
32 temperature=0.7, # randomness
33 top_p=0.9, # nucleus sampling
34 do_sample=True
35)
36
37# Decode and print
38print(tokenizer.decode(outputs[0]))
39
40# verify selected kernels
41from pprint import pprint
42pprint(xlstm.backbone.blocks[0].mlstm_layer.config)torch.cuda.graph and torch.compile optimizations on one NVIDIA H100:
lm_eval:| BBH | MMLU-Pro | Math | MUSR | GPQA | IfEval |
|---|---|---|---|---|---|
| 0.381 | 0.242 | 0.036 | 0.379 | 0.280 | 0.244 |
lighteval in the Leaderboard-v1 settings:| Arc-Challenge (25-shot) | MMLU (5-shot) | Hellaswag (10-shot) | Winogrande (5-shot) | TruthfulQA (0-shot) | GSM8k (5-shot) | OpenbookQA (5-shot) | PiQA (5-shot) |
|---|---|---|---|---|---|---|---|
| 0.584 | 0.589 | 0.710 | 0.742 | 0.420 | 0.004 | 0.443 | 0.817 |
LICENSE file)