Various SDXL quantization fixes

src/brevitas_examples/stable_diffusion/main.py

@@ -12,6 +12,7 @@
 
 from dependencies import value
 from diffusers import DiffusionPipeline
+from diffusers import EulerDiscreteScheduler
 from diffusers import StableDiffusionXLPipeline
 from diffusers.models.attention_processor import Attention
 from diffusers.models.attention_processor import AttnProcessor
@@ -37,7 +38,6 @@
 from brevitas.utils.torch_utils import KwargsForwardHook
 from brevitas_examples.common.generative.quantize import generate_quant_maps
 from brevitas_examples.common.generative.quantize import generate_quantizers
-from brevitas_examples.common.generative.quantize import quantize_model
 from brevitas_examples.common.parse_utils import add_bool_arg
 from brevitas_examples.common.parse_utils import quant_format_validator
 from brevitas_examples.llm.llm_quant.export import BlockQuantProxyLevelManager
@@ -152,7 +152,7 @@ def main(args):
 
     latents = None
     if args.path_to_latents is not None:
-        latents = torch.load(args.path_to_latents).to(torch.float16)
+        latents = torch.load(args.path_to_latents).to(dtype)
 
     # Create output dir. Move to tmp if None
     ts = datetime.fromtimestamp(time.time())
@@ -170,7 +170,11 @@ def main(args):
 
     # Load model from float checkpoint
     print(f"Loading model from {args.model}...")
-    pipe = DiffusionPipeline.from_pretrained(args.model, torch_dtype=dtype)
+    variant = 'fp16' if dtype == torch.float16 else None
+    pipe = DiffusionPipeline.from_pretrained(
+        args.model, torch_dtype=dtype, variant=variant, use_safetensors=True)
+    pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
+    pipe.vae.config.force_upcast = True
     print(f"Model loaded from {args.model}.")
 
     # Move model to target device
@@ -213,7 +217,7 @@ def main(args):
 
     if args.activation_equalization:
         pipe.set_progress_bar_config(disable=True)
-        with activation_equalization_mode(
+        with torch.no_grad(), activation_equalization_mode(
                 pipe.unet,
                 alpha=args.act_eq_alpha,
                 layerwise=True,
@@ -262,8 +266,6 @@ def input_bit_width(module):
                 return args.linear_input_bit_width
             elif isinstance(module, nn.Conv2d):
                 return args.conv_input_bit_width
-            elif isinstance(module, QuantIdentity):
-                return args.quant_identity_bit_width
             else:
                 raise RuntimeError(f"Module {module} not supported.")
 
@@ -346,7 +348,7 @@ def input_zp_stats_type():
                 weight_group_size=args.weight_group_size,
                 quantize_weight_zero_point=args.quantize_weight_zero_point,
                 quantize_input_zero_point=args.quantize_input_zero_point,
-                input_bit_width=input_bit_width,
+                input_bit_width=args.linear_output_bit_width,
                 input_quant_format='e4m3',
                 input_scale_type=args.input_scale_type,
                 input_scale_precision=args.input_scale_precision,
@@ -359,7 +361,6 @@ def input_zp_stats_type():
             # We generate all quantizers, but we are only interested in activation quantization for
             # the output of softmax and the output of QKV
             input_quant = float_sdpa_quantizers[0]
-            input_quant = input_quant.let(**{'bit_width': args.linear_output_bit_width})
             if args.quantize_sdp_2:
                 rewriter = ModuleToModuleByClass(
                     Attention,
@@ -375,14 +376,13 @@ def input_zp_stats_type():
                 config.IGNORE_MISSING_KEYS = False
                 pipe.unet = pipe.unet.to(args.device)
                 pipe.unet = pipe.unet.to(dtype)
-            quant_kwargs = layer_map[torch.nn.Linear][1]
+            quant_kwargs = layer_map['diffusers.models.lora.LoRACompatibleLinear'][1]
             what_to_quantize = []
             if args.quantize_sdp_1:
                 what_to_quantize.extend(['to_q', 'to_k'])
             if args.quantize_sdp_2:
                 what_to_quantize.extend(['to_v'])
             quant_kwargs['output_quant'] = lambda module, name: input_quant if any(ending in name for ending in what_to_quantize) else None
-            layer_map[torch.nn.Linear] = (layer_map[torch.nn.Linear][0], quant_kwargs)
 
             if args.override_conv_quant_config:
                 print(
@@ -420,8 +420,8 @@ def input_zp_stats_type():
                 print(f"Checkpoint loaded!")
             pipe = pipe.to(args.device)
         elif not args.dry_run:
-            if (args.linear_input_bit_width is not None or
-                    args.conv_input_bit_width is not None) and args.input_scale_type == 'static':
+            if (args.linear_input_bit_width > 0 or args.conv_input_bit_width > 0 or
+                    args.linear_output_bit_width > 0) and args.input_scale_type == 'static':
                 print("Applying activation calibration")
                 with torch.no_grad(), calibration_mode(pipe.unet):
                     run_val_inference(
@@ -459,7 +459,7 @@ def input_zp_stats_type():
                         torch.cuda.empty_cache()
             if args.bias_correction:
                 print("Applying bias correction")
-                with bias_correction_mode(pipe.unet):
+                with torch.no_grad(), bias_correction_mode(pipe.unet):
                     run_val_inference(
                         pipe,
                         args.resolution,

src/brevitas_examples/stable_diffusion/mlperf_evaluation/accuracy.py

@@ -615,7 +615,7 @@ def compute_mlperf_fid(
 
     if model_to_replace is not None:
         model.pipe = model_to_replace
-
+    model.pipe.vae.config.force_upcast = True
     ds = Coco(
         data_path=path_to_coco,
         name="coco-1024",

src/brevitas_examples/stable_diffusion/sd_quant/export.py

@@ -93,11 +93,13 @@ def export_quant_params(pipe, output_dir):
             elif isinstance(
                     module,
                     QuantWeightBiasInputOutputLayer) and id(module) not in handled_quant_layers:
+                full_name = name
                 layer_dict = dict()
                 layer_dict = handle_quant_param(module, layer_dict)
                 quant_params[full_name] = layer_dict
                 handled_quant_layers.add(id(module))
             elif isinstance(module, QuantNonLinearActLayer):
+                full_name = name
                 layer_dict = dict()
                 act_scale = module.act_quant.export_handler.symbolic_kwargs[
                     'dequantize_symbolic_kwargs']['scale'].data

tests/brevitas/graph/test_transforms.py

@@ -287,6 +287,6 @@ def forward(self, x):
     model = TestModel()
     assert model.conv.stride == (1, 1)
 
-    kwargs = {'stride': lambda module: 2 if module.in_channels == 3 else 1}
+    kwargs = {'stride': lambda module, name: 2 if module.in_channels == 3 else 1}
     model = ModuleToModuleByInstance(model.conv, nn.Conv2d, **kwargs).apply(model)
     assert model.conv.stride == (2, 2)