166 lines
5.2 KiB
Python
166 lines
5.2 KiB
Python
import marimo
|
|
|
|
__generated_with = "0.9.17"
|
|
app = marimo.App(width="medium")
|
|
|
|
|
|
@app.cell
|
|
def __():
|
|
import matplotlib.pyplot as plt
|
|
from sklearn.linear_model import LinearRegression
|
|
import numpy as np
|
|
return LinearRegression, np, plt
|
|
|
|
|
|
@app.cell
|
|
def __(plt):
|
|
plt.rcParams['font.family'] = 'Times New Roman'
|
|
plt.rcParams['font.size'] = 16
|
|
return
|
|
|
|
|
|
@app.cell(hide_code=True)
|
|
def __(LinearRegression, np, plt):
|
|
x = np.arange(4)
|
|
|
|
fig, ax = plt.subplots(figsize=(7, 4), ncols=3, nrows=2, layout="constrained", dpi=300)
|
|
|
|
c_1 = (230 / 255, 241 / 255, 243 / 255)
|
|
c_2 = (75 / 255, 116 / 155, 178 / 255)
|
|
c_3 = (255 / 255, 223 / 255, 146 / 255)
|
|
c_4 = (230 / 255, 109 / 255, 104 / 255)
|
|
|
|
# tp = 8 * l * b * h * (d-1)/d
|
|
# pp = 2 * b * h * (d - 1)
|
|
x = [2, 3, 4, 5, 6, 7, 8]
|
|
ticks = [2, 3, 4, 5, 6, 7, 8]
|
|
ticks_label = ["2", "3", "4", "5", "6", "7", "8"]
|
|
|
|
y = [5605.27, 8103.95, 10500.86, 12340.55, 15104, 16361.64, 19575.52]
|
|
model = LinearRegression()
|
|
model.fit(np.array(x).reshape(-1, 1), np.array(y).reshape(-1, 1))
|
|
py = [i * model.coef_[0] + model.intercept_ for i in x]
|
|
ax[0][0].plot(x, y, color=c_4, label="LoRAPP", marker="s")
|
|
ax[0][0].plot(x, py, color="black", label="Perfect Linear Scaling", linestyle=":")
|
|
ax[0][0].set_xticks(ticks)
|
|
ax[0][0].set_xticklabels(ticks_label)
|
|
ax[0][0].set_ylabel("Throughput (tokens/s)", fontsize=14)
|
|
ax[0][0].set_yticks([0, 10000, 20000])
|
|
ax[0][0].set_yticklabels(["0", "10k", "20k"], rotation=90, ha="center", va="center")
|
|
ax[0][0].tick_params(pad=7)
|
|
ax[0][0].set_title("(a) TinyLlama-1.1B", fontsize=14)
|
|
ax[0][0].tick_params(axis="y", which="major", pad=10)
|
|
ax[0][0].set_xlabel("Number of GPUs", fontsize=14)
|
|
y = [1214.44, 1813.77, 2363.19, 2813.11, 3300.74, 3923.94, 4460]
|
|
model = LinearRegression()
|
|
model.fit(np.array(x[1:]).reshape(-1, 1), np.array(y[1:]).reshape(-1, 1))
|
|
py = [i * model.coef_[0] + model.intercept_ for i in x]
|
|
ax[0][1].plot(x, y, color=c_4, marker="s")
|
|
ax[0][1].plot(x, py, color="black", linestyle=":")
|
|
ax[0][1].set_xticks(ticks)
|
|
ax[0][1].set_xticklabels(ticks_label)
|
|
ax[0][1].set_yticks([0, 2000, 4000])
|
|
ax[0][1].set_yticklabels(["0", "2k", "4k"], rotation=90, ha="center", va="center")
|
|
ax[0][1].tick_params(pad=7)
|
|
ax[0][1].set_title("(b) Llama-2-7B", fontsize=14)
|
|
ax[0][1].set_xlabel("Number of GPUs", fontsize=14)
|
|
|
|
y = [0, 977, 1280, 1600, 1918.28, 2260, 2513]
|
|
model = LinearRegression()
|
|
model.fit(np.array(x[1:]).reshape(-1, 1), np.array(y[1:]).reshape(-1, 1))
|
|
py = [i * model.coef_[0] + model.intercept_ for i in x]
|
|
ax[0][2].plot(x[1:], y[1:], color=c_4, marker="s")
|
|
ax[0][2].plot(x[1:], py[1:], color="black", linestyle=":")
|
|
ax[0][2].plot(x[0], y[0], color="r", marker="x")
|
|
ax[0][2].text(
|
|
2, 0, "OOM", fontsize=12, va="bottom", ha="left", color="r", style="italic"
|
|
)
|
|
ax[0][2].set_xticks(ticks)
|
|
ax[0][2].set_xticklabels(ticks_label)
|
|
ax[0][2].set_yticks([0, 1e3, 2e3, 3e3])
|
|
ax[0][2].set_yticklabels(["0", "1k", "2k", "3k"], rotation=90, ha="center", va="center")
|
|
ax[0][2].tick_params(pad=7)
|
|
ax[0][2].set_title("(c) Llama-2-13B", fontsize=14)
|
|
ax[0][2].set_xlabel("Number of GPUs", fontsize=14)
|
|
|
|
ax[1][0].set_title("(d) TinyLlama-1.1B", fontsize=14)
|
|
ax[1][1].set_title("(e) Llama-2-7B", fontsize=14)
|
|
ax[1][2].set_title("(f) Llama-2-13B", fontsize=14)
|
|
|
|
ax[1][0].set_xlabel("Rank of LoRA adapters", fontsize=14)
|
|
ax[1][1].set_xlabel("Rank of LoRA adapters", fontsize=14)
|
|
ax[1][2].set_xlabel("Rank of LoRA adapters", fontsize=14)
|
|
|
|
ax[1][0].set_ylabel("Throughput (tokens/s)", fontsize=14)
|
|
|
|
y_0 = [10525.52, 10467.51, 10315.85, 10286.17]
|
|
x_0 = [4, 8, 16, 32]
|
|
|
|
y_1 = [2309.40, 2258.73, 2252.43, 2242.80]
|
|
x_1 = [16, 32, 64, 128]
|
|
|
|
y_2 = [1245.79, 1244.60, 1224.91, 1207.40]
|
|
x_2 = [16, 32, 64, 128]
|
|
|
|
ax[1][0].plot(x_0, y_0, "s-", color=c_4)
|
|
ax[1][0].set_ylim(0, 11000)
|
|
ax[1][0].set_xticks(x_0)
|
|
ax[1][0].set_yticks([0, 5000, 10000])
|
|
ax[1][0].set_yticklabels(
|
|
["0", "5k", "10k"], rotation=90, ha="center", va="center"
|
|
)
|
|
ax[1][0].tick_params(pad=7)
|
|
|
|
ax[1][1].plot(x_1, y_1, "s-", color=c_4)
|
|
ax[1][1].set_ylim(0, 11000)
|
|
ax[1][1].set_xticks([32, 64, 128])
|
|
ax[1][1].set_yticks([0, 5000, 10000])
|
|
ax[1][1].set_yticklabels(
|
|
["0", "5k", "10k"], rotation=90, ha="center", va="center"
|
|
)
|
|
ax[1][1].tick_params(pad=7)
|
|
|
|
ax[1][2].plot(x_2, y_2, "s-", color=c_4)
|
|
ax[1][2].set_ylim(0, 11000)
|
|
ax[1][2].set_xticks([32, 64, 128])
|
|
ax[1][2].set_yticks([0, 5000, 10000])
|
|
ax[1][2].set_yticklabels(
|
|
["0", "5k", "10k"], rotation=90, ha="center", va="center"
|
|
)
|
|
ax[1][2].tick_params(pad=7)
|
|
|
|
fig.legend(
|
|
ncol=2,
|
|
bbox_to_anchor=(0.8, 1.1),
|
|
fancybox=False,
|
|
framealpha=0.0,
|
|
fontsize=14,
|
|
)
|
|
|
|
# plt.show()
|
|
plt.savefig("scalability.pdf", bbox_inches="tight", dpi=1000)
|
|
return (
|
|
ax,
|
|
c_1,
|
|
c_2,
|
|
c_3,
|
|
c_4,
|
|
fig,
|
|
model,
|
|
py,
|
|
ticks,
|
|
ticks_label,
|
|
x,
|
|
x_0,
|
|
x_1,
|
|
x_2,
|
|
y,
|
|
y_0,
|
|
y_1,
|
|
y_2,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app.run()
|