Issue for collecting gradients from hyper-scale LLMs

[SShock92]

Hi, thanks for sharing such a nice work!

I am trying to collect gradients for LLaMA2-13B with the released code, but out-of-memory issue occurs.
Currently, I am using a single NVIDIA H100 GPU.
How many GPUs do we need for models having more than 13B parameters?

Thanks!

[jusjinuk]

Hello, @SShock92.

Thanks for your interest in our work!

The gradient extraction step requires approximately:

Note

The gradient extraction step in scripts/run_sqllm.sh requires:

• 0.3 hrs on 1 A100 for Llama-2-7B
• 0.6 hrs on 2 A100s for Llama-2-13B
• 2.7 hrs on 6–7 A100s for Llama-2-70B

The script automatically splits the model across visible GPUs. To run gradient extraction separately, use the -m gradients option, then re-run scripts/run_sqllm.sh.

We've added the GPU hour estimates for both gradient extraction and Hessian construction to the README.md.

Feel free to reach out with any further questions!