add jsonl support

.gitignore

@@ -1,5 +1,6 @@
 corpus/wikipedia-cn-20230720-filtered.json
 corpus/wikipedia.txt
+corpus/*.jsonl
 dolly_env_local
 llm_dolly.egg-info
 src_llm/__pycache__

src_pretrain/pretrain_deepspeed.py

@@ -11,6 +11,7 @@ from datasets import load_dataset
 from accelerate import Accelerator
 import os
 import logging
+import glob
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -21,7 +22,7 @@ set_seed(42)
 model_name = "/mnt/han.luzhi/dolly_llm/weights_llm"  
 resume_option = None                     
 # resume_option = True                    
-# resume_option = "/mnt/han.luzhi/dolly_llm/result/checkpoint-59000"  
+# resume_option = "/mnt/han.luzhi/dolly_llm/checkpoints_ds/checkpoint-500"  
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
@@ -35,7 +36,11 @@ else:
     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'})
+data_files = glob.glob("/mnt/han.luzhi/dolly_llm/corpus/*.jsonl")
+logger.info(f"找到训练语料: {data_files}")
+
+dataset = load_dataset('json', data_files=data_files, split='train')
+# dataset = load_dataset('text', data_files={'train': '/mnt/han.luzhi/dolly_llm/corpus/wikipedia.txt'})
 
 def preprocess_function(examples):
     return tokenizer(
@@ -49,7 +54,8 @@ def preprocess_function(examples):
 tokenized_datasets = dataset.map(
     preprocess_function,
     batched=True,
-    remove_columns=['text'], 
+    # remove_columns=['text'], 
+    remove_columns=dataset.column_names,
     num_proc=16  
 )
 
@@ -65,13 +71,13 @@ training_args = TrainingArguments(
     save_strategy="steps",
     save_steps=500,
     learning_rate=1e-4,
-    per_device_train_batch_size=32,
-    num_train_epochs=5,
+    per_device_train_batch_size=36,
+    num_train_epochs=2,
     logging_dir="/mnt/han.luzhi/dolly_llm/logs",
     logging_steps=20,
     save_total_limit=2,
     fp16=True,  
-    dataloader_num_workers=6,
+    dataloader_num_workers=3,
     gradient_accumulation_steps=2,
     report_to="tensorboard",
     optim="adamw_torch",
@@ -89,7 +95,7 @@ training_args = TrainingArguments(
 trainer = Trainer(
     model=model,
     args=training_args,
-    train_dataset=tokenized_datasets['train'],
+    train_dataset=tokenized_datasets,
     processing_class=tokenizer, 
     data_collator=data_collator,
 )
@@ -99,8 +105,11 @@ 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}")
 
+if resume_option !=None:
+    trainer.train(resume_from_checkpoint=resume_option)
+else:
     trainer.train()
-# trainer.train(resume_from_checkpoint=resume_option)
+
 
 if accelerator.is_main_process: 
     trainer.save_model("/mnt/han.luzhi/dolly_llm/result/final")