NVIDIA / NVIDIA/Megatron-LM

Add per-layer measured logging for activation memory and forward/backward time

Open
#5,595 5 comments 0 reactions 0 assignees View on GitHub
community-request enhancement waiting-on-maintainers
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Is your feature request related to a problem? Please describe.**
This item is listed in the Q2 2026 roadmap (#4997) under Ease of Use ("Per-layer logging and memory estimation").

When planning a model configuration against a fixed memory budget, or debugging which layers dominate step latency, Megatron currently offers no on-line, per-layer view of measured resource usage. The existing tools each cover a different part of the space but leave a gap:

report_memory (megatron/training/utils/common_utils.py) reports process-global CUDA memory (memory_allocated / max_memory_allocated / reserved), not per-layer attribution.
timers.py (megatron/core/timers.py) provides operation-level timing (e.g. optimizer-copy-to-main-grad, params-all-gather), but nothing at transformer-layer granularity.
theoretical_memory_usage.py provides detailed analytical per-layer estimates, but they are never reconciled against measured values (the activation path relies on a hardcoded overhead_factor = 1.05).
nvtx instrumentation exists, but requires an external profiler (nsys) and offline analysis, and the existing ranges cover specific ops (e.g. gated_delta_net, fused swiglu) rather than systematically per transformer layer.

None of these give an in-framework, per-layer breakdown of measured memory and time during a normal training run.

**Describe the solution you'd like**
A config-gated, measured per-layer profiling utility, scoped to keep the first contribution reviewable:

1. Per-layer logging of allocated activation memory and forward/backward wall-clock time, collected via module hooks on TransformerLayer (megatron/core/transformer/transformer_layer.py), iterating over layers in TransformerBlock.
2. Reuse the existing CUDA memory-query pattern from report_memory and the timing conventions in timers.py rather than introducing a parallel path.
3. Gated behind a config flag, disabled by default, with negligible overhead when off; structured per-layer summary emitted at a configurable interval (consistent with log_memory_interval).

First PR scoped to single-device measured logging. A natural follow-up would be a theoretical-vs-measured comparison view that reconciles these numbers against theoretical_memory_usage.py — which could, in turn, help calibrate the hardcoded overhead factor.

A few questions before I start:
1. Preferred location — under megatron/core/ profiling utilities, or alongside the existing logging code?
2. Any prior design discussion or owner I should align with?
3. is there interest in a measured counterpart to theoretical_memory_usage.py, and would reconciling the two (to calibrate the overhead factor) be in scope?

**Describe alternatives you've considered**
1. External profiling via existing nvtx ranges + nsys: works, but requires external tooling and offline analysis, isn't integrated with training logs or the config system, and current ranges aren't systematically per-layer.
2. Existing report_memory: process-global only, no per-layer attribution.
3. Existing theoretical_memory_usage.py: analytical per-layer estimates, but purely theoretical and never validated against measurements.

An in-framework, config-gated, per-layer measured logger fills the gap none of these cover: on-line, per-layer, no external tooling required.

**Additional context**
A few questions:
1. Is there a preferred location for this — under megatron/core/ profiling utilities, or alongside the existing logging in megatron/training/?
2. Is a measured counterpart to theoretical_memory_usage.py (and reconciling the two) of interest, or should the first PR stay strictly to per-layer logging?
3. Any prior design discussion or owner I should align with?
Validation hardware available: single cloud GPU (A100 / 4090). Happy to open a draft PR once scope is confirmed.

Contributor guide

Open the contributing guide

Research direction

Start by reading megatron/core/transformer/transformer_layer.py and the TransformerBlock layer iteration, then compare the CUDA memory-query pattern in megatron/training/utils/common_utils.py with timing conventions in megatron/core/timers.py. Check how log_memory_interval and configuration flags are defined. Done means a disabled-by-default, single-device utility emits configurable-interval per-layer activation-memory and forward/backward timing summaries during training.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, observability, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.