executorch: persist external .ptd weights and preserve lifted constant dtype/device

[Conarnar]

Description

Reworked per review (previously "executorch: fall back to the CUDA backend..."):
this PR is trimmed to two standalone fixes and drops the cuda_fallback=True
flag. save(output_format="executorch") already appends caller partitioners
after the TensorRTPartitioner, so the supported way to get a coalesced
TensorRT + CUDA .pte is to pass partitioners=[CudaPartitioner(...)] (now
documented in saving_models.rst and the save() docstring). The flag only
duplicated that — and "fallback" wrongly implied a downgrade rather than a second
GPU delegate — so it's removed. This PR keeps the two pieces that are real fixes:

External weights (.ptd). The CUDA (AOTInductor) backend emits its weights as
external named data, which ExecuTorch serializes only via
write_tensor_data_to_file — write_to_file does not persist it. Without this, a
partition that carries external weights (e.g. a CudaPartitioner delegate) loses
its blob and the .pte can't load. save() now writes the .ptd next to the
.pte (extracted into _write_external_tensor_data, guarded on _tensor_data
directly so a future rename fails loud instead of silently skipping) and logs the
target directory, so a missing runtime data-file is an obvious breadcrumb rather
than a silent null-weights load. It's a no-op for TRT-only programs
(_tensor_data is empty). Co-locating the .ptd is a convention — the runtime
must still be pointed at the data file(s) to load them. The docs also warn that the
CUDA backend names its blob per-device, so each coalesced model must be saved into
its own directory.

Lifted constant dtype/device (_exporter.py). lift() built the lifted
placeholder's fake meta["val"] with torch.empty_strided(...) from shape only,
defaulting it to cpu/float32. A lifted non-fp32 or non-CPU parameter/buffer
therefore got the wrong meta (e.g. a FakeTensorDeviceMismatchError, cuda:0 vs
cpu, when a weighted op is excluded from TensorRT). Now it passes
dtype=lift_val.dtype, device=lift_val.device. Independent bug fix; can stand
alone.

Testing

• tests/py/dynamo/executorch/test_api.py — test_lift_preserves_constant_dtype_device,
  parametrized over (bfloat16, cpu) and (float32, cuda) (the CUDA case is
  runtime-skipped when no GPU), asserts lift() keeps the lifted constant's dtype
  and device.
• tests/py/dynamo/executorch/test_edge_cases.py — unit tests for
  _write_external_tensor_data: writes the .ptd when external data is present,
  no-op when empty, and raises (fails loud) if _tensor_data is absent.
• tests/py/dynamo/executorch/test_cuda_partitioner_composition.py (build-only;
  skips without a GPU / CUDA backend / nvcc):
  • test_erfinv_routes_to_cuda_backend — erfinv (no TRT converter) routes to the
    CUDA backend via partitioners=[CudaPartitioner(...)] while tanh/cos stay on
    TensorRT.
  • test_weighted_partition_persists_external_data — a weighted, TRT-excluded mm
    lands on the CUDA backend and its external .ptd is written next to the .pte
    (asserted non-empty).
  • test_trt_only_writes_no_ptd — a TRT-only program writes no .ptd and has no
    CUDA delegate (exercises the real write_to_file/_tensor_data path).

Composition tests are build-only (delegate composition + external-data persistence,
not load/run) and skip cleanly without a GPU/CUDA backend/nvcc; a load-run-allclose
test is deferred to the runtime follow-up below.

Follow-up

A coalesced ATen-mode .pte can be built but not yet loaded/run: memory-planned
CUDA buffers get a CPU data pointer (tensor_parser_aten.cpp hardcodes CPU), so
Method::init fails the CUDA backend's device check. That runtime fix (and the
load-run-allclose test) will land separately.

[Gasoonjia]

can we dynamic choose the method name, instead of statically use forward?
Also in the test lets add a multi-method one.

[Conarnar]

At the moment save only supports a single exported program which to_edge_transform_and_lower will default to the method name forward.

[Gasoonjia]

is the model runnable now? Can we handle the data transfer between different delegate blobs?

[Conarnar]

Should be runnable.

[shoumikhin]

A few points from a closer look, most important first.

1. Dropped external CUDA weights. _save_as_executorch only calls executorch_program.write_to_file(f), but the CUDA backend emits its weights as external named data (save_data_externally is True, tagged aoti_<device>_blob) that ExecuTorch serializes only through write_tensor_data_to_file. Since that is never called, any residual partition that carries weights loses its blob and the .pte cannot load. The new test does not catch this because cos(erfinv(tanh(x))) has no parameters, so please persist the external data next to the .pte and add a test model with weights in the fallback partition.

2. The import guard is too narrow. The try/except only catches ImportError on the module import, but the real CUDA compile runs later inside to_edge_transform_and_lower via CudaBackend.preprocess calling torch._inductor.aot_compile with max_autotune and the Triton backends, which needs nvcc/ptxas. On a host that has the executorch cuda wheel but no cuda toolkit, export now throws where before the residual op stayed non-delegated and save produced a working CPU .pte, so the docstring line about degrading gracefully only holds for the missing-import case. Consider gating the fallback on ptxas availability or making it opt-in, and narrowing the docstring.

3. Always on with no opt-out. Appending CudaPartitioner unconditionally means any model with even one non-TRT op now embeds a CudaBackend blob and requires a CUDA runtime at load, where before those ops stayed non-delegated and the .pte ran on CPU. Most real models are not fully TRT-convertible, and this also silently turns on max_autotune autotuning at save time, so please add a save() kwarg to keep the portable CPU-fallback behavior.

4. target_device is not propagated to the fallback. It is forwarded only to TensorRTPartitioner; the fallback is built as cuda_specs = [CudaBackend.generate_method_name_compile_spec("forward")] with no target_device, so CudaPartitioner defaults to cuda:0. If a user pins TRT to cuda:1 the two delegates end up on different devices, so derive the fallback target_device from the same compile_specs.

5. The test proves composition, not runnability. It only deserializes the program and checks that CudaBackend and TensorRTBackend appear in the delegate ids; it never builds a Method, runs, or compares to eager. The mixed ATen-mode roundtrip is also still gated on the tensor_parser_aten fix (planned CUDA buffers get a CPU data ptr so Method::init fails in getDeviceFromPtr), so scope the assertion to "produces both delegates" and add a real load-run-allclose once that lands.

[Gasoonjia]

Always enjoy reading Anthony's comment

1. for external weight lets make it as part of pte to in line with tensorrt type., instead of an individual ptd file like standard cuda-only delegated model.
2. we should make sure it is executable. I would encourage to add a pybinding test as guard.

[shoumikhin]

Thanks for working on this. The direction is right and I want to use it, but I
think the PR should be trimmed and split. Here is what I would like it to look
like and why.

Short version

There are three separate things in this PR:

1. Write the external weight .ptd next to the .pte. Keep this, it is the real
   fix. Two small changes below.
2. Fix _exporter.py so lifted params keep their dtype and device. Keep this too,
   it is a clean bug fix.
3. The new cuda_fallback=True flag. I would drop this. It duplicates something
   save() can already do, and the name is misleading.

I would land (1) and (2) as one small PR with tests, and document the CUDA path
instead of adding the flag.

1. The .ptd write: keep it, with two tweaks

This is the important part. The CUDA (AOTInductor) backend externalizes its weight
constants (package_constants_in_so=False), so without writing the .ptd the
runtime loads a null weights blob and the kernel reads out of bounds. write_to_file
alone does not persist that data, so this write is needed.

Two changes I would make:

a) Check _tensor_data directly instead of getattr(executorch_program, "_tensor_data", None). The attribute always exists on the program object, so
the None default never helps the real case. All it does is turn a future
rename or removal of _tensor_data into a silent no-op, which quietly brings
this exact crash back. A plain if executorch_program._tensor_data: fails loud
instead, which is what we want.

b) Say in a comment that the loader has to be pointed at the .ptd. Writing the
file next to the .pte is a co-location convention, not an automatic search.
The runtime still needs to be given the data files (for example via the Module
data-files argument). Right now the wording reads like the runtime finds them on
its own, which will confuse whoever wires up loading.

For reference, this is the shape I have been running:

if executorch_program._tensor_data:
    executorch_program.write_tensor_data_to_file(
        os.path.dirname(os.path.abspath(file_path))
    )

The guard makes it a no-op for TRT-only programs (_tensor_data is empty there),
so it is safe to always run.

2. _exporter.py dtype/device fix: keep it

Passing dtype=lift_val.dtype, device=lift_val.device into empty_strided is a
straight bug fix. A lifted constant that is not float32 or not on CPU gets the
wrong meta today. Please add a small unit test for it (a lifted bf16 or CUDA
constant, assert the placeholder meta dtype and device match the source). It is
independent of the CUDA work and can land on its own.

3. Drop cuda_fallback=True

save(output_format="executorch") already appends caller partitioners after the
TensorRT partitioner:

partitioners = [TensorRTPartitioner(...)] + list(extra_partitioners)

So this already gives "TRT engines plus a CUDA catch-all last", no new flag needed:

save(
    trt_gm, path, output_format="executorch", retrace=False, arg_inputs=...,
    partitioners=[
        CudaPartitioner([CudaBackend.generate_method_name_compile_spec("forward")])
    ],
)

The only extra thing the flag does is copy target_device onto the CudaPartitioner,
but CudaPartitioner already defaults to cuda:0, and a different device is one
line for the caller:

CudaPartitioner([
    CompileSpec("target_device", b"cuda:1"),
    CudaBackend.generate_method_name_compile_spec("forward"),
])

So the flag adds a bool, plumbing through three call sites, a docstring, and an
ImportError branch that couples torch_tensorrt.save to
executorch.backends.cuda, all to save one line at the call site. I would rather
keep save() free of CUDA-backend specifics and let CudaPartitioner own that.

Two more concrete problems with the flag as written:

• Double catch-all. If someone passes both cuda_fallback=True and
  partitioners=[CudaPartitioner(...)], you get two CUDA partitioners. The first
  greedily claims all non-TRT nodes, the second is dead or conflicts, and they can
  disagree on target_device silently. If the flag stays, it has to detect an
  existing CudaPartitioner and not append a second one.

• The name. "fallback" usually means "could not lower, drop to a CPU or reference
  kernel", which is a downgrade. Here the ops TRT will not take are promoted to a
  second GPU delegate. The name says the opposite of what happens. If we ever want
  sugar for this, name it after delegation, not fallback.

4. Please scope the end-to-end claim

The test only checks that both delegates are present, and as the PR notes it cannot
load or run the .pte. The reason is real: planned CUDA buffers currently get a CPU
data pointer in tensor_parser_aten.cpp, so Method::init fails in
getDeviceFromPtr and the CUDA backend rejects the pointer. So a coalesced .pte
can be built but not loaded yet.

I am fine landing a composition-only test as long as it is labeled that way and the
runtime gap is written down as a follow-up with a task. I do not want this PR framed
as "export works end to end", because it does not run yet. A real load-run-allclose
test should come after the ATen device-pointer path is fixed.

What I would like the final shape to be

• PR A: the .ptd write (with the two tweaks above) plus the _exporter.py
  dtype/device fix, each with a unit test. Small and landable now.
• Docs: show partitioners=[CudaPartitioner(...)] as the supported way to get a
  coalesced TRT + CUDA .pte, instead of a flag.
• Follow-up PR: the tensor_parser_aten / getDeviceFromPtr runtime fix that makes
  a coalesced .pte actually load and run, with an end-to-end test.

[shoumikhin]

Thanks for turning this around — trimming to the two fixes + documenting the partitioners=[CudaPartitioner(...)] recipe addresses my earlier review, and both fixes read as root-cause correct. A few small polish items left, all non-blocking:

Should-fix

1. The save() docstring recipe is missing retrace=False. The .rst examples have it, but the docstring copy (the torch_tensorrt.save(...) block) passes only arg_inputs=inputs, so a copy-paste from the docstring runs the default retrace=True path rather than the tested one.
2. _write_external_tensor_data writes the .ptd silently. Since co-location is only a convention and the runtime must be pointed at the data file(s) explicitly, a single logger.info naming the directory turns a likely silent null-weights load failure into an obvious breadcrumb:

if executorch_program._tensor_data:
    out_dir = os.path.dirname(os.path.abspath(file_path))
    executorch_program.write_tensor_data_to_file(out_dir)
    logger.info(
        "Wrote external delegate weights (.ptd) to %s; point the runtime's "
        "data-files at this directory to load them.", out_dir,
    )

3. Worth a one-line docs warning: the CUDA backend names its external blob per-device (e.g. aoti_cuda_blob.ptd), not per-model, so saving two different coalesced .pte into the same directory overwrites the blob and the first .pte won't load. "Save each coalesced model into its own directory" covers it.

Nice-to-have
4. test_api.py gates the CUDA lift test with a module-level @pytest.mark.skipif, whereas test_cuda_partitioner_composition.py deliberately uses a runtime autouse fixture and even documents why (module-level skipif resolves at collection time, which is fragile on remote-GPU runners). Suggest matching the runtime-gating pattern for consistency. The bf16 and cuda lift tests are also near-identical — a single @pytest.mark.parametrize over (dtype, device) would cover both and read cleaner.
5. The TRT-only no-op is verified via a stub; a small end-to-end check (a real save(...) without a CudaPartitioner asserting no .ptd next to the .pte) would exercise the real write_to_file/_tensor_data path, and asserting the written .ptd is non-empty (st_size > 0) in the composition test would make the persistence check airtight.

LGTM once the should-fix items are addressed.

[shoumikhin]

The docstring recipe is missing retrace=False (the .rst examples have it). Without it a copy-paste from here runs the default retrace=True path rather than the tested one — suggest adding retrace=False, next to arg_inputs=inputs. Since this ~15-line recipe is also duplicated in saving_models.rst (and the two already differ, gm vs trt_gm), consider trimming the docstring to a one-line summary plus a :ref: to the guide for a single source of truth.

[shoumikhin]

Consider a breadcrumb so the required runtime step is discoverable — co-location is only a convention, so a program that needs the .ptd but isn't given it fails at load with a null-weights dereference and no hint:

if executorch_program._tensor_data:
    out_dir = os.path.dirname(os.path.abspath(file_path))
    executorch_program.write_tensor_data_to_file(out_dir)
    logger.info(
        "Wrote external delegate weights (.ptd) to %s; point the runtime's "
        "data-files at this directory to load them.", out_dir,
    )

[shoumikhin]

detect_fake_mode() only retrieves the mode; it does not activate it, and cast(FakeTensor, ...) has no runtime effect. Consequently, this calls torch.empty_strided outside the detected mode and creates a real tensor on lift_val.device. For CUDA constants, saving can now allocate on or initialize the target GPU, potentially causing OOM or making an otherwise device-independent export require a live CUDA device. It also stores a real tensor in metadata associated with the graph's FakeTensorMode.

Could we construct this through the detected mode instead?

const_placeholder_node.meta["val"] = fake_mode.from_tensor(
    lift_val, static_shapes=True
)

This avoids real device allocation and preserves the source tensor's layout, unlike the synthetic all-ones stride. The regression test currently checks only dtype/device, which a real CUDA tensor also satisfies; please additionally assert isinstance(val, FakeTensor), that it belongs to fake_mode, and that its stride matches lift_val.stride().

[shoumikhin]

Minor follow-up: _traced_gm_with_buffer creates self.c as an nn.Parameter, so this exercises the PARAMETER branch rather than the buffer branch described by the helper name/docstring. Could we rename it to _traced_gm_with_parameter, or register an actual buffer if buffer coverage is intended?