Merge LoCo with Zero++

[XingyuXie]

Integration of LoCo Method into ZeRO++

Overview

This PR introduces the integration of the LoCo method, as outlined in this paper, into the ZeRO++ framework of DeepSpeed. The key enhancement involves applying error feedback compensation to 4-bit gradients before communication. This approach improves pre-training loss outcomes without additional time overhead, though it requires extra GPU memory. The extent of this memory increase depends on model size and training configuration.

Experimental Results

We conducted pre-training experiments using the Llama2 architecture, adjusting the number of layers and hidden size. The experiments included:

• A smaller-scale model with 0.8B parameters trained on 30B tokens.
• A larger-scale model with 8B parameters trained on 5B tokens.

The training data was sampled from Redpajama-V2.

Findings:

• Smaller Models (0.8B parameters): Significant gains were observed when applying the LoCo method.
• Larger Models (8B parameters): The gains were present but less pronounced. This could be due to:
  1. Relatively smaller data volume.
  2. Lower pre-training loss for larger models, making significant improvements harder to achieve.

However, even a smaller pre-training loss gap in larger models can translate to meaningful gains in downstream tasks.

Example Script

For reference, the run.sh script used for the 8B parameter, 5B tokens experiment is attached. The experiment was conducted using the DeepSpeed-Megatron platform.

Acknowledgments

Special thanks to cc @GuanhuaWang for ongoing communication and guidance throughout this work.

---

We appreciate your consideration of this PR and welcome any feedback or questions!

[XingyuXie]

@XingyuXie please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.

@microsoft-github-policy-service agree [company="{your company}"]

Options:

• (default - no company specified) I have sole ownership of intellectual property rights to my Submissions and I am not making Submissions in the course of work for my employer.

@microsoft-github-policy-service agree

• (when company given) I am making Submissions in the course of work for my employer (or my employer has intellectual property rights in my Submissions by contract or applicable law). I have permission from my employer to make Submissions and enter into this Agreement on behalf of my employer. By signing below, the defined term “You” includes me and my employer.

@microsoft-github-policy-service agree company="Microsoft"

Contributor License Agreement

@microsoft-github-policy-service agree

[GuanhuaWang]

@XingyuXie thx for this effort.

Overall looks good to me. Just left a few comments

[XingyuXie]

As required by cc @GuanhuaWang , we added the unit-test code to verify the logic of cuda kernels used in LoCo-Zero++.

[GuanhuaWang]

Thx @XingyuXie for the pr updates on unit-test. Overall, it Looks good to me.

cc @tjruwase @hwchen2017