calc v0

.github/workflows/sync_to_hub.yaml

@@ -17,4 +17,4 @@ jobs:
       - name: Push to hub
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
-        run: git push https://training-transformers-together:[email protected]/spaces/training-transformers-together/dashboard-embedded main --force
+        run: git push https://training-transformers-together:[email protected]/spaces/training-transformers-together/calc main --force

README.md

@@ -1,8 +1,8 @@
 ---
-title: Mini-dashboard
+title: Memory calculator
 emoji: ⚡
-colorFrom: gray
-colorTo: gray
+colorFrom: blue
+colorTo: blue
 sdk: streamlit
 app_file: app.py
 pinned: false

app.py

@@ -5,41 +5,47 @@ If you're not a hedgehog, you shouldn't reuse this code. Use this instead: https
 
 import streamlit as st
 
-from dashboard_utils.main_metrics import get_main_metrics
-
-st.set_page_config(page_title="Training Transformers Together - Mini-Dashboard", layout="wide")
+import mem_calc
+from models import models
+st.set_page_config(page_title="Memory calculator", layout="centered")
 st.markdown("""<style>
 .reportview-container {
     top: -80px;
 }
 </style>""", unsafe_allow_html=True)
-source = get_main_metrics()
-st.vega_lite_chart(
-    source, {
-        "height": 200,
-        "title": {"text": "Training DALL-E with volunteers", "dy": 7},
-        # ^-- WARNING: do not use long titles, otherwise vega collapses on small screens
-        "$schema": "https://vega.github.io/schema/vega-lite/v5.json",
-        "description": "Current training progress",
-        "encoding": {"x": {"field": "wall time", "type": "temporal"}},
-        "config": {"axisX": {"labelAngle": -40}},
-        "resolve": {"scale": {"y": "independent"}},
-        "layer": [
-            {
-                "mark": {"type": "line", "point": {"tooltip": True, "filled": False, "strokeOpacity": 0},
-                         "color": "#85A9C5"},
-                "encoding": {
-                    "y": {"field": "training loss", "type": "quantitative", "axis": {"titleColor": "#85A9C5"},
-                          "scale": {"zero": False}}},
-            },
-            {
-                "mark": {"type": "line", "point": {"tooltip": True, "filled": False, "strokeOpacity": 0.0},
-                         "color": "#85C5A6", "opacity": 0.5},
-                "encoding": {
-                    "y": {"field": "active participants", "type": "quantitative",
-                          "axis": {"titleColor": "#85C5A6"}}},
-            },
-        ],
-    },
-    use_container_width=True,  # breaks on <600px screens
-)
+
+models = list(models.keys())  # respect the original order because py37
+model = st.selectbox('Model architecture', models, index=models.index("gpt2-l"))
+
+optimizers_names = ('32-bit', '16-bit', '8-bit', 'factorized')
+optimizers_values = ['adam', '16-bit-adam', '8-bit-adam', 'adafactor']
+optimizer = st.radio('Adam / LAMB states', optimizers_names)
+checkpoint = st.checkbox("Gradient checkpointing", value=True)
+offload = st.checkbox("Offload optimizer", value=False)
+share_params = st.checkbox("Share parameters", value=False)
+
+with st.expander("More options"):
+
+    precisions_names = ('Full', 'Mixed ("O1")', 'Pure 16-bit')
+    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)]}
+    {'--checkpoint' if checkpoint else ''} {'--offload' if offload else ''} {'--albert' if share_params else ''}
+    --fp16-level {precisions_values[precisions_names.index(precision)]}
+""".split())
+
+
+memory = mem_calc.calculate_memory(args)
+
+cols = st.columns(3)
+cols[0].metric("Parameters (GPU)", f"{memory['model']:.2f} GB", f"{memory['model']/memory['total_mem'] * 100:.2f} %", delta_color="off")
+cols[1].metric(f"Optimizer ({'GPU' if offload else 'CPU'})", f"{memory['optim']:.2f} GB", f"{memory['optim']/memory['total_mem'] * 100:.2f} %", delta_color="off")
+cols[2].metric("Activations (GPU)", f"{memory['grad']:.2f} GB", f"{memory['grad']/memory['total_mem'] * 100:.2f} %", delta_color="off")
+cols = st.columns(3)
+cols[0].metric("GPU total", f"{memory['total_mem']:.2f} GB")
+cols[1].metric("Offloaded to RAM", f"{memory['cpu_mem']:.2f} GB")
+cols[2].metric("Communication overhead", f"{memory['overhead'] * 1000:.2f} ms")

dashboard_utils/main_metrics.py

@@ -1,33 +0,0 @@
-import datetime
-
-import streamlit as st
-import pandas as pd
-
-import wandb
-
-from dashboard_utils.time_tracker import _log, simple_time_tracker
-
-WANDB_REPO = "learning-at-home/dalle-hivemind"
-CACHE_TTL = 120  # note: in the text, we claim that this plot is updated every few minutes
-
-
-@st.cache(ttl=CACHE_TTL)
-@simple_time_tracker(_log)
-def get_main_metrics():
-    wandb.login(anonymous="must")
-    api = wandb.Api()
-    runs = api.runs(WANDB_REPO)
-    run = runs[0]
-    history = run.history(keys=["step", "loss", "alive peers", "_timestamp"])
-
-    steps = []
-    losses = []
-    alive_peers = []
-    dates = []
-    for _, row in history.iterrows():
-        steps.append(row["step"])
-        losses.append(row["loss"])
-        alive_peers.append(row["alive peers"])
-        dates.append(datetime.datetime.utcfromtimestamp(row["_timestamp"]))
-
-    return pd.DataFrame({"steps": steps, "training loss": losses, "active participants": alive_peers, "wall time": dates})

dashboard_utils/time_tracker.py

@@ -1,32 +0,0 @@
-from functools import wraps
-from time import time
-
-
-def simple_time_tracker(log_fun):
-    def _simple_time_tracker(fn):
-        @wraps(fn)
-        def wrapped_fn(*args, **kwargs):
-            start_time = time()
-
-            try:
-                result = fn(*args, **kwargs)
-            finally:
-                elapsed_time = time() - start_time
-
-                # log the result
-                log_fun(
-                    {
-                        "function_name": fn.__name__,
-                        "total_time": elapsed_time,
-                    }
-                )
-
-            return result
-
-        return wrapped_fn
-
-    return _simple_time_tracker
-
-
-def _log(message):
-    print("[SimpleTimeTracker] {function_name} {total_time:.3f}".format(**message))

mem_calc.py

@@ -0,0 +1,237 @@
+import argparse
+import math
+from models import models
+
+
+def get_GB(nbytes):
+    return nbytes/(1024**3)
+
+
+def vocab(bsz, seqlen, dmodel, vocab_dim):
+    # assumes tied embeddings
+
+    w = vocab_dim*dmodel
+    emb = seqlen*bsz*dmodel
+    emb_norm = seqlen*bsz*dmodel
+    pos_emb = seqlen*bsz*dmodel
+    out_emb = seqlen*bsz*vocab_dim
+    softmax_emb = seqlen*bsz*vocab_dim
+
+    model = w + dmodel
+    grad = emb + emb_norm + pos_emb + out_emb + softmax_emb
+    grad *= 1
+    return model, grad
+
+
+def transformer(bsz, seqlen, dmodel, nlayers, vocab_type, dhid=None,
+                checkpoint=False, albert=False):
+    if dhid is None: dhid = 4*dmodel
+    model = 0
+    grad = 0
+    for i in range(nlayers):
+        m, g = transformer_layer(bsz, seqlen, dmodel, dhid, checkpoint=checkpoint)
+        model += m
+        grad += g
+
+    if albert:
+        model = model / nlayers
+
+    m, g = vocab(bsz, seqlen, dmodel, vocab_type)
+    model += m
+    grad += g
+
+    return model, grad
+
+def layer_norm(bsz, seqlen, dmodel):
+    w = dmodel
+    x_grad = bsz*seqlen*dmodel
+    return w, x_grad
+
+
+def transformer_layer(bsz, seqlen, dmodel, dhid, checkpoint=False):
+    model = 0
+    grad = 0
+
+    m, g = ffn(bsz, seqlen, dmodel, dhid, 'gelu')
+    model += m
+    grad += g*3
+
+    m, g = attention_layer(bsz, seqlen, dmodel)
+    model += m
+    grad += g*5.0
+
+    m, g = layer_norm(bsz, seqlen, dmodel)
+    model += m
+    grad += g*1.0
+
+    if checkpoint:
+        grad = bsz * seqlen * dmodel
+
+    return model, grad
+
+def attention_layer(bsz, seqlen, dmodel):
+    w_proj = dmodel*3*dmodel
+    w_out = dmodel*dmodel
+
+    x_residual = bsz*seqlen*dmodel
+    x_proj = bsz*seqlen*dmodel*3
+    #x_proj_contiguous = bsz*seqlen*dmodel*3
+    x_proj_contiguous = 0
+
+    x_qscaled = bsz*seqlen*dmodel
+    x_qk = bsz*seqlen*seqlen*2 # we need to store both input sequence directions for gradient computation
+    x_softmax = bsz*seqlen*seqlen
+    x_softmax_v = bsz*seqlen*dmodel*2 # we need to store both input sequence directions for gradient computation
+    #x_out_contiguous = bsz*seqlen*dmodel
+    x_out_contiguous = 0
+    x_out = bsz*seqlen*dmodel
+
+    model = w_proj + w_out
+    grad = x_residual + x_proj + x_proj_contiguous + x_qscaled + x_qk + x_softmax + x_softmax_v + x_out_contiguous + x_out
+    return model, grad
+
+
+
+def ffn(bsz, seqlen, dmodel, dhid, func='relu'):
+    # out = linear(relu(linear(x), inplace=True)) + x
+    w1 = dmodel*dhid
+    w2 = dhid*dmodel
+    model = w1 + w2
+    wgrad = model
+    x1 = bsz*seqlen*dhid
+    if func != 'relu': x1 *= 2 # inplace not possible with most other functions
+    x2 = bsz*seqlen*dmodel
+    residual = bsz*seqlen*dmodel
+    grad = x1 + x2 + residual
+
+    return model, grad
+
+
+OPTIMIZERS = ['adam', 'adafactor', 'adafactor-fac-only', '8-bit-adam', '16-bit-adam']
+
+
+def parse_args(args=None):
+    parser = argparse.ArgumentParser('Memory calculator')
+
+    parser.add_argument('--nlayers', type=int, help='The number of transformer layers.')
+    parser.add_argument('--bsz', type=int, default=1, help='The batch size. Default: 2')
+    parser.add_argument('--seqlen', type=int, help='The sequence length.')
+    parser.add_argument('--dmodel', type=int, help='The core model size.')
+    parser.add_argument('--dhid', type=int, default=None,
+                        help='The hidden size of the FFN layer. Default: 4x model size.')
+    parser.add_argument('--fp16-level', type=str, default='O1',
+                        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('--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,
+                        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('--checkpoint', action='store_true', help='Use gradient checkpointing.')
+
+    return parser.parse_args(args)
+
+
+def calculate_memory(args):
+    if args.model != '':
+        if args.model not in models:
+            raise ValueError(f'{args.model} is not supported')
+        else:
+            for key, value in models[args.model].items():
+                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)
+    parameters = model
+
+    if args.optimizer == 'adam':
+        optim = 8*model
+    elif args.optimizer == '8-bit-adam':
+        optim = 2*model
+    elif args.optimizer in ['16-bit-adam', 'adafactor']:
+        optim = 4*model
+    elif args.optimizer in ['adafactor-fac-only']:
+        optim = math.log(model)
+
+    if args.fp16_level == 'O0':
+        # fp32 weights
+        wgrad = 4*model
+        model = 4*model
+        grad = 4*grad # fp32
+    elif args.fp16_level in ['O1', 'O2']:
+        # fp16 weights + fp32 master weights
+        wgrad = 2*model
+        model = 4*model + (2*model)
+        grad = 2*grad # fp16
+    elif args.fp16_level == 'O3':
+        wgrad = 2*model
+        model = 2*model #fp16
+        grad = 2*grad # fp32
+
+    model = get_GB(model)
+    grad = get_GB(grad)
+    optim = get_GB(optim)
+    wgrad = get_GB(wgrad)
+
+    cpu_mem = 0
+    overhead = 0
+
+    if args.zero == 1:
+        if not args.offload:
+            # assumes PCIe 4.0 infiniband (200 Gbit/s = 25 GB/s)
+            overhead += optim/25
+
+        optim = optim / args.ngpus
+    elif args.zero == 2:
+        if not args.offload:
+            # assumes PCIe 4.0 infiniband (200 Gbit/s = 25 GB/s)
+            overhead += optim/25
+            overhead += wgrad/25
+
+        optim = optim / args.ngpus
+        wgrad = wgrad / args.ngpus
+    elif args.zero == 3:
+        if not args.offload:
+            # assumes PCIe 4.0 infiniband (200 Gbit/s = 25 GB/s)
+            overhead += optim/25
+            overhead += model/25
+            overhead += wgrad/25
+
+        optim = optim / args.ngpus
+        model = model / args.ngpus
+        wgrad = wgrad / args.ngpus
+
+
+    if args.offload:
+        cpu_mem = optim + wgrad
+        optim = 0
+        wgrad = 0
+        if args.ngpus <= 2:
+            # 12 GB/s for PCIe 3.0 and 1-2x GPU setup (16 lanes, 16 GB/s theoretical)
+            overhead = cpu_mem/12
+        else:
+            # 6 GB/s for PCIe 3.0 and 4x GPU setup
+            overhead = cpu_mem/6
+
+
+    total_mem = model + grad + optim + wgrad
+    return locals()
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    mem = calculate_memory(args)
+    print('')
+    print(f'Model: {args.model} with batch size {args.bsz} and sequence length {args.seqlen} and a total of {mem["parameters"]/1e9:.4f}B parameters.')
+    print('='*80)
+    print('Weight memory:           {0:.2f} GB ({1:.2f}%)'.format(mem['model'], 100*mem['model']/mem['total_mem']))
+    print('Weight gradient memory:  {0:.2f} GB ({1:.2f}%)'.format(mem['wgrad'], 100*mem['wgrad']/mem['total_mem']))
+    print('Input gradient memory:   {0:.2f} GB ({1:.2f}%)'.format(mem['grad'], 100*mem['grad']/mem['total_mem']))
+    print('Optimizer memory:        {0:.2f} GB ({1:.2f}%)'.format(mem['optim'], 100*mem['optim']/mem['total_mem']))
+    print('Total GPU memory:        {0:.2f} GB'.format(mem['total_mem']))
+    if mem['cpu_mem'] > 0:
+        print('Total CPU memory:        {0:.2f} GB'.format(mem['cpu_mem']))
+    if mem['overhead'] > 0:
+        print('Overhead: {0:.2f} seconds per update (can be partially overlapped with compute)'.format(mem['overhead']))

models.py

@@ -0,0 +1,97 @@
+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['t5-3b'] = {}
+models['t5-3b']['seqlen'] = 512
+models['t5-3b']['dmodel'] = 1024
+models['t5-3b']['dhidden'] = 16384
+models['t5-3b']['nlayers'] = 48
+
+models['t5-11b'] = {}
+models['t5-11b']['seqlen'] = 512
+models['t5-11b']['dmodel'] = 1024
+models['t5-11b']['dhidden'] = 64*1024
+models['t5-11b']['nlayers'] = 48
+
+models['gpt2-s'] = {}
+models['gpt2-s']['seqlen'] = 1024
+models['gpt2-s']['dmodel'] = 768
+models['gpt2-s']['dhidden'] = 768*4
+models['gpt2-s']['nlayers'] = 12
+
+models['gpt2-m'] = {}
+models['gpt2-m']['seqlen'] = 1024
+models['gpt2-m']['dmodel'] = 1024
+models['gpt2-m']['dhidden'] = 1024*4
+models['gpt2-m']['nlayers'] = 24
+
+models['gpt2-l'] = {}
+models['gpt2-l']['seqlen'] = 1024
+models['gpt2-l']['dmodel'] = 1280
+models['gpt2-l']['dhidden'] = 1280*4
+models['gpt2-l']['nlayers'] = 36
+
+models['gpt2-xl'] = {}
+models['gpt2-xl']['seqlen'] = 1024
+models['gpt2-xl']['dmodel'] = 1600
+models['gpt2-xl']['dhidden'] = 1600*4
+models['gpt2-xl']['nlayers'] = 48
+
+
+models['gpt3-s'] = {}
+models['gpt3-s']['seqlen'] = 2048
+models['gpt3-s']['dmodel'] = 768
+models['gpt3-s']['dhidden'] = 768*4
+models['gpt3-s']['nlayers'] = 12
+
+models['gpt3-m'] = {}
+models['gpt3-m']['seqlen'] = 2048
+models['gpt3-m']['dmodel'] = 1024
+models['gpt3-m']['dhidden'] = 1024*4
+models['gpt3-m']['nlayers'] = 24
+
+models['gpt3-l'] = {}
+models['gpt3-l']['seqlen'] = 2048
+models['gpt3-l']['dmodel'] = 1536
+models['gpt3-l']['dhidden'] = 1536*4
+models['gpt3-l']['nlayers'] = 24
+
+models['gpt3-xl'] = {}
+models['gpt3-xl']['seqlen'] = 2048
+models['gpt3-xl']['dmodel'] = 2560
+models['gpt3-xl']['dhidden'] = 2560*4
+models['gpt3-xl']['nlayers'] = 24
+
+models['gpt3-3b'] = {}
+models['gpt3-3b']['seqlen'] = 2048
+models['gpt3-3b']['dmodel'] = 2560
+models['gpt3-3b']['dhidden'] = 2560*4
+models['gpt3-3b']['nlayers'] = 32
+
+models['gpt3-7b'] = {}
+models['gpt3-7b']['seqlen'] = 2048
+models['gpt3-7b']['dmodel'] = 4096
+models['gpt3-7b']['dhidden'] = 4096*4
+models['gpt3-7b']['nlayers'] = 32
+
+models['gpt3-13b'] = {}
+models['gpt3-13b']['seqlen'] = 2048
+models['gpt3-13b']['dmodel'] = 5120
+models['gpt3-13b']['dhidden'] = 5120*4
+models['gpt3-13b']['nlayers'] = 40
+
+models['gpt3-175b'] = {}
+models['gpt3-175b']['seqlen'] = 2048
+models['gpt3-175b']['dmodel'] = 12288
+models['gpt3-175b']['dhidden'] = 12288*4
+models['gpt3-175b']['nlayers'] = 96