This commit is contained in:
Disty0
2025-10-09 19:58:57 +03:00
parent e19fb2d833
commit a9de8ef152
+13 -10
View File
@@ -18,22 +18,25 @@ def get_module_names(model: ModelMixin) -> list:
def save_sdnq_model(model: ModelMixin, model_path: str, max_shard_size: str = "10GB", is_pipeline: bool = False, sdnq_config: SDNQConfig = None) -> None:
model.save_pretrained(model_path, max_shard_size=max_shard_size) # actual save
quantization_config_path = os.path.join(model_path, "quantization_config.json")
if sdnq_config is not None: # if provided, save global config
sdnq_config.to_json_file(os.path.join(model_path, "quantization_config.json"))
sdnq_config.to_json_file(quantization_config_path)
if is_pipeline:
for module_name in get_module_names(model): # save per-module config if available
module = getattr(model, module_name, None)
if (module is not None) and hasattr(module, "quantization_config") and isinstance(module.quantization_config, SDNQConfig):
module.quantization_config.to_json_file(os.path.join(model_path, module_name, "quantization_config.json"))
if module is None:
continue
module_quantization_config_path = os.path.join(model_path, module_name, "quantization_config.json")
if hasattr(module, "quantization_config") and isinstance(module.quantization_config, SDNQConfig):
module.quantization_config.to_json_file(module_quantization_config_path)
elif hasattr(module, "config"), hasattr(module.config, "quantization_config") and isinstance(module.config.quantization_config, SDNQConfig):
module.config.quantization_config.to_json_file(module_quantization_config_path)
elif sdnq_config is None:
quantization_config = None
if hasattr(model, "quantization_config"):
quantization_config = model.quantization_config
elif hasattr(model, "config") and hasattr(model.config, "quantization_config"):
quantization_config = model.config.quantization_config
if quantization_config is not None:
quantization_config.to_json_file(os.path.join(model_path, "quantization_config.json"))
if hasattr(model, "quantization_config") and isinstance(model.quantization_config, SDNQConfig):
model.quantization_config.to_json_file(quantization_config_path)
elif hasattr(model, "config") and hasattr(model.config, "quantization_config") and isinstance(model.config.quantization_config, SDNQConfig):
model.config.quantization_config.to_json_file(quantization_config_path)
def load_sdnq_model(model_path: str, model_cls: ModelMixin = None, file_name: str = None, dtype: torch.dtype = None, device: torch.device = 'cpu', dequantize_fp32: bool = None, use_quantized_matmul: bool = None, model_config: dict = None, quantization_config: dict = None) -> ModelMixin: