[Bugfix] IntermediatesCache nested model inputs

[kylesayrs]

Purpose

• Fix bug where nested model inputs were not being properly offloaded by intermediates cache
  • This was caused by this line where the value may be nested, but the IntermediateValue constructor does not handle nested values

- values[key] = IntermediateValue(value=value, device=model_device)
+ values[key] = cls._offload_value(value, offload_device, model_device)

Changes

• Use _offload_value when offloading model inputs from dataloader
  • This necessitates converting _offload_value and _onload_value to be class methods
  • _offload_value accepts an onload_device argument to override which device is used when the value in onloaded (in the case of values from the dataloader, this should be the device of the model)

Testing

• Tested with model with nested inputs e2e
• Added additional unit tests to cover this case

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[gemini-code-assist]

Summary of Changes

Hello @kylesayrs, 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 resolves a critical issue within the IntermediatesCache that prevented the correct offloading and onloading of nested model inputs, such as those contained within dataclasses or complex dictionary/list structures. The core fix involves refactoring the caching mechanism to recursively process these nested types, ensuring that all torch.Tensor components are moved to the appropriate offload device and correctly restored upon retrieval. This enhancement significantly improves the robustness and reliability of the cache when dealing with diverse and complex data formats commonly encountered in deep learning pipelines.

Highlights

• Bug Fix for Nested Inputs: Fixed a bug where the IntermediatesCache failed to properly offload nested model inputs, leading to incorrect device management.
• Refactored Offloading Logic: The from_dataloader method now correctly utilizes _offload_value for all inputs, ensuring nested structures are handled.
• Class Method Conversion: _offload_value and _onload_value were converted to class methods to facilitate recursive calls and better parameter passing.
• Enhanced _offload_value: The _offload_value method now accepts an onload_device argument, allowing explicit control over the target device when values are reloaded.
• Comprehensive Testing: New unit tests were added, including parametrized tests for various nested data types and a deep_equal utility, to ensure the fix works for diverse input structures.

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 effectively addresses a bug in IntermediatesCache where nested model inputs were not being offloaded correctly. The fix, which involves using a recursive _offload_value function, is well-implemented. Refactoring _onload_value and _offload_value into class methods is a good design choice that improves modularity and testability. The unit tests have also been significantly improved by using parametrization to cover a wider range of data types, including the nested structures that caused the original bug. I have one minor suggestion for the new test helper function to improve its clarity.

[brian-dellabetta]

nice!