Update

[ghstack-poisoned]

Author: vkuzo

.github/workflows/torchao_experimental_test.yml

@@ -33,7 +33,10 @@ jobs:
       - name: Install requirements
         run: |
           conda activate venv
-          pip install torch --index-url "https://download.pytorch.org/whl/nightly/cpu"
+          # Install executorch first because it installs its own version
+          # of torch and torchao, which we do not want to use
+          pip install executorch
+          pip install torch --index-url "https://download.pytorch.org/whl/nightly/cpu" --force-reinstall
           pip install numpy
           pip install pytest
           pip install parameterized
@@ -57,6 +60,12 @@ jobs:
           sh build_and_run_tests.sh
           rm -rf /tmp/cmake-out
           popd
+      - name: ET ops build
+        run: |
+          conda activate venv
+          pushd torchao/experimental
+          sh build_torchao_ops.sh executorch
+          popd
 
   test-mps-ops:
     strategy:

README.md

@@ -115,13 +115,13 @@ swap_linear_with_semi_sparse_linear(model, {"seq.0": SemiSparseLinear})
 ADAM takes 2x as much memory as the model params so we can quantize the optimizer state to either 8 or 4 bit effectively reducing the optimizer VRAM requirements by 2x or 4x respectively over an fp16 baseline
 
 ```python
-from torchao.prototype.low_bit_optim import AdamW8bit, AdamW4bit, AdamWFp8
+from torchao.optim import AdamW8bit, AdamW4bit, AdamWFp8
 optim = AdamW8bit(model.parameters()) # replace with Adam4bit and AdamFp8 for the 4 / fp8 versions
 ```
 
-In practice, we are a tiny bit slower than expertly written kernels but the implementations for these optimizers were written in a **few hundred lines of PyTorch code** and compiled so please use them or copy-paste them for your quantized optimizers. Benchmarks [here](https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim)
+In practice, we are a tiny bit slower than expertly written kernels but the implementations for these optimizers were written in a **few hundred lines of PyTorch code** and compiled so please use them or copy-paste them for your quantized optimizers. Benchmarks [here](https://github.com/pytorch/ao/tree/main/torchao/optim)
 
-We also have support for [single GPU CPU offloading](https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim#optimizer-cpu-offload) where both the gradients (same size as weights) and the optimizers will be efficiently sent to the CPU. This alone can **reduce your VRAM requirements by 60%**
+We also have support for [single GPU CPU offloading](https://github.com/pytorch/ao/tree/main/torchao/optim#optimizer-cpu-offload) where both the gradients (same size as weights) and the optimizers will be efficiently sent to the CPU. This alone can **reduce your VRAM requirements by 60%**
 
 ```python
 optim = CPUOffloadOptimizer(model.parameters(), torch.optim.AdamW, fused=True)

benchmarks/benchmark_low_bit_adam.py

@@ -34,7 +34,7 @@
 from torchvision.transforms import v2
 from tqdm import tqdm
 
-from torchao.prototype import low_bit_optim
+from torchao import optim
 from torchao.utils import get_available_devices
 
 _DEVICE = get_available_devices()[-1]
@@ -43,9 +43,9 @@
 OPTIM_MAP = dict(
     AdamW=partial(torch.optim.AdamW, fused=True),
     AdamW8bitBnb=bnb.optim.AdamW8bit,
-    AdamW8bitAo=low_bit_optim.AdamW8bit,
-    AdamWFp8Ao=low_bit_optim.AdamWFp8,
-    AdamW4bitAo=low_bit_optim.AdamW4bit,
+    AdamW8bitAo=optim.AdamW8bit,
+    AdamWFp8Ao=optim.AdamWFp8,
+    AdamW4bitAo=optim.AdamW4bit,
 )
 
 try:
@@ -249,12 +249,10 @@ def evaluate_model(model, args):
         optim_cls = OPTIM_MAP[args.optim]
 
         if args.optim_cpu_offload == "ao":
-            optim_cls = partial(
-                low_bit_optim.CPUOffloadOptimizer, optimizer_class=optim_cls
-            )
+            optim_cls = partial(optim.CPUOffloadOptimizer, optimizer_class=optim_cls)
         elif args.optim_cpu_offload == "ao_offload_grads":
             optim_cls = partial(
-                low_bit_optim.CPUOffloadOptimizer,
+                optim.CPUOffloadOptimizer,
                 optimizer_class=optim_cls,
                 offload_gradients=True,
             )

benchmarks/float8/float8_roofline.py

@@ -184,8 +184,11 @@ def get_gemm_times(
         elif float8_recipe_name in ("rowwise", "rowwise_with_gw_hp"):
             scale_a = torch.ones(M, 1, device=device)
             scale_b = torch.ones(1, N, device=device)
+        elif mx_recipe_name == "mxfp8_cublas":
+            scale_a = torch.ones(M, K // 32, device=device, dtype=torch.float8_e8m0fnu)
+            scale_b = torch.ones(N, K // 32, device=device, dtype=torch.float8_e8m0fnu)
         else:
-            assert False, "TODO add mx gemm here"
+            assert False, "TODO add cutlass mx gemm here"
 
         def do_matmul(A, B):
             return torch._scaled_mm(

benchmarks/float8/training/README.md

@@ -4,15 +4,15 @@ The `float8_training_benchmark.sh` script in this directory can be used to launc
 
 ## Usage
 
-Example: `TORCHTITAN_ROOT=${HOME}/torchtitan FLOAT8_RECIPE=rowwise ./float8_training_benchmark.sh`
+Example: `TORCHTITAN_ROOT=${HOME}/torchtitan FLOAT8_RECIPE_WITH_BEST_SETTINGS=rowwise ./float8_training_benchmark.sh`
 
 Training parameters can be configured via environment variables.
 
 - Required:
-    - `TORCHTITAN_ROOT`
+    - `TORCHTITAN_ROOT`: Root directory of torchtitan in your local filesystem
 - Optional:
-    - `RECIPE`: rowwise|tensorwise. defaults to tensorwise.
-    - `BATCH_SIZE`: defaults to 1.
-    - `STEPS`: defaults to 100.
+    - `FLOAT8_RECIPE_WITH_BEST_SETTINGS`: "rowwise" or "tensorwise". Applies float8 training with the specified scaling recipe, as well as additional training configs which are optimal for that scaling recipe. See `float8_training_benchmark.sh` for more details.
+    - `BATCH_SIZE`: Defaults to 1.
+    - `STEPS`: Defaults to 100.
 
 **NOTE**: `torch.compile` and FSDP2 are always used. Other forms of parallelism supported in torchtitan are not yet supported in this script.

benchmarks/quantized_training/pretrain_llama2.py

@@ -22,14 +22,13 @@
 from torch.utils.checkpoint import checkpoint
 from tqdm import tqdm
 
-from torchao import quantize_
+from torchao import optim, quantize_
 from torchao._models.llama.model import (
     ModelArgs,
     RMSNorm,
     Transformer,
     transformer_configs,
 )
-from torchao.prototype import low_bit_optim
 from torchao.prototype.quantized_training import (
     bitnet_training,
     int8_mixed_precision_training,
@@ -190,10 +189,10 @@ def insert_rmsnorm(module: torch.nn.Module):
     print(f"No. of buffers: {sum(p.numel() for p in model.buffers()):,}")
     torch.cuda.reset_peak_memory_stats()  # don't count memory occupied by unquantized weights
 
-    # only use optimizers from torchao.prototype.low_bit_optim to support quantized training
+    # only use optimizers from torchao.optim to support quantized training
     if args.optim == "AdamW":
         args.optim = "_AdamW"
-    optim = getattr(low_bit_optim, args.optim)(
+    optimizer = getattr(optim, args.optim)(
         model.parameters(),
         lr=args.lr,
         weight_decay=args.weight_decay,
@@ -228,15 +227,15 @@ def insert_rmsnorm(module: torch.nn.Module):
         if step % args.log_interval == 0:
             log_dict = dict(
                 loss=loss.item(),
-                lr=optim.param_groups[0]["lr"],
+                lr=optimizer.param_groups[0]["lr"],
                 max_memory_allocated=torch.cuda.max_memory_allocated() / 1e9,
                 max_memory_reserved=torch.cuda.max_memory_reserved() / 1e9,
             )
             run.log(log_dict, step=step)
             pbar.set_postfix(loss=log_dict["loss"])
 
-        optim.step()
-        optim.zero_grad()
+        optimizer.step()
+        optimizer.zero_grad()
 
         step += 1
         pbar.update()

test/dtypes/test_bitnet.py

@@ -67,7 +67,10 @@ def init_multi_module(self) -> nn.Module:
         return module
 
     def init_transformer(
-        self, weight_tying: bool, dtype: Optional[torch.dtype] = None
+        self,
+        weight_tying: bool,
+        dtype: Optional[torch.dtype] = None,
+        requires_grad: bool = True,
     ) -> nn.Module:
         torch.manual_seed(42)
         args = ModelArgs(
@@ -81,6 +84,13 @@ def init_transformer(
         module = Transformer(args).cuda()
         if dtype is not None:
             module = module.to(dtype=dtype)
+
+        # if requires_grad=False, just set requires_grad to False
+        # in the first layer to ensure we still train some params.
+        if requires_grad is False:
+            for param in module.layers[0].parameters():
+                param.requires_grad = requires_grad
+
         self.broadcast_module(module)
         return module
 
@@ -107,6 +117,7 @@ def test_transformer_parity(self):
                 ],
                 "compile_transformer_block": [False, True],
                 "dtype": [torch.float32, torch.bfloat16],
+                "requires_grad": [True, False],
             },
             self._test_transformer_parity,
         )
@@ -117,6 +128,7 @@ def _test_transformer_parity(
         precompute: bool,
         scaling_type_weight: ScalingType,
         compile_transformer_block: bool,
+        requires_grad: bool,
         dtype: Optional[torch.dtype] = None,
     ):
         if not enable_fsdp_float8_all_gather and precompute:
@@ -127,7 +139,10 @@ def _test_transformer_parity(
         # latter uses fp8 compute. With fp8 all-gather, FSDP would pre-cast to
         # fp8 for that tied weight, incorrectly using fp8 for the embedding.
         weight_tying = not enable_fsdp_float8_all_gather
-        module = self.init_transformer(weight_tying=weight_tying, dtype=dtype)
+        module = self.init_transformer(
+            weight_tying=weight_tying, dtype=dtype, requires_grad=requires_grad
+        )
+
         ref_module = copy.deepcopy(module)
         float8_linear_config1 = Float8LinearConfig(
             cast_config_weight=CastConfig(scaling_type=scaling_type_weight),