Skip to content

[Utils] Offloaded cache size #1714

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+43 −0
Open

[Utils] Offloaded cache size #1714

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 43 additions & 0 deletions src/llmcompressor/pipelines/cache.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
import sys
import warnings
from collections import defaultdict
from dataclasses import dataclass, fields, is_dataclass
from typing import Any, Dict, Generator, List, Optional, Union

Expand Down Expand Up @@ -132,10 +134,51 @@ def delete(self, batch_index: int, consumed_names: Optional[List[str]] = None):
del intermediates[name]

def append(self, values: Dict[str, Any]):
"""
Append new values to the cache. The new values will be assigned the next
available batch index

:param values: dictionary mapping keys to values used for update
"""
batch_index = len(self.batch_intermediates)
self.batch_intermediates.append({})
self.update(batch_index, values)

def size(self) -> Dict[torch.device, int]:
"""
Returns the memory used by cached values, keyed by device, in bytes

:return: dictionary mapping torch device to number of bytes in cache
"""
sizes = defaultdict(lambda: 0)

def _size_helper(intermediate: IntermediateValue) -> int:
value = intermediate.value

if isinstance(value, torch.Tensor):
sizes[value.device] += value.nbytes

elif is_dataclass(value):
for field in fields(value):
_size_helper(getattr(value, field.name))

elif isinstance(value, tuple):
for v in value:
_size_helper(v)

elif isinstance(value, dict):
for v in value.values():
_size_helper(v)

else:
sizes[torch.device("cpu")] += sys.getsizeof(value, 0)

for intermediates in self.batch_intermediates:
for value in intermediates.values():
_size_helper(value)

return dict(sizes)
Comment on lines +147 to +180
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The implementation of size() has a few issues that could be improved for robustness and correctness:

  1. Recursion Depth: The recursive _size_helper function can lead to a RecursionError if the cached data structures are deeply nested. An iterative approach using a queue would be more robust.
  2. Incorrect list size: The size of list objects is not calculated correctly. The current implementation falls through to the else block, where sys.getsizeof() is called on the list object itself, not its contents. This will significantly underestimate memory usage for lists. Lists should be traversed like tuples.
  3. Minor issues:
    • The type hint for _size_helper's return value is int, but it doesn't return anything. It should be None.
    • defaultdict(lambda: 0) can be simplified to defaultdict(int).
    • sys.getsizeof(value, 0) is a bit unusual; sys.getsizeof(value) is more idiomatic.

Here is a suggested iterative implementation that addresses these points:

    def size(self) -> Dict[torch.device, int]:
        """
        Returns the memory used by cached values, keyed by device, in bytes

        :return: dictionary mapping torch device to number of bytes in cache
        """
        sizes = defaultdict(int)
        queue = []
        for intermediates in self.batch_intermediates:
            queue.extend(intermediates.values())

        while queue:
            intermediate = queue.pop()
            value = intermediate.value

            if isinstance(value, torch.Tensor):
                sizes[value.device] += value.nbytes
            elif is_dataclass(value):
                for field in fields(value):
                    queue.append(getattr(value, field.name))
            elif isinstance(value, (list, tuple)):
                queue.extend(value)
            elif isinstance(value, dict):
                queue.extend(value.values())
            else:
                sizes[torch.device("cpu")] += sys.getsizeof(value)

        return dict(sizes)

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We've never encountered a list so far, but seems valid to add it


def iter(
self, input_names: Optional[List[str]] = None
) -> Generator[Any, None, None]:
Expand Down