Spaces:
Build error
Build error
| import numpy as np | |
| import gradio as gr | |
| import matplotlib.pyplot as plt | |
| description = """ | |
| Minimizing L under the constraint FLOPs(N, D) = C. | |
| The functions $N_{opt}(C)$, and $D_{opt}(C)$ describe the optimal allocation of a computational budget $C$. | |
| We use the following notation: | |
| • L – the cross entropy loss in nats. Typically it will be averaged over the tokens in a context, but in | |
| some cases we report the loss for specific tokens within the context. | |
| • N – the number of model parameters, excluding all vocabulary and positional embeddings | |
| • D – the dataset size in tokens | |
| • C ≈ 6ND – an estimate of the total non-embedding training compute | |
| $$E=1.69, A=406.4, B=410.7, \\alpha=0.34, \\beta=0.28$$ | |
| $$C\\approx6DN$$ | |
| $$L(N,D)=E+\\frac{A}{N^\\alpha}+\\frac{B}{D^\\beta}$$ | |
| $$N_{opt}(C),D_{opt}(C)={\\arg\\min}_{N,D\ s.t.\ FLOP/s(N,D)=C}\ L(N,D)$$ | |
| """ | |
| article = """ | |
| References | |
| - [Training Compute-Optimal Large Language Models](https://arxiv.org/pdf/2203.15556.pdf) | |
| - [Scaling Laws for Neural Language Models](https://arxiv.org/pdf/2001.08361.pdf) | |
| - [karpathy/nanoGPT](https://github.com/karpathy/nanoGPT/blob/master/scaling_laws.ipynb) | |
| """ | |
| def L(N, D): | |
| """ | |
| Approximates loss given N parameters and D dataset size (in tokens), | |
| per Chinchilla paper. | |
| """ | |
| E = 1.69 # entropy of natural language, limit of infinite model on infinite data | |
| A = 406.4 | |
| B = 410.7 | |
| alpha = 0.34 | |
| beta = 0.28 | |
| return A / (N ** alpha) + B / (D ** beta) + E | |
| def plot_pens(tflpos_card, utilization, num_gps, training_days): | |
| fig = plt.figure() | |
| tflpos_card = float(tflpos_card)*(10**12) | |
| utilization = float(utilization) | |
| num_gps = int(num_gps) | |
| training_days = float(training_days) | |
| # target compute budget (usually know this because we know how many GPU for how long go brrr) | |
| c = tflpos_card*num_gps*86400*training_days*utilization | |
| # (I got this flop number from row 1 of Table A3) | |
| # sweep model sizes from 10M to 100B | |
| ns = 10 ** np.arange(7, 11, step=2**-4) | |
| # using C = 6*N*D, solve for D that maintains the compute budget c | |
| ds = c / (6 * ns) | |
| # evaluate the loss in each case | |
| losses = L(ns, ds) | |
| # find the argmin | |
| best = np.argmin(losses) | |
| best_model_size = f"{ns[best]/1e6:.2f}M" | |
| best_dataset_size = f"{ds[best]/1e9:.2f}B" | |
| # plot the loss | |
| # plt.figure(figsize=(3,3)) | |
| plt.plot(ns, losses) | |
| plt.xscale('log') | |
| # plot a vertical bar at the best model size | |
| plt.axvline(ns[best], color='red') | |
| plt.xlabel('model size') | |
| plt.ylabel('loss') | |
| fig.savefig("/tmp/tmp.jpg") | |
| plt.close() | |
| return "/tmp/tmp.jpg", c, round(losses[best], 3), best_model_size, best_dataset_size | |
| if __name__ == "__main__": | |
| iface = gr.Interface( | |
| fn=plot_pens, | |
| inputs=[ | |
| gr.Textbox(label="TFLOP/s pre Card",value="40"), | |
| gr.Slider(label="GPU Utilization", minimum=0, maximum=1, step=0.01,value=0.25), | |
| gr.Textbox(label="Number of cards"), | |
| gr.Textbox(label="Training Days") | |
| ], | |
| outputs=[ | |
| gr.Image(label="Estimated Loss"), | |
| gr.Label(label="Total Compute Budget"), | |
| gr.Label(label="Estimated Final Loss"), | |
| gr.Label(label="Optimal Model Size"), | |
| gr.Label(label="Optimal Dataset Size (Tokens)") | |
| ], | |
| title="Compute-Optimal Model Estimator", | |
| description=description, | |
| article=article, | |
| live=False | |
| ).launch() | |