Upload train_sft_hf_cloud.py with huggingface_hub

train_sft_hf_cloud.py

@@ -70,8 +70,9 @@ LORA_DROPOUT = 0.05
 # MLP layers (gate_proj, up_proj, down_proj) are in 512 experts - too many params
 LORA_TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "o_proj"]
 
-# Quantization - 8-bit for 80B model on single A100
-USE_8BIT = True
+# Quantization - 4-bit QLoRA for 80B model (more stable than 8-bit on multi-GPU)
+USE_4BIT = True
+USE_8BIT = False
 
 # Trackio monitoring
 TRACKIO_PROJECT = "univer-api-training"
@@ -125,6 +126,7 @@ def main():
     print(f"LR:         {LEARNING_RATE}")
     print(f"LoRA:       r={LORA_R}, alpha={LORA_ALPHA}")
     print(f"Targets:    {LORA_TARGET_MODULES}")
+    print(f"4-bit:      {USE_4BIT}")
     print(f"8-bit:      {USE_8BIT}")
     print("=" * 60)
 
@@ -146,7 +148,21 @@ def main():
         tokenizer.pad_token = tokenizer.eos_token
 
     # Configure quantization
-    if USE_8BIT:
+    if USE_4BIT:
+        print("  Using 4-bit QLoRA quantization...")
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16,
+            bnb_4bit_use_double_quant=True,
+        )
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            quantization_config=bnb_config,
+            device_map="auto",
+            trust_remote_code=True,
+        )
+    elif USE_8BIT:
         print("  Using 8-bit quantization...")
         bnb_config = BitsAndBytesConfig(load_in_8bit=True)
         model = AutoModelForCausalLM.from_pretrained(