[Bug]: torch._inductor.exc.InductorError: TypeError: cannot pickle 'torch._C.DispatchKeySet' object

[zou3519]

Your current environment

vLLM main branch, PyTorch main branch

🐛 Describe the bug

Repro:
pytest -v -s tests/compile/piecewise/test_toy_llama.py

Gives:

>                           rv = reductor(4)
E                           torch._inductor.exc.InductorError: TypeError: cannot pickle 'torch._C.DispatchKeySet' object
E
E                           Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even m
ore developer context, set TORCH_LOGS="+dynamo"

../env/lib/python3.12/copy.py:151: InductorError
====================================================================== warnings summary =======================================================================

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[zou3519]

Also saw this with

pytest tests/compile/test_full_graph.py::test_full_graph[model_info1-3]
pytest tests/compile/test_full_graph.py::test_full_graph[model_info2-3]
pytest tests/compile/test_fusion.py