[Bugfix] Support offloaded parameters when initializing KV cache parameters (#261)

* use register_offload_parameter

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

* fix typo

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

---------

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

src/compressed_tensors/linear/compressed_linear.py

@@ -21,6 +21,7 @@
     QuantizationStatus,
     initialize_module_for_quantization,
 )
+from compressed_tensors.utils import register_offload_parameter
 from torch import Tensor
 from torch.nn import Parameter
 from torch.nn.functional import linear
@@ -68,7 +69,7 @@ def from_linear(
             param = Parameter(
                 torch.empty(shape, device=device, dtype=dtype), requires_grad=False
             )
-            module.register_parameter(name, param)
+            register_offload_parameter(module, name, param)
 
         # mark module as compressed
         module.quantization_status = QuantizationStatus.COMPRESSED

src/compressed_tensors/quantization/lifecycle/initialize.py

@@ -203,11 +203,10 @@ def _initialize_attn_scales(module: Module) -> None:
         torch.empty(expected_shape, dtype=scale_dtype, device=device),
         requires_grad=False,
     )
-
-    module.register_parameter(KVCacheScaleType.KEY.value, init_scale)
+    register_offload_parameter(module, KVCacheScaleType.KEY.value, init_scale)
 
     init_scale = Parameter(
         torch.empty(expected_shape, dtype=scale_dtype, device=device),
         requires_grad=False,
     )
-    module.register_parameter(KVCacheScaleType.VALUE.value, init_scale)
+    register_offload_parameter(module, KVCacheScaleType.VALUE.value, init_scale)