[25.05.20] update code support resume from checkpoints

src_pretrain/pretrain_deepspeed.py

@@ -19,9 +19,9 @@ accelerator = Accelerator()
 set_seed(42)
 
 model_name = "/mnt/han.luzhi/dolly_llm/weights_llm"  
-# resume_option = None                     
-resume_option = True                    
-# resume_option = "/mnt/han.luzhi/result/checkpoint-500"  
+resume_option = None                     
+# resume_option = True                    
+# resume_option = "/mnt/han.luzhi/dolly_llm/result/checkpoint-59000"  
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
@@ -29,18 +29,10 @@ if tokenizer.pad_token is None:
 if isinstance(resume_option, str) and os.path.exists(resume_option):
     logger.info(f"从指定检查点 {resume_option} 加载模型...")
     model = DollyForCausalLM.from_pretrained(resume_option)
-elif resume_option is True and os.path.exists("/mnt/han.luzhi/dolly_llm/result"):
-    latest_checkpoint = Trainer.find_latest_checkpoint("/mnt/han.luzhi/dolly_llm/result")
-    if latest_checkpoint:
-        logger.info(f"自动发现最新检查点 {latest_checkpoint}...")
-        model = DollyForCausalLM.from_pretrained(latest_checkpoint)
-    else:
-        logger.info("未找到检查点，从预训练模型加载...")
-        model = DollyForCausalLM.from_pretrained(model_name)
 else:
+    logger.info("未找到检查点，从预训练模型加载...")
     logger.info(f"从预训练模型 {model_name} 加载...")
     model = DollyForCausalLM.from_pretrained(model_name)
-
 model.gradient_checkpointing_enable()
 
 dataset = load_dataset('text', data_files={'train': '/mnt/han.luzhi/dolly_llm/corpus/wikipedia.txt'})
@@ -58,7 +50,7 @@ tokenized_datasets = dataset.map(
     preprocess_function,
     batched=True,
     remove_columns=['text'], 
-    num_proc=4  
+    num_proc=16  
 )
 
 data_collator = DataCollatorForLanguageModeling(
@@ -67,19 +59,19 @@ data_collator = DataCollatorForLanguageModeling(
 )
 
 training_args = TrainingArguments(
-    output_dir="/mnt/han.luzhi/dolly_llm/result",
+    output_dir="/mnt/han.luzhi/dolly_llm/checkpoints_ds",
     deepspeed="/mnt/han.luzhi/dolly_llm/deepspeed_config/ds_stage2.json",  
     do_eval=False,
     save_strategy="steps",
     save_steps=500,
-    learning_rate=5e-5,
-    per_device_train_batch_size=26,
+    learning_rate=1e-4,
+    per_device_train_batch_size=32,
     num_train_epochs=5,
     logging_dir="/mnt/han.luzhi/dolly_llm/logs",
-    logging_steps=10,
+    logging_steps=20,
     save_total_limit=2,
     fp16=True,  
-    dataloader_num_workers=4,
+    dataloader_num_workers=6,
     gradient_accumulation_steps=2,
     report_to="tensorboard",
     optim="adamw_torch",
@@ -107,7 +99,8 @@ logger.info(f" 进程数 = {accelerator.num_processes}")
 logger.info(f" 设备 = {accelerator.device}")
 logger.info(f" 总批大小 = {training_args.per_device_train_batch_size * accelerator.num_processes * training_args.gradient_accumulation_steps}")
 
-trainer.train(resume_from_checkpoint=resume_option)
+trainer.train()
+# trainer.train(resume_from_checkpoint=resume_option)
 
 if accelerator.is_main_process: 
     trainer.save_model("/mnt/han.luzhi/dolly_llm/result/final")