Update on "[FSDP] Add unsafe setattr gated by env var"

Andrew Gu · Andrew Gu · commit a68a119ffa8b · 2023-03-08T21:45:09.000Z
This adds the option to use an unsafe `setattr` for `_use_sharded_views()` and `_use_unsharded_views()` gated by the environment variable `FSDP_USE_UNSAFE_SETATTR`, where a value of `1` means to use the unsafe `setattr`. The unsafe option is disabled by default.

The unsafe `setattr` may be able to save CPU overhead and may be used to intentionally bypass `setattr` checks. Both `_use_sharded_views()` and `_use_unsharded_views()` must use the unsafe version or use the safe versions atomically.

[ghstack-poisoned]
diff --git a/test/distributed/fsdp/test_fsdp_misc.py b/test/distributed/fsdp/test_fsdp_misc.py
@@ -682,7 +682,6 @@ def __setattr__(self, name: str, value: Any) -> None:
         os.environ[_FSDP_USE_UNSAFE_SETATTR] = "1"
         module = SetattrLinear(5, 5, torch.device("cuda"))
         fsdp_module = FSDP(module, use_orig_params=use_orig_params)
-        inp = torch.randn((8, 5), device=torch.device("cuda"))
         called_setattr_override = False
         fsdp_module(inp)
         self.assertFalse(called_setattr_override)
@@ -691,7 +690,6 @@ def __setattr__(self, name: str, value: Any) -> None:
         os.environ[_FSDP_USE_UNSAFE_SETATTR] = "0"
         module = SetattrLinear(5, 5, torch.device("cuda"))
         fsdp_module = FSDP(module, use_orig_params=use_orig_params)
-        inp = torch.randn((8, 5), device=torch.device("cuda"))
         called_setattr_override = False
         fsdp_module(inp)
         self.assertTrue(called_setattr_override)
