dacorvo's picture
dacorvo HF Staff
Show NeuronCores, read the selection from selection_rank
07c358f verified
Raw History Blame Contribute Delete
16.9 kB
// Lists the support results of the aws-neuron/torch-neuronx-model-support
// dataset, one tab per library, 100 models per page: for the models the latest
// sweep selected, or for the Hub's models by trending score or downloads.
import { parquetReadObjects } from "https://cdn.jsdelivr.net/npm/hyparquet@1.31.2/src/index.js";
const HUB = "https://huggingface.co";
const DATASET = "aws-neuron/torch-neuronx-model-support";
const NEFF_CACHE = "hf://buckets/aws-neuron/torch-neuronx-neff-cache";
const LIBRARIES = ["transformers", "diffusers", "sentence-transformers"];
const VIEWS = { selected: "Selected", trending: "Trending", downloads: "Downloads" };
// The task Trending and Downloads open on: all tasks would open on models the sweep never tests.
const DEFAULT_TASK = "text-generation";
// The Hub's sort, and model field, of each view but "selected".
const HUB_SORTS = { trending: "trendingScore", downloads: "downloads" };
const PAGE = 100;
// The test instance: 4 Trainium2 devices of 96 GiB of HBM each.
const DEVICES = 4;
const DEVICE_HBM = 96 * 1024 ** 3;
const COLUMNS = ["model_id", "cores", "result", "wall (s)", "Memory (GB)", "notes"];
// An element with attributes and children. Strings become text nodes, never HTML.
function h(tag, attrs = {}, ...children) {
const el = document.createElement(tag);
for (const [key, value] of Object.entries(attrs)) {
if (key.startsWith("on")) el.addEventListener(key.slice(2), value);
else if (value === true) el.setAttribute(key, "");
else if (value !== false && value != null) el.setAttribute(key, value);
}
el.append(...children.flat(Infinity).filter((c) => c != null && c !== false));
return el;
}
const encodePath = (path) => path.split("/").map(encodeURIComponent).join("/");
async function fetchOk(url) {
const response = await fetch(url);
if (!response.ok) throw new Error(`${url}: HTTP ${response.status}`);
return response;
}
// Promises by key, dropped on failure so that the next render retries.
const cache = new Map();
function cached(key, load) {
if (!cache.has(key)) cache.set(key, load().catch((e) => { cache.delete(key); throw e; }));
return cache.get(key);
}
// The dataset's files at its latest commit: one per (neuronx-cc, library version)
// pair.
async function loadIndex() {
const info = await (await fetchOk(`${HUB}/api/datasets/${DATASET}`)).json();
const pairs = [];
for (const { rfilename: path } of info.siblings) {
const pair = path.match(/^data\/neuronxcc-([^/]+)\/(.+)-([^-/]+)\.parquet$/);
if (pair) pairs.push({ sdk: pair[1], library: pair[2], version: pair[3], path });
}
return { sha: info.sha, pairs };
}
const index = await loadIndex().catch((e) => e);
const fileUrl = (path) => `${HUB}/datasets/${DATASET}/resolve/${index.sha}/${encodePath(path)}`;
// A pair's rows: its runs, the models its sweep left out, and each model's rank
// in the selection of its last sweep, best-ranked first.
const loadRows = (path) => cached(path, async () => {
const rows = await parquetReadObjects({ file: await (await fetchOk(fileUrl(path))).arrayBuffer() });
const ranked = rows.filter((row) => row.selection_rank != null)
.sort((a, b) => a.selection_rank - b.selection_rank);
return {
runs: rows.filter((row) => row.result !== "EXCLUDED"),
exclusions: new Map(rows.filter((row) => row.result === "EXCLUDED").map((row) => [row.model_id, row])),
selection: new Map(ranked.map((row) => [row.model_id, row.selection_rank])),
};
});
// A model's task on the Hub, or "unknown" if the Hub no longer has it.
const loadTask = (id) => cached(`task:${id}`, async () => {
const query = new URLSearchParams({ "expand[]": "pipeline_tag" });
const response = await fetch(`${HUB}/api/models/${encodePath(id)}?${query}`);
if (response.status === 404) return "unknown";
if (!response.ok) throw new Error(`${id}: HTTP ${response.status}`);
return (await response.json()).pipeline_tag ?? "unknown";
}).catch(() => "unknown");
// The models of a library, of a task unless it is null, and whose id contains
// `name` unless it is empty, in the order of the Hub's model list, with their
// parameter count when the Hub knows it: at least `count` of them, or all if it
// has fewer. `next` is the Hub's link to the models after them, if any. The list
// grows one Hub page at a time, as deep as it was asked for.
const hubLists = new Map();
function hubList(library, task, name, view, count) {
const key = `${library}:${task}:${name}:${view}`;
const sort = HUB_SORTS[view];
if (!hubLists.has(key)) {
const query = new URLSearchParams({ library, sort, direction: "-1", limit: PAGE });
if (task) query.set("pipeline_tag", task);
if (name) query.set("search", name);
for (const field of [sort, "safetensors", "library_name"]) query.append("expand[]", field);
hubLists.set(key, { models: [], next: `${HUB}/api/models?${query}`, growing: Promise.resolve() });
}
const list = hubLists.get(key);
list.growing = list.growing.catch(() => {}).then(async () => {
while (list.models.length < count && list.next) {
const response = await fetchOk(list.next);
list.models.push(...(await response.json())
// The Hub's library filter matches a tag that models of other libraries
// carry too, such as diffusers' FLUX or sentence-transformers' MiniLM.
.filter((model) => model.library_name === library)
.map((model) => ({ id: model.id, score: model[sort] ?? 0, params: model.safetensors?.total })));
list.next = response.headers.get("Link")?.match(/<([^>]+)>;\s*rel="next"/)?.[1];
}
});
return list.growing.then(() => list);
}
// Newest first. A version that is not a release, such as a git commit pin,
// comes before the releases.
function newestFirst(a, b) {
const release = (v) => v.match(/^\d+(?:\.\d+)+/)?.[0].split(".").map(Number);
const [ra, rb] = [release(a), release(b)];
if (!ra || !rb) return Boolean(ra) - Boolean(rb) || a.localeCompare(b);
for (let i = 0; i < Math.max(ra.length, rb.length); i++) {
const d = (rb[i] ?? 0) - (ra[i] ?? 0);
if (d) return d;
}
// A pre-release, such as 5.21.0rc1, is older than its release.
const pre = (v) => /^\d+(?:\.\d+)+(?:a|b|rc|\.?dev)/.test(v);
return pre(a) - pre(b);
}
const badge = (result) => `${result === "PASS" ? "✅" : "❌"} ${result}`;
const memory = (run) => run.hbm_peak_gb == null ? ""
: `${run.hbm_peak_gb.toFixed(2)} / ${Number(run.hbm_capacity_gb.toFixed(1))}`;
const day = (date) => new Date(date).toISOString().slice(0, 10);
const tabs = document.getElementById("tabs");
const panel = document.getElementById("panel");
const dialog = document.getElementById("run");
const state = new Map(LIBRARIES.map((library) =>
[library, { view: "selected", sdk: null, version: null, tasks: {}, names: {}, page: 1 }]));
let active = LIBRARIES[0];
let rendering = 0;
// The pending list render of the name being typed.
let typing;
function renderTabs() {
tabs.replaceChildren(...LIBRARIES.map((library) => h("button", {
role: "tab",
"aria-selected": String(library === active),
onclick: () => { active = library; renderTabs(); renderPanel(); },
}, library)));
}
async function renderPanel() {
const render = ++rendering;
clearTimeout(typing);
if (index instanceof Error) {
panel.replaceChildren(h("p", { class: "status" }, `Could not read the dataset: ${index.message}`));
return;
}
const pairs = index.pairs.filter((p) => p.library === active);
if (!pairs.length) {
panel.replaceChildren(h("p", { class: "status" }, "No records yet."));
return;
}
const s = state.get(active);
const sdks = [...new Set(pairs.map((p) => p.sdk))].sort(newestFirst);
const versionsOf = (sdk) => pairs.filter((p) => p.sdk === sdk).map((p) => p.version).sort(newestFirst);
// The latest selection was swept with the newest library version of the newest compiler.
const latest = { sdk: sdks[0], version: versionsOf(sdks[0])[0] };
if (!sdks.includes(s.sdk)) s.sdk = latest.sdk;
if (!versionsOf(s.sdk).includes(s.version)) s.version = versionsOf(s.sdk)[0];
const selected = s.view === "selected";
const shown = selected ? latest : { sdk: s.sdk, version: s.version };
const pathOf = ({ sdk, version }) => pairs.find((p) => p.sdk === sdk && p.version === version).path;
const select = (label, key, options, onchange) => h("label", {}, label, h("select", {
onchange: (e) => { s[key] = e.target.value; s.page = 1; onchange?.(); renderPanel(); },
}, options.map((o) => h("option", { value: o, selected: o === s[key] }, o))));
const tasks = h("div", { class: "tasks", role: "group", "aria-label": "Task" }, h("span", {}, "Task"));
const results = h("div", { class: "results" }, h("p", { class: "status" }, "Loading…"));
// Typing renders the list again, not the panel, which would take the focus away.
let renderResults;
const nameField = h("label", {}, "Name", h("input", {
type: "search",
placeholder: "Filter by name",
value: s.names[s.view] ?? "",
oninput: (e) => {
s.names[s.view] = e.target.value.trim();
s.page = 1;
clearTimeout(typing);
typing = setTimeout(() => renderResults?.(), 300);
},
}));
panel.replaceChildren(
h("div", { class: "views", role: "tablist" }, Object.entries(VIEWS).map(([view, label]) => h("button", {
role: "tab",
"aria-selected": String(view === s.view),
onclick: () => { s.view = view; s.page = 1; renderPanel(); },
}, label))),
selected
? h("p", { class: "pair" }, `neuronx-cc ${latest.sdk} · ${active} ${latest.version}`)
: h("div", { class: "controls" },
select("SDK (neuronx-cc)", "sdk", sdks, () => { s.version = null; }),
select(active, "version", versionsOf(s.sdk)),
nameField),
...(selected ? [] : [tasks]),
results);
let runs, rank, exclusions, taskList;
try {
let latestRows;
[{ runs, exclusions }, latestRows] = await Promise.all([loadRows(pathOf(shown)), loadRows(pathOf(latest))]);
rank = latestRows.selection;
if (!selected) {
// The tasks the sweep covers: those of the latest pair's and the selection's models.
const covered = new Set([...latestRows.runs.map((run) => run.model_id), ...rank.keys()]);
taskList = [...new Set(await Promise.all([...covered].map(loadTask)))].sort();
}
} catch (e) {
if (render === rendering) results.replaceChildren(h("p", { class: "status" }, `Could not load the results: ${e.message}`));
return;
}
if (render !== rendering) return;
const runsOf = new Map();
for (const run of runs) runsOf.set(run.model_id, [...(runsOf.get(run.model_id) ?? []), run]);
// The task Trending or Downloads shows, or null for all of them.
if (!selected && !(s.view in s.tasks)) s.tasks[s.view] = taskList.includes(DEFAULT_TASK) ? DEFAULT_TASK : null;
const task = s.tasks[s.view] ?? null;
// The current page of the selection, or of the Hub's list of the chosen task or
// of all tasks: its models, each with its runs on the shown pair, the position
// of the first, and whether more follow.
async function listed() {
const start = (s.page - 1) * PAGE;
const end = start + PAGE;
let models, more;
if (selected) {
const ids = [...rank.keys()];
models = ids.slice(start, end).map((id) => ({ id }));
more = ids.length > end;
} else {
const list = await hubList(active, task, s.names[s.view] ?? "", s.view, end);
models = list.models.slice(start, end);
more = list.models.length > end || Boolean(list.next);
}
return { start, more, models: models.map((model) => ({ ...model, runs: runsOf.get(model.id) ?? [] })) };
}
// Why a model was not tested: weights larger than the test instance's HBM, or
// the reason the sweep's selection left it out.
function untestedReason({ id, params }) {
if (params * 2 > DEVICES * DEVICE_HBM) return `Requires more than ${DEVICES} Trainium2 devices`;
return exclusions.get(id)?.notes ?? "";
}
renderResults = async () => {
const render = ++rendering;
let page;
try {
page = await listed();
} catch (e) {
if (render === rendering) results.replaceChildren(h("p", { class: "status" }, `Could not load the Hub's list: ${e.message}`));
return;
}
if (render !== rendering) return;
if (!page.models.length) {
const name = s.names[s.view];
results.replaceChildren(h("p", { class: "status" }, selected ? "No selection yet."
: `No model of ${task ?? "any task"}${name ? ` matching "${name}"` : ""} is listed.`));
return;
}
// A model passes if one of its runs passed.
const passed = page.models.filter((m) => m.runs.some((run) => run.result === "PASS")).length;
const untested = page.models.filter((m) => !m.runs.length).length;
const go = async (pages) => {
s.page += pages;
await renderResults();
if (results.getBoundingClientRect().top < 0) results.scrollIntoView();
};
const pager = () => h("div", { class: "pager" },
h("button", { disabled: s.page === 1, onclick: () => go(-1) }, "‹ Previous"),
h("button", { disabled: !page.more, onclick: () => go(1) }, "Next ›"));
results.replaceChildren(
h("div", { class: "summary" },
h("p", {}, h("strong", {}, `${page.start + 1}-${page.start + page.models.length}`),
` · ✅ Pass ${passed} · ❌ Fail ${page.models.length - passed - untested} · Untested ${untested}`),
pager()),
h("div", { class: "scroll" }, h("table", {},
h("thead", {}, h("tr", {}, COLUMNS.map((c) => h("th", {}, c)))),
h("tbody", {}, page.models.map((model) => model.runs.length ? model.runs.map((run) => h("tr", {
tabindex: 0,
onclick: () => showRun(run),
onkeydown: (e) => { if (e.key === "Enter") showRun(run); },
},
h("td", {}, h("code", {}, model.id)),
h("td", { class: "num" }, run.neuron_cores),
h("td", { class: "result" }, badge(run.result)),
h("td", { class: "num" }, run.wall_s ?? ""),
h("td", { class: "num" }, memory(run)),
h("td", { class: "notes" }, run.notes ?? "")))
: h("tr", { class: "untested" },
h("td", {}, h("code", {}, model.id)), h("td"), h("td", { class: "result" }, "not tested"),
h("td"), h("td"), h("td", { class: "notes" }, untestedReason(model))))))),
pager());
};
// One task at a time, as on the Hub: clicking the chosen task again shows all of them.
if (!selected) {
tasks.append(...[null, ...taskList].map((t) => h("button", {
"aria-pressed": String(task === t),
onclick: () => { s.tasks[s.view] = task === t ? null : t; s.page = 1; renderPanel(); },
}, t ?? "All")));
}
renderResults();
}
let scriptUrl;
function showRun(run) {
URL.revokeObjectURL(scriptUrl);
scriptUrl = URL.createObjectURL(new Blob([run.run_py], { type: "text/x-python" }));
const pip = run.run_py.match(/^# (pip install .*)$/m)?.[1];
const stack = [
["neuronx-cc", run.neuronx_cc], ["torch-neuronx", run.torch_neuronx], ["torch", run.torch],
["nki", run.nki], ["Python", run.python], ["Neuron driver", run.neuron_driver],
].filter(([, version]) => version);
dialog.replaceChildren(h("div", { class: "dialog" },
h("form", { method: "dialog" }, h("button", { class: "close", "aria-label": "Close" }, "×")),
h("h2", {}, h("a", { href: `${HUB}/${run.model_id}`, target: "_blank", rel: "noopener" }, run.model_id)),
h("p", {}, `${badge(run.result)} · ${run.neuron_cores} NeuronCore${run.neuron_cores > 1 ? "s" : ""}`
+ ` · ${run.runner} · ${run.library} ${run.library_version}`
+ ` · tested ${day(run.tested_at)}`),
h("h3", {}, "Run it"),
h("p", {}, "On a Trainium2 (trn2) instance with this Neuron stack, for example in the AWS Neuron"
+ " Deep Learning Container:"),
h("dl", {}, stack.map(([name, version]) => [h("dt", {}, name), h("dd", {}, h("code", {}, version))])),
h("ol", {},
pip && h("li", {}, "Install the packages the script adds:", h("pre", {}, pip)),
h("li", {}, "Optionally, skip compilation with the graphs this compiler already built:",
h("pre", {}, `hf buckets sync ${NEFF_CACHE}/neuronxcc-${run.neuronx_cc} /tmp/neff_cache`)),
h("li", {}, "Download ", h("a", { href: scriptUrl, download: "run.py" }, "run.py"), " and run it",
run.neuron_cores > 1 ? ". It launches its own torch.distributed ranks:" : ":",
h("pre", {}, "python3 run.py"))),
run.error && [h("h3", {}, "Error"), h("pre", { class: "error" }, run.error)],
h("details", {}, h("summary", {}, "run.py"), h("pre", {}, run.run_py))));
dialog.showModal();
}
dialog.addEventListener("click", (e) => { if (e.target === dialog) dialog.close(); });
renderTabs();
renderPanel();