[autoparallel]integrate auto parallel feature with new tracer (#3408)

* [autoparallel] integrate new analyzer in module level

* unify the profiling method

* polish

* fix no codegen bug

* fix pass bug

* fix liveness test

* polish

colossalai/_analyzer/_subclasses/flop_tensor.py

@@ -235,7 +235,28 @@ def matmul_flop_jit(inputs: List[Any], outputs: List[Any]) -> Number:
     # Inputs contains the shapes of two matrices.
     input_shapes = [v.shape for v in inputs]
     assert len(input_shapes) == 2, input_shapes
+
+    # There are three cases: 1) gemm, 2) gemv, 3) dot
+    if all(len(shape) == 2 for shape in input_shapes):
+        # gemm
         assert input_shapes[0][-1] == input_shapes[1][-2], input_shapes
+    elif all(len(shape) == 1 for shape in input_shapes):
+        # dot
+        assert input_shapes[0][0] == input_shapes[1][0], input_shapes
+
+        # expand shape
+        input_shapes[0] = torch.Size([1, input_shapes[0][0]])
+        input_shapes[1] = torch.Size([input_shapes[1][0], 1])
+    else:
+        # gemv
+        if len(input_shapes[0]) == 1:
+            assert input_shapes[0][0] == input_shapes[1][-2], input_shapes
+            input_shapes.reverse()
+        else:
+            assert input_shapes[1][0] == input_shapes[0][-1], input_shapes
+
+        # expand the shape of the vector to [batch size, 1]
+        input_shapes[-1] = torch.Size([input_shapes[-1][-1], 1])
     flops = reduce(operator.mul, input_shapes[0]) * input_shapes[-1][-1]
     return flops
 

colossalai/_analyzer/fx/codegen.py

 from typing import Any, Callable, Dict, Iterable, List, Tuple
 
 import torch
+
+try:
+    from torch.fx.graph import CodeGen
+except:
+    pass
 from torch.fx.graph import (
-    CodeGen,
     PythonCode,
     _custom_builtins,
     _format_target,
@@ -48,8 +52,8 @@ def _end_of_ckpt(node: Node, ckpt_level: int) -> bool:
     """
     Check if the node could end the ckpt region at `ckpt_level`
     """
-    if len(node.meta['info'].to_recompute) > ckpt_level:
-        return node.meta['info'].to_recompute[ckpt_level] is not None
+    if len(node.meta['info'].activation_checkpoint) > ckpt_level:
+        return node.meta['info'].activation_checkpoint[ckpt_level] is not None
     return True
 
 
@@ -90,8 +94,8 @@ def _find_nested_ckpt_regions(node_list: List[Node], ckpt_level: int = 0):
     current_region = None
 
     for idx, node in enumerate(node_list):
-        if len(node.meta['info'].to_recompute) > ckpt_level:
-            act_ckpt_label = node.meta['info'].to_recompute[ckpt_level]
+        if len(node.meta['info'].activation_checkpoint) > ckpt_level:
+            act_ckpt_label = node.meta['info'].activation_checkpoint[ckpt_level]
 
             # this activation checkpoint label is not set yet
             # meaning this is the first node of the activation ckpt region
@@ -152,12 +156,12 @@ def emit_ckpt_func(body,
 
     # label given by each layer, e.g. if you are currently at level (0, 1, 1)
     # the label will be '0_1_1'
-    label = "_".join([str(idx) for idx in node_list[0].meta['info'].to_recompute[:ckpt_level + 1]])
+    label = "_".join([str(idx) for idx in node_list[0].meta['info'].activation_checkpoint[:ckpt_level + 1]])
     ckpt_fn_def = _gen_ckpt_fn_def(label, inputs)
     ckpt_func.append(f'{ckpt_fn_def}\n')
 
     # if there is more level to fetch
-    if ckpt_level + 1 < max(map(lambda node: len(node.meta['info'].to_recompute), node_list)):
+    if ckpt_level + 1 < max(map(lambda node: len(node.meta['info'].activation_checkpoint), node_list)):
         ckpt_regions = _find_nested_ckpt_regions(node_list, ckpt_level + 1)
         start_idx = [item[0] for item in ckpt_regions]
         end_idx = [item[1] for item in ckpt_regions]
@@ -215,7 +219,6 @@ def emit_code_with_activation_checkpoint(body, ckpt_func, nodes, emit_node_func,
     ckpt_regions = _find_nested_ckpt_regions(nodes, 0)
     start_idx = [item[0] for item in ckpt_regions]
     end_idx = [item[1] for item in ckpt_regions]
-
     node_list = list(nodes)
 
     node_idx = 0

colossalai/_analyzer/fx/node_util.py

@@ -112,7 +112,7 @@ class MetaInfo:
 
     # should keep the same whenever manipulated
     # ============================= Invariant ==================================
-    to_recompute: Tuple[torch.Tensor] = ()    # (region_0, region_1, ...) support nested codegen
+    activation_checkpoint: Tuple[torch.Tensor] = ()    # (region_0, region_1, ...) support nested codegen
     to_offload: Optional[bool] = False
     sharding_spec: str = 'RR'
 

colossalai/_analyzer/fx/passes/shape_prop.py

@@ -237,7 +237,14 @@ class ShapeProp(torch.fx.Interpreter):
         Returns:
             Any: The value returned from executing the Module
         """
-        wrap_fn = lambda elem: MetaTensor(elem, device=device)
+
+        # wrap_fn = lambda elem: MetaTensor(elem, device=device)
+        def wrap_fn(elem, device=device):
+            if isinstance(elem, torch.Tensor):
+                return MetaTensor(elem, device=device)
+            else:
+                return elem
+
         with self._mode:
             return super().run(*tree_map(wrap_fn, args))
 

colossalai/_analyzer/fx/tracer/bias_addition.py

@@ -21,69 +21,111 @@ def linear_impl(input, weight, bias=None):
 
 
 @register_tracer_impl(F.conv1d, name='_bias_addition_impl')
-def conv1d_impl(input, weight, **kwargs):
-    bias = getattr(kwargs, 'bias', None)
+def conv1d_impl(input, weight, bias=None, stride=_single(1), padding=_single(0), dilation=_single(1), groups=1):
     if bias is None:
-        return F.conv1d(input, weight, **kwargs)
+        return F.conv1d(input, weight, stride=stride, padding=padding, dilation=dilation, groups=groups)
     else:
-        new_kwargs = kwargs
-        new_kwargs['bias'] = None
-        return F.conv1d(input, weight, **kwargs) + bias.reshape((-1, 1))
+        return F.conv1d(input, weight, stride=stride, padding=padding, dilation=dilation, groups=groups) + bias.reshape(
+            (-1, 1))
 
 
 @register_tracer_impl(F.conv2d, name='_bias_addition_impl')
-def conv2d_impl(input, weight, **kwargs):
-    bias = getattr(kwargs, 'bias', None)
+def conv2d_impl(input, weight, bias=None, stride=_pair(1), padding=_pair(0), dilation=_pair(1), groups=1):
     if bias is None:
-        return F.conv2d(input, weight, **kwargs)
+        return F.conv2d(input, weight, stride=stride, padding=padding, dilation=dilation, groups=groups)
     else:
-        new_kwargs = kwargs
-        new_kwargs['bias'] = None
-        return F.conv2d(input, weight, **kwargs) + bias.reshape((-1, 1, 1))
+        return F.conv2d(input, weight, stride=stride, padding=padding, dilation=dilation, groups=groups) + bias.reshape(
+            (-1, 1, 1))
 
 
 @register_tracer_impl(F.conv3d, name='_bias_addition_impl')
-def conv3d_impl(input, weight, **kwargs):
-    bias = getattr(kwargs, 'bias', None)
+def conv3d_impl(input, weight, bias=None, stride=_triple(1), padding=_triple(0), dilation=_triple(1), groups=1):
     if bias is None:
-        return F.conv3d(input, weight, **kwargs)
+        return F.conv3d(input, weight, stride=stride, padding=padding, dilation=dilation, groups=groups)
     else:
-        new_kwargs = kwargs
-        new_kwargs['bias'] = None
-        return F.conv3d(input, weight, **new_kwargs) + bias.reshape((-1, 1, 1, 1))
+        return F.conv3d(input, weight, stride=stride, padding=padding, dilation=dilation, groups=groups) + bias.reshape(
+            (-1, 1, 1, 1))
 
 
 @register_tracer_impl(F.conv_transpose1d, name='_bias_addition_impl')
-def conv_transpose1d_impl(input, weight, **kwargs):
-    bias = getattr(kwargs, 'bias', None)
+def conv_transpose1d_impl(input,
+                          weight,
+                          bias=None,
+                          stride=_single(1),
+                          padding=_single(0),
+                          output_padding=_single(0),
+                          groups=1,
+                          dilation=_single(1)):
     if bias is None:
-        return F.conv_transpose1d(input, weight, **kwargs)
+        return F.conv_transpose1d(input,
+                                  weight,
+                                  stride=stride,
+                                  padding=padding,
+                                  output_padding=output_padding,
+                                  groups=groups,
+                                  dilation=dilation)
     else:
-        new_kwargs = kwargs
-        new_kwargs['bias'] = None
-        return F.conv_transpose1d(input, weight, **new_kwargs) + bias.reshape((-1, 1))
+        return F.conv_transpose1d(input,
+                                  weight,
+                                  stride=stride,
+                                  padding=padding,
+                                  output_padding=output_padding,
+                                  groups=groups,
+                                  dilation=dilation) + bias.reshape((-1, 1))
 
 
 @register_tracer_impl(F.conv_transpose2d, name='_bias_addition_impl')
-def conv_transpose2d_impl(input, weight, **kwargs):
-    bias = getattr(kwargs, 'bias', None)
+def conv_transpose2d_impl(input,
+                          weight,
+                          bias=None,
+                          stride=_pair(1),
+                          padding=_pair(0),
+                          output_padding=_pair(0),
+                          groups=1,
+                          dilation=_pair(1)):
     if bias is None:
-        return F.conv_transpose2d(input, weight, **kwargs)
+        return F.conv_transpose2d(input,
+                                  weight,
+                                  stride=stride,
+                                  padding=padding,
+                                  output_padding=output_padding,
+                                  groups=groups,
+                                  dilation=dilation)
     else:
-        new_kwargs = kwargs
-        new_kwargs['bias'] = None
-        return F.conv_transpose2d(input, weight, **new_kwargs) + bias.reshape((-1, 1, 1))
+        return F.conv_transpose2d(input,
+                                  weight,
+                                  stride=stride,
+                                  padding=padding,
+                                  output_padding=output_padding,
+                                  groups=groups,
+                                  dilation=dilation) + bias.reshape((-1, 1, 1))
 
 
 @register_tracer_impl(F.conv_transpose3d, name='_bias_addition_impl')
-def conv_transpose3d_impl(input, weight, **kwargs):
-    bias = getattr(kwargs, 'bias', None)
+def conv_transpose3d_impl(input,
+                          weight,
+                          bias=None,
+                          stride=_triple(1),
+                          padding=_triple(0),
+                          output_padding=_triple(0),
+                          groups=1,
+                          dilation=_triple(1)):
     if bias is None:
-        return F.conv_transpose3d(input, weight, **kwargs)
+        return F.conv_transpose3d(input,
+                                  weight,
+                                  stride=stride,
+                                  padding=padding,
+                                  output_padding=output_padding,
+                                  groups=groups,
+                                  dilation=dilation)
     else:
-        new_kwargs = kwargs
-        new_kwargs['bias'] = None
-        return F.conv_transpose3d(input, weight, **new_kwargs) + bias.reshape((-1, 1, 1, 1))
+        return F.conv_transpose3d(input,
+                                  weight,
+                                  stride=stride,
+                                  padding=padding,
+                                  output_padding=output_padding,
+                                  groups=groups,
+                                  dilation=dilation) + bias.reshape((-1, 1, 1, 1))
 
 
 @register_tracer_impl(torch.addmm, name='_bias_addition_impl')

colossalai/_analyzer/fx/tracer/tracer.py

@@ -155,7 +155,7 @@ class ColoTracer(Tracer):
 
     def create_node(self, *args, **kwargs) -> Node:
         node = super().create_node(*args, **kwargs)
-        n_info = MetaInfo(node, mod_dir=self.mod_dir, to_recompute=tuple(self.ckpt_regions))
+        n_info = MetaInfo(node, mod_dir=self.mod_dir, activation_checkpoint=tuple(self.ckpt_regions))
         return node
 
     def trace(self,

colossalai/auto_parallel/meta_profiler/__init__.py

 from .meta_registry import *
-from .metainfo import *
 from .registry import meta_register
+from .shard_metainfo import *

colossalai/auto_parallel/meta_profiler/meta_registry/activation.py

@@ -2,9 +2,9 @@ from typing import Callable, List, Tuple
 
 import torch
 
+from colossalai._analyzer._subclasses.flop_tensor import ewise_flop_counter as elementwise_flop_counter
+from colossalai._analyzer.fx.node_util import compute_size_in_bytes as activation_size
 from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem
-from colossalai.fx.profiler.memory_utils import activation_size
-from colossalai.fx.profiler.opcount import elementwise_flop_counter
 
 from ..registry import meta_register
 

colossalai/auto_parallel/meta_profiler/meta_registry/binary_elementwise_ops.py

@@ -2,9 +2,9 @@ from typing import List, Tuple
 
 import torch
 
+from colossalai._analyzer._subclasses.flop_tensor import flop_mapping
+from colossalai._analyzer.fx.node_util import compute_size_in_bytes as activation_size
 from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem
-from colossalai.fx.profiler.memory_utils import activation_size
-from colossalai.fx.profiler.opcount import flop_mapping
 
 from ..constants import BCAST_FUNC_OP, NO_SAVE_ACTIVATION
 from ..registry import meta_register
@@ -17,7 +17,7 @@ def binary_elementwise_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, Train
     """Meta information generator for binary elementwise operations
     NOTE: Some of the binary elementwise operations will discard the input activation after computation, as they
     don't need those tensors for back propagation, for example, if there are two tensors being sent for `torch.add`,
-    they will be discarded right after add operation is done. We create a simple API in `MetaInfo` class to identify
+    they will be discarded right after add operation is done. We create a simple API in `ShardMetaInfo` class to identify
     this behavior, it is critical for better memory estimation.
 
     Returns:

colossalai/auto_parallel/meta_profiler/meta_registry/conv.py

@@ -2,6 +2,8 @@ from typing import Callable, Dict, List, Tuple, Union
 
 import torch
 
+from colossalai._analyzer._subclasses.flop_tensor import flop_mapping
+from colossalai._analyzer.fx.node_util import compute_size_in_bytes
 from colossalai.auto_parallel.tensor_shard.sharding_strategy import (
     MemoryCost,
     OperationData,
@@ -10,8 +12,6 @@ from colossalai.auto_parallel.tensor_shard.sharding_strategy import (
     StrategiesVector,
     TrainCycleItem,
 )
-from colossalai.fx.profiler.memory_utils import activation_size
-from colossalai.fx.profiler.opcount import flop_mapping
 from colossalai.tensor.sharding_spec import ShardingSpec
 
 from ..registry import meta_register
@@ -110,16 +110,16 @@ def convnd_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L
     # calculate memory cost
     # TODO: use profiler to check conv temp memory
     # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward
-    fwd_memory_cost = MemoryCost(
-        activation=activation_size([input_tensor, output_tensor]),
-        parameter=activation_size([weight_tensor, bias_tensor]) if has_bias else activation_size(weight_tensor),
+    fwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor, output_tensor]),
+                                 parameter=compute_size_in_bytes([weight_tensor, bias_tensor])
+                                 if has_bias else compute_size_in_bytes(weight_tensor),
                                  temp=0,
                                  buffer=0)
 
-    bwd_memory_cost = MemoryCost(
-        activation=activation_size([input_tensor, weight_tensor, bias_tensor])
-        if has_bias else activation_size([input_tensor, weight_tensor]),
-        parameter=activation_size([weight_tensor, bias_tensor]) if has_bias else activation_size(weight_tensor),
+    bwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor, weight_tensor, bias_tensor])
+                                 if has_bias else compute_size_in_bytes([input_tensor, weight_tensor]),
+                                 parameter=compute_size_in_bytes([weight_tensor, bias_tensor])
+                                 if has_bias else compute_size_in_bytes(weight_tensor),
                                  temp=0,
                                  buffer=0)
 

colossalai/auto_parallel/meta_profiler/meta_registry/embedding.py

@@ -2,9 +2,9 @@ from typing import List, Tuple
 
 import torch
 
+from colossalai._analyzer._subclasses.flop_tensor import flop_mapping
+from colossalai._analyzer.fx.node_util import compute_size_in_bytes
 from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem
-from colossalai.fx.profiler.memory_utils import activation_size
-from colossalai.fx.profiler.opcount import flop_mapping
 
 from ..registry import meta_register
 
@@ -34,11 +34,11 @@ def embedding_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem
     # NOTE: during the backward phase of torch.nn.Embedding, it seems when the input is large enough, it will
     # have a temp memory which is kind of weird and we don't know the reason yet, so currently we just assume
     # that there will be no temp memory, as the temp memory is significantly smaller than the gradient memory
-    fwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, output_tensor]),
+    fwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor, output_tensor]),
                                  parameter=0,
                                  temp=0,
                                  buffer=0)
-    bwd_memory_cost = MemoryCost(activation=activation_size([weight_tensor]), parameter=0, temp=0, buffer=0)
+    bwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([weight_tensor]), parameter=0, temp=0, buffer=0)
 
     total_memory_cost = MemoryCost(activation=fwd_memory_cost.activation + bwd_memory_cost.activation)
 

colossalai/auto_parallel/meta_profiler/meta_registry/linear.py

@@ -3,6 +3,8 @@ from typing import Callable, Dict, List, Tuple, Union
 
 import torch
 
+from colossalai._analyzer._subclasses.flop_tensor import flop_mapping
+from colossalai._analyzer.fx.node_util import compute_size_in_bytes
 from colossalai.auto_parallel.tensor_shard.sharding_strategy import (
     MemoryCost,
     OperationData,
@@ -11,8 +13,6 @@ from colossalai.auto_parallel.tensor_shard.sharding_strategy import (
     StrategiesVector,
     TrainCycleItem,
 )
-from colossalai.fx.profiler.memory_utils import activation_size
-from colossalai.fx.profiler.opcount import flop_mapping
 from colossalai.tensor.sharding_spec import ShardingSpec
 
 from ..registry import meta_register
@@ -112,14 +112,14 @@ def linear_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L
         # NOTE: Linear don't have buffer and temp in forward and backward phase
         # the forward activation cost is the size of output_tensor, parameter cost is the size of weight_tensor and bias_tensor
         # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward
-        fwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, output_tensor]),
-                                     parameter=activation_size([weight_tensor, bias_tensor]),
+        fwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor, output_tensor]),
+                                     parameter=compute_size_in_bytes([weight_tensor, bias_tensor]),
                                      temp=0,
                                      buffer=0)
 
         # the backward activation cost is the size of input_tensor, weight_tensor and bias_tensor, parameter cost is 0
-        bwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, weight_tensor, bias_tensor]),
-                                     parameter=activation_size([weight_tensor, bias_tensor]),
+        bwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor, weight_tensor, bias_tensor]),
+                                     parameter=compute_size_in_bytes([weight_tensor, bias_tensor]),
                                      temp=0,
                                      buffer=0)
 
@@ -148,14 +148,14 @@ def linear_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L
         # NOTE: Linear don't have buffer and temp in forward and backward phase
         # the forward activation cost is the size of output_tensor, parameter cost is the size of weight_tensor
         # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward
-        fwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, output_tensor]),
-                                     parameter=activation_size(weight_tensor),
+        fwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor, output_tensor]),
+                                     parameter=compute_size_in_bytes(weight_tensor),
                                      temp=0,
                                      buffer=0)
 
         # the backward activation cost is the size of input_tensor and weight_tensor, parameter cost is 0
-        bwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, weight_tensor]),
-                                     parameter=activation_size(weight_tensor),
+        bwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor, weight_tensor]),
+                                     parameter=compute_size_in_bytes(weight_tensor),
                                      temp=0,
                                      buffer=0)
 
@@ -210,48 +210,48 @@ def matmul_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L
     # Check dimension
     if all(len(tensor.shape) == 1 for tensor in input_tensors):
         # Dot
-        fwd_compute_cost = flop_mapping[torch.ops.aten.dot.default](input_tensors, output_tensors)
+        fwd_compute_cost = flop_mapping[torch.ops.aten.matmul.default](input_tensors, output_tensors)
         bwd_compute_cost = flop_mapping[torch.ops.aten.mul.Tensor](input_tensors[0], output_tensors) * 2
 
-        fwd_mem_cost = MemoryCost(activation=activation_size(output_tensors), parameter=0, temp=0, buffer=0)
-        bwd_mem_cost = MemoryCost(activation=activation_size(input_tensors), parameter=0, temp=0, buffer=0)
+        fwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(output_tensors), parameter=0, temp=0, buffer=0)
+        bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(input_tensors), parameter=0, temp=0, buffer=0)
 
     elif len(input_tensors[0].shape) >= 2 and len(input_tensors[1].shape) == 1:
         # gemv case 1: matrix-vector multiplication
         # &
         # batched gemv case 1: batched matrix-vector multiplication
 
-        fwd_compute_cost = flop_mapping[torch.ops.aten.mv.default](
+        fwd_compute_cost = flop_mapping[torch.ops.aten.matmul.default](
             [input_tensors[0].reshape(-1, input_tensors[0].shape[-1]), input_tensors[1]], output_tensors)
 
         # combine the dimensions of output
         bwd_compute_cost = flop_mapping[torch.ops.aten.mul.Tensor](
                            [output_tensors[0].reshape(-1), input_tensors[1]],
                            output_tensors) + \
-                           flop_mapping[torch.ops.aten.mv.default](
+                           flop_mapping[torch.ops.aten.matmul.default](
                            [input_tensors[0].reshape(-1, input_tensors[0].shape[-1]).transpose(0, 1), output_tensors[0].reshape(-1)],
                            output_tensors)
 
-        fwd_mem_cost = MemoryCost(activation=activation_size(output_tensors), parameter=0, temp=0, buffer=0)
-        bwd_mem_cost = MemoryCost(activation=activation_size(input_tensors), parameter=0, temp=0, buffer=0)
+        fwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(output_tensors), parameter=0, temp=0, buffer=0)
+        bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(input_tensors), parameter=0, temp=0, buffer=0)
 
     elif len(input_tensors[0].shape) == 1 and len(input_tensors[1].shape) == 2:
         # gemv case 2: vector-matrix multiplication
-        fwd_compute_cost = flop_mapping[torch.ops.aten.mv.default](input_tensors, output_tensors)
+        fwd_compute_cost = flop_mapping[torch.ops.aten.matmul.default](input_tensors, output_tensors)
 
         bwd_compute_cost = flop_mapping[torch.ops.aten.mul.Tensor]([output_tensors[0], input_tensors[0]], output_tensors) + \
-                           flop_mapping[torch.ops.aten.mv.default]([input_tensors[1], output_tensors[0]], output_tensors)
+                           flop_mapping[torch.ops.aten.matmul.default]([input_tensors[1], output_tensors[0]], output_tensors)
 
-        fwd_mem_cost = MemoryCost(activation=activation_size(output_tensors), parameter=0, temp=0, buffer=0)
-        bwd_mem_cost = MemoryCost(activation=activation_size(input_tensors),
+        fwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(output_tensors), parameter=0, temp=0, buffer=0)
+        bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(input_tensors),
                                   parameter=0,
-                                  temp=activation_size(input_tensors[1]),
+                                  temp=compute_size_in_bytes(input_tensors[1]),
                                   buffer=0)
 
     elif len(input_tensors[0].shape) == 1 and len(input_tensors[1].shape) >= 3:
         # batched gemv case 2: vector-batched matrix multiplication
 
-        fwd_compute_cost = flop_mapping[torch.ops.aten.mv.default](
+        fwd_compute_cost = flop_mapping[torch.ops.aten.matmul.default](
             [input_tensors[1].transpose(-2, -1).reshape(-1, input_tensors[1].shape[-2]), input_tensors[0]],
             [output_tensors[0].reshape(-1)])
 
@@ -260,15 +260,15 @@ def matmul_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L
                            [output_tensors[0].reshape(-1), input_tensors[0]],
                            output_tensors
                            ) + \
-                           flop_mapping[torch.ops.aten.mv.default](
+                           flop_mapping[torch.ops.aten.matmul.default](
                            [input_tensors[1].transpose(-2, -1).reshape(-1, input_tensors[1].shape[-2]).transpose(0, 1), output_tensors[0].reshape(-1)],
                            output_tensors
                            )
 
-        fwd_mem_cost = MemoryCost(activation=activation_size(output_tensors + [input_tensors[1]]))
-        bwd_mem_cost = MemoryCost(activation=activation_size(input_tensors[0]),
+        fwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(output_tensors + [input_tensors[1]]))
+        bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(input_tensors[0]),
                                   parameter=0,
-                                  temp=activation_size(input_tensors[1]),
+                                  temp=compute_size_in_bytes(input_tensors[1]),
                                   buffer=0)
 
     elif len(input_tensors[0].shape) >= 2 and len(input_tensors[1].shape) == 2:
@@ -287,8 +287,8 @@ def matmul_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L
                            [input_tensors[0].reshape(-1, input_tensors[0].shape[-1])]
                            )
 
-        fwd_mem_cost = MemoryCost(activation=activation_size(output_tensors), parameter=0, temp=0, buffer=0)
-        bwd_mem_cost = MemoryCost(activation=activation_size(input_tensors), parameter=0, temp=0, buffer=0)
+        fwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(output_tensors), parameter=0, temp=0, buffer=0)
+        bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(input_tensors), parameter=0, temp=0, buffer=0)
 
     elif len(input_tensors[0].shape) == 2 and len(input_tensors[1].shape) >= 3:
         # batched gemm case 2: matrix-batched matrix multiplication
@@ -306,11 +306,12 @@ def matmul_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L
                            [input_tensors[1].transpose(-2, -1).reshape(-1, input_tensors[1].shape[-2])]
                            )
 
-        fwd_mem_cost = MemoryCost(activation=activation_size(output_tensors) + activation_size(input_tensors[1]),
-                                  temp=activation_size(output_tensors))
-        bwd_mem_cost = MemoryCost(activation=activation_size(input_tensors[0]),
+        fwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(output_tensors) +
+                                  compute_size_in_bytes(input_tensors[1]),
+                                  temp=compute_size_in_bytes(output_tensors))
+        bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(input_tensors[0]),
                                   parameter=0,
-                                  temp=activation_size(input_tensors[1]) + activation_size(output_tensors))
+                                  temp=compute_size_in_bytes(input_tensors[1]) + compute_size_in_bytes(output_tensors))
 
     elif all(len(tensor.shape) >= 3 for tensor in input_tensors):
         # Batched matrix-batched matrix multiplication
@@ -351,8 +352,8 @@ def matmul_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L
                                [input_tensors[0].reshape(-1, input_dim_00, input_dim_01)]
                                )
 
-            fwd_mem_cost = MemoryCost(activation=activation_size(output_tensors))
-            bwd_mem_cost = MemoryCost(activation=activation_size(input_tensors))
+            fwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(output_tensors))
+            bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(input_tensors))
 
         else:
             # Case 2: batch dimensions are different
@@ -381,10 +382,10 @@ def matmul_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem, L
                                )
 
             fwd_mem_cost = MemoryCost(
-                activation=activation_size([output_tensors[0], extended_input_0, extended_input_1]))
-            bwd_mem_cost = MemoryCost(activation=activation_size(input_tensors) -
-                                      activation_size([extended_input_0, extended_input_1]),
-                                      temp=activation_size([extended_input_0, extended_input_1]))
+                activation=compute_size_in_bytes([output_tensors[0], extended_input_0, extended_input_1]))
+            bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(input_tensors) -
+                                      compute_size_in_bytes([extended_input_0, extended_input_1]),
+                                      temp=compute_size_in_bytes([extended_input_0, extended_input_1]))
 
     # compute cost
     compute_cost = TrainCycleItem(fwd=fwd_compute_cost, bwd=bwd_compute_cost, total=fwd_compute_cost + bwd_compute_cost)

colossalai/auto_parallel/meta_profiler/meta_registry/non_spmd.py

@@ -4,8 +4,6 @@ from typing import List, Tuple
 import torch
 
 from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem
-from colossalai.fx.profiler.memory_utils import activation_size
-from colossalai.fx.profiler.opcount import flop_mapping
 
 from ..registry import meta_register
 

colossalai/auto_parallel/meta_profiler/meta_registry/norm.py

@@ -2,6 +2,8 @@ from typing import Callable, Dict, List, Tuple, Union
 
 import torch
 
+from colossalai._analyzer._subclasses.flop_tensor import flop_mapping
+from colossalai._analyzer.fx.node_util import compute_size_in_bytes
 from colossalai.auto_parallel.tensor_shard.sharding_strategy import (
     MemoryCost,
     OperationData,
@@ -10,8 +12,6 @@ from colossalai.auto_parallel.tensor_shard.sharding_strategy import (
     StrategiesVector,
     TrainCycleItem,
 )
-from colossalai.fx.profiler.memory_utils import activation_size
-from colossalai.fx.profiler.opcount import flop_mapping
 from colossalai.tensor.sharding_spec import ShardingSpec
 
 from ..registry import meta_register
@@ -77,17 +77,18 @@ def batchnormnd_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleIt
     # calculate memory cost
     # the fwd activation cost is output plus saved mean and saved inv std
     # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward
-    fwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, output_tensor, mean_tensor, var_tensor]),
-                                 parameter=activation_size([weight_tensor, bias_tensor]),
+    fwd_memory_cost = MemoryCost(activation=compute_size_in_bytes(
+        [input_tensor, output_tensor, mean_tensor, var_tensor]),
+                                 parameter=compute_size_in_bytes([weight_tensor, bias_tensor]),
                                  temp=0,
-                                 buffer=activation_size([mean_tensor, var_tensor]))
+                                 buffer=compute_size_in_bytes([mean_tensor, var_tensor]))
 
     # the bwd memory cost is quite tricky here, BatchNorm will remove saved mean
     # and saved inv std during backward phase
-    bwd_memory_cost = MemoryCost(activation=activation_size([input_tensor]),
-                                 parameter=activation_size([weight_tensor, bias_tensor]),
-                                 temp=activation_size([mean_tensor, var_tensor]),
-                                 buffer=activation_size([mean_tensor, var_tensor]))
+    bwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor]),
+                                 parameter=compute_size_in_bytes([weight_tensor, bias_tensor]),
+                                 temp=compute_size_in_bytes([mean_tensor, var_tensor]),
+                                 buffer=compute_size_in_bytes([mean_tensor, var_tensor]))
 
     # total cost is the sum of forward and backward cost
     total_cost = MemoryCost(activation=fwd_memory_cost.activation + bwd_memory_cost.activation,
@@ -131,15 +132,16 @@ def layernorm_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem
 
     # memory cost
     # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward
-    fwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, output_tensor, weight_tensor, bias_tensor]),
-                                 parameter=activation_size([weight_tensor, bias_tensor]),
+    fwd_memory_cost = MemoryCost(activation=compute_size_in_bytes(
+        [input_tensor, output_tensor, weight_tensor, bias_tensor]),
+                                 parameter=compute_size_in_bytes([weight_tensor, bias_tensor]),
                                  temp=0,
-                                 buffer=activation_size([running_mean, running_var]))
+                                 buffer=compute_size_in_bytes([running_mean, running_var]))
 
-    bwd_memory_cost = MemoryCost(activation=activation_size([input_tensor, weight_tensor, bias_tensor]),
-                                 parameter=activation_size([weight_tensor, bias_tensor]),
-                                 temp=activation_size([running_mean, running_var]),
-                                 buffer=activation_size([running_mean, running_var]))
+    bwd_memory_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor, weight_tensor, bias_tensor]),
+                                 parameter=compute_size_in_bytes([weight_tensor, bias_tensor]),
+                                 temp=compute_size_in_bytes([running_mean, running_var]),
+                                 buffer=compute_size_in_bytes([running_mean, running_var]))
 
     total_cost = MemoryCost(activation=fwd_memory_cost.activation + bwd_memory_cost.activation,
                             parameter=fwd_memory_cost.parameter + bwd_memory_cost.parameter,

colossalai/auto_parallel/meta_profiler/meta_registry/pooling.py

@@ -2,9 +2,9 @@ from typing import List, Tuple
 
 import torch
 
+from colossalai._analyzer._subclasses.flop_tensor import flop_mapping
+from colossalai._analyzer.fx.node_util import compute_size_in_bytes
 from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem
-from colossalai.fx.profiler.memory_utils import activation_size
-from colossalai.fx.profiler.opcount import flop_mapping
 
 from ..registry import meta_register
 
@@ -52,8 +52,8 @@ def avgpool_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem,
     compute_cost = TrainCycleItem(fwd=fwd_compute_cost, bwd=bwd_compute_cost, total=fwd_compute_cost + bwd_compute_cost)
 
     # calculate memory cost
-    fwd_mem_cost = MemoryCost() if is_inplace else MemoryCost(activation=activation_size(output_tensor))
-    bwd_mem_cost = MemoryCost() if is_inplace else MemoryCost(activation=activation_size(input_tensor))
+    fwd_mem_cost = MemoryCost() if is_inplace else MemoryCost(activation=compute_size_in_bytes(output_tensor))
+    bwd_mem_cost = MemoryCost() if is_inplace else MemoryCost(activation=compute_size_in_bytes(input_tensor))
 
     # total cost
     total_mem_cost = MemoryCost(activation=fwd_mem_cost.activation + bwd_mem_cost.activation)
@@ -114,11 +114,11 @@ def maxpool_meta_info(*args, **kwargs) -> Tuple[TrainCycleItem, TrainCycleItem,
     # calculate memory cost
     # NOTE: the index matrix will be discarded in backward phase
     # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward
-    fwd_mem_cost = MemoryCost(activation=activation_size([input_tensor, output_tensor, index_matrix]))
+    fwd_mem_cost = MemoryCost(activation=compute_size_in_bytes([input_tensor, output_tensor, index_matrix]))
 
     # temp memory for backward is the index matrix to be discarded
-    bwd_mem_cost = MemoryCost(activation=activation_size(input_tensor) - activation_size(index_matrix),
-                              temp=activation_size(index_matrix))
+    bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(input_tensor) - compute_size_in_bytes(index_matrix),
+                              temp=compute_size_in_bytes(index_matrix))
 
     # total cost
     total_mem_cost = MemoryCost(activation=fwd_mem_cost.activation + bwd_mem_cost.activation, temp=bwd_mem_cost.temp)

colossalai/auto_parallel/meta_profiler/meta_registry/tensor.py

@@ -2,9 +2,9 @@ from typing import Callable, List, Tuple
 
 import torch
 
+from colossalai._analyzer._subclasses.flop_tensor import flop_mapping
+from colossalai._analyzer.fx.node_util import compute_size_in_bytes
 from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem
-from colossalai.fx.profiler.memory_utils import activation_size
-from colossalai.fx.profiler.opcount import flop_mapping
 
 from ..registry import meta_register
 
@@ -35,11 +35,11 @@ def tensor_related_metainfo(bwd_mem_out_factor: float = 1, bwd_mem_tmp_factor: f
 
         # memory costs
         # NOTE: currently in SPMD solver we always believe that there will be a new tensor created in forward
-        fwd_mem_cost = MemoryCost(activation=activation_size(outputs) * 2, parameter=0, temp=0, buffer=0)
+        fwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(outputs) * 2, parameter=0, temp=0, buffer=0)
 
-        bwd_mem_cost = MemoryCost(activation=activation_size(outputs) * bwd_mem_out_factor,
+        bwd_mem_cost = MemoryCost(activation=compute_size_in_bytes(outputs) * bwd_mem_out_factor,
                                   parameter=0,
-                                  temp=activation_size(outputs) * bwd_mem_tmp_factor,
+                                  temp=compute_size_in_bytes(outputs) * bwd_mem_tmp_factor,
                                   buffer=0)
 
         total_mem_cost = MemoryCost(activation=fwd_mem_cost.activation + bwd_mem_cost.activation,

colossalai/auto_parallel/meta_profiler/meta_registry/where.py

@@ -2,9 +2,9 @@ from typing import List, Tuple
 
 import torch
 
+from colossalai._analyzer._subclasses.flop_tensor import flop_mapping
+from colossalai._analyzer.fx.node_util import compute_size_in_bytes as activation_size
 from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, OperationDataType, TrainCycleItem
-from colossalai.fx.profiler.memory_utils import activation_size
-from colossalai.fx.profiler.opcount import flop_mapping
 
 from ..registry import meta_register
 

colossalai/auto_parallel/meta_profiler/metainfo.py → colossalai/auto_parallel/meta_profiler/shard_metainfo.py

@@ -15,11 +15,11 @@ from colossalai.tensor.sharding_spec import ShardingSpec
 from .constants import INPLACE_MODULE, INPLACE_OPS, NO_SAVE_ACTIVATION
 from .registry import meta_register
 
-__all__ = ['MetaInfo']
+__all__ = ['ShardMetaInfo']
 
 
-class MetaInfo:
-    """MetaInfo class
+class ShardMetaInfo:
+    """ShardMetaInfo class
     This class is used to store meta info based on sharding strategy and the given
     target function.
     """
@@ -46,9 +46,9 @@ class MetaInfo:
         # target function
         self._target = target
 
-        # compute metainfo if possible
+        # compute shard_metainfo if possible
         if self._strategy is not None and self._target is not None:
-            self.compute_metainfo()
+            self.compute_shard_metainfo()
 
     @property
     def strategy(self) -> ShardingStrategy:
@@ -62,13 +62,13 @@ class MetaInfo:
     def strategy(self, strategy: ShardingStrategy) -> None:
         self._strategy = strategy
         if self._strategy is not None and self._target is not None:
-            self.compute_metainfo()
+            self.compute_shard_metainfo()
 
     @target.setter
     def target(self, target: Callable) -> None:
         self._target = target
         if self._strategy is not None and self._target is not None:
-            self.compute_metainfo()
+            self.compute_shard_metainfo()
 
     def compute_sharded_opdata(self, operation_data: OperationData, sharding_spec: ShardingSpec):
         """
@@ -93,7 +93,7 @@ class MetaInfo:
 
         return op_data
 
-    def compute_metainfo(self):
+    def compute_shard_metainfo(self):
         """
         Compute meta info based on sharding strategy and the given target function.
         """

colossalai/auto_parallel/passes/comm_metainfo_pass.py

@@ -4,7 +4,7 @@ import torch
 from torch.fx import GraphModule
 from torch.fx.node import Node
 
-from colossalai.auto_parallel.meta_profiler import MetaInfo
+from colossalai.auto_parallel.meta_profiler import ShardMetaInfo
 from colossalai.auto_parallel.passes.runtime_apply_pass import runtime_apply, runtime_comm_spec_apply
 from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, TrainCycleItem
 from colossalai.tensor.comm_spec import CommSpec
@@ -14,15 +14,15 @@ from colossalai.tensor.sharding_spec import ShardingSpec
 shape_consistency_manager = ShapeConsistencyManager()
 
 
-def _construct_meta_info(node: Node, origin_sharding_spec: ShardingSpec,
-                         target_sharding_spec: ShardingSpec) -> MetaInfo:
+def _construct_shard_meta_info(node: Node, origin_sharding_spec: ShardingSpec,
+                               target_sharding_spec: ShardingSpec) -> ShardMetaInfo:
     # get comm_action_sequence and total_cost from shape_consistency_manager
     _, comm_action_sequence, total_cost = shape_consistency_manager.shape_consistency(
         origin_sharding_spec, target_sharding_spec)
 
-    meta_info = MetaInfo()
+    meta_info = ShardMetaInfo()
     # NOTE: the cost in shape_consistency_manager.mem_cost is the count in number of numel
-    # get mem cost for MetaInfo
+    # get mem cost for ShardMetaInfo
     mem_cost = shape_consistency_manager.mem_cost(comm_action_sequence)
     # extract user that has _meta_data and extract element length
     input_node = next(n for n in node._input_nodes if hasattr(n, '_meta_data'))
@@ -36,12 +36,12 @@ def _construct_meta_info(node: Node, origin_sharding_spec: ShardingSpec,
 
     meta_info.memory_cost = mem_cost
 
-    # get computation cost for MetaInfo
+    # get computation cost for ShardMetaInfo
     meta_info.compute_cost = TrainCycleItem(total_cost['forward'] * element_length,
                                             total_cost['backward'] * element_length,
                                             total_cost['total'] * element_length)
 
-    # get tensor shape for MetaInfo
+    # get tensor shape for ShardMetaInfo
     origin_sharding_spec: ShardingSpec
     target_sharding_spec: ShardingSpec
     input_shape = origin_sharding_spec.get_sharded_shape_per_device()
@@ -54,7 +54,7 @@ def _construct_meta_info(node: Node, origin_sharding_spec: ShardingSpec,
     return meta_info
 
 
-def _runtime_apply_meta_info(node: Node, origin_spec_dict, sharding_spec_dict) -> MetaInfo:
+def _runtime_apply_meta_info(node: Node, origin_spec_dict, sharding_spec_dict) -> ShardMetaInfo:
     """
     This method is used to construct `MetaInto` for shape consistency node
     """
@@ -65,17 +65,17 @@ def _runtime_apply_meta_info(node: Node, origin_spec_dict, sharding_spec_dict) -
     origin_sharding_spec, target_sharding_spec = origin_spec_dict[node_index], sharding_spec_dict[node_index][
         user_node_index]
 
-    return _construct_meta_info(node, origin_sharding_spec, target_sharding_spec)
+    return _construct_shard_meta_info(node, origin_sharding_spec, target_sharding_spec)
 
 
-def _runtime_comm_spec_apply_meta_info(node: Node, comm_actions_dict: Dict) -> MetaInfo:
+def _runtime_comm_spec_apply_meta_info(node: Node, comm_actions_dict: Dict) -> ShardMetaInfo:
     # extract node_index and op_data_name
     node_index, op_data_name = node.args[2], node.args[3]
 
     comm_action = comm_actions_dict[node_index][op_data_name]
     if isinstance(comm_action.comm_spec, CommSpec):
         # this case is for all_reduce, there will be no memory cost
-        meta_info = MetaInfo()
+        meta_info = ShardMetaInfo()
         meta_info.memory_cost = TrainCycleItem(MemoryCost(), MemoryCost(), MemoryCost)
         output_node = next(n for n in node.users if hasattr(n, '_meta_data'))
         element_length = output_node._meta_data.element_size()
@@ -93,7 +93,7 @@ def _runtime_comm_spec_apply_meta_info(node: Node, comm_actions_dict: Dict) -> M
         # this case will be handled by shape consistency manager
         origin_sharding_spec, target_sharding_spec = comm_action.comm_spec['src_spec'], comm_action.comm_spec[
             'tgt_spec']
-        meta_info = _construct_meta_info(node, origin_sharding_spec, target_sharding_spec)
+        meta_info = _construct_shard_meta_info(node, origin_sharding_spec, target_sharding_spec)
 
     return meta_info
 
@@ -105,9 +105,9 @@ def comm_metainfo_pass(gm: GraphModule, sharding_spec_dict: Dict, origin_spec_di
     """
     for node in gm.graph.nodes:
         if node.target == runtime_apply:
-            setattr(node, 'best_metainfo', _runtime_apply_meta_info(node, origin_spec_dict, sharding_spec_dict))
+            setattr(node, 'best_strategy_info', _runtime_apply_meta_info(node, origin_spec_dict, sharding_spec_dict))
         elif node.target == runtime_comm_spec_apply:
-            setattr(node, 'best_metainfo', _runtime_comm_spec_apply_meta_info(node, comm_actions_dict))
+            setattr(node, 'best_strategy_info', _runtime_comm_spec_apply_meta_info(node, comm_actions_dict))
         else:
             pass
     return gm

colossalai/auto_parallel/passes/meta_info_prop.py

@@ -7,7 +7,7 @@ import torch.fx
 from torch.fx import GraphModule
 from torch.fx.node import Node
 
-from colossalai.auto_parallel.meta_profiler import MetaInfo
+from colossalai.auto_parallel.meta_profiler import ShardMetaInfo
 from colossalai.auto_parallel.passes.constants import OUTPUT_SAVED_MOD, OUTPUT_SAVED_OPS
 from colossalai.fx._compatibility import compatibility
 from colossalai.fx.profiler import GraphInfo
@@ -96,12 +96,12 @@ class MetaInfoProp:
         """
         Handle other kind of nodes
         """
-        assert hasattr(node, 'best_metainfo'), f"Cannot find best_metainfo in node {node}, {node.op}"
+        assert hasattr(node, 'best_strategy_info'), f"Cannot find best_strategy_info in node {node}, {node.op}"
         graph_info = GraphInfo()
-        meta_info = node.best_metainfo
-        meta_info: MetaInfo
+        meta_info = node.best_strategy_info
+        meta_info: ShardMetaInfo
 
-        # set data_ptr for input_tensor in MetaInfo class
+        # set data_ptr for input_tensor in ShardMetaInfo class
         input_tensors: List[torch.Tensor] = meta_info.fwd_in
         buffer_tensors: List[torch.Tensor] = meta_info.fwd_buffer
         output_tensors: List[torch.Tensor] = meta_info.fwd_out