code: def column_shard_linear_pass(gm: torch.fx.GraphModule):
      # Split all the linear module with column shard. Currently for testing only.
      mod_graph = gm.graph
      for node in mod_graph.nodes:
          if node.op == "call_module":
              target_module = node.graph.owning_module.get_submodule(
                  node.target
              )
              if isinstance(target_module, torch.nn.Linear):
                  target_module.weight = weight_split(
                      target_module.weight, dim=0, col_normal=False
                  )
                  if target_module.bias is not None:
                      target_module.bias.data = weight_split(
                          target_module.bias.data,
                          dim=0,
                          col_normal=False,
                      )
  
      gm.recompile()
      return gm
  
file_path: ColossalAI/colossalai/fx/passes/shard_1d_pass.py
imports: import operator
  import torch
  import torch.nn
  from colossalai.legacy.tensor import ProcessGroup
  from colossalai.legacy.tensor.compute_spec import ComputePattern, ComputeSpec
  from colossalai.legacy.tensor.distspec import ShardSpec
label_desc: Use image node_id 2 for calling a global function with example usage: column_shard_linear_pass(gm) and returns: gm