tico 0.2.0.dev260414__tar.gz → 0.2.0.dev260415__tar.gz

{tico-0.2.0.dev260414 → tico-0.2.0.dev260415}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tico
-Version: 0.2.0.dev260414
+Version: 0.2.0.dev260415
 Summary: Convert Exported Torch Module To Circle
 License: This file provides full text of licenses used in this project
         

tico-0.2.0.dev260415/tico/_version.py

@@ -0,0 +1 @@
+__version__ = "0.2.0.dev260415"

{tico-0.2.0.dev260414 → tico-0.2.0.dev260415}/tico/quantization/algorithm/gptq/quantizer.py

@@ -28,6 +28,11 @@ from tico.quantization.algorithm.gptq.utils import (
 from tico.quantization.config.gptq import GPTQConfig
 from tico.quantization.quantizer import BaseQuantizer
 from tico.quantization.quantizer_registry import register_quantizer
+from tico.utils.utils import move_to_device
+
+
+def move_to_cpu(obj):
+    return move_to_device(obj, "cpu")
 
 
 class StopForward(Exception):
@@ -118,12 +123,12 @@ class GPTQQuantizer(BaseQuantizer):
             for idx, item in enumerate(args):
                 if (idx + 1) > len(self.cache_args):
                     self.cache_args.append([])
-                self.cache_args[idx].append(item)
+                self.cache_args[idx].append(move_to_cpu(item))
             # Store keyword args
             for k, v in kwargs.items():
                 if self.cache_kwargs.get(k, None) is None:
                     self.cache_kwargs[k] = []
-                self.cache_kwargs[k].append(v)
+                self.cache_kwargs[k].append(move_to_cpu(v))
 
             self.num_batches += 1
             raise StopForward  # stop after the first layer
@@ -280,6 +285,7 @@ class GPTQQuantizer(BaseQuantizer):
 
                 # Run layer forward over all cached batches to build Hessian/statistics
                 batch_num = self.num_batches
+                device = next(model.parameters()).device
                 for batch_idx in tqdm(
                     range(batch_num),
                     desc=f"[L{l_idx}] collecting",
@@ -290,9 +296,13 @@ class GPTQQuantizer(BaseQuantizer):
                     cache_args_batch = gather_single_batch_from_list(
                         self.cache_args, batch_idx
                     )
+                    cache_args_batch = move_to_device(cache_args_batch, device)
+
                     cache_kwargs_batch = gather_single_batch_from_dict(
                         self.cache_kwargs, batch_idx
                     )
+                    cache_kwargs_batch = move_to_device(cache_kwargs_batch, device)
+
                     layer(*cache_args_batch, **cache_kwargs_batch)
 
                 # Remove handles
@@ -314,6 +324,7 @@ class GPTQQuantizer(BaseQuantizer):
                     gptq[name].free()
 
             # 4) After quantization, re-run the layer to produce outputs for the next layer
+            device = next(model.parameters()).device
             for batch_idx in tqdm(
                 range(batch_num),
                 desc=f"[L{l_idx}] re-forward",
@@ -324,9 +335,13 @@ class GPTQQuantizer(BaseQuantizer):
                 cache_args_batch = gather_single_batch_from_list(
                     self.cache_args, batch_idx
                 )
+                cache_args_batch = move_to_device(cache_args_batch, device)
+
                 cache_kwargs_batch = gather_single_batch_from_dict(
                     self.cache_kwargs, batch_idx
                 )
+                cache_kwargs_batch = move_to_device(cache_kwargs_batch, device)
+
                 outs = layer(*cache_args_batch, **cache_kwargs_batch)
                 # LLaMA's decoder layer return type differs across Transformers versions:
                 # some return a tuple (hidden_states, ...), others return just a tensor.
@@ -334,7 +349,14 @@ class GPTQQuantizer(BaseQuantizer):
                 outs = outs[0] if isinstance(outs, tuple) else outs
                 # Update inputs for next iteration.
                 if len(self.cache_args) > 0:
-                    self.cache_args[0][batch_idx] = outs
+                    if hasattr(outs, "to") and hasattr(
+                        self.cache_args[0][batch_idx], "device"
+                    ):
+                        self.cache_args[0][batch_idx] = outs.to(
+                            self.cache_args[0][batch_idx].device
+                        )
+                    else:
+                        self.cache_args[0][batch_idx] = outs
 
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()

{tico-0.2.0.dev260414 → tico-0.2.0.dev260415}/tico/quantization/wrapq/examples/quantize_full_qmodel_with_gptq.py

@@ -313,7 +313,7 @@ def main():
         "--save",
         nargs="*",
         type=str,
-        choices=["model_circle", "model_layers", "ptq_checkpoint", "sensitivity"],
+        choices=["circle_full", "circle_per_layer", "ptq_checkpoint", "sensitivity"],
         help="which artifacts should be saved to output_dir",
     )
     parser.add_argument(
@@ -512,10 +512,10 @@ def main():
     # after PTQ quantizer only fixed-length input sequences are valid
     evaluate(q_m, tokenizer, dataset_test, args)
 
-    if args.output_dir is not None and "model_layers" in args.save:
+    if args.output_dir is not None and "circle_per_layer" in args.save:
         save_layers_to(q_m, args.max_seq_len, args.output_dir)
 
-    if args.output_dir is not None and "model_circle" in args.save:
+    if args.output_dir is not None and "circle_full" in args.save:
         calib_inputs = list(torch.stack(calib_inputs).reshape(-1, 1, args.max_seq_len))
         save_model_to(q_m, calib_inputs, args.output_dir)
 

{tico-0.2.0.dev260414 → tico-0.2.0.dev260415}/tico/utils/utils.py

@@ -396,3 +396,24 @@ def is_target_node(
         return False
 
     return True
+
+
+def move_to_device(obj, device):
+    """
+    Recursively move tensors inside a nested structure to the given device.
+    Non-tensor objects are preserved as-is.
+    """
+    if isinstance(obj, torch.Tensor):
+        return obj.to(device)
+
+    elif isinstance(obj, tuple):
+        return tuple(move_to_device(x, device) for x in obj)
+
+    elif isinstance(obj, list):
+        return [move_to_device(x, device) for x in obj]
+
+    elif isinstance(obj, dict):
+        return {k: move_to_device(v, device) for k, v in obj.items()}
+
+    # preserve everything else (bool, int, None, custom objects, etc.)
+    return obj

{tico-0.2.0.dev260414 → tico-0.2.0.dev260415}/tico.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tico
-Version: 0.2.0.dev260414
+Version: 0.2.0.dev260415
 Summary: Convert Exported Torch Module To Circle
 License: This file provides full text of licenses used in this project
         

tico-0.2.0.dev260414/tico/_version.py

@@ -1 +0,0 @@
-__version__ = "0.2.0.dev260414"