nemo-mbridge-perf-memory-tuning

What problem does it solve?

GPU training with Megatron Bridge often hits out-of-memory due to memory fragmentation and suboptimal allocator behavior. The skill introduces fixes like using expandable memory segments and guidelines to tune parallelism, activation recompute, and CPU offloading to reduce peak memory without sacrificing throughput.

Core Features & Use Cases

• Expandable segments to reduce memory fragmentation and prevent OOMs with zero throughput cost.
• Memory planning guidance: theoretical memory estimator, and activation recompute as a secondary option.
• Compatibility constraints: CPU offloading is blocked when pipeline model parallel size > 1; CUDA graph considerations.

Quick Start

Before launching training, set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to enable memory fragmentation mitigation and validate that the OOM does not occur.