update trainer argument

scripts/run_simpo.py

@@ -218,6 +218,7 @@ def main():
         use_cache=False if training_args.gradient_checkpointing else True,
         device_map=get_kbit_device_map() if quantization_config is not None else None,
         quantization_config=quantization_config,
+        attn_implementation=training_args.attn_implementation,
     )
 
     model = model_args.model_name_or_path
@@ -245,6 +246,7 @@ def main():
     #     )
     #     model_kwargs = None
 
+    training_args.model_init_kwargs = model_kwargs
     #########################
     # Instantiate SimPO trainer
     #########################

scripts/simpo_config.py

@@ -68,3 +68,5 @@ class SimPOConfig(TrainingArguments):
     model_init_kwargs: Optional[Dict] = None
 
     dataset_num_proc: Optional[int] = None
+
+    attn_implementation: str = None

training_configs/llama-3-8b-base-sft.yaml

@@ -2,7 +2,7 @@
 model_name_or_path: meta-llama/Meta-Llama-3-8B
 model_revision: main
 torch_dtype: bfloat16
-use_flash_attention_2: true
+attn_implementation: flash_attention_2
 
 # Data training arguments
 chat_template: "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}"

training_configs/llama-3-8b-base-simpo.yaml

@@ -1,7 +1,7 @@
 # Model arguments
 model_name_or_path: princeton-nlp/Llama-3-Base-8B-SFT
 torch_dtype: null
-use_flash_attention_2: true
+attn_implementation: flash_attention_2
 
 # Data training arguments
 dataset_mixer:

training_configs/llama-3-8b-instruct-simpo-v2.yaml

@@ -1,7 +1,7 @@
 # Model arguments
 model_name_or_path: meta-llama/Meta-Llama-3-8B-Instruct
 torch_dtype: null
-use_flash_attention_2: true
+attn_implementation: flash_attention_2
 
 # Data training arguments
 dataset_mixer:

training_configs/llama-3-8b-instruct-simpo.yaml

@@ -1,7 +1,7 @@
 # Model arguments
 model_name_or_path: meta-llama/Meta-Llama-3-8B-Instruct
 torch_dtype: null
-use_flash_attention_2: true
+attn_implementation: flash_attention_2
 
 # Data training arguments
 dataset_mixer:

training_configs/mistral-7b-base-simpo.yaml

@@ -1,7 +1,7 @@
 # Model arguments
 model_name_or_path: alignment-handbook/zephyr-7b-sft-full
 torch_dtype: null
-use_flash_attention_2: true
+attn_implementation: flash_attention_2
 
 # Data training arguments
 dataset_mixer:

training_configs/mistral-7b-instruct-simpo.yaml

@@ -1,7 +1,7 @@
 # Model arguments
 model_name_or_path: mistralai/Mistral-7B-Instruct-v0.2
 torch_dtype: null
-use_flash_attention_2: true
+attn_implementation: flash_attention_2
 
 # Data training arguments
 dataset_mixer: