perf-memory-tuning

Configure PyTorch CUDA allocator expandable segments to reduce GPU memory fragmentation in Megatron Bridge training.

Updated May 23, 2026
One-click install
npx skills add https://github.com/yo-steven/skills-exploration-20260522 --skill perf-memory-tuning-yo-steven
Or copy as Structured Prompt for Agent▼
Please help me install this Agent Skill.
Skill: perf-memory-tuning
Source: https://github.com/yo-steven/skills-exploration-20260522/tree/main/skills/Megatron-Bridge/perf-memory-tuning
Command: npx skills add https://github.com/yo-steven/skills-exploration-20260522 --skill perf-memory-tuning-yo-steven

SYSTEM DOCUMENTATION & REQUIREMENTS

What problem does it solve?

Training Megatron Bridge models can fail with out-of-memory (OOM) errors due to GPU memory fragmentation or misbalanced memory usage across parameters, activations, and temporary workspaces.

Core Features & Use Cases

  • Allocator-first OOM fix (expandable segments): Reduce fragmentation by setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, often resolving borderline OOM without changing model or parallelism.
  • Memory-focused parallelism strategy: Adjust TP/PP/DP and distributed optimizer/FSDP choices to change parameter and optimizer memory placement when capacity truly doesn’t fit.
  • Activation memory reduction levers: Use selective activation recompute (referencing perf-activation-recompute) to trade compute for lower activation peak.
  • CPU offloading constraints: Enable cpu_offloading only when pipeline_model_parallel_size=1, since CPU offloading is blocked for PP > 1.

Quick Start

Set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True in your job environment before launching your training run to eliminate fragmentation-related OOM.

Frequently Asked Questions about perf-memory-tuning

High-intent search queries and answers about installing and using this skill.

FAQPage Schema
How do I fix GPU OOM errors in Megatron Bridge training runs?▼

To fix GPU OOM errors in Megatron Bridge, set the environment variable PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True before launching training to reduce memory fragmentation and resolve borderline out-of-memory crashes without changing model configuration.

What causes GPU memory fragmentation OOM on a single rank during distributed training?▼

GPU memory fragmentation occurs when allocator gaps prevent fitting tensors despite having free memory. Enabling PyTorch CUDA expandable segments addresses this fragmentation-induced OOM by allowing the allocator to manage memory regions dynamically.

How do I reduce activation memory peak in Megatron Bridge without changing parallelism?▼

To reduce activation memory peak, use selective activation recompute to trade extra compute for lower memory usage. This avoids changing parallelism configuration while lowering the activation memory footprint during training.

Can I use CPU offloading with pipeline model parallelism in Megatron Bridge?▼

You cannot use CPU offloading with pipeline model parallel size greater than one. CPU offloading in Megatron Bridge is blocked when pipeline_model_parallel_size is greater than one due to compatibility constraints.

When should I adjust TP, PP, or DP parallelism to resolve out-of-memory errors?▼

Adjust tensor, pipeline, or data parallelism to resolve out-of-memory errors when the model capacity truly does not fit in GPU memory. This shifts parameter and optimizer memory placement across devices after allocator tuning proves insufficient.