code: def _apply_tp_ffn(mlp: FeedForward) -> None:
      assert hasattr(mlp, "w1")
      assert hasattr(mlp, "w3")
      assert hasattr(mlp, "w2")
  
      _apply_tp_linear(mlp.w1, "colwise")
      _apply_tp_linear(mlp.w3, "colwise")
      _apply_tp_linear(mlp.w2, "rowwise")
  
      world_size = _get_world_size()
      mlp.register_forward_hook(
          lambda _module, _input, output: funcol.all_reduce(
              output, "sum", list(range(world_size))
          )
      )
  
file_path: gpt-fast/tp.py
imports: import os
  from typing import List, Optional
  import torch
  import torch.distributed
  from torch import nn
  from torch.distributed import _functional_collectives
  from model import Attention, FeedForward, Transformer
  from quantize import WeightOnlyInt4Linear
label_desc: Use image node_id 7 for calling a global function with example usage: _apply_tp_ffn(mlp) without return types