mamba-dpo-qlora.yaml

# Model arguments
model_name_or_path: state-spaces/mamba-2.8b-hf
torch_dtype: float32

# LoRA arguments
use_peft: true
load_in_4bit: true
lora_r: 128
lora_alpha: 128
lora_dropout: 0.05
lora_target_modules:
- x_proj
- embeddings
- in_proj
- out_proj

# Data training arguments

dataset_mixer:
  HuggingFaceH4/ultrafeedback_binarized: 1.0
dataset_splits:
- train_prefs
- test_prefs
preprocessing_num_workers: 12

# DPOTrainer arguments
bf16: false
beta: 0.01
do_eval: false
evaluation_strategy: steps
eval_steps: 100
gradient_accumulation_steps: 4
gradient_checkpointing: true
gradient_checkpointing_kwargs:
  use_reentrant: false
hub_model_id: mamba-2.8b-dpo-qlora
learning_rate: 2.0e-03
log_level: info
logging_steps: 10
lr_scheduler_type: cosine
max_length: 1024
max_prompt_length: 512
num_train_epochs: 1
optim: paged_adamw_32bit
adam_beta1: 0.9
adam_beta2: 0.95
weight_decay: 0.1
max_grad_norm: 1
output_dir: /mnt/fast/nobackup/users/nt00601/mamba-alignment-handbook/data/mamba-2.8b-hf-dpo-qlora # It is handy to append `hub_model_revision` to keep track of your local experiments
per_device_train_batch_size: 32
per_device_eval_batch_size: 8
push_to_hub: true
save_strategy: "steps"
save_steps: 100
save_total_limit: 1
seed: 42
warmup_ratio: 0.1