Fine-Tuning a Mamba Model with using Hugging Face Transformers

šŸ¤—Transformers
1

Hey community!

Iā€™m working on fine-tuning the Mamba model (specifically state-spaces/mamba-2.8b-hf ) for a multi-turn dialogue system , but Iā€™m hitting some roadblocks. My goal is to build a chatbot that retains context across conversations, like:

Input : Dialogue1: Hi! Can you recommend a pizza place?
Dialogue2: Sure! Are you looking for vegan options?
Dialogue3: Yes, preferably near downtown.

Output : [Bot]: [Expected Response]

My Setup:

  • Using Hugging Face Transformers and PEFT for LoRA.
  • Training on custom conversational multi-turn dialogue data.

Specific Questions:

  1. Data Formatting :
  • How should I structure multi-turn dialogues? Iā€™m using <|endoftext|> as a separator(eos token for state-spaces/mamba-2.8b-hf), but the model ignores past turns.
  • Should I prepend [User] /[Bot] labels or use special tokens?
  1. LoRA Targets :
  • Which Mamba layers should I adapt?

Code Snippet (Training Args):

pythontraining_args = TrainingArguments(
per_device_train_batch_size=2,
gradient_accumulation_steps=4,
learning_rate=3e-5,
fp16=True,
)

I am having hard time writing the code for mamba 2.8b, to fine-tune it. Either it doesnā€™t work or it doesnā€™t fine-tune properly.

Any tips on architecture tweaks, data prep, evaluation strategies or any code suggestions/documentations ?

2 Likes
2

Hmm, I donā€™t have any clues as to why the fine tuning isnā€™t working, so Iā€™ve gathered information that is easy to find.

Fine-tuning Mamba

Tips

Here are some tips and resources that might help you fine-tune the Mamba model for your multi-turn dialogue system:

Data Formatting:

  1. Input Structure: For multi-turn dialogues, itā€™s essential to include the entire conversation history in the input to maintain context. This means including all previous turns in the input sequence so the model can learn to retain context. For example:

    [User]: Hi! Can you recommend a pizza place?
[Bot]: Sure! Are you looking for vegan options?
[User]: Yes, preferably near downtown.

  2. Separator Tokens: While <|endoftext|> is a common separator, the model might not be leveraging it effectively for context retention. Instead, you can use [User] and [Bot] labels to explicitly differentiate between user and bot turns. This can help the model understand the flow of the conversation better.

  3. Data Formatting Example: A properly formatted input might look like this:

    [User]: Hi! Can you recommend a pizza place?
[Bot]: Sure! Are you looking for vegan options?
[User]: Yes, preferably near downtown.
[Bot]:

    Here, the model is prompted to generate a response for the bot after seeing the entire conversation history.

LoRA Targets:

  1. Layers to Adapt: For fine-tuning Mamba-2.8b-hf, focus on adapting layers that are responsible for attention and context processing. Since Mamba uses a hybrid SSM-Transformer architecture, prioritize adapting the attention layers and possibly the feed-forward layers.

  2. LoRA Configuration: Start with a relatively small number of LoRA adapters, such as 8, to see how the model performs. If the model underfits, you can increase the number of adapters.

Code and Architecture Tweaks:

  1. Training Arguments: Your current training arguments look reasonable, but you might want to experiment with the learning rate and number of training steps. For example:

    training_args = TrainingArguments(
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    learning_rate=3e-5,
    fp16=True,
    num_train_epochs=3,
    max_steps=1000,
)

  2. Model Configuration: Ensure youā€™re using the correct configuration for Mamba-2.8b-hf. You might need to adjust the modelā€™s attention mechanism or the context window if itā€™s not capturing the entire conversation history.

  3. Evaluation Strategy: Implement an evaluation strategy where the model is evaluated on multi-turn conversations to ensure it retains context. For example, you can evaluate the model by having it generate responses after seeing multiple turns of a conversation.

Additional Tips:

  • Fine-Tuning Script: If youā€™re struggling with writing the fine-tuning script, consider using the Hugging Face PEFT libraryā€™s provided scripts as a starting point. You can adapt these scripts to work with Mamba-2.8b-hf.
  • Documentation and Examples: Review the Mamba model documentation and any available examples or notebooks for fine-tuning Mamba models on dialogue tasks.

Concrete Code Suggestions:

Hereā€™s a rough outline of how your code might look:

from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM

# Load the Mamba model
model = AutoModelForCausalLM.from_pretrained("state-spaces/mamba-2.8b-hf")

# Define LoRA configuration
lora_config = LoraConfig(
    r=8,
    lora_alpha=32,
    target_modules=[
        "q_attn",
        "v_attn",
        # Add other attention layers as needed
    ],
)

# Apply LoRA to the model
model = get_peft_model(model, lora_config)

# Your training loop here

Final Thoughts:

Fine-tuning Mamba-2.8b-hf for multi-turn dialogue requires careful attention to how you structure your data and which parts of the model you adapt. Start with small adjustments and gradually experiment with more complex configurations as needed.