Upload train_sft_hf_cloud.py with huggingface_hub

train_sft_hf_cloud.py

@@ -70,9 +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 - 4-bit QLoRA for 80B model (more stable than 8-bit on multi-GPU)
-USE_4BIT = True
-USE_8BIT = False
+# Quantization - 8-bit for 80B model on multi-GPU (A100x4)
+USE_4BIT = False
+USE_8BIT = True
 
 # Trackio monitoring
 TRACKIO_PROJECT = "univer-api-training"
@@ -164,11 +164,28 @@ def main():
         )
     elif USE_8BIT:
         print("  Using 8-bit quantization...")
-        bnb_config = BitsAndBytesConfig(load_in_8bit=True)
+        bnb_config = BitsAndBytesConfig(
+            load_in_8bit=True,
+            llm_int8_enable_fp32_cpu_offload=False,  # Keep everything on GPU
+        )
+
+        # For multi-GPU, set max_memory per device
+        n_gpus = torch.cuda.device_count()
+        print(f"  Detected {n_gpus} GPUs")
+
+        if n_gpus > 1:
+            # Distribute model across GPUs with balanced memory
+            max_memory = {i: "70GiB" for i in range(n_gpus)}
+            max_memory["cpu"] = "0GiB"  # Don't offload to CPU
+            print(f"  Max memory per GPU: 70GiB")
+        else:
+            max_memory = None
+
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_NAME,
             quantization_config=bnb_config,
             device_map="auto",
+            max_memory=max_memory,
             trust_remote_code=True,
         )
     else: