[optim] Fix low-bit optim when used with FSDP2+CPUOffload (#2195)

Author: gau-nernst

test/test_low_bit_optim.py

@@ -11,7 +11,11 @@
 import pytest
 import torch
 from torch import nn
-from torch.distributed._composable.fsdp import fully_shard
+from torch.distributed._composable.fsdp import (
+    fully_shard,
+    CPUOffloadPolicy,
+    OffloadPolicy,
+)
 from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
 from torch.testing._internal.common_fsdp import FSDPTest
 from torch.testing._internal.common_utils import (
@@ -427,16 +431,21 @@ def world_size(self) -> int:
     @skip_if_lt_x_gpu(_FSDP_WORLD_SIZE)
     @skip_if_rocm("ROCm enablement in progress")
     def test_fsdp2(self):
-        optim_classes = [optim.AdamW8bit, optim.AdamW4bit]
+        # we do this to avoid all combinations
+        args_list = [
+            (optim.AdamW8bit, OffloadPolicy),
+            (optim.AdamW4bit, OffloadPolicy),
+            (optim.AdamW8bit, CPUOffloadPolicy),
+        ]
         if torch.cuda.get_device_capability() >= (8, 9):
-            optim_classes.append(optim.AdamWFp8)
+            args_list.append((optim.AdamWFp8, OffloadPolicy))
 
         self.run_subtests(
-            {"optim_cls": optim_classes},
+            {"args": args_list},
             self._test_fsdp2,
         )
 
-    def _test_fsdp2(self, optim_cls):
+    def _test_fsdp2(self, args):
         import torch.distributed as dist
         import torch.distributed.checkpoint as dcp
         import torch.utils._pytree as pytree
@@ -447,6 +456,8 @@ def _test_fsdp2(self, optim_cls):
             TransformerBlock,
         )
 
+        optim_cls, offload_policy = args
+
         batch_size = 3
         vocab_size = 1024
         seq_len = 64
@@ -466,8 +477,8 @@ def _test_fsdp2(self, optim_cls):
         fsdp_model = copy.deepcopy(base_model)
         for m in fsdp_model.modules():
             if isinstance(m, TransformerBlock):
-                fully_shard(m)
-        fully_shard(fsdp_model)
+                fully_shard(m, offload_policy=offload_policy)
+        fully_shard(fsdp_model, offload_policy=offload_policy)
         fsdp_optim = optim_cls(fsdp_model.parameters(), lr=1e-2)
 
         torch.manual_seed(42 + self.rank + 1)

torchao/optim/adam.py

@@ -83,6 +83,11 @@ def _new_buffer(self, p: Tensor, signed: bool):
                 stride=p.stride(),
             )
 
+        # when there is CPU offload, p.device is cpu, but device_mesh.device_type is cuda.
+        # DTensor.from_local() will move local_tensor to device_mesh.device_type.
+        # hence, we need to manually move it back to CPU.
+        # https://github.com/pytorch/pytorch/blob/bc4cf1c1/torch/distributed/tensor/_api.py#L410-L415
+        out = out.to(p.device)
         return out
 
     @torch.no_grad()