kyegomez/mambaformer

Implementation of MambaFormer in Pytorch ++ Zeta from the paper: "Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks"

MambaFormer

Implementation of MambaFormer in Pytorch ++ Zeta from the paper: "Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks"

install

pip3 install mamba-former

usage

import torch
from mamba_former.main import MambaFormer

# Forward pass example
x = torch.randint(1, 1000, (1, 100))  # Token
# Tokens are integers representing input data

# Model
model = MambaFormer(
    dim=512,  # Dimension of the model
    num_tokens=1000,  # Number of unique tokens in the input data
    depth=6,  # Number of transformer layers
    d_state=512,  # Dimension of the transformer state
    d_conv=128,  # Dimension of the convolutional layer
    heads=8,  # Number of attention heads
    dim_head=64,  # Dimension of each attention head
    return_tokens=True,  # Whether to return the tokens in the output
)

# Forward pass
out = model(x)  # Perform a forward pass through the model

# If training
# out = model(x, return_loss=True)  # Perform a forward pass and calculate the loss

# Print the output
print(out)  # Print the output tensor
print(out.shape)  # Print the shape of the output tensor

License

MIT

mambaformer's Issues

[BUG] IndexError: Invalid key: 4280187 is out of bounds for size 0

Describe the bug
A clear and concise description of what the bug is and what the main root cause error is. Test very thoroughly before submitting.

To Reproduce

from transformers import AutoTokenizer
from transformers import (AutoTokenizer, AutoConfig, LlamaForCausalLM, DataCollatorForLanguageModeling, Trainer, TrainingArguments)
from transformers import Trainer
from typing import Optional
import torch 
from mamba_former.main import MambaFormer
from datasets import *
from torch.utils.data import *
import os
import sys
import tempfile
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp

tokenizer = AutoTokenizer.from_pretrained("TheBloke/Llama-2-7B-fp16")

print(len(tokenizer))



# Forward pass example
x = torch.randint(1, 1000, (1, 100)) # Token
# Tokens are integrers

# Model
model = MambaFormer(
    dim = 128,
    num_tokens = len(tokenizer),
    depth = 2,
    d_state = 128,
    d_conv = 128,
    heads = 8,
    dim_head = 64,
    return_tokens = True
)

# Forward
out = model(x)
print(out)
print(out.shape)


# count parameters
model_size = sum(t.numel() for t in model.parameters())
print(f"parameter size: {model_size/1000**2:.1f}M parameters")

import datasets


#tokenizer.padding_side = "right"
tokenizer.pad_token = tokenizer.eos_token


tokenized_data = load_dataset("xz56/openwebtext-tokenized-small")

print(tokenized_data)

data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)


output_path = "outputs"
args = TrainingArguments(
    output_dir=output_path,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    evaluation_strategy="steps",
    eval_steps=0.05,
    logging_steps=100,
    gradient_accumulation_steps=2,
    num_train_epochs=1,
    weight_decay=0.01,
    warmup_steps=0.1,
    lr_scheduler_type="cosine",
    learning_rate=1.5e-3,
    save_steps=0.25,
    fp16=True,
    report_to="none"
)

print("Train dataset size:", len(tokenized_data['train']))
print("Test dataset size:", len(tokenized_data['test']))



from transformers import Trainer
from torch.utils.data import RandomSampler, SequentialSampler

class CustomTrainer(Trainer):
    def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
        if isinstance(self.train_dataset, torch.utils.data.IterableDataset):
            return None
        else:
            return RandomSampler(self.train_dataset)

    def _get_eval_sampler(self, eval_dataset: Dataset) -> Optional[torch.utils.data.Sampler]:
        if isinstance(eval_dataset, torch.utils.data.IterableDataset):
            return None
        else:
            return SequentialSampler(eval_dataset)

    def compute_loss(self, model, inputs, return_outputs=False):
        outputs = model(**inputs)
        loss = outputs.loss
        return (loss, outputs) if return_outputs else loss

trainer = CustomTrainer(
    model=model,
    tokenizer=tokenizer,
    args=args,
    data_collator=data_collator,
    train_dataset=tokenized_data["train"],
    eval_dataset=tokenized_data["test"],
)

trainer.train()

trainer.save_model(f"{output_path}/final_model")

Expected behavior
training

Screenshots

ERROR in plain text

Traceback (most recent call last):
  File "/notebooks/main.py", line 118, in <module>
    trainer.train()
  File "/usr/local/lib/python3.9/dist-packages/transformers/trainer.py", line 1624, in train
    return inner_training_loop(
  File "/usr/local/lib/python3.9/dist-packages/transformers/trainer.py", line 1928, in _inner_training_loop
    for step, inputs in enumerate(epoch_iterator):
  File "/usr/local/lib/python3.9/dist-packages/accelerate/data_loader.py", line 452, in __iter__
    current_batch = next(dataloader_iter)
  File "/usr/local/lib/python3.9/dist-packages/torch/utils/data/dataloader.py", line 631, in __next__
    data = self._next_data()
  File "/usr/local/lib/python3.9/dist-packages/torch/utils/data/dataloader.py", line 675, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/usr/local/lib/python3.9/dist-packages/torch/utils/data/_utils/fetch.py", line 49, in fetch
    data = self.dataset.__getitems__(possibly_batched_index)
  File "/usr/local/lib/python3.9/dist-packages/datasets/arrow_dataset.py", line 2814, in __getitems__
    batch = self.__getitem__(keys)
  File "/usr/local/lib/python3.9/dist-packages/datasets/arrow_dataset.py", line 2810, in __getitem__
    return self._getitem(key)
  File "/usr/local/lib/python3.9/dist-packages/datasets/arrow_dataset.py", line 2794, in _getitem
    pa_subtable = query_table(self._data, key, indices=self._indices)
  File "/usr/local/lib/python3.9/dist-packages/datasets/formatting/formatting.py", line 583, in query_table
    _check_valid_index_key(key, size)
  File "/usr/local/lib/python3.9/dist-packages/datasets/formatting/formatting.py", line 536, in _check_valid_index_key
    _check_valid_index_key(int(max(key)), size=size)
  File "/usr/local/lib/python3.9/dist-packages/datasets/formatting/formatting.py", line 526, in _check_valid_index_key
    raise IndexError(f"Invalid key: {key} is out of bounds for size {size}")
IndexError: Invalid key: 4280187 is out of bounds for size 0

[BUG] bug

Traceback (most recent call last):
File "/home/liguanting/project/AdaIR/net/mambaformer.py", line 24, in
out = model(x, return_loss=True) # Perform a forward pass and calculate the loss
File "/home/liguanting/anaconda3/envs/adairvim/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/home/liguanting/anaconda3/envs/adairvim/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "/home/liguanting/anaconda3/envs/adairvim/lib/python3.10/site-packages/mamba_former/main.py", line 126, in forward
x = self.input_mamba(x) + skip
File "/home/liguanting/anaconda3/envs/adairvim/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/home/liguanting/anaconda3/envs/adairvim/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
File "/home/liguanting/anaconda3/envs/adairvim/lib/python3.10/site-packages/zeta/nn/modules/simple_mamba.py", line 107, in forward
(b, l, d) = x.shape
ValueError: not enough values to unpack (expected 3, got 2)