code: def shard_checkpoint(
      state_dict: Dict[str, torch.Tensor],
      max_shard_size: Union[int, str] = "10GB",
      weights_name: str = WEIGHTS_NAME,
  ):
      """
      Splits a model state dictionary in sub-checkpoints so that the final size of each sub-checkpoint does not exceed a
      given size.
  
      The sub-checkpoints are determined by iterating through the `state_dict` in the order of its keys, so there is no
      optimization made to make each sub-checkpoint as close as possible to the maximum size passed. For example, if the
      limit is 10GB and we have weights of sizes [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] they will get sharded as [6GB], [6+2GB],
      [6+2+2GB] and not [6+2+2GB], [6+2GB], [6GB].
  
      <Tip warning={true}>
  
      If one of the model's weight is bigger that `max_sahrd_size`, it will end up in its own sub-checkpoint which will
      have a size greater than `max_shard_size`.
  
      </Tip>
  
      Args:
          state_dict (`Dict[str, torch.Tensor]`): The state dictionary of a model to save.
          max_shard_size (`int` or `str`, *optional*, defaults to `"10GB"`):
              The maximum size of each sub-checkpoint. If expressed as a string, needs to be digits followed by a unit
              (like `"5MB"`).
          weights_name (`str`, *optional*, defaults to `"pytorch_model.bin"`):
              The name of the model save file.
      """
      max_shard_size = convert_file_size_to_int(max_shard_size)
  
      sharded_state_dicts = [{}]
      last_block_size = 0
      total_size = 0
      storage_id_to_block = {}
  
      for key, weight in state_dict.items():
          # when bnb serialization is used the weights in the state dict can be strings
          # check: https://github.com/huggingface/transformers/pull/24416 for more details
          if isinstance(weight, str):
              continue
          else:
              storage_id = id_tensor_storage(weight)
  
          # If a `weight` shares the same underlying storage as another tensor, we put `weight` in the same `block`
          if storage_id in storage_id_to_block:
              block_id = storage_id_to_block[storage_id]
              sharded_state_dicts[block_id][key] = weight
              continue
  
          weight_size = weight.numel() * dtype_byte_size(weight.dtype)
  
          # If this weight is going to tip up over the maximal size, we split.
          if last_block_size + weight_size > max_shard_size:
              sharded_state_dicts.append({})
              last_block_size = 0
  
          sharded_state_dicts[-1][key] = weight
          last_block_size += weight_size
          total_size += weight_size
          storage_id_to_block[storage_id] = len(sharded_state_dicts) - 1
  
      # If we only have one shard, we return it
      if len(sharded_state_dicts) == 1:
          return {weights_name: sharded_state_dicts[0]}, None
  
      # Otherwise, let's build the index
      weight_map = {}
      shards = {}
      for idx, shard in enumerate(sharded_state_dicts):
          shard_file = weights_name.replace(
              ".bin",
              f"-{idx + 1:05d}-of-{len(sharded_state_dicts):05d}.bin",
          )
          shard_file = shard_file.replace(
              ".safetensors",
              f"-{idx + 1:05d}-of-{len(sharded_state_dicts):05d}.safetensors",
          )
          shards[shard_file] = shard
          for key in shard.keys():
              weight_map[key] = shard_file
  
      # Add the metadata
      metadata = {"total_size": total_size}
      index = {"metadata": metadata, "weight_map": weight_map}
      return shards, index
  
imports: import contextlib
  import gc
  import inspect
  import json
  import logging
  import os
  import re
  import shutil
  import tempfile
  from collections import OrderedDict, defaultdict
  from typing import Dict, List, Optional, Tuple, Union
  import torch
  import torch.nn
  from ..state import AcceleratorState
  from .constants import SAFE_WEIGHTS_NAME, WEIGHTS_NAME
  from .dataclasses import AutocastKwargs, CustomDtype, DistributedType
  from .imports import is_mps_available, is_npu_available, is_xpu_available
  from .offload import load_offloaded_weight, offload_weight, save_offload_index
  from .tqdm import is_tqdm_available, tqdm
  from safetensors import safe_open
  from safetensors.torch import load_file
label_desc: Use image node_id 5 for calling a global function with example usage: shard_checkpoint(state_dict, max_shard_size, weights_name) and returns: shards,  index, dict,  None