[feat] CPU mem: release and realloc gradient buffers when offload train engine

[yueming-yuan]

ci-megatron-pr: #17

core changes

Megatron: radixark/Megatron-LM#17
torch memory saver: fzyzcjy/torch_memory_saver#72

result

Green is the newest

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request optimizes memory utilization within the offload training engine by implementing a strategy to release and reallocate gradient buffers. By explicitly freeing gradient buffer GPU memory before offloading the model and then reallocating it upon resuming, the system avoids unnecessary transfer of gradient data to CPU memory, thereby reducing the overall memory footprint during the offload process.

Highlights

• Gradient Buffer Management: Introduced new methods (_get_grad_buffers, _free_grad_buffers, _realloc_grad_buffers) to explicitly manage the allocation and deallocation of gradient buffers.
• Memory Optimization during Offload: Integrated the _free_grad_buffers call into the sleep method to release GPU memory associated with gradients before offloading, preventing unnecessary CPU offload.
• Buffer Reallocation on Wake-up: Incorporated the _realloc_grad_buffers call into the wake_up method to restore gradient buffers after the model resumes, ensuring views remain valid.
• Reduced CPU Memory Footprint: The changes aim to save approximately 4 bytes per parameter of CPU memory during offload by not preserving gradient data across sleep/wake-up cycles.

Changelog

• miles/backends/megatron_utils/actor.py
  • Added _free_grad_buffers method to release GPU memory associated with gradient buffers before model offload.
  • Added _realloc_grad_buffers method to reallocate gradient buffers after the model resumes.
  • Added _get_grad_buffers helper method to retrieve all ParamAndGradBuffer objects from the DDP model.
  • Called _free_grad_buffers within the sleep method to optimize memory before torch_memory_saver.pause().
  • Called _realloc_grad_buffers within the wake_up method to restore buffers after torch_memory_saver.resume().

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a memory optimization by releasing and reallocating gradient buffers when the training engine is offloaded. The logic is sound and correctly placed within the sleep and wake_up methods. I have one suggestion to make the implementation more robust against potential errors during the buffer freeing process.

[gemini-code-assist]

To improve robustness, it's better to first collect all buffer sizes and then resize them in a separate loop. This ensures that if an error occurs while getting a buffer's size, no buffers are resized, leaving the system in a consistent state. The current implementation could lead to a partially freed state if an error occurs mid-loop.

Suggested change

	self._grad_buffer_nbytes = []
	for buffer in self._get_grad_buffers():
	self._grad_buffer_nbytes.append(buffer.grad_data.untyped_storage().nbytes())
	buffer.grad_data.untyped_storage().resize_(0)
	buffers = self._get_grad_buffers()
	self._grad_buffer_nbytes = [b.grad_data.untyped_storage().nbytes() for b in buffers]
	for buffer in buffers:
	buffer.grad_data.untyped_storage().resize_(0)

[yueming-yuan]

Do people think we need to keep this argument? Is there any senario that we need to keep cpu copy of gradient and distributed optimizer's parameter gather buffer? if no i will delete

[fzyzcjy]

I personally do not have an example to need this flag in my mind

[yueming-yuan]

okay

[fzyzcjy]

I personally do not have an example to need this flag in my mind