def _process_vocab_cache(self, slice_mode): """PS embeddingLookup cache check and process.""" self.cache_enable = False if self.vocab_cache_size > 0: if self.target == 'CPU': logger.warning("The configuration of 'vocab_cache_size' is valid only in 'DEVICE' target, " "current target is CPU, so it will be ignored.") return enable_ps = _get_ps_context("enable_ps") if not enable_ps: logger.warning("The configuration of 'vocab_cache_size' is valid only in parameter server trainning " "mode, current mode is not parameter server trainning mode, so it will be ignored.") return parallel_mode = _get_parallel_mode() is_auto_parallel = parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL) if is_auto_parallel: device_num = get_group_size() full_batch = _get_full_batch() if device_num > 1 and not (full_batch and slice_mode == "table_row_slice"): raise ValueError("The embeddingLookup cache of parameter server parallel only be used " "in 'full_batch' and 'table_row_slice' parallel strategy.") self.vocab_cache_size = self.vocab_cache_size * device_num self.cache_enable = True if _is_role_worker(): self.vocab_size = self.vocab_cache_size
def _process_vocab_cache(self, slice_mode): """PS embeddingLookup cache check and process.""" self.cache_enable = False if self.vocab_cache_size > 0: if self.target == 'CPU': logger.warning("The configuration of 'vocab_cache_size' is valid only in 'DEVICE' target, " "current target is CPU, so it will be ignored.") return enable_ps = _get_ps_context("enable_ps") if not enable_ps: logger.warning("The configuration of 'vocab_cache_size' is valid only in parameter server trainning " "mode, current mode is not parameter server trainning mode, so it will be ignored.") return parallel_mode = _get_parallel_mode() is_auto_parallel = parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL) if is_auto_parallel: rank_size = get_group_size() rank_id = get_rank() full_batch = _get_full_batch() if rank_size > 1 and not (full_batch and slice_mode == "table_row_slice"): raise ValueError("The embeddingLookup cache of parameter server parallel only be used " "in 'full_batch' and 'table_row_slice' parallel strategy.") self.vocab_cache_size = self.vocab_cache_size * rank_size _set_rank_id(rank_id) self.cache_enable = True if _is_role_worker(): self.vocab_size = self.vocab_cache_size if context.get_context("enable_sparse") != self.sparse: raise ValueError("The value of parameter 'sparse' must be same for all EmbeddingLookup " "kernels and equal the value of 'enable_sparse' in context setting in " "parameter server cache mode")
def get_ps_context(attr_key): """ Get parameter server training mode context attribute value according to the key. Args: attr_key (str): The key of the attribute. Returns: Returns attribute value according to the key. Raises: ValueError: If input key is not attribute in auto parallel context. """ return _get_ps_context(attr_key)
def get_fl_context(attr_key): """ Get federated learning mode context attribute value according to the key. Args: attr_key (str): The key of the attribute. Returns: Returns attribute value according to the key. Raises: ValueError: If input key is not attribute in federated learning mode context. """ return _get_ps_context(attr_key)
def get_fl_context(attr_key): """ Get federated learning mode context attribute value according to the key. Args: attr_key (str): The key of the attribute. Please refer to `set_fl_context`'s parameters to decide which key should passed. Returns: Returns attribute value according to the key. Raises: ValueError: If input key is not attribute in federated learning mode context. Examples: >>> context.get_fl_context("server_mode") """ return _get_ps_context(attr_key)
def __init__(self, vocab_size, embedding_size, param_init='normal', target='CPU', slice_mode='batch_slice', manual_shapes=None, max_norm=None, sparse=True, vocab_cache_size=0): super(EmbeddingLookup, self).__init__() validator.check_value_type('sparse', sparse, [bool], self.cls_name) self.vocab_size = validator.check_positive_int(vocab_size, 'vocab_size') self.vocab_cache_size = validator.check_non_negative_int( vocab_cache_size, 'vocab_cache_size') self.target = target self.sparse = sparse self.cache_enable = self.vocab_cache_size > 0 self.forward_unique = False if target not in ('CPU', 'DEVICE'): raise ValueError( 'Attr \'target\' of \'EmbeddingLookup\' Op passed ' + str(target) + ', should be one of values in \'CPU\', \'DEVICE\'.') if not sparse and target == 'CPU': raise ValueError( 'When target is CPU, embedding_lookup must be sparse.') if sparse: self.gatherv2 = P.SparseGatherV2() else: self.gatherv2 = P.Gather() self.embeddinglookup = P.EmbeddingLookup().add_prim_attr( 'primitive_target', 'CPU') enable_ps = _get_ps_context("enable_ps") if enable_ps: self._process_vocab_cache(slice_mode) self.embedding_size = validator.check_positive_int( embedding_size, 'embedding_size') self.embedding_table = Parameter(initializer( param_init, [self.vocab_size, self.embedding_size]), name='embedding_table') parallel_mode = _get_parallel_mode() is_auto_parallel = parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL) self.gather_revert = P.Gather() self.reshape_first = P.Reshape() self.reshape = P.Reshape() self.unique = P.Unique() self.shape = P.Shape() if is_auto_parallel: self.unique = P.Unique().shard(((1, ), )) if self.cache_enable and enable_ps: self._set_voacb_cache_enable_for_ps(vocab_cache_size, embedding_size, vocab_size) if is_auto_parallel: self.unique.add_prim_attr('cache_enable', True) indices_shape_size = 2 if slice_mode == "field_slice" and is_auto_parallel: if not manual_shapes: raise ValueError( "in slice field mode, the manual_shapes should not be none" ) if not isinstance(manual_shapes, tuple): raise TypeError( "manual_shapes type must be tuple(int) cannot be {}!". format(type(manual_shapes))) for dim in manual_shapes: validator.check_positive_int(dim, 'manual shape dim', self.cls_name) self.gatherv2.add_prim_attr("manual_split", manual_shapes) self.embeddinglookup.add_prim_attr("manual_split", manual_shapes) self.gatherv2.shard(((get_group_size(), 1), (1, get_group_size()))) self.embeddinglookup.shard( ((get_group_size(), 1), (1, get_group_size()))) elif slice_mode == "table_row_slice" and is_auto_parallel: full_batch = _get_full_batch() if (target == 'DEVICE' and not full_batch) or (self.cache_enable and enable_ps and sparse): indices_shape_size = 1 self.gather_revert.shard(((1, 1), (get_group_size(), ))) self.forward_unique = True indices_strategy = (1, ) * indices_shape_size self.gatherv2.shard(((get_group_size(), 1), indices_strategy)) self.embeddinglookup.shard( ((get_group_size(), 1), indices_strategy)) elif slice_mode == "table_column_slice" and is_auto_parallel: if target == 'DEVICE': indices_shape_size = 1 self.gather_revert.shard(((1, get_group_size()), (1, ))) self.forward_unique = True indices_strategy = (1, ) * indices_shape_size self.gatherv2.shard(((1, get_group_size()), indices_strategy)) self.embeddinglookup.shard( ((1, get_group_size()), indices_strategy)) elif slice_mode == "batch_slice" and is_auto_parallel: indices_strategy = [get_group_size()] indices_strategy.extend([1] * (indices_shape_size - 1)) indices_strategy = tuple(indices_strategy) self.gatherv2.shard(((1, 1), indices_strategy)) self.embeddinglookup.shard(((1, 1), indices_strategy)) else: if is_auto_parallel: raise ValueError( "slice_mode should support mode in nn.EmbeddingLookup, but get " + str(slice_mode)) if self.cache_enable and not enable_ps: if parallel_mode != ParallelMode.STAND_ALONE: raise ValueError( "parallel mode haven't supported cache enable yet.") self._set_cache_enable() self.embedding_table.unique = self.forward_unique self.max_norm = max_norm if self.max_norm is not None: self.max_norm = validator.check_positive_float( self.max_norm, 'max_norm', self.cls_name) self.max_norm = Tensor(self.max_norm, dtype=mstype.float32)