Add perf monitor integration (flagos-ai#1216)

### PR Category
Train
### PR Types
New Features, Improvements

### PR Description
This PR ports the perf monitor integration from flagos-ai#1180 onto the latest
main branch.

It includes:
- Perf monitor integration for the training flow
- Runner and launcher updates for perf monitor support
- Compatibility updates for the latest main branch

Based on:
- flagos-ai#1180

---------

Co-authored-by: wanglei19991004 <13453558035@163.com>
Co-authored-by: zhaoyingli <86812880+zhaoyinglia@users.noreply.github.com>

Author: 3 people

flagscale/runner/backend/backend_megatron.py

@@ -10,6 +10,18 @@
 )
 from flagscale.runner.utils import get_pkg_dir, logger, parse_hostfile, resolve_path
 
+PERF_MONITOR_RUNNER_KEYS = (
+    "enable_perf_monitor",
+    "perf_log_interval",
+    "perf_log_dir",
+    "perf_console_output",
+    "perf_log_format",
+    "perf_memory_tracking",
+    "perf_breakdown",
+    "perf_max_log_files",
+    "perf_model_type",
+)
+
 
 class MegatronBackend(BackendBase):
     def __init__(self, config: DictConfig):
@@ -20,6 +32,7 @@ def __init__(self, config: DictConfig):
 
     def _prepare(self):
         _update_config_train(self.config)
+        self._prepare_perf_monitor_config()
         self.user_args = _get_args_megatron(self.config)
         self.rdzv_id = datetime.now().strftime("%Y%m%d_%H%M%S.%f")
         self.user_envs = self.config.experiment.get("envs", {})
@@ -30,6 +43,22 @@ def _prepare(self):
         logger.info("\n************** configuration **************")
         logger.info(f"\n{OmegaConf.to_yaml(self.config)}")
 
+    def _prepare_perf_monitor_config(self):
+        system_config = self.config.train.system
+        runner_config = self.config.experiment.runner
+
+        OmegaConf.set_struct(system_config, False)
+        for key in PERF_MONITOR_RUNNER_KEYS:
+            if runner_config.get(key, None) is not None and system_config.get(key, None) is None:
+                system_config[key] = runner_config.get(key)
+
+        if system_config.get("perf_log_dir", None) is not None:
+            system_config.perf_log_dir = resolve_path(
+                system_config.perf_log_dir, "system.perf_log_dir"
+            )
+        elif system_config.get("enable_perf_monitor", False):
+            system_config.perf_log_dir = os.path.join(system_config.logging.log_dir, "perf_monitor")
+
     def generate_run_script(
         self,
         config,
@@ -78,6 +107,8 @@ def generate_run_script(
             f.write(f"mkdir -p {system_config.logging.details_dir}\n")
             f.write(f"mkdir -p {system_config.logging.tensorboard_dir}\n")
             f.write(f"mkdir -p {system_config.logging.wandb_save_dir}\n")
+            if system_config.get("perf_log_dir", None):
+                f.write(f"mkdir -p {system_config.perf_log_dir}\n")
             f.write("\n")
             f.write(f"cd {pkg_dir}\n")
             f.write("\n")

flagscale/runner/launcher/launcher_ssh.py

@@ -116,6 +116,24 @@ def _get_runner_cmd_train(
         del runner_args["nsys_rep_file_path"]
     if "deploy" in runner_args:
         del runner_args["deploy"]
+    if "enable_perf_monitor" in runner_args:
+        del runner_args["enable_perf_monitor"]
+    if "perf_log_interval" in runner_args:
+        del runner_args["perf_log_interval"]
+    if "perf_log_dir" in runner_args:
+        del runner_args["perf_log_dir"]
+    if "perf_console_output" in runner_args:
+        del runner_args["perf_console_output"]
+    if "perf_log_format" in runner_args:
+        del runner_args["perf_log_format"]
+    if "perf_memory_tracking" in runner_args:
+        del runner_args["perf_memory_tracking"]
+    if "perf_breakdown" in runner_args:
+        del runner_args["perf_breakdown"]
+    if "perf_max_log_files" in runner_args:
+        del runner_args["perf_max_log_files"]
+    if "perf_model_type" in runner_args:
+        del runner_args["perf_model_type"]
     runner_args["rdzv_id"] = rdzv_id
     # runner_args["master_addr"] = master_addr
     # runner_args["master_port"] = master_port

flagscale/train/megatron/training/arguments_fs.py

@@ -889,6 +889,73 @@ def _add_vision_args(parser):
     )
     return parser
 
+def _add_perf_monitor_args(parser):
+    group = parser.add_argument_group(title="flagscale perf monitor")
+
+    group.add_argument(
+        "--enable-perf-monitor",
+        action="store_true",
+        default=False,
+        help="Enable FlagScale performance monitoring during training.",
+    )
+    group.add_argument(
+        "--perf-log-interval",
+        type=int,
+        default=10,
+        help="Log performance metrics every N iterations.",
+    )
+    group.add_argument(
+        "--perf-log-dir",
+        type=str,
+        default=None,
+        help="Directory used to save performance monitor logs.",
+    )
+    group.add_argument(
+        "--perf-console-output",
+        action="store_true",
+        default=False,
+        help="Also emit performance monitor logs to stdout on rank 0.",
+    )
+    group.add_argument(
+        "--perf-log-format",
+        type=str,
+        choices=["text", "json", "both"],
+        default="both",
+        help="Output format for performance monitor files.",
+    )
+    group.add_argument(
+        "--perf-memory-tracking",
+        dest="perf_memory_tracking",
+        action="store_true",
+        help="Track CUDA memory usage in the performance monitor.",
+    )
+    group.add_argument(
+        "--no-perf-memory-tracking",
+        dest="perf_memory_tracking",
+        action="store_false",
+        help="Disable CUDA memory tracking in the performance monitor.",
+    )
+    group.set_defaults(perf_memory_tracking=True)
+    group.add_argument(
+        "--perf-breakdown",
+        action="store_true",
+        default=False,
+        help="Include estimated component breakdowns in performance logs.",
+    )
+    group.add_argument(
+        "--perf-max-log-files",
+        type=int,
+        default=10,
+        help="Maximum number of historical performance log files to keep.",
+    )
+    group.add_argument(
+        "--perf-model-type",
+        type=str,
+        choices=["auto", "gpt", "llama", "qwen", "mixtral", "aquila", "moe"],
+        default="auto",
+        help="Model type hint used for FLOPS estimation.",
+    )
+    return parser
 
 def _add_flagos_args(parser):
     group = parser.add_argument_group(title="flagscale fl")
@@ -977,4 +1044,5 @@ def add_flagscale_arguments(parser):
     parser = _add_peft_args(parser)
     parser = _add_flagos_args(parser)
     parser = _add_flagscale_specific_args(parser)
+    parser = _add_perf_monitor_args(parser)
     return parser

flagscale/train/megatron/training/training.py

@@ -257,6 +257,12 @@ def set_startup_timestamps(program_start=None, main_entry=None):
 from megatron.training.global_vars import get_spiky_loss_detector
 from megatron.training.peft import PEFT # Import PEFT from peft module
 from megatron.plugin.hetero.parallel_context import get_parallel_context
+from flagscale.train.perf_monitor.hooks import (
+    initialize_perf_monitor,
+    perf_monitor_end_iteration,
+    perf_monitor_end_training,
+    perf_monitor_start_iteration,
+)
 
 from megatron.plugin.platform import get_platform
 cur_platform = get_platform()
@@ -2987,6 +2993,7 @@ def train(
     timers('interval-time', log_level=0).start(barrier=True)
     print_datetime('before the start of training step')
     report_memory_flag = True
+    perf_callback = initialize_perf_monitor(args)
     pre_hook_enabled = False
     should_exit = False
     exit_code = 0
@@ -3019,6 +3026,7 @@ def train(
     num_microbatches = get_num_microbatches()
 
     ########## FlagScale Begin ##########
+    writer = get_tensorboard_writer()
     wandb_writer = get_wandb_writer()
     if wandb_writer and args.wandb_log_model:
         # wandb.watch's log_freg needs to take the accumulated number of microbatches into account
@@ -3278,6 +3286,8 @@ def trace_handler(p):
             num_zeros_in_grad = 0
             max_attention_logit = None
         else:
+            if perf_callback is not None:
+                perf_monitor_start_iteration(iteration)
             ft_integration.on_training_step_start()
             (
                 loss_dict,
@@ -3292,6 +3302,8 @@ def trace_handler(p):
                 forward_step_func, train_data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=iteration
             )
             ft_integration.on_training_step_end()
+            if perf_callback is not None:
+                perf_monitor_end_iteration(iteration, writer, wandb_writer)
         if should_checkpoint:
             save_checkpoint_and_time(
                 iteration,
@@ -3569,6 +3581,7 @@ def trace_handler(p):
 
     # Flush TensorBoard, WandB writers and one-logger.
     writer = get_tensorboard_writer()
+    perf_monitor_end_training(writer, wandb_writer)
     if writer:
         writer.flush()
 

flagscale/train/perf_monitor/__init__.py

@@ -0,0 +1,20 @@
+"""FlagScale performance monitor utilities."""
+
+from .hooks import (
+    get_perf_monitor,
+    initialize_perf_monitor,
+    perf_monitor_end_iteration,
+    perf_monitor_end_training,
+    perf_monitor_start_iteration,
+)
+from .perf_metrics import FLOPSMeasurementCallback, PerformanceMonitor
+
+__all__ = [
+    "FLOPSMeasurementCallback",
+    "PerformanceMonitor",
+    "get_perf_monitor",
+    "initialize_perf_monitor",
+    "perf_monitor_end_iteration",
+    "perf_monitor_end_training",
+    "perf_monitor_start_iteration",
+]

flagscale/train/perf_monitor/flops_calculator.py

@@ -0,0 +1,63 @@
+"""FLOPS estimation helpers for performance monitoring."""
+
+from __future__ import annotations
+
+
+class FLOPSFormulas:
+    """Collection of approximate transformer FLOPS formulas."""
+
+    @staticmethod
+    def attention_flops(batch_size, seq_length, hidden_size, num_attention_heads):
+        head_dim = hidden_size // max(num_attention_heads, 1)
+        qkv_flops = 3 * 2 * batch_size * seq_length * hidden_size * hidden_size
+        score_flops = 2 * batch_size * num_attention_heads * seq_length * seq_length * head_dim
+        value_flops = 2 * batch_size * num_attention_heads * seq_length * seq_length * head_dim
+        out_flops = 2 * batch_size * seq_length * hidden_size * hidden_size
+        return qkv_flops + score_flops + value_flops + out_flops
+
+    @staticmethod
+    def gqa_attention_flops(
+        batch_size, seq_length, hidden_size, num_attention_heads, num_query_groups
+    ):
+        head_dim = hidden_size // max(num_attention_heads, 1)
+        kv_hidden_size = head_dim * max(num_query_groups, 1)
+        q_flops = 2 * batch_size * seq_length * hidden_size * hidden_size
+        kv_flops = 2 * 2 * batch_size * seq_length * hidden_size * kv_hidden_size
+        score_flops = 2 * batch_size * num_attention_heads * seq_length * seq_length * head_dim
+        value_flops = 2 * batch_size * num_attention_heads * seq_length * seq_length * head_dim
+        out_flops = 2 * batch_size * seq_length * hidden_size * hidden_size
+        return q_flops + kv_flops + score_flops + value_flops + out_flops
+
+    @staticmethod
+    def ffn_flops(batch_size, seq_length, hidden_size, ffn_hidden_size, use_swiglu=False):
+        if use_swiglu:
+            gate_flops = 2 * batch_size * seq_length * hidden_size * ffn_hidden_size
+            up_flops = 2 * batch_size * seq_length * hidden_size * ffn_hidden_size
+            swiglu_flops = batch_size * seq_length * ffn_hidden_size
+            down_flops = 2 * batch_size * seq_length * ffn_hidden_size * hidden_size
+            return gate_flops + up_flops + swiglu_flops + down_flops
+
+        up_flops = 2 * batch_size * seq_length * hidden_size * ffn_hidden_size
+        down_flops = 2 * batch_size * seq_length * ffn_hidden_size * hidden_size
+        return up_flops + down_flops
+
+    @staticmethod
+    def moe_flops(
+        batch_size,
+        seq_length,
+        hidden_size,
+        ffn_hidden_size,
+        num_experts,
+        top_k,
+        use_swiglu=False,
+    ):
+        router_flops = 2 * batch_size * seq_length * hidden_size * num_experts
+        active_tokens = batch_size * seq_length * top_k
+        if use_swiglu:
+            expert_flops = (
+                3 * 2 * active_tokens * hidden_size * ffn_hidden_size
+                + active_tokens * ffn_hidden_size
+            )
+        else:
+            expert_flops = 4 * active_tokens * hidden_size * ffn_hidden_size
+        return router_flops + expert_flops

flagscale/train/perf_monitor/hooks.py

@@ -0,0 +1,47 @@
+"""Training-loop hooks for the performance monitor."""
+
+from __future__ import annotations
+
+import torch
+
+from flagscale.train.perf_monitor.perf_metrics import FLOPSMeasurementCallback
+
+_perf_monitor_callback = None
+
+
+def initialize_perf_monitor(args):
+    """Initialize the global performance monitor if enabled."""
+    global _perf_monitor_callback
+
+    if not getattr(args, "enable_perf_monitor", False):
+        _perf_monitor_callback = None
+        return None
+
+    log_interval = getattr(args, "perf_log_interval", 10)
+    _perf_monitor_callback = FLOPSMeasurementCallback(args, log_interval=log_interval)
+
+    if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
+        print(f"[Performance Monitor] Initialized with log interval: {log_interval}")
+
+    return _perf_monitor_callback
+
+
+def get_perf_monitor():
+    return _perf_monitor_callback
+
+
+def perf_monitor_start_iteration(iteration):
+    if _perf_monitor_callback is not None:
+        _perf_monitor_callback.on_train_batch_start(iteration)
+
+
+def perf_monitor_end_iteration(iteration, writer=None, wandb_writer=None):
+    if _perf_monitor_callback is not None:
+        _perf_monitor_callback.on_train_batch_end(iteration, writer, wandb_writer)
+
+
+def perf_monitor_end_training(writer=None, wandb_writer=None):
+    global _perf_monitor_callback
+    if _perf_monitor_callback is not None:
+        _perf_monitor_callback.on_train_end(writer, wandb_writer)
+        _perf_monitor_callback = None