parameters: (self, block, startup_block, offload=True)
docstring: (p_fp16) = cast(p)
          (p_fp16_recompute) = cast(p)
          (pout,) = adam(p)
          ===========================>
          rename(p_fp16_recompute, p_fp16)
  
          (p,) = prefetch(p@offload)
          (pout,) = adam(p)
          (p_fp16) = cast(p)
          (p@offload) = memcpy(p)
docstring_summary: (p_fp16) = cast(p)
          (p_fp16_recompute) = cast(p)
          (pout,) = adam(p)
          ===========================>
          rename(p_fp16_recompute, p_fp16)
docstring_tokens: ( p_fp16 ) = cast ( p ) ( p_fp16_recompute ) = cast ( p ) ( pout ) = adam ( p ) =========================== > rename ( p_fp16_recompute p_fp16 )
function: def offload_fp32param(self, block, startup_block, offload=True):
          """
          (p_fp16) = cast(p)
          (p_fp16_recompute) = cast(p)
          (pout,) = adam(p)
          ===========================>
          rename(p_fp16_recompute, p_fp16)
  
          (p,) = prefetch(p@offload)
          (pout,) = adam(p)
          (p_fp16) = cast(p)
          (p@offload) = memcpy(p)
          """
          param_to_idx = dict()
          param_to_fp16 = dict()
          # recompute_var which need rename to fp16_param
          fp16_param_to_recompute = dict()
          recompute_to_fp16 = dict()
  
          def remove_param(input_name):
              param_to_idx.pop(input_name)
              if input_name in param_to_fp16:
                  fp16_param = param_to_fp16.pop(input_name)
                  if fp16_param in fp16_param_to_recompute:
                      recompute = fp16_param_to_recompute.pop(fp16_param)
                      recompute_to_fp16.pop(recompute)
  
          # step1: record param
          for idx, op in reversed(list(enumerate(block.ops))):
              if is_update_op(op):
                  param = op.desc.input("Param")[0]
                  param_to_idx[param] = idx
  
          # step2: remove param which can't offload and
          #        record param->fp16param, fp16param->recompute_var
          for idx, op in enumerate(block.ops):
              if is_optimizer_op(op):
                  break
              # TODO (Yuang Liu): tmp solution for fuse_grad_merge + optimize_cast
              if not offload and op.type == 'coalesce_tensor':
                  continue
              for input_name in op.desc.input_arg_names():
                  if input_name not in param_to_idx:
                      continue
  
                  # param which will be used by fp32 op
                  if op.type != 'cast':
                      remove_param(input_name)
                      continue
  
                  # param is only used by cast op,
                  # which to cast fp32_param to fp16_param
                  output_name = op.output_arg_names[0]
                  if 'cast_fp16' not in output_name:
                      remove_param(input_name)
                      continue
  
                  if 'subprog' not in output_name:
                      assert output_name == input_name + '.cast_fp16'
                      assert input_name not in param_to_fp16, \
                          "There must be only one cast op from fp32 param to fp16 param."
                      param_to_fp16[input_name] = output_name
                  else:
                      # fp16-->recompute_var
                      assert input_name in param_to_fp16, \
                          "param must first be cast to fp16"
                      fp16_param = param_to_fp16[input_name]
                      fp16_param_to_recompute[fp16_param] = output_name
                      recompute_to_fp16[output_name] = fp16_param
  
          param_name_to_offload_name = dict()
          # step3: main_block add offload, cast op
          # change recompute to fp16, remove cast(param) to fp16
          for idx, op in reversed(list(enumerate(block.ops))):
              if is_update_op(op):
                  param = op.desc.input("Param")[0]
                  if param not in param_to_idx: continue
                  # step3.1: create offload_var
                  offload_var_name = self._get_offload_var_name(param)
                  param_name_to_offload_name[param] = offload_var_name
                  if offload:
                      self._create_offload_var(param, offload_var_name,
                                               [block, startup_block])
  
                      # step3.2: insert cast op and offload op
                      self._insert_offload_op(block, idx + 1, param,
                                              offload_var_name)
  
                  assert param in param_to_fp16
                  fp16_param_name = param_to_fp16[param]
                  fp16_param_var = block.var(fp16_param_name)
                  fp16_param_var.persistable = True
                  self._insert_cast_op(block, idx + 1, param,
                                       param_to_fp16[param])
  
                  if offload:
                      # step3.3: insert fetch op
                      self._insert_fetch_op(block, idx, offload_var_name, param)
                  continue
  
              # step3.4: remove cast op
              if op.type == 'cast':
                  input_name = op.desc.input_arg_names()[0]
                  if input_name in param_to_idx:
                      block._remove_op(idx, sync=False)
                      continue
  
              # step3.5: change recompute_param to fp16_param
              for input_name in op.desc.input_arg_names():
                  if input_name in recompute_to_fp16:
                      op._rename_input(input_name, recompute_to_fp16[input_name])
              for output_name in op.desc.output_arg_names():
                  if output_name in recompute_to_fp16:
                      op._rename_output(output_name,
                                        recompute_to_fp16[output_name])
  
          # step4: remove recompute_param
          for name in recompute_to_fp16.keys():
              block._remove_var(name, sync=False)
  
          # step5: startup_block add offload
          visited_vars = set()
          # FIXME(wangxi): should insert in idx, need move comm init to the head.
          insert_idx = len(startup_block.ops)
          for idx, op in reversed(list(enumerate(startup_block.ops))):
              for out_name in op.output_arg_names:
                  if out_name in visited_vars:
                      continue
  
                  if out_name in param_name_to_offload_name:
                      var_name = out_name
                      if offload:
                          offload_var_name = param_name_to_offload_name[var_name]
                          self._insert_offload_op(startup_block, insert_idx,
                                                  var_name, offload_var_name)
                      self._insert_cast_op(startup_block, insert_idx, var_name,
                                           param_to_fp16[var_name])
                      # NOTE(wangxi): cast and offload should insert after broadcast param.
                      # the insert op order is: {mp, dp}broadcast, cast, offload
                      self._insert_broadcast_op(startup_block, insert_idx,
                                                var_name)
  
                  visited_vars.add(out_name)
  
          block._sync_with_cpp()
          startup_block._sync_with_cpp()
function_tokens: def offload_fp32param ( self , block , startup_block , offload = True ) : param_to_idx = dict ( ) param_to_fp16 = dict ( ) # recompute_var which need rename to fp16_param fp16_param_to_recompute = dict ( ) recompute_to_fp16 = dict ( ) def remove_param ( input_name ) : param_to_idx . pop ( input_name ) if input_name in param_to_fp16 : fp16_param = param_to_fp16 . pop ( input_name ) if fp16_param in fp16_param_to_recompute : recompute = fp16_param_to_recompute . pop ( fp16_param ) recompute_to_fp16 . pop ( recompute ) # step1: record param for idx , op in reversed ( list ( enumerate ( block . ops ) ) ) : if is_update_op ( op ) : param = op . desc . input ( "Param" ) [ 0 ] param_to_idx [ param ] = idx # step2: remove param which can't offload and #        record param->fp16param, fp16param->recompute_var for idx , op in enumerate ( block . ops ) : if is_optimizer_op ( op ) : break # TODO (Yuang Liu): tmp solution for fuse_grad_merge + optimize_cast if not offload and op . type == 'coalesce_tensor' : continue for input_name in op . desc . input_arg_names ( ) : if input_name not in param_to_idx : continue # param which will be used by fp32 op if op . type != 'cast' : remove_param ( input_name ) continue # param is only used by cast op, # which to cast fp32_param to fp16_param output_name = op . output_arg_names [ 0 ] if 'cast_fp16' not in output_name : remove_param ( input_name ) continue if 'subprog' not in output_name : assert output_name == input_name + '.cast_fp16' assert input_name not in param_to_fp16 , "There must be only one cast op from fp32 param to fp16 param." param_to_fp16 [ input_name ] = output_name else : # fp16-->recompute_var assert input_name in param_to_fp16 , "param must first be cast to fp16" fp16_param = param_to_fp16 [ input_name ] fp16_param_to_recompute [ fp16_param ] = output_name recompute_to_fp16 [ output_name ] = fp16_param param_name_to_offload_name = dict ( ) # step3: main_block add offload, cast op # change recompute to fp16, remove cast(param) to fp16 for idx , op in reversed ( list ( enumerate ( block . ops ) ) ) : if is_update_op ( op ) : param = op . desc . input ( "Param" ) [ 0 ] if param not in param_to_idx : continue # step3.1: create offload_var offload_var_name = self . _get_offload_var_name ( param ) param_name_to_offload_name [ param ] = offload_var_name if offload : self . _create_offload_var ( param , offload_var_name , [ block , startup_block ] ) # step3.2: insert cast op and offload op self . _insert_offload_op ( block , idx + 1 , param , offload_var_name ) assert param in param_to_fp16 fp16_param_name = param_to_fp16 [ param ] fp16_param_var = block . var ( fp16_param_name ) fp16_param_var . persistable = True self . _insert_cast_op ( block , idx + 1 , param , param_to_fp16 [ param ] ) if offload : # step3.3: insert fetch op self . _insert_fetch_op ( block , idx , offload_var_name , param ) continue # step3.4: remove cast op if op . type == 'cast' : input_name = op . desc . input_arg_names ( ) [ 0 ] if input_name in param_to_idx : block . _remove_op ( idx , sync = False ) continue # step3.5: change recompute_param to fp16_param for input_name in op . desc . input_arg_names ( ) : if input_name in recompute_to_fp16 : op . _rename_input ( input_name , recompute_to_fp16 [ input_name ] ) for output_name in op . desc . output_arg_names ( ) : if output_name in recompute_to_fp16 : op . _rename_output ( output_name , recompute_to_fp16 [ output_name ] ) # step4: remove recompute_param for name in recompute_to_fp16 . keys ( ) : block . _remove_var ( name , sync = False ) # step5: startup_block add offload visited_vars = set ( ) # FIXME(wangxi): should insert in idx, need move comm init to the head. insert_idx = len ( startup_block . ops ) for idx , op in reversed ( list ( enumerate ( startup_block . ops ) ) ) : for out_name in op . output_arg_names : if out_name in visited_vars : continue if out_name in param_name_to_offload_name : var_name = out_name if offload : offload_var_name = param_name_to_offload_name [ var_name ] self . _insert_offload_op ( startup_block , insert_idx , var_name , offload_var_name ) self . _insert_cast_op ( startup_block , insert_idx , var_name , param_to_fp16 [ var_name ] ) # NOTE(wangxi): cast and offload should insert after broadcast param. # the insert op order is: {mp, dp}broadcast, cast, offload self . _insert_broadcast_op ( startup_block , insert_idx , var_name ) visited_vars . add ( out_name ) block . _sync_with_cpp ( ) startup_block . _sync_with_cpp ( )