Skip to content

Conversation

@mikaylagawarecki
Copy link
Contributor

@mikaylagawarecki mikaylagawarecki commented May 23, 2024

Tensor types that have been opted in to the swap_tensors path.

The following table lists tensor types that have been opted in by default to the swap_tensors path in either _apply or load_state_dict. All other types have not been opted in by default and will only use swap_tensors if torch.__future__.set_swap_module_params_on_conversion(True)

nn.Module._apply opted in _apply escape hatch nn.Module.load_state_dict opted in load escape hatch
traceable wrapper subclass silent incorrectness Y N (bugfix) need to overwrite copy_ to interop with other subclasses Y Y
meta device assignment (reference lost) Y Y N Y
xla device assignment (reference lost) Y Y N Y

This PR adds the escape hatches mentioned in the above table to torch.__future__

  • _{get/set}_swap_overwrite_escape_hatch: which allows tensor types that assigned new params in nn.Module._apply that were opted in to swap_tensors (right now just src/dest on xla/meta to fall back to this behavior
  • _{get/set}_swap_load_state_dict_escape_hatch: which allows tensor types that were opted in to swap_tensors path in nn.Module.load_state_dict (for now just to fall back to old behavior

This allows backward compatibility as it provides an option to get the old behavior. In the future, we will need one more escape hatch for tensors that used tensor.data = in nn.Module._apply

Can we let these be private for now?

Stack from ghstack (oldest at bottom):

[ghstack-poisoned]
@pytorch-bot
Copy link

pytorch-bot bot commented May 23, 2024

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/126984

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

✅ You can merge normally! (1 Unrelated Failure)

As of commit 9855948 with merge base 5196ef1 (image):

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

mikaylagawarecki added a commit that referenced this pull request May 23, 2024
…le._apply

ghstack-source-id: de44549
Pull Request resolved: #126984
[ghstack-poisoned]
[ghstack-poisoned]
@mikaylagawarecki mikaylagawarecki changed the title Add escape hatch to fall back to old behavior for meta/xla in nn.Module._apply Add private escape hatches to fall back to pre-swap tensors behavior May 29, 2024
@mikaylagawarecki mikaylagawarecki marked this pull request as ready for review May 29, 2024 19:19
[ghstack-poisoned]
[ghstack-poisoned]
mikaylagawarecki added a commit that referenced this pull request May 30, 2024
…le._apply

ghstack-source-id: c56eeca
Pull Request resolved: #126984
@github-actions
Copy link
Contributor

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.

@github-actions github-actions bot added the Stale label Jul 29, 2024
@github-actions github-actions bot closed this Aug 28, 2024
@github-actions github-actions bot deleted the gh/mikaylagawarecki/210/head branch September 28, 2024 02:05
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

release notes: nn release notes category Stale

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants