code: def mem_test_for_node_strategy(
      rank: int,
      model: torch.nn.Module,
      device_mesh: DeviceMesh,
      node_index: int,
      strategy_number: int,
      input_args: List[torch.Tensor],
      meta_arg_names: List[str],
      input_kwargs: Dict[str, torch.Tensor] = {},
  ):
      for strategy_index in range(strategy_number):
          # We need to copy the model to avoid do backward more than once in same graph
          model_to_shard, args_to_shard, kwargs_to_shard = (
              copy.deepcopy(model),
              copy.deepcopy(input_args),
              copy.deepcopy(input_kwargs),
          )
  
          tracer = ColoTracer(bias_addition_split=True)
          input_sample = {}
          for input_arg, meta_arg_name in zip(
              input_args, meta_arg_names
          ):
              input_sample[meta_arg_name] = torch.rand(
                  input_arg.shape
              ).to("meta")
          for meta_kwarg_name, input_kwarg in input_kwargs.items():
              input_sample[meta_kwarg_name] = torch.rand(
                  input_kwarg.shape
              ).to("meta")
          graph = tracer.trace(
              root=model_to_shard, meta_args=input_sample
          )
          gm = ColoGraphModule(
              model_to_shard, graph, model_to_shard.__class__.__name__
          )
          shape_prop_pass(gm, *input_sample.values())
          gm.recompile()
          solver_options = SolverOptions()
          strategies_constructor = StrategiesConstructor(
              graph, device_mesh, solver_options
          )
          strategies_constructor.build_strategies_and_cost()
          target_node = list(graph.nodes)[node_index]
  
          # solution construction
          # construct the strategy for the target node
          solution_len = len(strategies_constructor.leaf_strategies)
          solution = [0] * solution_len
          solution[node_index] = strategy_index
  
          # construct the strategy for the output node
          placeholder_strategy = list(graph.nodes)[
              -1
          ].strategies_vector[0]
  
          output_key = next(
              key
              for key in target_node.strategies_vector[
                  strategy_index
              ].sharding_specs.keys()
              if key.type == OperationDataType.OUTPUT
          )
          placeholder_strategy.sharding_specs[
              output_key
          ] = target_node.strategies_vector[
              strategy_index
          ].sharding_specs[
              output_key
          ]
  
          (
              gm,
              sharding_spec_dict,
              origin_spec_dict,
              comm_actions_dict,
          ) = runtime_preparation_pass(
              gm, solution, device_mesh, strategies_constructor
          )
          gm = runtime_apply_pass(gm)
          gm.recompile()
          gm: GraphModule
  
          num_of_strategies = len(target_node.strategies_vector)
          if rank == 0:
              print("=======================")
              print(
                  f"#strategy_index: {strategy_index + 1}/{num_of_strategies}"
              )
              pprint(target_node.strategies_vector[strategy_index])
  
          # warmup
          with torch.no_grad():
              output = gm(
                  *args_to_shard,
                  sharding_spec_convert_dict=sharding_spec_dict,
                  origin_node_sharding_spec_dict=origin_spec_dict,
                  comm_actions_dict=comm_actions_dict,
                  **kwargs_to_shard,
              )
  
          del output
          # forward memory compare
          if rank == 0:
              torch.cuda.reset_peak_memory_stats()
              mem_stamp0 = torch.cuda.memory_allocated()
          output = gm(
              *args_to_shard,
              sharding_spec_convert_dict=sharding_spec_dict,
              origin_node_sharding_spec_dict=origin_spec_dict,
              comm_actions_dict=comm_actions_dict,
              **kwargs_to_shard,
          )
  
          if rank == 0:
              # print forward memory allocated and peak memory stats in kb
              print(
                  f"forward memory allocated: {(torch.cuda.memory_allocated() - mem_stamp0) / 1024} kb, peak memory stats: {(torch.cuda.max_memory_allocated() - mem_stamp0) / 1024} kb"
              )
  
          # backward memory compare
          grad_tensors = torch.ones_like(output)
          torch.cuda.reset_peak_memory_stats()
          mem_stamp0 = torch.cuda.memory_allocated()
          torch.autograd.backward(output, grad_tensors)
  
          if rank == 0:
              # print backward memory allocated and peak memory stats in kb
              print(
                  f"backward memory allocated: {(torch.cuda.memory_allocated() - mem_stamp0) / 1024} kb, peak memory stats: {(torch.cuda.max_memory_allocated() - mem_stamp0) / 1024} kb"
              )
  
              # estimated memory
              if target_node.op == "call_module":
                  metainfo = ShardMetaInfo(
                      target_node.strategies_vector[strategy_index],
                      target_node.graph.owning_module.get_submodule(
                          target_node.target
                      ),
                  )
              else:
                  metainfo = ShardMetaInfo(
                      target_node.strategies_vector[strategy_index],
                      target_node.target,
                  )
  
              print("estimated memory:")
              print(
                  f"forward activation: {metainfo.memory_cost.fwd.activation / 1024} kb, forward param: {metainfo.memory_cost.fwd.parameter / 1024} kb"
              )
              print(
                  f"forward temp: {metainfo.memory_cost.fwd.temp / 1024} kb, forward buffer: {metainfo.memory_cost.fwd.buffer / 1024} kb"
              )
              print(
                  f"backward activation: {metainfo.memory_cost.bwd.activation / 1024} kb, backward param: {metainfo.memory_cost.bwd.parameter / 1024} kb"
              )
              print(
                  f"backward temp: {metainfo.memory_cost.bwd.temp / 1024} kb, backward buffer: {metainfo.memory_cost.bwd.buffer / 1024} kb"
              )
              print("=======================")
  
file_path: ColossalAI/tests/test_auto_parallel/test_tensor_shard/test_metainfo/utils.py
imports: import copy
  from pprint import pprint
  from typing import Dict, List
  import torch
  from torch.fx import GraphModule
  from colossalai._analyzer.fx.graph_module import ColoGraphModule
  from colossalai._analyzer.fx.passes import shape_prop_pass
  from colossalai._analyzer.fx.tracer.tracer import ColoTracer
  from colossalai.auto_parallel.passes.runtime_apply_pass import runtime_apply_pass
  from colossalai.auto_parallel.passes.runtime_preparation_pass import runtime_preparation_pass
  from colossalai.auto_parallel.tensor_shard.options import SolverOptions
  from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationDataType, TrainCycleItem
  from colossalai.auto_parallel.tensor_shard.solver import StrategiesConstructor
  from colossalai.device.device_mesh import DeviceMesh
label_desc: Use image node_id 1 for calling a global function with example usage: mem_test_for_node_strategy(rank, model, device_mesh, node_index, strategy_number, input_args, meta_arg_names, input_kwargs) without return types