Skip to content

Fix disk offload crash on FP8 tensors - #4151

Open
shoemoney wants to merge 1 commit into
huggingface:mainfrom
shoemoney:fix/fp8-disk-offload
Open

Fix disk offload crash on FP8 tensors#4151
shoemoney wants to merge 1 commit into
huggingface:mainfrom
shoemoney:fix/fp8-disk-offload

Conversation

@shoemoney

Copy link
Copy Markdown

What does this PR do?

Disk-offloading a model with FP8 weights crashes in offload_weight:

File "src/accelerate/utils/offload.py", line 32, in offload_weight
    array = weight.cpu().numpy()
TypeError: Got unsupported ScalarType Float8_e4m3fn

Same error for float8_e5m2. NumPy has no FP8 representation, so weight.cpu().numpy()
fails outright.

This is the identical problem bfloat16 already has, and offload_weight already has a
branch for it: view the tensor as int16, write that, and record the real dtype in the
index so load_offloaded_weight can view it back. FP8 hits the same NumPy gap but has no
equivalent branch.

The reason this is easy to miss: dtype_byte_size (src/accelerate/utils/modeling.py)
already understands FP8 (correctly returns 1 byte for float8_e4m3fn, float8_e5m2,
float8_e4m3fnuz, float8_e5m2fnuz, float8_e8m0fnu — the last three added in #4063). So
device-map planning happily sizes an FP8 model, decides some layers go to disk, and then
the write itself blows up. The planning path knows about FP8; the write path doesn't.

Fix

Mirror the bfloat16 idiom on both sides of offload.py:

  • offload_weight: if the dtype string starts with torch.float8_, view the tensor as
    int8 before handing it to NumPy, and record the real dtype name in the index (matching
    what's already done for bfloat16 -> int16).
  • load_offloaded_weight: if the recorded dtype starts with float8_, load the memmap as
    int8 and view it back to the original FP8 dtype via getattr(torch, dtype).

Judgment call worth flagging: I matched on the float8_ prefix rather than listing
the five variants explicitly. dtype_byte_size already tracks five FP8 dtypes (two of them
added by #4063 after the original three), and a hardcoded pair/list here would just be
another place for that list to drift out of sync as new FP8 variants land. Happy to switch
this to an explicit list (mirroring the dtype_byte_size set) if maintainers would rather
be conservative about it.

Precedent

The int-view-and-restore trick is the established pattern for this exact class of bug —
it's what closed issue #454 ("Disk offload fails with bfloat16 weights") produced for
bfloat16. This PR applies the same fix to FP8.

Test plan

Extended the existing dtype loop in test_offload_weight (tests/test_offload.py) to
include float8_e4m3fn / float8_e5m2, guarded with is_torch_version(">=", "2.1.0") +
hasattr(torch, name) — the same guard style used by #4063. torch.randn(..., dtype=float8_*)
raises NotImplementedError, so the FP8 tensors are produced the way they actually show up
in practice: cast down from a float32 tensor via .to(dtype).

Round-trip correctness is checked on raw bits (.view(int8/int16/int32)), not float
equality, since float equality can't be trusted to catch a byte-level corruption for these
dtypes anyway.

Manually verified (beyond what's in the test) with NaN/inf-adjacent bit patterns injected
directly, including a full sweep of all 256 possible FP8 byte values:

dtype byte_exact dtype_preserved
float8_e4m3fn (all 256 bit patterns incl. NaN) True True
float8_e5m2 (all 256 bit patterns incl. inf/NaN) True True
bfloat16 (inf/nan/zero bit patterns) True True
float16 (inf/nan/zero bit patterns) True True
float32 (inf/nan/zero bit patterns) True True

Confirmed the test actually discriminates: reverted only the source change (kept the new
test), and it fails with the original crash:

weight = tensor([[-1.0000,  0.4688,  1.0000],
        [-1.2500,  0.2500, -0.1719]], dtype=torch.float8_e4m3fn)
...
>       array = weight.cpu().numpy()
E       TypeError: Got unsupported ScalarType Float8_e4m3fn
src/accelerate/utils/offload.py:32: TypeError

Full suite results on this branch:

  • tests/test_offload.py: 4 passed
  • tests/test_modeling_utils.py: 42 passed, 2 skipped (GPU-only)
  • ruff check / ruff format --check (pinned 0.13.1): clean

Who can review?

@SunMarc @muellerzr

offload_weight sizes FP8 tensors fine via dtype_byte_size, but crashes on
write with 'Got unsupported ScalarType Float8_e4m3fn' since NumPy has no
FP8 representation. Mirror the existing bfloat16 workaround: reinterpret
the 1 byte of FP8 data as int8 on write and restore the original dtype
on read, keyed off the torch.float8_ prefix rather than an explicit list
so it stays in sync with the five variants dtype_byte_size already knows
about.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant