Implement straggler detection function, generating report periodically.

Author: Eric-Wei-10

.github/workflows/functional-tests-nvidia.yml

@@ -49,22 +49,26 @@ jobs:
 
           git config --global --add safe.directory /__w/FlagScale/FlagScale
           if [ "${{ inputs.type }}" = "train" ] || [ "${{ inputs.type }}" = "hetero_train" ]; then
+            source /root/miniconda3/bin/activate flagscale-train
             PYTHONPATH=./:$PYTHONPATH pip install . --no-build-isolation --verbose --config-settings=device="gpu" --config-settings=backend="Megatron-LM"
             if [ "${{ inputs.task }}" = "llava_onevision" ]; then
               PYTHONPATH=./:$PYTHONPATH pip install . --no-build-isolation --verbose --config-settings=device="gpu" --config-settings=backend="Megatron-Energon"
               cp -r third_party/Megatron-Energon/src/megatron/energon third_party/Megatron-LM/megatron
             fi
+            conda deactivate
           elif [ "${{ inputs.type }}" = "inference" ] || [ "${{ inputs.type }}" = "serve" ]; then
             source /root/miniconda3/bin/activate flagscale-inference
             pip install scikit-build scikit-build-core
             pip install git+https://github.com/FlagOpen/[email protected]
             PYTHONPATH=./:$PYTHONPATH pip install . --config-settings=backend="vllm" --verbose --no-build-isolation
             conda deactivate
           elif [ "${{ inputs.type }}" = "rl" ]; then
+            source /root/miniconda3/bin/activate flagscale-RL
             python tools/patch/unpatch.py --backend verl
             cd third_party/verl
             pip install --no-deps -e .
             cd ../..
+            conda deactivate
           else
             echo "Unknown backend type: ${{ inputs.type }}"
             exit 1

.github/workflows/unit-tests-nvidia.yml

@@ -55,6 +55,12 @@ jobs:
           git config --global --add safe.directory /__w/FlagScale/FlagScale
           if [ "${{ inputs.backend }}" = "megatron" ] || [ "${{ inputs.backend }}" = "flagscale" ]; then
             echo ""
+            source /root/miniconda3/bin/activate flagscale-train
+            git clone https://github.com/NVIDIA/nvidia-resiliency-ext
+            cd nvidia-resiliency-ext
+            pip install .
+            cd ..
+            conda deactivate
             # PYTHONPATH=./:$PYTHONPATH pip install . --config-settings=backend="Megatron-LM" --verbose --no-build-isolation
           elif [ "${{ inputs.backend }}" = "vllm" ]; then
             source /root/miniconda3/bin/activate flagscale-inference

flagscale/backends/Megatron-LM/megatron/core/pipeline_parallel/schedules.py

@@ -22,6 +22,7 @@
     nvtx_range_pop,
     nvtx_range_push,
 )
+from flagscale.train.straggler_detection import StragglerDetectionWrapper
 
 # Types
 Shape = Union[List[int], torch.Size]
@@ -184,6 +185,7 @@ def set_current_microbatch(model, microbatch_id):
             layer.current_microbatch = microbatch_id
 
 
+@StragglerDetectionWrapper(level=2, section_name="microbatch_forward")
 def forward_step(
     forward_step_func,
     data_iterator,
@@ -368,6 +370,7 @@ def forward_step(
     return [output_tensor], num_tokens
 
 
+@StragglerDetectionWrapper(level=2, section_name="microbatch_backward")
 def backward_step(input_tensor, output_tensor, output_tensor_grad, model_type, config):
     """Backward step through passed-in output tensor.
 
@@ -485,6 +488,9 @@ def forward_backward_no_pipelining(
         adjust_tensor_shapes_fn is None
     ), "adjust_tensor_shapes_fn is not supported for non-pipeline-parallel schedule"
 
+    from megatron.training.global_vars import get_args
+    args = get_args()
+
     config = get_model_config(model)
     if config.timers is not None:
         config.timers('forward-backward', log_level=1).start(barrier=config.barrier_with_L1_time)
@@ -511,10 +517,16 @@ def forward_backward_no_pipelining(
                 collect_non_loss_data,
                 is_first_microbatch=check_first_val_step(first_val_step, forward_only, i == 0),
                 current_microbatch=i,
+                user_specified_level=args.straggler_detection_level,
+                passed_warmup_stage=args.curr_iteration > args.straggler_detection_warmup_iterations,
             )
             total_num_tokens += num_tokens
             if not forward_only:
-                backward_step(input_tensor, output_tensor, output_tensor_grad, model_type, config)
+                backward_step(
+                    input_tensor, output_tensor, output_tensor_grad, model_type, config,
+                    user_specified_level=args.straggler_detection_level, 
+                    passed_warmup_stage=args.curr_iteration > args.straggler_detection_warmup_iterations,
+                )
 
     # Run computation for last microbatch out of context handler (want to
     # synchronize gradients).
@@ -531,11 +543,19 @@ def forward_backward_no_pipelining(
             first_val_step, forward_only, num_microbatches == 1
         ),
         current_microbatch=num_microbatches - 1,
+        user_specified_level=args.straggler_detection_level,
+        passed_warmup_stage=args.curr_iteration > args.straggler_detection_warmup_iterations,
+        generate_report=forward_only and (args.curr_iteration % args.straggler_detection_interval) == 0
     )
     total_num_tokens += num_tokens
 
     if not forward_only:
-        backward_step(input_tensor, output_tensor, output_tensor_grad, model_type, config)
+        backward_step(
+            input_tensor, output_tensor, output_tensor_grad, model_type, config,
+            user_specified_level=args.straggler_detection_level, 
+            passed_warmup_stage=args.curr_iteration > args.straggler_detection_warmup_iterations,
+            generate_report=not forward_only and (args.curr_iteration % args.straggler_detection_interval) == 0
+        )
 
     if config.finalize_model_grads_func is not None and not forward_only:
         # Finalize model grads (perform full grad all-reduce / reduce-scatter for

flagscale/backends/Megatron-LM/megatron/training/arguments.py

@@ -1666,6 +1666,16 @@ def _add_ft_package_args(parser):
     group.add_argument('--calc-ft-timeouts', action='store_true',
                        help='If set, FT package will try to automatically compute the timeouts. '
                        'Note: This feature is for Nvidia internal use only.')
+    group.add_argument('--straggler-detection-level', type=int, 
+                       default=0, choices=range(0,3),
+                       help='Granularity of straggler detection level.'
+                       '    0: off.'
+                       '    1: per train step.'
+                       '    2: per train section.')
+    group.add_argument('--straggler-detection-interval', type=int, default=10,
+                       help='Interval in iterations for generating detection report.')
+    group.add_argument('--straggler-detection-warmup-iterations', type=int, default=50,
+                       help='Interval in iterations for generating detection report.')
     return parser