Skip to content

Commit 995ef0c

Browse files
committed
docs(recipe): allow cheap selective recompute as a starting choice
Signed-off-by: Chen Cui <chcui@nvidia.com>
1 parent f6907f1 commit 995ef0c

4 files changed

Lines changed: 32 additions & 28 deletions

File tree

  • skills
    • create-model-verification-card
    • nemo-mbridge-perf-activation-recompute
    • nemo-mbridge-perf-moe-vlm-training
    • nemo-mbridge-recipe-recommender

‎skills/create-model-verification-card/SKILL.md‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -426,9 +426,11 @@ hardware. It may vary across models or be tuned for throughput:
426426
provided routing, capacity, token dropping, and auxiliary losses are
427427
unchanged.
428428

429-
For the final verification-card recipe, use no recompute when it fits;
430-
otherwise use the narrowest measured selective boundaries. Full recompute is
431-
acceptable temporarily for diagnostic bring-up, but can severely reduce
429+
For the final verification-card recipe, no recompute and selective recompute
430+
on cheap supported modules are both appropriate starting choices. Choose by
431+
measured memory headroom and steady-state performance; selective recompute
432+
does not require a no-recompute OOM first. Full recompute is acceptable
433+
temporarily for diagnostic bring-up, but can severely reduce
432434
throughput. Retain it only when target-scale measurements show no recompute
433435
and supported selective boundaries still fail to fit. Record memory/OOM
434436
evidence, boundaries tested, and steady-state performance cost in the recipe's

‎skills/nemo-mbridge-perf-activation-recompute/SKILL.md‎

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@ Activation recompute (activation checkpointing) trades additional forward work d
1919
## Quick Decision Guide
2020

2121
1. Confirm the pressure is real allocation, not allocator fragmentation. Compare `max_memory_allocated()` with `max_memory_reserved()` on every rank.
22-
2. Start final library, performance, and verification-card recipes with no recompute when the workload fits. Keep an explicit no-recompute control. Under selective granularity, `recompute_modules=[]` is valid and useful for this comparison.
23-
3. When no recompute does not fit, select the narrowest supported boundary from the architecture and observed peak:
22+
2. For final library, performance, and verification-card recipes, no recompute and selective recompute on cheap supported modules are both appropriate starting choices. Choose by measured memory headroom and steady-state performance; a no-recompute OOM is not a prerequisite for selective recompute. Keep an explicit no-recompute control when it fits. Under selective granularity, `recompute_modules=[]` is valid and useful for this comparison.
23+
3. For selective recompute, start with cheap supported boundaries selected from the architecture and observed peak:
2424
- **Standard attention:** `core_attn` is the common first candidate. It is strongest when unfused attention materializes score/probability tensors. With Transformer Engine fused or Flash Attention, compare it against `[]` because those backends already rematerialize attention internals.
2525
- **Multi-Latent Attention (MLA):** start with `mla_up_proj` when expanded Q/K/V projections dominate. Add `core_attn` only when the attention-core state still matters.
2626
- **Grouped MoE:** start with `moe_act` when the expert intermediate activation dominates; add `layernorm` when norm outputs are material. Use whole `moe` recompute only after accounting for the extra expert compute and communication it replays.
@@ -138,8 +138,8 @@ iteration 2. Time and throughput are means over iterations 11--20.
138138
| `moe_act` | 35.941 | 466.26 | 75.50 | -1.85% | +1.99% |
139139
| `layernorm`, `moe_act` | 35.949 | 506.53 | 70.27 | -1.83% | +10.79% |
140140

141-
For this exact workload, no recompute fits and is preferred. If additional
142-
headroom is needed, `moe_act` is the best first tested boundary: it recovered
141+
For this exact workload, both no recompute and cheap selective `moe_act` are
142+
reasonable choices depending on the desired memory headroom. `moe_act` recovered
143143
nearly as much allocated memory as `mla_up_proj` for less replay cost.
144144
`mla_up_proj` is the next candidate if its roughly 39 MB additional reduction
145145
matters. Adding `mlp` to `mla_up_proj` or `layernorm` to `moe_act` did not
@@ -182,12 +182,11 @@ is used here.
182182
### Cross-model conclusion
183183

184184
These measurements do not define one ranking. Moonlight fit with an empty
185-
control, so no recompute is preferred; narrow `moe_act` is a candidate when
186-
additional headroom is needed. Nemotron required broad whole-`moe` recompute;
187-
historical dense Llama evidence found whole-`mlp` replay costly and lacked an
188-
empty control. Start with no recompute when it fits, then the narrowest boundary
189-
implicated by the architecture and peak, broadening only when the narrow choice
190-
does not pass the complete step.
185+
control, and cheap selective `moe_act` offered additional headroom at modest
186+
cost. Nemotron required broad whole-`moe` recompute; historical dense Llama
187+
evidence found whole-`mlp` replay costly and lacked an empty control. Start with
188+
no recompute or cheap selective boundaries, choose by measured memory and
189+
performance, and broaden only as needed to pass the complete step.
191190

192191
## Compatibility and Validation
193192

‎skills/nemo-mbridge-perf-moe-vlm-training/SKILL.md‎

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -72,10 +72,11 @@ The smaller Qwen3.5-style multimodal experiments reinforce the same lessons:
7272

7373
### Choose recompute for the final recipe
7474

75-
For library, performance, and verification-card recipes, start with **no
76-
recompute when it fits**; otherwise use the narrowest supported selective
77-
boundaries measured to fit at target scale. Full recompute can severely reduce
78-
throughput. It is acceptable temporarily during memory-constrained diagnostic
75+
For library, performance, and verification-card recipes, **no recompute and
76+
selective recompute on cheap supported modules are both appropriate starting
77+
choices**. Choose using measured memory headroom and steady-state performance
78+
at target scale; selective recompute does not require a no-recompute OOM first.
79+
Full recompute can severely reduce throughput. It is acceptable temporarily during memory-constrained diagnostic
7980
bring-up, but a successful bring-up run does not justify publishing it as the
8081
default. Retain full recompute only when target-scale measurements show that
8182
no recompute and supported selective boundaries still fail to fit. Document
@@ -98,9 +99,9 @@ performance cost.
9899
or output-layer weights, and the final choice must respect PP/DP rank layout
99100
and the hardware topology.
100101

101-
3. **Use no recompute when it fits, then measured selective recompute**: if
102-
temporary full recompute was used for bring-up, test no recompute and the
103-
narrowest supported selective boundaries before recommending a recipe.
102+
3. **Use no recompute or cheap selective recompute**: if temporary full
103+
recompute was used for bring-up, evaluate these lighter choices before
104+
recommending a recipe; broaden selective boundaries only as needed for fit.
104105
Measure memory and steady-state throughput rather than carrying the
105106
diagnostic setting forward.
106107

@@ -140,9 +141,9 @@ Use this order so each measurement answers one question:
140141
2. Minimize TP while retaining enough memory for non-expert weights.
141142
3. Establish a stable real-data run. Temporary full recompute is allowed for
142143
diagnostic bring-up when memory fit is uncertain.
143-
4. Select the final recipe by testing no recompute first, then the narrowest
144-
supported selective boundaries needed to fit at target scale. Retain full
145-
recompute only under the documented exception above.
144+
4. Select the final recipe using no recompute or cheap supported selective
145+
boundaries, measuring their memory/performance tradeoff at target scale.
146+
Retain full recompute only under the documented exception above.
146147
5. Measure per-rank peak allocated/reserved memory and steady-state step time
147148
or throughput through optimizer initialization and multiple complete steps.
148149
6. If there is headroom, increase MBS; if TP was used for activation pressure,
@@ -163,7 +164,7 @@ fits is the baseline; it is not automatically the performance recipe.
163164
TP=1 CP=1 PP=1
164165
EP sized to the expert topology, often large
165166
Dispatcher: HybridEP on GB200-class systems
166-
Recompute: none when it fits; otherwise the narrowest measured selective boundaries
167+
Recompute: none or cheap supported selective modules; measure memory and throughput
167168
```
168169

169170
### 3D-parallel GB200 path
@@ -182,7 +183,7 @@ CUDA Graph: start narrow, then widen only after the real-data path is stable
182183
| HybridEP on GB200 | strong default | strong default once topology is stable |
183184
| CUDA graphs | useful after bring-up | useful, but more scope-sensitive |
184185
| Freeze vision | natural fit | possible, but less often used as the headline perf path |
185-
| Selective recompute | when no recompute does not fit | when no recompute does not fit |
186+
| Cheap selective recompute | appropriate alongside no recompute | appropriate alongside no recompute |
186187

187188
## Pitfalls
188189

‎skills/nemo-mbridge-recipe-recommender/SKILL.md‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -261,9 +261,11 @@ For an MoE recipe on a fixed allocation, use this practical order:
261261
is a common first candidate on one 8-GPU node, and EP=32 is a common first
262262
candidate on a 32-GPU allocation. These are starting points, not universal
263263
defaults: EP shards expert weights only, not dense layers or activations.
264-
3. For final library, benchmark, and verification-card recipes, start with no
265-
recompute when it fits; otherwise use the narrowest measured selective
266-
boundaries. Full recompute is acceptable temporarily for diagnostic bring-up,
264+
3. For final library, benchmark, and verification-card recipes, no recompute
265+
and selective recompute on cheap supported modules are both appropriate
266+
starting choices. Choose by measured memory headroom and steady-state
267+
performance; selective recompute does not require a no-recompute OOM first.
268+
Full recompute is acceptable temporarily for diagnostic bring-up,
267269
but can severely reduce throughput. Retain it only when target-scale
268270
measurements show no recompute and supported selective boundaries still fail
269271
to fit; document memory/OOM evidence and steady-state performance cost. For

0 commit comments

Comments
 (0)