update

examples/tutorials/by_deepspeed/step_2_train_model.py

@@ -11,12 +11,12 @@ import platform
 os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
 
 if platform.system() in ("Windows", "Darwin"):
-    from project_settings import project_path
+    from project_settings import project_path, temp_directory
 else:
     project_path = os.path.abspath("../../../")
     project_path = Path(project_path)
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
 
-from peft import LoraConfig
 # from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from modelscope import AutoConfig, AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from trl import SFTTrainer, SFTConfig
@@ -42,14 +42,12 @@ def get_args():
     parser.add_argument("--dataset_split", default=None, type=str),
     parser.add_argument(
         "--dataset_cache_dir",
-        # default=(project_path / "hub_datasets").as_posix(),
-        default="/root/autodl-tmp/OpenMiniMind/hub_datasets",
+        default=(temp_directory / "hub_datasets").as_posix(),
         type=str
     ),
     parser.add_argument(
         "--model_cache_dir",
-        # default=(project_path / "hub_models").as_posix(),
-        default="/root/autodl-tmp/OpenMiniMind/hub_models",
+        default=(temp_directory / "hub_models").as_posix(),
         type=str
     ),
     parser.add_argument("--dataset_streaming", default=None, type=str),
@@ -75,13 +73,13 @@ def main():
         quantization_config=None,
         # device_map="auto",
         trust_remote_code=True,
-        # cache_dir=args.model_cache_dir,
+        cache_dir=args.model_cache_dir,
     )
     print(model)
     tokenizer = AutoTokenizer.from_pretrained(
         pretrained_model_name_or_path=args.model_name,
         trust_remote_code=True,
-        # cache_dir=args.model_cache_dir,
+        cache_dir=args.model_cache_dir,
     )
     print(tokenizer)
 
@@ -128,7 +126,7 @@ def main():
         train_dataset=train_dataset,
         eval_dataset=None,  # Can set up evaluation!
         args=SFTConfig(
-            output_dir="/root/autodl-tmp/OpenMiniMind/trainer_output",  # 请替换为你想要的路径
+            output_dir=(temp_directory / "Qwen3-8B-sft-deepspeed/trainer_output"),  # 请替换为你想要的路径
             dataset_text_field="formated_text",
             deepspeed="./ds_config/deepspeed_stage_3_config.json",  # 添加deepspeed配置文件
             per_device_train_batch_size=1,
@@ -138,6 +136,8 @@ def main():
             # max_steps = 30,
             learning_rate=3e-5,  # Reduce to 2e-5 for long training runs
             logging_steps=1,
+            save_steps=100,           # 每500步保存一次检查点
+            save_total_limit=2,       # 最多只保留2个检查点，旧的自动清理
             optim="adamw_8bit",
             weight_decay=0,
             lr_scheduler_type="constant_with_warmup",
@@ -169,18 +169,10 @@ def main():
     print(f"Peak reserved memory % of max memory = {used_percentage} %.")
     print(f"Peak reserved memory for training % of max memory = {lora_percentage} %.")
 
-    # 只保存lora适配器参数
-    trained_models_dir = project_path / "trained_models" / "Qwen3-8B-sft-deepspeed"
+    trained_models_dir = temp_directory / "trained_models" / "Qwen3-8B-sft-deepspeed"
     trained_models_dir.mkdir(parents=True, exist_ok=True)
     trainer.model.save_pretrained(trained_models_dir.as_posix())
     tokenizer.save_pretrained(trained_models_dir.as_posix())
-
-    # trained_models_dir = project_path / "trained_models" / "Qwen3-8B-sft-fp16"
-    # trained_models_dir.mkdir(parents=True, exist_ok=True)
-    # trainer.model.save_pretrained_merged(trained_models_dir.as_posix(), tokenizer, save_method="merged_16bit",)
-    # trained_models_dir = project_path / "trained_models" / "Qwen3-8B-sft-int4"
-    # trained_models_dir.mkdir(parents=True, exist_ok=True)
-    # trainer.model.save_pretrained_merged(trained_models_dir.as_posix(), tokenizer, save_method="merged_4bit",)
     return