From 9e408dfd7ebe2535c3aebe57e494a9b0346680bc Mon Sep 17 00:00:00 2001
From: Kyle Sayers <kylesayrs@gmail.com>
Date: Thu, 25 Jun 2026 22:15:32 -0400
Subject: [PATCH 01/10] fix tensor get

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
---
 src/transformers/modeling_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 42cb3d974b0e..e915d6fefaa7 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -475,7 +475,7 @@ def remove_tied_weights_from_state_dict(
             # In offloaded cases, there may be meta tensors in the state_dict.
             # For these cases, key by the pointer of the original tensor object
             # (state_dict tensors are detached and therefore no longer shared)
-            tensor = model.get_parameter(name)
+            tensor = model.get_parameter_or_buffer(name)
             ptrs[id(tensor)].append(name)
 
         else:

From d38d3fa42e0fd5b7b7b26eaf61bdb6552ff8ca67 Mon Sep 17 00:00:00 2001
From: Kyle Sayers <kylesayrs@gmail.com>
Date: Fri, 26 Jun 2026 02:09:44 -0400
Subject: [PATCH 02/10] proper disk saving

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
---
 src/transformers/core_model_loading.py      |  2 +-
 src/transformers/integrations/accelerate.py | 46 ++++++++++++++++++---
 src/transformers/modeling_utils.py          | 12 ++++--
 3 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py
index 6c9fe775c8b5..105f934a287d 100644
--- a/src/transformers/core_model_loading.py
+++ b/src/transformers/core_model_loading.py
@@ -1508,7 +1508,7 @@ def convert_and_load_state_dict_in_model(
     return loading_info, disk_offload_index
 
 
-def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch.Tensor]):
+def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
     """
     Revert the conversion mapping that was used to load the model with `from_pretrained`, or the default one
     if the model was created in another way and is part of the default mappings.
diff --git a/src/transformers/integrations/accelerate.py b/src/transformers/integrations/accelerate.py
index 63a3ce22ea55..a7d4719ff8a0 100644
--- a/src/transformers/integrations/accelerate.py
+++ b/src/transformers/integrations/accelerate.py
@@ -519,11 +519,44 @@ def offload_weight(weight: torch.Tensor, weight_name: str, offload_folder: str |
     return offload_index
 
 
-def load_offloaded_parameter(model: "PreTrainedModel", param_name: str) -> torch.Tensor:
-    """Load `param_name` from disk, if it was offloaded due to the device_map, and thus lives as a meta parameter
-    inside `model`.
+def load_offloaded_parameter(model: "PreTrainedModel", param_name: str) -> dict[str, torch.Tensor]:
+    """Load source `param_name` from disk, if it was offloaded due to the device_map,
+    and thus lives as a meta parameter inside `model`.
+
     This is needed when resaving a model, when some parameters were offloaded (we need to load them from disk, to
-    then resave them to disk in the correct shard...)."""
+    then resave them to disk in the correct shard...).
+
+    Example flow:
+    ```
+    source_param_name -> "experts.0.w1.weight"
+    target_param_name -> "experts.gate_up_proj"
+    loaded_state_dict -> {
+        "experts.0.w1.weight": ...
+        "experts.1.w1.weight": ...
+        ...
+        "experts.0.w3.weight": ...
+        "experts.1.w3.weight": ...
+        ...
+    }
+    ```
+
+    Args:
+        model (`PreTrainedModel`): Model containing target offloaded weight to load
+        param_name (`str`): Name of source (checkpoint) weight to load from model
+
+    Returns:
+        `dict[str, torch.Tensor]`: Loaded state dict of all weights which map to the source weight
+    """
+    from ..core_model_loading import WeightRenaming, WeightConverter, rename_source_key, revert_weight_conversion
+
+    og = param_name
+
+    # Convert from source key in checkpoint to target key in model
+    meta_state_dict = model.state_dict()
+    renamings = [entry for entry in model._weight_conversions if isinstance(entry, WeightRenaming)]
+    converters = [entry for entry in model._weight_conversions if isinstance(entry, WeightConverter)]
+    param_name = rename_source_key(param_name, renamings, converters, model.base_model_prefix, meta_state_dict)[0]
+
     # Start from the most inner module, and try to find the hook that was used for offloading the param
     module_parts = param_name.split(".")
     modules_to_check = [".".join(module_parts[:-idx]) for idx in range(1, len(module_parts))] + [""]
@@ -542,7 +575,10 @@ def load_offloaded_parameter(model: "PreTrainedModel", param_name: str) -> torch
 
     # This call loads it from disk
     tensor = weights_map[truncated_param_name]
-    return tensor
+
+    # Convert from target key to source key(s)
+    loaded_state_dict = revert_weight_conversion(model, {param_name: tensor})
+    return loaded_state_dict
 
 
 def _init_infer_auto_device_map(
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index e915d6fefaa7..18105aaa8bd2 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -3511,8 +3511,10 @@ def save_pretrained(
         is_offloaded = False
         if (
             hasattr(self, "hf_device_map")
-            and len(set(self.hf_device_map.values())) > 1
-            and ("cpu" in self.hf_device_map.values() or "disk" in self.hf_device_map.values())
+            and (
+                len(set(self.hf_device_map.values())) > 1
+                or "disk" in self.hf_device_map.values()
+            )
         ):
             is_offloaded = True
             warnings.warn(
@@ -3593,9 +3595,11 @@ def save_pretrained(
 
                 # If the param was offloaded, we need to load it back from disk to resave it. It's a strange pattern,
                 # but it would otherwise not be contained in the saved shard if we were to simply move the file
-                # or something
+                # or something. Note that multiple weights may be loaded by `load_offloaded_parameter`
+                # but each weight is only loaded once
                 if is_offloaded and tensor.device.type == "meta":
-                    tensor = load_offloaded_parameter(model_to_save, tensor_name)
+                    state_dict.update(load_offloaded_parameter(model_to_save, tensor_name))
+                    tensor = state_dict.pop(tensor_name)
 
                 # only do contiguous after it's permuted correctly in case of TP
                 shard_state_dict[tensor_name] = tensor.contiguous()

From cd3dcec3d8532067a898f77090b54e53409b485f Mon Sep 17 00:00:00 2001
From: Kyle Sayers <kylesayrs@gmail.com>
Date: Fri, 26 Jun 2026 02:26:59 -0400
Subject: [PATCH 03/10] cleanup

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
---
 src/transformers/integrations/accelerate.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/transformers/integrations/accelerate.py b/src/transformers/integrations/accelerate.py
index a7d4719ff8a0..6dab49bb3a82 100644
--- a/src/transformers/integrations/accelerate.py
+++ b/src/transformers/integrations/accelerate.py
@@ -549,8 +549,6 @@ def load_offloaded_parameter(model: "PreTrainedModel", param_name: str) -> dict[
     """
     from ..core_model_loading import WeightRenaming, WeightConverter, rename_source_key, revert_weight_conversion
 
-    og = param_name
-
     # Convert from source key in checkpoint to target key in model
     meta_state_dict = model.state_dict()
     renamings = [entry for entry in model._weight_conversions if isinstance(entry, WeightRenaming)]

From c015d80eace1aebf175f9939de9c1177d30417b9 Mon Sep 17 00:00:00 2001
From: Kyle Sayers <kylesayrs@gmail.com>
Date: Fri, 26 Jun 2026 13:45:49 -0400
Subject: [PATCH 04/10] cache meta_data_dict

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
---
 src/transformers/integrations/accelerate.py | 11 ++++++++---
 src/transformers/modeling_utils.py          |  3 ++-
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/transformers/integrations/accelerate.py b/src/transformers/integrations/accelerate.py
index 6dab49bb3a82..bae4845937cd 100644
--- a/src/transformers/integrations/accelerate.py
+++ b/src/transformers/integrations/accelerate.py
@@ -22,7 +22,7 @@
 import re
 from collections import OrderedDict, defaultdict
 from collections.abc import Callable
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional, Any
 
 from safetensors import safe_open
 from safetensors.torch import save_file
@@ -519,7 +519,11 @@ def offload_weight(weight: torch.Tensor, weight_name: str, offload_folder: str |
     return offload_index
 
 
-def load_offloaded_parameter(model: "PreTrainedModel", param_name: str) -> dict[str, torch.Tensor]:
+def load_offloaded_parameter(
+    model: "PreTrainedModel",
+    param_name: str,
+    meta_state_dict: Optional[dict[str, torch.Tensor | Any]] = None,
+) -> dict[str, torch.Tensor]:
     """Load source `param_name` from disk, if it was offloaded due to the device_map,
     and thus lives as a meta parameter inside `model`.
 
@@ -550,7 +554,8 @@ def load_offloaded_parameter(model: "PreTrainedModel", param_name: str) -> dict[
     from ..core_model_loading import WeightRenaming, WeightConverter, rename_source_key, revert_weight_conversion
 
     # Convert from source key in checkpoint to target key in model
-    meta_state_dict = model.state_dict()
+    if meta_state_dict is None:
+        meta_state_dict = model.state_dict()  # this can be costly: please pass as arg when possible
     renamings = [entry for entry in model._weight_conversions if isinstance(entry, WeightRenaming)]
     converters = [entry for entry in model._weight_conversions if isinstance(entry, WeightConverter)]
     param_name = rename_source_key(param_name, renamings, converters, model.base_model_prefix, meta_state_dict)[0]
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 18105aaa8bd2..6fdf5003e60b 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -3584,6 +3584,7 @@ def save_pretrained(
                 os.remove(full_filename)
 
         # Save the model
+        meta_state_dict = model_to_save.state_dict()
         for shard_file, tensor_names in logging.tqdm(
             state_dict_split.filename_to_tensors.items(), desc="Writing model shards"
         ):
@@ -3598,7 +3599,7 @@ def save_pretrained(
                 # or something. Note that multiple weights may be loaded by `load_offloaded_parameter`
                 # but each weight is only loaded once
                 if is_offloaded and tensor.device.type == "meta":
-                    state_dict.update(load_offloaded_parameter(model_to_save, tensor_name))
+                    state_dict.update(load_offloaded_parameter(model_to_save, tensor_name, meta_state_dict))
                     tensor = state_dict.pop(tensor_name)
 
                 # only do contiguous after it's permuted correctly in case of TP

From 324a8b787263f6107a739894558d848babc4bd33 Mon Sep 17 00:00:00 2001
From: Kyle Sayers <kylesayrs@gmail.com>
Date: Fri, 26 Jun 2026 14:12:04 -0400
Subject: [PATCH 05/10] fix style

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
---
 src/transformers/integrations/accelerate.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/transformers/integrations/accelerate.py b/src/transformers/integrations/accelerate.py
index bae4845937cd..26d8820e211f 100644
--- a/src/transformers/integrations/accelerate.py
+++ b/src/transformers/integrations/accelerate.py
@@ -22,7 +22,10 @@
 import re
 from collections import OrderedDict, defaultdict
 from collections.abc import Callable
-from typing import TYPE_CHECKING, Optional, Any
+from typing import TYPE_CHECKING, Any
+
+from safetensors import safe_open
+from safetensors.torch import save_file
 
 from safetensors import safe_open
 from safetensors.torch import save_file
@@ -522,7 +525,7 @@ def offload_weight(weight: torch.Tensor, weight_name: str, offload_folder: str |
 def load_offloaded_parameter(
     model: "PreTrainedModel",
     param_name: str,
-    meta_state_dict: Optional[dict[str, torch.Tensor | Any]] = None,
+    meta_state_dict: dict[str, torch.Tensor | Any] | None = None,
 ) -> dict[str, torch.Tensor]:
     """Load source `param_name` from disk, if it was offloaded due to the device_map,
     and thus lives as a meta parameter inside `model`.
@@ -551,7 +554,7 @@ def load_offloaded_parameter(
     Returns:
         `dict[str, torch.Tensor]`: Loaded state dict of all weights which map to the source weight
     """
-    from ..core_model_loading import WeightRenaming, WeightConverter, rename_source_key, revert_weight_conversion
+    from ..core_model_loading import WeightConverter, WeightRenaming, rename_source_key, revert_weight_conversion
 
     # Convert from source key in checkpoint to target key in model
     if meta_state_dict is None:

From e1232e6ba4a273c09090f826eab947df91389940 Mon Sep 17 00:00:00 2001
From: Kyle Sayers <kylesayrs@gmail.com>
Date: Fri, 26 Jun 2026 15:01:28 -0400
Subject: [PATCH 06/10] better comments

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
---
 src/transformers/modeling_utils.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 6fdf5003e60b..6ee7404bb5cc 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -3519,7 +3519,8 @@ def save_pretrained(
             is_offloaded = True
             warnings.warn(
                 "Attempting to save a model with offloaded modules. Ensure that unallocated cpu memory "
-                "exceeds the `shard_size` (50GB default)"
+                "exceeds the `shard_size` (50GB default) and/or the largest model weight size (this can "
+                "be very large for MoE models with fused experts)."
             )
 
         # Translate state_dict from smp to hf if saving with smp >= 1.10
@@ -3590,15 +3591,16 @@ def save_pretrained(
         ):
             filename = os.path.join(save_directory, shard_file)
             shard_state_dict = {}
-            for tensor_name in tensor_names:
+            for tensor_name in sorted(tensor_names):
                 # Get the tensor, and remove it from state_dict to avoid keeping the ref
                 tensor = state_dict.pop(tensor_name)
 
-                # If the param was offloaded, we need to load it back from disk to resave it. It's a strange pattern,
-                # but it would otherwise not be contained in the saved shard if we were to simply move the file
-                # or something. Note that multiple weights may be loaded by `load_offloaded_parameter`
-                # but each weight is only loaded once
+                # If the param was offloaded, we need to load it back onto cpu from disk to resave it.
+                # It's a strange pattern, but is necessary to ensure saving into the proper file shard
                 if is_offloaded and tensor.device.type == "meta":
+                    # Note that `load_offloaded_parameter` may load multiple weights for a single tensor.
+                    # While it is possible to overload CPU memory by loading parameters in a bad order,
+                    # in practice `split_torch_state_dict_into_shards` preserves weight locality
                     state_dict.update(load_offloaded_parameter(model_to_save, tensor_name, meta_state_dict))
                     tensor = state_dict.pop(tensor_name)
 

From 52d3e4f51fb005012801c51820f2f9c3a70daf81 Mon Sep 17 00:00:00 2001
From: Kyle Sayers <kylesayrs@gmail.com>
Date: Mon, 29 Jun 2026 13:57:57 -0400
Subject: [PATCH 07/10] format

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
---
 src/transformers/modeling_utils.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 6ee7404bb5cc..45a1b323f8d4 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -3509,12 +3509,8 @@ def save_pretrained(
 
         # if any model parameters are offloaded, we need to know it for later
         is_offloaded = False
-        if (
-            hasattr(self, "hf_device_map")
-            and (
-                len(set(self.hf_device_map.values())) > 1
-                or "disk" in self.hf_device_map.values()
-            )
+        if hasattr(self, "hf_device_map") and (
+            len(set(self.hf_device_map.values())) > 1 or "disk" in self.hf_device_map.values()
         ):
             is_offloaded = True
             warnings.warn(

From d9281e39d07dcaa637b3abe1a71aaea8e46136fa Mon Sep 17 00:00:00 2001
From: Kyle Sayers <kylesayrs@gmail.com>
Date: Wed, 1 Jul 2026 11:49:12 -0400
Subject: [PATCH 08/10] break out load_offloaded_checkpoint_parameters

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
---
 src/transformers/integrations/accelerate.py | 54 ++++++++++++---------
 src/transformers/modeling_utils.py          |  4 +-
 2 files changed, 33 insertions(+), 25 deletions(-)

diff --git a/src/transformers/integrations/accelerate.py b/src/transformers/integrations/accelerate.py
index 26d8820e211f..09d5cd035ffc 100644
--- a/src/transformers/integrations/accelerate.py
+++ b/src/transformers/integrations/accelerate.py
@@ -27,9 +27,6 @@
 from safetensors import safe_open
 from safetensors.torch import save_file
 
-from safetensors import safe_open
-from safetensors.torch import save_file
-
 from ..utils import (
     is_accelerate_available,
     is_torch_available,
@@ -522,12 +519,39 @@ def offload_weight(weight: torch.Tensor, weight_name: str, offload_folder: str |
     return offload_index
 
 
-def load_offloaded_parameter(
+def load_offloaded_parameter(model: "PreTrainedModel", param_name: str) -> torch.Tensor:
+    """Load `param_name` from disk, if it was offloaded due to the device_map, and thus lives as a meta parameter
+    inside `model`.
+    This is needed when resaving a model, when some parameters were offloaded (we need to load them from disk, to
+    then resave them to disk in the correct shard...)."""
+    # Start from the most inner module, and try to find the hook that was used for offloading the param
+    module_parts = param_name.split(".")
+    modules_to_check = [".".join(module_parts[:-idx]) for idx in range(1, len(module_parts))] + [""]
+    for parent_name in modules_to_check:
+        parent = model.get_submodule(parent_name)
+        if hasattr(parent, "_hf_hook"):
+            weights_map = parent._hf_hook.weights_map
+            truncated_param_name = param_name.replace(f"{parent_name}." if parent_name != "" else parent_name, "")
+            break
+    # If we did not break the loop, something is wrong
+    else:
+        raise ValueError(
+            f"{param_name} is on the meta device because it was offloaded, but we could not find "
+            "the corresponding hook for it"
+        )
+
+    # This call loads it from disk
+    tensor = weights_map[truncated_param_name]
+    return tensor
+
+
+def load_offloaded_checkpoint_parameters(
     model: "PreTrainedModel",
     param_name: str,
     meta_state_dict: dict[str, torch.Tensor | Any] | None = None,
 ) -> dict[str, torch.Tensor]:
-    """Load source `param_name` from disk, if it was offloaded due to the device_map,
+    """
+    Load source `param_name` from disk, if it was offloaded due to the device_map,
     and thus lives as a meta parameter inside `model`.
 
     This is needed when resaving a model, when some parameters were offloaded (we need to load them from disk, to
@@ -563,24 +587,8 @@ def load_offloaded_parameter(
     converters = [entry for entry in model._weight_conversions if isinstance(entry, WeightConverter)]
     param_name = rename_source_key(param_name, renamings, converters, model.base_model_prefix, meta_state_dict)[0]
 
-    # Start from the most inner module, and try to find the hook that was used for offloading the param
-    module_parts = param_name.split(".")
-    modules_to_check = [".".join(module_parts[:-idx]) for idx in range(1, len(module_parts))] + [""]
-    for parent_name in modules_to_check:
-        parent = model.get_submodule(parent_name)
-        if hasattr(parent, "_hf_hook"):
-            weights_map = parent._hf_hook.weights_map
-            truncated_param_name = param_name.replace(f"{parent_name}." if parent_name != "" else parent_name, "")
-            break
-    # If we did not break the loop, something is wrong
-    else:
-        raise ValueError(
-            f"{param_name} is on the meta device because it was offloaded, but we could not find "
-            "the corresponding hook for it"
-        )
-
-    # This call loads it from disk
-    tensor = weights_map[truncated_param_name]
+    # load parameter from model
+    tensor = load_offloaded_parameter(param_name)
 
     # Convert from target key to source key(s)
     loaded_state_dict = revert_weight_conversion(model, {param_name: tensor})
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index 45a1b323f8d4..b7d083be5381 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -62,7 +62,7 @@
     check_and_set_device_map,
     expand_device_map,
     get_device,
-    load_offloaded_parameter,
+    load_offloaded_checkpoint_parameters,
 )
 from .integrations.deepspeed import _load_state_dict_into_zero3_model
 from .integrations.eager_paged import eager_paged_attention_forward
@@ -3597,7 +3597,7 @@ def save_pretrained(
                     # Note that `load_offloaded_parameter` may load multiple weights for a single tensor.
                     # While it is possible to overload CPU memory by loading parameters in a bad order,
                     # in practice `split_torch_state_dict_into_shards` preserves weight locality
-                    state_dict.update(load_offloaded_parameter(model_to_save, tensor_name, meta_state_dict))
+                    state_dict.update(load_offloaded_checkpoint_parameters(model_to_save, tensor_name, meta_state_dict))
                     tensor = state_dict.pop(tensor_name)
 
                 # only do contiguous after it's permuted correctly in case of TP

From c3caa61dbaf47a0d26da6280c12e3a7e775e09e0 Mon Sep 17 00:00:00 2001
From: Kyle Sayers <kylesayrs@gmail.com>
Date: Wed, 1 Jul 2026 23:50:36 -0400
Subject: [PATCH 09/10] fix typo

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
---
 src/transformers/integrations/accelerate.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/transformers/integrations/accelerate.py b/src/transformers/integrations/accelerate.py
index 09d5cd035ffc..0c258482b5fd 100644
--- a/src/transformers/integrations/accelerate.py
+++ b/src/transformers/integrations/accelerate.py
@@ -588,7 +588,7 @@ def load_offloaded_checkpoint_parameters(
     param_name = rename_source_key(param_name, renamings, converters, model.base_model_prefix, meta_state_dict)[0]
 
     # load parameter from model
-    tensor = load_offloaded_parameter(param_name)
+    tensor = load_offloaded_parameter(model, param_name)
 
     # Convert from target key to source key(s)
     loaded_state_dict = revert_weight_conversion(model, {param_name: tensor})

From dbc8c39761f1b21ff98c1b992f0e23fe6241c7b8 Mon Sep 17 00:00:00 2001
From: Kyle Sayers <kylesayrs@gmail.com>
Date: Thu, 2 Jul 2026 00:44:49 -0400
Subject: [PATCH 10/10] fix style

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
---
 src/transformers/modeling_utils.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
index b7d083be5381..02b0fcb58304 100644
--- a/src/transformers/modeling_utils.py
+++ b/src/transformers/modeling_utils.py
@@ -3597,7 +3597,9 @@ def save_pretrained(
                     # Note that `load_offloaded_parameter` may load multiple weights for a single tensor.
                     # While it is possible to overload CPU memory by loading parameters in a bad order,
                     # in practice `split_torch_state_dict_into_shards` preserves weight locality
-                    state_dict.update(load_offloaded_checkpoint_parameters(model_to_save, tensor_name, meta_state_dict))
+                    state_dict.update(
+                        load_offloaded_checkpoint_parameters(model_to_save, tensor_name, meta_state_dict)
+                    )
                     tensor = state_dict.pop(tensor_name)
 
                 # only do contiguous after it's permuted correctly in case of TP