pin memory bug fix

[tenderness-git]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @tenderness-git, 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!

I've addressed a potential bug related to the use of "pin_memory()" in the offloading utility. The core change ensures that memory pinning operations are now conditionally applied, specifically targeting Linux environments. This modification aims to improve cross-platform compatibility and prevent errors that could occur when these operations are performed on non-Linux systems, thereby making the memory offloading mechanism more robust.

Highlights

• Conditional Pin Memory Application: I've updated the "_forward_pre_hook" and "offload_model_to_dict" functions to ensure that "pin_memory()" is only called on tensors if the operating system is detected as Linux. This prevents potential issues or errors that might arise when running on other platforms where this operation might not be supported or behave as expected.
• Platform Module Import: To facilitate the conditional application of "pin_memory()", I've added an import for the "platform" module, which allows for checking the current operating system.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 or fill out our survey to provide feedback.

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 addresses a bug with pin_memory on non-Linux systems by conditionally calling it only when the OS is Linux. The changes are correct in principle. However, the implementation can be made more efficient by checking the platform only once outside the loops instead of on every iteration. I've added specific suggestions to refactor the code for better performance and style.

[gemini-code-assist]

Calling platform.system() inside loops is inefficient as the result will not change during execution. It's better to call it once before the loops and store the result in a variable. This will improve performance by avoiding redundant system calls.

Additionally, there are some minor style issues: there are trailing whitespaces and two spaces around == instead of one (platform.system() == 'Linux').

Suggested change

	for name, param in module.named_parameters(recurse=recurse):
	if platform.system() == 'Linux':
	param.data = param.data.pin_memory()
	offload_param_dict[name] = param.data
	param.data = param.data.to(device=device)
	for name, buffer in module.named_buffers(recurse=recurse):
	buffer.data = buffer.data.pin_memory()
	if platform.system() == 'Linux':
	buffer.data = buffer.data.pin_memory()
	offload_param_dict[name] = buffer.data
	buffer.data = buffer.data.to(device=device)
	setattr(module, "_offload_param_dict", offload_param_dict)
	is_linux = platform.system() == "Linux"
	for name, param in module.named_parameters(recurse=recurse):
	if is_linux:
	param.data = param.data.pin_memory()
	offload_param_dict[name] = param.data
	param.data = param.data.to(device=device)
	for name, buffer in module.named_buffers(recurse=recurse):
	if is_linux:
	buffer.data = buffer.data.pin_memory()
	offload_param_dict[name] = buffer.data
	buffer.data = buffer.data.to(device=device)
	setattr(module, "_offload_param_dict", offload_param_dict)

[gemini-code-assist]

Similar to the previous comment, platform.system() is being called inside loops, which is inefficient. It's better to perform this check once before the loops.

Also, there are some minor style issues: there are trailing whitespaces and two spaces around == instead of one (platform.system() == 'Linux').

Suggested change

	for name, param in module.named_parameters(recurse=True):
	param.data = param.data.pin_memory()
	if platform.system() == 'Linux':
	param.data = param.data.pin_memory()
	offload_param_dict[name] = param.data
	for name, buffer in module.named_buffers(recurse=True):
	buffer.data = buffer.data.pin_memory()
	if platform.system() == 'Linux':
	buffer.data = buffer.data.pin_memory()
	offload_param_dict[name] = buffer.data
	is_linux = platform.system() == "Linux"
	for name, param in module.named_parameters(recurse=True):
	if is_linux:
	param.data = param.data.pin_memory()
	offload_param_dict[name] = param.data
	for name, buffer in module.named_buffers(recurse=True):
	if is_linux:
	buffer.data = buffer.data.pin_memory()
	offload_param_dict[name] = buffer.data