Function add_model_config_args

SwissArmyTransformer/sat/arguments.py:32–97

Model arguments

(parser)

_GLOBAL_RANDOM_SEED = None

def add_model_config_args(parser):
    """Model arguments"""

    group = parser.add_argument_group('model', 'model configuration')

    # --------------- Core hyper-parameters --------------- 
    group.add_argument('--num-layers', type=int, default=6,
                       help='num of layers')
    group.add_argument('--hidden-size', type=int, default=1024,
                       help='transformer hidden size')
    group.add_argument('--num-attention-heads', type=int, default=16,
                       help='num of transformer attention heads')
    group.add_argument('--vocab-size', type=int, default=100,
                       help='vocab size for tokenization. the size of word_embeddings.')
    group.add_argument('--max-sequence-length', type=int, default=512,
                       help='maximum number of position embeddings to use')
    
    # --------------- Optional hyper-parameters --------------- 
    # for developers:
    #   if you want to add more optional arch hyper-parameters, please also add them in model_io.py, base_model.py, and of course, transformer.py and anywhere in use.

    group.add_argument('--layernorm-order', type=str, default='pre',
                       help='choose from "pre", "post", "sandwich".',
                       choices=['post', # In the original Transformer.
                                'pre', # Used by most current frameworks.
                                'sandwich' # More stable.
                                ])
    # The inner-hidden-size in MLP, default "None" means 4*hidden-size
    group.add_argument('--inner-hidden-size', type=int, default=None,
                       help='inner hidden size for transformer FFN, None means 4*hidden-size')
    # The hidden-size-per-attention-head in Self and Cross Attention, 
    # default "None" means hidden-size/num-attention-heads.
    group.add_argument('--hidden-size-per-attention-head', type=int, default=None)
    # TODO: fully test it, support the generation.
    group.add_argument('--model-parallel-size', type=int, default=1,
                       help='size of the model parallel. only use if you are an expert.')

    group.add_argument('--skip-init', action='store_true',
                       help='skip model initialization')
    
    group.add_argument('--use-gpu-initialization', action='store_true', 
                       help='initialize model on gpu')

    group.add_argument('--num-multi-query-heads', type=int, default=0, 
                       help='use multi-query attention, num of kv groups. 0 means multi-head attention.')
    group.add_argument('--is-gated-mlp', action='store_true',
                       help='use gated MLP (GLU), common in LLAMA etc.')
    group.add_argument('--is-rotary-emb', action='store_true',
                       help='use rotary embedding, common in LLAMA etc.')
    
    # ---------------  Inessential hyper-parameters --------------- 

    # Dropout and eps hyper-parameters
    group.add_argument('--layernorm-epsilon', type=float, default=1e-5,
                       help='layer norm epsilon')
    group.add_argument('--hidden-dropout', type=float, default=0.1,
                       help='dropout probability for hidden state transformer')
    group.add_argument('--attention-dropout', type=float, default=0.1,