sharing groups

app.py

@@ -27,15 +27,17 @@ share_params = col2.checkbox("Share parameters", value=False)
 
 with st.expander("More options"):
     batch_size = int(st.number_input('Microbatch size (sequences)', min_value=1, step=1, value=1, format="%i"))
-    seq_len = int(st.number_input('Sequence length (max. tokens)', min_value=1, step=1, value=1024, format="%i"))
     precisions_names = ('Full', 'Mixed ("O1")', 'Pure 16-bit')
     precisions_values = ('O0', 'O1', 'O3')
+    sharing_groups = int(st.number_input('Shared parameter groups (used if Share parameters is checked)',
+                                         min_value=1, step=1, value=1, format="%i"))
     precision = st.selectbox('Precision', precisions_names, index=1)
 
 args = mem_calc.parse_args(f"""
     --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}
+    {'--checkpoint' if checkpoint else ''} {'--offload' if offload else ''}
+    --fp16-level {precisions_values[precisions_names.index(precision)]} --bsz {batch_size}
+    {f'--shared_groups {sharing_groups}' if share_params else ''} 
 """.split())
 
 

mem_calc.py

@@ -24,7 +24,7 @@ def vocab(bsz, seqlen, dmodel, vocab_dim):
 
 
 def transformer(bsz, seqlen, dmodel, nlayers, vocab_type, dhid=None,
-                checkpoint=False, albert=False):
+                checkpoint=False, shared_groups=None):
     if dhid is None: dhid = 4*dmodel
     model = 0
     grad = 0
@@ -33,8 +33,8 @@ def transformer(bsz, seqlen, dmodel, nlayers, vocab_type, dhid=None,
         model += m
         grad += g
 
-    if albert:
-        model = model / nlayers
+    if shared_groups is not None:
+        model = model / nlayers * shared_groups
 
     m, g = vocab(bsz, seqlen, dmodel, vocab_type)
     model += m
@@ -128,7 +128,7 @@ def parse_args(args=None):
     parser.add_argument('--ngpus', type=int, default=1, help='The number of gpus. Default: 1')
     parser.add_argument('--zero', type=int, default=0,
                         help='The ZeRO level (1 optimizer, 2 optimizer+weights, 3 everything. Default: 1')
-    parser.add_argument('--albert', action='store_true', help='Use parameter sharing.')
+    parser.add_argument('--shared_groups', type=int, default=None, help='Number of shared layer groups (as in ALBERT). Defaults to no sharing.')
     parser.add_argument('--checkpoint', action='store_true', help='Use gradient checkpointing.')
 
     return parser.parse_args(args)
@@ -143,7 +143,7 @@ def calculate_memory(args):
                 if getattr(args, key, None) is None:
                     setattr(args, key, value)
 
-    model, grad = transformer(args.bsz, args.seqlen, args.dmodel, args.nlayers, args.vocab_size, args.dhid, args.checkpoint, args.albert)
+    model, grad = transformer(args.bsz, args.seqlen, args.dmodel, args.nlayers, args.vocab_size, args.dhid, args.checkpoint, args.shared_groups)
     parameters = model
 
     if args.optimizer == 'adam':

models.py

@@ -56,13 +56,6 @@ 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
@@ -118,3 +111,16 @@ models['gpt3-175b']['dmodel'] = 12288
 models['gpt3-175b']['dhid'] = 12288*4
 models['gpt3-175b']['nlayers'] = 96
 models['gpt3-175b']['vocab_size'] = 50257  # from public reimplementations
+
+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['dalle-12b'] = {}
+models['dalle-12b']['seqlen'] = 1024 + 256
+models['dalle-12b']['dmodel'] = 62 * 64
+models['dalle-12b']['nlayers'] = 64
+models['dalle-12b']['vocab_size'] = 8192 + 16384