Skip to content

fix(megatron): restore "compute per GPU" TFLOP label and make log parsers robust - #926

Open
amd-fuyuajin wants to merge 6 commits into
mainfrom
dev/fix-megatron-logging-format
Open

fix(megatron): restore "compute per GPU" TFLOP label and make log parsers robust#926
amd-fuyuajin wants to merge 6 commits into
mainfrom
dev/fix-megatron-logging-format

Conversation

@amd-fuyuajin

Copy link
Copy Markdown
Collaborator

Summary

Restores the compute per GPU (TFLOP/s/GPU) training-log label (from #842) that
was inadvertently reverted by the Flux MLPerf logging refactor (#820), and updates
the three Megatron log parsers so both the current and legacy log formats parse
correctly.

Background

What changed

primus/backends/megatron/patches/training_log/print_rank_last_patches.py

  • Restore the compute per GPU (TFLOP/s/GPU): X (avg Y) label in _inject_tflops().
  • Refresh the module/render docstrings to match the actual output
    (compute per GPU ... and tokens/s/GPU inst/harmonic mean: ...).

Log parserstools/daily/daily_report.py, tools/auto_benchmark/metrics.py,
benchmark/megatron/model/benchmark_report.py

  • Accept both the current labels (compute per GPU, tokens/s/GPU inst/harmonic mean)
    and the legacy ones (throughput per GPU, tokens per GPU (tokens/s/GPU)), using
    non-capturing alternations so positional group indices are unchanged.
  • benchmark_report: the memory segment is now appended last (as
    hip mem usage/free/total/usage_ratio: X GB/...) rather than sitting between
    elapsed and throughput as mem usages: X; reorder the pattern/group indices and
    accept both memory labels.
  • Tolerate the optional space before (avg (real logs can render X(avg Y)), and
    fall back to the instantaneous value when a running average is absent (raw
    warmup lines).

Why the parsers were already broken

The token label from #842 (tokens/s/GPU inst/harmonic mean) never matched the
parsers' tokens per GPU (tokens/s/GPU) regex, and benchmark_report expected a
mem usages: segment that current logs no longer emit. This PR realigns all three
with the current log format while staying backward-compatible.

Test plan

  • pytest tests/unit_tests/backends/megatron/test_training_log_patches.py — 7 passed
  • End-to-end: drove the real ThroughputAverageExtension/ElapsedAverageExtension
    across multiple iterations, rendered the line, and confirmed all three parsers
    extract correct inst/avg TFLOP, tokens, and memory values.
  • Validated all three parsers against real pre-train logs
    (llama2_7B-BF16-pretrain, MI300X): current + legacy + raw-warmup lines all parse.
  • No new lint errors.

amd-fuyuajin and others added 4 commits July 27, 2026 21:26
The Flux MLPerf logging refactor (#820) reverted the TFLOP label from
"compute per GPU (TFLOP/s/GPU): X (avg Y)" back to the ambiguous
"throughput per GPU (TFLOP/s/GPU): X/Y" while extracting the shared
_inject_tflops helper. Restore the clearer "compute" wording (per #842)
so the FLOP rate is not confused with the tokens/s throughput line, and
so the code matches the current examples/odc docs. Also refresh the
render-format docstrings to match the actual output.

Co-authored-by: Cursor <cursoragent@cursor.com>
The training-log patch now renders "compute per GPU (TFLOP/s/GPU): X
(avg Y)" and "tokens/s/GPU inst/harmonic mean: X/Y". Update the three
log scrapers (daily_report, auto_benchmark/metrics, benchmark_report) to
accept both these current labels and the legacy "throughput per GPU
(...): X/Y" / "tokens per GPU (tokens/s/GPU): X/Y" forms, so old and new
logs both parse. Label/separator alternations use non-capturing groups
to keep positional group indices stable. In daily_report the average is
now optional and falls back to the instantaneous value for raw
(un-averaged) Megatron lines.

Co-authored-by: Cursor <cursoragent@cursor.com>
The training log now appends memory last as
"hip mem usage/free/total/usage_ratio: X GB/..." (after tokens) rather
than "mem usages: X" between elapsed and throughput. Match the memory
segment at the end of the line in its current position, accept both the
hip-mem and legacy "mem usages:" labels (capturing only the used value),
and reorder the extracted group indices accordingly.

Co-authored-by: Cursor <cursoragent@cursor.com>
Real Megatron pre-train logs render the compute average as
"X(avg Y)" (the space before "(avg" can be dropped by the log sink),
which the previous " \(avg " literal did not match. Make the space
optional so benchmark_report parses real logs, matching the tolerance
already present in daily_report and auto_benchmark/metrics.

Co-authored-by: Cursor <cursoragent@cursor.com>
amd-fuyuajin and others added 2 commits July 27, 2026 22:40
Black joins the two comment-adjacent implicitly-concatenated regex
fragments onto one line (line-length 110). No functional change; the
concatenated pattern string is identical.

Co-authored-by: Cursor <cursoragent@cursor.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant