test(conversion): exercise scheduler path in verdict tests

Mock scheduler selection instead of bypassing the Megatron forward helper.
Drive the real inference wrapper, iterator, and forward-step callback with
a single-microbatch test scheduler, then assert the scheduling and model
input contracts. Preserve the existing offline AutoConfig fixture.

Assisted-by: OpenAI Codex
Signed-off-by: Yuchen Wang <eang333cms@gmail.com>

Author: yuchenwang3

tests/unit_tests/test_compare_mask_handling.py

@@ -20,6 +20,7 @@
 
 import os
 import sys
+from collections.abc import Callable
 from datetime import timedelta
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
@@ -160,14 +161,31 @@ def test_forward_comparison_verdict(self, monkeypatch, fail_on_mismatch, hf_valu
         hf_logits = torch.tensor(hf_values, dtype=torch.float32)
         megatron_output = torch.tensor(megatron_values, dtype=torch.float32).view(1, 1, -1)
         tokenizer = SimpleNamespace(decode=lambda ids: str(ids))
+        megatron_model = MagicMock(return_value=megatron_output)
+
+        def forward_backward(
+            *,
+            forward_step_func: Callable[..., tuple[torch.Tensor, Callable[..., torch.Tensor]]],
+            data_iterator: SingleBatchIterator,
+            model: list[MagicMock],
+            **kwargs: object,
+        ) -> list[torch.Tensor]:
+            assert not torch.is_grad_enabled()
+            output, collect_data = forward_step_func(data_iterator, model[0])
+            return [collect_data(output, non_loss_data=True)]
+
+        # Scheduler selection requires real model-parallel groups. Keep the
+        # inference wrapper, iterator, and forward step on their production path.
+        scheduler = MagicMock(side_effect=forward_backward)
+        get_scheduler = MagicMock(return_value=scheduler)
 
         monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
         monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None)
         monkeypatch.setattr(torch.Tensor, "cuda", lambda tensor: tensor)
         monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False)
         monkeypatch.setattr(compare, "is_vision_language_model", lambda *args: False)
         monkeypatch.setattr(compare.AutoConfig, "from_pretrained", lambda *args, **kwargs: SimpleNamespace())
-        monkeypatch.setattr(compare, "_load_megatron_model", lambda args: ([MagicMock()], MagicMock()))
+        monkeypatch.setattr(compare, "_load_megatron_model", lambda args: ([megatron_model], MagicMock()))
         monkeypatch.setattr(compare, "_load_hf_model", lambda *args: MagicMock())
         monkeypatch.setattr(compare, "_setup_tokenizer_and_processor", lambda *args: (tokenizer, None))
         monkeypatch.setattr(compare, "process_inputs", lambda *args: (input_ids, None, None, None, None))
@@ -176,14 +194,36 @@ def test_forward_comparison_verdict(self, monkeypatch, fail_on_mismatch, hf_valu
             "_run_hf_inference",
             lambda *args, **kwargs: (hf_logits, hf_logits.argmax(), {}, [], (1, 2, len(hf_values))),
         )
-        monkeypatch.setattr(compare, "_run_megatron_forward", lambda *args, **kwargs: megatron_output)
+        monkeypatch.setattr(compare, "get_forward_backward_func", get_scheduler)
 
         if fail_on_mismatch and not passes:
             with pytest.raises(ValueError, match="HF/Megatron comparison failed"):
                 compare.compare_models_one_step(args)
         else:
             compare.compare_models_one_step(args)
 
+        get_scheduler.assert_called_once_with()
+        iterator = scheduler.call_args.kwargs["data_iterator"]
+        assert isinstance(iterator, SingleBatchIterator)
+        scheduler.assert_called_once_with(
+            forward_step_func=inference_forward_step,
+            data_iterator=iterator,
+            model=[megatron_model],
+            num_microbatches=1,
+            forward_only=True,
+            seq_length=input_ids.size(1),
+            micro_batch_size=1,
+            collect_non_loss_data=True,
+        )
+        megatron_model.assert_called_once()
+        forward_args = megatron_model.call_args.kwargs
+        assert forward_args["input_ids"] is input_ids
+        torch.testing.assert_close(forward_args["position_ids"], torch.tensor([[0, 1]]))
+        assert forward_args["attention_mask"] is None
+        assert forward_args["inference_context"] is iterator.batch["inference_context"]
+        assert forward_args["inference_context"] is not None
+        assert forward_args["runtime_gather_output"] is True
+
 
 @pytest.mark.unit
 class TestCompareMaskHandling: