Allow NJT by default for weights_only torch.load (take 2) #140739

mikaylagawarecki · 2024-11-14T19:37:39Z

Per discussion with @malfet, only allow weights_only unpickler to load NJT if torch.nested and torch._dynamo are imported

(this is slightly weird as technically torch.nested is actually imported by default and torch._dynamo.decorators._DimRange is actually what needs to be imported)

we can't import this from torch.nested as this would

undo dynamo lazy import
cause circular import

===========================
Redo of #140304 caused issues as torch.nested._internal.foo needs to be imported, which causes issues like

torch/_weights_only_unpickler.py", line 339, in load
    if full_path in _get_allowed_globals():
torch/_weights_only_unpickler.py", line 188, in _get_allowed_globals
    torch.nested._internal.nested_tensor.NestedTensor
AttributeError: module 'torch.nested' has no attribute '_internal'

This likely wasn't caught in our CI because imports are global during unit tests(?), so we use subprocess to properly test this time

Stack from ghstack (oldest at bottom):

Differential Revision: D65961691

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames

@jbschlosser

[ghstack-poisoned]

pytorch-bot · 2024-11-14T19:37:43Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

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:

[DomainsOnly] Jobs fail with GLIBC version not found

✅ No Failures

As of commit c0c9843 with merge base b379a28 ():
💚 Looks good so far! There are no failures yet. 💚

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

mikaylagawarecki · 2024-11-14T19:58:22Z

@mikaylagawarecki has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

#140304 caused issues as `torch.nested._internal.foo` needs to be imported, which causes issues like ```python torch/_weights_only_unpickler.py", line 339, in load if full_path in _get_allowed_globals(): torch/_weights_only_unpickler.py", line 188, in _get_allowed_globals torch.nested._internal.nested_tensor.NestedTensor AttributeError: module 'torch.nested' has no attribute '_internal' ``` **This likely wasn't caught in our CI because imports are global during unit tests(?)** Further, I am wrapping all these imports in functions that will only be called if the GLOBAL is found in the checkpoint to avoid all users of `torch.load(weights_only=True)` from always paying the cost of importing these Differential Revision: [D65961691](https://our.internmc.facebook.com/intern/diff/D65961691) [ghstack-poisoned]

#140304 caused issues as `torch.nested._internal.foo` needs to be imported, which causes issues like ```python torch/_weights_only_unpickler.py", line 339, in load if full_path in _get_allowed_globals(): torch/_weights_only_unpickler.py", line 188, in _get_allowed_globals torch.nested._internal.nested_tensor.NestedTensor AttributeError: module 'torch.nested' has no attribute '_internal' ``` **This likely wasn't caught in our CI because imports are global during unit tests(?)** Further, I am wrapping all these imports in functions that will only be called if the GLOBAL is found in the checkpoint to avoid all users of `torch.load(weights_only=True)` from always paying the cost of importing these Differential Revision: [D65961691](https://our.internmc.facebook.com/intern/diff/D65961691) cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]

mikaylagawarecki · 2024-11-15T17:54:43Z

@mikaylagawarecki has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

…or DTensor" Same as #140739 but for DTensor (move safe globals for DTensor to `torch.distributed.tensor.__init__` and update error message to let user know `torch.distributed.tensor` must be imported to load DTensor) Differential Revision: [D65961690](https://our.internmc.facebook.com/intern/diff/D65961690) cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]

Same as #140739 but for DTensor (move safe globals for DTensor to `torch.distributed.tensor.__init__` and update error message to let user know `torch.distributed.tensor` must be imported to load DTensor) Differential Revision: [D65961690](https://our.internmc.facebook.com/intern/diff/D65961690) cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]

…or DTensor" Same as #140739 but for DTensor (move safe globals for DTensor to `torch.distributed.tensor.__init__` and update error message to let user know `torch.distributed.tensor` must be imported to load DTensor) Differential Revision: [D65961690](https://our.internmc.facebook.com/intern/diff/D65961690) cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]

Same as #140739 but for DTensor (move safe globals for DTensor to `torch.distributed.tensor.__init__` and update error message to let user know `torch.distributed.tensor` must be imported to load DTensor) Differential Revision: [D65961690](https://our.internmc.facebook.com/intern/diff/D65961690) cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]

…or DTensor" Same as #140739 but for DTensor (move safe globals for DTensor to `torch.distributed.tensor.__init__` and update error message to let user know `torch.distributed.tensor` must be imported to load DTensor) Differential Revision: [D65961690](https://our.internmc.facebook.com/intern/diff/D65961690) cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]

Same as #140739 but for DTensor (move safe globals for DTensor to `torch.distributed.tensor.__init__` and update error message to let user know `torch.distributed.tensor` must be imported to load DTensor) Differential Revision: [D65961690](https://our.internmc.facebook.com/intern/diff/D65961690) cc H-Huang awgu kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o [ghstack-poisoned]

torch/_weights_only_unpickler.py

test/test_serialization.py

Per discussion with malfet, only allow weights_only unpickler to load NJT if `torch.nested` and `torch._dynamo` are imported (this is slightly weird as technically `torch.nested` is actually imported by default and `torch._dynamo.decorators._DimRange` is actually what needs to be imported) we can't import this from `torch.nested` as this would - undo dynamo lazy import - cause circular import =========================== Redo of #140304 caused issues as `torch.nested._internal.foo` needs to be imported, which causes issues like ```python torch/_weights_only_unpickler.py", line 339, in load if full_path in _get_allowed_globals(): torch/_weights_only_unpickler.py", line 188, in _get_allowed_globals torch.nested._internal.nested_tensor.NestedTensor AttributeError: module 'torch.nested' has no attribute '_internal' ``` **This likely wasn't caught in our CI because imports are global during unit tests(?), so we use subprocess to properly test this time** Differential Revision: [D65961691](https://our.internmc.facebook.com/intern/diff/D65961691) cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]