parameters: (self)
docstring: Integrate the parameters into a continuous memory according to rank, and support the update of training parameters.
docstring_summary: Integrate the parameters into a continuous memory according to rank, and support the update of training parameters.
docstring_tokens: Integrate the parameters into a continuous memory according to rank and support the update of training parameters .
function: def _integration_params(self):
          """
          Integrate the parameters into a continuous memory according to rank, and support the update of training parameters.
          """
  
          for dtype, per_rank_params in self.dtype_rank_params.items():
              if dtype not in self.param_storages.keys():
                  self.param_storages[dtype] = {}
  
              for dst_rank, params in enumerate(per_rank_params):
                  if len(params) > 0:
  
                      # Merge all the trainable params in a single InternalStorage
                      trainable_params = list(
                          filter(lambda x: x.trainable, params))
                      if self._pfp16 and dst_rank == self.rank:
                          self._generate_master_params(trainable_params)
                      if trainable_params:
                          param_storage = ParamStorage(
                              size=self.rank_buffer_size[dtype][dst_rank],
                              dtype=dtype,
                              device=self._default_device)
  
                          param_storage.add_rank_params(trainable_params,
                                                        self._param2align)
                          self.param_storages[dtype][dst_rank] = param_storage
  
          # Clear the InternalStorage keys which are not in use anymore
          dtype_in_use = list(self.dtype_rank_params.keys())
          dtype_to_pop = list(
              filter(lambda x: x not in dtype_in_use, self.param_storages.keys()))
          for d in dtype_to_pop:
              self.param_storages.pop(d)
  
          if self.offload:
              self._optim._master_weights = self._master_params
              cpu_master_params = [p for p in self._master_params.values()]
              for param in cpu_master_params:
                  size = np.prod(param.shape) * align[Type.fp32.value]
                  remaining = size % alignment[self.offload_device]
                  ali = 0 if remaining == 0 else alignment[
                      self.offload_device] - remaining
                  align_ = ali // align[Type.fp32.value]
                  self.offload_buffer_size += np.prod(param.shape) + align_
                  self.offload_param2align[param.name] = align_
  
              if cpu_master_params:
                  with device_guard(self.rank, self.offload_device):
                      self.offload_params = ParamStorage(
                          size=self.offload_buffer_size,
                          dtype=Type.fp32.value,
                          device=self.offload_device)
                      self.offload_params.add_rank_params(
                          cpu_master_params, self.offload_param2align, False)
                      self.offload_params.buffer.stop_gradient = False
  
                      self.offload_grads = GradStorage(
                          size=self.offload_buffer_size,
                          dtype=Type.fp32.value,
                          device=self.offload_device,
                          destination=self.rank,
                          parm2align=self.offload_param2align,
                          convert_cpu=True)
                      for p in cpu_master_params:
                          self.offload_grads.add_grad(
                              p, self.offload_param2align[p.name])
  
                      self._optim._master_weights[
                          self.offload_params.buffer.
                          name] = self.offload_params.buffer
function_tokens: def _integration_params ( self ) : for dtype , per_rank_params in self . dtype_rank_params . items ( ) : if dtype not in self . param_storages . keys ( ) : self . param_storages [ dtype ] = { } for dst_rank , params in enumerate ( per_rank_params ) : if len ( params ) > 0 : # Merge all the trainable params in a single InternalStorage trainable_params = list ( filter ( lambda x : x . trainable , params ) ) if self . _pfp16 and dst_rank == self . rank : self . _generate_master_params ( trainable_params ) if trainable_params : param_storage = ParamStorage ( size = self . rank_buffer_size [ dtype ] [ dst_rank ] , dtype = dtype , device = self . _default_device ) param_storage . add_rank_params ( trainable_params , self . _param2align ) self . param_storages [ dtype ] [ dst_rank ] = param_storage # Clear the InternalStorage keys which are not in use anymore dtype_in_use = list ( self . dtype_rank_params . keys ( ) ) dtype_to_pop = list ( filter ( lambda x : x not in dtype_in_use , self . param_storages . keys ( ) ) ) for d in dtype_to_pop : self . param_storages . pop ( d ) if self . offload : self . _optim . _master_weights = self . _master_params cpu_master_params = [ p for p in self . _master_params . values ( ) ] for param in cpu_master_params : size = np . prod ( param . shape ) * align [ Type . fp32 . value ] remaining = size % alignment [ self . offload_device ] ali = 0 if remaining == 0 else alignment [ self . offload_device ] - remaining align_ = ali // align [ Type . fp32 . value ] self . offload_buffer_size += np . prod ( param . shape ) + align_ self . offload_param2align [ param . name ] = align_ if cpu_master_params : with device_guard ( self . rank , self . offload_device ) : self . offload_params = ParamStorage ( size = self . offload_buffer_size , dtype = Type . fp32 . value , device = self . offload_device ) self . offload_params . add_rank_params ( cpu_master_params , self . offload_param2align , False ) self . offload_params . buffer . stop_gradient = False self . offload_grads = GradStorage ( size = self . offload_buffer_size , dtype = Type . fp32 . value , device = self . offload_device , destination = self . rank , parm2align = self . offload_param2align , convert_cpu = True ) for p in cpu_master_params : self . offload_grads . add_grad ( p , self . offload_param2align [ p . name ] ) self . _optim . _master_weights [ self . offload_params . buffer . name ] = self . offload_params . buffer