Skip to content

GitLab

Unverified Commit ccf3c58c authored by ver217's avatar ver217 Committed by GitHub
Browse files

embedding op use gather_out (#1143)

parent e61dc31b
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
colossalai/nn/_ops/embedding.py
View file @ ccf3c58c
...@@ -30,7 +30,8 @@ def colo_embedding_1Dcol(input_tensor: ColoTensor, ...@@ -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()]), distspec.shard(weight.spec.get_process_group(), [-1], [weight.spec.get_process_group_size()]),
ParallelAction(ComputePattern.TP1D)) ParallelAction(ComputePattern.TP1D))
output = ColoTensor.from_torch_tensor(output_parallel, spec=output_spec) 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 return output
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment