[QUESTION] How will recompute-num-layers influence the gpu memory usage with uniform recompute-method.

[bugm]

Megatron-LM gtihub document says:

The uniform method uniformly divides the transformer layers into groups of layers (each group of size --recompute-num-layers) and stores the input activations of each group in memory. The baseline group size is 1 and, in this case, the input activation of each transformer layer is stored. When the GPU memory is insufficient, increasing the number of layers per group reduces the memory usage, enabling a bigger model to be trained. For example, when --recompute-num-layers is set to 4, only the input activation of each group of 4 transformer layers is stored.

I agree that with a bigger group size, it stores less input activations which take less gpu memory. But with a bigger group size, during the recompute process, which means it will also take more gpu memory for the temporary activations.

So the activations peak memory during the training process should be stored_input_activations_memory + M / num_groups (M is the activations memory with no recomputing for the model).
But the official doc says " When the GPU memory is insufficient, increasing the number of layers per group reduces the memory usage, enabling a bigger model to be trained", Is there any other technical applied to support it?

But the way, I have tested to set the --recompute-num-layers equal to the number of layers for the model with a uniform recompute method and found it actually saved some activations memory. But I am not clearly about how this is implemented since with whole layers in one group, which means it needs to store all the activations during the recomputed process, which should be the same as the activations in a no-recompute process.

[Advait251206]

This is a very insightful question — and you're actually correct to be suspicious of the simplified explanation in the docs. The real behavior of uniform activation recomputation in Megatron-LM is slightly subtle.

Let’s break it down properly.

---

🧠 1. What --recompute-num-layers REALLY means

With uniform recompute, Megatron does:

• Split transformer layers into groups of size K = recompute-num-layers
• Only store boundary activations (inputs to each group)
• Recompute everything inside the group during backward

So memory tradeoff is:

Component	Effect
Saved activations	↓ decreases with larger K
Recomputation compute	↑ increases with larger K

---

📌 2. Your key intuition is correct

You said:

Larger group size → less stored activations BUT more recomputation memory

Yes — but the important missing point is:

🔥 recomputation does NOT increase activation peak memory in the way forward-pass activations do

Because recomputation is executed in a checkpointed micro-phase, not all layers exist in memory simultaneously.

---

🧩 3. Why peak memory still decreases

Your assumption:

stored activations + recompute activations = may cancel out

❌ This is incorrect because:

✔️ Forward pass:

All activations exist simultaneously across layers → high peak memory

✔️ With recomputation:

During backward:

• Forward activations are discarded aggressively
• Only checkpoint boundaries are stored
• Each group is recomputed locally and sequentially

So at any moment:

memory = checkpoint activations + current recompute block

NOT:

memory = full model activations + recompute activations

---

🔥 4. The key missing optimization: “segmented recomputation”

Megatron does NOT recompute the whole model at once.

Instead:

For K = recompute-num-layers

Each group behaves like:

Forward:
  layer1 → layerK
  store only input of layer1

Backward:
  re-run layer1 → layerK
  discard intermediate activations again

So peak memory is bounded by:

O(K × hidden_activation_size)

not full model size.

---

🧠 5. Why increasing K still reduces memory (important answer)

This is the core of your confusion.

Case K = 1 (extreme checkpointing)

• Store every layer input
• No recompute inside groups
• High memory for saved checkpoints

Case K = large (e.g. full model)

• Store only input to entire model
• BUT recompute full forward during backward

👉 Still lower memory because:

❗ forward activations are never fully resident at once

They are:

• recomputed
• immediately consumed
• freed again

So peak memory becomes:

checkpoint input + one forward slice

not full graph activations.

---

⚠️ 6. Why your “it should be same as no recompute” intuition is close but not correct

You said:

If whole layers are in one group, it should behave like no recompute

Not quite, because:

❌ No recompute:

• Stores ALL activations for backward

✔️ Full-group recompute:

• Stores ONLY input boundary
• Recomputes everything during backward
• Does NOT store intermediate activations at all

So you replace:

storage memory (large)

with:

compute cost (large)

---

📉 7. Final memory model

A more accurate formula:

Peak Memory ≈
  checkpoint activations
+ activation of ONE recompute segment

NOT:

stored + recompute_total

---

🧭 8. Why docs say “larger groups reduce memory”

Because:

• fewer checkpoints stored
• smaller activation graph kept alive
• more aggressive freeing between recompute windows

So:

K value	Memory	Compute
small K	medium	low recompute
large K	low memory	high recompute

---

🧾 Final answer

• Your intuition about recompute cost increasing is correct
• But recomputation does NOT increase peak memory linearly
• Megatron reduces memory because:
  • it never keeps full activations alive
  • recomputation is segmented and sequential
• Therefore increasing --recompute-num-layers reduces memory at the cost of compute

---

If this answer helped or clarified things, I'd appreciate it if you could mark it as the accepted answer so it's easier for others with the same issue to find.

Also, if you found my contribution useful, I'd appreciate it if you could check out my GitHub profile, follow me, and star any repositories you find interesting.

GitHub: https://github.com/Advait251206