Update patch version v1.3.4

Co-authored-by: seungahdev <[email protected]>

friendli/cli/model.py

@@ -225,19 +225,14 @@ def convert(
         output_model_file_name or default_names[output_ckpt_file_type]
     )
 
-    model_output_path = os.path.join(output_dir, output_model_file_name)
-    tokenizer_output_dir = output_dir
-    attr_output_path = os.path.join(output_dir, output_attr_file_name)
-
     try:
         convert_checkpoint(
             model_name_or_path=model_name_or_path,
-            model_output_path=model_output_path,
+            output_model_file_name=output_model_file_name,
             output_ckpt_file_type=output_ckpt_file_type,
+            output_attr_file_name=output_attr_file_name,
             output_dir=output_dir,
             data_type=data_type,
-            tokenizer_output_dir=tokenizer_output_dir,
-            attr_output_path=attr_output_path,
             cache_dir=cache_dir,
             dry_run=dry_run,
             quantize=quantize,
@@ -344,14 +339,12 @@ def convert_adapter(
         output_adapter_filename or default_names[output_adapter_file_type]
     )
 
-    adapter_output_path = os.path.join(output_dir, output_adapter_filename)
-    attr_output_path = os.path.join(output_dir, output_attr_filename)
-
     try:
         convert_adapter_checkpoint(
             adapter_name_or_path=adapter_name_or_path,
-            adapter_output_path=adapter_output_path,
-            adapter_attr_output_path=attr_output_path,
+            output_attr_filename=output_attr_filename,
+            output_dir=output_dir,
+            output_adapter_filename=output_adapter_filename,
             base_model_name_or_path=base_model_name_or_path,
             data_type=data_type,
             output_adapter_file_type=output_adapter_file_type,
@@ -360,3 +353,10 @@ def convert_adapter(
         )
     except (NotFoundError, CheckpointConversionError, InvalidConfigError) as exc:
         secho_error_and_exit(str(exc))
+
+    msg = (
+        f"Checkpoint({adapter_name_or_path}) can be converted."
+        if dry_run
+        else f"Checkpoint({adapter_name_or_path}) has been converted successfully."
+    )
+    typer.secho(msg)

friendli/modules/converter/base.py

@@ -415,23 +415,30 @@ def check_config(self) -> None:
             for target_module in self.adapter_config.target_modules:
                 if (
                     target_module
-                    in MODEL_TYPE_TO_UNSUPPORTED_LORA_TARGET_MODULES_MAP[
+                    not in MODEL_TYPE_TO_SUPPORTED_LORA_TARGET_MODULES_MAP[
                         self.config.model_type
                     ]
                 ):
-                    raise NotSupportedCheckpointError(
-                        invalid_option=f"target_module={target_module}",
-                        valid_options=list(
-                            MODEL_TYPE_TO_SUPPORTED_LORA_TARGET_MODULES_MAP[
-                                self.config.model_type
-                            ]
-                        ),
+                    if (
+                        target_module
+                        in MODEL_TYPE_TO_UNSUPPORTED_LORA_TARGET_MODULES_MAP[
+                            self.config.model_type
+                        ]
+                    ):
+                        raise NotSupportedCheckpointError(
+                            invalid_option=f"target_module={target_module}",
+                            valid_options=list(
+                                MODEL_TYPE_TO_SUPPORTED_LORA_TARGET_MODULES_MAP[
+                                    self.config.model_type
+                                ]
+                            ),
+                        )
+
+                    logger.warn(
+                        "Target module %s does not exist in the base model (%s). Will be ignored.",
+                        target_module,
+                        self.adapter_config.base_model_name_or_path,
                     )
-                logger.warn(
-                    "Target module %s does not exist in the base model (%s). Will be ignored.",
-                    target_module,
-                    self.adapter_config.base_model_name_or_path,
-                )
 
         if (self.adapter_config.layers_to_transform is not None) and (
             self.adapter_config != list(range(self.converter.decoder_layer_num))

friendli/modules/converter/convert.py

@@ -41,13 +41,12 @@
 
 def convert_checkpoint(  # pylint: disable=too-many-branches
     model_name_or_path: str,
-    model_output_path: str,
-    output_ckpt_file_type: CheckpointFileType,
+    output_model_file_name: str,
+    output_attr_file_name: str,
     output_dir: str,
+    output_ckpt_file_type: CheckpointFileType,
     *,
     data_type: Optional[ModelDataType] = None,
-    tokenizer_output_dir: Optional[str] = None,
-    attr_output_path: Optional[str] = None,
     cache_dir: Optional[str] = None,
     dry_run: bool = False,
     quantize: bool = False,
@@ -57,11 +56,14 @@ def convert_checkpoint(  # pylint: disable=too-many-branches
 
     Args:
         model_name_or_path (str): Hugging Face model name or local path to the checkpoint.
-        model_output_path (str): Path to the converted checkpoint to save.
+        output_model_file_name (str): File name of converted checkpoint to save.
+        output_attr_file_name (str): File name of the attribute YAML file for
+            the converted checkpoint.
+        output_dir (str) : Directory path to save the converted checkpoint and the attribute YAML,
+            and tokenizer configuration file.
+        output_ckpt_file_type (CheckpointFileType): The file type of converted checkpoint.
         data_type (Optional[ModelDataType]): Converted checkpoint data type.
             Defaults to torch_dtype in 'config.json'
-        tokenizer_output_dir (Optional[str], optional): Directory path to save 'tokenizer.json'
-            for the converted checkpoint. Defaults to None.
         attr_output_path (Optional[str], optional): Path to create the attribute YAML file for
             the converted checkpoint. Defaults to None.
         cache_dir (Optional[str], optional): Path for downloading checkpoint. Defaults to None.
@@ -77,6 +79,7 @@ def convert_checkpoint(  # pylint: disable=too-many-branches
 
     """
     # pylint: disable=too-many-locals
+    model_output_path = os.path.join(output_dir, output_model_file_name)
     model_config = get_model_pretrained_config(
         model_name_or_path, model_output_path, cache_dir
     )
@@ -122,7 +125,9 @@ def convert_checkpoint(  # pylint: disable=too-many-branches
         )
 
         convert_info_list = converter.get_convert_info_list()
-        with get_saver(output_ckpt_file_type, output_dir) as saver:
+        with get_saver(
+            output_ckpt_file_type, output_dir, output_model_file_name
+        ) as saver:
             for name, w in converter.convert(model, convert_info_list):
                 saver.save_tensor(name, w)
 
@@ -131,63 +136,58 @@ def convert_checkpoint(  # pylint: disable=too-many-branches
             model_name_or_path,
         )
 
-    if attr_output_path is not None:
-        if (
-            quant_config
-            and quant_config.mode == QuantMode.FP8
-            and ModelDataType.FP8_E4M3
-        ):
-            model_config.torch_dtype = (
-                get_torch_data_type(data_type)
-                if data_type
-                else model_config.torch_dtype
-            )
-            setattr(model_config, "use_fp8_e4m3", True)
-            assert output_dir is not None
-            model_config.to_json_file(os.path.join(output_dir, "config.json"))
-        else:
-            attr = converter.get_attributes()
-            with open(attr_output_path, "w", encoding="utf-8") as file:
-                yaml.dump(attr, file, sort_keys=False)
-
-    if tokenizer_output_dir is not None:
-        try:
-            saved_tokenizer_file_paths = save_tokenizer(
-                model_name_or_path=model_name_or_path,
-                cache_dir=cache_dir,
-                save_dir=tokenizer_output_dir,
-            )
-        except TokenizerNotFoundError as exc:
-            logger.warn(str(exc))
-
-        if not (
-            quant_config
-            and quant_config.mode == QuantMode.FP8
-            and ModelDataType.FP8_E4M3
-        ):
-            for path in saved_tokenizer_file_paths:
-                if path != "tokenizer.json":
-                    try:
-                        os.remove(path)
-                    except FileNotFoundError:
-                        logger.warn(
-                            "Tried to delete unnecessary tokenizer file %s but the file "
-                            "is not found.",
-                            path,
-                        )
+    # Save attr.yaml
+    attr_output_path = os.path.join(output_dir, output_attr_file_name)
+    if quant_config and quant_config.mode == QuantMode.FP8 and ModelDataType.FP8_E4M3:
+        model_config.torch_dtype = (
+            get_torch_data_type(data_type) if data_type else model_config.torch_dtype
+        )
+        setattr(model_config, "use_fp8_e4m3", True)
+        model_config.to_json_file(os.path.join(output_dir, "config.json"))
+    else:
+        attr = converter.get_attributes()
+        with open(attr_output_path, "w", encoding="utf-8") as file:
+            yaml.dump(attr, file, sort_keys=False)
+
+    # Save tokenizer files.
+    tokenizer_output_dir = output_dir
+    try:
+        saved_tokenizer_file_paths = save_tokenizer(
+            model_name_or_path=model_name_or_path,
+            cache_dir=cache_dir,
+            save_dir=tokenizer_output_dir,
+        )
+    except TokenizerNotFoundError as exc:
+        logger.warn(str(exc))
+
+    if not (
+        quant_config and quant_config.mode == QuantMode.FP8 and ModelDataType.FP8_E4M3
+    ):
+        for path in saved_tokenizer_file_paths:
+            if "tokenizer.json" not in path:
+                try:
+                    os.remove(path)
+                except FileNotFoundError:
+                    logger.warn(
+                        "Tried to delete unnecessary tokenizer file %s but the file "
+                        "is not found.",
+                        path,
+                    )
 
 
 def convert_adapter_checkpoint(  # pylint: disable=too-many-locals, too-many-arguments
     adapter_name_or_path: str,
-    adapter_output_path: str,
-    adapter_attr_output_path: str,
+    output_attr_filename: str,
+    output_dir: str,
+    output_adapter_filename: str,
     base_model_name_or_path: Optional[str],
     data_type: Optional[ModelDataType],
     output_adapter_file_type: CheckpointFileType,
     cache_dir: Optional[str],
     dry_run: bool = False,
 ) -> None:
     """Convert HuggingFace model checkpoint to Friendli format."""
+    adapter_attr_output_path = os.path.join(output_dir, output_attr_filename)
     adapter_config = get_adapter_config(adapter_name_or_path, cache_dir)
     base_model_name_or_path = (
         base_model_name_or_path or adapter_config.base_model_name_or_path
@@ -232,16 +232,17 @@ def convert_adapter_checkpoint(  # pylint: disable=too-many-locals, too-many-arg
             adapter_name_or_path,
         )
         convert_dict = adapter_converter.get_convert_info_list()
-        with get_saver(output_adapter_file_type, adapter_output_path) as saver:
+        with get_saver(
+            output_adapter_file_type, output_dir, output_adapter_filename
+        ) as saver:
             for name, w in adapter_converter.convert(model, convert_dict):
                 saver.save_tensor(name, w)
 
-        if adapter_attr_output_path is not None:
-            attr = adapter_converter.get_attributes()
-            with open(adapter_attr_output_path, "w", encoding="utf-8") as file:
-                yaml.dump([attr], file, sort_keys=False)
-
         logger.info(
             "Hugging Face checkpoint (%s) is successfully converted to Friendli format!",
             adapter_name_or_path,
         )
+
+    attr = adapter_converter.get_attributes()
+    with open(adapter_attr_output_path, "w", encoding="utf-8") as file:
+        yaml.dump([attr], file, sort_keys=False)

friendli/modules/converter/saver.py

@@ -23,14 +23,13 @@
 
 
 def get_saver(
-    ckpt_file_type: CheckpointFileType,
-    output_dir: str,
+    ckpt_file_type: CheckpointFileType, output_dir: str, output_file_name: str
 ) -> CheckpointSaver:
     """Create a saver that corresponds to the file type."""
     if ckpt_file_type == CheckpointFileType.HDF5:
-        return HDF5Saver(output_dir)
+        return HDF5Saver(output_dir, output_file_name)
     if ckpt_file_type == CheckpointFileType.SAFETENSORS:
-        return SafetensorsSaver(output_dir)
+        return SafetensorsSaver(output_dir, output_file_name)
     raise CheckpointConversionError(
         f"Output file type {ckpt_file_type} is not supported."
     )
@@ -39,10 +38,13 @@ def get_saver(
 class CheckpointSaver(AbstractContextManager):
     """Abstract for savers."""
 
-    def __init__(self, output_dir: Union[str, os.PathLike]) -> None:
+    def __init__(
+        self, output_dir: Union[str, os.PathLike], output_file_name: str
+    ) -> None:
         """Check that the output file already exists."""
         super().__init__()
-        self.output_dir = output_dir
+        self._output_dir = output_dir
+        self._output_file_name = output_file_name
 
     @abstractmethod
     def save_tensor(self, tensor_id: str, t: Union[np.ndarray, torch.Tensor]) -> None:
@@ -66,10 +68,10 @@ def __exit__(self, *exc) -> None:
 class HDF5Saver(CheckpointSaver):
     """Saver for HDF5."""
 
-    def __init__(self, output_dir: str) -> None:
+    def __init__(self, output_dir: str, output_file_name: str) -> None:
         """Create a HDF5 file."""
-        super().__init__(output_dir)
-        self._out_f = h5py.File(output_dir + "model.h5", "w")
+        super().__init__(output_dir, output_file_name)
+        self._out_f = h5py.File(os.path.join(output_dir, output_file_name), "w")
 
     def save_tensor(self, tensor_id: str, t: Union[np.ndarray, torch.Tensor]) -> None:
         """Create a group if not exists, and save the tensor in the file."""
@@ -158,10 +160,11 @@ class SafetensorsSaver(CheckpointSaver):
     because Safetensors does not support stream saving.
     """
 
-    def __init__(self, output_dir: Union[str, os.PathLike]) -> None:
+    def __init__(
+        self, output_dir: Union[str, os.PathLike], output_file_name: str
+    ) -> None:
         """Initialize a saver."""
-        super().__init__(output_dir)
-        self._output_dir = output_dir
+        super().__init__(output_dir, output_file_name)
         self._tensors: Dict[str, Union[np.ndarray, torch.Tensor]] = {}
         self._saver: UnionSafetensorsSaverInterface = UnionSafetensorsSaverInterface()
 
@@ -191,12 +194,12 @@ def shard_checkpoint(self, max_shard_size: str):
             total_size += weight_size
 
         if len(sharded_tensors) == 1:
-            return {"model.safetensors": sharded_tensors[0]}, None
+            return {self._output_file_name: sharded_tensors[0]}, None
 
         weight_map = {}
         shards = {}
         for idx, shard in enumerate(sharded_tensors):
-            shard_file = "model.safetensors".replace(
+            shard_file = self._output_file_name.replace(
                 ".safetensors",
                 f"-{idx + 1:05d}-of-{len(sharded_tensors):05d}.safetensors",
             )
@@ -219,7 +222,7 @@ def _save_to_file(self) -> None:
             self._saver.save_file(shard, os.path.join(self._output_dir, shard_file))
 
         if index is None:
-            path_to_weights = os.path.join(self._output_dir, "model.safetensors")
+            path_to_weights = os.path.join(self._output_dir, self._output_file_name)
             logger.info("Model weights saved in (%s)", path_to_weights)
         else:
             save_index_file = os.path.join(

friendli/modules/converter/utils.py

@@ -206,7 +206,6 @@ def save_tokenizer(
 
     tokenizer = get_tokenizer(model_name_or_path, cache_dir=cache_dir)
     saved_file_paths = tokenizer.save_pretrained(save_directory=save_dir)
-
     tokenizer_json_path = None
     for path in saved_file_paths:
         if "tokenizer.json" == os.path.basename(path):

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "friendli-client"
-version = "1.3.3"
+version = "1.3.4"
 description = "Client of Friendli Suite."
 license = "Apache-2.0"
 authors = ["FriendliAI teams <[email protected]>"]