96 lines
4.5 KiB
JavaScript
96 lines
4.5 KiB
JavaScript
import assert from "node:assert/strict";
|
|
import test from "node:test";
|
|
import {
|
|
assessRecipe,
|
|
assessWeightsRecipe,
|
|
capacityHardware,
|
|
hostMemoryGb,
|
|
weightsFor,
|
|
} from "../lib/memory-assessment.ts";
|
|
import { buildCapacityCatalog } from "../../docs/js/build-cookbook-capacity.mjs";
|
|
import { assessCapacity } from "../../docs/js/cookbook-memory.mjs";
|
|
|
|
const hardware = { backend: "cuda", memory_gb: 80, memory_unit: "GiB", memory_type: "dedicated" };
|
|
const weights = weightsFor("fasth3-8step");
|
|
const config = {
|
|
generator: {
|
|
engine: { num_gpus: 4, offload: { text_encoder: true, vae: true } },
|
|
pipeline: { experimental: {} },
|
|
},
|
|
};
|
|
|
|
test("checkpoint data never certifies runtime fit or subtracts CPU-offloaded weights", () => {
|
|
for (const unified of [false, true]) {
|
|
const target = { ...hardware, memory_type: unified ? "unified" : "dedicated", usable_budget_gb: 60 };
|
|
const baseline = assessWeightsRecipe({ weights, config, hardware: target });
|
|
const edited = structuredClone(config);
|
|
edited.generator.engine.offload.text_encoder = false;
|
|
edited.generator.engine.num_gpus = 1;
|
|
const result = assessWeightsRecipe({ weights, config: edited, hardware: target });
|
|
assert.equal(result.status, "unknown");
|
|
assert.ok(Math.abs(result.capacityGiB - 60) < 1e-9);
|
|
assert.deepEqual(result.parts, weights.parts);
|
|
assert.equal(result.checkpointGiB, baseline.checkpointGiB);
|
|
assert.equal(result.requiredGiB, undefined);
|
|
}
|
|
});
|
|
|
|
test("a component swap preserves other known sizes and identifies missing coverage", () => {
|
|
for (const swap of ["taeh3", "nvfp4-decoder", "nvfp4-transformer", "encoder"]) {
|
|
const edited = structuredClone(config);
|
|
if (swap === "taeh3") edited.generator.pipeline.experimental.video_decode_backend = "taeh3";
|
|
if (swap === "nvfp4-decoder") edited.generator.pipeline.experimental.h3_vae_nvfp4_decoder = "decoder.safetensors";
|
|
if (swap === "nvfp4-transformer") edited.generator.engine.quantization = { transformer_quant: "NVFP4" };
|
|
if (swap === "encoder") edited.generator.pipeline.components = { text_encoder_weights: "replacement" };
|
|
const result = assessWeightsRecipe({ weights, config: edited, hardware });
|
|
const missingRole =
|
|
swap.includes("decoder") || swap === "taeh3" ? "vae" : swap === "encoder" ? "text_encoder" : "transformer";
|
|
assert.equal(result.status, "unknown");
|
|
assert.match(result.reason, /Partial checkpoint data/);
|
|
assert.equal(result.parts[missingRole], undefined);
|
|
assert.equal(result.parts.audio_vae, weights.parts.audio_vae);
|
|
assert.equal(
|
|
result.checkpointGiB,
|
|
Object.values(result.parts).reduce((sum, size) => sum + size, 0),
|
|
);
|
|
}
|
|
});
|
|
|
|
test("host RAM is rejected and unified memory gets one budget", () => {
|
|
assert.equal(
|
|
assessWeightsRecipe({ weights, config, hardware: { ...hardware, memory_type: "system" } }).parts,
|
|
undefined,
|
|
);
|
|
const result = assessWeightsRecipe({ weights, config, hardware: { ...hardware, memory_type: "unified" } });
|
|
assert.equal(result.capacityGiB, 60);
|
|
assert.ok(result.notes.some((note) => note.includes("75%")));
|
|
});
|
|
|
|
test("the UI adapter preserves GB, GiB, usable budgets and host memory through the analytical engine", () => {
|
|
const model = buildCapacityCatalog().models.find((model) => model.id === "fastwan21");
|
|
const example = { workload: "t2v", env: { FASTVIDEO_ATTENTION_BACKEND: "TORCH_SDPA" } };
|
|
const dense = structuredClone(model.examples[0].base_config);
|
|
dense.generator.engine.compile.enabled = false;
|
|
dense.generator.pipeline.experimental = {};
|
|
for (const unit of ["GB", "GiB"]) {
|
|
const target = { ...hardware, memory_unit: unit, usable_budget_gb: 70, host_ram_gb: 128 };
|
|
const result = assessRecipe({ model, example, config: dense, hardware: target });
|
|
const raw = assessCapacity({ model, example, config: dense, hardware: capacityHardware(target) });
|
|
assert.equal(result.status, raw.status);
|
|
assert.notEqual(result.status, "unknown");
|
|
assert.equal(result.requiredGiB, raw.requiredGiB);
|
|
assert.equal(result.capacityGiB, unit === "GiB" ? 70 : 70e9 / 1024 ** 3);
|
|
const hostGiB = unit === "GiB" ? 128 : 128e9 / 1024 ** 3;
|
|
assert.ok(Math.abs(raw.host.capacityGiB - hostGiB) < 1e-9);
|
|
const { host_ram_gb: _ram, ...withoutHost } = target;
|
|
const edited = assessCapacity({
|
|
model,
|
|
example,
|
|
config: dense,
|
|
hardware: capacityHardware(withoutHost),
|
|
hostMemoryGb: hostMemoryGb(target, 128),
|
|
});
|
|
assert.ok(Math.abs(edited.host.capacityGiB - hostGiB) < 1e-9);
|
|
}
|
|
});
|