Create dummy iters to satisy iter type len checks in core + update core commit

Dockerfile

@@ -46,7 +46,7 @@ WORKDIR /workspace/
 # Install Megatron-core
 RUN git clone https://github.com/NVIDIA/Megatron-LM.git && \
   cd Megatron-LM && \
-  git checkout 3db2063b1ff992a971ba18f7101eecc9c4e90f03 && \
+  git checkout 9f8bdeb4814ed61fbc9c7d5b39c7710e77b99754 && \
   pip install -e .
 
 WORKDIR /tmp/

Jenkinsfile

@@ -62,7 +62,7 @@ pipeline {
       steps {
         sh 'git clone https://github.com/NVIDIA/Megatron-LM.git && \
             cd Megatron-LM && \
-            git checkout 3db2063b1ff992a971ba18f7101eecc9c4e90f03 && \
+            git checkout 9f8bdeb4814ed61fbc9c7d5b39c7710e77b99754 && \
             pip install -e .'
       }
     }

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

@@ -114,7 +114,8 @@ def get_forward_output_only_func(self):
         Used for generate method only for now.
         """
 
-        def fwd_output_only_func(batch, model):
+        def fwd_output_only_func(dataloader_iter, model):
+            batch = next(dataloader_iter)
             extra_arg = {}
             (
                 tokens,

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

@@ -600,7 +600,8 @@ def loss_func(output_tensor):
         return fwd_output_and_loss_func
 
     def get_forward_output_only_func(self):
-        def fwd_output_only_func(batch, model):
+        def fwd_output_only_func(dataloader_iter, model):
+            batch = next(dataloader_iter)
             extra_arg = {}
             if len(batch) == 3:
                 batch = [x.cuda() for x in batch]

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

@@ -650,7 +650,8 @@ def get_forward_output_only_func(self):
         Used for generate method only for now.
         """
 
-        def fwd_output_only_func(batch, model):
+        def fwd_output_only_func(dataloader_iter, model):
+            batch = next(dataloader_iter)
             extra_arg = {}
             (
                 tokens,

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

@@ -513,7 +513,8 @@ def get_forward_output_only_func(self):
         Used for generate method only.
         """
 
-        def fwd_output_only_func(batch, model):
+        def fwd_output_only_func(dataloader_iter, model):
+            batch = next(dataloader_iter)
             extra_arg = {}
             (
                 tokens,

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

@@ -215,7 +215,8 @@ def get_forward_output_only_func(self):
         Used for generate method only for now.
         """
 
-        def fwd_output_only_func(batch, model):
+        def fwd_output_only_func(dataloader_iter, model):
+            batch = next(dataloader_iter)
             extra_arg = {}
             (
                 tokens,

nemo/collections/nlp/modules/common/text_generation_strategy.py

@@ -56,7 +56,7 @@ def forward_step(self, batch, tensor_shape):
 
         output_tensor = fwd_bwd_function(
             forward_step_func=self.model.get_forward_output_only_func(),
-            data_iterator=batch,
+            data_iterator=iter([batch,]),
             model=[self.forward_model],
             num_microbatches=get_num_microbatches(),
             forward_only=True,