From 5077a4786066241f62f29d0bcf17de4693738c6a Mon Sep 17 00:00:00 2001 From: Jhin Lee Date: Tue, 27 Oct 2026 09:54:50 -0400 Subject: [PATCH 1/3] fix: apply LoRAs at runtime where a merge crashes, and size FP8 on Metal The immediately LoRA apply mode, which auto selects for unquantized weights, merges each diff into the weight in place. Merging into FP8 weights segfaults ggml-cpu (the F32 cast is a MUL_MAT with no FP8 vec_dot, and there is no F32 to FP8 conversion), and merging an F32 diff into F16 weights aborts Metal (ggml-metal ADD needs equal types; the scheduler cannot move an op on a pre-allocated weight). lora_immediate_fallback.patch selects at_runtime for FP8 weights, and for F16 weights on a backend whose supports_op rejects that add, with a warning when immediately was asked. An FP8 LoRA tensor crashed ggml-cpu in every mode for the same missing kernel; fp8_metal_f16_load.patch now loads FP8 LoRA tensors as F16 on every backend. It also makes auto-fit count FP8 weights at F16 size when Metal is present, since Metal holds them as F16. tests/test_fp8_runtime.py now covers E5M2, applies a synthetic LoRA in each mode to F32, F16 and FP8 weights on every device, and checks the auto-fit plan on Metal. The retirement check also catches FP8 helpers and _fp8 kernel names. Refs #31 --- CMakeLists.txt | 3 +- README.md | 9 +- patches/README.md | 3 +- patches/fp8_metal_f16_load.patch | 53 +++++- patches/lora_immediate_fallback.patch | 80 +++++++++ tests/native/fp8_runtime_test.cpp | 55 +++++- tests/test_fp8_runtime.py | 232 +++++++++++++++++++------- 7 files changed, 358 insertions(+), 77 deletions(-) create mode 100644 patches/lora_immediate_fallback.patch 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..23f5ceb 100644 --- a/README.md +++ b/README.md @@ -610,10 +610,13 @@ 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 it cannot be merged: into +FP8 weights, and into F16 weights on Metal. 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..4a59e2c 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). Those weights 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..145f618 --- /dev/null +++ b/patches/lora_immediate_fallback.patch @@ -0,0 +1,80 @@ +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,62 @@ 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. ++static bool backend_adds_f32_into_f16(ggml_backend_t backend) { ++ 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; ++ bool have_f16_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; ++ } else if (load_type == GGML_TYPE_F16) { ++ have_f16_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. ++ const char* immediate_lora_blocker = nullptr; ++ if (have_fp8_weight) { ++ immediate_lora_blocker = "FP8 weights"; ++ } else if (have_f16_weight) { ++ 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 = "F16 weights on this 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 != nullptr; + for (const auto& [type, _] : wtype_stat) { + if (ggml_is_quantized(type)) { + have_quantized_weight = true; +@@ -1054,6 +1090,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 != nullptr) { ++ LOG_WARN("%s do not support the immediately LoRA apply mode; using at_runtime instead", ++ immediate_lora_blocker); ++ 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..24758c6 100644 --- a/tests/native/fp8_runtime_test.cpp +++ b/tests/native/fp8_runtime_test.cpp @@ -2,13 +2,17 @@ // 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. `plan` loads the model +// and prints the log, which holds upstream's auto-fit plan. #include "sd_dart_wrapper.h" +#include #include #include #include @@ -25,7 +29,7 @@ int list_devices() { return 0; } -int generate(const char* model, const char* backend, const char* output) { +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 +37,33 @@ 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); + } + 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()); + } + return context != nullptr ? 0 : 3; +} + +int generate(const char* model, const char* backend, const char* output, const char* lora, + const char* mode) { + 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; + 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 +80,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 +119,16 @@ 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]); + if (argc == 4 && std::strcmp(argv[1], "plan") == 0) { + return plan(argv[2], argv[3]); + } + if ((argc == 5 || argc == 7) && std::strcmp(argv[1], "generate") == 0) { + return generate(argv[2], argv[3], argv[4], argc == 7 ? argv[5] : nullptr, + argc == 7 ? argv[6] : nullptr); } - std::fprintf(stderr, "usage: %s devices | generate \n", argv[0]); + std::fprintf(stderr, + "usage: %s devices | generate [ ] | " + "plan \n", + argv[0]); return 2; } diff --git a/tests/test_fp8_runtime.py b/tests/test_fp8_runtime.py index b2f5709..360f4f9 100644 --- a/tests/test_fp8_runtime.py +++ b/tests/test_fp8_runtime.py @@ -1,14 +1,18 @@ -"""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 and F16 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 one small +adapter in every apply mode to F32, F16 and FP8 weights. 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 for F16 weights (#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 +24,10 @@ from __future__ import annotations import json +import math import os import platform +import re import shutil import struct import subprocess @@ -39,21 +45,46 @@ 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.05 on the CPU and 0 on Metal. MAX_MEAN_DIFFERENCE = 0.5 +# The LoRA changes the image by 1.6 per byte on an M4 Max. +MIN_LORA_EFFECT = 1.0 # 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_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))) +LORA_TARGET = "model.diffusion_model.final_layer.linear" +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))) + 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 +107,57 @@ 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, weight_shape: list[int], down_dtype: str) -> None: + out_features, in_features = weight_shape + down = f"{LORA_TARGET}.lora_down.weight" + up = f"{LORA_TARGET}.lora_up.weight" + header = {down: {"dtype": down_dtype, "shape": [LORA_RANK, in_features]}, + up: {"dtype": "F16", "shape": [out_features, LORA_RANK]}} + write_safetensors(path, header, { + down: encode(LORA_DOWN_VALUES, down_dtype, LORA_RANK * in_features), + 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 +187,13 @@ 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 ("F32", "F16", *FP8_TYPES)} + cls.converted = write_weight_copies(source_model, cls.models) + header, _ = read_safetensors(cls.models["F32"]) + cls.loras = {dtype: work / f"lora-{dtype.lower()}.safetensors" for dtype in ("F16", "F8_E4M3")} + for dtype, path in cls.loras.items(): + write_lora(path, header[f"{LORA_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,43 +209,104 @@ 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 = "") -> tuple[subprocess.CompletedProcess, bytes]: + 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) + command = [str(self.probe), "generate", str(self.models[dtype]), device, str(output)] + if lora is not None: + command += [str(self.loras[lora]), mode] + 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") - - 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()}) + self.assertGreaterEqual(len(set(image)), MIN_DISTINCT_BYTES) + return image - 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") + + +class LoraRuntimeTest(RuntimeTestCase): + def check_device(self, device: str) -> None: + self.generate_control(device) + for weights in ("F32", "F16", "F8_E4M3"): + plain = self.generate(weights, device)[1] + reference, expected = self.generate(weights, device, "F16", "at_runtime") + self.assertEqual(0, reference.returncode, + f"{weights} weights, F16 LoRA at_runtime on {device}:\n{reference.stderr[-4000:]}") + self.assertGreaterEqual(mean_difference(plain, expected), MIN_LORA_EFFECT, + f"the LoRA does not change the {weights} image on {device}") + for lora, mode in (("F16", "immediately"), ("F16", "auto"), + ("F8_E4M3", "at_runtime"), ("F8_E4M3", "immediately")): + with self.subTest(weights=weights, lora=lora, mode=mode): + result, image = self.generate(weights, device, lora, mode) + what = f"{weights} weights, {lora} LoRA {mode} on {device}" + self.assertEqual(0, result.returncode, f"{what}:\n{result.stderr[-4000:]}") + self.assert_matches(expected, image, f"{what} is not the at_runtime image") + + 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) @@ -210,9 +324,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") From dc48762827c9b9c210b385603c4ec6aba95ee323 Mon Sep 17 00:00:00 2001 From: Jhin Lee Date: Tue, 27 Oct 2026 11:44:32 -0400 Subject: [PATCH 2/3] fix: fall back from immediately LoRA on Metal whatever the stored weight type Upstream creates Conv2d, grouped Conv2d and Conv3d weights as F16 whatever the file holds, so an F32 or BF16 checkpoint with a LoRA on a convolution still aborted Metal under immediately or auto. The fallback no longer depends on stored F16 weights: it probes the TE, diffusion and VAE backends of every model without FP8 weights. A backend whose default buffers are host memory passes, since the scheduler runs the add on the CPU there. The probe prints the runtime log, so the tests check the fallback warning and that the LoRA applied. They add a convolution LoRA, BF16 and E5M2 weights, the text encoder and TAESD on Metal behind a CPU diffusion model, and FP8 weights a type rule loads as F16. Refs #31 --- README.md | 5 +- patches/README.md | 2 +- patches/lora_immediate_fallback.patch | 26 ++--- tests/native/fp8_runtime_test.cpp | 52 ++++++---- tests/test_fp8_runtime.py | 139 ++++++++++++++++++-------- 5 files changed, 147 insertions(+), 77 deletions(-) diff --git a/README.md b/README.md index 23f5ceb..b159511 100644 --- a/README.md +++ b/README.md @@ -615,8 +615,9 @@ the [`pickle_line_span.patch`](patches/README.md) for legacy checkpoints above 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 it cannot be merged: into -FP8 weights, and into F16 weights on Metal. +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 4a59e2c..a44241e 100644 --- a/patches/README.md +++ b/patches/README.md @@ -14,4 +14,4 @@ an upstream issue or pull request to retire it. | --- | --- | --- | --- | | `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`, `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). Those weights 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` | +| `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/lora_immediate_fallback.patch b/patches/lora_immediate_fallback.patch index 145f618..413462d 100644 --- a/patches/lora_immediate_fallback.patch +++ b/patches/lora_immediate_fallback.patch @@ -1,14 +1,19 @@ 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,62 @@ bool StableDiffusionGGML::tae_preview_only() const { +@@ -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. ++// 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) { @@ -28,7 +33,6 @@ diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.c bool have_int8_tensorwise = false; bool have_quant_type_rule = false; + bool have_fp8_weight = false; -+ bool have_f16_weight = false; for (const auto& [_, tensor_storage] : model_loader.get_tensor_storage_map()) { if (tensor_storage.is_int8_tensorwise) { have_int8_tensorwise = true; @@ -37,8 +41,6 @@ diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.c + : tensor_storage.type; + if (load_type == GGML_TYPE_F8_E4M3 || load_type == GGML_TYPE_F8_E5M2) { + have_fp8_weight = true; -+ } else if (load_type == GGML_TYPE_F16) { -+ have_f16_weight = true; + } if (tensor_storage.expected_type < GGML_TYPE_COUNT && ggml_is_quantized(tensor_storage.expected_type)) { have_quant_type_rule = true; @@ -48,14 +50,14 @@ diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.c + } + // 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. -+ const char* immediate_lora_blocker = nullptr; ++ std::string immediate_lora_blocker; + if (have_fp8_weight) { + immediate_lora_blocker = "FP8 weights"; -+ } else if (have_f16_weight) { ++ } 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 = "F16 weights on this backend"; ++ immediate_lora_blocker = std::string("F16 weights on ") + ggml_backend_name(backend); + break; + } } @@ -63,17 +65,17 @@ diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.c 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 != nullptr; ++ 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 +1090,10 @@ void StableDiffusionGGML::configure_weight_loading() { +@@ -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 != nullptr) { ++ } else if (!immediate_lora_blocker.empty()) { + LOG_WARN("%s do not support the immediately LoRA apply mode; using at_runtime instead", -+ immediate_lora_blocker); ++ immediate_lora_blocker.c_str()); + apply_lora_immediately = false; } else if (row_split_active()) { LOG_WARN( diff --git a/tests/native/fp8_runtime_test.cpp b/tests/native/fp8_runtime_test.cpp index 24758c6..c491acc 100644 --- a/tests/native/fp8_runtime_test.cpp +++ b/tests/native/fp8_runtime_test.cpp @@ -2,13 +2,14 @@ // 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. is a -// LoRA apply mode name: auto, immediately or at_runtime. `plan` loads the model -// and prints the log, which holds upstream's auto-fit plan. +// 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" @@ -29,6 +30,14 @@ int list_devices() { return 0; } +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; @@ -40,22 +49,18 @@ int plan(const char* model, const char* backend) { if (context != nullptr) { sd_dart_exit_free(context); } - 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()); - } return context != nullptr ? 0 : 3; } int generate(const char* model, const char* backend, const char* output, const char* lora, - const char* mode) { + 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.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) { @@ -119,16 +124,19 @@ int main(int argc, char** argv) { if (argc == 2 && std::strcmp(argv[1], "devices") == 0) { return list_devices(); } + int status = 2; if (argc == 4 && std::strcmp(argv[1], "plan") == 0) { - return plan(argv[2], argv[3]); - } - if ((argc == 5 || argc == 7) && std::strcmp(argv[1], "generate") == 0) { - return generate(argv[2], argv[3], argv[4], argc == 7 ? argv[5] : nullptr, - argc == 7 ? argv[6] : nullptr); + 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 [ ] | " - "plan \n", - argv[0]); - return 2; + print_log(); + return status; } diff --git a/tests/test_fp8_runtime.py b/tests/test_fp8_runtime.py index 360f4f9..14c5278 100644 --- a/tests/test_fp8_runtime.py +++ b/tests/test_fp8_runtime.py @@ -2,17 +2,18 @@ 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 that E4M3, E5M2 and F16 represent exactly, kept as F32 in a control +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 an FP8 generation aborts the process (#29). The LoRA tests apply one small -adapter in every apply mode to F32, F16 and FP8 weights. 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 for F16 weights (#31). 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 @@ -46,19 +47,26 @@ 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 two images of one device -# that should match. An M4 Max measured at most 0.05 on the CPU and 0 on Metal. +# that should match. An M4 Max measured at most 0.11 on the CPU and 0 on Metal. MAX_MEAN_DIFFERENCE = 0.5 -# The LoRA changes the image by 1.6 per byte on an M4 Max. -MIN_LORA_EFFECT = 1.0 # Fewer distinct bytes would mean a flat image that hides any difference. MIN_DISTINCT_BYTES = 16 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))) -LORA_TARGET = "model.diffusion_model.final_layer.linear" +# 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)] @@ -72,6 +80,12 @@ def encode(values: list[float], dtype: str, count: int) -> bytes: 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: @@ -132,14 +146,16 @@ def write_weight_copies(source: Path, copies: dict[str, Path]) -> int: return converted -def write_lora(path: Path, weight_shape: list[int], down_dtype: str) -> None: - out_features, in_features = weight_shape - down = f"{LORA_TARGET}.lora_down.weight" - up = f"{LORA_TARGET}.lora_up.weight" - header = {down: {"dtype": down_dtype, "shape": [LORA_RANK, in_features]}, - up: {"dtype": "F16", "shape": [out_features, LORA_RANK]}} +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 * in_features), + down: encode(LORA_DOWN_VALUES, down_dtype, LORA_RANK * math.prod(in_shape)), up: encode(LORA_UP_VALUES, "F16", out_features * LORA_RANK)}) @@ -187,13 +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.models = {dtype: work / f"{dtype.lower()}.safetensors" - for dtype in ("F32", "F16", *FP8_TYPES)} + 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 = {dtype: work / f"lora-{dtype.lower()}.safetensors" for dtype in ("F16", "F8_E4M3")} - for dtype, path in cls.loras.items(): - write_lora(path, header[f"{LORA_TARGET}.weight"]["shape"], dtype) + 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()] @@ -209,13 +224,15 @@ def setUp(self) -> None: self.fail(self.skip_reason) self.skipTest(self.skip_reason) - def generate(self, dtype: str, device: str, lora: str | None = None, - mode: str = "") -> tuple[subprocess.CompletedProcess, bytes]: + 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) - command = [str(self.probe), "generate", str(self.models[dtype]), device, str(output)] + 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] + 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"" @@ -284,23 +301,35 @@ def test_auto_fit_plans_fp8_as_f16_on_metal(self) -> None: "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 ("F32", "F16", "F8_E4M3"): + for weights in WEIGHT_TYPES: plain = self.generate(weights, device)[1] - reference, expected = self.generate(weights, device, "F16", "at_runtime") - self.assertEqual(0, reference.returncode, - f"{weights} weights, F16 LoRA at_runtime on {device}:\n{reference.stderr[-4000:]}") - self.assertGreaterEqual(mean_difference(plain, expected), MIN_LORA_EFFECT, - f"the LoRA does not change the {weights} image on {device}") - for lora, mode in (("F16", "immediately"), ("F16", "auto"), - ("F8_E4M3", "at_runtime"), ("F8_E4M3", "immediately")): - with self.subTest(weights=weights, lora=lora, mode=mode): - result, image = self.generate(weights, device, lora, mode) - what = f"{weights} weights, {lora} LoRA {mode} on {device}" - self.assertEqual(0, result.returncode, f"{what}:\n{result.stderr[-4000:]}") - self.assert_matches(expected, image, f"{what} is not the at_runtime image") + 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") @@ -310,6 +339,36 @@ def test_lora_applies_on_each_gpu(self) -> None: 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}") + 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 From 6c4d6998b1be55ae512834533a52ecb418584c8a Mon Sep 17 00:00:00 2001 From: Jhin Lee Date: Tue, 27 Oct 2026 12:28:12 -0400 Subject: [PATCH 3/3] test: skip the Metal module LoRA test where Metal cannot generate GitHub's macOS runners list MTL0 but cannot generate on it; the other Metal tests skip on the F32 control there. Refs #31 --- tests/test_fp8_runtime.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_fp8_runtime.py b/tests/test_fp8_runtime.py index 14c5278..07123a0 100644 --- a/tests/test_fp8_runtime.py +++ b/tests/test_fp8_runtime.py @@ -343,6 +343,7 @@ 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]}"