| 44 | |
| 45 | |
| 46 | def get_model_defaults(model_size: ModelSize): |
| 47 | # Define default model configurations for base and large sizes |
| 48 | default_model_configs = { |
| 49 | "base": { |
| 50 | "num_hidden_layers": 22, |
| 51 | "hidden_size": 768, |
| 52 | "intermediate_size": 1152, |
| 53 | "num_attention_heads": 12, # to have head size of 64 |
| 54 | }, |
| 55 | "large": { |
| 56 | "num_hidden_layers": 28, |
| 57 | "hidden_size": 1024, |
| 58 | "intermediate_size": 2624, |
| 59 | "num_attention_heads": 16, |
| 60 | }, |
| 61 | } |
| 62 | |
| 63 | # Select the default model config based on the model_size argument |
| 64 | model_config = default_model_configs[model_size.value] |
| 65 | |
| 66 | # Additional default configurations common to both sizes |
| 67 | default_model_config_common = { |
| 68 | "vocab_size": 50368, |
| 69 | "init_method": "full_megatron", |
| 70 | "attention_layer": "rope", |
| 71 | "attention_probs_dropout_prob": 0.0, |
| 72 | "attn_out_bias": False, |
| 73 | "attn_out_dropout_prob": 0.1, |
| 74 | "attn_qkv_bias": False, |
| 75 | "bert_layer": "prenorm", |
| 76 | "embed_dropout_prob": 0.0, |
| 77 | "embed_norm": True, |
| 78 | "final_norm": True, |
| 79 | "skip_first_prenorm": True, |
| 80 | "embedding_layer": "sans_pos", |
| 81 | "loss_function": "fa_cross_entropy", |
| 82 | "loss_kwargs": {"reduction": "mean"}, |
| 83 | "mlp_dropout_prob": 0.0, |
| 84 | "mlp_in_bias": False, |
| 85 | "mlp_layer": "glu", |
| 86 | "mlp_out_bias": False, |
| 87 | "normalization": "layernorm", |
| 88 | "norm_kwargs": {"eps": 1e-5, "bias": False}, |
| 89 | "hidden_act": "gelu", |
| 90 | "head_pred_act": "gelu", |
| 91 | "activation_function": "gelu", # better safe than sorry |
| 92 | "padding": "unpadded", |
| 93 | "rotary_emb_dim": None, |
| 94 | "rotary_emb_scale_base": None, |
| 95 | "rotary_emb_interleaved": False, |
| 96 | "local_attn_rotary_emb_base": 10000.0, |
| 97 | "local_attn_rotary_emb_dim": None, |
| 98 | "allow_embedding_resizing": True, |
| 99 | "sliding_window": 128, |
| 100 | "global_attn_every_n_layers": 3, |
| 101 | "unpad_embeddings": True, |
| 102 | "compile_model": True, |
| 103 | "use_fa2": True, |