Skip to content

Commit 5d88ef1

Browse files
authored
[shardformer] remove useless code (#5645)
1 parent 148506c commit 5d88ef1

File tree

1 file changed

+0
-14
lines changed
  • colossalai/shardformer/modeling

1 file changed

+0
-14
lines changed

colossalai/shardformer/modeling/opt.py

Lines changed: 0 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -58,20 +58,6 @@ class OPTPipelineForwards:
5858
under pipeline setting.
5959
"""
6060

61-
@staticmethod
62-
def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
63-
"""
64-
Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
65-
"""
66-
bsz, src_len = mask.size()
67-
tgt_len = tgt_len if tgt_len is not None else src_len
68-
69-
expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
70-
71-
inverted_mask = 1.0 - expanded_mask
72-
73-
return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
74-
7561
@staticmethod
7662
def opt_model_forward(
7763
self: OPTModel,

0 commit comments

Comments
 (0)