nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

flash-attn train and sample

Browse Source 
Former-commit-id: e336565bd5ff2aaf4d0b4ad5f42732e89cd216b3
This commit is contained in:
LinB203
2024-03-09 16:30:09 +08:00
parent 39e398e815
commit 6df13b07bb
12 changed files with 107 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
opensora/dataset/sky_datasets.py
+2 -1
View File
@@ -52,7 +52,8 @@ class Sky(data.Dataset):
for _, meta in enumerate(frame_list):
root = meta[0]
try:
frames = sorted(meta[2], key=lambda item: int(item.split('.')[0].split('_')[-1]))
frames = [i for i in meta[2] if is_image_file(i)]
frames = sorted(frames, key=lambda item: int(item.split('.')[0].split('_')[-1]))
except:
pass
# print(meta[0]) # root
opensora/models/diffusion/dit/dit.py
+57 -22
View File
@@ -18,6 +18,25 @@ from timm.layers import use_fused_attn, to_2tuple
from timm.models.vision_transformer import PatchEmbed, Mlp
from torch.nn import functional as F
# the xformers lib allows less memory, faster training and inference
try:
import xformers
import xformers.ops
except:
XFORMERS_IS_AVAILBLE = False
try:
# needs to have https://github.com/corl-team/rebased/ installed
from fla.ops.triton.rebased_fast import parallel_rebased
except:
REBASED_IS_AVAILABLE = False
try:
# needs to have https://github.com/lucidrains/ring-attention-pytorch installed
from ring_attention_pytorch.ring_flash_attention_cuda import ring_flash_attn_cuda
except:
RING_ATTENTION_IS_AVAILABLE = False
class Attention(nn.Module):
fused_attn: Final[bool]
@@ -31,13 +50,14 @@ class Attention(nn.Module):
attn_drop: float = 0.,
proj_drop: float = 0.,
norm_layer: nn.Module = nn.LayerNorm,
attention_mode='math'
) -> None:
super().__init__()
assert dim % num_heads == 0, 'dim should be divisible by num_heads'
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.scale = self.head_dim ** -0.5
self.fused_attn = use_fused_attn()
self.attention_mode = attention_mode
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
@@ -48,28 +68,43 @@ class Attention(nn.Module):
def forward(self, x: torch.Tensor, attention_mask) -> torch.Tensor:
B, N, C = x.shape
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
q, k, v = qkv.unbind(0)
q, k = self.q_norm(q), self.k_norm(k)
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4).contiguous()
q, k, v = qkv.unbind(0) # make torchscript happy (cannot use tensor as tuple) b h n c
# if self.fused_attn:
# x = F.scaled_dot_product_attention(
# q, k, v,
# attn_mask=attention_mask,
# dropout_p=self.attn_drop.p if self.training else 0.,
# )
# else:
attn = (q @ k.transpose(-2, -1)) * self.scale
if attention_mask is not None:
attn = attn + attention_mask
attn = attn.softmax(dim=-1)
if torch.any(torch.isnan(attn)):
print('torch.any(torch.isnan(attn))')
attn = attn.masked_fill(torch.isnan(attn), float(0.))
attn = self.attn_drop(attn)
x = attn @ v
if self.attention_mode == 'xformers': # cause loss nan while using with amp
assert q.ndim == 4 and (attention_mask is None or torch.all(attention_mask.bool()))
q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
x = xformers.ops.memory_efficient_attention(q, k, v, p=self.attn_drop.p, scale=self.scale).reshape(B, N, C)
elif self.attention_mode == 'flash':
assert attention_mask is None or torch.all(attention_mask.bool())
# cause loss nan while using with amp
# Optionally use the context manager to ensure one of the fused kerenels is run
with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False):
x = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop.p,
scale=self.scale).reshape(B, N,
C) # require pytorch 2.0
elif self.attention_mode == 'math':
attn = (q @ k.transpose(-2, -1)) * self.scale
if attention_mask is not None:
attn = attn + attention_mask
attn = attn.softmax(dim=-1)
if torch.any(torch.isnan(attn)):
print('torch.any(torch.isnan(attn))')
attn = attn.masked_fill(torch.isnan(attn), float(0.))
attn = self.attn_drop(attn)
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
elif self.attention_mode == 'rebased':
x = parallel_rebased(q, k, v, self.eps, True, True).reshape(B, N, C)
elif self.attention_mode == 'ring':
x = ring_flash_attn_cuda(q, k, v, causal=self.causal, bucket_size=self.ring_bucket_size).reshape(B, N, C)
else:
raise NotImplemented
x = x.transpose(1, 2).reshape(B, N, C)
x = self.proj(x)
x = self.proj_drop(x)
return x
@@ -245,7 +280,7 @@ class DiT(nn.Module):
self.temp_embed = nn.Parameter(torch.zeros(1, num_frames, hidden_size), requires_grad=False)
self.blocks = nn.ModuleList([
DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio) for _ in range(depth)
DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, attention_mode=attention_mode) for _ in range(depth)
])
self.final_layer = FinalLayer(hidden_size, patch_size_t, patch_size, self.out_channels)
self.initialize_weights()
opensora/models/diffusion/latte/latte.py
+5 -4
View File
@@ -68,14 +68,15 @@ class Attention(nn.Module):
def forward(self, x, attention_mask):
B, N, C = x.shape
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4).contiguous()
q, k, v = qkv.unbind(0) # make torchscript happy (cannot use tensor as tuple)
q, k, v = qkv.unbind(0) # make torchscript happy (cannot use tensor as tuple) b h n c
if self.attention_mode == 'xformers': # cause loss nan while using with amp
assert q.ndim == 4
q, k, v = q.transpose(1,2), k.transpose(1,2), v.transpose(1,2)
x = xformers.ops.memory_efficient_attention(q, k, v).reshape(B, N, C)
assert q.ndim == 4 and (attention_mask is None or torch.all(attention_mask.bool()))
q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
x = xformers.ops.memory_efficient_attention(q, k, v, p=self.attn_drop.p, scale=self.scale).reshape(B, N, C)
elif self.attention_mode == 'flash':
assert attention_mask is None or torch.all(attention_mask.bool())
# cause loss nan while using with amp
# Optionally use the context manager to ensure one of the fused kerenels is run
with torch.backends.cuda.sdp_kernel(enable_math=False, enable_flash=True, enable_mem_efficient=False):
opensora/sample/sample.py
+25 -4
View File
@@ -9,6 +9,8 @@ Sample new images from a pre-trained Latte.
import os
import sys
from accelerate import Accelerator
from opensora.dataset import ae_denorm
from opensora.models.ae import ae_channel_config, getae, ae_stride_config
from opensora.models.diffusion import Diffusion_models
@@ -28,7 +30,11 @@ def main(args):
# Setup PyTorch:
# torch.manual_seed(args.seed)
torch.set_grad_enabled(False)
device = "cuda" if torch.cuda.is_available() else "cpu"
assert torch.cuda.is_available(), "Training currently requires at least one GPU."
# Setup accelerator:
accelerator = Accelerator()
device = accelerator.device
using_cfg = args.cfg_scale > 1.0
@@ -40,7 +46,8 @@ def main(args):
num_classes=args.num_classes,
in_channels=ae_channel_config[args.ae],
extras=args.extras,
num_frames=args.num_frames // ae_stride_config[args.ae][0]
num_frames=args.num_frames // ae_stride_config[args.ae][0],
attention_mode=args.attention_mode
).to(device)
if args.use_compile:
@@ -52,14 +59,23 @@ def main(args):
model.load_state_dict(state_dict)
model.eval() # important!
model = accelerator.prepare(model)
diffusion = create_diffusion(str(args.num_sampling_steps))
ae = getae(args).to(device)
# if args.attention_mode == 'flash':
# ae.to(dtype=torch.bfloat16)
# model.to(dtype=torch.bfloat16)
# Labels to condition the model with (feel free to change):
# Create sampling noise:
z = torch.randn(1, args.num_frames // ae_stride_config[args.ae][0], model.in_channels, latent_size, latent_size, device=device)
# if args.attention_mode == 'flash':
# z = torch.randn(1, args.num_frames // ae_stride_config[args.ae][0], model.in_channels, latent_size, latent_size, dtype=torch.bfloat16, device=device)
# else:
z = torch.randn(1, args.num_frames // ae_stride_config[args.ae][0], model.module.in_channels, latent_size, latent_size, device=device)
# Setup classifier-free guidance:
if using_cfg:
@@ -83,7 +99,11 @@ def main(args):
sample_fn, z.shape, z, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device
)
samples = ae.decode(samples)
# if args.attention_mode == 'flash':
# samples = samples.to(dtype=torch.bfloat16)
with torch.no_grad():
samples = ae.decode(samples)
# Save and display images:
if not os.path.exists(args.save_video_path):
@@ -112,5 +132,6 @@ if __name__ == "__main__":
parser.add_argument("--cfg-scale", type=float, default=1.0)
parser.add_argument("--use-compile", action="store_true")
parser.add_argument("--sample-method", type=str, default='ddpm')
parser.add_argument("--attention_mode", type=str, choices=['xformers', 'math', 'flash'], default="math")
args = parser.parse_args()
main(args)
opensora/train/train.py
+4 -2
View File
@@ -55,7 +55,6 @@ def main(args):
# Setup an experiment folder:
if accelerator.is_main_process:
print(args)
os.makedirs(args.results_dir, exist_ok=True) # Make results folder (holds all experiment subfolders)
experiment_index = len(glob(f"{args.results_dir}/*"))
model_string_name = args.model.replace("/", "-") # e.g., Latte-XL/2 --> Latte-XL-2 (for naming folders)
@@ -67,6 +66,7 @@ def main(args):
logger = create_logger(experiment_dir)
tb_writer = create_tensorboard(experiment_dir)
logger.info(f"Experiment directory created at {experiment_dir}")
logger.info(f'{args}')
else:
logger = create_logger(None)
tb_writer = None
@@ -91,7 +91,8 @@ def main(args):
num_classes=args.num_classes,
in_channels=ae_channel_config[args.ae],
extras=args.extras,
num_frames=args.num_frames // ae_stride_t
num_frames=args.num_frames // ae_stride_t,
attention_mode=args.attention_mode
)
model.gradient_checkpointing = args.gradient_checkpointing
@@ -319,6 +320,7 @@ if __name__ == "__main__":
parser.add_argument("--use-compile", action="store_true")
parser.add_argument("--lr", type=float, default=1e-4)
parser.add_argument("--lr-warmup-steps", type=int, default=0)
parser.add_argument("--attention_mode", type=str, choices=['xformers', 'math', 'flash'], default="math")
parser.add_argument("--clip-grad-norm", default=None, type=float, help="the maximum gradient norm (default None)")
# --------------------------------------
opensora/utils/utils.py
+1 -1
View File
@@ -158,7 +158,7 @@ def get_experiment_dir(root_dir, args):
if args.gradient_checkpointing:
root_dir += '-Gc'
if args.mixed_precision:
root_dir += '-Amp'
root_dir += f'-{args.mixed_precision}'
if args.max_image_size == 512:
root_dir += '-512'
return root_dir
scripts/land_scope/train_feature.sh
+1 -1
View File
@@ -1,6 +1,6 @@
accelerate launch --multi_gpu --num_processes 7 opensora/train/train_with_feature.py \
accelerate launch --multi_gpu --num_processes 7 --mixed_precision bf16 opensora/train/train_with_feature.py \
--model Latte-XL/122 --dataset landscope_feature \
--ae stabilityai/sd-vae-ft-mse \
--data-path "/remote-home/yeyang/train_movie/*/*256.npy" --extras 1 \
scripts/sample.sh
+2 -2
View File
@@ -1,5 +1,5 @@
python opensora/sample/sample.py \
accelerate launch --num_processes 1 --main_process_port 29501 --mixed_precision bf16 opensora/sample/sample.py \
--model Latte-XL/122 --ae stabilityai/sd-vae-ft-mse \
--ckpt results/011-Latte-XL-122-F128S3-landscope_feature-Gc-Amp/checkpoints/0005000.pt --extras 1 \
--fps 10 --num-frames 128 --image-size 256 --num-sampling-steps 250
--fps 10 --num-frames 128 --image-size 256 --num-sampling-steps 250 --attention_mode flash
scripts/sky/train.sh
+3 -3
View File
@@ -4,6 +4,6 @@ accelerate launch --multi_gpu --num_processes 8 --mixed_precision bf16 opensora/
--model Latte-XL/122 --dataset sky \
--ae stabilityai/sd-vae-ft-mse \
--data-path "/remote-home/yeyang/sky_timelapse/sky_train/" --extras 1 \
--sample-rate 3 --num-frames 16 --max-image-size 256 \
--max-train-steps 1000000 --local-batch-size 5 --lr 1e-4 \
--ckpt-every 500 --log-every 50 --gradient-checkpointing
--sample-rate 1 --num-frames 128 --max-image-size 256 \
--max-train-steps 1000000 --local-batch-size 2 --lr 1e-4 --num-workers 10 \
--ckpt-every 500 --log-every 50 --gradient-checkpointing --attention_mode flash
scripts/sky/train_512.sh
+3 -3
View File
@@ -1,8 +1,8 @@
accelerate launch --multi_gpu --num_processes 8 opensora/train/train.py \
accelerate launch --multi_gpu --num_processes 8 --mixed_precision bf16 opensora/train/train.py \
--model Latte-XL/122 --dataset sky \
--ae stabilityai/sd-vae-ft-mse \
--data-path "/remote-home/yeyang/sky_timelapse/sky_train/" --extras 1 \
--sample-rate 3 --num-frames 32 --max-image-size 512 \
--sample-rate 3 --num-frames 16 --max-image-size 512 \
--max-train-steps 1000000 --local-batch-size 2 --lr 1e-4 \
--ckpt-every 500 --log-every 50 --gradient-checkpointing
--ckpt-every 500 --log-every 50 --gradient-checkpointing --attention_mode flash
scripts/sky/train_feature.sh
+2 -2
View File
@@ -1,9 +1,9 @@
accelerate launch --num_processes 8 opensora/train/train_with_feature.py \
accelerate launch --multi_gpu --num_processes 8 --mixed_precision bf16 opensora/train/train_with_feature.py \
--model Latte-XL/122 --dataset sky_feature \
--ae stabilityai/sd-vae-ft-mse \
--data-path "/remote-home/yeyang/sky_timelapse/sky_train/" --extras 1 \
--sample-rate 3 --num-frames 128 --max-image-size 256 \
--max-train-steps 1000000 --local-batch-size 5 --lr 1e-4 \
--ckpt-every 500 --log-every 50 --num-workers 10 --gradient-checkpointing
--ckpt-every 500 --log-every 50 --gradient-checkpointing --attention_mode flash
scripts/sky/train_feature_512.sh
+2 -2
View File
@@ -1,9 +1,9 @@
accelerate launch --multi_gpu --num_processes 8 opensora/train/train_with_feature.py \
accelerate launch --multi_gpu --num_processes 8 --mixed_precision bf16 opensora/train/train_with_feature.py \
--model Latte-XL/122 --dataset landscope_feature \
--ae stabilityai/sd-vae-ft-mse \
--data-path "/remote-home/yeyang/sky_timelapse/sky_train/" --extras 1 \
--sample-rate 3 --num-frames 64 --max-image-size 512 \
--max-train-steps 1000000 --local-batch-size 5 --lr 1e-4 \
--ckpt-every 500 --log-every 50 --gradient-checkpointing
--ckpt-every 500 --log-every 50 --gradient-checkpointing --attention_mode flash