nicolabs/Open-Sora
1
0
Fork 0
mirror of https://github.com/hpcaitech/Open-Sora.git synced 2026-05-21 11:59:01 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

change flashattn to flash_attn (#387)

Browse source
This commit is contained in:
Zheng Zangwei (Alex Zheng) 2024-05-10 14:36:31 +08:00 committed by GitHub
parent 07f35a48e4
commit a6a0325b47
30 changed files with 52 additions and 52 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
README.md
View file

@ -161,7 +161,7 @@ conda activate opensora
pip install torch torchvision
# install flash attention (optional)
# set enable_flashattn=False in config to avoid using flash attention
# set enable_flash_attn=False in config to avoid using flash attention
pip install packaging ninja
pip install flash-attn --no-build-isolation

2
configs/dit/train/16x256x256.py
View file

@ -18,7 +18,7 @@ sp_size = 1
model = dict(
type="DiT-XL/2",
from_pretrained="DiT-XL-2-256x256.pt",
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
)
vae = dict(

2
configs/dit/train/1x256x256.py
View file

@ -19,7 +19,7 @@ sp_size = 1
model = dict(
type="DiT-XL/2",
no_temporal_pos_emb=True,
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
)
vae = dict(

2
configs/latte/train/16x256x256.py
View file

@ -17,7 +17,7 @@ sp_size = 1
# Define model
model = dict(
type="Latte-XL/2",
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
)
vae = dict(

2
configs/opensora-v1-1/train/benchmark.py
View file

@ -65,7 +65,7 @@ model = dict(
from_pretrained=None,
input_sq_size=512, # pretrained model is trained on 512x512
qk_norm=True,
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
)
vae = dict(

2
configs/opensora-v1-1/train/image.py
View file

@ -29,7 +29,7 @@ model = dict(
from_pretrained=None,
input_sq_size=512, # pretrained model is trained on 512x512
qk_norm=True,
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
)
vae = dict(

2
configs/opensora-v1-1/train/stage1.py
View file

@ -41,7 +41,7 @@ model = dict(
from_pretrained=None,
input_sq_size=512, # pretrained model is trained on 512x512
qk_norm=True,
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
)
vae = dict(

2
configs/opensora-v1-1/train/stage2.py
View file

@ -43,7 +43,7 @@ model = dict(
from_pretrained=None,
input_sq_size=512, # pretrained model is trained on 512x512
qk_norm=True,
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
)
vae = dict(

2
configs/opensora-v1-1/train/stage3.py
View file

@ -43,7 +43,7 @@ model = dict(
from_pretrained=None,
input_sq_size=512, # pretrained model is trained on 512x512
qk_norm=True,
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
)
vae = dict(

2
configs/opensora-v1-1/train/video.py
View file

@ -31,7 +31,7 @@ model = dict(
from_pretrained=None,
input_sq_size=512, # pretrained model is trained on 512x512
qk_norm=True,
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
)
vae = dict(

2
configs/opensora/inference/16x256x256.py
View file

@ -7,7 +7,7 @@ model = dict(
type="STDiT-XL/2",
space_scale=0.5,
time_scale=1.0,
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
from_pretrained="PRETRAINED_MODEL",
)

2
configs/opensora/inference/16x512x512.py
View file

@ -7,7 +7,7 @@ model = dict(
type="STDiT-XL/2",
space_scale=1.0,
time_scale=1.0,
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
from_pretrained="PRETRAINED_MODEL",
)

2
configs/opensora/inference/64x512x512.py
View file

@ -7,7 +7,7 @@ model = dict(
type="STDiT-XL/2",
space_scale=1.0,
time_scale=2 / 3,
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
from_pretrained="PRETRAINED_MODEL",
)

2
configs/opensora/train/16x256x256-mask.py
View file

@ -20,7 +20,7 @@ model = dict(
space_scale=0.5,
time_scale=1.0,
from_pretrained="PixArt-XL-2-512x512.pth",
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
)
mask_ratios = {

2
configs/opensora/train/16x256x256-spee.py
View file

@ -20,7 +20,7 @@ model = dict(
space_scale=0.5,
time_scale=1.0,
from_pretrained="PixArt-XL-2-512x512.pth",
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
)
mask_ratios = {

2
configs/opensora/train/16x256x256.py
View file

@ -20,7 +20,7 @@ model = dict(
space_scale=0.5,
time_scale=1.0,
from_pretrained="PixArt-XL-2-512x512.pth",
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
)
vae = dict(

2
configs/opensora/train/16x512x512.py
View file

@ -20,7 +20,7 @@ model = dict(
space_scale=1.0,
time_scale=1.0,
from_pretrained=None,
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
)
vae = dict(

2
configs/opensora/train/360x512x512.py
View file

@ -26,7 +26,7 @@ model = dict(
space_scale=1.0,
time_scale=2 / 3,
from_pretrained=None,
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
enable_sequence_parallelism=True, # enable sq here
)

2
configs/opensora/train/64x512x512-sp.py
View file

@ -20,7 +20,7 @@ model = dict(
space_scale=1.0,
time_scale=2 / 3,
from_pretrained=None,
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
enable_sequence_parallelism=True, # enable sq here
)

2
configs/opensora/train/64x512x512.py
View file

@ -20,7 +20,7 @@ model = dict(
space_scale=1.0,
time_scale=2 / 3,
from_pretrained=None,
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
)
vae = dict(

2
configs/pixart/train/16x256x256.py
View file

@ -20,7 +20,7 @@ model = dict(
space_scale=0.5,
time_scale=1.0,
from_pretrained="PixArt-XL-2-512x512.pth",
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
)
vae = dict(

2
configs/pixart/train/1x512x512.py
View file

@ -21,7 +21,7 @@ model = dict(
time_scale=1.0,
no_temporal_pos_emb=True,
from_pretrained="PixArt-XL-2-512x512.pth",
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
)
vae = dict(

2
configs/pixart/train/64x512x512.py
View file

@ -21,7 +21,7 @@ model = dict(
space_scale=1.0,
time_scale=2 / 3,
from_pretrained=None,
enable_flashattn=True,
enable_flash_attn=True,
enable_layernorm_kernel=True,
)
vae = dict(

6
docs/config.md
View file

@ -26,8 +26,8 @@ model = dict(
from_pretrained="PRETRAINED_MODEL", # (Optional) Load from pretrained model
input_sq_size=512, # Base spatial position embedding size
qk_norm=True, # Normalize query and key in attention
enable_flashattn=True, # (Optional) Speed up training and inference with flash attention
# Turn enable_flashattn to False if you skip flashattn installation
enable_flash_attn=True, # (Optional) Speed up training and inference with flash attention
# Turn enable_flash_attn to False if you skip flashattn installation
enable_layernorm_kernel=True, # (Optional) Speed up training and inference with fused kernel
# Turn enable_layernorm_kernel to False if you skip apex installation
)
@ -177,7 +177,7 @@ model = dict(
from_pretrained=None, # Load from pretrained model
input_sq_size=512, # Base spatial position embedding size
qk_norm=True, # Normalize query and key in attention
enable_flashattn=True, # (Optional) Speed up training and inference with flash attention
enable_flash_attn=True, # (Optional) Speed up training and inference with flash attention
enable_layernorm_kernel=True, # (Optional) Speed up training and inference with fused kernel
)
vae = dict(

4
docs/zh_CN/structure.md
View file

@ -89,7 +89,7 @@ model = dict(
type="STDiT-XL/2", # Select model type (STDiT-XL/2, DiT-XL/2, etc.)
space_scale=1.0, # (Optional) Space positional encoding scale (new height / old height)
time_scale=2 / 3, # (Optional) Time positional encoding scale (new frame_interval / old frame_interval)
enable_flashattn=True, # (Optional) Speed up training and inference with flash attention
enable_flash_attn=True, # (Optional) Speed up training and inference with flash attention
enable_layernorm_kernel=True, # (Optional) Speed up training and inference with fused kernel
from_pretrained="PRETRAINED_MODEL", # (Optional) Load from pretrained model
no_temporal_pos_emb=True, # (Optional) Disable temporal positional encoding (for image)
@ -144,7 +144,7 @@ model = dict(
space_scale=1.0,
time_scale=2 / 3,
from_pretrained="YOUR_PRETRAINED_MODEL",
enable_flashattn=True, # Enable flash attention
enable_flash_attn=True, # Enable flash attention
enable_layernorm_kernel=True, # Enable layernorm kernel
)
vae = dict(

14
opensora/models/dit/dit.py
View file

@ -44,13 +44,13 @@ class DiTBlock(nn.Module):
hidden_size,
num_heads,
mlp_ratio=4.0,
enable_flashattn=False,
enable_flash_attn=False,
enable_layernorm_kernel=False,
):
super().__init__()
self.hidden_size = hidden_size
self.num_heads = num_heads
self.enable_flashattn = enable_flashattn
self.enable_flash_attn = enable_flash_attn
mlp_hidden_dim = int(hidden_size * mlp_ratio)
self.norm1 = get_layernorm(hidden_size, eps=1e-6, affine=False, use_kernel=enable_layernorm_kernel)
@ -58,7 +58,7 @@ class DiTBlock(nn.Module):
hidden_size,
num_heads=num_heads,
qkv_bias=True,
enable_flashattn=enable_flashattn,
enable_flash_attn=enable_flash_attn,
)
self.norm2 = get_layernorm(hidden_size, eps=1e-6, affine=False, use_kernel=enable_layernorm_kernel)
self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
@ -93,7 +93,7 @@ class DiT(nn.Module):
caption_channels=512,
model_max_length=77,
dtype=torch.float32,
enable_flashattn=False,
enable_flash_attn=False,
enable_layernorm_kernel=False,
enable_sequence_parallelism=False,
):
@ -111,7 +111,7 @@ class DiT(nn.Module):
self.num_heads = num_heads
self.dtype = dtype
self.use_text_encoder = not condition.startswith("label")
if enable_flashattn:
if enable_flash_attn:
assert dtype in [
torch.float16,
torch.bfloat16,
@ -143,7 +143,7 @@ class DiT(nn.Module):
hidden_size,
num_heads,
mlp_ratio=mlp_ratio,
enable_flashattn=enable_flashattn,
enable_flash_attn=enable_flash_attn,
enable_layernorm_kernel=enable_layernorm_kernel,
)
for _ in range(depth)
@ -152,7 +152,7 @@ class DiT(nn.Module):
self.final_layer = FinalLayer(hidden_size, np.prod(self.patch_size), self.out_channels)
self.initialize_weights()
self.enable_flashattn = enable_flashattn
self.enable_flash_attn = enable_flash_attn
self.enable_layernorm_kernel = enable_layernorm_kernel
def get_spatial_pos_embed(self):

12
opensora/models/pixart/pixart.py
View file

@ -61,13 +61,13 @@ class PixArtBlock(nn.Module):
num_heads,
mlp_ratio=4.0,
drop_path=0.0,
enable_flashattn=False,
enable_flash_attn=False,
enable_layernorm_kernel=False,
enable_sequence_parallelism=False,
):
super().__init__()
self.hidden_size = hidden_size
self.enable_flashattn = enable_flashattn
self.enable_flash_attn = enable_flash_attn
self._enable_sequence_parallelism = enable_sequence_parallelism
if enable_sequence_parallelism:
@ -82,7 +82,7 @@ class PixArtBlock(nn.Module):
hidden_size,
num_heads=num_heads,
qkv_bias=True,
enable_flashattn=enable_flashattn,
enable_flash_attn=enable_flash_attn,
)
self.cross_attn = self.mha_cls(hidden_size, num_heads)
self.norm2 = get_layernorm(hidden_size, eps=1e-6, affine=False, use_kernel=enable_layernorm_kernel)
@ -130,7 +130,7 @@ class PixArt(nn.Module):
freeze=None,
space_scale=1.0,
time_scale=1.0,
enable_flashattn=False,
enable_flash_attn=False,
enable_layernorm_kernel=False,
enable_sequence_parallelism=False,
):
@ -152,7 +152,7 @@ class PixArt(nn.Module):
self.no_temporal_pos_emb = no_temporal_pos_emb
self.depth = depth
self.mlp_ratio = mlp_ratio
self.enable_flashattn = enable_flashattn
self.enable_flash_attn = enable_flash_attn
self.enable_layernorm_kernel = enable_layernorm_kernel
self.space_scale = space_scale
self.time_scale = time_scale
@ -179,7 +179,7 @@ class PixArt(nn.Module):
num_heads,
mlp_ratio=mlp_ratio,
drop_path=drop_path[i],
enable_flashattn=enable_flashattn,
enable_flash_attn=enable_flash_attn,
enable_layernorm_kernel=enable_layernorm_kernel,
)
for i in range(depth)

14
opensora/models/stdit/stdit.py
View file

@ -37,13 +37,13 @@ class STDiTBlock(nn.Module):
d_t=None,
mlp_ratio=4.0,
drop_path=0.0,
enable_flashattn=False,
enable_flash_attn=False,
enable_layernorm_kernel=False,
enable_sequence_parallelism=False,
):
super().__init__()
self.hidden_size = hidden_size
self.enable_flashattn = enable_flashattn
self.enable_flash_attn = enable_flash_attn
self._enable_sequence_parallelism = enable_sequence_parallelism
if enable_sequence_parallelism:
@ -58,7 +58,7 @@ class STDiTBlock(nn.Module):
hidden_size,
num_heads=num_heads,
qkv_bias=True,
enable_flashattn=enable_flashattn,
enable_flash_attn=enable_flash_attn,
)
self.cross_attn = self.mha_cls(hidden_size, num_heads)
self.norm2 = get_layernorm(hidden_size, eps=1e-6, affine=False, use_kernel=enable_layernorm_kernel)
@ -82,7 +82,7 @@ class STDiTBlock(nn.Module):
hidden_size,
num_heads=num_heads,
qkv_bias=True,
enable_flashattn=self.enable_flashattn,
enable_flash_attn=self.enable_flash_attn,
)
def t_mask_select(self, x, masked_x, x_mask):
@ -174,7 +174,7 @@ class STDiT(nn.Module):
space_scale=1.0,
time_scale=1.0,
freeze=None,
enable_flashattn=False,
enable_flash_attn=False,
enable_layernorm_kernel=False,
enable_sequence_parallelism=False,
):
@ -194,7 +194,7 @@ class STDiT(nn.Module):
self.no_temporal_pos_emb = no_temporal_pos_emb
self.depth = depth
self.mlp_ratio = mlp_ratio
self.enable_flashattn = enable_flashattn
self.enable_flash_attn = enable_flash_attn
self.enable_layernorm_kernel = enable_layernorm_kernel
self.space_scale = space_scale
self.time_scale = time_scale
@ -221,7 +221,7 @@ class STDiT(nn.Module):
self.num_heads,
mlp_ratio=self.mlp_ratio,
drop_path=drop_path[i],
enable_flashattn=self.enable_flashattn,
enable_flash_attn=self.enable_flash_attn,
enable_layernorm_kernel=self.enable_layernorm_kernel,
enable_sequence_parallelism=enable_sequence_parallelism,
d_t=self.num_temporal,

4
tests/test_attn.py
View file

@ -11,7 +11,7 @@ B, S, H = 128, 32, 1152
N, D = 16, 72
def run_attn(enable_flashattn: bool):
def run_attn(enable_flash_attn: bool):
get_accelerator().reset_peak_memory_stats()
rope = RotaryEmbedding(D).to(device=get_current_device(), dtype=torch.bfloat16)
attn = Attention(
@ -19,7 +19,7 @@ def run_attn(enable_flashattn: bool):
N,
qkv_bias=True,
rope=rope.rotate_queries_or_keys,
enable_flashattn=enable_flashattn,
enable_flash_attn=enable_flash_attn,
).to(device=get_current_device(), dtype=torch.bfloat16)
x = torch.randn(B, S, H, device=get_current_device(), dtype=torch.bfloat16).requires_grad_()
y = attn(x)

4
tests/test_seq_parallel_attention.py
View file

@ -18,14 +18,14 @@ def run_attention(rank, world_size):
torch.manual_seed(1024)
set_sequence_parallel_group(dist.group.WORLD)
seq_parallel_attention = SeqParallelAttention(dim=256, num_heads=4, qkv_bias=True, enable_flashattn=False).cuda()
seq_parallel_attention = SeqParallelAttention(dim=256, num_heads=4, qkv_bias=True, enable_flash_attn=False).cuda()
torch.manual_seed(1024)
attention = Attention(
dim=256,
num_heads=4,
qkv_bias=True,
enable_flashattn=False,
enable_flash_attn=False,
).cuda()
# create inputs