attempt to fix unit tests

Signed-off-by: ashors1 <ashors@nvidia.com>

Author: ashors1

tests/unit/algorithms/test_sequence_packing_gradients.py

@@ -44,7 +44,10 @@ def test_sequence_packing_gradients(self):
         from nemo_rl.models.megatron.common import (
             forward_step_arbitrary_loss,
         )
-        from nemo_rl.models.megatron.data import _pack_sequences_for_megatron
+        from nemo_rl.models.megatron.data import (
+            _pack_sequences_for_megatron,
+            make_processed_microbatch_iterator,
+        )
 
         # Initialize process group
         torch.distributed.init_process_group(backend="nccl")
@@ -325,13 +328,27 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             MockMcoreState(),
             global_valid_seqs,
             global_valid_toks,
-            data_iterator=iter([packed_data_dict]),
+            data_iterator=make_processed_microbatch_iterator(
+                iter([packed_data_dict]),
+                cfg={
+                    "sequence_packing": {"enabled": True},
+                    "dynamic_batching": {"enabled": False},
+                    "megatron_cfg": {
+                        "tensor_model_parallel_size": 1,
+                        "sequence_parallel": False,
+                        "pipeline_model_parallel_size": 1,
+                        "context_parallel_size": cp_size,
+                    },
+
+
+                },
+                seq_length_key="input_lengths",
+                pad_individual_seqs_to_multiple_of=pad_to_multiple,
+                pad_full_seq_to=max_seq_len * batch_size if cp_size > 1 else None,
+            ),
             model=MockModel(),
             loss_fn=base_loss_fn,
             pack_sequences=True,
-            seq_length_key="input_lengths",
-            pad_individual_seqs_to_multiple_of=pad_to_multiple,
-            pad_full_seq_to=max_seq_len * batch_size if cp_size > 1 else None,
             cp_normalize=True,
         )
         loss, metrics = wrapped_loss_fn(output_tensor)