Writes the state dict of the LoRA layers (optionally with metadata) to disk.
(
state_dict: dict[str, torch.Tensor],
save_directory: str,
is_main_process: bool,
weight_name: str,
save_function: Callable,
safe_serialization: bool,
lora_adapter_metadata: dict | None = None,
)
| 1008 | |
| 1009 | @staticmethod |
| 1010 | def write_lora_layers( |
| 1011 | state_dict: dict[str, torch.Tensor], |
| 1012 | save_directory: str, |
| 1013 | is_main_process: bool, |
| 1014 | weight_name: str, |
| 1015 | save_function: Callable, |
| 1016 | safe_serialization: bool, |
| 1017 | lora_adapter_metadata: dict | None = None, |
| 1018 | ): |
| 1019 | """Writes the state dict of the LoRA layers (optionally with metadata) to disk.""" |
| 1020 | if os.path.isfile(save_directory): |
| 1021 | logger.error(f"Provided path ({save_directory}) should be a directory, not a file") |
| 1022 | return |
| 1023 | |
| 1024 | if lora_adapter_metadata and not safe_serialization: |
| 1025 | raise ValueError("`lora_adapter_metadata` cannot be specified when not using `safe_serialization`.") |
| 1026 | if lora_adapter_metadata and not isinstance(lora_adapter_metadata, dict): |
| 1027 | raise TypeError("`lora_adapter_metadata` must be of type `dict`.") |
| 1028 | |
| 1029 | if save_function is None: |
| 1030 | if safe_serialization: |
| 1031 | |
| 1032 | def save_function(weights, filename): |
| 1033 | # Inject framework format. |
| 1034 | metadata = {"format": "pt"} |
| 1035 | if lora_adapter_metadata: |
| 1036 | for key, value in lora_adapter_metadata.items(): |
| 1037 | if isinstance(value, set): |
| 1038 | lora_adapter_metadata[key] = list(value) |
| 1039 | metadata[LORA_ADAPTER_METADATA_KEY] = json.dumps( |
| 1040 | lora_adapter_metadata, indent=2, sort_keys=True |
| 1041 | ) |
| 1042 | |
| 1043 | return safetensors.torch.save_file(weights, filename, metadata=metadata) |
| 1044 | |
| 1045 | else: |
| 1046 | save_function = torch.save |
| 1047 | |
| 1048 | os.makedirs(save_directory, exist_ok=True) |
| 1049 | |
| 1050 | if weight_name is None: |
| 1051 | if safe_serialization: |
| 1052 | weight_name = LORA_WEIGHT_NAME_SAFE |
| 1053 | else: |
| 1054 | weight_name = LORA_WEIGHT_NAME |
| 1055 | |
| 1056 | save_path = Path(save_directory, weight_name).as_posix() |
| 1057 | save_function(state_dict, save_path) |
| 1058 | logger.info(f"Model weights saved in {save_path}") |
| 1059 | |
| 1060 | @classmethod |
| 1061 | def _save_lora_weights( |
no test coverage detected