def load(
    sharded_state_dict: ShardedStateDict,
    checkpoint_dir: str,
    sharded_strategy: Union[LoadShardedStrategy, None] = None,
    common_strategy: Union[LoadCommonStrategy, None] = None,
) -> StateDict:
    """Loading entrypoint.

    Arguments:
        sharded_state_dict: state dict of the existing model populated with
            ShardedTensors. Used as a mapping to determine which parts of
            global tensors stored in the checkpoint should be loaded.
        checkpoint_dir: directory with the checkpoint
        sharded_strategy: configures loading behavior for sharded tensors
        common_strategy: configures loading behavior for common data
    """
    if common_strategy is not None:
        raise NotImplementedError('The only supported common strategy is torch')

    checkpoint_dir = Path(checkpoint_dir)
    common_state_dict = load_common_state_dict(checkpoint_dir)
    if not sharded_state_dict:
        return common_state_dict

    sharded_objects, sharded_state_dict = load_sharded_objects(sharded_state_dict, checkpoint_dir)
    merge(common_state_dict, sharded_objects)

    saved_config = maybe_load_config(checkpoint_dir)
    if saved_config is None:
        raise CheckpointingException(f'{checkpoint_dir} is not a distributed checkpoint')

    sh_ten_factories, _ = extract_matching_values(
        sharded_state_dict, lambda x: isinstance(x, ShardedTensorFactory)
    )
    apply_factories(sharded_state_dict)
    sharded_state_dict, _ = extract_sharded_tensors_or_nonpersistent(sharded_state_dict)
    sharded_state_dict, nonpersistent_state_dict = extract_sharded_tensors(sharded_state_dict)
    dict_list_map_inplace(lambda o: o.unwrap(), nonpersistent_state_dict)
    merge(common_state_dict, nonpersistent_state_dict)

    validate_sharding_integrity(nested_values(sharded_state_dict))

    if sharded_strategy is None:
        sharded_strategy = get_default_strategy(
            StrategyAction.LOAD_SHARDED,
            saved_config.sharded_backend,
            saved_config.sharded_backend_version,
        )
    else:
        # TODO: implement consistency checks here
        pass
    loaded_state_dict = sharded_strategy.load(sharded_state_dict, checkpoint_dir)

    loaded_state_dict = apply_factory_merges(loaded_state_dict, sh_ten_factories)

    merge(common_state_dict, loaded_state_dict)
    return common_state_dict
Loading entrypoint.

    Arguments:
        sharded_state_dict: state dict of the existing model populated with
            ShardedTensors. Used as a mapping to determine which parts of
            global tensors stored in the checkpoint should be loaded.
        checkpoint_dir: directory with the checkpoint
        sharded_strategy: configures loading behavior for sharded tensors
        common_strategy: configures loading behavior for common data
    
training/DeepSpeed-Domino/megatron/core/dist_checkpointing/serialization.py