[BUG]: all output logits are 0 after one optimizer step (gemini & grad accumulate).

[bestbzw]

🐛 Describe the bug

I add gradient accumulation into https://github.com/hpcaitech/ColossalAI/blob/e327e95144f4db8875531699e5b048f77cb80eba/examples/language/gpt/gemini/train_gpt_demo.py
and find that all output logits are 0 (loss is 10.828) after train one optimizer step .

If use "zero2", it runs well.

Environment

Colossal-AI version: 0.1.12

PyTorch Version: 1.12.1
PyTorch Version required by Colossal-AI: 1.12
PyTorch version match: ✓

System CUDA Version: 11.3
CUDA Version required by PyTorch: 11.3
CUDA Version required by Colossal-AI: 11.3
CUDA Version Match: ✓

CUDA Extension: ✓

[feifeibear]

gemini can not work with grad accumulation.

[bestbzw]

@feifeibear
Thanks for your quick reply.
does colossalai provide any other zero-offload api that works with grad accumulation.

[feifeibear]

ZeRO1 and ZeRO2 work with grad accumulation. In my opinion, if you can use a large batch size with the help of Gemini, you might not need grad accumulation. Because we do some aggressive memory optimization for Gemini, which conflicts with grad acc.

[bestbzw]

I see.
Unfortunately, I train a gpt model with 10B parameters, it is too large to use a large batch size, even with the help of Gemini.

[feifeibear]

How much CPU and GPU memory do you have? 10B is not a big model for gemini.

[haofanwang]

@feifeibear Just to ask, if gemini cannot support grad accumulation, does that mean the gradient_accumulation_steps in train_dreambooth_colossalai.py actually is useless? It makes no difference whatever value is set. The batch size doesn't increase with accumulation.

[feifeibear]

@haofanwang I am afraid yes. I add an assertation in dreambooth example

[bestbzw]

@feifeibear I have 128 A100