Skip to content

fix(dynamo): preserve explicit device placement across partitions - #4629

Open
narendasan wants to merge 1 commit into
narendasan/fix-engine-scalar-iofrom
narendasan/fix-device-placement
Open

fix(dynamo): preserve explicit device placement across partitions#4629
narendasan wants to merge 1 commit into
narendasan/fix-engine-scalar-iofrom
narendasan/fix-device-placement

Conversation

@narendasan

Copy link
Copy Markdown
Collaborator

Summary

  • keep Torch-executed submodules on their source devices unless compilation explicitly offloaded them
  • reject _to_copy conversions that move between devices, since TensorRT can represent the dtype cast but not the placement change
  • preserve CPU placement across partition boundaries

Issues

Closes #4613
Closes #4615

Tests

Stack

3 of 7. Base: #4628. Follow-up: narendasan/fix-truncate-double-metadata.

Keep torch-executed submodules on their source devices unless compilation explicitly offloaded the module, and reject _to_copy conversions that move between devices because TensorRT can represent the dtype cast but not the placement change.

Fixes #4613
Fixes #4615

Tests:
- pytest tests/py/dynamo/conversion/test_casts.py -n0 (8 passed)
- pytest tests/py/dynamo/models/test_device_placement.py -n0 (2 passed)
- original issue reproductions for #4613 and #4615
@github-actions github-actions Bot added component: tests Issues re: Tests component: conversion Issues re: Conversion stage component: core Issues re: The core compiler component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Aug 28, 2026
@meta-cla meta-cla Bot added the cla signed label Aug 28, 2026
@github-actions
github-actions Bot requested a review from cehongwang August 28, 2026 02:59
@narendasan
narendasan requested a review from apbose August 28, 2026 03:05
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [Python] Issues re: Python API component: conversion Issues re: Conversion stage component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: tests Issues re: Tests

Projects

None yet

1 participant