Upload scripts/train_orpo_n8n_thinking.py with huggingface_hub

scripts/train_orpo_n8n_thinking.py

@@ -9,6 +9,7 @@
 #     "bitsandbytes",
 #     "sentencepiece",
 #     "protobuf",
+#     "flash-attn",
 # ]
 # ///
 
@@ -23,9 +24,10 @@ in a single training objective, making it more efficient than DPO for this use c
 """
 
 import trackio
+import torch
 from datasets import load_dataset
 from peft import LoraConfig
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from trl import ORPOTrainer, ORPOConfig
 
 
@@ -56,11 +58,20 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
-print(f"Loading model from {MODEL_NAME}...")
+# 4-bit quantization config to reduce memory
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_use_double_quant=True,
+)
+
+print(f"Loading model from {MODEL_NAME} with 4-bit quantization...")
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
-    torch_dtype="auto",
+    quantization_config=bnb_config,
     device_map="auto",
+    attn_implementation="flash_attention_2",
 )
 
 # LoRA configuration for efficient training on 7B model
@@ -87,10 +98,10 @@ config = ORPOConfig(
     # Training parameters
     num_train_epochs=2,
     per_device_train_batch_size=1,
-    gradient_accumulation_steps=16,  # Effective batch size = 16
+    gradient_accumulation_steps=32,  # Effective batch size = 32
     learning_rate=5e-5,
-    max_length=4096,  # Long context for workflows + thinking
-    max_prompt_length=512,
+    max_length=2048,  # Reduced for memory
+    max_prompt_length=256,
 
     # Memory optimization
     gradient_checkpointing=True,