Patrick C. Toulme d46768db04 [MTIA] Allow users who know what they are doing to ignore all device mismatches in tracing and take a preferred device. (#159931) ...

Summary:
Device mismatches in tracing can most often be ignored. These are only logical mismatches not physical.

Take any intermediate computation, and that computation will not actually materialize in a compiled binary execution. So a device mismatch in the middle of the program is not real. The runtime will never materialize those tensors on CPU device during the execution, as they are temporary allocations.

If a user knows his tensors at graph input are all on the correct device, then he can ignore all tracing errors.

Users who know what they are doing should have an escape hatch to ignore any device mismatch in tracing.

Users can set
```
  torch._functorch.config.fake_tensor_prefer_device_type = 'mtia'
```
to forcefully override any mismatch and prefer the non cpu device. This unblocks vLLM graph mode for MTIA.

Test Plan:
Added two unit tests.

Rollback Plan:

Differential Revision: D79698438

Pull Request resolved: https://github.com/pytorch/pytorch/pull/159931
Approved by: https://github.com/jansel

2025-08-07 22:37:15 +00:00

..

_activation_checkpointing

[BE][PYFMT] migrate PYFMT for torch/_[a-h]*/ to ruff format (#144551)

2025-06-25 06:16:06 +00:00

_aot_autograd

Extract some HOP utils to be importable (#159705)

2025-08-05 23:59:47 +00:00

__init__.py

Revert "Replace _dynamo.config with an object instead of module (#96455)"

2023-06-05 07:06:14 +00:00

aot_autograd.py

Pass fw/bw compilers to aot_export_joint_with_descriptors (#159814)

2025-08-06 04:50:56 +00:00

apis.py

[BE][PYFMT] migrate PYFMT for torch/_[a-h]*/ to ruff format (#144551)

2025-06-25 06:16:06 +00:00

autograd_function.py

Typo fixes for "overridden" in comments and function names (#155944)

2025-06-14 03:37:38 +00:00

batch_norm_replacement.py

remove allow-untyped-defs for torch/_functorch/batch_norm_replacement.py (#143438)

2024-12-18 09:01:06 +00:00

benchmark_utils.py

[BE][PYFMT] migrate PYFMT for torch/_[a-h]*/ to ruff format (#144551)

2025-06-25 06:16:06 +00:00

compile_utils.py

[BE][2/16] fix typos in torch/ (torch/_*/) (#156312)

2025-07-12 05:47:06 +00:00

compilers.py

[BE][2/16] fix typos in torch/ (torch/_*/) (#156312)

2025-07-12 05:47:06 +00:00

config.py

[MTIA] Allow users who know what they are doing to ignore all device mismatches in tracing and take a preferred device. (#159931)

2025-08-07 22:37:15 +00:00

deprecated.py

PEP585 update - torch/_functorch (#145139)

2025-01-19 07:06:10 +00:00

eager_transforms.py

[BE][PYFMT] migrate PYFMT for torch/_[a-h]*/ to ruff format (#144551)

2025-06-25 06:16:06 +00:00

functional_call.py

[BE][PYFMT] migrate PYFMT for torch/_[a-h]*/ to ruff format (#144551)

2025-06-25 06:16:06 +00:00

fx_minifier.py

PEP585 update - torch/_functorch (#145139)

2025-01-19 07:06:10 +00:00

make_functional.py

[BE][PYFMT] migrate PYFMT for torch/_[a-h]*/ to ruff format (#144551)

2025-06-25 06:16:06 +00:00

partitioners.py

Add expanded_def option for FX printing, render descriptor, update tests (#158708)

2025-07-25 13:22:32 +00:00

pyfunctorch.py

[dynamo] Guard serialization for FUNCTORCH_STACK_MATCH (#152616)

2025-05-05 18:05:56 +00:00

python_key.py

[BE][Easy][14/19] enforce style for empty lines in import segments in torch/_[a-c]*/ and torch/_[e-h]*/ and torch/_[j-z]*/ (#129765)

2024-07-31 10:42:50 +00:00

pytree_hacks.py

[BE] wrap deprecated function/class with typing_extensions.deprecated (#127689)

2024-06-02 12:30:43 +00:00

top_operators_github_usage.py

[BE][PYFMT] migrate PYFMT for torch/_[a-h]*/ to ruff format (#144551)

2025-06-25 06:16:06 +00:00

utils.py

PEP585 update - torch/_functorch (#145139)

2025-01-19 07:06:10 +00:00

vmap.py

[dynamo] Support functools.partial variables through inspect.signature (#146339)

2025-02-04 04:39:39 +00:00