mention our friends

app.py

@@ -32,10 +32,8 @@ with st.expander("More options"):
     precisions_values = ('O0', 'O1', 'O3')
     precision = st.selectbox('Precision', precisions_names, index=1)
 
-    vocab_size = int(st.number_input('Vocabulary size', min_value=1, step=1, value=50257, format="%i"))
-
 args = mem_calc.parse_args(f"""
-    --model {model} --vocab_size {vocab_size} --optimizer {optimizers_values[optimizers_names.index(optimizer)]}
+    --model {model} --optimizer {optimizers_values[optimizers_names.index(optimizer)]}
     {'--checkpoint' if checkpoint else ''} {'--offload' if offload else ''} {'--albert' if share_params else ''}
     --fp16-level {precisions_values[precisions_names.index(precision)]} --bsz {batch_size} --seqlen {seq_len}
 """.split())

mem_calc.py

@@ -123,7 +123,7 @@ def parse_args(args=None):
                         help='FP16-level to use. O0 = FP32; O1 = mixed-precision (16+32); O3 = fp16. Default: O1.')
     parser.add_argument('--model', default='', choices=list(models.keys()), help='Predefined NLP transformer models')
     parser.add_argument('--optimizer', default='adam', choices=OPTIMIZERS, help='The optimizer to use.')
-    parser.add_argument('--vocab_size', type=int, default=50257, help='The vocabulary to use.')
+    parser.add_argument('--vocab_size', type=int, default=None, help='The vocabulary to use.')
     parser.add_argument('--offload', action='store_true', help='Whether to use optimizer offload.')
     parser.add_argument('--ngpus', type=int, default=1, help='The number of gpus. Default: 1')
     parser.add_argument('--zero', type=int, default=0,

models.py

@@ -1,97 +1,120 @@
 models = {}
-models['bert-s'] = {}
-models['bert-s']['seqlen'] = 512
-models['bert-s']['dmodel'] = 768
-models['bert-s']['dhidden'] = 3072
-models['bert-s']['nlayers'] = 12
-
-models['bert-l'] = {}
-models['bert-l']['seqlen'] = 512
-models['bert-l']['dmodel'] = 1024
-models['bert-l']['dhidden'] = 4096
-models['bert-l']['nlayers'] = 24
+models['bert-base'] = {}
+models['bert-base']['seqlen'] = 512
+models['bert-base']['dmodel'] = 768
+models['bert-base']['dhid'] = 3072
+models['bert-base']['nlayers'] = 12
+models['bert-base']['vocab_size'] = 30522
+
+
+models['bert-large'] = {}
+models['bert-large']['seqlen'] = 512
+models['bert-large']['dmodel'] = 1024
+models['bert-large']['dhid'] = 4096
+models['bert-large']['nlayers'] = 24
+models['bert-large']['vocab_size'] = 30522
 
 models['t5-3b'] = {}
 models['t5-3b']['seqlen'] = 512
 models['t5-3b']['dmodel'] = 1024
-models['t5-3b']['dhidden'] = 16384
+models['t5-3b']['dhid'] = 16384
 models['t5-3b']['nlayers'] = 48
+models['t5-3b']['vocab_size'] = 32128
 
 models['t5-11b'] = {}
 models['t5-11b']['seqlen'] = 512
 models['t5-11b']['dmodel'] = 1024
-models['t5-11b']['dhidden'] = 64*1024
+models['t5-11b']['dhid'] = 64*1024
 models['t5-11b']['nlayers'] = 48
+models['t5-11b']['vocab_size'] = 32128
 
 models['gpt2-s'] = {}
 models['gpt2-s']['seqlen'] = 1024
 models['gpt2-s']['dmodel'] = 768
-models['gpt2-s']['dhidden'] = 768*4
+models['gpt2-s']['dhid'] = 768*4
 models['gpt2-s']['nlayers'] = 12
+models['gpt2-s']['vocab_size'] = 50257
 
 models['gpt2-m'] = {}
 models['gpt2-m']['seqlen'] = 1024
 models['gpt2-m']['dmodel'] = 1024
-models['gpt2-m']['dhidden'] = 1024*4
+models['gpt2-m']['dhid'] = 1024*4
 models['gpt2-m']['nlayers'] = 24
+models['gpt2-m']['vocab_size'] = 50257
 
 models['gpt2-l'] = {}
 models['gpt2-l']['seqlen'] = 1024
 models['gpt2-l']['dmodel'] = 1280
-models['gpt2-l']['dhidden'] = 1280*4
+models['gpt2-l']['dhid'] = 1280*4
 models['gpt2-l']['nlayers'] = 36
+models['gpt2-l']['vocab_size'] = 50257
 
 models['gpt2-xl'] = {}
 models['gpt2-xl']['seqlen'] = 1024
 models['gpt2-xl']['dmodel'] = 1600
-models['gpt2-xl']['dhidden'] = 1600*4
+models['gpt2-xl']['dhid'] = 1600*4
 models['gpt2-xl']['nlayers'] = 48
+models['gpt2-xl']['vocab_size'] = 50257
 
+models['gpt-j-6b'] = {}
+models['gpt-j-6b']['seqlen'] = 2048
+models['gpt-j-6b']['dmodel'] = 4096
+models['gpt-j-6b']['dhid'] = 4096 * 4
+models['gpt-j-6b']['nlayers'] = 28
+models['gpt-j-6b']['vocab_size'] = 50400
 
 models['gpt3-s'] = {}
 models['gpt3-s']['seqlen'] = 2048
 models['gpt3-s']['dmodel'] = 768
-models['gpt3-s']['dhidden'] = 768*4
+models['gpt3-s']['dhid'] = 768*4
 models['gpt3-s']['nlayers'] = 12
+models['gpt3-s']['vocab_size'] = 50257  # from public reimplementations
 
 models['gpt3-m'] = {}
 models['gpt3-m']['seqlen'] = 2048
 models['gpt3-m']['dmodel'] = 1024
-models['gpt3-m']['dhidden'] = 1024*4
+models['gpt3-m']['dhid'] = 1024*4
 models['gpt3-m']['nlayers'] = 24
+models['gpt3-m']['vocab_size'] = 50257  # from public reimplementations
 
 models['gpt3-l'] = {}
 models['gpt3-l']['seqlen'] = 2048
 models['gpt3-l']['dmodel'] = 1536
-models['gpt3-l']['dhidden'] = 1536*4
+models['gpt3-l']['dhid'] = 1536*4
 models['gpt3-l']['nlayers'] = 24
+models['gpt3-l']['vocab_size'] = 50257  # from public reimplementations
 
 models['gpt3-xl'] = {}
 models['gpt3-xl']['seqlen'] = 2048
 models['gpt3-xl']['dmodel'] = 2560
-models['gpt3-xl']['dhidden'] = 2560*4
+models['gpt3-xl']['dhid'] = 2560*4
 models['gpt3-xl']['nlayers'] = 24
+models['gpt3-xl']['vocab_size'] = 50257  # from public reimplementations
 
 models['gpt3-3b'] = {}
 models['gpt3-3b']['seqlen'] = 2048
 models['gpt3-3b']['dmodel'] = 2560
-models['gpt3-3b']['dhidden'] = 2560*4
+models['gpt3-3b']['dhid'] = 2560*4
 models['gpt3-3b']['nlayers'] = 32
+models['gpt3-3b']['vocab_size'] = 50257  # from public reimplementations
 
 models['gpt3-7b'] = {}
 models['gpt3-7b']['seqlen'] = 2048
 models['gpt3-7b']['dmodel'] = 4096
-models['gpt3-7b']['dhidden'] = 4096*4
+models['gpt3-7b']['dhid'] = 4096*4
 models['gpt3-7b']['nlayers'] = 32
+models['gpt3-7b']['vocab_size'] = 50257  # from public reimplementations
 
 models['gpt3-13b'] = {}
 models['gpt3-13b']['seqlen'] = 2048
 models['gpt3-13b']['dmodel'] = 5120
-models['gpt3-13b']['dhidden'] = 5120*4
+models['gpt3-13b']['dhid'] = 5120*4
 models['gpt3-13b']['nlayers'] = 40
+models['gpt3-13b']['vocab_size'] = 50257  # from public reimplementations
 
 models['gpt3-175b'] = {}
 models['gpt3-175b']['seqlen'] = 2048
 models['gpt3-175b']['dmodel'] = 12288
-models['gpt3-175b']['dhidden'] = 12288*4
+models['gpt3-175b']['dhid'] = 12288*4
 models['gpt3-175b']['nlayers'] = 96
+models['gpt3-175b']['vocab_size'] = 50257  # from public reimplementations