[Tensor ] Add 1Drow weight reshard by spec (#854)

tests/components_to_test/simple_net.py

@@ -1,10 +1,11 @@
+from zmq import device
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from colossalai.nn import CheckpointModule
 from .utils.dummy_data_generator import DummyDataGenerator
 from .registry import non_distributed_component_funcs
-
+from colossalai.utils.cuda import get_current_device
 
 class SimpleNet(CheckpointModule):
     """
@@ -25,8 +26,8 @@ class SimpleNet(CheckpointModule):
 class DummyDataLoader(DummyDataGenerator):
 
     def generate(self):
-        data = torch.rand(16, 4)
-        label = torch.randint(low=0, high=2, size=(16,))
+        data = torch.rand(16, 4, device=get_current_device())
+        label = torch.randint(low=0, high=2, size=(16,), device=get_current_device())
         return data, label
 
 

tests/test_tensor/test_linear_tp.py

@@ -35,7 +35,7 @@ def run_linear_tp1d_row_test():
 
     W_shape = (out_features, in_features)
     W_master = torch.randn(W_shape, dtype=dtype, device=device)
-    W = broadcast_tensor_chunk(W_master, chunk_size=DEPTH, local_rank=local_rank)
+    W = broadcast_tensor_chunk(W_master, chunk_size=1)
     W.requires_grad = True
 
     B_shape = (out_features)
@@ -45,7 +45,7 @@ def run_linear_tp1d_row_test():
 
     # replace the torch nn.Parameters with ColoTensor
     sharded_weight = ColoTensor.init_from_torch_tensor(W)
-    sharded_weight._shard_spec = "1Drow"
+    sharded_weight.set_spec(spec="1Drow") # reshard
     sharded_bias = ColoTensor.init_from_torch_tensor(B)
     replace_parameter_add_grad(layer, sharded_weight, sharded_bias)
     out = layer(A)

tests/test_tensor/test_net_tp.py

@@ -23,9 +23,9 @@ def run_simple_net():
     with ColoInitContext():
         model = model_builder(checkpoint=True)
 
-    # TODO(jzy) we set the Specs for weight of each linear.
-    # model.proj1.weight.set_spec('1Drow')
-    # model.proj2.weight.set_spec('1Drow')
+    # we set the Specs for weight of each linear.
+    model.proj1.weight.set_spec('1Drow')
+    model.proj2.weight.set_spec('1Drow')
 
     for i, (data, label) in enumerate(train_dataloader):
         output = model(data)