parameters: (self, loss)
docstring: something should be done before step, such as calc gradients, mpi sync, and so on.
  
          Example::
  
              class MyOptimizer(Optimizer):
                  def step(self, loss):
                      self.post_step(loss)
                      ...
docstring_summary: something should be done before step, such as calc gradients, mpi sync, and so on.
docstring_tokens: something should be done before step such as calc gradients mpi sync and so on .
function: def pre_step(self, loss):
          """ something should be done before step, such as calc gradients, mpi sync, and so on.
  
          Example::
  
              class MyOptimizer(Optimizer):
                  def step(self, loss):
                      self.post_step(loss)
                      ...
          """
          # clean prev grads
          params = []
          params_has_grad = []
          for pg in self.param_groups:
              for p in pg['params']:
                  params.append(p)
                  if not p.is_stop_grad():
                      params_has_grad.append(p)
  
          # get gradient
          grads = jt.grad(loss, params_has_grad)
  
          # sync grads and model if in mpi
          if jt.in_mpi:
              dep = []
              def add_dep(v):
                  nonlocal dep
                  v._add_dependency(dep)
                  dep = [v]
  
              for g in grads:
                  g.assign(g.mpi_all_reduce("mean"))
                  add_dep(g._input(0))
              if self.n_step % self.param_sync_iter == 0:
                  for p in params:
                      p.assign(p.mpi_broadcast())
                      add_dep(p)
          self.n_step += 1
  
          # set up grads in param_groups
          pid = 0
          for pg in self.param_groups:
              if "grads" not in pg:
                  pg["grads"] = [ jt.zeros_like(p).stop_grad().stop_fuse() for p in pg['params'] ]
              pg_grads = pg["grads"]
              for i, p in enumerate(pg['params']):
                  if not p.is_stop_grad():
                      # accumulate grad and stop grad of grad
                      g = grads[pid].stop_grad()
                      if not self.__zero_grad:
                          g = g + pg_grads[i]
                      pg_grads[i].update(g)
                      pid += 1
          self.__zero_grad = False
function_tokens: def pre_step ( self , loss ) : # clean prev grads params = [ ] params_has_grad = [ ] for pg in self . param_groups : for p in pg [ 'params' ] : params . append ( p ) if not p . is_stop_grad ( ) : params_has_grad . append ( p ) # get gradient grads = jt . grad ( loss , params_has_grad ) # sync grads and model if in mpi if jt . in_mpi : dep = [ ] def add_dep ( v ) : nonlocal dep v . _add_dependency ( dep ) dep = [ v ] for g in grads : g . assign ( g . mpi_all_reduce ( "mean" ) ) add_dep ( g . _input ( 0 ) ) if self . n_step % self . param_sync_iter == 0 : for p in params : p . assign ( p . mpi_broadcast ( ) ) add_dep ( p ) self . n_step += 1 # set up grads in param_groups pid = 0 for pg in self . param_groups : if "grads" not in pg : pg [ "grads" ] = [ jt . zeros_like ( p ) . stop_grad ( ) . stop_fuse ( ) for p in pg [ 'params' ] ] pg_grads = pg [ "grads" ] for i , p in enumerate ( pg [ 'params' ] ) : if not p . is_stop_grad ( ) : # accumulate grad and stop grad of grad g = grads [ pid ] . stop_grad ( ) if not self . __zero_grad : g = g + pg_grads [ i ] pg_grads [ i ] . update ( g ) pid += 1 self . __zero_grad = False