Improve Hugging Face support

Author: cbalioglu

src/fairseq2/checkpoint/_metadata_provider.py

@@ -20,15 +20,15 @@
 )
 from fairseq2.file_system import FileMode, FileSystem
 from fairseq2.gang import GangError, Gangs
-from fairseq2.models.llama import LLAMA_MODEL_FAMILY, LLaMAConfig
-from fairseq2.models.llama.integ import convert_to_hg_llama_config
 from fairseq2.utils.structured import unstructure
 from fairseq2.utils.yaml import YamlDumper
 
 
 class CheckpointMetadataSaver(ABC):
     @abstractmethod
-    def save(self, model_family: str, model_config: object) -> None: ...
+    def save(
+        self, model_family: str, model_config: object, hg_model_config: object = None
+    ) -> None: ...
 
 
 @final
@@ -50,39 +50,14 @@ def __init__(
         self._file_system = file_system
         self._yaml_dumper = yaml_dumper
 
-    def save(self, model_family: str, model_config: object) -> None:
+    def save(
+        self, model_family: str, model_config: object, hg_model_config: object = None
+    ) -> None:
         if self._gangs.root.rank == 0:
-            unstructured_config = unstructure(model_config)
-
-            metadata: dict[str, object] = {
-                "name": "checkpoint",
-                "model_family": model_family,
-                "model_config": {
-                    "_set_": unstructured_config,
-                },
-            }
+            self._save_asset_card(model_family, model_config)
 
-            if self._gangs.tp.size != 1:
-                metadata["num_shards"] = self._gangs.tp.size
-
-            metadata_file = self._checkpoint_dir.joinpath("model.yaml")
-
-            def save_error() -> AssetMetadataSaveError:
-                return AssetMetadataSaveError(
-                    f"The checkpoint metadata cannot be saved to the '{metadata_file}' file. See the nested exception for details."
-                )
-
-            try:
-                self._file_system.make_directory(metadata_file.parent)
-            except OSError as ex:
-                raise save_error() from ex
-
-            try:
-                self._yaml_dumper.dump(metadata, metadata_file)
-            except OSError as ex:
-                raise save_error() from ex
-
-            self._save_huggingface_config(model_family, model_config)
+            if hg_model_config is not None:
+                self._save_hg_config(hg_model_config)
 
         try:
             self._gangs.root.barrier()
@@ -91,17 +66,38 @@ def save_error() -> AssetMetadataSaveError:
                 "The collective barrier after the checkpoint metadata save operation has failed. See the nested exception for details."
             ) from ex
 
-    def _save_huggingface_config(self, model_family: str, model_config: object) -> None:
-        if model_family != LLAMA_MODEL_FAMILY:
-            return
+    def _save_asset_card(self, model_family: str, model_config: object) -> None:
+        unstructured_model_config = unstructure(model_config)
+
+        metadata: dict[str, object] = {
+            "name": "checkpoint",
+            "model_family": model_family,
+            "model_config": {
+                "_set_": unstructured_model_config,
+            },
+        }
 
-        if not isinstance(model_config, LLaMAConfig):
-            raise TypeError(
-                f"`model_config` must be of type `{LLaMAConfig}`, but is of type `{type(model_config)}` instead."
+        if self._gangs.tp.size != 1:
+            metadata["num_shards"] = self._gangs.tp.size
+
+        metadata_file = self._checkpoint_dir.joinpath("model.yaml")
+
+        def save_error() -> AssetMetadataSaveError:
+            return AssetMetadataSaveError(
+                f"The checkpoint metadata cannot be saved to the '{metadata_file}' file. See the nested exception for details."
             )
 
-        hg_config = convert_to_hg_llama_config(model_config)
+        try:
+            self._file_system.make_directory(metadata_file.parent)
+        except OSError as ex:
+            raise save_error() from ex
+
+        try:
+            self._yaml_dumper.dump(metadata, metadata_file)
+        except OSError as ex:
+            raise save_error() from ex
 
+    def _save_hg_config(self, hg_model_config: object) -> None:
         hg_config_file = self._checkpoint_dir.joinpath("cc/config.json")
 
         def save_error() -> AssetMetadataSaveError:
@@ -120,7 +116,7 @@ def save_error() -> AssetMetadataSaveError:
             raise save_error() from ex
 
         try:
-            json.dump(hg_config, fp, indent=2, sort_keys=True)
+            json.dump(hg_model_config, fp, indent=2, sort_keys=True)
         except OSError as ex:
             raise save_error() from ex
         finally:

src/fairseq2/cli/_setup.py

@@ -9,10 +9,7 @@
 from fairseq2.chatbots import UnknownChatbotError
 from fairseq2.cli.commands.assets import ListAssetsHandler, ShowAssetHandler
 from fairseq2.cli.commands.chatbot import RunChatbotHandler
-from fairseq2.cli.commands.llama import (
-    ConvertLLaMACheckpointHandler,
-    WriteHFLLaMAConfigHandler,
-)
+from fairseq2.cli.commands.llama import ConvertLLaMACheckpointHandler
 from fairseq2.cli.commands.recipe import RecipeCommandHandler
 from fairseq2.context import RuntimeContext
 from fairseq2.data.text.tokenizers import (
@@ -179,12 +176,6 @@ def _register_llama_cli(cli: Cli) -> None:
         help="convert fairseq2 LLaMA checkpoints to reference checkpoints",
     )
 
-    group.add_command(
-        name="write_hf_config",
-        handler=WriteHFLLaMAConfigHandler(),
-        help="write fairseq2 LLaMA configurations in Hugging Face format",
-    )
-
 
 def _register_lm_cli(cli: Cli) -> None:
     group = cli.add_group("lm", help="language model recipes")

src/fairseq2/cli/commands/llama/__init__.py

@@ -9,6 +9,3 @@
 from fairseq2.cli.commands.llama._convert_checkpoint import (
     ConvertLLaMACheckpointHandler as ConvertLLaMACheckpointHandler,
 )
-from fairseq2.cli.commands.llama._write_hf_config import (
-    WriteHFLLaMAConfigHandler as WriteHFLLaMAConfigHandler,
-)

src/fairseq2/cli/commands/llama/_convert_checkpoint.py

@@ -232,7 +232,7 @@ def file_write_error() -> CliCommandError:
             "dim": model_config.model_dim,
             "n_layers": model_config.num_layers,
             "n_heads": model_config.num_attn_heads,
-            "multiple_of": model_config.ffn_inner_dim_to_multiple,
+            "multiple_of": model_config.ffn_inner_dim_multiple_of,
             "rope_theta": model_config.rope_theta,
             "norm_eps": 1e-5,
         }

src/fairseq2/cli/commands/llama/_write_hf_config.py

@@ -25,6 +25,7 @@
 from fairseq2.models._handler import CheckpointConverter as CheckpointConverter
 from fairseq2.models._handler import DelegatingModelHandler as DelegatingModelHandler
 from fairseq2.models._handler import FsdpApplier as FsdpApplier
+from fairseq2.models._handler import HGConfigConverter as HGConfigConverter
 from fairseq2.models._handler import ModelCompiler as ModelCompiler
 from fairseq2.models._handler import ModelFactory as ModelFactory
 from fairseq2.models._handler import ModelHandler as ModelHandler