prefetch num microbatches

Signed-off-by: eharper <[email protected]>

nemo/collections/nlp/models/language_modeling/megatron_base_model.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import gc
+import itertools
 import os
 import re
 from dataclasses import fields
@@ -783,3 +784,20 @@ def build_model_parallel_config(self):
             )
 
         return model_parallel_config
+
+    def _prefetch(self, iterator):
+        """Checks if the iterator still has elements to return.
+        Used in models using dataloader_iter to prefetch the next batch before fwd_bwd func
+        is called to avoid PP rank 2 from wait indefinitely to get outpits from PP 1
+        """
+        elements = []
+        num_microbatches = get_num_microbatches()
+        for _ in range(num_microbatches):
+            try:
+                element = next(iterator)
+                elements.append(element)
+            except StopIteration:
+                return iterator, True
+
+        # return a new iterator with the prefetched element reinserted at the front
+        return itertools.chain(elements, iterator), False

nemo/collections/nlp/models/nlp_model.py

@@ -14,7 +14,6 @@
 
 import copy
 import hashlib
-import itertools
 import json
 import os
 from typing import Any, Mapping, Optional
@@ -135,18 +134,6 @@ def __init__(self, cfg: DictConfig, trainer: Trainer = None, no_lm_init=False):
             # register encoder config
             self.register_bert_model()
 
-    def _prefetch(self, iterator):
-        """Checks if the iterator still has elements to return.
-        Used in models using dataloader_iter to prefetch the next batch before fwd_bwd func
-        is called to avoid PP rank 2 from wait indefinitely to get outpits from PP 1
-        """
-        try:
-            element = next(iterator)
-        except StopIteration:
-            return iterator, True
-
-        # return a new iterator with the prefetched element reinserted at the front
-        return itertools.chain([element], iterator), False
 
     def register_artifact(
         self, config_path: str, src: str, verify_src_exists: bool = False,