fixes to move to fbcode

Summary: removing unused code

Test Plan: python test/test.py

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: f6a92d1
Pull Request resolved: #25

Author: HDCharles

test/test.py

@@ -48,10 +48,10 @@
     Int4WeightOnlyQuantizedLinearWeight
 )
 from torchao.quantization.utils import (
-    apply_logging_hook,
+    _apply_logging_hook,
     compute_error,
     compute_error as SQNR,
-    fqn_to_op_to_shape_to_count,
+    _fqn_to_op_to_shape_to_count,
     LoggingTensorMode,
 )
 from torch.ao.quantization.quantize_fx import convert_to_reference_fx, prepare_fx
@@ -1111,12 +1111,12 @@ def test_shape_logger(self):
             ),
         )
 
-        apply_logging_hook(m)
+        _apply_logging_hook(m)
         with LoggingTensorMode():
             m(x)
             m(x)
 
-        for fqn, d1 in fqn_to_op_to_shape_to_count.items():  # noqa: PERF102
+        for fqn, d1 in _fqn_to_op_to_shape_to_count.items():  # noqa: PERF102
             for op, d2 in d1.items():  # noqa: PERF102
                 for shape, count in d2.items():  # noqa: PERF102
                     # print(fqn, op, shape, count)

torchao/quantization/__init__.py

@@ -38,11 +38,7 @@
     "Int8DynamicallyQuantizedLinearWeight",
     "Int8WeightOnlyQuantizedLinearWeight",
     "Int4WeightOnlyQuantizedLinearWeight",
-    "log_with_rank",
-    "clear_logs",
     "compute_error",
-    "forward_hook",
-    "apply_logging_hook",
     "get_model_size_in_bytes",
     "WeightOnlyInt8QuantLinear",
 ]

torchao/quantization/utils.py

@@ -3,19 +3,15 @@
 
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
-
-import os
 from typing import Dict, Optional
 
 import torch
 from torch.utils._python_dispatch import TorchDispatchMode
 
 __all__ = [
     "find_multiple",
-    "log_with_rank",
-    "clear_logs",
     "compute_error",
-    "apply_logging_hook",
+    "_apply_logging_hook",
     "get_model_size_in_bytes",
 ]
 
@@ -26,25 +22,6 @@ def find_multiple(n: int, k: int) -> int:
     return n + k - (n % k)
 
 
-def log_with_rank(*args):
-    # append
-    #
-    #   {thing_to_log}
-    #
-    # to {file}_{rank}.txt, for printing stuff from multiple GPUs
-    if not os.path.exists(log_dir):
-        os.makedirs(log_dir)
-    with open(log_fname, "a") as f:
-        f.write(" ".join([str(s) for s in args]) + "\n")
-    if local_rank == 0:
-        print(*args)
-
-
-def clear_logs():
-    if os.path.isfile(log_fname):
-        os.remove(log_fname)
-
-
 # basic SQNR
 def compute_error(x, y):
     Ps = torch.linalg.norm(x)
@@ -65,13 +42,13 @@ def forward_hook(module, input):
     return forward_hook
 
 
-def apply_logging_hook(model):
+def _apply_logging_hook(model):
     for name, mod in model.named_modules():
         mod.register_forward_pre_hook(_get_logging_hook(name))
 
 
 # collections.defaultdict printing is weird with lambdas, so hand writing for now
-fqn_to_op_to_shape_to_count: Dict[
+_fqn_to_op_to_shape_to_count: Dict[
     Optional[str], Dict[Optional[str], Dict[Optional[str], int]]
 ] = {}
 
@@ -90,13 +67,13 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
         if shape_str != "":
             shape_str = shape_str[:-2]
 
-        if _cur_fqn not in fqn_to_op_to_shape_to_count:
-            fqn_to_op_to_shape_to_count[_cur_fqn] = {}
-        if op_name not in fqn_to_op_to_shape_to_count[_cur_fqn]:
-            fqn_to_op_to_shape_to_count[_cur_fqn][op_name] = {}
-        if shape_str not in fqn_to_op_to_shape_to_count[_cur_fqn][op_name]:
-            fqn_to_op_to_shape_to_count[_cur_fqn][op_name][shape_str] = 0
-        fqn_to_op_to_shape_to_count[_cur_fqn][op_name][shape_str] += 1
+        if _cur_fqn not in _fqn_to_op_to_shape_to_count:
+            _fqn_to_op_to_shape_to_count[_cur_fqn] = {}
+        if op_name not in _fqn_to_op_to_shape_to_count[_cur_fqn]:
+            _fqn_to_op_to_shape_to_count[_cur_fqn][op_name] = {}
+        if shape_str not in _fqn_to_op_to_shape_to_count[_cur_fqn][op_name]:
+            _fqn_to_op_to_shape_to_count[_cur_fqn][op_name][shape_str] = 0
+        _fqn_to_op_to_shape_to_count[_cur_fqn][op_name][shape_str] += 1
 
         return rs