[SW-228570] support FP8 GaudiFluxPipeline save and load (#254)

* [SW-228570] support FP8 GaudiFluxPipeline save and load
---------

Signed-off-by: Xin He <[email protected]>
Co-authored-by: Xin He <[email protected]>

Author: xin3he

neural_compressor/torch/algorithms/fp8_quant/save_load.py

@@ -36,6 +36,7 @@
     SaveLoadFormat,
     get_enum_from_format,
     UNIT_MAPPING,
+    write_json_file,
 )
 
 
@@ -390,8 +391,13 @@ def save(model, checkpoint_dir="saved_results", format="huggingface", **kwargs):
         # Ensure those codes run on a single rank.
         configs_mapping = model.qconfig
         config_object = configs_mapping[next(iter(configs_mapping))]
-        update_model_config(model, format, config_object)
-        model.config.save_pretrained(checkpoint_dir)
+        config_object.mode = "LOAD"
+        config_object.world_size = world_size  # record world_size for loading
+        # Flux pipeline has FrozenDict as config
+        if not isinstance(model.config, dict):
+            update_model_config(model, format, config_object)
+            model.config.save_pretrained(checkpoint_dir)
+        write_json_file(os.path.join(checkpoint_dir, "quantization_config.json"), config_object.to_dict())
 
         if hasattr(model, "generation_config") and model.generation_config is not None:
             model.generation_config.save_pretrained(checkpoint_dir)
@@ -405,16 +411,31 @@ def load_empty_raw_model(model_name_or_path, **kwargs):
     """Initialize BF16 model with meta tensor."""
     import transformers
     from accelerate import init_empty_weights
-    config = transformers.AutoConfig.from_pretrained(model_name_or_path, **kwargs)
+
+    # Handling model objects not in AutoModelForCausalLM
+    model = kwargs.get("original_model", None)
+    # Handle Flux pipeline without AutoConfig
+    try:
+        config = transformers.AutoConfig.from_pretrained(model_name_or_path, **kwargs)
+        quantization_config = config.quantization_config if hasattr(config, "quantization_config") else None
+        hp_dtype = config.torch_dtype
+    except:
+        config, hp_dtype = model.config, torch.bfloat16
+        quantization_config = kwargs.get("quantization_config", None)
+        setattr(model.config, "quantization_config", quantization_config)
+
+    if quantization_config is not None and "hp_dtype" in quantization_config:
+        hp_dtype = HpDtype[quantization_config["hp_dtype"].upper()].value
+
     # fp8 model provided by neuralmagic.
     if (
-        "quant_method" in config.quantization_config
-        and config.quantization_config["quant_method"] in ["fp8", "compressed-tensors"]
+        "quant_method" in quantization_config
+        and quantization_config["quant_method"] in ["fp8", "compressed-tensors"]
     ):
         from_neuralmagic = True
         if (
-            "kv_cache_scheme" in config.quantization_config
-            and config.quantization_config["kv_cache_scheme"] is not None
+            "kv_cache_scheme" in quantization_config
+            and quantization_config["kv_cache_scheme"] is not None
         ):
             from_neuralmagic_with_kv = True
         else:
@@ -431,16 +452,13 @@ def load_empty_raw_model(model_name_or_path, **kwargs):
         else:
             raise ValueError("Please install optimum-habana to load fp8 kv cache model.")
 
-    from neural_compressor.torch.utils import get_non_persistent_buffers, load_non_persistent_buffers
-
-    hp_dtype = config.torch_dtype
-    if hasattr(config, "quantization_config") and "hp_dtype" in config.quantization_config:
-        hp_dtype = HpDtype[config.quantization_config["hp_dtype"].upper()].value
+    if model is None:
+        with init_empty_weights(include_buffers=False):
+            model = transformers.AutoModelForCausalLM.from_config(config, torch_dtype=hp_dtype)
     if world_size > 1:
         import deepspeed
+        from neural_compressor.torch.utils import get_non_persistent_buffers, load_non_persistent_buffers
 
-        with init_empty_weights(include_buffers=False):
-            model = transformers.AutoModelForCausalLM.from_config(config, torch_dtype=hp_dtype)
         # TODO: [SW-199728] [DeepSpeed] Buffers initialized by model are not correct after tensor parallel
         # get_non_persistent_buffers and load_non_persistent_buffers are workarounds of [SW-199728]
         non_persistent_buffers = get_non_persistent_buffers(model)
@@ -451,16 +469,13 @@ def load_empty_raw_model(model_name_or_path, **kwargs):
         model = deepspeed.init_inference(model, **ds_inference_kwargs)
         model = model.module
         load_non_persistent_buffers(model, non_persistent_buffers)
-    else:
-        with init_empty_weights(include_buffers=False):
-            model = transformers.AutoModelForCausalLM.from_config(config, torch_dtype=hp_dtype)
     model.to(hp_dtype)
 
     try:
         generation_config = transformers.GenerationConfig.from_pretrained(model_name_or_path, **kwargs)
         model.generation_config = generation_config
     except:  # Since model.generation_config is optional, relaxed exceptions can handle more situations.
-        logger.warning("model.generation_config is not loaded correctly.")
+        logger.warning("model.generation_config may not be loaded correctly.")
     return model, from_neuralmagic, from_neuralmagic_with_kv
 
 
@@ -635,7 +650,8 @@ def load(model_name_or_path, format="huggingface", device="hpu", **kwargs):
         model.load_state_dict(rank_state_dict, assign=True, strict=False)
         load_scale_params(model, rank_state_dict)  # ensure per-channel scale is loaded correctly
     clear_quantized_func_wrapper_factory()
-    model.tie_weights()
+    if hasattr(model, "tie_weights"):
+        model.tie_weights()
     model = model.to(cur_accelerator.name())
     model = model.eval()
     cur_accelerator.synchronize()
@@ -745,8 +761,6 @@ def update_model_config(model, format, config_object):
         quantization_config = convert_config_to_vllm_compatible(config_object)
         model.config.quantization_config = quantization_config
     else:
-        config_object.mode = "LOAD"
-        config_object.world_size = world_size  # record world_size for loading
         model.config.quantization_config = config_object
 
 

neural_compressor/torch/quantization/save_load_entry.py

@@ -26,7 +26,7 @@
     RTNConfig,
     TEQConfig,
 )
-from neural_compressor.torch.utils import SaveLoadFormat, get_enum_from_format
+from neural_compressor.torch.utils import SaveLoadFormat, get_enum_from_format, read_json_file
 
 config_name_mapping = {
     FP8_QUANT: FP8Config,
@@ -111,6 +111,8 @@ def load(model_name_or_path, original_model=None, format="default", device="cpu"
         from neural_compressor.common.base_config import ConfigRegistry
 
         qconfig_file_path = os.path.join(os.path.abspath(os.path.expanduser(model_name_or_path)), "qconfig.json")
+        if not os.path.exists(qconfig_file_path):
+            raise ValueError(f"qconfig.json file is necessary for the default format.")
         with open(qconfig_file_path, "r") as f:
             per_op_qconfig = json.load(f)
 
@@ -138,22 +140,46 @@ def load(model_name_or_path, original_model=None, format="default", device="cpu"
                 return qmodel.to(device)
     elif format == SaveLoadFormat.HUGGINGFACE:
         import transformers
-        config = transformers.AutoConfig.from_pretrained(model_name_or_path, **kwargs)
+
+        try:
+            config = transformers.AutoConfig.from_pretrained(model_name_or_path, **kwargs)
+            quantization_config = config.quantization_config
+        except:
+            quantization_config_file = "quantization_config.json"
+            # for Flux pipeline
+            if os.path.exists(model_name_or_path):
+                # If the model_name_or_path is a local path, try to load the config from there
+                quantization_config_path = os.path.join(model_name_or_path, quantization_config_file)
+            else:
+                # If the model_name_or_path is a Hugging Face model ID, try to download the config
+                from huggingface_hub import hf_hub_download
+
+                quantization_config_path = hf_hub_download(
+                    repo_id=model_name_or_path,
+                    filename=quantization_config_file,
+                    revision=kwargs.get("revision", "main"),
+                )
+            quantization_config = read_json_file(quantization_config_path)
+            kwargs["quantization_config"] = quantization_config
+
+        if original_model is not None:
+            kwargs["original_model"] = original_model
         # use config to check which algorithm is used.
         if (
-            "fp8_config" in config.quantization_config or
+            "fp8_config" in quantization_config or
             # for FP8 LLMs for vLLM (https://huggingface.co/neuralmagic).
             (
-                "quant_method" in config.quantization_config and
-                config.quantization_config["quant_method"] in ["fp8", "compressed-tensors"]
+                "quant_method" in quantization_config and
+                quantization_config["quant_method"] in ["fp8", "compressed-tensors"]
             )
         ):
             from neural_compressor.torch.algorithms import fp8_quant
+
             return fp8_quant.load(model_name_or_path, format=format, device=device, **kwargs)
         else:
             from neural_compressor.torch.algorithms import weight_only
 
             qmodel = weight_only.load(model_name_or_path, format=SaveLoadFormat.HUGGINGFACE, device=device, **kwargs)
             return qmodel.to(device)
     else:
-        assert False, "This code path should never be reached."
+        assert False, "Unexpected format: {} occurred during model loading".format(format)

neural_compressor/torch/utils/utility.py

@@ -13,7 +13,8 @@
 # limitations under the License.
 """Intel Neural Compressor PyTorch utilities."""
 
-
+import os
+import json
 import enum
 import importlib
 from collections import UserDict
@@ -311,8 +312,6 @@ def get_processor_type_from_user_config(user_processor_type: Optional[Union[str,
 
 def dowload_hf_model(repo_id, cache_dir=None, repo_type=None, revision=None):
     """Download hugging face model from hf hub."""
-    import os
-
     from huggingface_hub.constants import DEFAULT_REVISION, HUGGINGFACE_HUB_CACHE
     from huggingface_hub.file_download import REGEX_COMMIT_HASH, repo_folder_name
     from huggingface_hub.utils import EntryNotFoundError
@@ -356,8 +355,6 @@ def dowload_hf_model(repo_id, cache_dir=None, repo_type=None, revision=None):
 
 def load_empty_model(pretrained_model_name_or_path, cls=None, **kwargs):
     """Load a empty model."""
-    import os
-
     from accelerate import init_empty_weights
     from transformers import AutoConfig, AutoModelForCausalLM
     from transformers.models.auto.auto_factory import _BaseAutoModelClass
@@ -741,3 +738,20 @@ def get_enum_from_format(format):
             return obj
     raise ValueError(
         f"Invalid format value ('{format}'). Enter one of [{[m.name for m in SaveLoadFormat]}]")
+
+
+def read_json_file(file_path):
+    """Read a JSON file and return its content."""
+    if not file_path or not os.path.exists(file_path):
+        raise FileNotFoundError(f"File {file_path} does not exist.")
+    with open(file_path, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def write_json_file(file_path, data):
+    """Write data to a JSON file."""
+    if not file_path:
+        raise ValueError("File path cannot be empty.")
+    with open(file_path, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=4, ensure_ascii=False)
+    logger.info(f"Data written to {file_path} successfully.")