class_docstr: Loads one checkpoint replica from storage and broadcasts to other nodes.
  
  This strategy loads checkpoint from storage on minimal set of nodes
  and distributes the checkpoint to other nodes with torch.distributed.
  Loading is performed with tensorstore.
  
  Steps:
  0. (optional) create Gloo distributed groups
  1. Exchange ShardedTensors metadata between all nodes
  2. Align needed tensors within DP groups
  3. For each globally unique tensor:
  3.a) on one of the ranks load it from storage to CPU and move to CUDA
  3.b) allocate CUDA tensor on other ranks
  3.c) broadcast within DP group
  3.d) copy tensor content to the model param location
  3.e) free tensor buffers from a) and b)
  
  Notes:
  1. Loading and broadcasting is done sequentially to avoid both host and device OOMs
  2. There is a lot of overlap potential between all three steps done for each tensor:
  2.a) loading from storage to numpy
  2.b) moving CPU tensors to CUDA
  2.c) broadcast
class_docstr_tok: ["Loads","one","checkpoint","replica","from","storage","and","broadcasts","to","other","nodes",".","This","strategy","loads","checkpoint","from","storage","on","minimal","set","of","nodes","and","distributes","the","checkpoint","to","other","nodes","with","torch.distributed",".","Loading","is","performed","with","tensorstore",".","Steps",":","0",".","(","optional",")","create","Gloo","distributed","groups","1",".","Exchange","ShardedTensors","metadata","between","all","nodes","2",".","Align","needed","tensors","within","DP","groups","3",".","For","each","globally","unique","tensor",":","3.a",")","on","one","of","the","ranks","load","it","from","storage","to","CPU","and","move","to","CUDA","3.b",")","allocate","CUDA","tensor","on","other","ranks","3.c",")","broadcast","within","DP","group","3.d",")","copy","tensor","content","to","the","model","param","location","3.e",")","free","tensor","buffers","from","a",")","and","b",")","Notes",":","1",".","Loading","and","broadcasting","is","done","sequentially","to","avoid","both","host","and","device","OOMs","2",".","There","is","a","lot","of","overlap","potential","between","all","three","steps","done","for","each","tensor",":","2.a",")","loading","from","storage","to","numpy","2.b",")","moving","CPU","tensors","to","CUDA","2.c",")","broadcast"]
code: def load(
      self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Path
  ):
      self.maybe_init_gloo_group()
      all_tensors_sorted = self._build_load_plan(sharded_state_dict)
      self._exchange_loaded_tensors(
          all_tensors_sorted, sharded_state_dict, checkpoint_dir
      )
      self.summarize_load_times()
      return sharded_state_dict
  
imports: import os
  import time
  from collections import defaultdict
  from dataclasses import dataclass
  from functools import partial, wraps
  from itertools import chain
  from logging import DEBUG, INFO, StreamHandler, getLogger
  from operator import attrgetter, itemgetter
  from pathlib import Path
  from typing import Iterable, List, NamedTuple, Optional, Tuple, Union
  import torch
  from ..dict_utils import dict_list_map_inplace, map_reduce, nested_values
  from ..mapping import ShardedStateDict, ShardedTensor, StateDict
  from .base import LoadShardedStrategy
  from .tensorstore import TensorStoreLoadShardedStrategy, _load_from_array, open_ts_array
  from .zarr import flatten_range, load_zarr_based_sharded_metadata
label_desc: Use image node_id 2 for calling the TwoStageDataParallelLoadShardedStrategy obj's underlying member method code with example usage: obj.load(sharded_state_dict, checkpoint_dir) and returns: sharded_state_dict