rename to more user-friendly

Author: gau-nernst

benchmarks/benchmark_adam_8bit.py

@@ -6,7 +6,7 @@
 #   --amp bf16 \
 #   --optim Adam
 # 
-# To use bnb 8-bit optimizer, set --optim AdamBnb8bit. To use 8-bit optimizer implemented in torchao, set --optim AdamDTQ8bit
+# To use bnb 8-bit optimizer, set --optim Adam8bitBnb. To use 8-bit optimizer implemented in torchao, set --optim Adam8bitAo
 # To profile and export chrome trace, set --profile
 # To enable cosine learning rate scheduler, set --cosine_lr_scheduler
 
@@ -25,7 +25,7 @@
 from torchvision.transforms import v2
 from tqdm import tqdm
 
-from torchao.prototype.optim_8bit import AdamDTQ8bit
+from torchao.prototype.optim_8bit import Adam8bit
 
 
 class CosineSchedule:
@@ -161,8 +161,8 @@ def evaluate_model(model, args):
 
     OPTIM_MAP = dict(
         Adam=torch.optim.Adam,
-        AdamBnb8bit=bnb.optim.Adam8bit,
-        AdamDTQ8bit=AdamDTQ8bit,
+        Adam8bitBnb=bnb.optim.Adam8bit,
+        Adam8bitAo=Adam8bit,
     )
     optim = OPTIM_MAP[args.optim](model.parameters(), args.lr, weight_decay=args.weight_decay)
     lr_schedule = CosineSchedule(args.lr, len(dloader) * args.n_epochs)

test/prototype/test_optim_8bit.py

@@ -9,7 +9,7 @@
     parametrize,
     run_tests,
 )
-from torchao.prototype.optim_8bit import AdamDTQ8bit, AdamWDTQ8bit
+from torchao.prototype import optim_8bit
 from torchao.prototype.optim_8bit.subclass import quantize_8bit_with_qmap, QMAP_SIGNED
 from torchao.utils import TORCH_VERSION_AFTER_2_3
 
@@ -50,17 +50,14 @@ def test_quantize_8bit_with_qmap_compile(self, device):
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="bitsandbytes 8-bit Adam only works for CUDA")
 @pytest.mark.xfail(not TORCH_VERSION_AFTER_2_3, reason="torch.compile() fails for PyTorch < 2.3")
 class TestOptim8bit(TestCase):
-    @parametrize("optim_cls,bnb_optim_cls", [
-        (AdamDTQ8bit, bnb.optim.Adam8bit),
-        (AdamWDTQ8bit, bnb.optim.AdamW8bit),
-    ])
-    def test_adam_8bit_correctness(self, optim_cls, bnb_optim_cls):
+    @parametrize("optim_name", ["Adam8bit", "AdamW8bit"])
+    def test_adam_8bit_correctness(self, optim_name):
         device = "cuda"
         model1 = nn.Sequential(nn.Linear(32, 1024), nn.ReLU(), nn.Linear(1024, 128)).to(device)
         model2 = copy.deepcopy(model1)
 
-        optim1 = bnb_optim_cls(model1.parameters())
-        optim2 = optim_cls(model2.parameters())
+        optim1 = getattr(bnb.optim, optim_name)(model1.parameters())
+        optim2 = getattr(optim_8bit, optim_name)(model2.parameters())
 
         for _ in range(2):
             x = torch.randn(4, 32, device=device)

torchao/prototype/optim_8bit/README.md

@@ -7,15 +7,15 @@ This folder implements 8-bit optimizers using dynamic tree quantization as outli
 This is a drop-in replacement for `torch.optim.Adam`
 
 ```python
-from torchao.prototype.optim_8bit import AdamDTQ8bit
+from torchao.prototype.optim_8bit import Adam8bit
 
 model = ...
-optim = AdamDTQ8bit(model.parameters())
+optim = Adam8bit(model.parameters())
 ```
 
 You can also change quantization block size (default 2048) by passing `block_size=value` to the optimizer.
 
-**Other optimizers**: AdamW is also available as `AdamWDTQ8bit`.
+**Other optimizers**: AdamW is also available as `AdamW8bit`.
 
 NOTE: this requires PyTorch >= 2.3
 

torchao/prototype/optim_8bit/__init__.py

@@ -1,2 +1,2 @@
-from .adam import AdamDTQ8bit
-from .adamw import AdamWDTQ8bit
+from .adam import Adam8bit
+from .adamw import AdamW8bit

torchao/prototype/optim_8bit/adam.py

@@ -7,7 +7,7 @@
 from .subclass import maybe_new_zero_buffer
 
 
-class AdamDTQ8bit(Optimizer):
+class Adam8bit(Optimizer):
     def __init__(
         self,
         params,

torchao/prototype/optim_8bit/adamw.py

@@ -7,7 +7,7 @@
 from .subclass import maybe_new_zero_buffer
 
 
-class AdamWDTQ8bit(Optimizer):
+class AdamW8bit(Optimizer):
     def __init__(
         self,
         params,