Add perf monitor integration - #1216
Merged
Merged
Conversation
lzy-edu
requested review from
aoyulong,
heavyrain-lzy and
zhaoyinglia
as code owners
June 9, 2026 04:28
Contributor
There was a problem hiding this comment.
Pull request overview
Ports the perf monitor integration from #1180 onto the current training stack, wiring performance logging/summary generation into the Megatron training loop and ensuring runner/launcher configs propagate correctly.
Changes:
- Added a new
flagscale.train.perf_monitormodule (FLOPS estimation, step timing, logging, and training-loop hooks), plus a standalone smoke script. - Integrated perf monitor hooks into
flagscale/train/megatron/training/training.pyand exposed CLI args for Megatron training. - Updated runner/backend/launcher handling and added unit tests to validate config propagation and command generation.
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| tools/perf_monitor/perf_smoke.py | Standalone smoke runner for perf monitor logging/summary output. |
| tests/unit_tests/train/test_perf_monitor.py | Smoke unit test to assert summary/log artifacts are written. |
| tests/unit_tests/runner/test_launcher_ssh_perf.py | Ensures perf-monitor runner keys are stripped from torchrun args. |
| tests/unit_tests/runner/test_backend_megatron_perf.py | Validates runner→train.system propagation and script mkdir for perf log dir. |
| flagscale/train/perf_monitor/READ_PERF_MONITOR.md | Added documentation note (currently describes legacy path). |
| flagscale/train/perf_monitor/perf_metrics.py | Core metric collection + callback for train-loop integration. |
| flagscale/train/perf_monitor/perf_logger.py | Text/JSON file logging and summary emission. |
| flagscale/train/perf_monitor/hooks.py | Global init + iteration/training hooks used by training loop/tests. |
| flagscale/train/perf_monitor/flops_calculator.py | FLOPS formula helpers for different transformer variants. |
| flagscale/train/perf_monitor/init.py | Exposes perf monitor public API. |
| flagscale/train/megatron/training/training.py | Calls perf monitor hooks around each training step and at train end. |
| flagscale/train/megatron/training/arguments_fs.py | Adds Megatron CLI flags for perf monitor configuration. |
| flagscale/runner/launcher/launcher_ssh.py | Strips perf monitor keys so torchrun doesn’t receive unknown args. |
| flagscale/runner/backend/backend_megatron.py | Copies perf monitor config from runner→train.system and creates log dir in scripts. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
40
to
46
| import inspect | ||
| import json | ||
| import logging | ||
| import math | ||
| import os | ||
| import socket | ||
| import sys |
| forward_step_func, train_data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=iteration | ||
| ) | ||
| ft_integration.on_training_step_end() | ||
|
|
Comment on lines
+43
to
+45
| def perf_monitor_end_training(writer=None, wandb_writer=None): | ||
| if _perf_monitor_callback is not None: | ||
| _perf_monitor_callback.on_train_end(writer, wandb_writer) |
Comment on lines
+117
to
+129
| def save_summary(self, final_stats=None): | ||
| if not self.enabled: | ||
| return | ||
|
|
||
| summary = { | ||
| "session_info": { | ||
| "start_time": self.session_timestamp, | ||
| "end_time": datetime.now().astimezone().isoformat(), | ||
| "total_iterations": len(self.json_data), | ||
| }, | ||
| "final_statistics": final_stats or {}, | ||
| "iteration_logs": self.json_data, | ||
| } |
Comment on lines
+285
to
+297
| metrics = self.monitor.calculate_metrics() | ||
| self.monitor.file_logger.save_summary( | ||
| { | ||
| "avg_tflops_per_gpu": metrics.tflops_per_gpu, | ||
| "avg_tflops_total": metrics.tflops_total, | ||
| "avg_step_time_ms": metrics.avg_step_time * 1000, | ||
| "min_step_time_ms": metrics.min_step_time * 1000 | ||
| if metrics.min_step_time != float("inf") | ||
| else 0.0, | ||
| "max_step_time_ms": metrics.max_step_time * 1000, | ||
| "peak_memory_gb": self.monitor.peak_memory_gb, | ||
| } | ||
| ) |
Comment on lines
+1
to
+13
| # Perf Monitor on Metax C550 | ||
|
|
||
| ## Scope | ||
|
|
||
| This note documents the current `perf_monitor` path for Metax C550 on `main-legacy`. | ||
|
|
||
| Important: | ||
|
|
||
| - The current runnable branch is `main-legacy`. | ||
| - The active code path is the legacy runner: | ||
| - `flagscale/runner/runner_train.py` | ||
| - The new runner launcher path is not active in this branch. | ||
|
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR Category
Train
PR Types
New Features, Improvements
PR Description
This PR ports the perf monitor integration from #1180 onto the latest main branch.
It includes:
Based on: