import from banana

app_to_share/optimal_training/app_requirements.txt

@@ -0,0 +1,17 @@
+bokeh==2.0.2
+Jinja2==2.11.2
+MarkupSafe==1.1.1
+numpy==1.18.4
+packaging==20.4
+pandas==1.0.3
+Pillow==7.1.2
+pkg-resources==0.0.0
+pyparsing==2.4.7
+python-dateutil==2.8.1
+pytz==2020.1
+PyYAML==5.3.1
+randomcolor==0.4.4.5
+scipy==1.4.1
+six==1.15.0
+tornado==6.0.4
+typing-extensions==3.7.4.2

app_to_share/optimal_training/bokeh_test.py

@@ -0,0 +1,99 @@
+''' Present an interactive function explorer with slider widgets.
+Scrub the sliders to change the properties of the ``sin`` curve, or
+type into the title text box to update the title of the plot.
+Use the ``bokeh serve`` command to run the example by executing:
+    bokeh serve sliders.py
+at your command prompt. Then navigate to the URL
+    http://localhost:5006/sliders
+in your browser.
+'''
+import numpy as np
+
+from bokeh.io import curdoc
+from bokeh.layouts import column, row
+from bokeh.models import ColumnDataSource, Slider, TextInput
+from bokeh.plotting import figure
+
+# Set up data
+N = 200
+x = np.linspace(0, 4 * np.pi, N)
+y = np.sin(x)
+source = ColumnDataSource(data=dict(x=x, y=y))
+
+# Set up plot
+plot = figure(plot_height=400, plot_width=400, title="my sine wave",
+              tools="crosshair,pan,reset,save,wheel_zoom",
+              x_range=[0, 4 * np.pi], y_range=[-2.5, 2.5])
+
+plot.line('x', 'y', source=source, line_width=3, line_alpha=0.6)
+
+# Set up widgets
+text = TextInput(title="title", value='my sine wave')
+offset = Slider(title="offset", value=0.0, start=-5.0, end=5.0, step=0.1)
+amplitude = Slider(title="amplitude", value=1.0, start=-5.0, end=5.0, step=0.1)
+phase = Slider(title="phase", value=0.0, start=0.0, end=2 * np.pi)
+freq = Slider(title="frequency", value=1.0, start=0.1, end=5.1, step=0.1)
+slider_moves = {"offset": 0, "amplitude": 0, "phase": 0, "freq": 0}
+
+
+# Set up callbacks
+def update_title(attrname, old, new):
+    plot.title.text = text.value
+
+
+text.on_change('value', update_title)
+
+
+def update_data(attrname, old, new):
+    # Get the current slider values
+    a = amplitude.value
+    b = offset.value
+    w = phase.value
+    k = freq.value
+
+    # Generate the new curve
+    x = np.linspace(0, 4 * np.pi, N)
+    y = a * np.sin(k * x + w) + b
+
+    source.data = dict(x=x, y=y)
+
+
+def offset_force(attrname, old, new):
+    slider_moves["offset"] += 1
+
+    if slider_moves["amplitude"] < slider_moves["offset"]:
+        a = amplitude.value = offset.value
+        w = phase.value = offset.value
+        k = freq.value = offset.value
+        b = offset.value
+        x = np.linspace(0, 4 * np.pi, N)
+        y = a * np.sin(k * x + w) + b
+
+        source.data = dict(x=x, y=y)
+
+
+def amp_force(attrname, old, new):
+    slider_moves["amplitude"] += 1
+
+    if slider_moves["offset"] < slider_moves["amplitude"]:
+        b = offset.value = amplitude.value * 2
+        w = phase.value = amplitude.value * 2
+        k = freq.value = amplitude.value * 2
+        a = amplitude.value
+        x = np.linspace(0, 4 * np.pi, N)
+        y = a * np.sin(k * x + w) + b
+
+        source.data = dict(x=x, y=y)
+
+
+for w in [phase, freq]:
+    w.on_change('value', update_data)
+
+offset.on_change('value', offset_force)
+amplitude.on_change('value', amp_force)
+
+# Set up layouts and add to document
+inputs = column(text, offset, amplitude, phase, freq)
+
+curdoc().add_root(row(inputs, plot, width=800))
+curdoc().title = "Sliders"

app_to_share/optimal_training/conversions.py

@@ -0,0 +1,230 @@
+import math
+import numpy as np
+from scipy.optimize import root
+
+day_ratio = 24 * 3600
+
+depth_width_ratio = 128
+
+constants_per_gpu = {
+    "V100": [2.21527743e+07, 1.18538628e+00, 1.43150104e+00, 1.66015023e+00,
+             1.32808220e+00, 5.91503856e+00],
+    "V100 (without tensor cores and cudnn.benchmark)": [1.82997989e+07, 1.05349588e+00, 1.25312127e+00, 1.67071294e+00,
+                                                        1.44610885e+00, 5.55824273e+00],
+    "P100": [6.01863899e+07, 9.23656025e-01, 1.03230702e+00, 1.46733667e+00,
+             1.03031298e+00, 5.38021875e+00],
+    "P4": [4.84472202e+07, 9.86822195e-01, 1.23474901e+00, 1.38493518e+00,
+           1.04630858e+00, 1.03572754e+01],
+    "K80": [2.58592374e+07, 6.42050890e-01, 7.06115162e-01, 1.44360777e+00,
+            7.50695980e-01, 6.25951436e+00]
+
+}
+
+price_per_gpu = {
+    "K80": 0.584,
+    "P4": 0.689,
+    "V100": 2.005,
+    "V100 (without tensor cores and cudnn.benchmark)": 2.005,
+    "P100": 1.416,
+}
+
+optimal_batch_size_per_gpu = {
+    "P4": 16,
+    "V100": 64,
+    "V100 (without tensor cores and cudnn.benchmark)": 64,
+    "P100": 64,
+    "K80": 16
+}
+
+features_per_amp_mode = {
+    "O0": (1, 0, 0),
+    "O1": (0, 1, 0),
+    "O2": (0, 0, 1)
+}
+
+gpu_consumption = {
+    "V100": 119.3495934959e-3,
+    "V100 (without tensor cores and cudnn.benchmark)": 119.3495934959e-3,
+    "K80": 142.42e-3,
+    "P4": 55.27e-3,
+    "P100": 139.65e-3
+}
+
+co2_intensity = 534 * 1e-3
+
+
+def flo_speed(features, constants):
+    k, k1, k2, b, c, layer_base = constants
+    o0, o1, o2, x, y, z = features
+    return k * np.power(k1, o1) * np.power(k2, o2) * x / (x + layer_base) * np.power(y, b) * np.power(np.log(z + 1), c)
+
+
+def param_polynomial(width, depth=None, inner=None):
+    if depth is not None:
+        if inner is not None:
+            return 5 * depth * (width ** 2) + 2 * depth * (width * inner) + 7 * depth * width + depth * inner + 3 * width + 3
+        else:
+            return 7 * depth * (width ** 2) + 8 * depth * width + 3 * width + 3
+    else:
+        if inner is not None:
+            return 5 * depth_width_ratio * (width ** 3) + 2 * depth_width_ratio * (width ** 2 * inner) + 7 * depth_width_ratio * width ** 2 + depth_width_ratio * width * inner + 3 * width + 3
+        else:
+            return 7 / depth_width_ratio * (width ** 3) + 8 / depth_width_ratio * (width ** 2) + 3 * width + 3
+
+
+def optimal_model_shape(width, param_number, base=8):
+    depth = max(1, math.floor(width / depth_width_ratio))
+    poly_params = np.array([depth * 7, depth * 8 + 3, 3 - param_number])
+    roots = np.roots(poly_params)
+    corresponding_width = int(base * round(max(roots) / base))
+    return depth, corresponding_width
+
+
+def alternate_model_shape(width, param_number, base=8):
+    linear_depth = max(1, math.floor(width / depth_width_ratio))
+    depth = max(linear_depth + 1, math.floor(0.3 * width ** 1.25 / depth_width_ratio))
+    poly_params = np.array([depth * 7, depth * 8 + 3, 3 - param_number])
+    roots = np.roots(poly_params)
+    corresponding_width = int(base * round(max(roots) / base))
+    return depth, corresponding_width
+
+
+def hours_to_width(hours, gpu, amp_mode, param_popt):
+    seconds = hours * 3600
+    d, e, f = param_popt
+    constants = constants_per_gpu[gpu]
+    amp_features = features_per_amp_mode[amp_mode]
+
+    def equation_function(width):
+        return np.power((param_polynomial(width) - f) / d, 1 / e) / flo_speed(
+            (*amp_features, width / depth_width_ratio, width, optimal_batch_size_per_gpu[gpu]),
+            constants) * day_ratio - seconds
+
+    width = iterative_solutions(equation_function, initial_guess=128)
+    # print("width: {}".format(math.floor(width)))
+    # print("depth: {}".format(width / depth_width_ratio))
+    # print("param number: {:.4e}".format(param_polynomial(width)))
+    speed = flo_speed((*amp_features, width / depth_width_ratio, width, optimal_batch_size_per_gpu[gpu]), constants)
+    # print("speed: {:.4e}".format(speed))
+    # print("flos from speed: {:.4e}".format(seconds * speed))
+    # print("flos from params: {:.4e}".format(np.power((param_polynomial(width) - f) / d, 1 / e) * day_ratio))
+    # print("params from flos: {:.4e}".format(np.exp(param_fit(speed * seconds / day_ratio, *param_popt))))
+    return width
+
+
+def iterative_solutions(equation_function, initial_guess):
+    while initial_guess > 16:
+        solution_array = root(equation_function, np.array([initial_guess]), method="hybr").x
+        width = solution_array[0]
+        should_be_zero = equation_function(width)
+        if np.abs(should_be_zero) < 1e0:
+            return width
+        else:
+            initial_guess *= 0.5
+    return width
+
+
+def width_to_flo(width, d, e, f):
+    return np.power((param_polynomial(width) - f) / d, 1 / e) * day_ratio
+
+
+def loss_fit(x, a, b, c):
+    return a * np.power(x, -b) + c
+
+
+def param_fit(x, d, e, f):
+    return np.log(d * np.power(x, e) + f)
+
+
+def hours_to_dollars(hours, gpu):
+    return hours * price_per_gpu[gpu]
+
+
+def dollars_to_hours(dollars, gpu):
+    return dollars / price_per_gpu[gpu]
+
+
+def hours_to_kWh(hours, gpu):
+    return hours * gpu_consumption[gpu]
+
+
+def hours_to_co2(hours, gpu):
+    return hours * gpu_consumption[gpu] * co2_intensity
+
+
+def loss_to_flo(loss, a, b, c):
+    return ((loss - c) / a) ** (-1 / b)
+
+
+def param_to_flo(param_number, d, e, f):
+    return ((param_number - f) / d) ** (1 / e)
+
+
+def safe_flo_to_param(flo, d, e, f):
+    return d * np.power(flo, e) + f
+
+
+def param_to_width(param_number):
+    poly_params = np.array([7 / depth_width_ratio, 8 / depth_width_ratio, 3, 3 - param_number])
+    roots = np.roots(poly_params)
+    real_roots = [np.real(candidate) for candidate in roots if np.imag(candidate) < 1e-5]
+    width = max(real_roots)
+    return width
+
+
+def safe_param_to_width(param_number):
+    try:
+        return param_to_width(param_number)
+    except np.linalg.LinAlgError:
+        return safe_param_to_width(1.5 * param_number)
+
+
+def width_to_hours(width, gpu, amp_mode, param_popt):
+    d, e, f = param_popt
+    constants = constants_per_gpu[gpu]
+    amp_features = features_per_amp_mode[amp_mode]
+    flos_from_params = np.power((param_polynomial(width) - f) / d, 1 / e) * day_ratio
+    speed = flo_speed((*amp_features, width / depth_width_ratio, width, optimal_batch_size_per_gpu[gpu]), constants)
+    seconds = flos_from_params / speed
+
+    return seconds / 3600
+
+
+def param_prime(width, depth=None):
+    if depth is not None:
+        return 14 * depth * (width ** 2) + 8 * depth + 3
+    else:
+        return 21 / depth_width_ratio * (width ** 2) + 16 / depth_width_ratio * width + 3
+
+
+def flo_speed_prime(width, gpu, amp_mode):
+    k, k1, k2, b, c, layer_base = constants_per_gpu[gpu]
+    o0, o1, o2 = features_per_amp_mode[amp_mode]
+    mult_constant = k * np.power(k1, o1) * np.power(k2, o2) * np.power(np.log(optimal_batch_size_per_gpu[gpu] + 1), c)
+    return mult_constant * ((b + 1) * np.power(width, b) / (width + layer_base * depth_width_ratio)
+                            - np.power(width, b + 1) / (width + layer_base * depth_width_ratio) ** 2)
+
+
+# awful equation; we're trying to find the width for which lowering width actually makes the model less efficient
+def tipping_point(gpu, amp_mode, param_popt):
+    d, e, f = param_popt
+    o0, o1, o2 = features_per_amp_mode[amp_mode]
+
+    def equation_function(width):
+        return np.power((param_polynomial(width) - f) / d, -1) / e * param_prime(width) / d \
+               * flo_speed((o0, o1, o2, width / depth_width_ratio, width, optimal_batch_size_per_gpu[gpu]),
+                           constants_per_gpu[gpu]) - \
+               flo_speed_prime(width, gpu, amp_mode)
+
+    tipping_width = iterative_solutions(equation_function, initial_guess=100)
+    return tipping_width
+
+
+def update_tip(tip, width, gpu, amp_mode, loss_popt, param_popt):
+    a, b, c = loss_popt
+    d, e, f = param_popt
+    tip["width"] = width
+    tip["param_number"] = param_polynomial(width)
+    tip["flo"] = np.power((param_polynomial(tip["param_number"]) - f) / d, 1 / e)
+    tip["loss"] = loss_fit(tip["flo"], a, b, c)
+    tip["hours"] = width_to_hours(width, gpu, amp_mode, param_popt)

app_to_share/optimal_training/main.py

@@ -0,0 +1,519 @@
+from bokeh.io import curdoc
+from bokeh.layouts import column, row
+from bokeh.models import Slider, Select, ColumnDataSource, Span, Div, Button, LogColorMapper, ColorBar, LogTicker
+from bokeh.models.tools import CrosshairTool
+from bokeh.plotting import figure
+from bokeh.events import Tap
+from bokeh.transform import log_cmap
+import pandas as pd
+from scipy.spatial import ConvexHull
+from scipy.optimize import curve_fit
+from time import sleep
+
+from utils import *
+from conversions import *
+
+########################################################################################################################
+# Basic dimensions
+########################################################################################################################
+
+plot_width = 1200
+plot_height = 400
+sidebar_width = 400
+in_text_plot_width = 800
+in_text_plot_height = 300
+
+########################################################################################################################
+# Set up data
+########################################################################################################################
+
+df = pd.read_csv("optimal_training/static/loss_vs_compute.csv")
+loss_keys = [key for key in df.keys() if "loss" in key]
+
+losses_per_run = {key: np.array(clean_run(list(zip(df["global_step"], df[key])))) for key in loss_keys}
+losses_per_run = {k: v for k, v in losses_per_run.items() if len(v) > 5}
+bounds_per_run = {key: [min(value[:, 0]), max(value[:, 0])] for key, value in losses_per_run.items()}
+params_per_run = {key: param_count(run) for key, run in losses_per_run.items()}
+ordered_keys = sorted(losses_per_run, key=lambda x: params_per_run[x])
+losses_per_run = [losses_per_run[key] for key in ordered_keys]
+bounds_per_run = [bounds_per_run[key] for key in ordered_keys]
+params_per_run = [params_per_run[key] for key in ordered_keys]
+palette = "Viridis256"
+color_mapper = LogColorMapper(palette=palette, low=min(params_per_run), high=max(params_per_run))
+general_bounds = bounds_per_run[2][0], bounds_per_run[-2][1]
+print("{:.4e}, {:.4e}".format(general_bounds[0] * day_ratio, general_bounds[1] * day_ratio))
+color_list = ["#000000" in params_per_run]
+# there's a bogus point of small coordinates at position 0 to get the ConvexHull facing the origin
+# hacky, but it's the syntax here, qhull_options=QG0 means the ConvexHull facing point 0
+bounded_points = np.array([(10e8, 3, -1)] + [(a, b, i) for i, run in enumerate(losses_per_run) for a, b in run if
+                                             general_bounds[0] < a < general_bounds[1]])
+all_points = np.array([(a, b, i) for i, run in enumerate(losses_per_run) for a, b in run])
+all_hull = ConvexHull(bounded_points[:, :2], qhull_options='QG0')
+log_points = np.array([(np.log(a), b) for a, b, i in bounded_points])
+log_hull = ConvexHull(log_points, qhull_options='QG0')
+indexed_runs = [np.array([(a, b) for a, b in run]) for run in losses_per_run]
+
+########################################################################################################################
+# Set up loss_plot
+########################################################################################################################
+
+color_bar = ColorBar(color_mapper=color_mapper, ticker=LogTicker(), label_standoff=12,
+                     border_line_color=None, location=(0, 0), title="Num of params")
+loss_plot = figure(plot_height=plot_height, plot_width=plot_width,
+                   title="Validation loss during training for an array of models of different sizes",
+                   tools="pan,reset,save,wheel_zoom,tap", active_scroll="wheel_zoom",
+                   x_range=[min(all_points[:, 0]) * day_ratio, max(all_points[:, 0]) * day_ratio],
+                   y_range=[min(all_points[:, 1]), max(all_points[:, 1])],
+                   x_axis_type="log", y_axis_type="log",
+                   x_axis_label="Floating-point operations (excluding embeddings & softmax)",
+                   y_axis_label="Validation loss on Wikitext-103", output_backend="webgl")
+loss_plot.add_tools(CrosshairTool(dimensions="width", line_alpha=0.2))
+loss_plot.add_layout(color_bar, "left")
+# for i, run in indexed_runs.items():
+#     source = ColumnDataSource(data=dict(x=run[:, 0] * day_ratio, y=run[:, 1]))
+#     loss_plot.line('x', 'y', source=source, line_width=1, line_alpha=0.6, color=color_list[i])
+#     loss_plot.scatter('x', 'y', source=source, line_width=1, line_alpha=0.6, color=color_list[i])
+
+source = ColumnDataSource(data=dict(
+    xs=[run[:, 0] * day_ratio for run in indexed_runs],  # x coords for each line (list of lists)
+    ys=[run[:, 1] for run in indexed_runs],  # y coords for each line (list of lists)
+    params=params_per_run  # data to use for colormapping
+))
+loss_plot.multi_line('xs', 'ys', source=source,
+                     color=log_cmap('params', palette, min(params_per_run), max(params_per_run)))
+source = ColumnDataSource(data=dict(
+    x=[compute for run in indexed_runs for compute in run[:, 0] * day_ratio],  # x coords for each line (list of lists)
+    y=[loss for run in indexed_runs for loss in run[:, 1]],  # y coords for each line (list of lists)
+    params=[repeated_params for i, params in enumerate(params_per_run)
+            for repeated_params in [params] * len(indexed_runs[i])]  # data to use for colormapping
+))
+loss_plot.scatter('x', 'y', source=source,
+                  color=log_cmap('params', palette, min(params_per_run), max(params_per_run)), size=3)
+
+hull_indices = set(index for pair in all_hull.simplices[all_hull.good] for index in pair)
+hull_indices = sorted(hull_indices, key=lambda x: bounded_points[x, 0])
+
+########################################################################################################################
+# Fit frontier
+########################################################################################################################
+
+hull_points = np.array([bounded_points[index] for index in hull_indices])
+loss_popt, loss_pcov = curve_fit(loss_fit, hull_points[:, 0], hull_points[:, 1])
+a, b, c = loss_popt
+print(a, b, c)
+display_abscisses = np.array([min(all_points[:, 0]) / 1.25] + sorted(list(all_points[:, 0])) +
+                             [max(all_points[:, 0]) * 1.25])
+source = ColumnDataSource(
+    data=dict(x=sorted(display_abscisses * day_ratio), y=loss_fit(sorted(display_abscisses), *loss_popt)))
+loss_plot.line('x', 'y', source=source, line_width=1, line_alpha=0.8, color="red")
+
+########################################################################################################################
+# Set up param_plot
+########################################################################################################################
+
+param_plot = figure(plot_height=plot_height, plot_width=plot_width,
+                    title="Optimal number of non-embedding parameters per floating-point operations budget",
+                    tools="pan,reset,save,wheel_zoom,tap", active_scroll="wheel_zoom",
+                    x_range=loss_plot.x_range,
+                    y_range=[min(params_per_run), max(params_per_run)],
+                    x_axis_type="log", y_axis_type="log",
+                    x_axis_label="Floating-point operations (excluding embeddings & softmax)",
+                    y_axis_label="Optimal number of non-embedding parameters", output_backend="webgl")
+param_plot.add_tools(CrosshairTool(dimensions="width", line_alpha=0.2))
+param_plot.add_layout(color_bar, "left")
+
+logspace_points = convert_to_logspace(bounded_points, *loss_popt)
+logspace_losses_per_run = [convert_to_logspace(run, *loss_popt) for run in losses_per_run]
+passing_points = []
+for run_index, log_run in enumerate(logspace_losses_per_run):
+    current_point = None
+    passed = False
+    difference = log_run[:, 1] - log_run[:, 0]
+    passing_points.append(np.argmax(difference))
+compute_at_passing_points = np.array([(losses_per_run[i][passing_point, 0], params_per_run[i])
+                                      for i, passing_point in enumerate(passing_points)])
+compute_at_hull = np.array([(losses_per_run[i][passing_point, 0], params_per_run[i])
+                            for i, passing_point in enumerate(passing_points) if i in set(hull_points[:, 2])])
+run_indices_at_hull = [i for i, passing_point in enumerate(passing_points) if i in set(hull_points[:, 2])]
+
+param_popt, param_pcov = curve_fit(param_fit, compute_at_hull[:, 0], np.log(compute_at_hull[:, 1]))
+d, e, f = param_popt
+
+source = ColumnDataSource(data=dict(x=compute_at_hull[:, 0] * day_ratio,
+                                    y=compute_at_hull[:, 1],
+                                    params=[params for i, params in enumerate(params_per_run) if
+                                            i in set(hull_points[:, 2])]))
+param_plot.scatter('x', 'y', source=source,
+                   color=log_cmap('params', palette, min(params_per_run), max(params_per_run)))
+display_abscisses = np.array([min(compute_at_hull[:, 0]) / 1.25] + sorted(list(compute_at_hull[:, 0])) +
+                             [max(compute_at_hull[:, 0]) * 1.25])
+source = ColumnDataSource(data=dict(x=display_abscisses * day_ratio,
+                                    y=safe_flo_to_param(display_abscisses, d, e, f)))
+param_plot.line('x', 'y', source=source, line_width=1, line_alpha=0.8, color="orange")
+
+########################################################################################################################
+# Set up widgets
+########################################################################################################################
+
+hours_end = 24
+hours_initial = 3.23
+gpu_dropdown = Select(title="GPU",
+                      options=["V100", "P100", "P4", "K80", ],
+                      value="V100", width=sidebar_width, sizing_mode="stretch_width")
+amp_mode_dropdown = Select(title="AMP mode", options=["O0", "O1", "O2"], value="O0", width=sidebar_width,
+                           sizing_mode="stretch_width")
+tipping_width = tipping_point(gpu_dropdown.value, amp_mode_dropdown.value, param_popt)
+tip = {}
+update_tip(tip, tipping_width, gpu_dropdown.value, amp_mode_dropdown.value, loss_popt, param_popt)
+hours_slider = Slider(title="Wall time (hours)", value=hours_initial, start=tip["hours"], end=hours_end, step=1 / 100,
+                      width=sidebar_width, sizing_mode="stretch_width")
+dollars_slider = Slider(title="Budget (dollars)", value=hours_to_dollars(hours_initial, gpu_dropdown.value),
+                        start=dollars_to_hours(tip["hours"], gpu_dropdown.value),
+                        end=hours_to_dollars(hours_end, gpu_dropdown.value),
+                        step=1 / 100, width=sidebar_width, sizing_mode="stretch_width")
+input_buffer = Div(text="", width=sidebar_width, height=10,
+                   style={"display": "block", "margin": "0 auto", "width": f"{sidebar_width}px",
+                          "text-align": 'center'})
+top_sidebar_div_style = {"display": "block", "margin": "0 auto", 'font-size': "125%",
+                         "width": f"{sidebar_width}px", "text-align": 'center'}
+energy_text = Div(text=energy_fill(hours_to_kWh(hours_slider.value, gpu_dropdown.value),
+                                   hours_to_co2(hours_slider.value, gpu_dropdown.value)),
+                  width=sidebar_width, height=45,
+                  style=top_sidebar_div_style)
+slider_moves = {"hours": 0, "dollars": 0, "kWh": 0, "co2": 0}
+n_sliders = len(slider_moves)
+
+width = hours_to_width(hours_slider.value, gpu_dropdown.value, amp_mode_dropdown.value, param_popt)
+flo = width_to_flo(width, *param_popt)
+optimal_params = safe_flo_to_param(flo / 24 / 3600, *param_popt)
+final_loss = loss_fit(flo / 24 / 3600, *loss_popt)
+example_shape = {}
+example_shape['example_depth'], example_shape['example_width'] = optimal_model_shape(width, optimal_params)
+example_shape['alternate_depth'], example_shape['alternate_width'] = alternate_model_shape(width, optimal_params)
+
+flo_line = Span(location=flo, line_alpha=0.7,
+                dimension='height', line_color='purple',
+                line_dash='dashed', line_width=1)
+loss_line = Span(location=final_loss, line_alpha=0.7,
+                 dimension='width', line_color='red',
+                 line_dash='dashed', line_width=1)
+param_line = Span(location=optimal_params, line_alpha=0.7,
+                  dimension='width', line_color='orange',
+                  line_dash='dashed', line_width=1)
+loss_plot.add_layout(flo_line)
+loss_plot.add_layout(loss_line)
+param_plot.add_layout(flo_line)
+param_plot.add_layout(param_line)
+
+sidebar_div_style = {"display": "block", "margin": "0 auto", "width": f"{sidebar_width}px", "text-align": 'center'}
+big_sidebar_div_style = {"display": "block", "margin": "0 auto", "width": f"{sidebar_width}px",
+                         "text-align": 'center', 'font-size': "200%", 'font-weight': "bold"}
+static_loss_text = Div(text="Expected wt-103 validation loss:", width=sidebar_width, height=10, style=sidebar_div_style)
+optimal_loss_text = Div(text="{:.2f}".format(final_loss), width=sidebar_width, height=45,
+                        style={"display": "block", "margin": "0 auto", 'font-size': "200%",
+                               'font-weight': "bold", "width": f"{sidebar_width}px", "text-align": 'center'})
+static_param_text = Div(text="Optimal number of non-embedding parameters:", width=sidebar_width, height=10,
+                        style=sidebar_div_style)
+optimal_param_text = Div(text="{:.2e}".format(optimal_params), width=sidebar_width, height=45,
+                         style=big_sidebar_div_style)
+static_shape_text = Div(text="For example, this could be a model of", width=sidebar_width, height=10,
+                        style=sidebar_div_style)
+optimal_shape_text = Div(text=f"{example_shape['example_depth']} layers of {example_shape['example_width']} dimensions",
+                         width=sidebar_width, height=30, style=big_sidebar_div_style)
+static_altshape_text = Div(text="Or a model of", width=sidebar_width, height=10, style=sidebar_div_style)
+optimal_altshape_text = Div(
+    text=f"{example_shape['alternate_depth']} layers of {example_shape['alternate_width']} dimensions",
+    width=sidebar_width, height=30, style=big_sidebar_div_style)
+
+
+def compare_and_update(width):
+    if width >= tip["width"]:
+        update_width(width)
+        hours = width_to_hours(width, gpu_dropdown.value, amp_mode_dropdown.value, param_popt)
+        hours_slider.value = hours
+    else:
+        width = min(tip["width"], width + 5)
+        update_width(width)
+        compare_and_update(width)
+
+
+def update_width(width):
+    flo = width_to_flo(width, *param_popt)
+    flo_line.location = flo
+    optimal_params = safe_flo_to_param(flo / 24 / 3600, *param_popt)
+    final_loss = loss_fit(flo / 24 / 3600, *loss_popt)
+    loss_line.location = final_loss
+    param_line.location = optimal_params
+    example_shape['example_depth'], example_shape['example_width'] = optimal_model_shape(width, optimal_params)
+    example_shape['alternate_depth'], example_shape['alternate_width'] = alternate_model_shape(width, optimal_params)
+    optimal_shape_text.text = f"{example_shape['example_depth']} layers of {example_shape['example_width']} dimensions"
+    optimal_altshape_text.text = f"{example_shape['alternate_depth']} layers of {example_shape['alternate_width']} dimensions"
+    optimal_param_text.text = "{:.2e}".format(optimal_params)
+    optimal_loss_text.text = "{:.2f}".format(final_loss)
+
+
+def hours_update(attrname, old, new):
+    slider_moves["hours"] += 1
+
+    # if hours was the first updated slider
+    if sum(slider_moves.values()) <= n_sliders * slider_moves["hours"] - n_sliders + 1:
+        dollars_slider.value = hours_to_dollars(hours_slider.value, gpu_dropdown.value)
+        energy_text.text = energy_fill(hours_to_kWh(hours_slider.value, gpu_dropdown.value),
+                                           hours_to_co2(hours_slider.value, gpu_dropdown.value))
+
+    width = hours_to_width(hours_slider.value, gpu_dropdown.value, amp_mode_dropdown.value, param_popt)
+    update_width(width)
+
+
+def dollars_update(attrname, old, new):
+    slider_moves["dollars"] += 1
+
+    # if hours was the first updated slider
+    if sum(slider_moves.values()) <= n_sliders * slider_moves["dollars"] - n_sliders + 1:
+        hours_slider.value = dollars_to_hours(dollars_slider.value, gpu_dropdown.value)
+        energy_text.text = energy_fill(hours_to_kWh(hours_slider.value, gpu_dropdown.value),
+                                           hours_to_co2(hours_slider.value, gpu_dropdown.value))
+
+
+def gpu_update(attrname, old, new):
+    update_tip(tip, tipping_point(gpu_dropdown.value, amp_mode_dropdown.value, param_popt), gpu_dropdown.value,
+               amp_mode_dropdown.value, loss_popt, param_popt)
+    hours_slider.start = tip["hours"]
+    dollars_slider.start = hours_to_dollars(tip["hours"], gpu_dropdown.value)
+    if dollars_to_hours(dollars_slider.value, gpu_dropdown.value) == hours_slider.value:
+        width = hours_to_width(hours_slider.value, gpu_dropdown.value, amp_mode_dropdown.value, param_popt)
+        compare_and_update(width)
+    else:
+        dollars_slider.end = hours_to_dollars(hours_end, new)
+        hours_slider.value = dollars_to_hours(dollars_slider.value, gpu_dropdown.value)
+    energy_text.text = energy_fill(hours_to_kWh(hours_slider.value, gpu_dropdown.value),
+                                       hours_to_co2(hours_slider.value, gpu_dropdown.value))
+
+
+def amp_update(attrname, old, new):
+    update_tip(tip, tipping_point(gpu_dropdown.value, amp_mode_dropdown.value, param_popt), gpu_dropdown.value,
+               amp_mode_dropdown.value, loss_popt, param_popt)
+    width = hours_to_width(hours_slider.value, gpu_dropdown.value, amp_mode_dropdown.value, param_popt)
+    hours_slider.start = tip["hours"]
+    dollars_slider.start = hours_to_dollars(tip["hours"], gpu_dropdown.value)
+    compare_and_update(width)
+    energy_text.text = energy_fill(hours_to_kWh(hours_slider.value, gpu_dropdown.value),
+                                       hours_to_co2(hours_slider.value, gpu_dropdown.value))
+
+
+def loss_tap(event):
+    _, loss = event.x, event.y
+    flo = loss_to_flo(loss, *loss_popt)
+    param_number = safe_flo_to_param(flo, *param_popt)
+    width = param_to_width(param_number)
+    compare_and_update(width)
+
+
+loss_plot.on_event(Tap, loss_tap)
+
+
+def param_tap(event):
+    _, param_number = event.x, event.y
+    width = param_to_width(param_number)
+    hours = width_to_hours(width, gpu_dropdown.value, amp_mode_dropdown.value, param_popt)
+    hours_slider.value = hours
+
+
+param_plot.on_event(Tap, param_tap)
+
+hours_slider.on_change('value', hours_update)
+dollars_slider.on_change('value', dollars_update)
+gpu_dropdown.on_change("value", gpu_update)
+amp_mode_dropdown.on_change("value", amp_update)
+
+
+########################################################################################################################
+# Buttons
+########################################################################################################################
+
+def on_optimal_click():
+    code_box.text = hf_code(example_shape['example_width'], example_shape['example_depth'])
+
+
+def on_alternate_click():
+    code_box.text = hf_code(example_shape['alternate_width'], example_shape['alternate_depth'])
+
+
+input_text = Div(text="Choose a GPU, AMP mode, and budget:", width=sidebar_width, height=30,
+                 style={"display": "block", "margin": "0 auto", 'font-size': "125%",
+                        'font-weight': "bold", "width": f"{sidebar_width}px", "text-align": 'center'})
+initialize_optimal = Button(width=175, label="Initialize in 🤗transformers!")
+initialize_optimal.align = "center"
+initialize_optimal.on_click(on_optimal_click)
+results_buffer = Div(text="", width=sidebar_width, height=5, style=sidebar_div_style)
+initialize_alternate = Button(width=175, label="Initialize in 🤗transformers!")
+initialize_alternate.align = "center"
+initialize_alternate.on_click(on_alternate_click)
+
+code_box_style = {"display": "block", "margin": "0 auto", "width": f"{sidebar_width + plot_width}px",
+                  "text-align": 'center',
+                  "white-space": "pre-wrap", "background": "#f4f4f4",
+                  "border": "1px solid #ddd",
+                  "border-left": "3px solid #f36d33",
+                  "color": "#666",
+                  "page-break-inside": "avoid",
+                  "font-family": "monospace",
+                  "font-size": "15px",
+                  "line-height": "1.6",
+                  "max-width": "100%",
+                  "overflow": "hidden",
+                  "min-height": "30px",
+                  "word-wrap": "break-word"}
+code_box = Div(text="Find the right model for you with the curves and sliders then click the buttons to display the "
+                    "corresponding 🤗transformers code here!", width=sidebar_width + plot_width, style=code_box_style,
+               sizing_mode="scale_width")
+code_box.align = "center"
+
+########################################################################################################################
+# Add write-up text
+########################################################################################################################
+
+text_width = "800px"
+main_text_style = {"min-height": "100px",
+                   "overflow": "hidden",
+                   "display": "block",
+                   "margin": "auto",
+                   "width": text_width,
+                   "font-size": "18px"}
+
+formula_img_style_1 = {"min-height": "25px",
+                       "display": "block",
+                       "margin": "0 auto",
+                       "width": text_width,
+                       "height": "auto",
+                       "max-width": "100%",
+                       "max-height": "100%"}
+
+formula_img_style_2 = {"min-height": "50px",
+                       "display": "block",
+                       "margin": "0 auto",
+                       "width": text_width,
+                       "height": "auto",
+                       "max-width": "100%",
+                       "max-height": "100%"}
+
+text_1 = Div(text=md1, style=main_text_style)
+text_2 = Div(text=md2, style=main_text_style)
+text_3 = Div(text=md3, style=main_text_style)
+text_4 = Div(text=md4, style=main_text_style)
+
+########################################################################################################################
+# Loss plot in write-up
+########################################################################################################################
+
+in_text_loss_plot = figure(plot_height=in_text_plot_height, plot_width=in_text_plot_width,
+                           title="Validation loss during training for an array of models of different sizes",
+                           tools="pan,reset,save,wheel_zoom,tap", active_scroll="wheel_zoom",
+                           x_range=[min(all_points[:, 0]) * day_ratio, max(all_points[:, 0]) * day_ratio],
+                           y_range=[min(all_points[:, 1]), max(all_points[:, 1])],
+                           x_axis_type="log", y_axis_type="log",
+                           x_axis_label="Floating-point operations (excluding embeddings & softmax)",
+                           y_axis_label="Validation loss on Wikitext-103", output_backend="webgl")
+in_text_loss_plot.add_layout(color_bar, "left")
+in_text_loss_plot.align = "center"
+
+source = ColumnDataSource(data=dict(
+    xs=[run[:, 0] * day_ratio for run in indexed_runs],  # x coords for each line (list of lists)
+    ys=[run[:, 1] for run in indexed_runs],  # y coords for each line (list of lists)
+    params=params_per_run  # data to use for colormapping
+))
+in_text_loss_plot.multi_line('xs', 'ys', source=source,
+                             color=log_cmap('params', palette, min(params_per_run), max(params_per_run)))
+source = ColumnDataSource(data=dict(
+    x=[compute for run in indexed_runs for compute in run[:, 0] * day_ratio],  # x coords for each line (list of lists)
+    y=[loss for run in indexed_runs for loss in run[:, 1]],  # y coords for each line (list of lists)
+    params=[repeated_params for i, params in enumerate(params_per_run)
+            for repeated_params in [params] * len(indexed_runs[i])]  # data to use for colormapping
+))
+in_text_loss_plot.scatter('x', 'y', source=source,
+                          color=log_cmap('params', palette, min(params_per_run), max(params_per_run)), size=3)
+# for i, run in indexed_runs.items():
+#     source = ColumnDataSource(data=dict(x=run[:, 0] * day_ratio, y=run[:, 1]))
+#     in_text_loss_plot.line('x', 'y', source=source, line_width=1, line_alpha=0.6, color=color_list[i])
+#     in_text_loss_plot.scatter('x', 'y', source=source, line_width=1, line_alpha=0.6, color=color_list[i])
+
+in_text_param_plot = figure(plot_height=in_text_plot_height, plot_width=in_text_plot_width,
+                            title="Optimal number of non-embedding parameters per floating-point operations budget",
+                            tools="pan,reset,save,wheel_zoom,tap", active_scroll="wheel_zoom",
+                            x_range=in_text_loss_plot.x_range,
+                            y_range=[min(params_per_run), max(params_per_run)],
+                            x_axis_type="log", y_axis_type="log",
+                            x_axis_label="Floating-point operations (excluding embeddings & softmax)",
+                            y_axis_label="Optimal number of non-embedding parameters", output_backend="webgl")
+in_text_param_plot.add_layout(color_bar, "left")
+in_text_param_plot.align = "center"
+# for i, run_apex in enumerate(compute_at_hull):
+#     source = ColumnDataSource(data=dict(x=[compute_at_hull[i, 0] * day_ratio], y=[compute_at_hull[i, 1]]))
+#     in_text_param_plot.scatter('x', 'y', source=source, color=color_list[run_indices_at_hull[i]])
+
+source = ColumnDataSource(data=dict(x=compute_at_hull[:, 0] * day_ratio, y=compute_at_hull[:, 1],
+                                    params=[params for i, params in enumerate(params_per_run) if
+                                            i in set(hull_points[:, 2])]))
+in_text_param_plot.scatter('x', 'y', source=source,
+                           color=log_cmap('params', palette, min(params_per_run), max(params_per_run)))
+
+training_button = Button(width=175, label="Fit!")
+training_button.align = "center"
+fit_button = Button(width=175, label="Fit!")
+fit_button.align = "center"
+
+
+def on_train_click():
+    display_abscisses = np.array([min(all_points[:, 0]) / 1.25] + sorted(list(all_points[:, 0])) +
+                                 [max(all_points[:, 0]) * 1.25])
+    source = ColumnDataSource(
+        data=dict(x=sorted(display_abscisses * day_ratio), y=loss_fit(sorted(display_abscisses), *loss_popt)))
+    in_text_loss_plot.line('x', 'y', source=source, line_width=1, line_alpha=1, color="red")
+
+
+def on_fit_click():
+    display_abscisses = np.array([min(compute_at_hull[:, 0]) / 1.25] + sorted(list(compute_at_hull[:, 0])) +
+                                 [max(compute_at_hull[:, 0]) * 1.25])
+    source = ColumnDataSource(data=dict(x=display_abscisses * day_ratio,
+                                        y=safe_flo_to_param(display_abscisses, d, e, f)))
+    in_text_param_plot.line('x', 'y', source=source, line_width=1, line_alpha=0.8, color="orange")
+
+
+training_button.on_click(on_train_click)
+fit_button.on_click(on_fit_click)
+
+before_text = column(text_1, training_button, in_text_loss_plot, text_2, fit_button, in_text_param_plot, text_3)
+after_text = column(text_4)
+
+########################################################################################################################
+# Set up layouts and add to document
+########################################################################################################################
+
+inputs = column(input_text, gpu_dropdown, amp_mode_dropdown, hours_slider, dollars_slider, input_buffer, energy_text,
+                sizing_mode="scale_width", width=sidebar_width, height=plot_height)
+
+results = column(static_loss_text,
+                 optimal_loss_text,
+                 static_param_text,
+                 optimal_param_text,
+                 static_shape_text,
+                 optimal_shape_text,
+                 initialize_optimal,
+                 results_buffer,
+                 static_altshape_text,
+                 optimal_altshape_text,
+                 initialize_alternate, sizing_mode="scale_width", width=sidebar_width, height=plot_height)
+
+# app = column(row(inputs, loss_plot, sizing_mode="scale_width"), row(results, param_plot, sizing_mode="scale_width"),
+#              code_box, sizing_mode="scale_width")
+app = column(row(column(inputs, results, sizing_mode="fixed"),
+                 column(loss_plot, param_plot, sizing_mode="stretch_width", )),
+             code_box, sizing_mode="scale_width")
+before_text.align = "center"
+app.align = "center"
+after_text.align = "center"
+
+main_body = column(before_text, app, after_text, sizing_mode="scale_width")
+
+curdoc().add_root(main_body)
+curdoc().title = "How big should my language model be ?"

app_to_share/optimal_training/utils.py

@@ -0,0 +1,98 @@
+import copy
+import numpy as np
+
+from conversions import day_ratio
+
+
+def clean_run(run):
+    return [(a, float(b)) for a, b in run if b != "undefined"]
+
+
+def param_count(run):
+    compute_per_eval = run[0][0]
+    return round(compute_per_eval / 4000 / 150 / 60 / 6 * day_ratio)
+
+
+def convert_to_logspace(run, a, b, c):
+    logspace_run = copy.deepcopy(run)
+    logspace_run[:, 0] = b * np.log(run[:, 0])
+    logspace_run[:, 1] = -np.log(run[:, 1] - c) + np.log(a)
+    return logspace_run
+
+
+# OpenAI used another unit for floating-point operations with a ratio of the number of seconds in a day; we'll display
+# the raw number, but do the calculations with the ratio as it can overflow without it (convex hull notably fails)
+
+
+def hf_code(width, depth):
+
+    return f"""<span style="color: #008800; font-weight: bold">import</span> <span style="color: #0e84b5; font-weight: bold">transformers</span>
+config <span style="color: #333333">=</span> transformers<span style="color: #333333">.</span>TransfoXLConfig(d_model<span style="color: #333333">=</span><span style="color: #40a070">{width}</span>, d_embed<span style="color: #333333">=</span><span style="color: #40a070">{width}</span>, n_head<span style="color: #333333">=</span><span style="color: #40a070">8</span>, d_head<span style="color: #333333">=</span><span style="color: #40a070">{int(width / 8)}</span>, d_inner<span style="color: #333333">=</span><span style="color: #40a070">{width}</span>, n_layer<span style="color: #333333">=</span><span style="color: #40a070">{depth}</span>, tgt_len<span style="color: #333333">=</span><span style="color: #40a070">152</span>, mem_len<span style="color: #333333">=</span><span style="color: #40a070">152</span>)
+ model <span style="color: #333333">=</span> transformers<span style="color: #333333">.</span>TransfoXLModel(config)"""
+
+
+def co2_to_trees(co2):
+    return co2 / 60 * 3650
+
+
+def co2_to_kms(co2):
+    return co2 / 0.403 * 1.60934
+
+
+def energy_fill(kWh, co2):
+    return 'This will consume about <span style="font-weight: bold">{:.2f}</span> ' \
+           'kWh, releasing <span style="font-weight: bold">{:.2f}</span> ' \
+           'kgs of CO2. That is equivalent to <span style="font-weight: bold">{:.2f}</span> ' \
+           'kms with an average American passenger car and could be offset ' \
+           'by growing a tree for <span style="font-weight: bold">{:.2f}</span> ' \
+           'days.<sup><a href=' \
+           '"https://www.epa.gov/energy/greenhouse-gases-equivalencies-calculator-calculations-and-references#miles"' \
+           '>1</a></sup>'.format(kWh, co2, co2_to_kms(co2), co2_to_trees(co2))
+
+
+md1 = """<h1 id="how-big-should-my-language-model-be">How Big Should My Language Model Be?</h1>
+<img class='center' style='height: 5em; float: right;' src='https://raw.githubusercontent.com/TevenLeScao/transformer-xl/master/pytorch/assets/avatar_logo_joint.png' alt='avatar'>
+<h4>Published on June 08, 2020.</h4> 
+<h4>Teven Le Scao, researcher at Hugging Face • <a href="https://twitter.com/Fluke_Ellington">@Fluke_Ellington</a> </h4> 
+<p>Natural Language Processing can sometimes feel like model size is optimized for headlines. <a href="https://arxiv.org/abs/2005.14165">175 billion parameters</a> is certainly an eye-catching number! Why not just train more efficiently with a smaller model? One surprising scaling effect of deep learning is that <strong>bigger neural networks are actually compute-efficient.</strong> This is something OpenAI in particular has explored in papers like <em><a href="https://arxiv.org/abs/2001.08361">Scaling Laws for Neural Language Models</a></em>. Research at Hugging Face also leverages this phenomenon, and we&#39;ve combined it with GPU speed estimations to ensure model size is just right for the compute budget of the experiment (when in doubt, it&#39;s bigger than you think!). This blog post will show how this impacts architecture decisions on a standard language modeling benchmark: <strong>we replicate the 14-layer state-of-the-art result from <a href="https://arxiv.org/pdf/1901.02860.pdf">Zhang et al.&#39;s Transformer-XL paper</a> without any hyper-parameter optimization and saving 25% of training time</strong>. We also estimate that <strong>the 18-layer model from the same paper trained for an order of magnitude too many training steps.</strong> <a name="start">Wanna</a> play with our demo before reading? Just click <a href="#demo">here</a>!</p>
+<h2 id="1-there-is-an-optimal-time-to-stop-training-and-its-earlier-than-you-think">1. There is an optimal time to stop training (and it&#39;s earlier than you think)</h2>
+<p>Let&#39;s look at some loss curves. For our example, the task will be training Transformer-XL, the state-of-the-art in language modeling, on Wikitext-103, a standard, medium-size benchmark. GPT-2 doesn&#39;t perform well on this dataset scale. As training progresses, we&#39;ll look at the performance of the model (as measured by validation loss) depending on compute cost (as measured by floating point operations). Let&#39;s run a few experiments! In the following plot, every line of colour corresponds to a Transformer-XL run of 200000 steps with a different number and size of layers, with all other hyperparameters kept the same. This spans models from a mere thousand to a hundred million parameters (excluding embeddings). Bigger models are on the right as they require more compute for every step. Don&#39;t worry, we&#39;ve already run them so you don&#39;t have to. All graphs are interactive, play with them!</p>"""
+md2 = """
+<p>As introduced in <em>Scaling Laws</em>, we plot the validation loss against non-embedding floating-point operations (neFLOs). There seems to be a frontier of performance for a given neFLO budget that no model manages to beat, depicted here in red. In <em>Scaling Laws</em>, it is referred to as the compute frontier. Every run reaches it, or comes close, after an initial phase of quick loss improvement, then tapers away as the end of training is not as efficient. This has a very practical implication: if you have a certain budget in floating-point operations, to reach the best performance, you should choose a model size that reaches the compute frontier after that many operations and stop it at that moment. This is actually way before model convergence, which usually happens around 10 times later! In essence, if you have extra compute budget, you should invest most of it in a bigger model, and only a small part in more training steps. In <em>Scaling Laws</em>, the OpenAI team fitted a power law to the compute frontier on GPT-2 training. This still seems to be a good fit in our task. In addition, we also fitted a power law between the compute budget and the number of parameters of the model that is optimal for that budget. It is pictured in the following plot.</p>
+
+"""
+md3 = """
+<p>As good models tend to spend considerable time tangent on the compute frontier, there is a bit of noise in the relationship. However, this also means that there is more tolerance in the estimation even if the model size we predict is a bit off, as the imperfect model will still be very close to optimal. We find that <strong>if the compute budget is multiplied by 10, the optimal model size is multiplied by 7.41 and the number of optimal training steps by only 1.35</strong>. Extrapolating with this rule to the much-bigger 18-layer SOTA model from Zhang et al., we find that <strong>its optimal number of training steps was around 250000</strong>. Even if this number is imprecise due to the change of scale, <strong>it is much smaller than the 4 million steps from their replication script</strong>. Starting from an even bigger model and stopping earlier would have yielded a better loss for that (huge) compute budget.</p>
+<h2 id="2-gpus-are-optimized-for-large-wide-models">2. GPUs are optimized for large, wide models</h2>
+<p>We now have a rule connecting performance and optimal size with neFLOs. However, neFLOs are a bit hard to picture. Can we translate that into a more immediate resource, like training time? Whether you are constrained by temporal or financial constraints, the main resource is GPU time. In order to establish a connection between neFLOs and GPU time, we benchmarked different Transformer-XL model sizes on 4 different GPUs available on Google Cloud Platform across tens of thousands of runs, taking into account mixed precision training. Here are our findings:</p>
+<h5 id="speed-estimation">Speed estimation</h5>
+<p>neFLOs per second speed can be modeled as a factorized multivariate function (sounds scary, but this just means the equation can be written simply as below) of model width (the number of neurons per layer), depth (the number of layers) and batch size, by increasing order of importance. In our estimations, the maximum prediction error was 15% of the observed speed.</p>
+<img class='center' style='height: 1.25em;' src='https://raw.githubusercontent.com/TevenLeScao/transformer-xl/master/pytorch/assets/formula_1.png' alt='formula_1'>
+<h5 id="width">Width</h5>
+<p>GPUs are optimized for the large feed-forward layers of wide transformers. In all of our experiments, neFLOs per second depended on model width as <strong>a power law of exponent around 1.6</strong>. This means that a model that&#39;s twice as wide, which requires 4 times more operations, also goes through those operations around 3.16 times faster, <strong>nearly offsetting the additional compute cost</strong>.</p>
+<h5 id="depth">Depth</h5>
+<p>neFLOs per second were also positively correlated with depth. Our best results were attained by modeling this connection as proportional to depth * (depth + additive constant). This is coherent with the fact that Transformers must process layers serially. In essence, <strong>deeper models aren&#39;t actually faster, but they appear to be so as their overhead is smaller relative to the more productive operations</strong>. The additive constant, which represents this overhead, was consistently around 5 in our experiments, which essentially means that data loading to the GPU, embeddings, and softmax operations, represent around 5 transformer layers&#39; worth of time.</p>
+<h5 id="batch-size">Batch size</h5>
+<p>Batch size played the least role. It was <strong>positively correlated with speed for small values, but quickly saturated</strong> (and even seemed to hurt at high values, after 64 on the V100 and P100 and 16 on the K80 and P4). We modeled its contribution as a logarithmic function to keep things simple as it was also the variable for which the factorized independence assumption was the weakest. We ran all our experiments at size 64 on a single GPU. This is another perk of big models: <strong>as bigger batch sizes don&#39;t seem to help much, if your model is too big to fit on a GPU, you could just use a smaller batch size and gradient accumulation.</strong></p>
+<h5 id="powers-of-2-still-matter-in-2020">Powers of 2 still matter in 2020!</h5>
+<p>Finally, one surprising takeaway was that <strong>hyperparameters whose width or batch size were powers of 2 out-performed the others</strong>. That was the case on GPUs with and without Tensor Core capability. On Tensor Core GPUs like the V100, NVIDIA recommends tensor shapes that are multiples of 8; however, we kept seeing improvements beyond that, up to multiples of 512. In the end, we only fitted on powers of 2 as fitting on all data points meant a poor fit quality that consistently under-estimated speed for powers of 2 points, and one might as well choose the fastest parameters.</p>
+<p>In the end, our final estimation of operation speed was as follows:</p>
+<img class='center' style='height: 2.5em;' src='https://raw.githubusercontent.com/TevenLeScao/transformer-xl/master/pytorch/assets/formula_2.png' alt='formula_2'>
+<p>with, for example on a V100 GPU without mixed precision, k=2.21*10<sup>7</sup>, a=1.66, b=5.92, and c=1.33. Different GPUs had close results with a different multiplicative constant.</p>
+<h2 id="3-demonstration-on-a-language-modeling-task-wikitext-103">3. <a name="demo">Demonstration on a language modeling task: Wikitext-103</a></h2>
+<p>Now that we have obtained a relation between model size and training speed, we can predict, for a certain GPU time or price budget, the optimal model size on the task and the performance it will achieve.</p>
+
+"""
+md4 = """<p>Prices are indicated for Google Cloud Platform. The energy consumption was estimated thanks to Peter Henderson&#39;s <a href="https://github.com/Breakend/experiment-impact-tracker">Experiment impact tracker</a> and the CO2 emissions with <a href="https://www.electricitymap.org/zone/NL">Electricity map</a> Netherlands data (where Google&#39;s European servers are located). Even though huge training costs make headlines, it is still possible to replicate a state-of-the-art result on a medium-size dataset for thirty bucks! A single V100 with properly optimized training is already quite a powerful weapon.</p>
+<p>Data shown is for single-GPU training at batch size 60 on Wikitext-103 for a target and memory length of 150, following CMU&#39;s Transformer-XL <a href="https://github.com/kimiyoung/transformer-xl">repo</a>. In order to leverage the Tensor Core capability of the V100, we set batch size 64 and sequence length 152 on that GPU. In our model size and speed predictions, we assumed that the inner feed-forward layer dimension was the same as the embedding and attention dimensions, and that the width-to-depth ratio was constant. This is a good way to save memory, as <em><a href="https://arxiv.org/abs/2001.04451">Reformer</a></em> has shown. <em><a href="https://arxiv.org/abs/2001.08361">Scaling Laws</a></em> has observed that shape doesn&#39;t impact performance significantly in GPT-2. However, for large scales, we found that the final performance of taller models with a bigger feed-forward layer was consistently better, which is why we give two possible model shapes. </p>
+<p>In order to replicate the result of the medium-size Transformer-XL pre-trained model (3.15 loss), we tweaked our example model size to add a bigger feed-forward dimension and have high powers of 2 while keeping the same number of parameters. This gave us a model of 14 layers with 768 hidden dimensions and 1024 feed-forward dimensions. In comparison, the CMU pre-trained model was found through aggressive hyper-parameter search with a much more unusual shape of 16 layers of 410 hidden dimensions and 2100 feed-forward dimensions. In our experiment, even though it was 50% bigger, our model was actually 20% faster per batch on an NVIDIA RTX Titan as its shapes were high powers of 2, and it was a shorter, wider model. For that model, the script provided by the CMU team was already very close to optimal stopping time; in the end, we obtained the same performance with <strong>25% less training time</strong>. Most importantly, this was the case even though the pre-trained model&#39;s hyper-parameter tuning gave it a much more optimized shape, and we had also kept the same random seed it was tuned with. Since we calculated our scaling laws with much smaller-scale trainings, saving on parameter search might actually be the bigger gain here. If you took the shortcut to the demo before reading, you can come back the start <a href="#start">here</a>!</p>
+<h2 id="4-takeaways">4. Takeaways</h2>
+<ul>
+<li>Big models are surprisingly efficient!</li>
+<li>Training until convergence is not efficient at all.</li>
+<li>Benchmarking smaller-scale runs allows us to predict model performance and optimal stopping time for production-scale models.</li>
+<li>Using larger models stopped earlier and optimizing model size for speed lowers training costs.</li>
+</ul>
+<p>I built this tool automatically using the data from our Transformer-XL runs. If you are interested in having this feature available for other NLP tasks as part of the Hugging Face repository, you can contact me on Twitter at <a href="https://twitter.com/Fluke_Ellington">@Fluke_Ellington</a>, drop me a mail at <code>teven@huggingface.co</code>, or add a reaction on <a href="https://github.com/huggingface/transformers/issues/4847">our Github issue</a>!</p>
+
+"""

requirements.txt

@@ -0,0 +1,17 @@
+bokeh==2.0.2
+Jinja2==2.11.2
+MarkupSafe==1.1.1
+numpy==1.18.4
+packaging==20.4
+pandas==1.0.3
+Pillow==7.1.2
+pyparsing==2.4.7
+python-dateutil==2.8.1
+pytz==2020.1
+PyYAML==5.3.1
+randomcolor==0.4.4.5
+scipy==1.4.1
+six==1.15.0
+tornado==6.0.4
+typing-extensions==3.7.4.2
+