tmp

Author: cbalioglu

src/fairseq2/checkpoint/_manager.py

@@ -6,7 +6,6 @@
 
 from __future__ import annotations
 
-import json
 from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Iterable, final
@@ -18,10 +17,8 @@
     AssetMetadataSaveError,
     CachedAssetMetadataProvider,
 )
-from fairseq2.file_system import FileMode, FileSystem
+from fairseq2.file_system import 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
 
@@ -52,37 +49,7 @@ def __init__(
 
     def save(self, model_family: str, model_config: object) -> 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,
-                },
-            }
-
-            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)
+            self._save_asset_card(model_family, model_config)
 
         try:
             self._gangs.root.barrier()
@@ -91,40 +58,36 @@ 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)
 
-        if not isinstance(model_config, LLaMAConfig):
-            raise TypeError(
-                f"`model_config` must be of type `{LLaMAConfig}`, but is of type `{type(model_config)}` instead."
-            )
+        metadata: dict[str, object] = {
+            "name": "checkpoint",
+            "model_family": model_family,
+            "model_config": {
+                "_set_": unstructured_model_config,
+            },
+        }
 
-        hg_config = convert_to_hg_llama_config(model_config)
+        if self._gangs.tp.size != 1:
+            metadata["num_shards"] = self._gangs.tp.size
 
-        hg_config_file = self._checkpoint_dir.joinpath("cc/config.json")
+        metadata_file = self._checkpoint_dir.joinpath("model.yaml")
 
         def save_error() -> AssetMetadataSaveError:
             return AssetMetadataSaveError(
-                f"The Hugging Face model configuration cannot be saved to the '{hg_config_file}' file. See the nested exception for details."
+                f"The checkpoint metadata cannot be saved to the '{metadata_file}' file. See the nested exception for details."
             )
 
         try:
-            self._file_system.make_directory(hg_config_file.parent)
-        except OSError as ex:
-            raise save_error() from ex
-
-        try:
-            fp = self._file_system.open_text(hg_config_file, mode=FileMode.WRITE)
+            self._file_system.make_directory(metadata_file.parent)
         except OSError as ex:
             raise save_error() from ex
 
         try:
-            json.dump(hg_config, fp, indent=2, sort_keys=True)
+            self._yaml_dumper.dump(metadata, metadata_file)
         except OSError as ex:
             raise save_error() from ex
-        finally:
-            fp.close()
 
 
 @final
@@ -170,23 +133,10 @@ def _load_cache(self) -> dict[str, dict[str, object]]:
                 "The checkpoint metadata does not have a 'checkpoint@' entry."
             ) from None
 
-        num_shards = metadata.get("num_shards", 1)
-
-        if not isinstance(num_shards, int) or num_shards < 1:
-            raise AssetMetadataLoadError(
-                "The 'num_shards' value in the checkpoint metadata is not a positive integer."
-            )
-
-        if num_shards == 1:
-            filename = "model.pt"
-        else:
-            # TODO: Fix once DownloadManager refactoring complete!
-            filename = "model.0{shard_idx}.pt"
-
         def add_checkpoint_metadata(name: str, step_nr: int) -> None:
-            file = self._checkpoint_dir.joinpath(f"step_{step_nr}/{filename}")
+            path = self._checkpoint_dir.joinpath(f"step_{step_nr}")
 
-            cache[name] = {"base": "checkpoint", "checkpoint": str(file)}
+            cache[name] = {"base": "checkpoint", "checkpoint": str(path)}
 
         max_step_nr = -1
 

src/fairseq2/checkpoint/_metadata_provider.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 (
@@ -37,7 +34,6 @@
 from fairseq2.metrics.text import UnknownBleuTokenizerError
 from fairseq2.models import (
     InvalidModelTypeError,
-    ShardedModelLoadError,
     UnknownModelArchitectureError,
     UnknownModelError,
     UnknownModelFamilyError,
@@ -120,7 +116,7 @@ def setup_cli(context: RuntimeContext) -> Cli:
 
     signature = "extension_function(context: RuntimeContext, cli: Cli) -> None"
 
-    run_extensions("fairseq2.cli", signature, context, cli)
+    run_extensions("fairseq2.cli", signature, cli, context)
 
     return cli
 
@@ -179,12 +175,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")
@@ -357,7 +347,6 @@ def _register_user_error_types(cli: Cli) -> None:
     cli.register_user_error_type(ModelCompilationNotSupportedError)
     cli.register_user_error_type(ModelParallelismNotSupportedError)
     cli.register_user_error_type(ModelPathNotFoundError)
-    cli.register_user_error_type(ShardedModelLoadError)
     cli.register_user_error_type(UnknownBeamSearchAlgorithmError)
     cli.register_user_error_type(UnknownBleuTokenizerError)
     cli.register_user_error_type(UnknownChatbotError)

src/fairseq2/cli/_setup.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/__init__.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/_convert_checkpoint.py

@@ -591,7 +591,7 @@ def setup_parallel_gangs(root_gang: Gang, *, tp_size: int = 1) -> Gangs:
     mesh = torch.arange(root_gang.size).view(dp_size, tp_size)
 
     # Get the coordinate of this process in the mesh.
-    rank_coords = [x.item() for x in torch.where(mesh == root_gang.rank)]
+    rank_coord = [x.item() for x in torch.where(mesh == root_gang.rank)]
 
     dp_gang: Gang | None = None
 
@@ -619,7 +619,7 @@ def setup_parallel_gangs(root_gang: Gang, *, tp_size: int = 1) -> Gangs:
             else:
                 for i in range(tp_size):
                     sub_gang = root_gang.create_gang(mesh[:, i].tolist())
-                    if i == rank_coords[1]:
+                    if i == rank_coord[1]:
                         dp_gang = sub_gang
 
     if dp_gang is None:
@@ -651,7 +651,7 @@ def setup_parallel_gangs(root_gang: Gang, *, tp_size: int = 1) -> Gangs:
             else:
                 for i in range(dp_size):
                     sub_gang = root_gang.create_gang(mesh[i, :].tolist())
-                    if i == rank_coords[0]:
+                    if i == rank_coord[0]:
                         tp_gang = sub_gang
 
     if tp_gang is None:
@@ -701,7 +701,7 @@ def setup_fsdp_gangs(gangs: Gangs, intra_node_size: int | None = None) -> Gangs:
     mesh = torch.arange(dp_gang.size).view(inter_node_size, intra_node_size)
 
     # Get the coordinate of this process in the mesh.
-    rank_coords = [x.item() for x in torch.where(mesh == dp_gang.rank)]
+    rank_coord = [x.item() for x in torch.where(mesh == dp_gang.rank)]
 
     inter_gang: Gang | None = None
 
@@ -729,7 +729,7 @@ def setup_fsdp_gangs(gangs: Gangs, intra_node_size: int | None = None) -> Gangs:
             else:
                 for i in range(intra_node_size):
                     sub_gang = dp_gang.create_gang(mesh[:, i].tolist())
-                    if i == rank_coords[1]:
+                    if i == rank_coord[1]:
                         inter_gang = sub_gang
 
     if inter_gang is None:
@@ -761,7 +761,7 @@ def setup_fsdp_gangs(gangs: Gangs, intra_node_size: int | None = None) -> Gangs:
             else:
                 for i in range(inter_node_size):
                     sub_gang = dp_gang.create_gang(mesh[i, :].tolist())
-                    if i == rank_coords[0]:
+                    if i == rank_coord[0]:
                         intra_gang = sub_gang
 
     if intra_gang is None:

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 HuggingFaceExporter as HuggingFaceExporter
 from fairseq2.models._handler import ModelCompiler as ModelCompiler
 from fairseq2.models._handler import ModelFactory as ModelFactory
 from fairseq2.models._handler import ModelHandler as ModelHandler