Load data at the correct position when resuming from a checkpoint

vince62s · Thai Chau Truong · commit d1913926cb6b · 2024-03-31T16:23:29.000-04:00
diff --git a/onmt/inputters/dynamic_iterator.py b/onmt/inputters/dynamic_iterator.py
@@ -129,6 +129,7 @@ def __init__(
         batch_type,
         batch_size,
         batch_size_multiple,
+        resume_corpora_info={},
         data_type="text",
         bucket_size=2048,
         bucket_size_init=-1,
@@ -144,6 +145,7 @@ def __init__(
         self.transforms = transforms
         self.vocabs = vocabs
         self.corpora_info = corpora_info
+        self.resume_corpora_info = resume_corpora_info
         self.task = task
         self.init_iterators = False
         self.batch_size = batch_size
@@ -171,7 +173,17 @@ def __init__(
 
     @classmethod
     def from_opt(
-        cls, corpora, transforms, vocabs, opt, task, copy, device, stride=1, offset=0
+        cls,
+        corpora,
+        transforms,
+        vocabs,
+        opt,
+        task,
+        copy,
+        device,
+        resume_corpora_info={},
+        stride=1,
+        offset=0,
     ):
         """Initilize `DynamicDatasetIter` with options parsed from `opt`."""
         corpora_info = {}
@@ -206,6 +218,7 @@ def from_opt(
             opt.batch_type,
             batch_size,
             batch_size_multiple,
+            resume_corpora_info=resume_corpora_info,
             data_type=opt.data_type,
             bucket_size=bucket_size,
             bucket_size_init=bucket_size_init,
@@ -388,6 +401,7 @@ def build_dynamic_dataset_iter(
     vocabs,
     copy=False,
     task=CorpusTask.TRAIN,
+    resume_corpora_info={},
     stride=1,
     offset=0,
     src=None,
@@ -412,7 +426,14 @@ def build_dynamic_dataset_iter(
     advance to avoid the GPU waiting during the refilling of the bucket.
     """
     transforms = make_transforms(opt, transforms_cls, vocabs)
-    corpora = get_corpora(opt, task, src=src, tgt=tgt, align=align)
+    corpora = get_corpora(
+        opt,
+        task,
+        src=src,
+        tgt=tgt,
+        align=align,
+        resume_corpora_info=resume_corpora_info,
+    )
     if corpora is None:
         assert task != CorpusTask.TRAIN, "only valid corpus is ignorable."
         return None
@@ -442,6 +463,7 @@ def build_dynamic_dataset_iter(
             vocabs,
             opt,
             task,
+            resume_corpora_info=resume_corpora_info,
             copy=copy,
             stride=stride,
             offset=offset,
diff --git a/onmt/inputters/text_corpus.py b/onmt/inputters/text_corpus.py
@@ -99,7 +99,14 @@ class ParallelCorpus(object):
     """A parallel corpus file pair that can be loaded to iterate."""
 
     def __init__(
-        self, name, src, tgt, align=None, n_src_feats=0, src_feats_defaults=None
+        self,
+        name,
+        src,
+        tgt,
+        align=None,
+        n_src_feats=0,
+        src_feats_defaults=None,
+        line_number_to_resume=0,
     ):
         """Initialize src & tgt side file path."""
         self.id = name
@@ -108,6 +115,12 @@ def __init__(
         self.align = align
         self.n_src_feats = n_src_feats
         self.src_feats_defaults = src_feats_defaults
+        self.line_number_to_resume = line_number_to_resume
+        self.can_read_file = False
+
+    def activate_reading_mode(self, line_number):
+        if line_number >= self.line_number_to_resume:
+            self.can_read_file = True
 
     def load(self, offset=0, stride=1):
         """
@@ -116,7 +129,7 @@ def load(self, offset=0, stride=1):
         `stride` example, starting from `offset`.
         """
 
-        def make_ex(sline, tline, align):
+        def make_ex(sline, tline, align, line_number):
             sline, sfeats = parse_features(
                 sline,
                 n_feats=self.n_src_feats,
@@ -131,6 +144,7 @@ def make_ex(sline, tline, align):
                 "tgt": tline,
                 "src_original": sline,
                 "tgt_original": tline,
+                "cid_line_number": line_number,
             }
             if align is not None:
                 example["align"] = align
@@ -145,19 +159,25 @@ def make_ex(sline, tline, align):
             for i, (sline, tline, align) in enumerate(
                 itertools.zip_longest(fs, ft, fa)
             ):
+                self.activate_reading_mode(line_number=i)
+                if not self.can_read_file:
+                    continue
                 if (i // stride) % stride == offset:
-                    yield make_ex(sline, tline, align)
+                    yield make_ex(sline, tline, align, i)
         else:
             with exfile_open(self.src, mode="rb") as fs, exfile_open(
                 self.tgt, mode="rb"
             ) as ft, exfile_open(self.align, mode="rb") as fa:
                 for i, (sline, tline, align) in enumerate(zip(fs, ft, fa)):
+                    self.activate_reading_mode(line_number=i)
+                    if not self.can_read_file:
+                        continue
                     if (i // stride) % stride == offset:
                         if tline is not None:
                             tline = tline.decode("utf-8")
                         if align is not None:
                             align = align.decode("utf-8")
-                        yield make_ex(sline.decode("utf-8"), tline, align)
+                        yield make_ex(sline.decode("utf-8"), tline, align, i)
 
     def __str__(self):
         cls_name = type(self).__name__
@@ -169,19 +189,25 @@ def __str__(self):
         )
 
 
-def get_corpora(opts, task=CorpusTask.TRAIN, src=None, tgt=None, align=None):
+def get_corpora(
+    opts, task=CorpusTask.TRAIN, src=None, tgt=None, align=None, resume_corpora_info={}
+):
     corpora_dict = {}
     if task == CorpusTask.TRAIN:
         for corpus_id, corpus_dict in opts.data.items():
             if corpus_id != CorpusName.VALID:
                 if corpus_dict.get("path_txt", None) is None:
+                    resume_line = 0
+                    if corpus_id in resume_corpora_info:
+                        resume_line = resume_corpora_info[corpus_id]
                     corpora_dict[corpus_id] = ParallelCorpus(
                         corpus_id,
                         corpus_dict["path_src"],
                         corpus_dict["path_tgt"],
                         corpus_dict["path_align"],
                         n_src_feats=opts.n_src_feats,
                         src_feats_defaults=opts.src_feats_defaults,
+                        line_number_to_resume=resume_line,
                     )
                 else:
                     corpora_dict[corpus_id] = BlockwiseCorpus(
@@ -244,8 +270,6 @@ def _process(self, stream):
                 example["src_feats"] = [
                     feat.strip().split(" ") for feat in example["src_feats"]
                 ]
-            line_number = i * self.stride + self.offset
-            example["cid_line_number"] = line_number
             example["cid"] = self.cid
             if "align" in example:
                 example["align"] = example["align"].strip().split(" ")
@@ -258,6 +282,7 @@ def _process(self, stream):
                     or ("align" in example and example["align"] == 0)
                 ):
                     # empty example: skip
+                    line_number = example["cid_line_number"]
                     empty_msg = f"Empty line  in {self.cid}#{line_number}."
                     if self.skip_empty_level == "error":
                         raise IOError(empty_msg)
@@ -282,7 +307,12 @@ def __iter__(self):
 
 
 def build_corpora_iters(
-    corpora, transforms, corpora_info, skip_empty_level="warning", stride=1, offset=0
+    corpora,
+    transforms,
+    corpora_info,
+    skip_empty_level="warning",
+    stride=1,
+    offset=0,
 ):
     """Return `ParallelCorpusIterator` for all corpora defined in opts."""
     corpora_iters = dict()
diff --git a/onmt/models/model_saver.py b/onmt/models/model_saver.py
@@ -1,13 +1,17 @@
 import os
 import torch
 import re
+import subprocess
 from collections import deque
+import onmt.utils
 from onmt.utils.logging import logger
 from onmt.inputters.inputter import vocabs_to_dict
 from onmt.modules.lora import lora_state_dict
 
 
-def build_model_saver(model_opt, opt, model, vocabs, optim, device_id):
+def build_model_saver(
+    model_opt, opt, model, vocabs, optim, resume_corpora_info, device_id
+):
     # _check_save_model_path
     save_model_path = os.path.abspath(opt.save_model)
     os.makedirs(os.path.dirname(save_model_path), exist_ok=True)
@@ -20,6 +24,7 @@ def build_model_saver(model_opt, opt, model, vocabs, optim, device_id):
         optim,
         opt.keep_checkpoint,
         opt.save_format,
+        resume_corpora_info,
         device_id,
     )
     return model_saver
@@ -81,6 +86,65 @@ def fix_key(s):
     return checkpoint
 
 
+def load_corpora_info(opts, checkpoint):
+    message_resume_from_beginning = (
+        "The training will resume from the beginning of each corpus."
+    )
+    # Check if resume_from_corpora is True
+    if not opts.resume_from_corpora:
+        logger.info(
+            "No resume from corpora is specified. " + message_resume_from_beginning
+        )
+        return {}
+
+    # Check if the corpus list from the last training
+    # and in the new training are identical.
+    checkpoint_corpora = checkpoint.get("corpus_info", None)
+    if checkpoint_corpora is None:
+        logger.info(
+            "Incoherent info: Some corpora in the last training "
+            + "and in the new list do not match. "
+            + message_resume_from_beginning
+        )
+        return {}
+
+    checkpoint_corpus_names = [name for name in checkpoint_corpora]
+    new_corpus_names = [name for name in opts.data]
+    if set(checkpoint_corpus_names) != set(new_corpus_names):
+        logger.info(
+            "Incoherent info: Some corpora in the last training "
+            + "and in the new list do not match. "
+            + message_resume_from_beginning
+        )
+        return {}
+
+    # For each corpus, check if the last line number to resume
+    # is smaller than or equal to the number of text lines.
+    message_incoherent_line_number = (
+        "Incoherent info: text line numbers "
+        + "to resume in some corpora exceed their total numbers of lines. "
+        + message_resume_from_beginning
+    )
+    for c_name in checkpoint_corpora:
+        number_of_text_lines = int(
+            subprocess.getoutput(
+                "wc -l " + opts.data[c_name]["path_src"] + " | awk '{print $1}'"
+            )
+        )
+        if checkpoint_corpora[c_name] > number_of_text_lines - 1:
+            logger.info(message_incoherent_line_number)
+            return {}
+
+        # To set the text lines to resume, we increase all text lines by 1
+        # (and return to the beginning if the end is reached)
+        checkpoint_corpora[c_name] = (
+            checkpoint_corpora[c_name] + 1
+        ) % number_of_text_lines
+
+    logger.info("The training will resume from the saved text line in each corpus.")
+    return checkpoint_corpora
+
+
 class ModelSaverBase(object):
     """Base class for model saving operations
 
@@ -98,6 +162,7 @@ def __init__(
         optim,
         keep_checkpoint=-1,
         save_format="pytorch",
+        resume_corpora_info={},
         device_id=0,
     ):
         self.base_path = base_path
@@ -108,14 +173,35 @@ def __init__(
         self.last_saved_step = None
         self.keep_checkpoint = keep_checkpoint
         self.save_format = save_format
+        self.corpus_info = resume_corpora_info
         self.device_id = device_id
 
         if keep_checkpoint > 0:
             self.checkpoint_queue = deque([], maxlen=keep_checkpoint)
             if save_format == "safetensors":
                 self.model_queue = deque([], maxlen=keep_checkpoint)
 
-    def save(self, step, moving_average=None):
+    def update_corpus_info_from_batches(self, batches, distributed=False):
+        # Update the last text line of each corpus
+        if batches is not None:
+            # Gather corpus line numbers to save to checkpoints
+            batch_cids = sum([batch["cid"] for batch in batches], [])
+            batch_cid_line_numbers = sum(
+                [batch["cid_line_number"] for batch in batches], []
+            )
+            if distributed:
+                batch_cids = sum(onmt.utils.distributed.all_gather_list(batch_cids), [])
+                batch_cid_line_numbers = sum(
+                    onmt.utils.distributed.all_gather_list(batch_cid_line_numbers), []
+                )
+            # Save the last processed line number of each corpus
+            new_corpus_info = {
+                c_name: cid_line_number
+                for c_name, cid_line_number in zip(batch_cids, batch_cid_line_numbers)
+            }
+            self.corpus_info = {**self.corpus_info, **new_corpus_info}
+
+    def save(self, step, moving_average=None, batches=None, distributed=False):
         """Main entry point for model saver
 
         It wraps the `_save` method with checks and apply `keep_checkpoint`
@@ -266,6 +352,7 @@ def _save(self, step, model):
             "vocab": vocabs_to_dict(self.vocabs),
             "opt": self.model_opt,
             "optim": self.optim.state_dict(),
+            "corpus_info": self.corpus_info,
         }
         if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
             logger.info("Saving checkpoint %s_step_%d.pt" % (self.base_path, step))
@@ -355,6 +442,7 @@ def _st_save(self, step, model):
             "vocab": vocabs_to_dict(self.vocabs),
             "opt": self.model_opt,
             "optim": self.optim.state_dict(),
+            "corpus_info": self.corpus_info,
         }
 
         if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
diff --git a/onmt/opts.py b/onmt/opts.py
@@ -1263,6 +1263,13 @@ def _add_train_general_opts(parser):
         help="If training from a checkpoint then this is the "
         "path to the pretrained model's state_dict.",
     )
+    group.add(
+        "--resume_from_corpora",
+        "-resume_from_corpora",
+        action="store_true",
+        help="If training from a checkpoint and this is set to True "
+        " then the data generator will resume from the last line of each corpora.",
+    )
     group.add(
         "--reset_optim",
         "-reset_optim",
diff --git a/onmt/train_single.py b/onmt/train_single.py
diff --git a/onmt/trainer.py b/onmt/trainer.py
