lint

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

Author: ashors1

nemo_rl/models/megatron/common.py

@@ -29,6 +29,7 @@
     get_moe_layer_wise_logging_tracker,
     reduce_aux_losses_tracker_across_ranks,
 )
+
 from nemo_rl.algorithms.loss_functions import LossFunction, SequencePackingLossWrapper
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 
@@ -40,6 +41,7 @@ def _round_up_to_multiple(value: int, multiple: int) -> int:
         else value
     )
 
+
 def forward_step_arbitrary_loss(
     state: GlobalState,
     global_valid_seqs: torch.Tensor,

nemo_rl/models/megatron/data.py

@@ -16,17 +16,17 @@
 from typing import Any, Iterator, Optional, Tuple
 
 import torch
-
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.parallel_state import (
     get_context_parallel_rank,
     get_context_parallel_world_size,
 )
 from megatron.training.utils import get_ltor_masks_and_position_ids
-from nemo_rl.models.megatron.common import _round_up_to_multiple
+
 from nemo_rl.algorithms.interfaces import LossFunction, LossType
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.distributed.model_utils import _get_tokens_on_this_cp_rank
+from nemo_rl.models.megatron.common import _round_up_to_multiple
 
 
 @dataclass
@@ -45,6 +45,7 @@ class ProcessedMicrobatch:
         packed_seq_params: PackedSeqParams for sequence packing (None if not packing)
         cu_seqlens_padded: Padded cumulative sequence lengths (None if not packing)
     """
+
     data_dict: BatchedDataDict[Any]
     input_ids: torch.Tensor
     input_ids_cp_sharded: torch.Tensor
@@ -192,6 +193,7 @@ def get_microbatch_iterator(
         padded_seq_length,
     )
 
+
 def process_microbatch(
     data_dict: BatchedDataDict[Any],
     seq_length_key: Optional[str] = None,
@@ -200,7 +202,7 @@ def process_microbatch(
     pad_full_seq_to: Optional[int] = None,
     pack_sequences: bool = False,
 ):
-    #with straggler_timer(bdata=True):
+    # with straggler_timer(bdata=True):
     input_ids = data_dict["input_ids"]
     attention_mask = None
     position_ids = None
@@ -217,9 +219,7 @@ def process_microbatch(
         assert seq_length_key is not None, (
             "seq_length_key must be provided for packed sequences"
         )
-        assert seq_length_key in data_dict, (
-            f"{seq_length_key} not found in data_dict"
-        )
+        assert seq_length_key in data_dict, f"{seq_length_key} not found in data_dict"
 
         # Get sequence lengths and context parallel size
         seq_lengths = data_dict[seq_length_key]
@@ -240,7 +240,7 @@ def process_microbatch(
             cp_rank=get_context_parallel_rank(),
             cp_size=get_context_parallel_world_size(),
         )
-    
+
         # For packed sequences, position_ids and attention_mask are typically None
         # The PackedSeqParams handles all necessary sequence information
         position_ids = None
@@ -265,6 +265,7 @@ def process_microbatch(
         cu_seqlens_padded,
     )
 
+
 def process_global_batch(
     data: BatchedDataDict[Any],
     batch_idx: int,
@@ -301,6 +302,7 @@ def process_global_batch(
         global_valid_toks,
     )
 
+
 def _pack_sequences_for_megatron(
     input_ids: torch.Tensor,
     seq_lengths: torch.Tensor,
@@ -605,6 +607,7 @@ def _unpack_sequences_from_megatron(
 
     return unpacked_output
 
+
 def check_sequence_dim(data: BatchedDataDict[Any]):
     # dim 1 is always assumed to be the sequence dim, sanity check this here
     sequence_dim = 1
@@ -614,4 +617,4 @@ def check_sequence_dim(data: BatchedDataDict[Any]):
             assert v.shape[sequence_dim] == seq_dim_size, (
                 f"Dim 1 must be the sequence dim, expected dim 1={seq_dim_size} but got shape {v.shape}"
             )
-    return sequence_dim, seq_dim_size
+    return sequence_dim, seq_dim_size

nemo_rl/models/policy/workers/megatron_policy_worker.py

@@ -113,11 +113,11 @@
     forward_step_arbitrary_loss,
     get_moe_metrics,
 )
+from nemo_rl.models.megatron.community_import import import_model_from_hf_name
 from nemo_rl.models.megatron.data import (
     get_microbatch_iterator,
     process_global_batch,
 )
-from nemo_rl.models.megatron.community_import import import_model_from_hf_name
 from nemo_rl.models.policy import PolicyConfig
 from nemo_rl.models.policy.interfaces import (
     ColocatablePolicyInterface,

tests/unit/algorithms/test_sequence_packing_gradients.py

@@ -339,8 +339,6 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                         "pipeline_model_parallel_size": 1,
                         "context_parallel_size": cp_size,
                     },
-
-
                 },
                 seq_length_key="input_lengths",
                 pad_individual_seqs_to_multiple_of=pad_to_multiple,

tests/unit/models/megatron/test_megatron_data.py

@@ -73,6 +73,7 @@ def test_processed_microbatch_fields(self):
         assert microbatch.packed_seq_params == mock_packed_seq_params
         assert torch.equal(microbatch.cu_seqlens_padded, mock_cu_seqlens_padded)
 
+
 class TestCheckSequenceDim:
     """Tests for check_sequence_dim function."""
 
@@ -154,7 +155,9 @@ def test_process_microbatch_no_packing(self, mock_get_masks):
 
         # Create test data
         data_dict = MagicMock()
-        input_ids = torch.tensor([[1, 2, 3, 4, 5, 0, 0, 0, 0, 0], [6, 7, 8, 9, 10, 11, 12, 0, 0, 0]])
+        input_ids = torch.tensor(
+            [[1, 2, 3, 4, 5, 0, 0, 0, 0, 0], [6, 7, 8, 9, 10, 11, 12, 0, 0, 0]]
+        )
         data_dict.__getitem__ = MagicMock(return_value=input_ids)
 
         (
@@ -178,7 +181,9 @@ def test_process_microbatch_no_packing(self, mock_get_masks):
         mock_get_masks.assert_called_once()
 
     @patch("nemo_rl.models.megatron.data.get_context_parallel_rank", return_value=0)
-    @patch("nemo_rl.models.megatron.data.get_context_parallel_world_size", return_value=1)
+    @patch(
+        "nemo_rl.models.megatron.data.get_context_parallel_world_size", return_value=1
+    )
     @patch("nemo_rl.models.megatron.data._pack_sequences_for_megatron")
     def test_process_microbatch_with_packing(
         self, mock_pack, mock_cp_world, mock_cp_rank
@@ -226,7 +231,7 @@ def test_process_microbatch_with_packing(
         assert attention_mask is None
         assert position_ids is None
         assert cu_seqlens_padded is not None
-  
+
         # Verify pack was called
         mock_pack.assert_called_once()
 
@@ -323,6 +328,7 @@ def test_process_global_batch_requires_sample_mask(self):
 
         assert "sample_mask must be present" in str(exc_info.value)
 
+
 class TestGetMicrobatchIterator:
     """Tests for get_microbatch_iterator function."""
 
@@ -383,8 +389,13 @@ def test_get_microbatch_iterator_sequence_packing(
         mock_make_iterator.return_value = mock_iterator
 
         mock_data = MagicMock()
-        mock_data.make_microbatch_iterator_for_packable_sequences.return_value = iter([])
-        mock_data.get_microbatch_iterator_for_packable_sequences_len.return_value = (10, 512)
+        mock_data.make_microbatch_iterator_for_packable_sequences.return_value = iter(
+            []
+        )
+        mock_data.get_microbatch_iterator_for_packable_sequences_len.return_value = (
+            10,
+            512,
+        )
 
         cfg = {
             "dynamic_batching": {"enabled": False},
@@ -473,8 +484,13 @@ def test_get_microbatch_iterator_auto_detects_seq_length_key(
         mock_make_iterator.return_value = mock_iterator
 
         mock_data = MagicMock()
-        mock_data.make_microbatch_iterator_for_packable_sequences.return_value = iter([])
-        mock_data.get_microbatch_iterator_for_packable_sequences_len.return_value = (5, 256)
+        mock_data.make_microbatch_iterator_for_packable_sequences.return_value = iter(
+            []
+        )
+        mock_data.get_microbatch_iterator_for_packable_sequences_len.return_value = (
+            5,
+            256,
+        )
 
         cfg = {
             "dynamic_batching": {"enabled": False},
@@ -1677,4 +1693,3 @@ def test_get_pack_sequence_parameters_for_megatron(get_pack_sequence_parameters_
     # Check that all workers succeeded
     for i, result in enumerate(results):
         assert result["success"], f"Worker {i} failed: {result['error']}"
-