github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[NFC] polish colossalai/nn/_ops/embedding_bag.py code style (#1552)

pull/1550/head
Maruyama_Aya 2022-09-08 14:42:02 +08:00 committed by Frank Lee
parent 73e9eb13b7
commit bd2d789832
1 changed files with 14 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -90,22 +90,21 @@ def colo_embedding_bag(input_tensor: GeneralTensor,
# Handle differen parallel actions.
if not weight.has_compute_spec(): # No Model Parallel Applied
if not weight.has_compute_spec(): # No Model Parallel Applied
assert weight.is_replicate(), 'Invalid weight spec for native embedding op'
return ColoTensor.from_torch_tensor(
tensor=F.embedding_bag(input_tensor,
weight,
offsets=offsets,
max_norm=max_norm,
norm_type=norm_type,
scale_grad_by_freq=scale_grad_by_freq,
mode=mode,
sparse=sparse,
per_sample_weights=per_sample_weights,
include_last_offset=include_last_offset,
padding_idx=padding_idx),
spec=ColoTensorSpec(weight.get_process_group()))
elif weight.has_compute_pattern(ComputePattern.TP1D): # Single Model Parallel Applied
return ColoTensor.from_torch_tensor(tensor=F.embedding_bag(input_tensor,
weight,
offsets=offsets,
max_norm=max_norm,
norm_type=norm_type,
scale_grad_by_freq=scale_grad_by_freq,
mode=mode,
sparse=sparse,
per_sample_weights=per_sample_weights,
include_last_offset=include_last_offset,
padding_idx=padding_idx),
spec=ColoTensorSpec(weight.get_process_group()))
elif weight.has_compute_pattern(ComputePattern.TP1D): # Single Model Parallel Applied
if weight.is_shard_1dcol():
tp_mode = 'col'
else: