code: def get_bert_components():
      vocab_size = 1024
      seq_len = 64
      batchSize = 64
  
      def bert_model_builder():
          model = BertLMModel(
              hidden_size=8192,
              num_layers=4,
              num_attention_heads=32,
              vocab_size=vocab_size,
          )
          return model
  
      def bert_data_gen(device="meta"):
          input_ids = torch.randint(
              0, vocab_size, (batchSize, seq_len), device=device
          )
          attention_mask = torch.ones_like(input_ids, device=device)
          kwargs = dict(
              input_ids=input_ids, attention_mask=attention_mask
          )
          return kwargs
  
      return bert_model_builder, bert_data_gen
  
file_path: colossalai/tests/test_auto_parallel/test_offload/model_utils.py
imports: import torch
  import torch.nn
  from transformers import BertConfig, BertLMHeadModel, GPT2Config, GPT2LMHeadModel
  from tests.components_to_test.registry import non_distributed_component_funcs
label_desc: Use image node_id 1 for calling a global function with example usage: get_bert_components() and returns: bert_model_builder,  bert_data_gen, model, kwargs