github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-12 12:47:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Tensor] rename some APIs in TensorSpec and Polish view unittest (#1176)

Browse Source
This commit is contained in:
Jiarui Fang
2022-06-27 15:56:11 +08:00
committed by GitHub
parent dd0420909f
commit 0dd4e2bbfb
9 changed files with 26 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
colossalai/nn/_ops/addmm.py
View File

@@ -72,10 +72,10 @@ def colo_addmm(input_tensor: GeneralTensor,
assert input_tensor.tensor_spec.is_replicate(), 'Invalid input spec for native addmm op'
ret_tensor = ColoTensor.from_torch_tensor(torch.addmm(input_tensor, mat1, mat2, beta=beta, alpha=alpha))
elif mat2.tensor_spec.has_compute_pattern(ComputePattern.TP1D): # Single Model Parallel Applied
if mat2.tensor_spec.is_1D_row() and input_tensor.tensor_spec.is_replicate():
if mat2.tensor_spec.is_shard_1drow() and input_tensor.tensor_spec.is_replicate():
mode = 'row'
elif mat2.tensor_spec.is_1D_col() and (input_tensor.tensor_spec.is_1D_col()
or input_tensor.tensor_spec.is_1D_row()):
elif mat2.tensor_spec.is_shard_1dcol() and (input_tensor.tensor_spec.is_shard_1dcol()
or input_tensor.tensor_spec.is_shard_1drow()):
mode = 'col'
else:
raise NotImplementedError

5
colossalai/nn/_ops/embedding.py
View File

@@ -32,6 +32,7 @@ def colo_embedding_1Dcol(input_tensor: ColoTensor,
output = ColoTensor.from_torch_tensor(output_parallel, spec=output_spec)
compute_spec = weight.tensor_spec.compute_spec
if compute_spec.output_replicate:
return output.to_replicate()
else:
@@ -125,9 +126,9 @@ def colo_embedding(input_tensor: GeneralTensor,
scale_grad_by_freq=scale_grad_by_freq,
sparse=sparse))
elif weight.tensor_spec.has_compute_pattern(ComputePattern.TP1D): # Single Model Parallel Applied
if weight.tensor_spec.is_1D_row():
if weight.tensor_spec.is_shard_1drow():
mode = 'row'
elif weight.tensor_spec.is_1D_col():
elif weight.tensor_spec.is_shard_1dcol():
mode = 'col'
else:
raise NotImplementedError

2
colossalai/nn/_ops/embedding_bag.py
View File

@@ -104,7 +104,7 @@ def colo_embedding_bag(input_tensor: GeneralTensor,
include_last_offset=include_last_offset,
padding_idx=padding_idx))
elif weight.tensor_spec.has_compute_pattern(ComputePattern.TP1D): # Single Model Parallel Applied
if weight.tensor_spec.is_1D_col():
if weight.tensor_spec.is_shard_1dcol():
tp_mode = 'col'
else:
raise NotImplementedError

6
colossalai/nn/_ops/linear.py
View File

@@ -71,10 +71,10 @@ def colo_linear_imp(input_tensor: GeneralTensor,
assert bias is None or bias.tensor_spec.is_replicate(), 'Invalid bias spec for native Linear op'
ret_tensor = ColoTensor.from_torch_tensor(F.linear(input_tensor, weight, bias))
elif weight.tensor_spec.has_compute_pattern(ComputePattern.TP1D): # Single Model Parallel Applied
if weight.tensor_spec.is_1D_col() and (bias is None or bias.tensor_spec.is_replicate()):
if weight.tensor_spec.is_shard_1dcol() and (bias is None or bias.tensor_spec.is_replicate()):
mode = 'row'
elif weight.tensor_spec.is_1D_row() and (bias is None or bias.tensor_spec.is_1D_row()
or bias.tensor_spec.is_1D_col()):
elif weight.tensor_spec.is_shard_1drow() and (bias is None or bias.tensor_spec.is_shard_1drow()
or bias.tensor_spec.is_shard_1dcol()):
mode = 'col'
else:
raise RuntimeError(f"the weight or bias tensor spec is not valid, weight {weight.tensor_spec}, bias {bias}")

2
colossalai/nn/_ops/loss.py
View File

@@ -29,7 +29,7 @@ def colo_cross_entropy(input_tensor: GeneralTensor,
label_smoothing=label_smoothing)
return ColoTensor.from_torch_tensor(output)
elif input_tensor.has_compute_spec(): # Single Model Parallel Applied
if input_tensor.tensor_spec.is_1D_col():
if input_tensor.tensor_spec.is_shard_1dcol():
output = VocabParallelCrossEntropyLoss1D()(input_tensor, target)
return ColoTensor.from_torch_tensor(output)
else: