embedding op use gather_out (#1143)

2 years ago · ccf3c58c89
1 changed files with 2 additions and 1 deletions
--- a/colossalai/nn/_ops/embedding.py
+++ b/colossalai/nn/_ops/embedding.py
@ -30,7 +30,8 @@ def colo_embedding_1Dcol(input_tensor: ColoTensor,
        distspec.shard(weight.spec.get_process_group(), [-1], [weight.spec.get_process_group_size()]),
        ParallelAction(ComputePattern.TP1D))
    output = ColoTensor.from_torch_tensor(output_parallel, spec=output_spec)
-    output = output.convert_to_dist_spec(distspec.replicate(weight.spec.get_process_group()))
+    if weight.spec.parallel_action.gather_out:
+        output = output.convert_to_dist_spec(distspec.replicate(weight.spec.get_process_group()))
    return output