code: def model_cfg(test_data_dir):
      model_cfg = {
          "mcore_gpt": True,
          "micro_batch_size": 4,
          "global_batch_size": 8,
          "rampup_batch_size": None,
          "tensor_model_parallel_size": 1,
          "pipeline_model_parallel_size": 1,
          "virtual_pipeline_model_parallel_size": None,
          "encoder_seq_length": 512,
          "max_position_embeddings": 512,
          "num_layers": 1,
          "hidden_size": 128,
          "ffn_hidden_size": 512,
          "num_attention_heads": 2,
          "num_query_groups": 1,
          "init_method_std": 0.02,
          "use_scaled_init_method": True,
          "hidden_dropout": 0.0,
          "attention_dropout": 0.0,
          "ffn_dropout": 0,
          "kv_channels": None,
          "apply_query_key_layer_scaling": False,
          "normalization": "layernorm",
          "layernorm_epsilon": 1e-05,
          "do_layer_norm_weight_decay": False,
          "make_vocab_size_divisible_by": 128,
          "pre_process": True,
          "post_process": True,
          "persist_layer_norm": True,
          "bias": False,
          "activation": "gelu",
          "headscale": False,
          "transformer_block_type": "pre_ln",
          "openai_gelu": False,
          "normalize_attention_scores": True,
          "position_embedding_type": "rope",
          "rotary_percentage": 1.0,
          "attention_type": "multihead",
          "share_embeddings_and_output_weights": False,
          "overlap_p2p_comm": False,
          "batch_p2p_comm": True,
          "seq_len_interpolation_factor": None,
          "tokenizer": {
              "library": "huggingface",
              "type": "tiiuae/falcon-40b",
              "use_fast": True,
          },
          "native_amp_init_scale": 4294967296,
          "native_amp_growth_interval": 1000,
          "hysteresis": 2,
          "fp32_residual_connection": False,
          "fp16_lm_cross_entropy": False,
          "megatron_amp_O2": False,
          "grad_allreduce_chunk_size_mb": 125,
          "grad_div_ar_fusion": True,
          "gradient_accumulation_fusion": False,
          "bias_activation_fusion": False,
          "bias_dropout_add_fusion": False,
          "masked_softmax_fusion": True,
          "get_attention_mask_from_fusion": True,
          "seed": 1234,
          "resume_from_checkpoint": None,
          "use_cpu_initialization": False,
          "onnx_safe": False,
          "apex_transformer_log_level": 30,
          "gradient_as_bucket_view": True,
          "sync_batch_comm": False,
          "activations_checkpoint_granularity": None,
          "activations_checkpoint_method": None,
          "activations_checkpoint_num_layers": None,
          "num_micro_batches_with_partial_activation_checkpoints": None,
          "activations_checkpoint_layers_per_pipeline": None,
          "sequence_parallel": False,
          "transformer_engine": True,
          "fp8": False,
          "fp8_e4m3": False,
          "fp8_hybrid": False,
          "fp8_margin": 0,
          "fp8_interval": 1,
          "fp8_amax_history_len": 1,
          "fp8_amax_compute_algo": "most_recent",
          "reduce_amax": True,
          "use_emha": False,
          "ub_tp_comm_overlap": False,
          "ub_tp_comm_overlap_cfg": None,
          "use_flash_attention": False,
          "nsys_profile": {
              "enabled": False,
              "start_step": 10,
              "end_step": 10,
              "ranks": [0],
              "gen_shape": False,
          },
          "optim": {
              "name": "distributed_fused_adam",
              "lr": "2e-4",
              "weight_decay": 0.01,
              "betas": [0.9, 0.98],
              "sched": {
                  "name": "CosineAnnealing",
                  "warmup_steps": 500,
                  "constant_steps": 50000,
                  "min_lr": "2e-5",
              },
          },
          "gc_interval": 0,
          "precision": "bf16",
          "new_decoder_architecture": False,
          "parallel_attention": True,
          "name": "megatron_falcon_gpt",
          "target": "nemo.collections.nlp.models.language_modeling.megatron_gpt_model.MegatronGPTModel",
      }
      return model_cfg
  
file_path: NeMo/tests/collections/nlp/test_falcon_model.py
imports: import pytest
  import torch
  from omegaconf import DictConfig
  from pytorch_lightning import Trainer
  from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer
  from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel
  from nemo.collections.nlp.modules.common.megatron.utils import get_ltor_masks_and_position_ids
  from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategy
label_desc: Use image node_id 1 for calling a global function with example usage: model_cfg(test_data_dir) and returns: model_cfg