|
| 1 | +""" |
| 2 | +Reproducible Parallelized Reduction is difficult |
| 3 | +================================================ |
| 4 | +
|
| 5 | +A reduction is a frequent operation in neural network. It appears in layer normalization, |
| 6 | +softmax. Because of the float precision, the result of the computation |
| 7 | +changes based on the order of the elements. The following examples show the variation |
| 8 | +based on different hypothesis on the vector distribution. |
| 9 | +We consider a vector :math:`X = (x_1, ..., x_n)`. |
| 10 | +It computes the average: |
| 11 | +
|
| 12 | +.. math:: |
| 13 | +
|
| 14 | + mean(X) = \\frac{\\sum_{i=1}^n x_i}{n} |
| 15 | +
|
| 16 | +Or the normalization of the vector: |
| 17 | +
|
| 18 | +.. math:: |
| 19 | +
|
| 20 | + norm(X)_i = \\frac{ X_i - \\mathbb{E}X}{ \\sqrt{ \\mathbb{V}X}} |
| 21 | +
|
| 22 | +We draw 128 random permutation of X. The average or mean should not change. |
| 23 | +And the normalized vector should have the same value. In the first case, we compute |
| 24 | +the difference between the highest and the lowest values obtained for the average. |
| 25 | +In the second case, we look for the maximum difference between the original normalized |
| 26 | +vector and the permuted one (both sorted). |
| 27 | +
|
| 28 | +The computation code |
| 29 | +++++++++++++++++++++ |
| 30 | +""" |
| 31 | + |
| 32 | +import itertools |
| 33 | +from tqdm import tqdm |
| 34 | +import numpy as np |
| 35 | +import pandas |
| 36 | + |
| 37 | +DATA = [] |
| 38 | + |
| 39 | + |
| 40 | +def str_dtype(dtype): |
| 41 | + """Displays numpy dtype in a nicer way.""" |
| 42 | + if dtype == np.float64: |
| 43 | + return "fp64" |
| 44 | + if dtype == np.float32: |
| 45 | + return "fp32" |
| 46 | + if dtype == np.float16: |
| 47 | + return "fp16" |
| 48 | + raise ValueError(f"Unexpected value {dtype}") |
| 49 | + |
| 50 | + |
| 51 | +def layer_norm(a, eps=1e-6): |
| 52 | + """ |
| 53 | + Normalized the vector a. |
| 54 | + The computation is done in float32 or float64. |
| 55 | + """ |
| 56 | + ctype = np.float32 if a.dtype == np.float16 else a.dtype |
| 57 | + a32 = a.astype(ctype) |
| 58 | + m = a32.mean(axis=-1, keepdims=True) |
| 59 | + c = a32 - m |
| 60 | + va = np.sqrt((c * c).mean(axis=-1, keepdims=True)) |
| 61 | + va += eps |
| 62 | + return (c / va).astype(a.dtype) |
| 63 | + |
| 64 | + |
| 65 | +def compute(values, fct): |
| 66 | + """ |
| 67 | + Compare the results of function ``fct`` on a sample. |
| 68 | + Loops over multiple sizes, dtypes. Tries 128 times. |
| 69 | + """ |
| 70 | + |
| 71 | + def make_value(base, value): |
| 72 | + if value.size > 1: |
| 73 | + return np.abs(np.sort(base) - np.sort(value)).max() |
| 74 | + return value |
| 75 | + |
| 76 | + sizes = [2, 4, 8, 16, 512, 1024, 2048, 4096, 8192] |
| 77 | + dtypes = [np.float64, np.float32, np.float16] |
| 78 | + N = list(range(128)) |
| 79 | + exps = list(itertools.product(sizes, dtypes, N)) |
| 80 | + data = [] |
| 81 | + ech = None |
| 82 | + for size, dtype, n in tqdm(exps): |
| 83 | + if n == 0: |
| 84 | + ech = values[:size].astype(dtype) |
| 85 | + base = fct(ech) |
| 86 | + assert base.dtype == ech.dtype |
| 87 | + obs = dict( |
| 88 | + n=n, size=size, dtype=str_dtype(ech.dtype), value=make_value(base, fct(ech)) |
| 89 | + ) |
| 90 | + data.append(obs) |
| 91 | + |
| 92 | + if n == 1: |
| 93 | + new_ech = np.sort(ech) |
| 94 | + elif n == 2: |
| 95 | + new_ech = np.sort(ech)[::-1] |
| 96 | + else: |
| 97 | + new_ech = np.random.permutation(ech) |
| 98 | + assert new_ech.dtype == ech.dtype |
| 99 | + assert new_ech.shape == ech.shape |
| 100 | + obs = dict( |
| 101 | + n=n + 1, |
| 102 | + size=size, |
| 103 | + dtype=str_dtype(new_ech.dtype), |
| 104 | + value=make_value(base, fct(new_ech)), |
| 105 | + ) |
| 106 | + data.append(obs) |
| 107 | + |
| 108 | + df = pandas.DataFrame(data) |
| 109 | + agg = df.drop("n", axis=1).groupby(["dtype", "size"], as_index=False).agg(["min", "max"]) |
| 110 | + agg["value", "delta"] = agg["value", "max"] - agg["value", "min"] |
| 111 | + piv = agg.pivot(index="size", columns="dtype", values=("value", "delta")) |
| 112 | + return piv |
| 113 | + |
| 114 | + |
| 115 | +# %% |
| 116 | +# Normal Law |
| 117 | +# ++++++++++ |
| 118 | +# |
| 119 | +# Let's see what it returns an on random sample following a normal law. |
| 120 | +# First the average. |
| 121 | + |
| 122 | +values = np.random.randn(4096) |
| 123 | +mean = compute(values, lambda x: np.mean(x).astype(x.dtype)) |
| 124 | +mean["name"] = "normal" |
| 125 | +print(mean) |
| 126 | + |
| 127 | +# %% |
| 128 | +# Then the layer normalization. |
| 129 | + |
| 130 | +ln = compute(values, layer_norm) |
| 131 | +ln["name"] = "normal" |
| 132 | +DATA.append(ln.reset_index(drop=True).max(axis=0)) |
| 133 | +print(ln) |
| 134 | + |
| 135 | +# %% |
| 136 | +# Fixed values |
| 137 | +# ++++++++++++ |
| 138 | +# |
| 139 | +# We try a fixed vector with one very high value and all the others are small. |
| 140 | + |
| 141 | +values[:] = -1e-4 |
| 142 | +values[::128] = 100 |
| 143 | +mean = compute(values, lambda x: np.mean(x).astype(x.dtype)) |
| 144 | +mean["name"] = "fixed" |
| 145 | +print(mean) |
| 146 | + |
| 147 | + |
| 148 | +ln = compute(values, layer_norm) |
| 149 | +ln["name"] = "fixed" |
| 150 | +DATA.append(ln.reset_index(drop=True).max(axis=0)) |
| 151 | +print(ln) |
| 152 | + |
| 153 | +# %% |
| 154 | +# Pareto Distribution |
| 155 | +# +++++++++++++++++++ |
| 156 | +# |
| 157 | +# A law with a long tail. |
| 158 | + |
| 159 | +values = np.random.pareto(1, (4096,)) |
| 160 | +print(values) |
| 161 | + |
| 162 | +mean = compute(values, lambda x: np.mean(x).astype(x.dtype)) |
| 163 | +mean["name"] = "normal" |
| 164 | +print(mean) |
| 165 | + |
| 166 | + |
| 167 | +ln = compute(values, layer_norm) |
| 168 | +ln["name"] = "pareto" |
| 169 | +DATA.append(ln.reset_index(drop=True).max(axis=0)) |
| 170 | +print(ln) |
| 171 | + |
| 172 | +# %% |
| 173 | +# Summary |
| 174 | +# +++++++ |
| 175 | +# |
| 176 | +# We consider the maximum difference obtained for any sample size. |
| 177 | + |
| 178 | +print(DATA) |
| 179 | +df = pandas.DataFrame(DATA).set_index("name") |
| 180 | +print(df) |
| 181 | + |
| 182 | +# %% |
| 183 | +# Visually. |
| 184 | + |
| 185 | +ax = df.plot.bar(logy=True) |
| 186 | +fig = ax.get_figure() |
| 187 | +fig.savefig("plot_parallelized_reduction.png") |
| 188 | + |
| 189 | +# %% |
| 190 | +# In a deep neural network |
| 191 | +# ++++++++++++++++++++++++ |
| 192 | +# |
| 193 | +# Some of the vector have 500 values, 16x32x1024x1024. A layer normalization |
| 194 | +# does 16x32x1024 ~ 2M reductions, over 20 layers. |
| 195 | +# When a deep neural network is computed with a difference code, |
| 196 | +# doing a different parallelization (GPU/CPU for example), |
| 197 | +# the order of the reduction may change and therefore, |
| 198 | +# some errors will appear and propagate. |
0 commit comments