[25.05.20] update code support resume from checkpoints

README.md

@@ -5,7 +5,7 @@ LLM-Dolly is a custom LLM built from scratch
 ## 🗓️ 更新日志
 
 #### 2025-05-19
-- 📝 添加支持使用deepspeed训练代码.
+- 📝 添加支持使用deepspeed训练代码，测试训练中最大batch_size提升38%，训练速度提升9.6%.
 
 #### 2025-05-16
 - 📝 添加dolly_llm的预训练代码，进行一次预训练测试：模型0.6B，语料500M，46GB*4显卡.

deepspeed_config/ds_stage2.json

@@ -1,40 +1,34 @@
 {
-    "fp16": {
+  "fp16": {
-      "enabled": true,
+    "enabled": true,
-      "loss_scale": 0,
+    "loss_scale": 0,
-      "loss_scale_window": 1000,
+    "loss_scale_window": 1000,
-      "initial_scale_power": 16,
+    "initial_scale_power": 16,
-      "hysteresis": 2,
+    "hysteresis": 2,
-      "min_loss_scale": 1
+    "min_loss_scale": 1
+  },
+  "optimizer": {
+    "type": "AdamW",
+    "params": {
+      "lr": "auto",
+      "weight_decay": 0.01
+    }
+  },
+  "zero_optimization": {
+    "stage": 2,
+    "offload_optimizer": {
+      "device": "none",
+      "pin_memory": true
     },
-    "optimizer": {
+    "allgather_partitions": true,
-      "type": "AdamW",
+    "allgather_bucket_size": 2e8,
-      "params": {
+    "overlap_comm": true,
-        "lr": 5e-5,
+    "reduce_scatter": true,
-        "weight_decay": 0.01
+    "reduce_bucket_size": 2e8,
-      }
+    "contiguous_gradients": true
-    },
+  },
-    "scheduler": {
+  "gradient_accumulation_steps": 2,
-      "type": "CosineAnnealing",
+  "train_batch_size": "auto",
-      "params": {
+  "gradient_clipping": 1.0,
-        "warmup_steps": 500
+  "steps_per_print": 20
-      }
+}
-    },
-    "zero_optimization": {
-      "stage": 2,
-      "offload_optimizer": {
-        "device": "cpu",
-        "pin_memory": true
-      },
-      "allgather_partitions": true,
-      "allgather_bucket_size": 2e8,
-      "overlap_comm": true,
-      "reduce_scatter": true,
-      "reduce_bucket_size": 2e8,
-      "contiguous_gradients": true
-    },
-    "gradient_accumulation_steps": 2,
-    "train_batch_size": "auto",
-    "gradient_clipping": 1.0,
-    "steps_per_print": 10
-  }

requirements.txt

@@ -1,5 +1,10 @@
 torch==2.3.1
+datasets==3.6.0
 sentencepiece==0.2.0
-protobuf==3.20.0 #important
+protobuf==3.20.0
 transformers==4.51.3
 tokenizers==0.21.1
+setuptools==70.3.0
+accelerate==1.7.0
+tensorboardX==2.6.2.2
+deepspeed==0.16.7

requirements_version.txt

@@ -1,4 +1,4 @@
-<<<<<<< HEAD
+accelerate==1.7.0
 aiohappyeyeballs==2.6.1
 aiohttp==3.11.18
 aiosignal==1.3.2
@@ -20,6 +20,7 @@ comm==0.2.2
 datasets==3.6.0
 debugpy==1.8.2
 decorator==5.1.1
+deepspeed==0.16.7
 defusedxml==0.7.1
 dill==0.3.8
 distlib==0.3.9
@@ -54,8 +55,6 @@ jupyterlab==4.2.3
 jupyterlab-language-pack-zh-CN==4.2.post1
 jupyterlab_pygments==0.3.0
 jupyterlab_server==2.27.2
-# Editable install with no version control (llm_dolly==0.0.1)
--e /mnt/han.luzhi/dolly_llm
 MarkupSafe==2.1.5
 matplotlib-inline==0.1.7
 mistune==3.0.2
@@ -71,39 +70,6 @@ networkx==3.3
 notebook==6.5.2
 notebook_shim==0.2.4
 numpy==2.0.0
-=======
-<<<<<<< HEAD
-certifi==2025.4.26
-charset-normalizer==3.4.2
-filelock==3.18.0
-fsspec==2025.3.2
-=======
-accelerate==1.7.0
-aiohappyeyeballs==2.6.1
-aiohttp==3.11.18
-aiosignal==1.3.2
-attrs==25.3.0
-certifi==2025.4.26
-charset-normalizer==3.4.2
-datasets==3.6.0
-dill==0.3.8
-filelock==3.18.0
-frozenlist==1.6.0
-fsspec==2025.3.0
->>>>>>> fc9eefc (add pretrain code)
-huggingface-hub==0.30.2
-idna==3.10
-Jinja2==3.1.6
-MarkupSafe==3.0.2
-mpmath==1.3.0
-<<<<<<< HEAD
-=======
-multidict==6.4.3
-multiprocess==0.70.16
->>>>>>> fc9eefc (add pretrain code)
-networkx==3.4.2
-numpy==2.2.5
->>>>>>> 11473a798f338efdf3f70d665a00810099f51174
 nvidia-cublas-cu12==12.1.3.1
 nvidia-cuda-cupti-cu12==12.1.105
 nvidia-cuda-nvrtc-cu12==12.1.105
@@ -114,7 +80,6 @@ nvidia-curand-cu12==10.3.2.106
 nvidia-cusolver-cu12==11.4.5.107
 nvidia-cusparse-cu12==12.1.0.106
 nvidia-nccl-cu12==2.20.5
-<<<<<<< HEAD
 nvidia-nvjitlink-cu12==12.5.82
 nvidia-nvtx-cu12==12.1.105
 overrides==7.7.0
@@ -156,12 +121,13 @@ sniffio==1.3.1
 soupsieve==2.5
 stack-data==0.6.3
 sympy==1.13.0
+tensorboardX==2.6.2.2
 terminado==0.18.1
 tinycss2==1.3.0
 tokenizers==0.21.1
-torch @ file:///workspace/torch-2.3.1%2Bcu121-cp312-cp312-linux_x86_64.whl#sha256=b3c586f4ab25e83efffccfb97079e91325329bc228166555c4bb93957753d4ea
+torch==torch-2.3.1
-torchaudio @ file:///workspace/torchaudio-2.3.1%2Bcu121-cp312-cp312-linux_x86_64.whl#sha256=b00773dde3aa30d9b1b98875650428fbfb69fe30b98ebb5ceffbd8fc61de8125
+torchaudio==-2.3.1
-torchvision @ file:///workspace/torchvision-0.18.1%2Bcu121-cp312-cp312-linux_x86_64.whl#sha256=ce8d5b992056f0640a39ef5734342e43ca4a801547de27fb8dbc3055d9345947
+torchvision==-0.18.1
 tornado==6.4.1
 tqdm==4.67.1
 traitlets==5.14.3
@@ -178,46 +144,3 @@ webencodings==0.5.1
 websocket-client==1.8.0
 xxhash==3.5.0
 yarl==1.20.0
-=======
-nvidia-nvjitlink-cu12==12.9.41
-nvidia-nvtx-cu12==12.1.105
-packaging==25.0
-<<<<<<< HEAD
-protobuf==3.20.0
-=======
-pandas==2.2.3
-propcache==0.3.1
-protobuf==3.20.0
-psutil==7.0.0
-pyarrow==20.0.0
-python-dateutil==2.9.0.post0
-pytz==2025.2
->>>>>>> fc9eefc (add pretrain code)
-PyYAML==6.0.2
-regex==2024.11.6
-requests==2.32.3
-safetensors==0.5.3
-sentencepiece==0.2.0
-<<<<<<< HEAD
-=======
-setuptools==78.1.1
-six==1.17.0
->>>>>>> fc9eefc (add pretrain code)
-sympy==1.14.0
-tokenizers==0.21.1
-torch==2.3.1
-tqdm==4.67.1
-transformers==4.51.3
-<<<<<<< HEAD
-triton==2.3.1
-typing_extensions==4.13.2
-urllib3==2.4.0
-=======
-typing_extensions==4.13.2
-tzdata==2025.2
-urllib3==2.4.0
-wheel==0.45.1
-xxhash==3.5.0
-yarl==1.20.0
->>>>>>> fc9eefc (add pretrain code)
->>>>>>> 11473a798f338efdf3f70d665a00810099f51174

setup.py

@@ -6,26 +6,15 @@ setup(
     packages=find_packages(include=["src_llm"]),  
     install_requires=[
         "torch==2.3.1",
-<<<<<<< HEAD
         "datasets==3.6.0",
-=======
-<<<<<<< HEAD
-=======
-        "datasets==3.6.0",
->>>>>>> fc9eefc (add pretrain code)
->>>>>>> 11473a798f338efdf3f70d665a00810099f51174
         "sentencepiece==0.2.0",
         "protobuf==3.20.0",
         "transformers==4.51.3",
         "tokenizers==0.21.1",
-<<<<<<< HEAD
+        "setuptools==70.3.0",
-        "setuptools",
-=======
-<<<<<<< HEAD
-=======
         "accelerate==1.7.0",
->>>>>>> fc9eefc (add pretrain code)
+        "tensorboardX==2.6.2.2",
->>>>>>> 11473a798f338efdf3f70d665a00810099f51174
+        "deepspeed==0.16.7"
         ],  
     author="IP127000",  
     author_email="hanluzhi@outlook.com", 

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 = None                     
-resume_option = True                    
+# resume_option = True                    
-# resume_option = "/mnt/han.luzhi/result/checkpoint-500"  
+# 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,
+    learning_rate=1e-4,
-    per_device_train_batch_size=26,
+    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")

src_pretrain/pretrain_transformers.py

@@ -11,7 +11,6 @@ from datasets import load_dataset
 from accelerate import Accelerator
 import os
 import logging
-
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -20,8 +19,9 @@ set_seed(42)
 
 model_name = "/mnt/han.luzhi/dolly_llm/weights_llm"  
 # resume_option = None                     
-resume_option = True                    
+# resume_option = True                    
-# resume_option = "/mnt/han.luzhi/result/checkpoint-500"  
+resume_option = "/mnt/han.luzhi/dolly_llm/result/checkpoint-59000"  
+# resume_option = "/mnt/han.luzhi/dolly_llm/result/final" 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
@@ -29,15 +29,8 @@ 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)
 
@@ -58,7 +51,7 @@ tokenized_datasets = dataset.map(
     preprocess_function,
     batched=True,
     remove_columns=['text'], 
-    num_proc=4  
+    num_proc=16  
 )
 
 data_collator = DataCollatorForLanguageModeling(
@@ -78,7 +71,7 @@ training_args = TrainingArguments(
     logging_steps=10,
     save_total_limit=2,
     fp16=True, 
-    dataloader_num_workers=4,
+    dataloader_num_workers=3,
     gradient_accumulation_steps=2, 
     report_to="tensorboard",  
     optim="adamw_torch",  

utils/dolly_pretrain_deepspeed.sh

@@ -1,2 +1,2 @@
 deepspeed --num_gpus=4 --master_port=12345 pretrain_transformers.py
-#   deepspeed --include="localhost:1,2,3" --master_port=12345 pretrain_transformers.py
+# CUDA_VISIBLE_DEVICES=1,2,3 deepspeed --include="localhost:1,2,3" --master_port=12345 src_pretrain/pretrain_deepspeed.py

utils/dolly_pretrain_torch.sh

@@ -1,3 +1,4 @@
 python3 -m torch.distributed.launch --nproc_per_node=4 --master_port=12345 src_pretrain/pretrain_transformers.py
 # torchrun --nproc_per_node=4 src_pretrain/pretrain_transformers.py
 # CUDA_VISIBLE_DEVICES=1,2,3 torchrun --nproc_per_node=3 src_pretrain/pretrain_transformers.py
+# CUDA_VISIBLE_DEVICES=1,2,3 python3 -m torch.distributed.launch --nproc_per_node=4 --master_port=12345 src_pretrain/pretrain_transformers.py