Update train.py

train.py

@@ -15,13 +15,24 @@ if hf_token:
 else:
     raise ValueError("Hugging Face token 未设置")
 
-# 加载基础模型（例如：DeepSeek-R1）
-model_name = "deepseek-ai/DeepSeek-R1"  # 你可以根据需要调整基础模型
-# ✅ 允许远程执行自定义代码
-model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, use_auth_token=hf_token)
 
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, use_auth_token=hf_token)
+# ✅ 设置量化配置，避免 fp8 量化错误
+quantization_config = BitsAndBytesConfig(load_in_4bit=True)  # 使用 4-bit 量化，避免 fp8 错误
 
+# 加载基础模型（例如：DeepSeek-R1）
+model_name = "deepseek-ai/DeepSeek-R1"  # 你可以根据需要调整基础模型
+# ✅ 允许远程执行自定义代码，并启用 4-bit 量化
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    trust_remote_code=True,
+    token=hf_token,  # 旧版 `use_auth_token` 被弃用
+    quantization_config=quantization_config  # 指定 4-bit 量化
+)
+tokenizer = AutoTokenizer.from_pretrained(
+    model_name,
+    trust_remote_code=True,
+    token=hf_token
+)
 # 读取数据文件
 with open('data.json', 'r', encoding='utf-8') as f:
     data = json.load(f)