From 8688746a123e4c0f62ca30c2fb8cb84e3d0c1318 Mon Sep 17 00:00:00 2001 From: Sam McLeod Date: Sat, 6 Jun 2026 09:57:32 +1000 Subject: [PATCH 1/5] Add Apple Silicon (MPS) support to the FP8 inference path MPS can neither store nor cast float8_e4m3fn, and supports neither float64 nor ndtri/expit, so the weight-only FP8 path crashed twice on MPS: first dequantizing Linear weights, then in the logit-normal sampler. Both are now handled, with the CUDA/CPU paths unchanged by construction. - quantized_loading.py: add device_supports_fp8(); on MPS, Fp8Linear holds an already-dequantized bf16 weight (no scale buffer) and load_fp8_state_dict dequantizes each FP8 weight via its per-row scale at load time, dropping the now-unused .weight_scale keys. CUDA/CPU keep the float8 storage path (store_fp8=True) byte-for-byte. - scheduler.py: run the float64 warp (ndtri/expit) on CPU for MPS only, returning float32 on the caller's device; CUDA/CPU keep the original on-device path. - pipeline_ideogram4.py: thread the target device into swap_linears_to_fp8 at both call sites. --- src/ideogram4/pipeline_ideogram4.py | 4 +- src/ideogram4/quantized_loading.py | 73 ++++++++++++++++++++++++----- src/ideogram4/scheduler.py | 12 ++++- 3 files changed, 74 insertions(+), 15 deletions(-) diff --git a/src/ideogram4/pipeline_ideogram4.py b/src/ideogram4/pipeline_ideogram4.py index 0f16e12..148e9f2 100644 --- a/src/ideogram4/pipeline_ideogram4.py +++ b/src/ideogram4/pipeline_ideogram4.py @@ -83,7 +83,7 @@ def _load_fp8_text_encoder( ) model = AutoModel.from_config(config, trust_remote_code=True) state_dict = _load_subfolder_state_dict(repo_id, text_encoder_subfolder, "model") - swap_linears_to_fp8(model, state_dict, compute_dtype=dtype) + swap_linears_to_fp8(model, state_dict, compute_dtype=dtype, device=device) # assign=True so unquantized params take the loaded dtype and the computed # rotary buffers (absent from the checkpoint) survive; tied weights, if any, # surface as benign missing keys. @@ -168,7 +168,7 @@ def _build_transformer( # Weight-only FP8: cast the unquantized params to the compute dtype first, # then swap in Fp8Linear layers (which keep their weights as float8). model.to(dtype) - swap_linears_to_fp8(model, state_dict, compute_dtype=dtype) + swap_linears_to_fp8(model, state_dict, compute_dtype=dtype, device=device) load_fp8_state_dict(model, state_dict, device=device, dtype=dtype) else: model.load_state_dict(state_dict) diff --git a/src/ideogram4/quantized_loading.py b/src/ideogram4/quantized_loading.py index 99d1ce2..9207791 100644 --- a/src/ideogram4/quantized_loading.py +++ b/src/ideogram4/quantized_loading.py @@ -161,16 +161,32 @@ def is_fp8_state_dict(state_dict: dict[str, torch.Tensor]) -> bool: ) +def device_supports_fp8(device: torch.device) -> bool: + """Whether ``device`` can store and cast ``float8_e4m3fn`` tensors. + + PyTorch's MPS (Apple Silicon) backend has no float8 support: it can neither + hold the dtype nor convert it, so any ``.to(mps)`` / ``.to(dtype)`` on a float8 + tensor raises. On MPS the FP8 weights must be dequantized to the compute dtype + at load time instead. CUDA and CPU both handle float8 storage and casting. + """ + return torch.device(device).type != "mps" + + class Fp8Linear(nn.Module): """Linear layer holding an e4m3 float8 weight + per-row float32 scale. The weight and scale are registered as buffers (not parameters) so they load via ``load_state_dict`` and are excluded from optimizer/grad machinery. The dequantized matmul runs in ``compute_dtype``. + + When ``store_fp8`` is False the layer holds an already-dequantized weight in + ``compute_dtype`` and no scale buffer; this is the path for devices that can't + store float8 (MPS). The half-size checkpoint download is still used; only the + in-memory weight is expanded. """ weight: torch.Tensor - weight_scale: torch.Tensor + weight_scale: torch.Tensor | None bias: torch.Tensor | None def __init__( @@ -179,23 +195,39 @@ def __init__( out_features: int, bias: bool, compute_dtype: torch.dtype, + *, + store_fp8: bool = True, ) -> None: super().__init__() self.in_features = in_features self.out_features = out_features self.compute_dtype = compute_dtype - self.register_buffer( - "weight", - torch.empty(out_features, in_features, dtype=FP8_WEIGHT_DTYPE), - ) - self.register_buffer("weight_scale", torch.empty(out_features, dtype=torch.float32)) + self.store_fp8 = store_fp8 + if store_fp8: + self.register_buffer( + "weight", + torch.empty(out_features, in_features, dtype=FP8_WEIGHT_DTYPE), + ) + self.register_buffer( + "weight_scale", torch.empty(out_features, dtype=torch.float32) + ) + else: + self.register_buffer( + "weight", + torch.empty(out_features, in_features, dtype=compute_dtype), + ) + self.weight_scale = None if bias: self.register_buffer("bias", torch.empty(out_features, dtype=compute_dtype)) else: self.bias = None def forward(self, x: torch.Tensor) -> torch.Tensor: - w = self.weight.to(x.dtype) * self.weight_scale.to(x.dtype).unsqueeze(1) + if self.store_fp8: + assert self.weight_scale is not None + w = self.weight.to(x.dtype) * self.weight_scale.to(x.dtype).unsqueeze(1) + else: + w = self.weight.to(x.dtype) bias = self.bias.to(x.dtype) if self.bias is not None else None return F.linear(x, w, bias) @@ -204,6 +236,7 @@ def swap_linears_to_fp8( module: nn.Module, state_dict: dict[str, torch.Tensor], compute_dtype: torch.dtype, + device: torch.device, *, prefix: str = "", ) -> None: @@ -211,8 +244,10 @@ def swap_linears_to_fp8( Gating on the presence of ``.weight_scale`` means only layers that were actually quantized at save time are swapped; everything else loads normally in - the compute dtype. + the compute dtype. On devices that can't store float8 (MPS) the swapped layers + hold a dequantized ``compute_dtype`` weight instead (see ``Fp8Linear``). """ + store_fp8 = device_supports_fp8(device) for name, child in list(module.named_children()): child_prefix = f"{prefix}{name}" if ( @@ -226,10 +261,13 @@ def swap_linears_to_fp8( child.out_features, bias=child.bias is not None, compute_dtype=compute_dtype, + store_fp8=store_fp8, ), ) else: - swap_linears_to_fp8(child, state_dict, compute_dtype, prefix=f"{child_prefix}.") + swap_linears_to_fp8( + child, state_dict, compute_dtype, device, prefix=f"{child_prefix}." + ) def load_fp8_state_dict( @@ -255,13 +293,26 @@ def load_fp8_state_dict( ``strict=False`` downgrades missing keys to a warning (e.g. tied weights that a ``transformers`` model resolves itself); unexpected keys always raise. + + On devices that can't store float8 (MPS) the FP8 weights are dequantized to + ``dtype`` here using their per-row scale, and the now-unused ``.weight_scale`` + entries are dropped to match the dequantized ``Fp8Linear`` layout. """ + store_fp8 = device_supports_fp8(device) prepared: dict[str, torch.Tensor] = {} for k, v in state_dict.items(): if v.dtype == FP8_WEIGHT_DTYPE: - prepared[k] = v.to(device=device) + if store_fp8: + prepared[k] = v.to(device=device) + else: + # MPS can't cast float8, so dequantize on CPU before moving across. + scale = state_dict[k[: -len(".weight")] + FP8_SCALE_SUFFIX] + w = v.to(torch.float32) * scale.to(torch.float32).unsqueeze(1) + prepared[k] = w.to(device=device, dtype=dtype) elif k.endswith(FP8_SCALE_SUFFIX): - prepared[k] = v.to(device=device, dtype=torch.float32) + if store_fp8: + prepared[k] = v.to(device=device, dtype=torch.float32) + # else: folded into the dequantized weight above; drop the scale key. elif v.is_floating_point(): prepared[k] = v.to(device=device, dtype=dtype) else: diff --git a/src/ideogram4/scheduler.py b/src/ideogram4/scheduler.py index d84b46b..f2682b4 100644 --- a/src/ideogram4/scheduler.py +++ b/src/ideogram4/scheduler.py @@ -16,14 +16,22 @@ class LogitNormalSchedule: logsnr_max: float = 18.0 def __call__(self, t: torch.Tensor) -> torch.Tensor: - t = t.to(torch.float64) + device = t.device + # The float64 warp (ndtri/expit) needs precision at the tails. MPS supports + # neither float64 nor these special functions, so there the warp runs on CPU + # (a fused `.to(cpu, float64)` would still cast on the MPS side first and + # raise, hence the split). Other devices keep the original on-device path. + if device.type == "mps": + t = t.cpu().to(torch.float64) + else: + t = t.to(torch.float64) z = torch.special.ndtri(t) y = self.mean + self.std * z t_ = torch.special.expit(y) t_ = 1 - t_ t_min = 1.0 / (1 + math.exp(0.5 * self.logsnr_max)) t_max = 1.0 / (1 + math.exp(0.5 * self.logsnr_min)) - return t_.clamp(t_min, t_max).to(torch.float32) + return t_.clamp(t_min, t_max).to(device=device, dtype=torch.float32) def get_schedule_for_resolution( From df8f08bf5fbf746c0390a99e4ee0582f416cbf19 Mon Sep 17 00:00:00 2001 From: Sam McLeod Date: Sat, 6 Jun 2026 10:09:45 +1000 Subject: [PATCH 2/5] Address PR review: runtime fp8 capability probe and explicit CPU dequant - device_supports_fp8 now probes the backend at runtime (cached per device type) instead of hard-coding "not mps", so it adapts to other backends that may lack float8 storage/casting. - load_fp8_state_dict dequant path moves the fp8 weight and scale to CPU explicitly before the float32 cast and multiply, matching the comment and guarding against non-CPU state-dict tensors. Scheduler CPU-roundtrip review comments were left as-is: the warp runs on a single-element tensor that the caller immediately scalarises via .item(), so the transfer is negligible (measured ~6.8s/step, transformer-bound, at V4_QUALITY_48). --- src/ideogram4/quantized_loading.py | 26 +++++++++++++++++++------- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/src/ideogram4/quantized_loading.py b/src/ideogram4/quantized_loading.py index 9207791..251fc01 100644 --- a/src/ideogram4/quantized_loading.py +++ b/src/ideogram4/quantized_loading.py @@ -1,5 +1,6 @@ from __future__ import annotations +import functools import warnings import bitsandbytes as bnb @@ -161,15 +162,25 @@ def is_fp8_state_dict(state_dict: dict[str, torch.Tensor]) -> bool: ) +@functools.lru_cache(maxsize=None) +def _device_type_supports_fp8(device_type: str) -> bool: + try: + torch.zeros(1, dtype=FP8_WEIGHT_DTYPE, device=device_type).to(torch.float32) + except Exception: + return False + return True + + def device_supports_fp8(device: torch.device) -> bool: """Whether ``device`` can store and cast ``float8_e4m3fn`` tensors. - PyTorch's MPS (Apple Silicon) backend has no float8 support: it can neither - hold the dtype nor convert it, so any ``.to(mps)`` / ``.to(dtype)`` on a float8 - tensor raises. On MPS the FP8 weights must be dequantized to the compute dtype - at load time instead. CUDA and CPU both handle float8 storage and casting. + Probed at runtime (cached per device type) rather than hard-coded, so the + answer matches the actual backend. PyTorch's MPS (Apple Silicon) backend has no + float8 support - it can neither hold the dtype nor convert it - so there the + FP8 weights must be dequantized to the compute dtype at load time instead. CUDA + and CPU both pass. """ - return torch.device(device).type != "mps" + return _device_type_supports_fp8(torch.device(device).type) class Fp8Linear(nn.Module): @@ -305,9 +316,10 @@ def load_fp8_state_dict( if store_fp8: prepared[k] = v.to(device=device) else: - # MPS can't cast float8, so dequantize on CPU before moving across. + # MPS can't cast float8, so dequantize on CPU (where the fp8 weights are + # loaded) before moving the result across to the target device. scale = state_dict[k[: -len(".weight")] + FP8_SCALE_SUFFIX] - w = v.to(torch.float32) * scale.to(torch.float32).unsqueeze(1) + w = v.cpu().to(torch.float32) * scale.cpu().to(torch.float32).unsqueeze(1) prepared[k] = w.to(device=device, dtype=dtype) elif k.endswith(FP8_SCALE_SUFFIX): if store_fp8: From 9ad4c69cb012cb682ae12c454c31ff4752718765 Mon Sep 17 00:00:00 2001 From: Sam McLeod Date: Sat, 6 Jun 2026 10:47:50 +1000 Subject: [PATCH 3/5] Precompute schedule values once per generation instead of per step step_intervals is fixed for the whole sampling loop, so warp it through LogitNormalSchedule once before the loop and index the results, rather than calling the schedule twice per step. Output is byte-identical; this hoists the loop-invariant scalar work (and its per-step host syncs) out of the inner loop. Addresses the scheduler.py review comments on the MPS CPU roundtrip. --- src/ideogram4/pipeline_ideogram4.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/ideogram4/pipeline_ideogram4.py b/src/ideogram4/pipeline_ideogram4.py index 148e9f2..1551e6b 100644 --- a/src/ideogram4/pipeline_ideogram4.py +++ b/src/ideogram4/pipeline_ideogram4.py @@ -584,9 +584,16 @@ def __call__( device=self.device, ) + # step_intervals is fixed for the whole generation, so warp it through the + # schedule once here rather than twice per step inside the loop. + schedule_values = [ + float(schedule(step_intervals[j].unsqueeze(0)).item()) + for j in range(num_steps + 1) + ] + for i in range(num_steps - 1, -1, -1): - t_val = float(schedule(step_intervals[i + 1].unsqueeze(0)).item()) - s_val = float(schedule(step_intervals[i].unsqueeze(0)).item()) + t_val = schedule_values[i + 1] + s_val = schedule_values[i] t = torch.full((batch_size,), t_val, dtype=torch.float32, device=self.device) pos_z = torch.cat([text_z_padding, z], dim=1) From bae952f9455dbf876da5ab11075b59a0268fcbda Mon Sep 17 00:00:00 2001 From: Sam McLeod Date: Sat, 6 Jun 2026 10:52:49 +1000 Subject: [PATCH 4/5] Address PR review: vectorize schedule precompute, per-device fp8 probe cache, explicit missing-scale error - pipeline_ideogram4: warp the whole step_intervals tensor through the schedule in one vectorized call (.tolist()) instead of per-element calls; the values are byte-identical so output is unchanged. - quantized_loading: key the fp8 capability probe cache by concrete device (e.g. cuda:0) rather than device type, so a heterogeneous multi-device setup is probed individually. - quantized_loading: raise a clear RuntimeError naming the weight and the missing scale key when an FP8 weight has no matching .weight_scale, instead of a bare KeyError. --- src/ideogram4/pipeline_ideogram4.py | 10 ++++------ src/ideogram4/quantized_loading.py | 23 ++++++++++++++--------- 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/src/ideogram4/pipeline_ideogram4.py b/src/ideogram4/pipeline_ideogram4.py index 1551e6b..210dd4f 100644 --- a/src/ideogram4/pipeline_ideogram4.py +++ b/src/ideogram4/pipeline_ideogram4.py @@ -584,12 +584,10 @@ def __call__( device=self.device, ) - # step_intervals is fixed for the whole generation, so warp it through the - # schedule once here rather than twice per step inside the loop. - schedule_values = [ - float(schedule(step_intervals[j].unsqueeze(0)).item()) - for j in range(num_steps + 1) - ] + # step_intervals is fixed for the whole generation, so warp the whole tensor + # through the schedule in one vectorized call (one device transfer) rather + # than calling it per element, twice per step, inside the loop. + schedule_values = schedule(step_intervals).tolist() for i in range(num_steps - 1, -1, -1): t_val = schedule_values[i + 1] diff --git a/src/ideogram4/quantized_loading.py b/src/ideogram4/quantized_loading.py index 251fc01..0d884f3 100644 --- a/src/ideogram4/quantized_loading.py +++ b/src/ideogram4/quantized_loading.py @@ -163,9 +163,9 @@ def is_fp8_state_dict(state_dict: dict[str, torch.Tensor]) -> bool: @functools.lru_cache(maxsize=None) -def _device_type_supports_fp8(device_type: str) -> bool: +def _probe_fp8_support(device_str: str) -> bool: try: - torch.zeros(1, dtype=FP8_WEIGHT_DTYPE, device=device_type).to(torch.float32) + torch.zeros(1, dtype=FP8_WEIGHT_DTYPE, device=device_str).to(torch.float32) except Exception: return False return True @@ -174,13 +174,13 @@ def _device_type_supports_fp8(device_type: str) -> bool: def device_supports_fp8(device: torch.device) -> bool: """Whether ``device`` can store and cast ``float8_e4m3fn`` tensors. - Probed at runtime (cached per device type) rather than hard-coded, so the - answer matches the actual backend. PyTorch's MPS (Apple Silicon) backend has no - float8 support - it can neither hold the dtype nor convert it - so there the - FP8 weights must be dequantized to the compute dtype at load time instead. CUDA - and CPU both pass. + Probed at runtime (cached per concrete device, e.g. ``cuda:0``, so a + heterogeneous multi-device setup is checked individually) rather than + hard-coded. PyTorch's MPS (Apple Silicon) backend has no float8 support - it + can neither hold the dtype nor convert it - so there the FP8 weights must be + dequantized to the compute dtype at load time instead. CUDA and CPU both pass. """ - return _device_type_supports_fp8(torch.device(device).type) + return _probe_fp8_support(str(torch.device(device))) class Fp8Linear(nn.Module): @@ -318,7 +318,12 @@ def load_fp8_state_dict( else: # MPS can't cast float8, so dequantize on CPU (where the fp8 weights are # loaded) before moving the result across to the target device. - scale = state_dict[k[: -len(".weight")] + FP8_SCALE_SUFFIX] + scale_key = k[: -len(".weight")] + FP8_SCALE_SUFFIX + if scale_key not in state_dict: + raise RuntimeError( + f"FP8 weight {k!r} has no matching scale {scale_key!r} in the checkpoint" + ) + scale = state_dict[scale_key] w = v.cpu().to(torch.float32) * scale.cpu().to(torch.float32).unsqueeze(1) prepared[k] = w.to(device=device, dtype=dtype) elif k.endswith(FP8_SCALE_SUFFIX): From ef5ffdedd82d7c03453f60361112c47369a9899b Mon Sep 17 00:00:00 2001 From: Sam McLeod Date: Sat, 6 Jun 2026 18:13:56 +1000 Subject: [PATCH 5/5] Address PR review round 3: schedule CPU input, probe docstring, robust scale-key guard - pipeline_ideogram4: pass a CPU tensor to the schedule precompute so the warp result isn't bounced back to the device just for tolist() to pull it off again; values are byte-identical. - scheduler: document __call__ as an inference-time, non-differentiable helper (the sampler runs under no_grad and reads the values out as Python scalars). - quantized_loading: reword device_supports_fp8 docstring so the runtime probe is authoritative rather than asserting CUDA/CPU always pass. - quantized_loading: guard that an FP8 weight key ends with '.weight' before deriving the scale key, raising a clear RuntimeError otherwise. --- src/ideogram4/pipeline_ideogram4.py | 10 ++++++---- src/ideogram4/quantized_loading.py | 12 +++++++++--- src/ideogram4/scheduler.py | 5 +++++ 3 files changed, 20 insertions(+), 7 deletions(-) diff --git a/src/ideogram4/pipeline_ideogram4.py b/src/ideogram4/pipeline_ideogram4.py index 210dd4f..2eab10e 100644 --- a/src/ideogram4/pipeline_ideogram4.py +++ b/src/ideogram4/pipeline_ideogram4.py @@ -584,10 +584,12 @@ def __call__( device=self.device, ) - # step_intervals is fixed for the whole generation, so warp the whole tensor - # through the schedule in one vectorized call (one device transfer) rather - # than calling it per element, twice per step, inside the loop. - schedule_values = schedule(step_intervals).tolist() + # step_intervals is fixed for the whole generation, so warp it through the + # schedule once and read the values out as a Python list, rather than calling + # the schedule per element (twice per step) inside the loop. Pass a CPU tensor + # so the warp (which runs on CPU anyway) doesn't bounce the result back to the + # device just for tolist() to pull it off again. + schedule_values = schedule(step_intervals.cpu()).tolist() for i in range(num_steps - 1, -1, -1): t_val = schedule_values[i + 1] diff --git a/src/ideogram4/quantized_loading.py b/src/ideogram4/quantized_loading.py index 0d884f3..8c44337 100644 --- a/src/ideogram4/quantized_loading.py +++ b/src/ideogram4/quantized_loading.py @@ -176,9 +176,11 @@ def device_supports_fp8(device: torch.device) -> bool: Probed at runtime (cached per concrete device, e.g. ``cuda:0``, so a heterogeneous multi-device setup is checked individually) rather than - hard-coded. PyTorch's MPS (Apple Silicon) backend has no float8 support - it - can neither hold the dtype nor convert it - so there the FP8 weights must be - dequantized to the compute dtype at load time instead. CUDA and CPU both pass. + hard-coded, since float8 support can vary by PyTorch build and device. In + practice PyTorch's MPS (Apple Silicon) backend has no float8 support - it can + neither hold the dtype nor convert it - and fails the probe, so there the FP8 + weights must be dequantized to the compute dtype at load time; CUDA and CPU + pass on current builds. The live probe is authoritative. """ return _probe_fp8_support(str(torch.device(device))) @@ -318,6 +320,10 @@ def load_fp8_state_dict( else: # MPS can't cast float8, so dequantize on CPU (where the fp8 weights are # loaded) before moving the result across to the target device. + if not k.endswith(".weight"): + raise RuntimeError( + f"unexpected FP8 tensor key {k!r} (expected it to end with '.weight')" + ) scale_key = k[: -len(".weight")] + FP8_SCALE_SUFFIX if scale_key not in state_dict: raise RuntimeError( diff --git a/src/ideogram4/scheduler.py b/src/ideogram4/scheduler.py index f2682b4..72b9442 100644 --- a/src/ideogram4/scheduler.py +++ b/src/ideogram4/scheduler.py @@ -16,6 +16,11 @@ class LogitNormalSchedule: logsnr_max: float = 18.0 def __call__(self, t: torch.Tensor) -> torch.Tensor: + """Map step positions to noise levels (inference-time, non-differentiable). + + Called from the sampler loop under ``torch.no_grad`` and read out as Python + scalars, so the CPU detour taken on MPS never participates in autograd. + """ device = t.device # The float64 warp (ndtri/expit) needs precision at the tails. MPS supports # neither float64 nor these special functions, so there the warp runs on CPU