ParScale/Parallel_Scaling_Law

1import gradio as gr2import matplotlib.pyplot as plt3import numpy as np4import math5from datetime import datetime6from matplotlib.ticker import FuncFormatter7 8# Predefined hyperparameter sets9PARAM_SETS = {10    "Stack-V2-Python": {"E": 0.69123678, "A": 0.01130616 * 1e9, "k": 0.393463, "alpha": 0.18937067},11    "Pile": {"E": 1.28254036, "A": 0.2035367 * 1e9, "k": 0.33027934, "alpha": 0.19479807}12}13 14def pred_loss(E, A, k, alpha, n, p):15    return E + (A / (n * (1 + np.log(p) * k))) ** alpha16 17def generate_plot(E, A, k, alpha):18    plt.clf()19    colors = ['#2B83BA', '#7BB7D6', '#ED7D5F', '#D7191C']20    ax = plt.gca()21    for i, p in enumerate([1, 2, 4, 8]):22        x_plot = np.linspace(535813376 * 0.9, 4353203200 * 1.1, 100)23        y_plot = pred_loss(E, A, k, alpha, x_plot, p)24        ax.plot(x_plot, y_plot, marker=None, markersize=1, linewidth=3, color=colors[int(math.log(p, 2))], label=f"$P={p}$")25 26    ax.legend(fontsize=12)27    # ax.set_xscale("log")28    # ax.set_yscale("log")29 30    def billions(x, pos):31        if x < 1e9:32            result = ""33        else:34            result = f'{x * 1e-9:.1f}B'35        return result36 37    ax.xaxis.set_major_formatter(FuncFormatter(billions))38    ax.xaxis.set_minor_formatter(FuncFormatter(billions))39    ax.yaxis.set_major_formatter(FuncFormatter(lambda x, pos: f"{x:.2f}"))40    ax.yaxis.set_minor_formatter(FuncFormatter(lambda x, pos: f"{x:.2f}"))41    ax.set_xlim(535813376 * 0.9, 4353203200 * 1.1)42    ax.set_ylim(ax.get_ylim()[0] * 1, ax.get_ylim()[1] * 1.01)43 44    ax.text(0.03, 0.03, f"$E={E}$\n$A={A}$\n$k={k}$\n$\\alpha={alpha}$", transform=ax.transAxes, fontsize=10, verticalalignment='bottom', multialignment='left')45 46    ax.spines['top'].set_visible(False)47    ax.spines['right'].set_visible(False)48 49    ax.set_xlabel('Parameters (Non-Embedding)', fontsize=12)50    ax.set_ylabel(f'Loss', fontsize=12)51    return plt52 53 54OUTPUT_TEMPLATE = """Loss for a {n}B model when P={p} is: **{loss:.5f}**. It is equivalant to:55 56- A **{n1}B** model with **P=1**;57- A **{n2}B** model with **P=2**;58- A **{n4}B** model with **P=4**;59- A **{n8}B** model with **P=8**;60 61Note: The equivalent parameters are for reference only. In some reasoning tasks, scaling the parallel streams will obtain more performance gains than the loss benefits!62 63Enjoy it! 😊"""64 65def process_inputs(E, A, k, alpha, n, p):66    """Process inputs and return results"""67    n = n * 1e968    plot = generate_plot(E, A, k, alpha)69    loss = pred_loss(E, A, k, alpha, n, p)70 71    n1 = n * (k * np.log(p) + 1) / (k * np.log(1) + 1) / 1e972    n2 = n * (k * np.log(p) + 1) / (k * np.log(2) + 1) / 1e973    n4 = n * (k * np.log(p) + 1) / (k * np.log(4) + 1) / 1e974    n8 = n * (k * np.log(p) + 1) / (k * np.log(8) + 1) / 1e975 76    print(f"[{datetime.now()}] {E = }, {A = }, {k = }, {alpha = }, {n = }, {p = }")77    78    return plot, OUTPUT_TEMPLATE.format(n=round(n / 1e9, 2), p=p, n1=round(n1, 2), n2=round(n2, 2), n4=round(n4, 2), n8=round(n8, 2), loss=loss)79 80# Create interface81 82HEAD = """<div align="center">83 84# Parallel Scaling Law Visualization85 86[![Paper](https://img.shields.io/badge/arXiv-2505.10475-red)](https://arxiv.org/abs/2505.10475)87</div>88"""89 90with gr.Blocks() as demo:91    gr.Markdown(HEAD)92    93    with gr.Row():94        with gr.Column():95 96            gr.Markdown("""$$97\\text{Loss}=E+\\left(98    \\frac{A}{\\text{Parameters}\\times (1+k\\log P)}99\\right)^{\\alpha}100$$""")101            102            # Input values103            N = gr.Number(value=2.8, label="N: Number of Non-Embedding Model Parameters (in Billion)")104            P = gr.Number(value=4, label="P: Number of Parallel Streams")105 106            gr.Markdown("---")107 108            # Hyperparameter selection section109            param_set = gr.Dropdown(110                choices=["Custom"] + list(PARAM_SETS.keys()),111                value=list(PARAM_SETS.keys())[0],112                label="Select our pre-fitted parameters for two datasets"113            )114            115            # Custom parameter inputs116            param_E = gr.Number(value=PARAM_SETS["Stack-V2-Python"]['E'], label="E")117            param_A = gr.Number(value=PARAM_SETS["Stack-V2-Python"]['A'], label="A")118            param_k = gr.Number(value=PARAM_SETS["Stack-V2-Python"]['k'], label="k")119            param_alpha = gr.Number(value=PARAM_SETS["Stack-V2-Python"]['alpha'], label="alpha")120            121        122 123        plot, output = process_inputs(PARAM_SETS["Stack-V2-Python"]['E'], PARAM_SETS["Stack-V2-Python"]['A'], PARAM_SETS["Stack-V2-Python"]['k'], PARAM_SETS["Stack-V2-Python"]['alpha'], 2.8, 4)124        with gr.Column():125 126            submit_btn = gr.Button("Calculate")127            # Output section128            plot_output = gr.Plot(label="Scaling Law Curve", value=plot)129            result_output = gr.Markdown(label="Result", value=output)130            131    132    # Auto-fill parameters when selecting predefined sets133    def update_params(param_set):134        if param_set in PARAM_SETS:135            params = PARAM_SETS[param_set]136            return [params["E"], params["A"], params["k"], params["alpha"]]137        return [gr.skip(), gr.skip(), gr.skip(), gr.skip()]138    139    param_set.change(140        update_params,141        inputs=[param_set],142        outputs=[param_E, param_A, param_k, param_alpha]143    )144    145    # Submit button event146    click_event = submit_btn.click(147        process_inputs,148        inputs=[param_E, param_A, param_k, param_alpha,149                N, P],150        outputs=[plot_output, result_output]151    )152 153 154demo.launch()