lint

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

Author: ananthsub

nemo_rl/models/megatron/data.py

@@ -342,7 +342,6 @@ def process_global_batch(
     }
 
 
-
 def _pack_sequences_for_megatron(
     input_ids: torch.Tensor,
     seq_lengths: torch.Tensor,

nemo_rl/models/megatron/pipeline_parallel.py

@@ -27,6 +27,7 @@
 
 def broadcast_obj_from_pp_rank(obj: Any) -> Any:
     """Broadcast an object across pipeline parallel ranks.
+
     This utility function handles broadcasting an object from the rank that owns it
     to all other pipeline parallel ranks. If only one rank has the object (non-None),
     it will be broadcast to all other ranks.

tests/unit/algorithms/test_sequence_packing_gradients.py

@@ -42,14 +42,14 @@ def __init__(self, cp_size):
 
     def test_sequence_packing_gradients(self):
         from nemo_rl.distributed.model_utils import _get_tokens_on_this_cp_rank
-        from nemo_rl.models.megatron.train import (
-            forward_with_post_processing_fn,
-            LossPostProcessor,
-        )
         from nemo_rl.models.megatron.data import (
             _pack_sequences_for_megatron,
             make_processed_microbatch_iterator,
         )
+        from nemo_rl.models.megatron.train import (
+            LossPostProcessor,
+            forward_with_post_processing_fn,
+        )
 
         # Initialize process group
         torch.distributed.init_process_group(backend="nccl")