You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Commit d06c130
Browse filesBrowse the repository at this point in the historyBrowse files
> 3. LoRA training with tensor and sequence parallelism can retain the sequence-local LoRA-A input for
324
+
> column-parallel `linear_qkv` and `linear_fc1` adapters. Set `LoRA(sequence_parallel_input_regather=True)` to
325
+
> gather the full LayerNorm output temporarily in forward, release it after LoRA-A, and gather it again during
326
+
> backward for the LoRA-A weight gradient. MCore overlaps the backward all-gather with dgrad computation when
327
+
> possible. The option has no effect when sequence parallelism is disabled and falls back to the existing path
328
+
> for row-parallel or expert adapters, CPU activation offload, CUDA graphs, full-layer recompute, and `linear_fc1`
329
+
> covered by selective MLP recompute.
330
+
323
331
2. Activation offloading to host memory
324
332
325
333
> 1. Megatron-Bridge supports offloading activation memory to host memory, essential for training tasks constrained by activation memory. This is particularly useful for scenarios like (1) FSDP, where model state memory is minimized through sharding but activation memory remains high, (2) LoRA, which has frozen parameters but significant activation memory demands, and (3) the training with a large sequence length. The efficiency of activation offloading depends on both the interconnect bandwidth between the GPU and host and the host memory bandwidth. From this perspective, Grace-based systems like the GB200 enhance offloading performance by optimizing these bandwidths.
Copy file name to clipboardExpand all lines: skills/nemo-mbridge-perf-memory-tuning/SKILL.md
+87-8Lines changed: 87 additions & 8 deletions
Display the source diff
Display the rich diff
Original file line number
Diff line number
Diff line change
@@ -1,8 +1,8 @@
1
1
---
2
2
name: nemo-mbridge-perf-memory-tuning
3
-
description: Techniques for reducing peak GPU memory in Megatron Bridge — expandable segments, parallelism resizing, activation recompute, CPU offloading constraints, and common OOM fixes.
3
+
description: 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.
4
4
license: Apache-2.0
5
-
when_to_use: GPU OOM errors, reducing peak memory, or tracing an OOM regression to a specific commit or config change; 'out of memory', 'OOM', 'memory fragmentation', 'expandable_segments', 'reduce GPU memory', 'PYTORCH_CUDA_ALLOC_CONF'.
5
+
when_to_use: GPU OOM errors, reducing peak memory, reducing LoRA or PEFT activation memory with sequence parallelism, or tracing an OOM regression to a specific commit or config change; 'out of memory', 'OOM', 'memory fragmentation', 'expandable_segments', 'reduce GPU memory', 'LoRA memory', 'PEFT memory', 'sequence_parallel_input_regather', 'PYTORCH_CUDA_ALLOC_CONF'.
6
6
---
7
7
8
8
# Memory Tuning
@@ -29,7 +29,7 @@ Beyond fragmentation, actual peak memory is determined by:
29
29
-**Parameter + optimizer state memory** — controlled by TP, PP, DP sharding
30
30
(distributed optimizer, FSDP)
31
31
-**Activation memory** — controlled by activation recompute, sequence length,
32
-
micro-batch size
32
+
micro-batch size, and PEFT-specific retention of gathered inputs
33
33
-**Temporary / workspace memory** — CUDA kernels, NCCL buffers, CUDA graphs
34
34
35
35
For configuration planning, use the Bridge theoretical estimator before launching
@@ -53,15 +53,19 @@ When a training run OOMs or is close to the memory limit:
53
53
1.**Set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` first.** This fixes
54
54
fragmentation-induced OOM with zero performance cost. Most Slurm launch
55
55
templates already include it.
56
-
2.**Add selective activation recompute** (`recompute_modules=[core_attn]`) if
56
+
2.**For LoRA with sequence parallelism, enable input re-gather**
57
+
(`LoRA(sequence_parallel_input_regather=True)`). This avoids retaining the
58
+
full gathered LoRA-A input in every eligible layer; it has no effect when SP
59
+
is disabled.
60
+
3.**Add selective activation recompute** (`recompute_modules=[core_attn]`) if
57
61
not already enabled. See @skills/nemo-mbridge-perf-activation-recompute/SKILL.md.
58
-
3.**Avoid increasing TP** as a memory fix — doubling TP dramatically increases
62
+
4.**Avoid increasing TP** as a memory fix — doubling TP dramatically increases
59
63
NVLink all-reduce volume and often kills throughput (-28% on Llama3 70B).
60
-
4.**Avoid increasing PP at the cost of DP** — halving DP doubles gradient
64
+
5.**Avoid increasing PP at the cost of DP** — halving DP doubles gradient
61
65
accumulation steps and hurts throughput (~6%).
62
-
5. Consider `mlp` recompute if still OOM. Saves ~3 GB but costs ~16% GPU
66
+
6. Consider `mlp` recompute if still OOM. Saves ~3 GB but costs ~16% GPU
63
67
utilization on large dense models (Llama3 70B).
64
-
6. CPU offloading is **blocked when PP > 1**.
68
+
7. CPU offloading is **blocked when PP > 1**.
65
69
66
70
## Enablement
67
71
@@ -98,6 +102,36 @@ If the model genuinely does not fit (not fragmentation), adjust parallelism:
98
102
99
103
See @skills/nemo-mbridge-perf-activation-recompute/SKILL.md for full details.
100
104
105
+
### PEFT + sequence-parallel input re-gather
106
+
107
+
For `LoRA` training with sequence parallelism, eligible column-parallel
108
+
`linear_qkv` and `linear_fc1` adapters consume a gathered LayerNorm output.
109
+
Because LoRA-A is trainable, the default path retains that full gathered input
110
+
until backward for the LoRA-A weight gradient.
111
+
112
+
Enable input re-gather when constructing the PEFT config:
113
+
114
+
```python
115
+
from megatron.bridge.peft.lora import LoRA
116
+
117
+
cfg.peft = LoRA(
118
+
# Keep the recipe's existing LoRA settings here.
119
+
sequence_parallel_input_regather=True,
120
+
)
121
+
```
122
+
123
+
With this option, forward still materializes the full input temporarily for the
124
+
LoRA-A GEMM, but MCore autograd retains only its sequence-local shard. Backward
125
+
asynchronously gathers the full input again, overlaps the collective with
126
+
dgrad when possible, computes the LoRA-A weight gradient, and then reuses the
127
+
temporary communication buffer.
128
+
129
+
This is a memory-for-communication tradeoff, not conventional activation
130
+
checkpointing: no LayerNorm, attention, MLP, or LoRA GEMM is rerun. Some
131
+
throughput degradation is expected, and the benefit grows with the amount of
132
+
eligible LoRA-A activation retained. The option has no effect when sequence
133
+
parallelism is disabled.
134
+
101
135
### CPU offloading
102
136
103
137
```python
@@ -132,6 +166,10 @@ not as a memory fix.
132
166
- CPU offloading requires `pipeline_model_parallel_size = 1`.
133
167
- Distributed optimizer requires `use_distributed_optimizer = True` in the
134
168
optimizer config.
169
+
-`sequence_parallel_input_regather` applies only to eligible non-expert
| 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`|
225
289
| OOM with `expandable_segments` already set | Genuine capacity limit | check `nvidia-smi` for param/optimizer memory | increase PP, use distributed optimizer, or add recompute |
226
290
| 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 |
291
+
| 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 |
227
292
|`ValueError: PP + CPU offloading`| using cpu_offloading with PP > 1 | check PP config | disable CPU offloading or set PP=1 |
228
293
|`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`|
0 commit comments