Let a serialized fake-tensor state_dict be read back

[ThomasJannaud]

Summary:
serialize_torch_artifact already supports fake tensors: it installs
_reduce_fake_tensor, which writes a FakeTensor as its TensorMeta -- shape,
dtype and stride, no storage. Two gaps stopped the matching read path from
working, so an artifact whose state_dict holds fake tensors could be written
and then not read back. Both are latent and independent of any caller; they are
split out of the Helios stack that found them so they can be reviewed on their
own.

1. weights_only=True refused ExecuTorch's own reconstructor.
deserialize_torch_artifact loads with weights_only=True, which rejects any
global that is not allowlisted:

WeightsUnpickler error: Unsupported global: GLOBAL
executorch.exir.serde.export_serialize._reconstruct_fake_tensor
was not an allowed global by default

Fixed by allowlisting the reconstructor for the load with
torch.serialization.safe_globals. That is narrower than add_safe_globals,
which registers process-globally and permanently -- but narrower in time only,
not per-call: the allowlist is a single process-global set, so a concurrent
torch.load(weights_only=True) on another thread during the window accepts this
global too.

Entering it is conditional, and that is not defensive coding. safe_globals
unions that set on entry and subtracts on exit, with no refcount. Entering it
unconditionally therefore revokes the registration of any caller who has already
added this global with add_safe_globals -- precisely the workaround people
applied before this function handled it. Measured rather than argued: with the
unconditional form, a probe reading the allowlist around one load reported
before=True after=False.

2. Rebuilding a parameter raised on any integer dtype.
_reconstruct_fake_tensor restored nn.Parameter entries as
torch.nn.Parameter(fake_tensor), and that default of requires_grad=True
raises:

RuntimeError: only Tensors of floating point dtype can require gradients

nn.Parameter accepts an integer tensor as long as requires_grad is False,
and quantized weights -- uint8 codes and their zero points -- are exactly that.
So this hit the whole class of quantized models, not an edge case.

requires_grad is derived from the dtype. Recording the flag in the reducer
instead was written first, then backed out: the reducer's return value is the
on-disk format, and two other copies of this module in fbsource
(xplat/spacecraft/sample_apps, third-party/mediatek/release-0.7) still have
the two-argument _reconstruct_fake_tensor, so widening the payload would
produce records they cannot read. The cost is fidelity, not correctness: a float
parameter that had requires_grad=False comes back True. That was already the
behaviour before integer parameters worked at all, and no reader of a serialized
fake state_dict consumes the flag.
test_requires_grad_is_derived_from_the_dtype pins it, so it stays a known
trade-off rather than a surprise.

The reasoning behind both choices is left as inline comments on this diff rather
than in the source, which keeps only the two constraints a future reader needs
in order not to undo them.

Scope, so this is not over-read: torch._export.serde.serialize carries the
same requires_grad defect and is untouched here. It is a different module
path, so it can never be resolved for a record written by this one -- the same
bug in another codebase, not a compatibility hazard for this one.

Differential Revision: D115918112

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22111

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures, 1 Unrelated Failure

As of commit 1aaaded with merge base 4d3b46d ():

NEW FAILURES - The following jobs have failed:

• Cadence Build & Test / hifi-build / hifi4 (gh)
  ##[error]Refusing to check out fork pull request code from a 'pull_request_target' workflow. This workflow runs with the base repository's GITHUB_TOKEN, secrets, default-branch cache scope, and runner access. Fetching and executing a fork's code in that trusted context commonly leads to "pwn request" vulnerabilities. To opt in, review the risks at https://gh.io/securely-using-pull_request_target and set 'allow-unsafe-pr-checkout: true' on the actions/checkout step.
• Cadence Build & Test / vision-build / vision (gh)
  ##[error]Refusing to check out fork pull request code from a 'pull_request_target' workflow. This workflow runs with the base repository's GITHUB_TOKEN, secrets, default-branch cache scope, and runner access. Fetching and executing a fork's code in that trusted context commonly leads to "pwn request" vulnerabilities. To opt in, review the risks at https://gh.io/securely-using-pull_request_target and set 'allow-unsafe-pr-checkout: true' on the actions/checkout step.

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

• ❌ - login: @ThomasJannaud / name: Thomas Jannaud. The commit (1aaaded) is not authorized under a signed CLA. Please click here to be authorized. For further assistance with EasyCLA, please visit our EasyCLA portal and chat with our support bot.

[meta-codesync]

@ThomasJannaud has exported this pull request. If you are a Meta employee, you can view the originating Diff in D115918112.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 1m 14s —— View job

---

Code Review

• Read the PR diff
• Analyze export_serialize.py changes
• Analyze test coverage
• Post review feedback

This is a careful, well-scoped fix. Both gaps are real: weights_only=True does reject non-allowlisted globals, and nn.Parameter(int_tensor) with the default requires_grad=True does raise. The tests pin exactly the two behaviors described, including the deliberately-lossy requires_grad derivation. The inline comments capture the two constraints a future reader needs without over-explaining. Nice work.

A few small notes, none blocking:

1. import contextlib ordering (export_serialize.py:20)
It's inserted between import io and import json, but the stdlib block is otherwise alphabetical (contextlib sorts before copy). lintrunner / isort will likely want it moved up. Also, line 28 already has from contextlib import contextmanager — you could add nullcontext there instead of importing the module and dropping a stray contextlib. call into an out-of-order slot:

from contextlib import contextmanager, nullcontext

Worth running lintrunner -a to be safe.

2. Thread-safety of the conditional guard is broader than the exit race (export_serialize.py:355-361)
The PR body already acknowledges that safe_globals is a single process-global set, so a concurrent torch.load(weights_only=True) during the window transiently accepts this global — accepted trade-off. Just flagging the read-then-enter check has a second race: if another thread calls add_safe_globals([_reconstruct_fake_tensor]) after this branch reads get_safe_globals() (sees it absent) but the entering thread's safe_globals(...) context then removes it on exit, you'd revoke a registration that the guard was meant to protect. It's a narrow TOCTOU window and out of scope for this fix, but since the whole point of the guard is to not stomp a caller's registration, it's worth a one-line comment noting the guard is best-effort, not airtight.

3. is_complex handling — correct, minor question
requires_grad = fake_tensor.dtype.is_floating_point or fake_tensor.dtype.is_complex correctly mirrors what nn.Parameter will accept (complex tensors can require grad). Good that you covered complex and didn't just special-case float. No change needed.

Correctness / tests

• test_reading_leaves_a_callers_own_allowlist_entry_alone correctly saves/restores global state in finally — good hygiene, avoids leaking allowlist mutations into other tests in the process.
• test_program_with_a_fake_state_dict_round_trips is the valuable one: it exercises _verify_exported_program_signature's InputKind.PARAMETER requirement, which is where a plain-FakeTensor (non-Parameter) reconstruction would actually surface for callers. Good that the is_parameter branch preserves the nn.Parameter type.

Non-code: the EasyCLA check is currently failing (commit not covered by a signed CLA) — that'll need resolving before merge, independent of the code.

Overall LGTM — just the import-ordering nit is worth addressing before landing.
· export-D115918112