replay_base: direct-evidence logs for R3 wrapper + overlap check

The previous --enable-r3-correctness-check flag turned on the overlap
check but produced no log output unless an actual mismatch happened,
making it impossible to distinguish "check passed" from "check never
ran." Add two unconditional logs gated on enable_check_replay_result:

1. get_topk_fn / new_topk_fn replay_forward + replay_backward branches:
   log when the wrapper actually returns replay indices rather than
   falling through to old_topk_fn. Direct evidence megatron's MoE
   forward used the rollout indices (vs recomputing them).

2. check_replay_result: log n_tokens and mismatch_count on every call,
   including the mismatch_count==0 case (which previously returned
   silently). Direct evidence the check ran, plus the actual overlap
   number for cross-step / cross-rank comparison.

Both logs gated on enable_check_replay_result so production training
runs (which leave it False) stay quiet. Adds no overhead when off.

Intended to make the LLM360/RL360 R3 regression E2E able to assert
directly that R3 worked end-to-end, rather than inferring from absence
of failure messages.

Author: DavidBellamy

miles/utils/replay_base.py

@@ -151,10 +151,32 @@ def new_topk_fn(scores, topk, *args, **kwargs):
                 return result
 
             elif stage == "replay_forward":
-                return _get_replay_result(replay.pop_forward(), scores, topk, *args, **kwargs)
+                replay_idx = replay.pop_forward()
+                if manager.enable_check_replay_result:
+                    # Direct evidence the wrapper's replay_forward branch ran in
+                    # the megatron MoE forward path (vs falling through to
+                    # old_topk_fn). Gated on the correctness-check toggle so
+                    # production training stays quiet. replay_idx_sum is a cheap
+                    # fingerprint of the indices for cross-rank / cross-step
+                    # comparison.
+                    logger.info(
+                        f"R3 wrapper: replay_forward branch taken "
+                        f"(rank {_get_rank()}, n_tokens={replay_idx.shape[0]}, "
+                        f"topk={replay_idx.shape[1]}, "
+                        f"replay_idx_sum={int(replay_idx.sum().item())})"
+                    )
+                return _get_replay_result(replay_idx, scores, topk, *args, **kwargs)
 
             elif stage == "replay_backward":
-                return _get_replay_result(replay.pop_backward(), scores, topk, *args, **kwargs)
+                replay_idx = replay.pop_backward()
+                if manager.enable_check_replay_result:
+                    logger.info(
+                        f"R3 wrapper: replay_backward branch taken "
+                        f"(rank {_get_rank()}, n_tokens={replay_idx.shape[0]}, "
+                        f"topk={replay_idx.shape[1]}, "
+                        f"replay_idx_sum={int(replay_idx.sum().item())})"
+                    )
+                return _get_replay_result(replay_idx, scores, topk, *args, **kwargs)
 
             else:
                 return old_topk_fn(scores, topk, *args, **kwargs)
@@ -196,6 +218,15 @@ def check_replay_result(self, old_topk_fn, scores, topk, top_indices, *args, **k
         is_mismatch = ~has_overlap & ~is_padding
 
         mismatch_count = is_mismatch.sum().item()
+        n_tokens = orig_flat.shape[0]
+        # Unconditional log so we have direct evidence the check actually
+        # ran (its silent return on mismatch_count==0 is otherwise
+        # indistinguishable from never being called).
+        logger.info(
+            f"R3 check (rank {_get_rank()}, stage {self.stage}): "
+            f"n_tokens={n_tokens} mismatch={mismatch_count} "
+            f"({100 * mismatch_count / max(n_tokens, 1):.2f}%)"
+        )
         if mismatch_count == 0:
             return