Spaces:
Running
on
Zero
Running
on
Zero
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| from collections import defaultdict | |
| import logging | |
| import typing as tp | |
| import flashy | |
| import torch | |
| from ..optim import ModuleDictEMA | |
| from .utils import copy_state | |
| logger = logging.getLogger(__name__) | |
| class BestStateDictManager(flashy.state.StateDictSource): | |
| """BestStateDictManager maintains a copy of best state_dict() for registered sources. | |
| BestStateDictManager has two main attributes: | |
| states (dict): State dict of the registered StateDictSource. | |
| param_ids (dict): Dict of parameter ids for registered states from ModuleDictEMA and other sources. | |
| When registering new sources, the BestStateDictManager will ensure two conflicting sources between | |
| ModuleDictEMA and original modules are not both registered as it would otherwise create ambiguity about | |
| what to consider for best state. | |
| Args: | |
| device (torch.device or str): Device on which we keep the copy. | |
| dtype (torch.dtype): Data type for the state parameters. | |
| """ | |
| def __init__(self, device: tp.Union[torch.device, str] = 'cpu', | |
| dtype: tp.Optional[torch.dtype] = None): | |
| self.device = device | |
| self.states: dict = {} | |
| self.param_ids: dict = defaultdict(dict) | |
| self.dtype = dtype | |
| def _get_parameter_ids(self, state_dict): | |
| return {id(p): name for name, p in state_dict.items() if isinstance(p, torch.Tensor)} | |
| def _validate_no_parameter_ids_overlap(self, name: str, param_ids: dict): | |
| for registered_name, registered_param_ids in self.param_ids.items(): | |
| if registered_name != name: | |
| overlap = set.intersection(registered_param_ids.keys(), param_ids.keys()) | |
| assert len(overlap) == 0, f"Found {len(overlap)} / {len(param_ids.keys())} overlapping parameters" | |
| f" in {name} and already registered {registered_name}: {' '.join(overlap)}" | |
| def update(self, name: str, source: flashy.state.StateDictSource): | |
| if name not in self.states: | |
| raise ValueError(f"{name} missing from registered states.") | |
| self.states[name] = copy_state(source.state_dict(), device=self.device, dtype=self.dtype) | |
| def register(self, name: str, source: flashy.state.StateDictSource): | |
| if name in self.states: | |
| raise ValueError(f"{name} already present in states.") | |
| # Registering parameter ids for EMA and non-EMA states allows us to check that | |
| # there is no overlap that would create ambiguity about how to handle the best state | |
| param_ids = self._get_parameter_ids(source.state_dict()) | |
| if isinstance(source, ModuleDictEMA): | |
| logger.debug(f"Registering to best state: ModuleDictEMA '{name}' with {len(param_ids)} params") | |
| self._validate_no_parameter_ids_overlap(name, param_ids) | |
| self.param_ids[name] = param_ids | |
| else: | |
| logger.debug(f"Registering to best state: StateDictSource '{name}' with {len(param_ids)} params") | |
| self._validate_no_parameter_ids_overlap('base', param_ids) | |
| self.param_ids['base'].update(param_ids) | |
| # Register state | |
| self.states[name] = copy_state(source.state_dict(), device=self.device, dtype=self.dtype) | |
| def state_dict(self) -> flashy.state.StateDict: | |
| return self.states | |
| def load_state_dict(self, state: flashy.state.StateDict): | |
| for name, sub_state in state.items(): | |
| for k, v in sub_state.items(): | |
| self.states[name][k].copy_(v) | |