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

[example] add llama2 example (#4527)

Browse Source 
* [example] transfer llama-1 example

* [example] fit llama-2

* [example] refactor scripts folder

* [example] fit new gemini plugin

* [cli] fix multinode runner

* [example] fit gemini optim checkpoint

* [example] refactor scripts

* [example] update requirements

* [example] update requirements

* [example] rename llama to llama2

* [example] update readme and pretrain script

* [example] refactor scripts
This commit is contained in:
Hongxin Liu
2023-08-28 17:59:11 +08:00
committed by GitHub
parent 839847b7d7
commit 0b00def881
17 changed files with 1087 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
colossalai/kernel/cuda_native/mha/mem_eff_attn.py
View File

@@ -2,7 +2,13 @@ import warnings
HAS_MEM_EFF_ATTN = False
try:
from xformers.ops.fmha import memory_efficient_attention
from xformers.ops.fmha import MemoryEfficientAttentionCutlassOp, memory_efficient_attention
from xformers.ops.fmha.attn_bias import (
BlockDiagonalCausalMask,
BlockDiagonalMask,
LowerTriangularMask,
LowerTriangularMaskWithTensorBias,
)
HAS_MEM_EFF_ATTN = True
except ImportError:
warnings.warn('please install xformers from https://github.com/facebookresearch/xformers')
@@ -16,13 +22,6 @@ if HAS_MEM_EFF_ATTN:
from typing import Optional
import torch
from xformers.ops.fmha import MemoryEfficientAttentionCutlassOp
from xformers.ops.fmha.attn_bias import (
BlockDiagonalCausalMask,
BlockDiagonalMask,
LowerTriangularMask,
LowerTriangularMaskWithTensorBias,
)
from .utils import SeqLenInfo