code: def initialize_gpt_model(seed, use_te=True, **config_kwargs):
      torch.manual_seed(seed)
      model_parallel_cuda_manual_seed(seed)
  
      default_config_kwargs = dict(
          num_layers=8,
          hidden_size=16,
          num_attention_heads=8,
          use_cpu_initialization=True,
      )
      default_config_kwargs.update(**config_kwargs)
      transformer_config = TransformerConfig(**default_config_kwargs)
      pre_process = ps.is_pipeline_first_stage()
      post_process = ps.is_pipeline_last_stage()
      layer_spec = (
          get_gpt_layer_with_transformer_engine_spec()
          if use_te
          else get_gpt_layer_local_spec()
      )
      model = GPTModel(
          config=transformer_config,
          transformer_layer_spec=layer_spec,
          vocab_size=128,
          max_sequence_length=4,
          pre_process=pre_process,
          post_process=post_process,
      )
  
      with torch.no_grad():
          for p in model.parameters():
              p.random_()
      return model
  
file_path: megatron-lm/tests/unit_tests/dist_checkpointing/models/test_gpt_model.py
imports: import pytest
  import torch
  from torch.distributed._tensor import DeviceMesh
  from megatron.core.dist_checkpointing import save, load, load_plain_tensors
  from megatron.core import parallel_state
  from megatron.core.dist_checkpointing.dict_utils import diff
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.models.gpt.gpt_model import GPTModel
  from tests.unit_tests.dist_checkpointing import TempNamedDir
  from tests.unit_tests.test_utilities import Utils
  from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
  from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec, get_gpt_layer_local_spec
label_desc: Use image node_id 1 for calling a global function with example usage: initialize_gpt_model(seed, use_te) and returns: model