[dynamo][fsdp] Consistent behavior of int attributes (#157262)

Reimpl of https://github.com/pytorch/pytorch/pull/150954

Pull Request resolved: https://github.com/pytorch/pytorch/pull/157262
Approved by: https://github.com/bdhirsh
This commit is contained in:
Animesh Jain
2025-06-30 12:04:14 -07:00
committed by PyTorch MergeBot
parent a9352bd25e
commit 42b48ee672
4 changed files with 98 additions and 11 deletions

View File

@ -155,17 +155,6 @@ class GuardSource(enum.Enum):
return self in (GuardSource.GLOBAL_FSDP_MODULE, GuardSource.LOCAL_FSDP_MODULE)
def is_specialized_nn_module(self) -> bool:
import torch._dynamo.config as config
if config._unsafe_skip_fsdp_module_guards:
return (
self
in (
GuardSource.GLOBAL_SPECIALIZED_NN_MODULE,
GuardSource.LOCAL_SPECIALIZED_NN_MODULE,
)
or self.is_fsdp_module()
)
return self in (
GuardSource.GLOBAL_SPECIALIZED_NN_MODULE,
GuardSource.LOCAL_SPECIALIZED_NN_MODULE,