code: def initialize_mlp(glu=True):
      model_parallel_cuda_manual_seed(123)
      pp_size = parallel_state.get_pipeline_model_parallel_world_size()
      transformer_config = TransformerConfig(
          num_layers=pp_size,
          hidden_size=12,
          num_attention_heads=4,
          use_cpu_initialization=True,
          gated_linear_unit=glu,
      )
      return MLP(
          transformer_config,
          get_gpt_layer_with_transformer_engine_spec().submodules.mlp.submodules,
      )
  
file_path: megatron-lm/tests/unit_tests/dist_checkpointing/models/test_mlp_glu.py
imports: import pytest
  import torch
  from torch.optim import Adam
  from megatron.core import parallel_state
  from megatron.core.dist_checkpointing.dict_utils import diff, nested_values
  from megatron.core.dist_checkpointing.optimizer import get_param_id_to_sharded_param_map, optim_state_to_sharding_state
  from megatron.core.transformer.mlp import MLP
  from tests.unit_tests.dist_checkpointing import TempNamedDir
  from tests.unit_tests.test_utilities import Utils
  from megatron.core.dist_checkpointing import save, load, load_plain_tensors, ShardedTensor
  from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
label_desc: Use image node_id 1 for calling a global function with example usage: initialize_mlp(glu) and returns: MLP