feat(executorch): forward ExecuTorch lowering kwargs through save(output_format="executorch")#4433
Open
Conarnar wants to merge 3 commits into
Open
feat(executorch): forward ExecuTorch lowering kwargs through save(output_format="executorch")#4433Conarnar wants to merge 3 commits into
Conarnar wants to merge 3 commits into
Conversation
lift() built the placeholder's fake meta["val"] with torch.empty_strided from shape only, so a lifted non-fp32 or non-CPU constant silently got cpu/float32 meta (e.g. a FakeTensorDeviceMismatchError, cuda:0 vs cpu, when a weighted op is excluded from TensorRT). Fakify the source constant through the graph's own fake_mode instead -- fake_mode.from_tensor(lift_val, static_shapes=True). The placeholder meta now preserves the source dtype, device and stride and belongs to the graph's FakeTensorMode, and no real tensor is allocated on lift_val.device (so a CUDA constant no longer needs a live GPU, or risks OOM, at export time). Adds unit tests (bf16 dtype + CUDA device) asserting dtype, device, stride, and that the meta is a FakeTensor.
… delegation
The CUDA (AOTInductor) backend externalizes its weight constants, so save() must
write the .ptd next to the .pte or the runtime loads a null weights blob and reads
out of bounds. Keep that write, extracted into _write_external_tensor_data() and
guarded on _tensor_data directly (no getattr default) so a future rename fails loud
instead of silently skipping. Document that the runtime must be pointed at the .ptd
(co-location is a convention, not an automatic search).
Drop the cuda_fallback=True flag: save(partitioners=[CudaPartitioner(...)]) already
appends a CUDA catch-all after the TensorRT partitioner, so the flag only duplicated
that, and 'fallback' wrongly implied a downgrade rather than a second GPU delegate.
Document the partitioners= recipe in saving_models.rst and the save() docstring.
Composition-only test added (build, not load/run); a load-run test follows the ATen
device-pointer runtime fix.
Test-Plan: pytest tests/py/dynamo/executorch/{test_api,test_edge_cases}.py -m unit
…put_format="executorch")
torch_tensorrt.save already forwards partitioners= and backend_config= to the
ExecuTorch lowering. Forward the remaining to_edge_transform_and_lower kwargs the
same way -- pop them in save(), thread them through the three _save_as_executorch
call sites, and pass them to to_edge_transform_and_lower(...):
- constant_methods: bake constant methods (e.g. LLM-runner metadata) into the .pte
- transform_passes: run caller-supplied edge transform passes
- compile_config: override the EdgeCompileConfig. Defaults to
_check_ir_validity=False (the TRT execute_engine placeholder
graph fails edge IR validation); a caller-supplied config is
forwarded verbatim so an explicit _check_ir_validity is honored.
- generate_etrecord: generate + persist an ETRecord as "<base>_etrecord.bin" next
to the .pte, following ExecuTorch's example convention
(e.g. examples/cuda/scripts/export.py)
Tests: test_save_executorch_forwards_lowering_kwargs and
test_save_executorch_defaults_when_lowering_kwargs_omitted assert each kwarg is
forwarded, that an omitted compile_config defaults to _check_ir_validity=False
while a caller-supplied one is forwarded verbatim, and that the ETRecord is written.
github-actions
Bot
added
documentation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
torch_tensorrt.save already forwards partitioners= and backend_config= to the
ExecuTorch lowering. Forward the remaining to_edge_transform_and_lower kwargs the
same way -- pop them in save(), thread them through the three _save_as_executorch
call sites, and pass them to to_edge_transform_and_lower(...):
_check_ir_validity=False (the TRT execute_engine placeholder
graph fails edge IR validation); a caller-supplied config is
forwarded verbatim so an explicit _check_ir_validity is honored.
to the .pte, following ExecuTorch's example convention
(e.g. examples/cuda/scripts/export.py)
Tests
test_save_executorch_forwards_lowering_kwargs and
test_save_executorch_defaults_when_lowering_kwargs_omitted assert each kwarg is
forwarded, that an omitted compile_config defaults to _check_ir_validity=False
while a caller-supplied one is forwarded verbatim, and that the ETRecord is written.
Requires #4404