github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-09 13:00:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

[hotfix] quick fixes to make legacy tutorials runnable (#5559)

Browse Source 
Co-authored-by: Edenzzzz <wtan45@wisc.edu>
This commit is contained in:
Edenzzzz
2024-04-07 12:06:27 +08:00
committed by GitHub
parent 8e412a548e
commit 15055f9a36
8 changed files with 20 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
colossalai/nn/layer/scaled_softmax.py
View File

@@ -8,6 +8,14 @@ import torch.nn as nn
from colossalai.kernel.kernel_loader import ScaledMaskedSoftmaxLoader, ScaledUpperTriangleMaskedSoftmaxLoader
# NOTE: These kernels are compiled on specific GPU arch and not widely applicable.
# try:
# from colossalai._C import scaled_masked_softmax as scaled_masked_softmax, scaled_upper_triangle_masked_softmax_cuda as scaled_upper_triang_masked_softmax
# except ImportError:
scaled_masked_softmax = None
scaled_upper_triang_masked_softmax = None
class AttnMaskType(enum.Enum):
padding = 1