Skip to content

fix(utils): skip device_context guard under torch.compile (#1590) - #1610

Open
YuvrajSinghBhadoria2 wants to merge 1 commit into
linkedin:mainfrom
YuvrajSinghBhadoria2:fix/device-context-compile
Open

YuvrajSinghBhadoria2 wants to merge 1 commit into
linkedin:mainfrom
YuvrajSinghBhadoria2:fix/device-context-compile

Conversation

@YuvrajSinghBhadoria2

Copy link
Copy Markdown

Fixes #1590

Skip the torch.cuda.device()/torch.xpu.device() context-manager guard
when torch.is_compiling() is true. Dynamo emits incorrect device-restore
nodes during tracing, which leads to wrong output or illegal memory
access (#1590).

Use getattr(torch, 'is_compiling', lambda: False)() so the check is
robust on torch versions that do not define torch.is_compiling.

Add CPU-runnable unit tests covering the compile-skip path, the
non-compile path, CPU device, and backends without a .device attribute.

All tests verified passing in a CPU-only Docker container (no GPU
required).

)

Skip the torch.cuda.device()/torch.xpu.device() context-manager guard
when torch.is_compiling() is true. Dynamo emits incorrect device-restore
nodes during tracing, which leads to wrong output or illegal memory
access (linkedin#1590).

Use getattr(torch, 'is_compiling', lambda: False)() so the check is
robust on torch versions that do not define torch.is_compiling.

Add CPU-runnable unit tests covering the compile-skip path, the
non-compile path, CPU device, and backends without a .device attribute.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant