[Sequential Onloading] Support onloading and offloading frozen dataclasses

src/llmcompressor/pipelines/cache.py

@@ -160,8 +160,8 @@ def _size_helper(intermediate: IntermediateValue) -> int:
                     for v in value.values():
                         _size_helper(v)
                 case _ if is_dataclass(value):
-                    for field in fields(value):
-                        _size_helper(getattr(value, field.name))
+                    for f in fields(value):
+                        _size_helper(getattr(value, f.name))
                 case _:
                     # this handles primitive values that don't match any other cases
                     sizes[torch.device("cpu")] += sys.getsizeof(value, 0)
@@ -205,10 +205,12 @@ def _onload_value(cls, intermediate: IntermediateValue) -> Any:
             case dict():
                 return {k: cls._onload_value(v) for k, v in value.items()}
             case _ if is_dataclass(value):
-                for field in fields(value):
-                    v = getattr(value, field.name)
-                    setattr(value, field.name, cls._onload_value(v))
-                return value
+                return type(value)(
+                    **{
+                        f.name: cls._onload_value(getattr(value, f.name))
+                        for f in fields(value)
+                    }
+                )
             case _:
                 # handles primitive values that should be returned as is.
                 # without this, a MatchError would be raised for unhandled types.
@@ -255,10 +257,15 @@ def _offload_value(
                     device=None,
                 )
             case _ if is_dataclass(value):
-                for field in fields(value):
-                    v = getattr(value, field.name)
-                    setattr(value, field.name, cls._offload_value(v, **kwargs))
-                return IntermediateValue(value=value, device=None)
+                return IntermediateValue(
+                    value=type(value)(
+                        **{
+                            f.name: cls._offload_value(getattr(value, f.name), **kwargs)
+                            for f in fields(value)
+                        }
+                    ),
+                    device=None,
+                )
             case _:
                 # handles primitive values and provides a warning for unsupported types.
                 # without this, values trigger a MatchError exception.

tests/llmcompressor/pipelines/test_cache.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass, fields, is_dataclass
+from typing import Optional
 
 import pytest
 import torch
@@ -7,10 +8,11 @@
 from llmcompressor.pipelines.cache import IntermediatesCache
 
 
-@dataclass
+@dataclass(frozen=True)
 class SampleDataclass:
-    a: torch.Tensor
-    b: int
+    a: int
+    b: Optional[torch.Tensor] = None
+    c: Optional["SampleDataclass"] = None
 
 
 @pytest.fixture
@@ -35,7 +37,7 @@ def sample_cache(sample_dataloader):
 
 values_to_test = [
     torch.randn(2, 3).to("cpu"),
-    SampleDataclass(a=torch.randn(2, 3), b=42),
+    SampleDataclass(a=42, b=torch.randn(2, 3), c=SampleDataclass(a=64)),
     torch.float32,
     [1, 2, 3],
 ]
@@ -156,8 +158,8 @@ def deep_equal(a, b) -> bool:
                 return False
             return all(deep_equal(a[key], b[key]) for key in a.keys())
         case _ if is_dataclass(a):
-            a_dict = {field.name: getattr(a, field.name) for field in fields(a)}
-            b_dict = {field.name: getattr(b, field.name) for field in fields(b)}
+            a_dict = {field: getattr(a, field.name) for field in fields(a)}
+            b_dict = {field: getattr(b, field.name) for field in fields(b)}
 
             return deep_equal(a_dict, b_dict)
         case _: