fix straggler detection reporting

Signed-off-by: Ananth Subramaniam <[email protected]>

Author: ananthsub

src/megatron/bridge/training/nvrx_straggler.py

@@ -94,42 +94,46 @@ def wrap_train_step_function(self, train_step_func: Callable) -> Callable:
         """
         Wrap the training step function with straggler detection monitoring.
 
+        The NVRx straggler detector instruments functions to measure CUDA kernel
+        execution times. This method wraps the train_step function so that each
+        call is profiled for straggler detection.
+
         Args:
             train_step_func: The actual training step function to wrap for monitoring.
 
         Returns:
-            The wrapped training step function.
+            The wrapped training step function that should be used instead of the original.
+            If wrapping fails or is disabled, returns the original function.
         """
 
         if not self.initialized or not self.config.enabled:
             return train_step_func
 
         if self.wrapped_function is not None:
             self.logger.warning("Train step function already wrapped. Skipping.")
-            return train_step_func
+            return self.wrapped_function
 
         try:
-            # Create a wrapper object with train_step method for nvidia-resiliency-ext
-            # TODO: See if NVRx can support functions directly without needing them attached to a class
+            # Create a wrapper object with train_step method for nvidia-resiliency-ext.
+            # NVRx requires a method on an object, not a standalone function.
+            # We store the wrapper object to prevent garbage collection.
             class TrainStepWrapper:
                 def __init__(self, func):
                     self.train_step = func
-                    self._original_func = func
 
                 def __call__(self, *args, **kwargs):
-                    return self._original_func(*args, **kwargs)
+                    return self.train_step(*args, **kwargs)
 
             wrapper_obj = TrainStepWrapper(train_step_func)
 
             # Create a callable ID for the training step function
+            # wrap_callables modifies wrapper_obj.train_step in-place to add instrumentation
             callable_id = straggler.CallableId(wrapper_obj, "train_step")
             straggler.Detector.wrap_callables(callable_ids=[callable_id])
 
-            self.wrapped_function = train_step_func
-            self.logger.debug("Train step function wrapped for NVRx straggler detection.")
-
-            # Return the original function since the wrapper is just for nvidia-resiliency-ext
-            return train_step_func
+            # Store the wrapped function to prevent garbage collection and enable reuse
+            self.wrapped_function = wrapper_obj
+            return wrapper_obj
 
         except Exception as e:
             self.logger.warning(f"Failed to wrap train step function with NVRx: {e}. Continuing without wrapping.")
@@ -145,7 +149,6 @@ def check_stragglers(self, global_rank: int) -> bool:
         Returns:
             True if stragglers were detected and stop_if_detected is True, False otherwise.
         """
-
         if not self.initialized or not self.config.enabled:
             return False
 
@@ -242,7 +245,7 @@ def _print_gpu_scores(self, report) -> None:
                 num_best=self.config.num_gpu_perf_scores_to_print,
                 num_worst=self.config.num_gpu_perf_scores_to_print,
             )
-            self.logger.info(f"\nGPU relative performance:\n{rel_perf_str}")
+            self.logger.info(f"GPU relative performance:\n{rel_perf_str}")
 
         if self.config.calc_individual_gpu_perf:
             indiv_perf_str = self._format_gpu_scores(
@@ -251,7 +254,7 @@ def _print_gpu_scores(self, report) -> None:
                 num_best=self.config.num_gpu_perf_scores_to_print,
                 num_worst=self.config.num_gpu_perf_scores_to_print,
             )
-            self.logger.info(f"\nGPU individual performance:\n{indiv_perf_str}")
+            self.logger.info(f"GPU individual performance:\n{indiv_perf_str}")
 
     def _log_gpu_scores(self, report) -> None:
         """Log GPU performance scores as structured data."""
@@ -299,11 +302,9 @@ def _gather_flag_from_rank0(self, flag: bool) -> bool:
     def shutdown(self) -> None:
         """Shutdown the straggler detector."""
         if self.initialized and self.config.enabled:
-            self.logger.info("Shutting down NVRx straggler detection...")
             straggler.Detector.shutdown()
             self.initialized = False
             self.wrapped_function = None
-            self.logger.info("NVRx straggler detection shutdown complete.")
 
 
 def check_nvrx_straggler_detection(nvrx_straggler_manager: Optional["NVRxStragglerDetectionManager"]) -> bool:

src/megatron/bridge/training/train.py

@@ -184,13 +184,14 @@ def train(
 
     # Initialize NVRx straggler detection if enabled
     nvrx_straggler_manager = global_state.nvrx_straggler_manager
+    wrapped_train_step = train_step  # Default to original function
     if nvrx_straggler_manager is not None:
         try:
             # Initialize the straggler detector first
             nvrx_straggler_manager.initialize()
             # Wrap the train_step function for monitoring
-            # Note: The nvidia-resiliency-ext library will monitor the actual train_step calls
-            nvrx_straggler_manager.wrap_train_step_function(train_step)
+            # The wrapped function must be used instead of the original to collect profiling data
+            wrapped_train_step = nvrx_straggler_manager.wrap_train_step_function(train_step)
         except Exception as e:
             print_rank_0(f"Failed to initialize NVRx straggler detection: {e}")
             # Set to None to disable further checks
@@ -335,7 +336,7 @@ def train(
             grad_norm,
             num_zeros_in_grad,
             log_max_attention_logit,
-        ) = train_step(
+        ) = wrapped_train_step(
             wrapped_forward_step_func,
             train_data_iterator,
             model,

tests/unit_tests/training/test_nvrx_straggler.py

@@ -280,41 +280,66 @@ def dummy_func():
     def test_wrap_train_step_function_already_wrapped(self, manager, mock_straggler_module):
         """Test wrapping when already wrapped."""
         manager.initialized = True
-        manager.wrapped_function = Mock()
+        existing_wrapper = Mock()
+        manager.wrapped_function = existing_wrapper
 
         def dummy_func():
             pass
 
         result = manager.wrap_train_step_function(dummy_func)
 
-        assert result is dummy_func
+        # Should return the already-wrapped function, not the new dummy_func
+        assert result is existing_wrapper
         mock_straggler_module.Detector.wrap_callables.assert_not_called()
 
     def test_wrap_train_step_function_success(self, manager, mock_straggler_module):
         """Test successful wrapping."""
         manager.initialized = True
 
         def dummy_func():
-            pass
+            return "original_result"
 
         result = manager.wrap_train_step_function(dummy_func)
 
-        assert result is dummy_func
-        assert manager.wrapped_function is dummy_func
+        # The result should be a callable wrapper object, not the original function
+        assert result is not dummy_func
+        assert callable(result)
 
-        # Verify CallableId was called with a TrainStepWrapper object, not the original function
+        # The wrapped_function should be the wrapper object
+        assert manager.wrapped_function is result
+
+        # Verify CallableId was called with the wrapper object
         mock_straggler_module.CallableId.assert_called_once()
         call_args = mock_straggler_module.CallableId.call_args[0]
         assert len(call_args) == 2
         wrapper_obj, method_name = call_args
         assert method_name == "train_step"
 
-        # Verify the wrapper object has the train_step method that wraps our dummy_func
+        # Verify the wrapper object has the train_step method
         assert hasattr(wrapper_obj, "train_step")
-        assert wrapper_obj.train_step is dummy_func
 
         mock_straggler_module.Detector.wrap_callables.assert_called_once()
 
+    def test_wrap_train_step_function_callable(self, manager, mock_straggler_module):
+        """Test that the wrapped function is callable and routes through train_step."""
+        manager.initialized = True
+
+        call_count = [0]
+
+        def dummy_func(*args, **kwargs):
+            call_count[0] += 1
+            return "result"
+
+        result = manager.wrap_train_step_function(dummy_func)
+
+        # The wrapper should be callable
+        assert callable(result)
+
+        # Calling the wrapper should route through train_step (which is the original func before wrap_callables)
+        # Note: In real usage, wrap_callables modifies train_step in-place, but in tests it's mocked
+        result()
+        assert call_count[0] == 1
+
     def test_check_stragglers_disabled(self, manager, mock_straggler_module):
         """Test check_stragglers when disabled."""
         manager.config.enabled = False
@@ -469,6 +494,7 @@ def test_shutdown_disabled(self, manager, mock_straggler_module):
     def test_shutdown_success(self, manager, mock_straggler_module):
         """Test successful shutdown."""
         manager.initialized = True
+        manager.wrapped_function = Mock()  # Simulate a wrapped function
 
         manager.shutdown()
 
@@ -584,7 +610,11 @@ def train_step():
                 return "training"
 
             wrapped_func = manager.wrap_train_step_function(train_step)
-            assert wrapped_func is train_step
+            # The wrapped function should NOT be the same as the original - it should be a wrapper
+            assert wrapped_func is not train_step
+            assert callable(wrapped_func)
+            # The wrapper should be stored
+            assert manager.wrapped_function is wrapped_func
 
             # Test straggler detection function
             with patch("torch.distributed.is_initialized", return_value=False):