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

[GuanhuaWang]

maybe having a brief introduction of two params loco needed and what are prefered default values

[XingyuXie]

Thanks for the advice, we already add the basic description here in the update.

[GuanhuaWang]

why loco_quant has params_to_reduce as first argument, not full_grads_for_rank

[XingyuXie]

1. This is because we want to use param.xxx within the all_to_all_loco_quant_reduce function to store, access, and modify the error feedback buffer corresponding to each parameter. For example, in the function, we use p.inter_ef_buf[0]. Since full_grads_for_rank is a list of tensors, assigning or maintaining attributes directly to it is not convenient.

2. Another reason is to minimize modifications to stage3.py. By restricting changes to our own functions, we reduce the risk of introducing unintended bugs.

[hwchen2017]

You use both quantize::h_per_load and quantize::h2_per_load in this kernel code. Can you double check to make sure that both of them were used correctly? This one doesn't match the code line 245 and 246.

[XingyuXie]

Fixed. And test on Llama2-8B. The loss curve seems okay.