diff --git a/CMakeLists.txt b/CMakeLists.txt index 2ea0751..ab27d3c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -76,7 +76,8 @@ function(sd_native_carry_patch patch) endfunction() sd_native_carry_patch(pickle_line_span.patch src/model_io/pickle_io.cpp) -sd_native_carry_patch(fp8_metal_f16_load.patch src/model_manager.cpp) +sd_native_carry_patch(fp8_metal_f16_load.patch src/model_manager.cpp src/core/backend_fit.cpp) +sd_native_carry_patch(lora_immediate_fallback.patch src/pipeline/diffusion_engine.cpp) # Wrapper exports are compiled into the upstream library target, so they share # its SD_API export macro and reach upstream without patching it. diff --git a/README.md b/README.md index 71ccb4e..b159511 100644 --- a/README.md +++ b/README.md @@ -610,10 +610,14 @@ the macOS probe and builds the iOS device and simulator slices. Upstream 1017 has its [ABI audit, wrapper review, FP8 evidence and outstanding gates](docs/upstream_1017_qualification.md); the carry inventory is in the [upstream 982 record](docs/upstream_982_qualification.md). -It carries three fixes the pin lacks: unbatched CFG in Metal builds (below), +It carries four fixes the pin lacks: unbatched CFG in Metal builds (below), the [`pickle_line_span.patch`](patches/README.md) for legacy checkpoints above -2 GiB, and the [`fp8_metal_f16_load.patch`](patches/README.md), which loads FP8 -weights as F16 on Metal instead of aborting. +2 GiB, the [`fp8_metal_f16_load.patch`](patches/README.md), which loads FP8 +weights as F16 on Metal instead of aborting, so they take twice their file size +there, and the [`lora_immediate_fallback.patch`](patches/README.md), which +applies a LoRA at runtime instead of crashing where the `immediately` mode +cannot merge it: into FP8 weights, and on Metal, which cannot add into the F16 +convolution weights every model has. Upstream's `generate_image` now returns an owned result with per-image generation metadata; `sd_dart_generate_image` keeps returning an image array that `free_sd_images` frees and drops that metadata. Since upstream 982 a null diff --git a/patches/README.md b/patches/README.md index c24be11..a44241e 100644 --- a/patches/README.md +++ b/patches/README.md @@ -13,4 +13,5 @@ an upstream issue or pull request to retire it. | Patch | Sources | Fixes | Test | | --- | --- | --- | --- | | `pickle_line_span.patch` | `src/model_io/pickle_io.cpp` | A legacy PyTorch checkpoint above 2 GiB is rejected: the pickle scanner narrows the rest of the file to `int` before looking for a newline ([#23](https://github.com/leehack/stable-diffusion-native/issues/23)). | `tests/test_pickle_span.py` | -| `fp8_metal_f16_load.patch` | `src/model_manager.cpp` | An FP8 (E4M3 or E5M2) checkpoint aborts the process on Metal: ggml-metal claims FP8 `MUL_MAT` but has no FP8 kernel. Weights a Metal backend computes with now load as F16, which holds every FP8 value; other backends keep FP8 ([#29](https://github.com/leehack/stable-diffusion-native/issues/29)). | `tests/test_fp8_runtime.py` | +| `fp8_metal_f16_load.patch` | `src/model_manager.cpp`, `src/core/backend_fit.cpp` | An FP8 (E4M3 or E5M2) checkpoint aborts the process on Metal: ggml-metal claims FP8 `MUL_MAT` but has no FP8 kernel. Weights a Metal backend computes with now load as F16, which holds every FP8 value, so they take twice their file size; other backends keep FP8 ([#29](https://github.com/leehack/stable-diffusion-native/issues/29)). Auto-fit plans them at that size, and FP8 LoRA tensors, which crashed ggml-cpu, load as F16 on every backend ([#31](https://github.com/leehack/stable-diffusion-native/issues/31)). | `tests/test_fp8_runtime.py` | +| `lora_immediate_fallback.patch` | `src/pipeline/diffusion_engine.cpp` | The `immediately` LoRA apply mode, which `auto` selects for unquantized weights, crashes the CPU merging into FP8 weights (ggml-cpu can neither multiply nor write FP8) and aborts Metal merging into F16 weights (ggml-metal adds only equal types), which every model has: convolution weights load as F16. FP8 weights, and every model on a backend that cannot run that add outside host memory, now take the LoRA `at_runtime`, with a warning when `immediately` was requested ([#31](https://github.com/leehack/stable-diffusion-native/issues/31)). | `tests/test_fp8_runtime.py` | diff --git a/patches/fp8_metal_f16_load.patch b/patches/fp8_metal_f16_load.patch index 20a03b5..fcff2e3 100644 --- a/patches/fp8_metal_f16_load.patch +++ b/patches/fp8_metal_f16_load.patch @@ -1,3 +1,38 @@ +diff --git a/src/core/backend_fit.cpp b/src/core/backend_fit.cpp +--- a/src/core/backend_fit.cpp ++++ b/src/core/backend_fit.cpp +@@ -96,7 +96,19 @@ namespace sd::backend_fit { + return false; + } + ++ // model_manager.cpp loads FP8 weights that Metal computes with as F16, at ++ // twice their size. A component the plan keeps off Metal is overestimated. ++ static bool metal_loads_fp8_as_f16() { ++#ifndef SD_USE_UPSTREAM_GGML ++ ggml_backend_reg_t reg = ggml_backend_reg_by_name("MTL"); ++ return reg != nullptr && ggml_backend_reg_dev_count(reg) > 0; ++#else ++ return false; ++#endif ++ } ++ + static std::vector estimate_components(ModelLoader& loader, ggml_type override_wtype) { ++ const bool fp8_as_f16 = metal_loads_fp8_as_f16(); + int64_t bytes[3] = {0, 0, 0}; + int64_t largest_tensor[3] = {0, 0, 0}; + for (const auto& [name, stored_tensor] : loader.get_tensor_storage_map()) { +@@ -110,6 +122,11 @@ namespace sd::backend_fit { + } else if (override_wtype != GGML_TYPE_COUNT && loader.tensor_should_be_converted(ts, override_wtype)) { + ts.type = override_wtype; + } ++#ifndef SD_USE_UPSTREAM_GGML ++ if (fp8_as_f16 && (ts.type == GGML_TYPE_F8_E4M3 || ts.type == GGML_TYPE_F8_E5M2)) { ++ ts.type = GGML_TYPE_F16; ++ } ++#endif + const int64_t tensor_bytes = (int64_t)ts.nbytes() + 64; + bytes[int(kind)] += tensor_bytes; + largest_tensor[int(kind)] = std::max(largest_tensor[int(kind)], tensor_bytes); diff --git a/src/model_manager.cpp b/src/model_manager.cpp --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -9,20 +44,23 @@ diff --git a/src/model_manager.cpp b/src/model_manager.cpp #include #include #include -@@ -185,6 +186,36 @@ size_t estimate_tensors_size(const std::map& tensors) +@@ -185,6 +186,40 @@ size_t estimate_tensors_size(const std::map& tensors) return size; } +// ggml-metal reports FP8 MUL_MAT as supported, so Linear keeps FP8 weights, but +// it has no FP8 matrix kernel and aborts the process encoding one. Weights a +// Metal backend computes with are loaded as F16 instead, which holds every FP8 -+// value exactly. -+static void load_fp8_params_as_f16_on_metal(const std::map& tensors, -+ ggml_backend_t compute_backend) { ++// value exactly. So are FP8 LoRA tensors on every backend: LoRA multiplies ++// them directly, and ggml-cpu has no FP8 matrix kernel either. ++static void load_fp8_params_as_f16(ModelComponent component, ++ const std::map& tensors, ++ ggml_backend_t compute_backend) { +#ifndef SD_USE_UPSTREAM_GGML + ggml_backend_dev_t device = compute_backend != nullptr ? ggml_backend_get_device(compute_backend) : nullptr; + ggml_backend_reg_t reg = device != nullptr ? ggml_backend_dev_backend_reg(device) : nullptr; -+ if (reg == nullptr || std::strcmp(ggml_backend_reg_name(reg), "MTL") != 0) { ++ if (component != ModelComponent::LoRA && ++ (reg == nullptr || std::strcmp(ggml_backend_reg_name(reg), "MTL") != 0)) { + return; + } + for (const auto& pair : tensors) { @@ -38,6 +76,7 @@ diff --git a/src/model_manager.cpp b/src/model_manager.cpp + } + } +#else ++ (void)component; + (void)tensors; + (void)compute_backend; +#endif @@ -46,11 +85,11 @@ diff --git a/src/model_manager.cpp b/src/model_manager.cpp void ModelManager::set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft, const std::vector>& device_limits) { if (compute_backend == nullptr) { return; -@@ -228,6 +259,7 @@ bool ModelManager::register_param_tensors(ModelComponent component, +@@ -228,6 +263,7 @@ bool ModelManager::register_param_tensors(ModelComponent component, LOG_ERROR("model manager tensor component is invalid"); return false; } -+ load_fp8_params_as_f16_on_metal(tensors, compute_backend); ++ load_fp8_params_as_f16(component, tensors, compute_backend); if (registered_tensor_size != nullptr) { *registered_tensor_size += estimate_tensors_size(tensors); } diff --git a/patches/lora_immediate_fallback.patch b/patches/lora_immediate_fallback.patch new file mode 100644 index 0000000..413462d --- /dev/null +++ b/patches/lora_immediate_fallback.patch @@ -0,0 +1,82 @@ +diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp +--- a/src/pipeline/diffusion_engine.cpp ++++ b/src/pipeline/diffusion_engine.cpp +@@ -1012,26 +1012,64 @@ bool StableDiffusionGGML::tae_preview_only() const { + return config_->params.tae_preview_only && version != VERSION_SDXS_512_DS && version != VERSION_SDXS_09; + } + ++// Applying a LoRA immediately adds an F32 diff into each F16 weight in place, on ++// the backend holding the weight. ggml-metal adds only equal types, and the ++// scheduler aborts on an op a pre-allocated tensor's backend cannot run, unless ++// the tensor is in host memory, where the CPU runs it. Every model has F16 ++// weights: convolutions load as F16 whatever the file holds. ++static bool backend_adds_f32_into_f16(ggml_backend_t backend) { ++ if (ggml_backend_buft_is_host(ggml_backend_get_default_buffer_type(backend))) { ++ return true; ++ } ++ ggml_init_params params = {8 * ggml_tensor_overhead(), nullptr, true}; ++ ggml_context* ctx = ggml_init(params); ++ if (ctx == nullptr) { ++ return false; ++ } ++ ggml_tensor* weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, 16); ++ ggml_tensor* diff = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 16); ++ const bool supported = ggml_backend_supports_op(backend, ggml_add_inplace(ctx, weight, diff)); ++ ggml_free(ctx); ++ return supported; ++} ++ + void StableDiffusionGGML::configure_weight_loading() { + const auto* sd_ctx_params = &config_->params; + const auto& model_loader = model_manager->loader(); + const auto wtype_stat = model_loader.get_wtype_stat(); + bool have_int8_tensorwise = false; + bool have_quant_type_rule = false; ++ bool have_fp8_weight = false; + for (const auto& [_, tensor_storage] : model_loader.get_tensor_storage_map()) { + if (tensor_storage.is_int8_tensorwise) { + have_int8_tensorwise = true; + } ++ const ggml_type load_type = tensor_storage.expected_type < GGML_TYPE_COUNT ? tensor_storage.expected_type ++ : tensor_storage.type; ++ if (load_type == GGML_TYPE_F8_E4M3 || load_type == GGML_TYPE_F8_E5M2) { ++ have_fp8_weight = true; ++ } + if (tensor_storage.expected_type < GGML_TYPE_COUNT && ggml_is_quantized(tensor_storage.expected_type)) { + have_quant_type_rule = true; + } +- if (have_int8_tensorwise && have_quant_type_rule) { +- break; ++ } ++ // Merging into an FP8 weight crashes ggml-cpu, which has neither the FP8 ++ // matrix kernel the F32 cast uses nor an F32 to FP8 conversion. ++ std::string immediate_lora_blocker; ++ if (have_fp8_weight) { ++ immediate_lora_blocker = "FP8 weights"; ++ } else { ++ for (SDBackendModule module : {SDBackendModule::TE, SDBackendModule::DIFFUSION, SDBackendModule::VAE}) { ++ ggml_backend_t backend = backend_manager.runtime_backend(module); ++ if (backend != nullptr && !backend_adds_f32_into_f16(backend)) { ++ immediate_lora_blocker = std::string("F16 weights on ") + ggml_backend_name(backend); ++ break; ++ } + } + } + + if (sd_ctx_params->lora_apply_mode == LORA_APPLY_AUTO) { +- bool have_quantized_weight = have_int8_tensorwise; ++ bool have_quantized_weight = have_int8_tensorwise || !immediate_lora_blocker.empty(); + for (const auto& [type, _] : wtype_stat) { + if (ggml_is_quantized(type)) { + have_quantized_weight = true; +@@ -1054,6 +1092,10 @@ void StableDiffusionGGML::configure_weight_loading() { + "INT8 tensorwise weights do not support the immediately LoRA apply mode; " + "using at_runtime instead"); + apply_lora_immediately = false; ++ } else if (!immediate_lora_blocker.empty()) { ++ LOG_WARN("%s do not support the immediately LoRA apply mode; using at_runtime instead", ++ immediate_lora_blocker.c_str()); ++ apply_lora_immediately = false; + } else if (row_split_active()) { + LOG_WARN( + "row-split tensors do not support the immediately LoRA apply mode; " diff --git a/tests/native/fp8_runtime_test.cpp b/tests/native/fp8_runtime_test.cpp index 09912ee..c491acc 100644 --- a/tests/native/fp8_runtime_test.cpp +++ b/tests/native/fp8_runtime_test.cpp @@ -2,13 +2,18 @@ // as llamadart does, and writes its RGB bytes. // // Usage: fp8_runtime_test devices -// fp8_runtime_test generate +// fp8_runtime_test generate [ []] +// fp8_runtime_test plan // // The model is one that `exit_teardown_runtime_test make-model` wrote, or an -// FP8 copy of it: it holds its TAESD decoder, so it needs no VAE. +// FP8 copy of it: it holds its TAESD decoder, so it needs no VAE. is a +// LoRA apply mode name: auto, immediately or at_runtime, and are +// tensor type rules. Both commands print the log, which holds upstream's +// auto-fit plan and LoRA apply mode warnings; `plan` only loads the model. #include "sd_dart_wrapper.h" +#include #include #include #include @@ -25,7 +30,15 @@ int list_devices() { return 0; } -int generate(const char* model, const char* backend, const char* output) { +void print_log() { + std::vector text(4096); + for (uint64_t sequence = 0; + (sequence = sd_dart_log_read(sequence, text.data(), text.size(), nullptr, nullptr)) != 0;) { + std::printf("%s", text.data()); + } +} + +int plan(const char* model, const char* backend) { sd_dart_log_enable(); sd_ctx_params_t context_params; sd_ctx_params_init(&context_params); @@ -33,6 +46,29 @@ int generate(const char* model, const char* backend, const char* output) { context_params.taesd_path = model; context_params.backend = backend; sd_ctx_t* context = sd_dart_new_sd_ctx(&context_params); + if (context != nullptr) { + sd_dart_exit_free(context); + } + return context != nullptr ? 0 : 3; +} + +int generate(const char* model, const char* backend, const char* output, const char* lora, + const char* mode, const char* rules) { + sd_dart_log_enable(); + sd_ctx_params_t context_params; + sd_ctx_params_init(&context_params); + context_params.model_path = model; + context_params.taesd_path = model; + context_params.backend = backend; + context_params.tensor_type_rules = rules; + if (lora != nullptr) { + context_params.lora_apply_mode = str_to_lora_apply_mode(mode); + if (context_params.lora_apply_mode == LORA_APPLY_MODE_COUNT) { + std::fprintf(stderr, "unknown LoRA apply mode %s\n", mode); + return 2; + } + } + sd_ctx_t* context = sd_dart_new_sd_ctx(&context_params); if (context == nullptr) { std::vector error(4096); sd_dart_last_error(error.data(), error.size()); @@ -49,6 +85,11 @@ int generate(const char* model, const char* backend, const char* output) { params.seed = 7; params.sample_params.sample_steps = 2; params.sample_params.guidance.txt_cfg = 1.0f; + sd_lora_t adapter = {false, 1.0f, lora}; + if (lora != nullptr) { + params.loras = &adapter; + params.lora_count = 1; + } sd_image_t* images = nullptr; int count = 0; if (!sd_dart_generate_image(context, ¶ms, &images, &count)) { @@ -83,9 +124,19 @@ int main(int argc, char** argv) { if (argc == 2 && std::strcmp(argv[1], "devices") == 0) { return list_devices(); } - if (argc == 5 && std::strcmp(argv[1], "generate") == 0) { - return generate(argv[2], argv[3], argv[4]); + int status = 2; + if (argc == 4 && std::strcmp(argv[1], "plan") == 0) { + status = plan(argv[2], argv[3]); + } else if ((argc == 5 || argc == 7 || argc == 8) && std::strcmp(argv[1], "generate") == 0) { + status = generate(argv[2], argv[3], argv[4], argc >= 7 ? argv[5] : nullptr, + argc >= 7 ? argv[6] : nullptr, argc == 8 ? argv[7] : nullptr); + } else { + std::fprintf(stderr, + "usage: %s devices | generate [ []] | " + "plan \n", + argv[0]); + return status; } - std::fprintf(stderr, "usage: %s devices | generate \n", argv[0]); - return 2; + print_log(); + return status; } diff --git a/tests/test_fp8_runtime.py b/tests/test_fp8_runtime.py index b2f5709..07123a0 100644 --- a/tests/test_fp8_runtime.py +++ b/tests/test_fp8_runtime.py @@ -1,14 +1,19 @@ -"""Generates from FP8 E4M3 weights on every device of a built runtime. +"""Generates from FP8 weights, and with LoRAs, on every device of a built runtime. The model is the 12 MB PixArt model `exit_teardown_runtime_test make-model` writes, with its weights replaced by fixed patterns. Its 2-D weights hold -values E4M3 represents exactly, kept as F32 in a control copy and stored as -F8_E4M3 in a second copy, so on each device the two images differ only by -matrix precision. Nothing is downloaded. +values that E4M3, E5M2, F16 and BF16 represent exactly, kept as F32 in a control +copy and stored in each of those types in further copies, so on each device +the images differ only by matrix precision. Nothing is downloaded. On Metal `patches/fp8_metal_f16_load.patch` loads FP8 weights as F16; without -it the FP8 generation aborts the process (#29). GitHub's macOS runners list -no Metal device; run this on a Mac with `SD_REQUIRE_METAL_GENERATION=1`. +it an FP8 generation aborts the process (#29). The LoRA tests apply small +linear and convolution adapters in every apply mode to each weight type. +Without the patches an FP8 LoRA tensor crashes the CPU in every mode, and the +immediately mode, which auto selects for these weights, crashes the CPU for +FP8 weights and aborts Metal, whose F16 convolution weights it cannot add +into (#31). GitHub's macOS runners list no Metal device; run this on a Mac +with `SD_REQUIRE_METAL_GENERATION=1`. SD_FP8_TARGET= a built Linux target; macOS defaults to the host's @@ -20,8 +25,10 @@ from __future__ import annotations import json +import math import os import platform +import re import shutil import struct import subprocess @@ -39,21 +46,59 @@ MODEL_WRITER = REPO_ROOT / "tests" / "native" / "exit_teardown_runtime_test.cpp" PROBE = REPO_ROOT / "tests" / "native" / "fp8_runtime_test.cpp" GGML_METAL = REPO_ROOT / "third_party" / "stable-diffusion.cpp" / "ggml" / "src" / "ggml-metal" -# Mean absolute difference per channel byte between the FP8 and F32 images of -# one device. An M4 Max measured 0.03 on the CPU and 0 on Metal. +# Mean absolute difference per channel byte between two images of one device +# that should match. An M4 Max measured at most 0.11 on the CPU and 0 on Metal. MAX_MEAN_DIFFERENCE = 0.5 # Fewer distinct bytes would mean a flat image that hides any difference. MIN_DISTINCT_BYTES = 16 -# 16 E4M3 values, +-2^-6 * (1 + m/8): sign, exponent field 1, mantissa m. -E4M3_PATTERN = bytes(sign | 0x08 | m for sign in (0, 0x80) for m in range(8)) +FP8_TYPES = {"F8_E4M3": (7, 3), "F8_E5M2": (15, 2)} # exponent bias, mantissa bits +# 2-D weights: +-2^-6 * (1 + m/4), which every type here holds exactly. +WEIGHT_TYPES = ("F32", "F16", "BF16", *FP8_TYPES) +WEIGHT_VALUES = [sign * 2.0 ** -6 * (1 + m / 4) for sign in (1, -1) for m in range(4)] # Every other tensor: values in [-0.125, 0.125). Larger ones make the tiny # network chaotic, so that matrix precision alone changes the image. OTHER_PATTERN = struct.pack("<61f", *(((i * 37) % 61 - 30.5) / 61 * 0.25 for i in range(61))) +# Upstream loads a convolution weight as F16 whatever the file holds. +LINEAR_TARGET = "model.diffusion_model.final_layer.linear" +CONV_TARGET = "model.diffusion_model.x_embedder.proj" +# Applied immediately, these merge on the text encoder's and the TAESD's own +# backend. Upstream applies no LoRA to the TAESD at runtime. +TE_TARGET = "text_encoders.t5xxl.transformer.encoder.block.0.layer.0.SelfAttention.q" +TAE_TARGET = "tae.decoder.layers.11" +LORAS = {"linear": (LINEAR_TARGET, "F16"), "linear-fp8": (LINEAR_TARGET, "F8_E4M3"), + "conv": (CONV_TARGET, "F16"), "te": (TE_TARGET, "F16"), "tae": (TAE_TARGET, "F16")} +LORA_RANK = 4 +LORA_DOWN_VALUES = [((i * 5) % 9 - 4) / 32 for i in range(9)] +LORA_UP_VALUES = [((i * 7) % 11 - 5) / 8 for i in range(11)] +# Auto-fit prints whole MiB, too coarse for the model's own weights. The model +# ignores this extra DiT tensor, and auto-fit counts it: 16 MiB as F16. +AUTO_FIT_PAD = ("model.diffusion_model.auto_fit_pad.weight", [8192, 1024]) -def e4m3_value(byte: int) -> float: - value = 2.0 ** -6 * (1 + (byte & 0x07) / 8) - return -value if byte & 0x80 else value +def encode(values: list[float], dtype: str, count: int) -> bytes: + """Repeats `values` to `count` elements of `dtype`, which must hold them exactly.""" + if dtype in ("F32", "F16"): + code = "f" if dtype == "F32" else "e" + return struct.pack(f"<{count}{code}", *(values[i % len(values)] for i in range(count))) + if dtype == "BF16": + f32 = encode(values, "F32", len(values)) + if any(f32[i:i + 2] != b"\0\0" for i in range(0, len(f32), 4)): + raise ValueError(f"BF16 does not hold every value of {values}") + encoded = b"".join(f32[i + 2:i + 4] for i in range(0, len(f32), 4)) + return (encoded * -(-count // len(values)))[:2 * count] + bias, bits = FP8_TYPES[dtype] + encoded = [] + for value in values: + if value == 0: + encoded.append(0) + continue + fraction, exponent = math.frexp(abs(value)) + field = exponent - 1 + bias + mantissa = (fraction * 2 - 1) * (1 << bits) + if field < 1 or mantissa != int(mantissa): + raise ValueError(f"{dtype} does not hold {value} exactly") + encoded.append((0x80 if value < 0 else 0) | field << bits | int(mantissa)) + return bytes(encoded[i % len(encoded)] for i in range(count)) def read_safetensors(path: Path) -> tuple[dict, bytes]: @@ -76,39 +121,59 @@ def write_safetensors(path: Path, header: dict, blobs: dict[str, bytes]) -> None file.write(blobs[name]) -def write_fp8_pair(source: Path, control: Path, fp8: Path) -> int: - """Writes the F32 control and FP8 copies of `source`; returns the FP8 count.""" +def write_weight_copies(source: Path, copies: dict[str, Path]) -> int: + """Writes a copy of `source` per 2-D weight type; returns the 2-D weight count.""" header, _ = read_safetensors(source) header.pop("__metadata__", None) - control_header, fp8_header, control_blobs, fp8_blobs = {}, {}, {}, {} - pattern_f32 = struct.pack(f"<{len(E4M3_PATTERN)}f", *map(e4m3_value, E4M3_PATTERN)) + outputs = {dtype: ({}, {}) for dtype in copies} converted = 0 for name, info in header.items(): if info["dtype"] != "F32": raise ValueError(f"{name} is {info['dtype']}, not F32") begin, end = info["data_offsets"] count = (end - begin) // 4 - control_header[name] = dict(info) - fp8_header[name] = dict(info) - if len(info["shape"]) == 2 and name.endswith(".weight"): - repeats = -(-count // len(E4M3_PATTERN)) - control_blobs[name] = (pattern_f32 * repeats)[:4 * count] - fp8_blobs[name] = (E4M3_PATTERN * repeats)[:count] - fp8_header[name]["dtype"] = "F8_E4M3" - converted += 1 - else: - repeats = -(-count // 61) - control_blobs[name] = fp8_blobs[name] = (OTHER_PATTERN * repeats)[:4 * count] - write_safetensors(control, control_header, control_blobs) - write_safetensors(fp8, fp8_header, fp8_blobs) + weight = len(info["shape"]) == 2 and name.endswith(".weight") + converted += int(weight) + for dtype, (copy_header, blobs) in outputs.items(): + copy_header[name] = dict(info) + if weight: + copy_header[name]["dtype"] = dtype + blobs[name] = encode(WEIGHT_VALUES, dtype, count) + else: + blobs[name] = (OTHER_PATTERN * -(-count // 61))[:4 * count] + for dtype, path in copies.items(): + write_safetensors(path, *outputs[dtype]) return converted +def write_lora(path: Path, target: str, weight_shape: list[int], down_dtype: str) -> None: + """Writes a LoRA for the linear or convolution weight `target`.""" + out_features, *in_shape = weight_shape + down = f"{target}.lora_down.weight" + up = f"{target}.lora_up.weight" + up_shape = [out_features, LORA_RANK] + [1] * (len(in_shape) - 1) + header = {down: {"dtype": down_dtype, "shape": [LORA_RANK, *in_shape]}, + up: {"dtype": "F16", "shape": up_shape}} + write_safetensors(path, header, { + down: encode(LORA_DOWN_VALUES, down_dtype, LORA_RANK * math.prod(in_shape)), + up: encode(LORA_UP_VALUES, "F16", out_features * LORA_RANK)}) + + +def write_padded(source: Path, padded: Path) -> None: + header, data = read_safetensors(source) + blobs = {name: data[info["data_offsets"][0]:info["data_offsets"][1]] for name, info in header.items()} + name, shape = AUTO_FIT_PAD + dtype = next(info["dtype"] for info in header.values() if len(info["shape"]) == 2) + header[name] = {"dtype": dtype, "shape": shape} + blobs[name] = encode([0.0], dtype, shape[0] * shape[1]) + write_safetensors(padded, header, blobs) + + def mean_difference(first: bytes, second: bytes) -> float: return sum(abs(a - b) for a, b in zip(first, second)) / len(first) -class Fp8RuntimeTest(unittest.TestCase): +class RuntimeTestCase(unittest.TestCase): @classmethod def setUpClass(cls) -> None: cls.skip_reason = "" @@ -138,9 +203,12 @@ def setUpClass(cls) -> None: source_model = work / "model.safetensors" subprocess.run([str(binaries[MODEL_WRITER]), "make-model", str(source_model)], check=True, timeout=600, capture_output=True) - cls.control = work / "control.safetensors" - cls.fp8 = work / "fp8.safetensors" - cls.converted = write_fp8_pair(source_model, cls.control, cls.fp8) + cls.models = {dtype: work / f"{dtype.lower()}.safetensors" for dtype in WEIGHT_TYPES} + cls.converted = write_weight_copies(source_model, cls.models) + header, _ = read_safetensors(cls.models["F32"]) + cls.loras = {name: work / f"lora-{name}.safetensors" for name in LORAS} + for name, (target, dtype) in LORAS.items(): + write_lora(cls.loras[name], target, header[f"{target}.weight"]["shape"], dtype) listing = subprocess.run([str(cls.probe), "devices"], check=True, capture_output=True, text=True, timeout=120).stdout cls.devices = [line.split("\t")[0] for line in listing.splitlines() if line.strip()] @@ -156,46 +224,152 @@ def setUp(self) -> None: self.fail(self.skip_reason) self.skipTest(self.skip_reason) - def generate(self, model: Path, device: str) -> tuple[subprocess.CompletedProcess, bytes]: - output = Path(self.directory.name) / f"{model.stem}-{device}.rgb" + def generate(self, dtype: str, device: str, lora: str | None = None, mode: str = "", + rules: str = "", model: Path | None = None) -> tuple[subprocess.CompletedProcess, bytes]: + """Generates one image; the result's stdout holds the runtime's log.""" + output = Path(self.directory.name) / "image.rgb" output.unlink(missing_ok=True) - result = subprocess.run([str(self.probe), "generate", str(model), device, str(output)], - capture_output=True, text=True, timeout=600) + model = model or self.models[dtype] + command = [str(self.probe), "generate", str(model), device, str(output)] + if lora is not None: + command += [str(self.loras[lora]), mode] + ([rules] if rules else []) + result = subprocess.run(command, capture_output=True, text=True, timeout=600) return result, output.read_bytes() if output.is_file() else b"" - def check_device(self, device: str) -> None: - control, control_image = self.generate(self.control, device) + def generate_control(self, device: str) -> bytes: + control, image = self.generate("F32", device) if control.returncode != 0: if device.startswith("MTL") and os.environ.get("SD_REQUIRE_METAL_GENERATION"): self.fail(f"the F32 control does not generate on {device}:\n{control.stderr}") self.skipTest(f"{device} cannot generate the F32 control:\n{control.stderr[-2000:]}") - fp8, fp8_image = self.generate(self.fp8, device) - self.assertEqual(0, fp8.returncode, f"FP8 weights on {device}:\n{fp8.stderr[-4000:]}") - self.assertGreaterEqual(len(set(control_image)), MIN_DISTINCT_BYTES) - self.assertEqual(len(control_image), len(fp8_image)) - difference = mean_difference(control_image, fp8_image) - self.assertLessEqual(difference, MAX_MEAN_DIFFERENCE, - f"the FP8 image on {device} is not the F32 control's") + self.assertGreaterEqual(len(set(image)), MIN_DISTINCT_BYTES) + return image - def test_model_carries_fp8_weights(self) -> None: - self.assertGreater(self.converted, 0) - header, _ = read_safetensors(self.fp8) - self.assertIn("F8_E4M3", {info.get("dtype") for info in header.values()}) - - def test_fp8_generates_on_the_cpu(self) -> None: - self.check_device("CPU") + def assert_matches(self, expected: bytes, actual: bytes, what: str) -> None: + self.assertEqual(len(expected), len(actual), what) + self.assertLessEqual(mean_difference(expected, actual), MAX_MEAN_DIFFERENCE, what) - def test_fp8_generates_on_each_gpu(self) -> None: + def gpus(self) -> list[str]: gpus = [device for device in self.devices if device != "CPU"] if sys.platform == "darwin" and not any(d.startswith("MTL") for d in gpus) and \ os.environ.get("SD_REQUIRE_METAL_GENERATION"): self.fail(f"no Metal device in {self.devices}") if not gpus: self.skipTest(f"no GPU device in {self.devices}") - for device in gpus: + return gpus + + +class Fp8RuntimeTest(RuntimeTestCase): + def check_device(self, device: str) -> None: + control = self.generate_control(device) + for dtype in FP8_TYPES: + with self.subTest(dtype=dtype): + fp8, image = self.generate(dtype, device) + self.assertEqual(0, fp8.returncode, f"{dtype} weights on {device}:\n{fp8.stderr[-4000:]}") + self.assert_matches(control, image, f"the {dtype} image on {device} is not the F32 control's") + + def test_model_carries_fp8_weights(self) -> None: + self.assertGreater(self.converted, 0) + for dtype in FP8_TYPES: + header, _ = read_safetensors(self.models[dtype]) + self.assertIn(dtype, {info.get("dtype") for info in header.values()}) + + def test_fp8_generates_on_the_cpu(self) -> None: + self.check_device("CPU") + + def test_fp8_generates_on_each_gpu(self) -> None: + for device in self.gpus(): + with self.subTest(device=device): + self.check_device(device) + + def test_auto_fit_plans_fp8_as_f16_on_metal(self) -> None: + metal = [device for device in self.gpus() if device.startswith("MTL")] + if not metal: + self.skipTest(f"no Metal device in {self.devices}") + planned = {} + for dtype in ("F16", *FP8_TYPES): + padded = Path(self.directory.name) / f"padded-{dtype.lower()}.safetensors" + write_padded(self.models[dtype], padded) + result = subprocess.run([str(self.probe), "plan", str(padded), metal[0]], + capture_output=True, text=True, timeout=600) + self.assertEqual(0, result.returncode, f"{dtype} on {metal[0]}:\n{result.stderr[-4000:]}") + match = re.search(r"DiT +params +(\d+) MiB", result.stdout) + self.assertIsNotNone(match, f"no auto-fit plan for {dtype}:\n{result.stdout[-4000:]}") + planned[dtype] = int(match.group(1)) + self.assertGreaterEqual(planned["F16"], 16) + self.assertEqual({dtype: planned["F16"] for dtype in planned}, planned, + "auto-fit plans FP8 weights below the F16 size Metal loads them at") + + +FALLBACK_WARNING = "do not support the immediately LoRA apply mode; using at_runtime instead" +APPLIED = re.compile(r"\((\d+) / (\d+)\) LoRA tensors have been applied") + + +class LoraRuntimeTest(RuntimeTestCase): + def generate_with_lora(self, what: str, *args, **kwargs) -> tuple[str, bytes]: + result, image = self.generate(*args, **kwargs) + self.assertEqual(0, result.returncode, f"{what}:\n{result.stderr[-4000:]}") + self.assertTrue(any(int(applied) > 0 and applied == total + for applied, total in APPLIED.findall(result.stdout)), + f"{what} applies no LoRA tensor:\n{result.stdout[-4000:]}") + return result.stdout, image + + def check_device(self, device: str) -> None: + self.generate_control(device) + for weights in WEIGHT_TYPES: + plain = self.generate(weights, device)[1] + for lora in ("linear", "linear-fp8", "conv"): + what = f"{weights} weights, {lora} LoRA on {device}" + _, expected = self.generate_with_lora(f"{what} at_runtime", weights, device, lora, "at_runtime") + self.assertNotEqual(plain, expected, f"{what} does not change the image") + for mode in ("immediately", "auto"): + with self.subTest(weights=weights, lora=lora, mode=mode): + log, image = self.generate_with_lora(f"{what} {mode}", weights, device, lora, mode) + self.assert_matches(expected, image, f"{what} {mode} is not the at_runtime image") + # Metal cannot add into the F16 weights every model has. + falls_back = mode == "immediately" and (weights in FP8_TYPES or device.startswith("MTL")) + if device == "CPU" or device.startswith("MTL"): + self.assertEqual(falls_back, FALLBACK_WARNING in log, f"{what} {mode}:\n{log[-4000:]}") + + def test_lora_applies_on_the_cpu(self) -> None: + self.check_device("CPU") + + def test_lora_applies_on_each_gpu(self) -> None: + for device in self.gpus(): with self.subTest(device=device): self.check_device(device) + def test_lora_falls_back_for_a_module_on_metal(self) -> None: + metal = [device for device in self.gpus() if device.startswith("MTL")] + if not metal: + self.skipTest(f"no Metal device in {self.devices}") + self.generate_control(metal[0]) + for module, lora in (("te", "te"), ("vae", "tae")): + with self.subTest(module=module): + backend = f"CPU,{module}={metal[0]}" + result, _ = self.generate("F16", backend, lora, "immediately") + self.assertEqual(0, result.returncode, f"{lora} LoRA on {backend}:\n{result.stderr[-4000:]}") + self.assertIn(FALLBACK_WARNING, result.stdout) + + def test_lora_merges_fp8_weights_a_type_rule_loads_as_f16(self) -> None: + # Upstream converts only some weights; FP8 elsewhere would keep the fallback. + fp8_header, fp8_data = read_safetensors(self.models["F8_E4M3"]) + control_header, control_data = read_safetensors(self.models["F32"]) + header, blobs = {}, {} + for name, info in fp8_header.items(): + source, data = (info, fp8_data) if ".blocks." in name else (control_header[name], control_data) + header[name] = dict(source) + blobs[name] = data[source["data_offsets"][0]:source["data_offsets"][1]] + model = Path(self.directory.name) / "fp8-blocks.safetensors" + write_safetensors(model, header, blobs) + self.assertIn("F8_E4M3", {info["dtype"] for info in header.values()}) + _, expected = self.generate_with_lora("at_runtime", "", "CPU", "linear", "at_runtime", + rules="blocks=f16", model=model) + log, image = self.generate_with_lora("immediately", "", "CPU", "linear", "immediately", + rules="blocks=f16", model=model) + self.assertNotIn(FALLBACK_WARNING, log) + self.assert_matches(expected, image, "the merged image is not the at_runtime image") + class PinnedGgmlMetalTest(unittest.TestCase): # The patch can go once ggml-metal multiplies FP8 matrices or stops @@ -210,9 +384,9 @@ def test_pinned_ggml_metal_still_needs_the_patch(self) -> None: start = supports.index("static bool ggml_metal_supports_mul_mat_op(") claim = supports[start:supports.index("\n}\n", start)] kernels = "".join(path.read_text() for path in (metal / "kernels").glob("*.metal")) - self.assertNotIn("F8_E", claim, "ggml-metal now excludes FP8 from MUL_MAT; " - "re-evaluate patches/fp8_metal_f16_load.patch") - self.assertNotRegex(kernels, r"kernel_mul_m[mv]\w*_f8", + self.assertNotRegex(claim, r"(?i)f8_e|fp8", "ggml-metal now excludes FP8 from MUL_MAT; " + "re-evaluate patches/fp8_metal_f16_load.patch") + self.assertNotRegex(kernels, r"(?i)kernel_mul_m[mv]\w*_fp?8", "ggml-metal now has an FP8 matrix kernel; " "re-evaluate patches/fp8_metal_f16_load.patch")