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

[shardformer] remove useless code (#5645)

pull/5654/head
flybird11111 2024-04-25 13:46:39 +08:00 committed by GitHub
parent 148506c828
commit 5d88ef1aaf
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 0 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
colossalai/shardformer/modeling/opt.py
View File

@ -58,20 +58,6 @@ class OPTPipelineForwards:
under pipeline setting.
"""
@staticmethod
def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
"""
Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
"""
bsz, src_len = mask.size()
tgt_len = tgt_len if tgt_len is not None else src_len
expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
inverted_mask = 1.0 - expanded_mask
return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
@staticmethod
def opt_model_forward(
self: OPTModel,