code: def replace_arguments(config, world_size):
      return {
          # 1. reduce hidden size
          "attn.embed_dim": config.hidden_size // world_size,
          "attn.split_size": config.hidden_size // world_size,
          "crossattention.embed_dim": config.hidden_size // world_size,
          "crossattention.split_size": config.hidden_size // world_size,
          # 2. reduce number of heads
          "attn.num_heads": config.num_attention_heads // world_size,
          "crossattention.num_heads": config.num_attention_heads
          // world_size,
      }
  
file_path: parallelformers/parallelformers/policies/gpt2.py
imports: from transformers.models.gpt2.modeling_gpt2 import GPT2Block
  from parallelformers.policies.base import Policy
  from parallelformers.policies.base.policy import Layer
  from parallelformers.utils import AllReduceConv1D
label_desc: Use image node_id 1 for calling the GPT2Policy obj's underlying member method code with example usage: obj.replace_arguments(config, world_size) and returns: dict