Update on "[RFC][FSDP2] Renamed FSDP to FSDPModule"

Andrew Gu · Andrew Gu · commit 93cf3c1b12b8 · 2024-04-25T17:18:23.000-07:00
cc mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse H-Huang kwen2501 penguinwu fegin XilunWu wanchaol fduwjj wz337 tianyu-l wconstab yf225 chauhang d4l3k

[ghstack-poisoned]
diff --git a/test/distributed/_composable/fsdp/test_fully_shard_comm.py b/test/distributed/_composable/fsdp/test_fully_shard_comm.py
@@ -12,7 +12,7 @@
 
 from torch.distributed._composable import checkpoint, replicate
 from torch.distributed._composable.fsdp import (
-    FSDP,
+    FSDPModule,
     fully_shard,
     MixedPrecisionPolicy,
     OffloadPolicy,
@@ -630,13 +630,13 @@ def __init__(self, dim: int, mesh: DeviceMesh):
 
             def forward(self, x: torch.Tensor):
                 y1, work1 = self.reduce_module1(x)
-                if isinstance(self.mlps.mlp1, FSDP):
+                if isinstance(self.mlps.mlp1, FSDPModule):
                     self.mlps.mlp1.unshard(async_op=True)
                 y2, work2 = self.reduce_module2(x)
-                if isinstance(self.mlps.mlp2, FSDP):
+                if isinstance(self.mlps.mlp2, FSDPModule):
                     self.mlps.mlp2.unshard(async_op=True)
                 y3, work3 = self.reduce_module3(x)
-                if isinstance(self.mlps.mlp3, FSDP):
+                if isinstance(self.mlps.mlp3, FSDPModule):
                     self.mlps.mlp3.unshard(async_op=True)
                 return self.mlps([y1, y2, y3], [work1, work2, work3])
 
