Make optim lazily intialize global state (#2277)

stack-info: PR: #2277, branch: drisspg/stack/60

Author: drisspg

torchao/optim/subclass_4bit.py

@@ -29,8 +29,19 @@
 # https://github.com/thu-ml/low-bit-optimizers/blob/e3e2854728e498c2a606e3fdb88daa27ae94f9a6/lpmm/configs/2nd_moment_group_128.yml
 # NOTE: power-1 is linear
 # TODO: since QMAP_UNSIGNED is linear, perhaps doing affine quantize is faster?
-QMAP_SIGNED = create_dynamic_map(True, 3, 4)
-QMAP_UNSIGNED = torch.linspace(0, 1, 17)[1:].tolist()  # no zero
+
+# Lazy initialization to avoid meta device issues during import
+from functools import lru_cache
+
+
+@lru_cache(maxsize=1)
+def get_qmap_signed():
+    return create_dynamic_map(True, 3, 4)
+
+
+@lru_cache(maxsize=1)
+def get_qmap_unsigned():
+    return torch.linspace(0, 1, 17, device="cpu")[1:].tolist()  # no zero
 
 
 class OptimState4bit(TorchAOBaseTensor):
@@ -90,7 +101,9 @@ def zeros(cls, shape, signed: bool = True, block_size: int = 128, device=None):
 
         codes = torch.zeros(n_elems // 2, dtype=torch.uint8, device=device)
         scale = torch.zeros(n_elems // block_size, device=device)
-        qmap = torch.tensor(QMAP_SIGNED if signed else QMAP_UNSIGNED, device=device)
+        qmap = torch.tensor(
+            get_qmap_signed() if signed else get_qmap_unsigned(), device=device
+        )
         return cls(codes, scale, qmap, signed, shape)
 
     def __repr__(self):

torchao/optim/subclass_8bit.py

@@ -26,8 +26,18 @@
 c10d_functional = torch.ops.c10d_functional
 _c10d_functional = torch.ops._c10d_functional
 
-QMAP_SIGNED = create_dynamic_map(signed=True)
-QMAP_UNSIGNED = create_dynamic_map(signed=False)
+# Lazy initialization to avoid meta device issues during import
+from functools import lru_cache
+
+
+@lru_cache(maxsize=1)
+def get_qmap_signed():
+    return create_dynamic_map(signed=True)
+
+
+@lru_cache(maxsize=1)
+def get_qmap_unsigned():
+    return create_dynamic_map(signed=False)
 
 
 class OptimState8bit(TorchAOBaseTensor):
@@ -79,7 +89,9 @@ def dequantize(self, output_dtype=None):
     def zeros(cls, shape, signed: bool = True, block_size: int = 256, device=None):
         codes = torch.zeros(shape, dtype=torch.uint8, device=device)
         scale = torch.zeros(codes.numel() // block_size, device=device)
-        qmap = torch.tensor(QMAP_SIGNED if signed else QMAP_UNSIGNED, device=device)
+        qmap = torch.tensor(
+            get_qmap_signed() if signed else get_qmap_unsigned(), device=device
+        )
         return cls(codes, scale, qmap, signed)
 
     def __repr__(self):