trainer commit (#1)

Dockerfile

@@ -0,0 +1,15 @@
+FROM python:3.9
+
+WORKDIR /code
+
+COPY . .
+
+# RUN pip install --no-cache-dir torch==2.2.1 --index-url https://download.pytorch.org/whl/cu121
+
+RUN pip install --no-cache-dir -r /code/requirements.txt
+
+COPY . .
+
+# CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]
+
+CMD ["python", "train_mamba.py", "--model", "state-spaces/mamba-130m", "--tokenizer", "EleutherAI/gpt-neox-20b", "--learning_rate", "5e-5", "--batch_size", "1", "--gradient_accumulation_steps", "1", "--optim paged_adamw_8bit", "--data_path", "./data/ultrachat_small.jsonl", "--num_epochs", "1"]

app.py

@@ -0,0 +1,89 @@
+import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
+from argparse import ArgumentParser
+
+def get_args():
+    parser = ArgumentParser()
+    parser.add_argument("--port", type=int, default=7860)
+    parser.add_argument("--device", type=str, default='cuda', help='Device to run the model on')
+    parser.add_argument("--model", type=str, default='havenhq/mamba-chat', help='Model to use')
+    parser.add_argument(
+        "--share",
+        action="store_true",
+        default=False,
+        help="share your instance publicly through gradio",
+    )
+    try:
+        args = parser.parse_args()
+    except:
+        parser.print_help()
+        exit(0)
+    return args
+
+
+if __name__ == "__main__":
+    args = get_args()
+
+    device = args.device
+    model_name = args.model
+    eos = "<|endoftext|>"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    tokenizer.eos_token = eos
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.chat_template = AutoTokenizer.from_pretrained(
+        "HuggingFaceH4/zephyr-7b-beta"
+    ).chat_template
+
+    model = MambaLMHeadModel.from_pretrained(
+        model_name, device=device, dtype=torch.float16
+    )
+
+    def chat_with_mamba(
+        user_message,
+        history: list[list[str]],
+        temperature: float = 0.9,
+        top_p: float = 0.7,
+        max_length: int = 2000,
+    ):
+        history_dict: list[dict[str, str]] = []
+        for user_m, assistant_m in history:
+            history_dict.append(dict(role="user", content=user_m))
+            history_dict.append(dict(role="assistant", content=assistant_m))
+        history_dict.append(dict(role="user", content=user_message))
+
+        input_ids = tokenizer.apply_chat_template(
+            history_dict, return_tensors="pt", add_generation_prompt=True
+        ).to(device)
+
+        out = model.generate(
+            input_ids=input_ids,
+            max_length=max_length,
+            temperature=temperature,
+            top_p=top_p,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+
+        decoded = tokenizer.batch_decode(out)
+        assistant_message = (
+            decoded[0].split("<|assistant|>\n")[-1].replace(eos, "")
+        )
+        return assistant_message
+    
+
+    demo = gr.ChatInterface(
+        fn=chat_with_mamba,
+        # examples=[
+        #     "Explain what is state space model",
+        #     "Nice to meet you!",
+        #     "'Mamba is way better than ChatGPT.' Is this statement correct?",
+        # ],
+        additional_inputs=[
+            gr.Slider(minimum=0, maximum=1, step=0.1, value=0.9, label="temperature"),
+            gr.Slider(minimum=0, maximum=1, step=0.1, value=0.7, label="top_p"),
+            gr.Number(value=2000, label="max_length"),
+        ],
+        title="Mamba Chat",
+    )
+    demo.launch(server_port=args.port, share=args.share)

chat.py

@@ -0,0 +1,31 @@
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
+
+device = "cuda"
+tokenizer = AutoTokenizer.from_pretrained("havenhq/mamba-chat")
+tokenizer.eos_token = "<|endoftext|>"
+tokenizer.pad_token = tokenizer.eos_token
+tokenizer.chat_template = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta").chat_template
+
+model = MambaLMHeadModel.from_pretrained("havenhq/mamba-chat", device="cuda", dtype=torch.float16)
+
+messages = []
+while True:
+    user_message = input("\nYour message: ")
+    messages.append(dict(
+        role="user",
+        content=user_message
+    ))
+
+    input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True).to("cuda")
+
+    out = model.generate(input_ids=input_ids, max_length=2000, temperature=0.9, top_p=0.7, eos_token_id=tokenizer.eos_token_id)
+
+    decoded = tokenizer.batch_decode(out)
+    messages.append(dict(
+        role="assistant",
+        content=decoded[0].split("<|assistant|>\n")[-1])
+    )
+
+    print("Model:", decoded[0].split("<|assistant|>\n")[-1])

requirements.txt

@@ -0,0 +1,8 @@
+packaging
+torch==2.2.1+cu121
+transformers
+causal-conv1d
+mamba-ssm
+accelerate
+bitsandbytes
+scipy==1.11.4

scripts/download_ultrachat.py

@@ -0,0 +1,10 @@
+import json
+from datasets import load_dataset
+
+data = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft")
+
+
+with open("../data/ultrachat.jsonl", "w") as f:
+    for d in data:
+        f.write(json.dumps(dict(messages=d["messages"]))+"\n")
+

train_mamba.py

@@ -0,0 +1,62 @@
+import torch
+import argparse
+
+from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
+from transformers import AutoTokenizer, TrainingArguments
+from trainer.data import ChatDataModule
+from trainer.mamba_trainer import MambaTrainer
+
+
+def run(args):
+    
+    print("Loading Mamba {} model".format(args.model))
+    model = MambaLMHeadModel.from_pretrained(args.model, dtype=torch.bfloat16, device="cuda")
+    print("Loading tokenizer {}".format(args.tokenizer))
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
+    tokenizer.eos_token = "<|endoftext|>"
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.chat_template = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta").chat_template
+
+    print("Loading data from {}".format(args.data_path))
+    data_module = ChatDataModule(
+        tokenizer=tokenizer,
+        data_path=args.data_path,
+        conversation_template=tokenizer.chat_template,
+        max_tokens=2048
+    )
+
+    print("Initializing trainer...")
+    trainer = MambaTrainer(
+        model=model,
+        train_dataset=data_module.dataset,
+        tokenizer=tokenizer,
+        args=TrainingArguments(
+            learning_rate=args.learning_rate,
+            num_train_epochs=args.num_epochs,
+            per_device_train_batch_size=args.batch_size,
+            gradient_accumulation_steps=args.gradient_accumulation_steps,
+            optim=args.optim,
+            output_dir="mamba-chat",
+            logging_steps=50,
+            save_steps=500,
+        ),
+        data_collator=data_module.data_collator,
+    )
+    print("Training started...")
+    trainer.train()
+    print("Training finished!")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", type=str, default="state-spaces/mamba-130m")
+    parser.add_argument("--tokenizer", type=str, default="EleutherAI/gpt-neox-20b")
+    parser.add_argument("--learning_rate", type=float, default=5e-5)
+    parser.add_argument("--batch_size", type=int, default=4)
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
+    parser.add_argument("--optim", type=str, default="adamw_torch")
+    parser.add_argument("--data_path", type=str, default="./data/ultrachat_small.jsonl")
+    parser.add_argument("--num_epochs", type=int, default=1)
+    args = parser.parse_args()
+
+    run(args)

trainer/data.py

@@ -0,0 +1,83 @@
+import torch
+import transformers
+import json
+
+from dataclasses import dataclass
+from typing import Dict, Sequence
+from tqdm import tqdm
+from torch.utils.data import Dataset
+
+
+class ChatDataset(Dataset):
+    def __init__(self, data_path: str, tokenizer: transformers.AutoTokenizer, conversation_template: str, max_tokens: int):
+        super(ChatDataset, self).__init__()
+        data = []
+        with open(data_path, "r") as file:
+            for line in file:  
+                try:
+                    data.append(json.loads(line))
+                except Exception as e:
+                    print("json processing exception", e)
+                    continue
+
+
+        data_dict = preprocess(data, tokenizer, conversation_template, max_tokens)
+
+        self.input_ids = data_dict["input_ids"]
+        self.labels = data_dict["labels"]
+
+    def __len__(self):
+        return len(self.input_ids)
+
+    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
+        return dict(input_ids=self.input_ids[i], labels=self.labels[i])
+
+
+@dataclass
+class DataCollatorForChatDataset(object):
+    """
+    Collate examples for supervised fine-tuning.
+    """
+
+    tokenizer: transformers.PreTrainedTokenizer
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        input_ids, labels = tuple([instance[key] for instance in instances] for key in ("input_ids", "input_ids"))
+        input_ids = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id)
+        labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=-100)
+
+        return dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+        )
+    
+
+class ChatDataModule():
+    def __init__(self, tokenizer: transformers.PreTrainedTokenizer, data_path: str, conversation_template, max_tokens: int):
+
+        self.dataset = ChatDataset(tokenizer=tokenizer, data_path=data_path, conversation_template=conversation_template, max_tokens=max_tokens)
+        self.data_collator = DataCollatorForChatDataset(tokenizer=tokenizer)
+        
+
+def preprocess(conversations: Sequence[Sequence[dict]], tokenizer: transformers.PreTrainedTokenizer, conversation_template: str, max_tokens: int) -> Dict:
+    """
+    Preprocess the data by tokenizing.
+    """
+    all_input_ids = []
+    all_label_ids = []
+    tokenizer.use_default_system_prompt = False
+
+    print("Tokenizing dataset...")
+    for conv in tqdm(conversations):
+        current_conv = conv["messages"]
+        tokenized_responses = []
+        for msg in current_conv:
+            if msg["role"] == "assistant":
+                tokenized_responses.append(tokenizer.encode(msg["content"], add_special_tokens=False))
+
+        tokenized_conv = tokenizer.apply_chat_template(current_conv, chat_template=conversation_template, max_length=max_tokens, truncation=True)
+        all_input_ids.append(torch.LongTensor(tokenized_conv))
+
+
+    return dict(input_ids=all_input_ids, labels=all_input_ids)

trainer/mamba_trainer.py

@@ -0,0 +1,39 @@
+from transformers import Trainer
+import torch
+import os
+
+
+class MambaTrainer(Trainer):
+    def compute_loss(self, model, inputs, return_outputs=False):
+        input_ids = inputs.pop("input_ids")
+        lm_logits = model(input_ids).logits
+
+        labels = input_ids.to(lm_logits.device)
+        shift_logits = lm_logits[:, :-1, :].contiguous()
+        labels = labels[:, 1:].contiguous()
+
+        loss_fct = torch.nn.CrossEntropyLoss()
+        lm_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1))
+
+        return lm_loss
+
+    def save_model(self, output_dir, _internal_call):
+        if not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+            
+        torch.save(self.model.state_dict(), f"{output_dir}/pytorch_model.bin")
+        self.tokenizer.save_pretrained(output_dir)
+        
+        json_str = """
+{
+    "d_model": 768,
+    "n_layer": 24,
+    "vocab_size": 50277,
+    "ssm_cfg": {},
+    "rms_norm": true,
+    "residual_in_fp32": true,
+    "fused_add_norm": true,
+    "pad_vocab_size_multiple": 8
+}"""
+        with open(f"{output_dir}/config.json", 'w') as f:
+            f.write(json_str)