Upload train_survival.py with huggingface_hub

train_survival.py

@@ -10,10 +10,9 @@ import torch
 import os
 
 # Configuration
-# Switching to Llama 3.2 3B Instruct to avoid BFloat16 issues
-MODEL_ID = "Qwen/Qwen2.5-3B-Instruct"
+MODEL_ID = "meta-llama/Llama-3.2-3B-Instruct"
 DATASET_ID = "sunkencity/survival-instruct"
-OUTPUT_MODEL_ID = "sunkencity/survival-expert-3b"
+OUTPUT_MODEL_ID = "sunkencity/survival-expert-llama-3b"
 
 # Load Dataset
 dataset = load_dataset(DATASET_ID, split="train")
@@ -30,10 +29,11 @@ def filter_empty(example):
 dataset = dataset.filter(filter_empty)
 
 # Load Model
+# We keep 4-bit loading for memory efficiency, but compute in float32 to avoid kernel issues
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_compute_dtype=torch.float32, # Changed to float32
 )
 
 model = AutoModelForCausalLM.from_pretrained(
@@ -41,7 +41,7 @@ model = AutoModelForCausalLM.from_pretrained(
     quantization_config=bnb_config,
     device_map="auto",
     use_cache=False,
-    torch_dtype=torch.float16 # Explicitly set float16
+    torch_dtype=torch.float32 # Changed to float32
 )
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 tokenizer.pad_token = tokenizer.eos_token
@@ -50,7 +50,7 @@ tokenizer.pad_token = tokenizer.eos_token
 def format_row(example):
     instruction = example['instruction']
     response = example['response']
-    text = f"<|im_start|>user\n{instruction}<|im_end|>\n<|im_start|>assistant\n{response}<|im_end|>{tokenizer.eos_token}"
+    text = f"Instruction: {instruction}\nResponse: {response}{tokenizer.eos_token}"
     return {"text": text}
 
 dataset = dataset.map(format_row)
@@ -69,14 +69,14 @@ peft_config = LoraConfig(
 training_args = SFTConfig(
     output_dir="./results",
     num_train_epochs=3,
-    per_device_train_batch_size=4,
-    gradient_accumulation_steps=4,
+    per_device_train_batch_size=1, # Reduced batch size for FP32
+    gradient_accumulation_steps=16, # Increased accumulation to compensate
     learning_rate=2e-4,
     logging_steps=10,
     push_to_hub=True,
     hub_model_id=OUTPUT_MODEL_ID,
-    fp16=True, # Force FP16
-    bf16=False, # Disable BF16 explicitly
+    fp16=False, # Disable Mixed Precision
+    bf16=False, # Disable BF16
     packing=False,
     max_length=1024,
     dataset_text_field="text"
@@ -96,4 +96,4 @@ trainer.train()
 
 print("Pushing to hub...")
 trainer.push_to_hub()
-print("Done!")
+print("Done!")