Resume sample (#3460)

---------

Co-authored-by: Jintao Huang <[email protected]>

Author: BC-A

swift/llm/argument/sampling_args.py

@@ -25,6 +25,7 @@ class SamplingArguments(BaseArguments):
     sampler_engine: Literal['pt', 'lmdeploy', 'vllm', 'no', 'client'] = 'pt'
     output_dir: str = 'sample_output'
     output_file: Optional[str] = None
+    resume: bool = False
     override_exist_file: bool = False
     num_return_sequences: int = 64
     num_sampling_per_gpu_batch_size: int = 1

swift/llm/sampling/sampling.py

@@ -4,6 +4,8 @@
 import time
 from typing import List, Union
 
+import json
+
 from swift.llm import SamplingArguments, SwiftPipeline, load_dataset
 from swift.utils import get_logger
 
@@ -49,28 +51,53 @@ def _get_dataset(self):
     def run(self):
         os.makedirs(self.args.output_dir, exist_ok=True)
         iter_file = os.path.join(self.args.output_dir, self.args.output_file)
+        resume_file = os.path.join(self.args.output_dir, self.args.output_file + '.resume')
         tmp_file = os.path.join(self.args.output_dir, self.args.output_file + '.tmp')
+        ckpt_state_file = os.path.join(self.args.output_dir, 'ckpt_state.json')
         if os.path.exists(iter_file) and not self.args.override_exist_file:
             return
-        if os.path.exists(tmp_file):
-            os.remove(tmp_file)
+
+        index_resume = -1
+        write_mode = 'w'
+        if self.args.resume:
+            write_mode = 'a'
+            if os.path.exists(resume_file):
+                shutil.copyfile(resume_file, tmp_file)
+
+            if os.path.exists(ckpt_state_file):
+                with open(ckpt_state_file, 'r') as ckpt_state:
+                    data = json.load(ckpt_state)
+                    index_resume = data.get('index', -1)
+                    logger.info(f'Loaded index_resume: {index_resume}')
+        else:
+            if os.path.exists(tmp_file):
+                os.remove(tmp_file)
+
         dataset = self._get_dataset()
         dataset_len = len(dataset)
         total_iters = int(dataset_len // self.args.num_sampling_per_gpu_batch_size)
+
         if self.args.num_sampling_per_gpu_batches is None or self.args.num_sampling_per_gpu_batches > total_iters:
             self.args.num_sampling_per_gpu_batches = total_iters
 
-        with open(tmp_file, 'w') as f:
+        with open(tmp_file, write_mode) as f:
             for _index in range(self.args.num_sampling_per_gpu_batches):
+                if _index <= index_resume:
+                    continue
                 logger.info(f' Sampling index:{_index}')
                 slices = dataset[self.args.num_sampling_per_gpu_batch_size
                                  * _index:self.args.num_sampling_per_gpu_batch_size * (_index + 1)]
                 slices = self.sampler.truncate_input(slices)
                 generated = self.sampler.do_sample(slices)
                 f.writelines(generated)
+                f.flush()
+                shutil.copy(tmp_file, resume_file)
+                with open(ckpt_state_file, 'w') as ckpt_state:
+                    json.dump({'index': _index}, ckpt_state)
+
         if os.path.exists(iter_file):
             shutil.move(iter_file, iter_file + '.' + str(int(time.time())))
-        shutil.move(tmp_file, iter_file)
+        shutil.move(resume_file, iter_file)
         logger.info(f'Sample file {iter_file} generated.')