[BUG] Unexpected GPU memory consumption when using transformers PEFT in Zero3

Description

Llama30B with Lora adapters cannot fit into 8 x A100 (80GB).

Suspected Cause

The possible reason for this issue might be that Zero3 does not partition non-trainable weights across GPUs. The basis for this assumption is:

• The memory consumption is consistent with predicted values when Lora is not used.
• When training the model with both Zero2 and Zero3 using Lora, I observe nearly the same memory consumption.
• A code examination of the Zero Runtime sources also suggests this could be the case.

Demonstration of Problem and Experiment Setups

I will illustrate this issue using various experiment setups on smaller models:

1. 7b+lora+stage 3

2. 7b+stage 3 3. 7b+lora+stage 2 4. 7b + stage 2

All other parameters remain consistent in the experiments below.

Question

Is it accurate that in Zero3, only trainable parameters are partitioned? If so, why is it designed this way?

Start by reading deepspeed/runtime/zero/stage3.py around the linked line and compare the reported Zero2 and Zero3 setups with and without LoRA. Reproduce the four experiment variants described in the issue, then verify whether non-trainable weights are partitioned and document the observed behavior and design rationale.