File size: 2,342 Bytes
ad5f26a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
# Copyright (c) Meta Platforms, Inc. and affiliates
from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE
from . import _version
from ._traverse import (
OBJ_PATH,
set_element,
STATE_DICT_ITEM,
traverse_state_dict,
traverse_state_dict_v_2_3,
)
"""
TODO:
Need to add ability to handle tuple, OrderedDict, NamedTuple.
Update mappings from dict to a class.
Change set_element to recreate the right type for tuple, OrderedDict, and NamedTuple.
"""
FLATTEN_MAPPING = dict[str, OBJ_PATH]
# TODO: Update Docstring for nested_dict.py
def flatten_state_dict(
state_dict: STATE_DICT_TYPE,
) -> tuple[STATE_DICT_TYPE, FLATTEN_MAPPING]:
"""
Flatten ``state_dict`` made of nested dicts and lists into a top level dictionary.
Use ``unflatten_state_dict`` to revert this process.
Returns:
A tuple with the flatten state_dict and a mapping from original to new state_dict.
N.B. The new keys are derived from the object paths, joined by dot.
For example: ``{ 'a': {'b':...}}`` results in the key `a.b`.
"""
flattened: STATE_DICT_TYPE = {}
mappings: FLATTEN_MAPPING = {}
def flat_copy(path: OBJ_PATH, value: STATE_DICT_ITEM) -> None:
new_fqn = ".".join(map(str, path))
if new_fqn in flattened:
raise ValueError(f"duplicated flatten key {new_fqn}")
flattened[new_fqn] = value
mappings[new_fqn] = path
# We started to flatten dictionary since v2.4. But in order to not break
# the checkpoints that were saved before v2.4, we need to keep the old
# traversal so that we can reconstruct those checkpoints.
use_v_2_3 = (
_version._derived_version is not None and _version._derived_version == "2_3"
)
if use_v_2_3:
traverse_state_dict_v_2_3(state_dict, flat_copy)
else:
traverse_state_dict(state_dict, flat_copy)
return flattened, mappings
def unflatten_state_dict(
state_dict: STATE_DICT_TYPE, mapping: FLATTEN_MAPPING
) -> STATE_DICT_TYPE:
"""Restore the original nested state_dict according to ``mapping`` and the flattened ``state_dict``."""
nested: STATE_DICT_TYPE = {}
for key, value in state_dict.items():
set_element(nested, mapping[key], value)
return nested
|