[Low-bit optim] Support for dcp.save() and dcp.load() (pytorch#1217)

* support dcp.save

* add test for dcp.load()

* fix test

* typo

* implement aten.slice

* skip test

* fix checks

* run ruff

* fix formatting

* remove add safe globals in test

* sort some imports

---------

Co-authored-by: Mark Saroufim <[email protected]>

test/prototype/test_low_bit_optim.py

@@ -197,9 +197,7 @@ def test_optim_8bit_correctness(self, optim_name):
         block_size = 256 if Version(bnb.__version__) >= Version("0.44.0") else 2048
 
         optim1 = getattr(bnb.optim, optim_name)(model1.parameters())
-        optim2 = getattr(low_bit_optim, optim_name)(
-            model2.parameters(), block_size=block_size
-        )
+        optim2 = getattr(low_bit_optim, optim_name)(model2.parameters(), block_size=block_size)
 
         for _ in range(2):
             x = torch.randn(4, 32, device=device)
@@ -315,9 +313,7 @@ def test_optim_cpu_offload_save_load(self):
 
         # resume training
         model2 = copy.deepcopy(model1)
-        optim2 = low_bit_optim.CPUOffloadOptimizer(
-            model2.parameters(), torch.optim.AdamW
-        )
+        optim2 = low_bit_optim.CPUOffloadOptimizer(model2.parameters(), torch.optim.AdamW)
         optim2.load_state_dict(state_dict)
 
         for _ in range(2):
@@ -425,6 +421,9 @@ def test_optim_exclude_low_bit_params(self, optim_name, dtype, device):
                 self.assertTrue(exp_avg.__class__ != torch.Tensor)
                 self.assertTrue(exp_avg_sq.__class__ != torch.Tensor)
 
+_FSDP_WORLD_SIZE = 2
+
+
 class TestFSDP2(FSDPTest):
     @property
     def world_size(self) -> int: