Verified & Integrated MXINT Implementation and Evaluation on NeRF (ADLS Group 0)

configs/examples/nerf_tpe_search.toml

@@ -0,0 +1,153 @@
+# basics
+model = "nerfvision"
+dataset = "nerf-lego"
+
+max_epochs = 5
+batch_size = 128
+learning_rate = 1e-2
+accelerator = "gpu"
+project = "nerf_lego_quantizations"
+seed = 42
+
+[search.search_space]
+name = "graph/quantize/mixed_precision_ptq"
+
+[search.search_space.setup]
+by = "name"
+
+[search.search_space.seed.mxint.config]
+name = ['mxint']
+
+data_in_width = [4, 8]
+data_in_exponent_width = [3, 4, 5, 6, 7, 8, 9]
+data_in_block_size = [3, 4, 5, 6, 7, 8, 9]
+
+weight_width = [2, 4, 8]
+weight_exponent_width = [3, 4, 5, 6, 7, 8, 9]
+weight_block_size = [3, 4, 5, 6, 7, 8, 9]
+
+bias_width = [2, 4, 8]
+bias_exponent_width = [3, 4, 5, 6, 7, 8, 9]
+bias_block_size = [3, 4, 5, 6, 7, 8, 9]
+
+
+[search.search_space.seed.integer.config]
+name = ["integer"]
+
+data_in_width = [4, 8]
+data_in_frac_width = [3, 4, 5, 6, 7, 8, 9]
+
+weight_width = [2, 4, 8]
+weight_frac_width = [3, 4, 5, 6, 7, 8, 9]
+
+bias_width = [2, 4, 8]
+bias_frac_width = [3, 4, 5, 6, 7, 8, 9]
+
+
+[search.search_space.seed.log.config]
+name = ["log"]
+
+data_in_width = [4, 8]
+data_in_exponent_bias = [3, 4, 5, 6, 7, 8, 9]
+
+weight_width = [2, 4, 8]
+weight_exponent_bias = [3, 4, 5, 6, 7, 8, 9]
+
+bias_width = [2, 4, 8]
+bias_exponent_bias = [3, 4, 5, 6, 7, 8, 9]
+
+
+[search.search_space.seed.minifloat_denorm.config]
+name = ["minifloat_denorm"]
+
+data_in_width = [4, 8]
+data_in_exponent_width = [3, 4, 5, 6, 7, 8, 9]
+data_in_exponent_bias = [3, 4, 5, 6, 7, 8, 9]
+
+weight_width = [2, 4, 8]
+weight_exponent_width = [3, 4, 5, 6, 7, 8, 9]
+weight_exponent_bias = [3, 4, 5, 6, 7, 8, 9]
+
+bias_width = [2, 4, 8]
+bias_exponent_width = [3, 4, 5, 6, 7, 8, 9]
+bias_exponent_bias = [3, 4, 5, 6, 7, 8, 9]
+
+
+[search.search_space.seed.minifloat_ieee.config]
+name = ["minifloat_ieee"]
+
+data_in_width = [4, 8]
+data_in_exponent_width = [3, 4, 5, 6, 7, 8, 9]
+data_in_exponent_bias = [3, 4, 5, 6, 7, 8, 9]
+
+weight_width = [2, 4, 8]
+weight_exponent_width = [3, 4, 5, 6, 7, 8, 9]
+weight_exponent_bias = [3, 4, 5, 6, 7, 8, 9]
+
+bias_width = [2, 4, 8]
+bias_exponent_width = [3, 4, 5, 6, 7, 8, 9]
+bias_exponent_bias = [3, 4, 5, 6, 7, 8, 9]
+
+
+[search.search_space.seed.block_fp.config]
+name = ["block_fp"]
+
+data_in_width = [4, 8]
+data_in_exponent_width = [3, 4, 5, 6, 7, 8, 9]
+data_in_exponent_bias = [3, 4, 5, 6, 7, 8, 9]
+data_in_block_size = [3, 4, 5, 6, 7, 8, 9]
+
+weight_width = [2, 4, 8]
+weight_exponent_width = [3, 4, 5, 6, 7, 8, 9]
+weight_exponent_bias = [3, 4, 5, 6, 7, 8, 9]
+weight_block_size = [3, 4, 5, 6, 7, 8, 9]
+
+bias_width = [2, 4, 8]
+bias_exponent_width = [3, 4, 5, 6, 7, 8, 9]
+bias_exponent_bias = [3, 4, 5, 6, 7, 8, 9]
+bias_block_size = [3, 4, 5, 6, 7, 8, 9]
+
+
+[search.search_space.seed.binary.config]
+name = ["binary"]
+
+data_in_width = [4, 8]
+data_in_stochastic = [3, 4, 5, 6, 7, 8, 9]
+data_in_bipolar = [3, 4, 5, 6, 7, 8, 9]
+
+weight_width = [2, 4, 8]
+weight_stochastic = [3, 4, 5, 6, 7, 8, 9]
+weight_bipolar = [3, 4, 5, 6, 7, 8, 9]
+
+bias_width = [2, 4, 8]
+bias_stochastic = [3, 4, 5, 6, 7, 8, 9]
+bias_bipolar = [3, 4, 5, 6, 7, 8, 9]
+
+
+[search.strategy]
+name = "optuna"
+eval_mode = true
+
+[search.strategy.sw_runner.basic_evaluation]
+data_loader = "val_dataloader"
+num_samples = 256
+
+[search.strategy.hw_runner.average_bitwidth]
+compare_to = 32  # compare to FP32
+
+[search.strategy.setup]
+n_jobs = 1
+n_trials = 20
+timeout = 20000
+sampler = "tpe"
+# sum_scaled_metrics = true # single objective
+# direction = "maximize"
+sum_scaled_metrics = false # multi objective
+
+[search.strategy.metrics]
+loss.scale = 1.0
+loss.direction = "minimize"
+psnr.scale = 1.0
+psnr.direction = "maximize"
+average_bitwidth.scale = 0.2
+average_bitwidth.direction = "minimize"

src/chop/actions/search/search.py

@@ -1,5 +1,6 @@
 import logging
 from os import PathLike
+import json
 
 import toml
 import torch
@@ -68,44 +69,56 @@ def search(
 
     # search preparation
     accelerator = parse_accelerator(accelerator)
-    strategy_config, search_space_config = parse_search_config(search_config)
+    strategy_config, orig_search_space_config = parse_search_config(search_config)
     save_path.mkdir(parents=True, exist_ok=True)
 
-    # load model if the save_name is provided
-    if load_name is not None and load_type in ["pl", "mz", "pt"]:
-        model = load_model(load_name=load_name, load_type=load_type, model=model)
-        logger.info(f"Loaded model from {load_name}.")
-    model.to(accelerator)
-    # set up data module
-    data_module.prepare_data()
-    data_module.setup()
-
-    # construct the search space
-    logger.info("Building search space...")
-    search_space_cls = get_search_space_cls(search_space_config["name"])
-    search_space = search_space_cls(
-        model=model,
-        model_info=model_info,
-        config=search_space_config,
-        dummy_input=get_dummy_input(model_info, data_module, task, device=accelerator),
-        accelerator=accelerator,
-        data_module=data_module,
-    )
-    search_space.build_search_space()
-
-    # construct a search strategy
-    strategy_cls = get_search_strategy_cls(strategy_config["name"])
-    strategy = strategy_cls(
-        model_info=model_info,
-        task=task,
-        dataset_info=dataset_info,
-        data_module=data_module,
-        config=strategy_config,
-        accelerator=accelerator,
-        save_dir=save_path,
-        visualizer=visualizer,
-    )
-
-    logger.info("Search started...")
-    # perform search and save the results
-    strategy.search(search_space)
+    for quantization_name in orig_search_space_config["seed"]:
+        search_space_config = json.loads(json.dumps(orig_search_space_config))
+        search_space_config["seed"]["default"] = orig_search_space_config["seed"][
+            quantization_name
+        ]
+
+        current_save_path = save_path / quantization_name
+        current_save_path.mkdir(parents=True, exist_ok=True)
+
+        # load model if the save_name is provided
+        if load_name is not None and load_type in ["pl", "mz", "pt"]:
+            model = load_model(load_name=load_name, load_type=load_type, model=model)
+            logger.info(f"Loaded model from {load_name}.")
+        model.to(accelerator)
+        # set up data module
+        data_module.prepare_data()
+        data_module.setup()
+
+        # construct the search space
+        logger.info("Building search space...")
+        search_space_cls = get_search_space_cls(search_space_config["name"])
+        search_space = search_space_cls(
+            model=model,
+            model_info=model_info,
+            config=search_space_config,
+            dummy_input=get_dummy_input(
+                model_info, data_module, task, device=accelerator
+            ),
+            accelerator=accelerator,
+            data_module=data_module,
+        )
+        search_space.build_search_space()
+
+        # construct a search strategy
+        strategy_cls = get_search_strategy_cls(strategy_config["name"])
+        strategy = strategy_cls(
+            model_info=model_info,
+            task=task,
+            dataset_info=dataset_info,
+            data_module=data_module,
+            config=strategy_config,
+            accelerator=accelerator,
+            save_dir=current_save_path,
+            visualizer=visualizer,
+            quantization_name=quantization_name,
+        )
+
+        logger.info(f"Search started... ...{quantization_name}")
+        # perform search and save the results
+        strategy.search(search_space)

src/chop/actions/search/strategies/base.py

@@ -41,6 +41,7 @@ def __init__(
         accelerator,
         save_dir,
         visualizer,
+        quantization_name: str = "",
     ):
         self.dataset_info = dataset_info
         self.task = task
@@ -50,6 +51,10 @@ def __init__(
         self.data_module = data_module
         self.visualizer = visualizer
 
+        if quantization_name != "":
+            quantization_name = f"{quantization_name}/"
+        self.quantization_name = quantization_name
+
         self.sw_runner = []
         self.hw_runner = []
         # the software runner's __call__ will use the rebuilt model to calculate the software metrics like accuracy, loss, ...

src/chop/actions/search/strategies/optuna.py

@@ -110,7 +110,12 @@ def objective(self, trial: optuna.trial.Trial, search_space):
         trial.set_user_attr("scaled_metrics", scaled_metrics)
         trial.set_user_attr("sampled_config", sampled_config)
 
-        self.visualizer.log_metrics(metrics=scaled_metrics, step=trial.number)
+        log_scaled_metrics = {
+            self.quantization_name + k: v
+            for k, v in scaled_metrics.items()
+            if v is not None
+        }
+        self.visualizer.log_metrics(metrics=log_scaled_metrics, step=trial.number)
 
         if not self.sum_scaled_metrics:
             return list(scaled_metrics.values())

src/chop/actions/search/strategies/runners/software/eval.py

@@ -4,6 +4,10 @@
 from torchmetrics.text import Perplexity
 from torchmetrics import MeanMetric
 
+from chop.tools.plt_wrapper.nerf.losses import (
+    ColorLoss,
+    NerfPsnr,
+)
 
 from .base import SWRunnerBase
 
@@ -42,6 +46,9 @@ def _setup_metric(self):
                     ).to(self.accelerator)
                 case _:
                     raise ValueError(f"task {self.task} is not supported.")
+        elif self.model_info.is_nerf_vision_model:
+            self.metric = NerfPsnr().to(self.accelerator)
+            self.color_loss = ColorLoss().to(self.accelerator)
         elif self.model_info.is_nlp_model:
             match self.task:
                 case "classification" | "cls":
@@ -62,6 +69,8 @@ def forward(self, batch: dict[str, torch.Tensor], model):
                     return self.vision_cls_forward(batch, model)
                 case _:
                     raise ValueError(f"task {self.task} is not supported.")
+        elif self.model_info.is_nerf_vision_model:
+            return self.nerf_vision_cls_forward(batch, model)
         elif self.model_info.is_nlp_model:
             match self.task:
                 case "classification" | "cls":
@@ -79,6 +88,19 @@ def vision_cls_forward(self, batch, model):
         self.loss(loss)
         return {"loss": loss, "accuracy": acc}
 
+    def nerf_vision_cls_forward(self, batch, model):
+        for k, v in batch.items():
+            batch[k] = v.to(self.accelerator)
+        if len(batch["pts"].shape) == 4:
+            for k in batch:
+                batch[k] = batch[k].squeeze(0)
+        logits = model(batch["pts"], batch["viewdirs"])
+        loss = self.color_loss(logits, batch)
+        psnr = self.metric(logits, batch)
+        self.loss(loss)
+
+        return {"loss": loss, "psnr": psnr}
+
     def nlp_cls_forward(self, batch, model):
         batch = {
             k: v.to(self.accelerator) if isinstance(v, torch.Tensor) else v
@@ -109,6 +131,8 @@ def compute(self) -> dict[str, float]:
             reduced["perplexity"] = self.metric.compute().item()
         elif isinstance(self.metric, MulticlassAccuracy):
             reduced["accuracy"] = self.metric.compute().item()
+        elif isinstance(self.metric, NerfPsnr):
+            reduced["psnr"] = self.metric.compute().item()
         else:
             raise ValueError(f"metric {self.metric} is not supported.")
         return reduced

src/chop/actions/simulate.py

@@ -60,6 +60,7 @@ def simulate(
                 "-Wno-fatal",
                 "-Wno-lint",
                 "-Wno-style",
+                "--assert",
                 "--trace-fst",
                 "--trace-structs",
                 "--trace-depth",

src/chop/cli.py

@@ -863,7 +863,7 @@ def _setup_model_and_dataset(self):
             name=self.args.model,
             task=self.args.task,
             dataset_info=dataset_info,
-            checkpoint=checkpoint,
+            checkpoint=checkpoint or self.args.model,
             pretrained=self.args.is_pretrained,
             quant_config=quant_config,
         )