fix unit test

Signed-off-by: Yuki Huang <yukih@nvidia.com>

Author: yuki-97

nemo_rl/algorithms/loss/__init__.py

@@ -45,7 +45,7 @@
     "NLLLossFn",
     "PreferenceLossDataDict",
     "PreferenceLossFn",
-    "SequencePackingLossWrapper",
     "prepare_loss_input",
+    "SequencePackingLossWrapper",
     "wrap_loss_fn_with_input_preparation",
 ]

tests/unit/models/automodel/test_automodel_train.py

@@ -24,6 +24,7 @@
 except ImportError:
     pytest.skip("nemo_automodel not available", allow_module_level=True)
 
+from nemo_rl.algorithms.loss.interfaces import LossInputType
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.models.automodel.data import (
     ProcessedInputs,
@@ -63,6 +64,7 @@ def mock_model():
 def mock_loss_fn():
     loss_fn = MagicMock()
     loss_fn.return_value = (torch.tensor(0.5), {"loss": 0.5})
+    loss_fn.input_type = LossInputType.LOGIT
     return loss_fn
 
 
@@ -310,10 +312,10 @@ def test_basic_loss_computation(
 
         # Verify loss function was called
         mock_loss_fn.assert_called_once()
-        call_args = mock_loss_fn.call_args[0]
-        assert torch.is_tensor(call_args[0])  # logits
-        assert call_args[2] == global_valid_seqs  # global_valid_seqs
-        assert call_args[3] == global_valid_toks  # global_valid_toks
+        call_kwargs = mock_loss_fn.call_args[1]
+        assert torch.is_tensor(call_kwargs["logits"])
+        assert call_kwargs["global_valid_seqs"] == global_valid_seqs
+        assert call_kwargs["global_valid_toks"] == global_valid_toks
 
     @patch("nemo_rl.models.automodel.train.SequencePackingLossWrapper")
     def test_loss_with_sequence_packing(
@@ -1896,10 +1898,12 @@ def forward(self, input_ids, **kwargs):
         )
 
         # Create loss function that returns requires_grad tensor
-        def loss_fn(logits, mb, global_valid_seqs, global_valid_toks):
+        def loss_fn(logits, data, global_valid_seqs, global_valid_toks):
             loss = logits.mean()
             return loss, {"loss": loss.item()}
 
+        loss_fn.input_type = LossInputType.LOGIT
+
         # Create loss post-processor
         loss_post_processor = LossPostProcessor(
             loss_fn=loss_fn,

tests/unit/models/generation/test_vllm_generation.py

@@ -894,13 +894,13 @@ async def run_hf_train_process(
 @pytest.mark.parametrize(
     ("async_engine", "cpu_offload", "vllm_precision", "enable_lora"),
     [
-        # (True, False, "bfloat16", False),
-        # (False, True, "bfloat16", False),
-        # (True, False, "fp8", False),
-        # (False, True, "fp8", False),
+        (True, False, "bfloat16", False),
+        (False, True, "bfloat16", False),
+        (True, False, "fp8", False),
+        (False, True, "fp8", False),
         # LoRA tests
         (False, False, "bfloat16", True),
-        # (True, False, "bfloat16", True),
+        (True, False, "bfloat16", True),
     ],
 )
 async def test_vllm_generation_with_hf_training_colocated(