[Training] feat: add memory estimator breakdown

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

docs/index.md

@@ -47,6 +47,7 @@ training/training-loop-settings.md
 training/optimizer-scheduler.md
 training/logging.md
 training/profiling.md
+training/memory-estimator.md
 training/checkpointing.md
 training/megatron-fsdp.md
 training/resiliency.md

docs/training/README.md

@@ -52,6 +52,7 @@ This directory contains comprehensive documentation for training and customizing
 |----------|---------|--------------|
 | **[Logging](logging.md)** | Logging configuration and TensorBoard/WandB integration | Monitoring training progress |
 | **[Profiling](profiling.md)** | Performance profiling and analysis | Identifying bottlenecks |
+| **[Theoretical Memory Estimator](memory-estimator.md)** | Formula-based per-GPU memory planning | Sizing dense and MoE training configs |
 | **[Resiliency](resiliency.md)** | Handling failures and recovery | Building robust training pipelines |
 
 ### Advanced Features

docs/training/memory-estimator.md

@@ -0,0 +1,73 @@
+# Theoretical Memory Estimator
+
+Megatron Bridge includes a formula-based training memory estimator for GPT-like
+model providers. It is intended for early configuration planning and for the
+training-time theoretical memory report printed alongside CUDA memory metrics.
+
+The estimator is available as a Python API:
+
+```python
+from megatron.bridge.training.utils.theoretical_memory_utils import (
+    estimate_training_memory,
+    format_training_memory_estimate,
+)
+
+estimate = estimate_training_memory(cfg, num_microbatches=num_microbatches)
+print(format_training_memory_estimate(estimate, unit="GB"))
+```
+
+`estimate_training_memory` returns a structured result with:
+
+- dense and embedding parameter memory on the most-loaded GPU shard
+- routed MoE expert parameter memory when `num_moe_experts` is set
+- activation memory when requested
+- total global parameter count covered by the estimator
+- assumptions attached to the estimate
+
+The training loop continues to use the same report hook:
+
+```text
+Theoretical memory footprints: weight and optimizer=..., activation=..., total=...
+```
+
+## What It Covers
+
+The model-state estimate covers weights, gradients, FP32 master weights, and
+Adam optimizer states. It uses the same byte model as the legacy utility:
+
+- `18` bytes per parameter when the distributed optimizer is disabled
+- `6 + 12 / shard_size` bytes per parameter when the distributed optimizer is enabled
+
+For dense parameters, `shard_size` is `data_parallel_size * context_parallel_size`.
+For routed MoE experts, expert parameters are divided by
+`expert_model_parallel_size * expert_tensor_parallel_size`, and optimizer state
+uses the expert data-parallel shard size, including context parallel ranks.
+
+The activation estimate follows the existing Megatron-LM theoretical formula and
+adds Bridge-aware accounting for MoE active expert width and context parallel
+partitioning. It is most meaningful when sequence parallelism and selective
+activation recomputation are enabled, which is also the condition used by the
+training-time report.
+
+## Assumptions
+
+The estimate reports the most-loaded GPU shard. Embeddings are assigned to the
+first and last pipeline stages; untied input/output embeddings both land on the
+same GPU only when `pipeline_model_parallel_size == 1`.
+
+The estimator does not model runtime allocator fragmentation, CUDA kernel
+workspace, NCCL buffers, CUDA graph static buffers, token-routing imbalance,
+dispatcher workspace, CPU offloading, or full activation recomputation. Use
+profiling and CUDA memory reports to validate final launch configurations.
+
+## MoE Notes
+
+MoE layer count follows `moe_layer_freq`, including list patterns. Routed expert
+parameters are counted separately from dense attention, dense MLP, shared expert,
+normalization, and embedding parameters. This keeps expert parallelism effects
+visible in the returned component breakdown.
+
+For memory tuning decisions, use this estimator as a planning signal together
+with the runtime guidance in [CPU Offloading](cpu-offloading.md),
+[Activation Recomputation](activation-recomputation.md), and
+[Megatron FSDP](megatron-fsdp.md).

skills/perf-memory-tuning/SKILL.md

@@ -31,6 +31,20 @@ Beyond fragmentation, actual peak memory is determined by:
   micro-batch size
 - **Temporary / workspace memory** — CUDA kernels, NCCL buffers, CUDA graphs
 
+For configuration planning, use the Bridge theoretical estimator before launching
+large jobs:
+
+```python
+from megatron.bridge.training.utils.theoretical_memory_utils import estimate_training_memory
+
+estimate = estimate_training_memory(cfg, num_microbatches=num_microbatches)
+```
+
+The estimator reports the most-loaded GPU shard and separates dense/embedding,
+routed MoE expert, and activation components. It does not include allocator
+fragmentation, CUDA/NCCL workspace, CUDA graph buffers, token imbalance, or
+dispatcher workspace, so validate final configs with runtime memory metrics.
+
 ## Quick Decision
 
 When a training run OOMs or is close to the memory limit:
@@ -208,14 +222,16 @@ model_config = GPTModelProvider(
 |---|---|---|---|
 | OOM on a single rank despite headroom on others | Memory fragmentation | check if `expandable_segments:True` is set | set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` |
 | OOM with `expandable_segments` already set | Genuine capacity limit | check `nvidia-smi` for param/optimizer memory | increase PP, use distributed optimizer, or add recompute |
+| Estimated memory exceeds GPU capacity before launch | model state or activations genuinely too large | run `estimate_training_memory` and inspect the largest component | adjust PP/TP/CP/EP, distributed optimizer, or recompute before launching |
 | `ValueError: PP + CPU offloading` | using cpu_offloading with PP > 1 | check PP config | disable CPU offloading or set PP=1 |
 | `RuntimeError` with `--use-nccl-ub` + expandable segments | NCCL UB incompatible with expandable allocator | check env vars | remove `expandable_segments:True` or disable `--use-nccl-ub` |
 
 ## Known Limitations
 
 - CPU offloading is blocked when PP > 1
 - Parallelism resizing (TP/PP) often has significant throughput costs
-- No automatic memory profiling to recommend the optimal strategy
+- The theoretical estimator is formula-based and does not replace runtime
+  profiling or CUDA memory reports
 
 ## Verification