Update on "[RFC][FSDP] Don't move ignored params / buffers to device"

rohan-varma · rohan-varma · commit 9a5eae4acacc · 2023-09-05T10:44:53.000-07:00
Since these are ignored by FSDP, don't move them. Differential Revision: [D48727044](https://our.internmc.facebook.com/intern/diff/D48727044/) [ghstack-poisoned]
diff --git a/torch/distributed/fsdp/_init_utils.py b/torch/distributed/fsdp/_init_utils.py
@@ -508,11 +508,9 @@ def _init_param_handle_from_module(
         )
 
     ignored_buffers = {
-        [
-            buffer
-            for ignored_module in state._ignored_modules
-            for buffer in ignored_module.buffers()
-        ]
+        buffer
+        for ignored_module in state._ignored_modules
+        for buffer in ignored_module.buffers()
     }
 
     _move_module_to_device(
