[ColoTensor] rename APIs and add output_replicate to ComputeSpec (#1168)

tests/test_tensor/test_op.py

@@ -36,10 +36,10 @@ def test_layernorm():
 
 def check_spec_eq(tensor, other):
     assert isinstance(tensor, ColoTensor) and isinstance(other, ColoTensor)
-    for k in dir(tensor.spec.dist_spec):
+    for k in dir(tensor.tensor_spec.dist_spec):
         if not k.startswith('__'):
-            assert hasattr(other.spec.dist_spec, k)
-            assert getattr(tensor.spec.dist_spec, k) == getattr(other.spec.dist_spec, k)
+            assert hasattr(other.tensor_spec.dist_spec, k)
+            assert getattr(tensor.tensor_spec.dist_spec, k) == getattr(other.tensor_spec.dist_spec, k)
 
 
 def check_element_wise_ops():

tests/test_tensor/test_tensor.py

@@ -66,7 +66,7 @@ def _run_tensor_shard_init(world_size):
     shard_spec = distspec.shard(process_group=gpc.get_group(ParallelMode.DATA), dims=[0], num_partitions=[world_size])
     tensor_spec = TensorSpec(shard_spec)
     t = ColoTensor.from_torch_tensor(t_ref.clone(), tensor_spec)
-    t.set_spec(TensorSpec(dist_spec=distspec.replicate()))
+    t.set_tensor_spec(TensorSpec(dist_spec=distspec.replicate()))
     assert t.shape == torch.Size((4 * world_size, 5))
 
 

tests/test_tensor/test_zero_optim.py

@@ -51,7 +51,7 @@ def init_1d_row_spec(model):
     with DistSpecManager.no_grad():
         for n, p in model.named_parameters():
             if 'weight' in n and 'ln' not in n:
-                p.set_spec(spec)
+                p.set_tensor_spec(spec)
 
 
 def init_1d_col_spec(model):
@@ -61,7 +61,7 @@ def init_1d_col_spec(model):
     with DistSpecManager.no_grad():
         for n, p in model.named_parameters():
             if 'ln' not in n and ('weight' in n or 'bias' in n):
-                p.set_spec(spec)
+                p.set_tensor_spec(spec)
 
 
 @parameterize('use_chunk', [False, True])