Precision loss in MRoPE `inv_freq` during FP8 pipeline loading causes spatial artifacts

[jrd-rocks]

Disclosure: I used an AI assistant to help isolate and document this bug. The bug is very real, however, see the image comparisons at the bottom; before and after the fix in my toy code

Description

When loading an FP8 or NF4 checkpoint (e.g., ideogram-ai/ideogram-4-fp8 or ideogram-ai/ideogram-4-nf4) via Ideogram4Pipeline.from_pretrained(..., dtype=torch.bfloat16), the _build_transformer function downcasts the inv_freq buffer inside the Ideogram4MRoPE module from float32 to bfloat16.

Because bfloat16 only has 7 bits of mantissa precision, and the pipeline uses a large position offset for image tokens (IMAGE_POSITION_OFFSET = 10000), computing pos * inv_freq with a bfloat16 inverse frequency results in massive rounding errors. Neighboring spatial tokens end up receiving the exact same positional encodings, which completely corrupts the model's understanding of space and results in severe visual and spatial mismatch artifacts.

Note: I personally experienced and debugged this on the FP8 branch, but upon reviewing the code, it appears the exact same bug exists in the BNB 4-bit (NF4) loading path as well.

The Root Cause

In src/ideogram4/pipeline_ideogram4.py (and quantized_loading.py), the quantized loading branches execute .to(dtype) casts that inadvertently degrade the non-persistent inv_freq buffer.

For FP8:

  elif is_fp8_state_dict(state_dict):
    model.to(dtype) # <--- BUG: Downcasts the float32 inv_freq buffer to bfloat16
    swap_linears_to_fp8(model, state_dict, compute_dtype=dtype)
    load_fp8_state_dict(model, state_dict, device=device, dtype=dtype)

For BNB 4-bit (NF4):
In src/ideogram4/quantized_loading.py, load_bnb4bit_state_dict iterates through all buffers and downcasts them:

  for name, b in list(model.named_buffers()):
    if b.is_floating_point() and b.dtype != dtype:
      # ...
      parent.register_buffer(
        leaf,
        b.to(device=device, dtype=dtype), # <--- BUG: Downcasts the float32 inv_freq buffer to bfloat16
        persistent=leaf not in parent._non_persistent_buffers_set,
      )

Symptoms

Because spatial attention is broken, the model's understanding of position is highly degraded.

1. The canvas appears disjointed, and the positions of rendered objects do not match the positions specified in bounding boxes or positional prompts (e.g., the canvas feels "bigger" than the expected positions).
2. The overall generated images look visually off and distorted.

Proposed Fix

The fix is to simply re-calculate and overwrite inv_freq in float32 immediately after the .to(dtype) casts in src/ideogram4/pipeline_ideogram4.py.

def _build_transformer(
  transformer_config: "Ideogram4Config",
  state_dict: dict[str, torch.Tensor],
  device: torch.device,
  dtype: torch.dtype,
) -> "Ideogram4Transformer":
  model = Ideogram4Transformer(transformer_config)
  
  if is_bnb4bit_state_dict(state_dict):
    if device.type != "cuda":
      raise ValueError(f"bnb 4-bit weights require a CUDA device, got device={device}")
    swap_linears_to_bnb4bit(model, compute_dtype=dtype)
    load_bnb4bit_state_dict(model, state_dict, device=device, dtype=dtype)
    
    # FIX: Restore inv_freq precision lost during load_bnb4bit_state_dict
    inv_freq = 1.0 / (
        transformer_config.rope_theta ** (torch.arange(0, model.rotary_emb.head_dim, 2, dtype=torch.float32) / model.rotary_emb.head_dim)
    )
    model.rotary_emb.register_buffer("inv_freq", inv_freq, persistent=False)
    
  elif is_fp8_state_dict(state_dict):
    model.to(dtype)
    swap_linears_to_fp8(model, state_dict, compute_dtype=dtype)
    load_fp8_state_dict(model, state_dict, device=device, dtype=dtype)

    # FIX: Restore inv_freq precision lost during model.to(dtype)
    inv_freq = 1.0 / (
        transformer_config.rope_theta ** (torch.arange(0, model.rotary_emb.head_dim, 2, dtype=torch.float32) / model.rotary_emb.head_dim)
    )
    model.rotary_emb.register_buffer("inv_freq", inv_freq, persistent=False)
    
  else:
    model.load_state_dict(state_dict)
    model.to(device=device, dtype=dtype)
    
  model.eval()
  return model

[PierrunoYT]

Confirmed: this bug is also present in the diffusers Ideogram4Pipeline (nf4), and the affected offset is larger there.

I verified this directly against an installed diffusers build (Ideogram4Pipeline, torch_dtype=torch.bfloat16, CUDA/Windows).

In diffusers/models/transformers/transformer_ideogram4.py:

• inv_freq is built in float32 and registered with persistent=False.
• forward() does self.inv_freq.to(dtype=torch.float32), but that runs after the buffer was already downcast at load time, so it cannot recover lost mantissa bits.
• IMAGE_POSITION_OFFSET = 65536 here — 6.5x larger than the 10000 in this repo, so the precision loss is worse.

Live test on the actual env:

• Ideogram4MRoPE(...).to(torch.bfloat16) downcasts inv_freq float32 -> bfloat16 (since from_pretrained(torch_dtype=torch.bfloat16) casts the whole model).
• At position 65536 the bf16 rounding gives a max phase error of ~58.8 rad (>9 full rotations).
• 12 / 64 frequencies have > pi rad error; 28 / 64 have > 0.1 rad error at the image offset.
• Mean relative rounding error in inv_freq is ~1.5e-3, which is enormous once multiplied by a 65536 offset.

So the float32 inv_freq fix proposed here should be applied on the diffusers loading path as well (re-register inv_freq in float32 after the dtype cast).

I also see smooth-gradient banding / vertical streaks that may be separate 4-bit quantization noise; sample here: https://huggingface.co/ideogram-ai/ideogram-4-nf4/discussions/5