Introduce weights sharding

Author: james77777778

keras/src/models/model.py

@@ -313,47 +313,72 @@ def save(self, filepath, overwrite=True, zipped=None, **kwargs):
         )
 
     @traceback_utils.filter_traceback
-    def save_weights(self, filepath, overwrite=True):
-        """Saves all layer weights to a `.weights.h5` file.
+    def save_weights(self, filepath, overwrite=True, max_shard_size=None):
+        """Saves all weights to a single file or sharded files.
+
+        By default, the weights will be saved in a single `.weights.h5` file.
+        If sharding is enabled (`max_shard_size` is not `None`), the weights
+        will be saved in multiple files, each with a size at most
+        `max_shard_size` (in GB). Additionally, a configuration file
+        `.weights.json` will contain the metadata for the sharded files.
 
         Args:
-            filepath: `str` or `pathlib.Path` object.
-                Path where to save the model. Must end in `.weights.h5`.
-            overwrite: Whether we should overwrite any existing model
-                at the target location, or instead ask the user
-                via an interactive prompt.
+            filepath: `str` or `pathlib.Path` object. Path where the weights
+                will be saved.  When sharding, the filepath must end in
+                `.weights.json`. If `.weights.h5` is provided, it will be
+                overridden.
+            overwrite: Whether to overwrite any existing weights at the target
+                location or instead ask the user via an interactive prompt.
+            max_shard_size: `int` or `float`. Maximum size in GB for each
+                sharded file. If `None`, no sharding will be done. Defaults to
+                `None`.
         """
-        return saving_api.save_weights(self, filepath, overwrite=overwrite)
+        return saving_api.save_weights(
+            self, filepath, overwrite=overwrite, max_shard_size=max_shard_size
+        )
 
     @traceback_utils.filter_traceback
-    def load_weights(self, filepath, skip_mismatch=False, **kwargs):
-        """Load weights from a file saved via `save_weights()`.
+    def load_weights(
+        self, filepath, skip_mismatch=False, sharded=False, **kwargs
+    ):
+        """Load the weights from a single file or sharded files.
 
-        Weights are loaded based on the network's
-        topology. This means the architecture should be the same as when the
-        weights were saved. Note that layers that don't have weights are not
-        taken into account in the topological ordering, so adding or removing
-        layers is fine as long as they don't have weights.
+        Weights are loaded based on the network's topology. This means the
+        architecture should be the same as when the weights were saved. Note
+        that layers that don't have weights are not taken into account in the
+        topological ordering, so adding or removing layers is fine as long as
+        they don't have weights.
 
         **Partial weight loading**
 
         If you have modified your model, for instance by adding a new layer
-        (with weights) or by changing the shape of the weights of a layer,
-        you can choose to ignore errors and continue loading
-        by setting `skip_mismatch=True`. In this case any layer with
-        mismatching weights will be skipped. A warning will be displayed
-        for each skipped layer.
+        (with weights) or by changing the shape of the weights of a layer, you
+        can choose to ignore errors and continue loading by setting
+        `skip_mismatch=True`. In this case any layer with mismatching weights
+        will be skipped. A warning will be displayed for each skipped layer.
+
+        **Sharding**
+
+        When loading sharded weights, it is important to set `sharded=True` and
+        specify `filepath` that ends with `.weights.json`.
 
         Args:
-            filepath: String, path to the weights file to load.
-                It can either be a `.weights.h5` file
-                or a legacy `.h5` weights file.
+            filepath: `str` or `pathlib.Path` object. Path where the weights
+                will be saved.  When sharding, the filepath must end in
+                `.weights.json`. If `.weights.h5` is provided, it will be
+                overridden.
             skip_mismatch: Boolean, whether to skip loading of layers where
                 there is a mismatch in the number of weights, or a mismatch in
                 the shape of the weights.
+            sharded: Whether the saved file(s) are sharded. Defaults to
+                `False`.
         """
         saving_api.load_weights(
-            self, filepath, skip_mismatch=skip_mismatch, **kwargs
+            self,
+            filepath,
+            skip_mismatch=skip_mismatch,
+            sharded=sharded,
+            **kwargs,
         )
 
     def quantize(self, mode, **kwargs):

keras/src/saving/saving_api.py

@@ -219,32 +219,45 @@ def load_model(filepath, custom_objects=None, compile=True, safe_mode=True):
 
 
 @keras_export("keras.saving.save_weights")
-def save_weights(model, filepath, overwrite=True, **kwargs):
-    if not str(filepath).endswith(".weights.h5"):
+def save_weights(
+    model, filepath, overwrite=True, max_shard_size=None, **kwargs
+):
+    filepath_str = str(filepath)
+    if max_shard_size is None and not filepath_str.endswith(".weights.h5"):
         raise ValueError(
             "The filename must end in `.weights.h5`. "
-            f"Received: filepath={filepath}"
+            f"Received: filepath={filepath_str}"
+        )
+    elif max_shard_size is not None and not filepath_str.endswith(
+        ("weights.h5", "weights.json")
+    ):
+        raise ValueError(
+            "The filename must end in `.weights.json` when `max_shard_size` is "
+            f"specified. Received: filepath={filepath_str}"
         )
     try:
         exists = os.path.exists(filepath)
     except TypeError:
         exists = False
     if exists and not overwrite:
-        proceed = io_utils.ask_to_proceed_with_overwrite(filepath)
+        proceed = io_utils.ask_to_proceed_with_overwrite(filepath_str)
         if not proceed:
             return
-    saving_lib.save_weights_only(model, filepath, **kwargs)
+    saving_lib.save_weights_only(model, filepath, max_shard_size, **kwargs)
 
 
 @keras_export("keras.saving.load_weights")
-def load_weights(model, filepath, skip_mismatch=False, **kwargs):
-    if str(filepath).endswith(".keras"):
+def load_weights(model, filepath, skip_mismatch=False, sharded=False, **kwargs):
+    filepath_str = str(filepath)
+    if filepath_str.endswith(".keras"):
         if kwargs:
             raise ValueError(f"Invalid keyword arguments: {kwargs}")
         saving_lib.load_weights_only(
-            model, filepath, skip_mismatch=skip_mismatch
+            model, filepath, skip_mismatch=skip_mismatch, sharded=sharded
         )
-    elif str(filepath).endswith(".weights.h5"):
+    elif filepath_str.endswith(".weights.h5") or filepath_str.endswith(
+        ".weights.json"
+    ):
         objects_to_skip = kwargs.pop("objects_to_skip", None)
         if kwargs:
             raise ValueError(f"Invalid keyword arguments: {kwargs}")
@@ -253,8 +266,9 @@ def load_weights(model, filepath, skip_mismatch=False, **kwargs):
             filepath,
             skip_mismatch=skip_mismatch,
             objects_to_skip=objects_to_skip,
+            sharded=sharded,
         )
-    elif str(filepath).endswith(".h5") or str(filepath).endswith(".hdf5"):
+    elif filepath_str.endswith(".h5") or filepath_str.endswith(".hdf5"):
         by_name = kwargs.pop("by_name", False)
         if kwargs:
             raise ValueError(f"Invalid keyword arguments: {kwargs}")

keras/src/saving/saving_api_test.py

@@ -1,4 +1,5 @@
 import os
+import pathlib
 import unittest.mock as mock
 
 import numpy as np
@@ -237,6 +238,20 @@ def test_load_weights_invalid_extension(self):
         with self.assertRaisesRegex(ValueError, "File format not supported"):
             model.load_weights("invalid_extension.pkl")
 
+    def test_load_sharded_weights(self):
+        src_model = self.get_model()
+        temp_filepath = pathlib.Path(
+            os.path.join(self.get_temp_dir(), "test_weights.weights.json")
+        )
+        src_model.save_weights(temp_filepath, max_shard_size=1)
+        self.assertLen(os.listdir(temp_filepath.parent), 2)
+        src_weights = src_model.get_weights()
+        dest_model = self.get_model()
+        dest_model.load_weights(temp_filepath, sharded=True)
+        dest_weights = dest_model.get_weights()
+        for orig, loaded in zip(src_weights, dest_weights):
+            self.assertAllClose(orig, loaded)
+
 
 class SaveModelTestsWarning(test_case.TestCase):
     def get_model(self):