a h6@sUddlZddlZddlmZmZddlmZddlmZm Z m Z m Z m Z ddl Z ddlmZddl mZddlmZddlmZgZeeed<eeZGd d d ejZeeed d d ZdS)N) CollectionMapping)deepcopy)AnyCallableOptionaloverloadUnion)optim) ShardedTensor)FullyShardedDataParallel__all__c @sxeZdZdZdeeeeje ffe j e e eeefe ejeedfeeefddddZdddd Zeeefdd d Zed ddd d dZeegefed ddZd!e egefe ed ddZeeejefdddZeeefddddZeeefddddZddddZeeefeeefdddZeeefeeefdddZ dS)"_NamedOptimizera ``_NamedOptimizer`` takes a dict of parameters and exposes ``state_dict`` by parameter key. We replace the original key (number) in an optim to the fully qualified name (FQN) string. User can initialize the optim as they initialize a PyTorch optim, the only difference is that they also need to pass in the FQN of each parameters. Args: named_parameters (Mapping[str, Union[torch.Tensor, ShardedTensor]]): Mapping from FQN to parameter. optimizer_class (optim.Optimizer): The class of optimizer to instantiate. param_groups (Collection[Mapping[str, Any]]): `param_groups` to pass to optimizer if specified. The key of the inner map needs to be FQNs. Default: None module (nn.Module): the module whose parameters to updated by the optimizer. args: arguments to pass to the optimizer constructor. kwargs: arguments to pass to the optimizer constructor. Example:: >>> # xdoctest: +SKIP("distributed") >>> from torch import optim >>> from torch.distributed.optim import _NamedOptimizer >>> >>> # Define the named optimizer. >>> m = Model(...) >>> named_optim = _NamedOptimizer(m.named_parameters(), optim.SGD) >>> # Forward pass + backward pass. >>> named_optim.step() >>> ... >>> # Call state_dict for the named optimizer returns a FQN state_dict. >>> named_optim.state_dict() Warning: This API is still in development and subject to change. TODO: Add tutorial for _NamedOptimizer. TODO: Add documentation in the docstring for the public attributes like self.param_groups and self.named_parameters. N.)named_parametersoptimizer_class param_groupsmoduleargskwargsreturnc Ostjd||_|t||_|dur6|jn|}||g|Ri||_||_ |durrt |j |_ nft ddd|jD}g} |D]8} | dD]*} | |vrtd| d| || qq| |_ |jj|_dS)Nz'torch.distributed.optim._NamedOptimizerzvSince we pass in param_groups, we will use param_groups to initialize the optimizer, not all parameters of the module.cSsi|]\}}||qSr.0keyparamrrU/var/www/auris/lib/python3.9/site-packages/torch/distributed/optim/named_optimizer.py \z,_NamedOptimizer.__init__..paramszExpect param name z% found in param group but is missing.)torchZ_CZ_log_api_usage_oncer_param_groups_checkdictrvalues _optimizerrlistkeysordered_param_keyswarningswarnitems ValueErrorappend) selfrrrrrrZparams_for_optimizer param_to_keyr&grouprrrr__init__?s<    z_NamedOptimizer.__init__)rcCs|jdur|jD]x}t|ts&Jdd|vs6Jd|d}t|tjrP|g}t|}|D]"}t|tjs\tdt|q\||d<qdS)Nparam group must be a dictrz#param group must contain key paramsz>optimizer can only optimize Tensors, but one of the params is )r isinstancer!rTensorr$ TypeErrortypename)r, param_grouprrrrrr is     z#_NamedOptimizer._param_groups_checkc sj}|d}fdd|dD}g}|D]V}fdd|dD}dt|i}|D]\}} |dkrbt| ||<qb||q4||dS) z Return the ``state_dict`` of the optimizer. Instead of using number to index parameters, we will use module fully qualified name (FQN) as the key. rcsi|]\}}j||qSrr&)rZst_key state_valr,rrrsz._NamedOptimizer.state_dict..statecsg|]}j|qSrr6)rrr8rr rz._NamedOptimizer.state_dict..r)r9r)r# state_dictr)sortedrr+_post_state_dict) r,r;rZ ret_stateZ ret_groupsr. param_keysZ ret_groupkvrr8rr;zs     z_NamedOptimizer.state_dict)closurercCsdSNrr,rArrrstepsz_NamedOptimizer.stepcCsdSrBrrCrrrrDscCs|jj|dS)z Perform a single optimization step. This will call :meth:`torch.optim.Optimizer.step` on the wrapped optimizer. rA)r#rDrCrrrrDscCs|jjSrB)r#r9r8rrrr9sz_NamedOptimizer.state)r;rc Cs|j}||}|d}|d}t|dkr8tdt|jD]l\}}||vrZqBt||t||krtdt||d|dt||||D]\}}|||vrtd|d|d|||} t |t rnt | t sJt| } t| } | | krZ new_group_mapZ new_groupZ group_keyZ src_groupr?rrrrNs    $            z_NamedOptimizer.load_state_dict)r5rcCst|tsJd|d}t|tjr2|g|d<n t||d<dd|jD}|dD]$}||vrntd|j ||qZ|j ||j j |_ dS)z Add a param group to the :class:`_NamedOptimizer` s `param_groups`. Warning: This API is still in development and subject to change. r0rcSsi|]\}}||qSrrrrrrrrz3_NamedOptimizer.add_param_group..z%some parameters are not in the moduleN) r1r!rr2r$rr)r*r&r+r#add_param_groupr)r,r5rr-rrrrrPs     z_NamedOptimizer.add_param_groupcCs>|jD]"}|jr t|}tj||_q |jdddS)z Run a dummy optimizer step, which allows to initialize optimizer state because we do lazy init for most optimizers. This allows doing in-place loading of optimizer state from a checkpoint. NrE) rr"Z requires_gradrZ zeros_likeZautogradVariableZgradrD)r,rtrrr init_state's  z_NamedOptimizer.init_statecCs&t|jtr"tj|j|j|ddS|S)NT)Zis_named_optimizer)r1rFSDPZoptim_state_dict_to_loadr#r,r;rrrrH4s   z$_NamedOptimizer._pre_load_state_dictcCs"t|jtrt|j|j||SrB)r1rrTZoptim_state_dictr#rUrrrr==s z _NamedOptimizer._post_state_dict)NN)N)N)!__name__ __module__ __qualname____doc__rstrr rr2r r OptimizerrrrnnModuletupler!r/r r;rrDrfloatpropertyr9rNrPrSrHr=rrrrrs4/   *" h r)r>rcCsdt|S)zFConcatenate all param keys as a unique identifier for one param group.rF)joinr<)r>rrrrMEsrM) loggingr'collections.abcrrcopyrtypingrrrrr rZtorch.nnr\r Z'torch.distributed._shard.sharded_tensorr Ztorch.distributed.fsdpr rTr r$rZ__annotations__ getLoggerrVloggerr[rrMrrrrs       4