Nemo Mbridge Perf Memory Tuning

NVIDIAOfficialPublished by NVIDIA

Skill
Guidelines
Software development
Apache-2.0
Medium risk

From the author

Techniques for reducing peak GPU memory in Megatron Bridge — expandable segments, PEFT + SP input re-gather, parallelism resizing, activation recompute, CPU offloading constraints, and common OOM fixes.

— the description by NVIDIA

AutoGineer summary

Written by AutoGineer from the skill’s own files, not by the author

This skill provides a systematic approach to troubleshooting and resolving GPU memory issues in Megatron Bridge training environments. It distinguishes between memory fragmentation and genuine capacity limits, offering specific configuration adjustments to optimize memory usage. Key techniques include enabling expandable CUDA segments, configuring LoRA sequence-parallel input re-gathering, and understanding the constraints of parallelism resizing and CPU offloading.

What it does

Identifies memory fragmentation as a primary cause of OOM errors.
Provides environment variables to optimize the PyTorch CUDA allocator.
Details configuration changes for LoRA training to reduce activation memory.
Offers guidance on adjusting parallelism strategies to balance memory and throughput.
Includes a theoretical memory estimator for pre-launch configuration planning.
Documents compatibility constraints for CPU offloading and NCCL user-buffer registration.

When to use it

When encountering GPU out-of-memory (OOM) errors during training.
When attempting to reduce peak GPU memory usage for large model training.
When training LoRA models with sequence parallelism enabled.
When investigating OOM regressions following configuration changes.

Advantages

Reduces OOM errors caused by fragmentation without performance loss.
Provides memory-saving options for PEFT/LoRA workflows.
Offers clear decision-making criteria to avoid unnecessary throughput degradation.
Includes validation steps to ensure memory optimizations are correctly applied.

Usage examples

“How do I fix a persistent OOM error in my Megatron Bridge training job?”: The agent recommends setting the environment variable PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True as the first step to resolve potential memory fragmentation.
“I am training a LoRA model with sequence parallelism and running out of memory. What can I do?”: The agent suggests enabling sequence_parallel_input_regather=True in the LoRA configuration to reduce the retention of gathered inputs.

Install

Download the skill as a ZIP:
https://autogineer.com/agent-skills/nvidia/nemo-mbridge-perf-memory-tuning/download

Unzip it into the cross-client skills folder of your project:

mkdir -p .agents/skills && unzip nemo-mbridge-perf-memory-tuning.zip -d .agents/skills

Documentation by NVIDIA

Memory Tuning

Stable docs: @docs/parallelisms.md Card: @skills/nemo-mbridge-perf-memory-tuning/card.yaml

What It Is

GPU OOM failures during training often stem from memory fragmentation rather than raw capacity. PyTorch's default CUDA allocator can leave unusable gaps between allocations. The single most effective fix is:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

This tells PyTorch to use expandable (non-fixed-size) memory segments, which dramatically reduces fragmentation and often eliminates borderline OOM without any model or parallelism changes.

Beyond fragmentation, actual peak memory is determined by:

Parameter + optimizer state memory — controlled by TP, PP, DP sharding (distributed optimizer, FSDP)
Activation memory — controlled by activation recompute, sequence length, micro-batch size, and PEFT-specific retention of gathered inputs
Temporary / workspace memory — CUDA kernels, NCCL buffers, CUDA graphs

For configuration planning, use the Bridge theoretical estimator before launching large jobs:

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:

Set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True first. This fixes fragmentation-induced OOM with zero performance cost. Most Slurm launch templates already include it.
For LoRA with sequence parallelism, enable input re-gather (LoRA(sequence_parallel_input_regather=True)). This avoids retaining the full gathered LoRA-A input in every eligible layer; it has no effect when SP is disabled.
Add selective activation recompute (recompute_modules=[core_attn]) if not already enabled. See @skills/nemo-mbridge-perf-activation-recompute/SKILL.md.
Avoid increasing TP as a memory fix — doubling TP dramatically increases NVLink all-reduce volume and often kills throughput (-28% on Llama3 70B).
Avoid increasing PP at the cost of DP — halving DP doubles gradient accumulation steps and hurts throughput (~6%).
Consider mlp recompute if still OOM. Saves ~3 GB but costs ~16% GPU utilization on large dense models (Llama3 70B).
CPU offloading is blocked when PP > 1.

Enablement

Expandable segments (recommended first step)

Set in the job's environment before launching:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

In Slurm scripts this is typically placed alongside other env vars:

export CUDA_DEVICE_MAX_CONNECTIONS=1
export NVTE_ALLOW_NONDETERMINISTIC_ALGO=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

No model config changes needed. Zero throughput cost.

Parallelism resizing

If the model genuinely does not fit (not fragmentation), adjust parallelism:

Strategy	Memory effect	Throughput cost	Notes
Increase PP (keeping DP)	Fewer layers per stage	Moderate (~6% if DP halved)	Only if GPU count allows
Increase TP	Fewer params per GPU	Severe (-28% on 70B)	Last resort
Distributed optimizer	Shards optimizer state across DP ranks	~1-2%	Recommended for large models
FSDP	Shards params + grads + optimizer	Varies	See @skills/nemo-mbridge-perf-megatron-fsdp/SKILL.md

Activation recompute

See @skills/nemo-mbridge-perf-activation-recompute/SKILL.md for full details.

PEFT + sequence-parallel input re-gather

For LoRA training with sequence parallelism, eligible column-parallel linear_qkv and linear_fc1 adapters consume a gathered LayerNorm output. Because LoRA-A is trainable, the default path retains that full gathered input until backward for the LoRA-A weight gradient.

Enable input re-gather when constructing the PEFT config:

from megatron.bridge.peft.lora import LoRA

cfg.peft = LoRA(
    # Keep the recipe's existing LoRA settings here.
    sequence_parallel_input_regather=True,
)

With this option, forward still materializes the full input temporarily for the LoRA-A GEMM, but MCore autograd retains only its sequence-local shard. Backward asynchronously gathers the full input again, overlaps the collective with dgrad when possible, computes the LoRA-A weight gradient, and then reuses the temporary communication buffer.

This is a memory-for-communication tradeoff, not conventional activation checkpointing: no LayerNorm, attention, MLP, or LoRA GEMM is rerun. Some throughput degradation is expected, and the benefit grows with the amount of eligible LoRA-A activation retained. The option has no effect when sequence parallelism is disabled.

CPU offloading

cfg.model.cpu_offloading = True

Incompatible with PP > 1. Only usable when pipeline_model_parallel_size = 1.

A Note on VPP

Virtual pipeline parallelism (VPP) is primarily a throughput optimization that reduces pipeline bubble overhead by interleaving smaller model chunks. Its effect on peak memory is minimal — changing VPP does not meaningfully change the total activation, parameter, or optimizer memory on a GPU.

In earlier experiments we incorrectly attributed an OOM fix to VPP tuning (VPP 5→10). The actual fix was PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True which eliminated memory fragmentation. The VPP=10 run actually used slightly more peak memory (60.2 GB vs 58.8 GB) but did not OOM because expandable segments prevented fragmentation.

VPP should be tuned for pipeline bubble reduction (see @docs/parallelisms.md), not as a memory fix.

Compatibility and Constraints

expandable_segments:True is incompatible with --use-nccl-ub (NCCL user-buffer registration). See Megatron-FSDP docs.
When using CUDA graphs with expandable_segments:True, set NCCL_GRAPH_REGISTER=0 (required on pre-Blackwell GPUs, enforced by MCore CudaGraphManager).
CPU offloading requires pipeline_model_parallel_size = 1.
Distributed optimizer requires use_distributed_optimizer = True in the optimizer config.
sequence_parallel_input_regather applies only to eligible non-expert column-parallel LoRA-A projections. Row-parallel adapters, expert adapters, TP=1, CUDA graphs, CPU activation offload, and overlapping full-layer or selective MLP activation recompute fall back to the existing path.

Measured Results

Llama3 70B SFT on 32x H100 80GB, FP8 (Current Scaling):

Baseline: TP=4, PP=4, VPP=5, DP=2, MBS=1, GBS=32, seq_len=4096
Golden GPU utilization: 709.93 TFLOP/s/GPU
Regression threshold: 5%

Strategy comparison: parallelism changes for memory reduction

Experiment	TP	PP	VPP	DP	TFLOP/s/GPU	vs Golden	Peak Mem (GB)	Result
Baseline	4	4	5	2	~704	-0.8%	58.8	OOM (fragmentation)
More PP	4	8	5	1	668.0	-5.9%	53.2	Borderline perf
More TP	8	4	5	1	508.7	-28.4%	50.2	Severe regression
Baseline + expandable_segments	4	4	5	2	~704	-0.8%	~59	Passed

Key takeaways:

expandable_segments:True is the winner. The baseline OOM was caused by memory fragmentation, not insufficient capacity. Setting this env var eliminated the OOM with zero throughput cost and no parallelism changes.
PP=8 works for memory but loses DP (2→1), meaning 32 gradient accumulation steps per batch, which hurts throughput by ~6%.
TP=8 is catastrophic (-28%) because doubling TP increases all-reduce communication volume proportionally across NVLink, and DP=1 means no micro-batch overlap.

CPU offloading: blocked

Experiment	offload_layers	Result
Exp 4	2	Incompatible (PP > 1)
Exp 5	4	Incompatible (PP > 1)
Exp 6	6	Incompatible (PP > 1)

ValueError: Currently there is no support for Pipeline parallelism with CPU offloading. This approach is blocked for any model using PP > 1.

Activation recompute: expensive alternative

Selective activation recompute with mlp saved ~3 GB peak memory but cost ~16% GPU utilization on this workload. See @skills/nemo-mbridge-perf-activation-recompute/SKILL.md for full results.

LoRA + SP input re-gather

Real-checkpoint H100 training with SQuAD showed lower peak memory in all tested configurations, with workload-dependent throughput cost:

Model/config	Baseline peak	Input re-gather peak	Memory saved	Throughput change
Qwen3-8B, TP2, seq 8192	47.545 GB	42.814 GB	4.731 GB (10.0%)	-6.74%
Qwen3-30B-A3B, TP4/EP4	29.890 GB	28.321 GB	1.569 GB (5.2%)	-2.89%
GPT-OSS-120B, TP2/EP8	52.185 GB	51.371 GB	0.814 GB (1.6%)	-0.34%

All runs had finite losses with zero skipped or NaN iterations. Two-rank BF16 and FP32 checks matched the baseline for outputs, input gradients, LoRA-A and LoRA-B gradients, and two-microbatch fused main_grad accumulation.

Code Anchors

LoRA sequence-parallel input re-gather

src/megatron/bridge/peft/lora.py
    LoRA.sequence_parallel_input_regather

src/megatron/bridge/peft/utils.py
    ParallelLinearAdapter._sequence_parallel_input_regather_eligibility()
    ParallelLinearAdapter.forward()

CPU offloading PP incompatibility (MCore)

        if self.cpu_offloading and self.pipeline_model_parallel_size > 1:
            raise ValueError(
                "Currently there is no support for Pipeline parallelism with CPU offloading"
            )

VPP config and layer divisibility validation (MCore)

            if pipeline_parallel_size and self.virtual_pipeline_model_parallel_size is not None:
                num_layers_per_middle_pipeline_rank = num_layers // pipeline_parallel_size
                if (
                    not num_layers_per_middle_pipeline_rank
                    % self.virtual_pipeline_model_parallel_size
                    == 0
                ):
                    raise ValueError(
                        f"number of layers on each middle pipeline rank:"
                        f"{num_layers_per_middle_pipeline_rank} must be divisible by virtual"
                        f"pipeline parallel degree {self.virtual_pipeline_model_parallel_size}"
                    )

Parallelism docs on interleaved pipeline schedule

To minimize the pipeline bubble, the computation on each GPU can be divided into multiple subsets of layers (referred to as model chunks), rather than a single contiguous block. Enable this by setting `virtual_pipeline_model_parallel_size`:

model_config = GPTModelProvider(
    pipeline_model_parallel_size=4,
    virtual_pipeline_model_parallel_size=2,  # 2 model chunks per pipeline stage
    # ... other model parameters
)

Failure Diagnosis

Symptom	Cause	Confirm	Fix
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
LoRA + SP retains unexpectedly high activation memory	full gathered LoRA-A inputs are retained until backward	check whether `cfg.peft.sequence_parallel_input_regather` is enabled and the target is eligible	set `LoRA(sequence_parallel_input_regather=True)`; verify fallback constraints
`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
The theoretical estimator is formula-based and does not replace runtime profiling or CUDA memory reports
LoRA input re-gather does not cover row-parallel or expert adapters and may have negligible benefit when few eligible LoRA-A activations dominate memory

Verification

Quick check that expandable_segments:True is active:

import os
assert "expandable_segments:True" in os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")

For Slurm jobs, verify the env var is exported before the training command in the launch script.

For LoRA + SP input re-gather, run the focused configuration tests and the real two-rank MCore backward-parity test:

uv run python -m pytest \
  tests/unit_tests/peft/test_utils.py -k "sequence_parallel_input_regather" \
  tests/unit_tests/peft/test_lora.py -k "sequence_parallel_input_regather"

uv run python -m torch.distributed.run --nproc_per_node=2 -m pytest \
  tests/unit_tests/peft/test_lora_sp_input_regather_distributed.py

Files

ATTRIBUTION.mdAdded by AutoGineer861 B
BENCHMARK.md3.9 KiB
LICENSELicence file added by AutoGineer11.2 KiB
SKILL.md14.1 KiB
card.yaml9.8 KiB
evals/evals.json1.4 KiB
skill-card.md4.1 KiB
skill.oms.sig4.7 KiB

card.yaml

title: memory_tuning
validated_on: "2026-06-29"
summary: >
  Techniques for reducing peak GPU memory to fix OOM or increase headroom.
  Current coverage: expandable segments (fragmentation fix), parallelism
  resizing, activation recompute, LoRA sequence-parallel input re-gather, and
  CPU offloading constraints. The most
  common OOM fix is setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,
  which eliminates memory fragmentation at zero throughput cost. Measured on
  Llama3 70B SFT (32x H100, FP8 CS): expandable_segments eliminated the
  baseline OOM without any parallelism changes. By contrast, doubling PP
  (4→8) cost ~6% and halved DP, while doubling TP (4→8) caused -28%
  regression. CPU offloading is blocked when PP > 1.
validation_status:
  expandable_segments:
    - verified  # standard PyTorch allocator option
  cpu_offloading_pp_incompatibility:
    - code_verified  # MCore transformer_config.py
  llama3_70b_sft_fp8_cs_experiment:
    - measured  # PR #3107, 32x H100, TP4 PP4 DP2
  lora_sequence_parallel_input_regather:
    - measured  # PR #4581, real checkpoints and SQuAD on H100
    - parity_verified  # TP2 BF16/FP32 and fused two-microbatch main_grad
training_dimensions:
  speed:
    effect: "zero cost for expandable_segments; LoRA input re-gather and parallelism changes vary"
    confidence: high
    rationale: >
      expandable_segments:True has no throughput impact. Parallelism resizing
      (TP, PP) can cost 6-28% depending on the strategy. Measured LoRA input
      re-gather cost ranged from 0.34% to 6.74% depending on the workload.
  memory:
    effect: "eliminates fragmentation-induced OOM; reduces retained LoRA + SP activations"
    confidence: high
    rationale: >
      The baseline OOM at 58.8 GB was caused by memory fragmentation, not
      insufficient capacity. expandable_segments:True eliminated the OOM
      without any model or parallelism changes. VPP changes do not
      meaningfully affect peak memory. LoRA input re-gather saved 0.8-4.7 GB
      in measured TP+SP workloads.
  scale:
    effect: "neutral for expandable_segments; LoRA input re-gather requires SP"
    confidence: high
    rationale: >
      expandable_segments is a per-process allocator setting with no
      distributed implications. LoRA input re-gather has no effect when
      sequence parallelism is disabled.
  convergence:
    effect: "no change expected"
    confidence: high
    rationale: >
      Allocator settings do not affect computation or numerics. LoRA input
      re-gather matched baseline outputs and gradients in TP2 parity tests.
  stability:
    effect: "improved — reduces fragmentation-induced OOM"
    confidence: high
    rationale: >
      expandable_segments reduces the likelihood of OOM from memory
      fragmentation, which is the most common cause of borderline OOM.
enable_when:
  - training OOMs or is close to the GPU memory limit
  - OOM occurs on a single rank while others have headroom (fragmentation)
  - PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True is not yet set
  - LoRA or PEFT training uses sequence parallelism and needs lower activation memory
avoid_when:
  - using --use-nccl-ub (NCCL user-buffer registration is incompatible with expandable allocator)
interactions:
  required: []
  conditional:
    - "with CUDA graphs: set NCCL_GRAPH_REGISTER=0 when using expandable_segments (required on pre-Blackwell GPUs)"
    - "with NCCL UB: expandable_segments is incompatible with --use-nccl-ub"
    - "with LoRA + SP: sequence_parallel_input_regather trades a backward all-gather for lower retained activation memory"
  incompatible:
    - "--use-nccl-ub with expandable_segments:True"
feature_meaning:
  PYTORCH_CUDA_ALLOC_CONF: >
    PyTorch CUDA allocator configuration. expandable_segments:True uses
    expandable memory segments that reduce fragmentation by allowing the
    allocator to grow segments rather than allocating fixed-size blocks.
  sequence_parallel_input_regather: >
    LoRA option that retains the sequence-local LoRA-A input instead of the
    full gathered input, then re-gathers it during backward for the LoRA-A
    weight gradient. It has no effect when sequence parallelism is disabled.
config_keys:
  - "env: PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True"
  - "peft: LoRA.sequence_parallel_input_regather=True"
recommended_path:
  first_try: "set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True"
  second: "for LoRA + SP, enable sequence_parallel_input_regather"
  third: "add selective activation recompute (core_attn)"
  fourth: "increase PP or use distributed optimizer"
  last_resort: "increase TP (severe throughput cost)"
expected_metric_change:
  - metric: peak_memory
    direction: down
    magnitude: "eliminates fragmentation overhead (variable, often 2-10 GB effective)"
    conditions: any model, any GPU
    evidence: measured_pr_3107
  - metric: gpu_utilization
    direction: unchanged
    magnitude: "0%"
    conditions: expandable_segments only
    evidence: measured_pr_3107
  - metric: peak_memory
    direction: down
    magnitude: "0.8-4.7 GB measured"
    conditions: LoRA with eligible column-parallel adapters and sequence parallelism
    evidence: measured_pr_4581
  - metric: throughput
    direction: down
    magnitude: "0.34-6.74% measured"
    conditions: LoRA sequence-parallel input re-gather
    evidence: measured_pr_4581
measured_results:
  - model: Llama3 70B
    task: sft
    precision: FP8_CS
    gpus: 32
    gpu: H100_80GB
    seq_length: 4096
    mbs: 1
    gbs: 32
    golden_tflops: 709.93
    regression_threshold_pct: 5
    experiments:
      - name: baseline_without_expandable_segments
        tp: 4
        pp: 4
        vpp: 5
        dp: 2
        tflops: 704
        vs_golden_pct: -0.8
        peak_mem_gb: 58.8
        status: OOM
        note: "OOM caused by memory fragmentation"
      - name: baseline_with_expandable_segments
        tp: 4
        pp: 4
        vpp: 5
        dp: 2
        tflops: 704
        vs_golden_pct: -0.8
        peak_mem_gb: ~59
        status: passed
        note: "expandable_segments eliminated fragmentation-induced OOM"
      - name: more_pp
        tp: 4
        pp: 8
        vpp: 5
        dp: 1
        tflops: 668.0
        vs_golden_pct: -5.9
        peak_mem_gb: 53.2
        status: passed_mem_borderline_perf
        note: "halved DP (2→1) means 32 gradient accumulation steps"
      - name: more_tp
        tp: 8
        pp: 4
        vpp: 5
        dp: 1
        tflops: 508.7
        vs_golden_pct: -28.4
        peak_mem_gb: 50.2
        status: severe_perf_regression
        note: "doubling TP increases all-reduce comm volume, DP=1 means no micro-batch overlap"
  - name: cpu_offloading_blocked
    note: >
      CPU activation offloading (cpu_offloading=True) was tested as an
      alternative but is incompatible with PP > 1. All experiments raised
      ValueError. This approach is blocked for any model using PP > 1.
  - name: lora_sequence_parallel_input_regather
    task: peft_sft
    gpu: H100
    dataset: SQuAD
    experiments:
      - model: Qwen3-8B
        parallelism: TP2
        seq_length: 8192
        baseline_peak_mem_gb: 47.545
        regather_peak_mem_gb: 42.814
        throughput_change_pct: -6.74
      - model: Qwen3-30B-A3B
        parallelism: TP4_EP4
        baseline_peak_mem_gb: 29.890
        regather_peak_mem_gb: 28.321
        throughput_change_pct: -2.89
      - model: GPT-OSS-120B
        parallelism: TP2_EP8
        baseline_peak_mem_gb: 52.185
        regather_peak_mem_gb: 51.371
        throughput_change_pct: -0.34
failure_modes:
  - name: fragmentation_oom
    symptom: "OOM on a single rank despite headroom on other ranks"
    likely_cause: memory fragmentation from default allocator
    fix: "set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True"
  - name: expandable_segments_nccl_ub_conflict
    symptom: "RuntimeError with --use-nccl-ub"
    likely_cause: expandable allocator incompatible with NCCL UB registration
    fix: "remove expandable_segments:True or disable --use-nccl-ub"
  - name: pp_increase_loses_dp
    symptom: throughput drop despite lower memory
    likely_cause: increasing PP reduces DP, increasing gradient accumulation steps
    fix: try expandable_segments first before resizing parallelism
  - name: lora_sp_gathered_input_retention
    symptom: "LoRA + SP activation memory remains high"
    likely_cause: full gathered LoRA-A inputs are retained until backward
    fix: "set LoRA(sequence_parallel_input_regather=True) and verify the target is eligible"
known_constraints:
  - expandable_segments is incompatible with --use-nccl-ub
  - CPU offloading requires pipeline_model_parallel_size = 1
  - with CUDA graphs, set NCCL_GRAPH_REGISTER=0 alongside expandable_segments
  - sequence_parallel_input_regather has no effect when sequence parallelism is disabled
  - sequence_parallel_input_regather falls back for row-parallel, expert, TP1, CUDA graph, CPU-offload, and overlapping recompute paths
known_limitations:
  - expandable_segments fixes fragmentation but not genuine capacity limits
  - parallelism resizing (TP/PP) has significant throughput costs
  - no automatic memory profiling to recommend the optimal strategy
  - LoRA input re-gather adds a backward all-gather and may not improve throughput
  - LoRA input re-gather applies only to eligible non-expert column-parallel LoRA-A projections
evidence:
  - docs/parallelisms.md
  - docs/performance-guide.md
  - "PR #3107 (Llama3 70B SFT OOM fix experiment)"
  - "PR #4581 (LoRA sequence-parallel input re-gather experiments)"
  - tests/unit_tests/peft/test_lora_sp_input_regather_distributed.py
  - 3rdparty/Megatron-LM/megatron/core/transformer/moe/README.md
follow_up_validation:
  - Quantify expandable_segments memory savings across different model sizes.
  - Measure interaction between expandable_segments and CUDA graph memory overhead.
  - Measure LoRA input re-gather across more sequence lengths, TP sizes, and target-module sets.

Requirements and compatibility

Works with: Claude Code, Claude apps, OpenAI Codex, GitHub Copilot, Cursor, Gemini CLI, Kiro, OpenCode, Goose, Amp, Roo Code, JetBrains Junie, Factory Droid, Devin (ex-Windsurf), Mistral Vibe, OpenHands, Letta

“Works with” means the tool reads the open Agent Skills format; a skill written for one tool may use features others lack.

Needs: Needs Python

Security

What our scanner noticed· Medium risk

Large base64-encoded blobs

A static scan is not a guarantee. Skills are third-party content: review a skill before you let an agent use it.

Attribution and licence

Published by

NVIDIA

Publisher

NVIDIA (Official)

Source

NVIDIA/skills · skills/nemo-mbridge-perf-memory-tuning · commit fce3cca

Licence

Apache License 2.0 (redistribution allowed)

LICENSE (added by AutoGineer from LICENSE-APACHE)
ATTRIBUTION.md (written by AutoGineer)

Synced

2026-11-09

AutoGineer is not affiliated with or endorsed by NVIDIA or NVIDIA.

Download ZIP

The ZIP holds the skill’s files, its licence files and an ATTRIBUTION.md. Download nemo-mbridge-perf-memory-tuning.zip