From 6df13b07bb75d7993e26ffdeb3377f8b1a930b49 Mon Sep 17 00:00:00 2001 From: LinB203 <2267330597@qq.com> Date: Sat, 9 Mar 2024 16:30:09 +0800 Subject: [PATCH] flash-attn train and sample Former-commit-id: e336565bd5ff2aaf4d0b4ad5f42732e89cd216b3 --- opensora/dataset/sky_datasets.py | 3 +- opensora/models/diffusion/dit/dit.py | 79 +++++++++++++++++------- opensora/models/diffusion/latte/latte.py | 9 +-- opensora/sample/sample.py | 29 +++++++-- opensora/train/train.py | 6 +- opensora/utils/utils.py | 2 +- scripts/land_scope/train_feature.sh | 2 +- scripts/sample.sh | 4 +- scripts/sky/train.sh | 6 +- scripts/sky/train_512.sh | 6 +- scripts/sky/train_feature.sh | 4 +- scripts/sky/train_feature_512.sh | 4 +- 12 files changed, 107 insertions(+), 47 deletions(-) diff --git a/opensora/dataset/sky_datasets.py b/opensora/dataset/sky_datasets.py index f491b4466..6410e0d68 100644 --- a/opensora/dataset/sky_datasets.py +++ b/opensora/dataset/sky_datasets.py @@ -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 diff --git a/opensora/models/diffusion/dit/dit.py b/opensora/models/diffusion/dit/dit.py index 4d08787c6..ddc6b4894 100644 --- a/opensora/models/diffusion/dit/dit.py +++ b/opensora/models/diffusion/dit/dit.py @@ -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() diff --git a/opensora/models/diffusion/latte/latte.py b/opensora/models/diffusion/latte/latte.py index c0f616a4e..ed6824f5b 100644 --- a/opensora/models/diffusion/latte/latte.py +++ b/opensora/models/diffusion/latte/latte.py @@ -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): diff --git a/opensora/sample/sample.py b/opensora/sample/sample.py index 3ff60d4aa..34d485334 100644 --- a/opensora/sample/sample.py +++ b/opensora/sample/sample.py @@ -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) diff --git a/opensora/train/train.py b/opensora/train/train.py index 7c908e85c..bd4b851e1 100644 --- a/opensora/train/train.py +++ b/opensora/train/train.py @@ -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)") # -------------------------------------- diff --git a/opensora/utils/utils.py b/opensora/utils/utils.py index 563c43cdd..1ea1ef392 100644 --- a/opensora/utils/utils.py +++ b/opensora/utils/utils.py @@ -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 diff --git a/scripts/land_scope/train_feature.sh b/scripts/land_scope/train_feature.sh index 4e1ef7d1a..6433b6256 100644 --- a/scripts/land_scope/train_feature.sh +++ b/scripts/land_scope/train_feature.sh @@ -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 \ diff --git a/scripts/sample.sh b/scripts/sample.sh index dd105577e..c6ddd8064 100644 --- a/scripts/sample.sh +++ b/scripts/sample.sh @@ -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 diff --git a/scripts/sky/train.sh b/scripts/sky/train.sh index 168dbbd0e..738f6fb91 100644 --- a/scripts/sky/train.sh +++ b/scripts/sky/train.sh @@ -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 \ No newline at end of file + --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 \ No newline at end of file diff --git a/scripts/sky/train_512.sh b/scripts/sky/train_512.sh index 76c2aadff..4450c8d13 100644 --- a/scripts/sky/train_512.sh +++ b/scripts/sky/train_512.sh @@ -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 diff --git a/scripts/sky/train_feature.sh b/scripts/sky/train_feature.sh index 8a77aaab0..1af1d67ec 100644 --- a/scripts/sky/train_feature.sh +++ b/scripts/sky/train_feature.sh @@ -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 \ No newline at end of file + --ckpt-every 500 --log-every 50 --gradient-checkpointing --attention_mode flash \ No newline at end of file diff --git a/scripts/sky/train_feature_512.sh b/scripts/sky/train_feature_512.sh index 902e6c58d..8c3dc09b7 100644 --- a/scripts/sky/train_feature_512.sh +++ b/scripts/sky/train_feature_512.sh @@ -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 \ No newline at end of file + --ckpt-every 500 --log-every 50 --gradient-checkpointing --attention_mode flash \ No newline at end of file