feat(conversion): add a failure exit mode to HF/Megatron comparison

examples/conversion/README.md

@@ -442,6 +442,21 @@ uv run python examples/conversion/compare_hf_and_megatron/compare.py \
   --enable_debug_hooks
 ```
 
+**Use the comparison as a validation gate:**
+```bash
+uv run python -m torch.distributed.run --nproc_per_node=2 examples/conversion/compare_hf_and_megatron/compare.py \
+  --hf_model_path Qwen/Qwen3-1.7B \
+  --prompt "Hello world" \
+  --tp 2 \
+  --fail-on-mismatch
+```
+
+With `--fail-on-mismatch`, next-token predictions must match within the HF vocabulary
+and cosine similarity must be at least 0.99. A failed comparison raises an error on
+every rank, including earlier pipeline stages. Maximum and mean absolute logit
+differences remain diagnostic metrics. Without this flag, the script only reports
+the comparison metrics.
+
 **Example Output:**
 ```
 Processing inputs - Prompt: 'Hello, how are you?', Image: None

examples/conversion/compare_hf_and_megatron/compare.py

@@ -932,6 +932,19 @@ def _broadcast_hf_results(hf_logits, hf_next_token, device):
     return hf_logits, hf_next_token
 
 
+def _check_comparison_result(comparison_passed: bool, device: torch.device) -> None:
+    """Raise on every rank if any comparison rank failed the correlation gate."""
+    if torch.distributed.is_initialized():
+        result = torch.tensor(int(comparison_passed), dtype=torch.int32, device=device)
+        torch.distributed.all_reduce(result, op=torch.distributed.ReduceOp.MIN)
+        comparison_passed = bool(result.item())
+    if not comparison_passed:
+        raise ValueError(
+            "HF/Megatron comparison failed: next tokens must match and cosine similarity "
+            f"must be at least {SIMILARITY_THRESHOLD}. See the comparison output for details."
+        )
+
+
 def compare_models_one_step(args) -> None:
     """Compare 1-step generation between HF and Megatron models with debugging.
 
@@ -1067,6 +1080,8 @@ def compare_models_one_step(args) -> None:
         hf_logits, hf_next_token = _broadcast_hf_results(hf_logits, hf_next_token, input_ids.device)
         print_rank_0("HF results broadcast complete.")
 
+    comparison_passed = True
+
     # Run Megatron model forward pass
     print_rank_0("=== RUNNING MEGATRON MODEL (1-STEP) ===")
     with torch.no_grad():
@@ -1175,6 +1190,7 @@ def compare_models_one_step(args) -> None:
                 print(f"Logits diff - max: {diff.max():.6f}, mean: {diff.mean():.6f}")
                 cosine_sim = torch.cosine_similarity(hf_logits.unsqueeze(0), megatron_logits_cmp.unsqueeze(0))
                 cos_val = cosine_sim.item()
+                comparison_passed = token_match and cos_val >= SIMILARITY_THRESHOLD
                 percent = cos_val * 100.0
                 status_emoji = "✅" if cos_val >= SIMILARITY_THRESHOLD else "❌"
                 limit_text = "within" if cos_val >= SIMILARITY_THRESHOLD else "outside"
@@ -1193,6 +1209,9 @@ def compare_models_one_step(args) -> None:
         if torch.distributed.is_initialized():
             torch.distributed.broadcast(megatron_next_token, get_last_rank())
 
+        if getattr(args, "fail_on_mismatch", False):
+            _check_comparison_result(comparison_passed, input_ids.device)
+
 
 def build_parser() -> argparse.ArgumentParser:
     """Build the CLI parser."""
@@ -1244,6 +1263,11 @@ def build_parser() -> argparse.ArgumentParser:
     parser.add_argument("--pp", type=int, default=1, help="Pipeline parallelism size")
     parser.add_argument("--ep", type=int, default=1, help="Expert parallelism size")
     parser.add_argument("--etp", type=int, default=1, help="Expert tensor parallelism size")
+    parser.add_argument(
+        "--fail-on-mismatch",
+        action="store_true",
+        help="Exit with an error on all ranks if next tokens differ or cosine similarity is below 0.99.",
+    )
     parser.add_argument(
         "--hf-device",
         default="cuda",
@@ -1285,7 +1309,8 @@ def build_parser() -> argparse.ArgumentParser:
     args = build_parser().parse_args()
 
     maybe_initialize_distributed()
-    compare_models_one_step(args)
-
-    if torch.distributed.is_initialized():
-        torch.distributed.destroy_process_group()
+    try:
+        compare_models_one_step(args)
+    finally:
+        if torch.distributed.is_initialized():
+            torch.distributed.destroy_process_group()

tests/unit_tests/test_compare_mask_handling.py

@@ -20,6 +20,8 @@
 
 import os
 import sys
+from collections.abc import Callable
+from datetime import timedelta
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
@@ -96,6 +98,133 @@
     sys.path.remove(_compare_dir)
 
 
+def _comparison_result_worker(rank: int, rendezvous_path: str) -> None:
+    """Exercise verdict propagation with real CPU collectives."""
+    torch.distributed.init_process_group(
+        "gloo",
+        init_method=f"file://{rendezvous_path}",
+        rank=rank,
+        world_size=2,
+        timeout=timedelta(seconds=30),
+    )
+    try:
+        compare._check_comparison_result(True, torch.device("cpu"))
+        for failing_rank in (0, 1):
+            with pytest.raises(ValueError, match="HF/Megatron comparison failed"):
+                compare._check_comparison_result(rank != failing_rank, torch.device("cpu"))
+    finally:
+        torch.distributed.destroy_process_group()
+
+
+@pytest.mark.unit
+class TestComparisonFailureHandling:
+    def test_failure_mode_is_opt_in(self):
+        arguments = ["--hf_model_path", "org/model", "--prompt", "hello"]
+        assert compare.build_parser().parse_args(arguments).fail_on_mismatch is False
+        assert compare.build_parser().parse_args([*arguments, "--fail-on-mismatch"]).fail_on_mismatch is True
+
+    @pytest.mark.parametrize("passed", [True, False])
+    def test_single_process_verdict(self, passed):
+        with patch.object(torch.distributed, "is_initialized", return_value=False):
+            if passed:
+                compare._check_comparison_result(True, torch.device("cpu"))
+            else:
+                with pytest.raises(ValueError, match="HF/Megatron comparison failed"):
+                    compare._check_comparison_result(False, torch.device("cpu"))
+
+    @pytest.mark.skipif(not torch.distributed.is_gloo_available(), reason="Gloo is not available")
+    def test_failure_on_either_rank_reaches_every_rank(self, tmp_path):
+        torch.multiprocessing.spawn(
+            _comparison_result_worker,
+            args=(str(tmp_path / "comparison-rendezvous"),),
+            nprocs=2,
+        )
+
+    @pytest.mark.parametrize("fail_on_mismatch", [False, True])
+    @pytest.mark.parametrize(
+        "hf_values, megatron_values, passes",
+        [
+            ([0, 0, 0, 0, 1], [0, 0, 0, 0, 1], True),
+            # High cosine similarity alone must not hide a different next token.
+            ([100, 100, 100, 100, 100.01], [100, 100, 100, 100.01, 100], False),
+            # A matching next token alone must not hide poor correlation.
+            ([0, 0, 0, 0, 1], [-5, -5, -5, -5, 1], False),
+            ([0, 0, 0, 0, 1], [0, 0, 0, 0, float("nan")], False),
+            # Padded Megatron vocabulary entries do not participate in the gate.
+            ([0, 0, 0, 0, 1], [0, 0, 0, 0, 1, 1000], True),
+        ],
+    )
+    def test_forward_comparison_verdict(self, monkeypatch, fail_on_mismatch, hf_values, megatron_values, passes):
+        args = compare.build_parser().parse_args(["--hf_model_path", "org/model", "--prompt", "hello"])
+        args.fail_on_mismatch = fail_on_mismatch
+        input_ids = torch.tensor([[1, 2]])
+        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: ([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))
+        monkeypatch.setattr(
+            compare,
+            "_run_hf_inference",
+            lambda *args, **kwargs: (hf_logits, hf_logits.argmax(), {}, [], (1, 2, len(hf_values))),
+        )
+        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:
     """Tests for attention_mask handling in compare.py Megatron and HF paths."""