diff --git a/fastvideo/attention/backends/sage_attn3.py b/fastvideo/attention/backends/sage_attn3.py index 8c4da63d4..2bfb5a283 100644 --- a/fastvideo/attention/backends/sage_attn3.py +++ b/fastvideo/attention/backends/sage_attn3.py @@ -116,18 +116,19 @@ def sage_attn_qk_torch(q, k, per_block_mean=False): p = torch.softmax(qk, dim=-1) return p.unsqueeze(0) -def attn_backward_16bit(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, p, is_causal=False): +def attn_backward_16bit(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_causal=False): assert is_causal == False q, k, v, o, do = q_BLHD.transpose(1 ,2).contiguous(), k_BLHD.transpose(1 ,2).contiguous(), v_BLHD.transpose(1 ,2).contiguous(), out_BLHD.transpose(1 ,2).contiguous(), grad_out_BLHD.transpose(1 ,2).contiguous() B, H, _, D = q.shape - # p = sage_attn_qk_torch(q, k, per_block_mean=False).to(do.dtype) + p = sage_attn_qk_torch(q, k, per_block_mean=False).to(do.dtype) dV = p.transpose(-2, -1) @ do # [B,H,L,D] dP = do @ v.transpose(-2, -1) # [B,H,L,L] # do_dot_o = (do.float() * o.float()).sum(dim=-1, keepdim=True).to(dP.dtype) # [B,H,L,1] - do_dot_o = (do * o).sum(dim=-1, keepdim=True) # [B,H,L,1] - dS = p * (dP - do_dot_o) # [B,H,L,L] + # do_dot_o = (do * o).sum(dim=-1, keepdim=True) # [B,H,L,1] + dp_dot_p = (dP * p).sum(dim=-1, keepdim=True) # [B,H,L,1] + dS = p * (dP - dp_dot_p) # [B,H,L,L] dQ = (dS @ k) / torch.sqrt(torch.tensor(D, device=q.device, dtype=q.dtype)) # [B,H,L,D] @@ -185,7 +186,7 @@ def flash_attn_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_caus ) return dq_BLHD, dk_BLHD, dv_BLHD, None, None -USE_TORCH_4BIT_FWD = True +USE_TORCH_4BIT_FWD = False USE_TORCH_16BIT_BWD = True class _SageAttnBlackwellWith16bitBwd(torch.autograd.Function): @staticmethod @@ -215,17 +216,17 @@ class _SageAttnBlackwellWith16bitBwd(torch.autograd.Function): # Back to BLHD; keep for FA bwd out_BLHD = out_BHLD.permute(0, 2, 1, 3).contiguous() - ctx.save_for_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD, p) + ctx.save_for_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD) return out_BLHD @staticmethod def backward(ctx, grad_out_BLHD): - q_BLHD, k_BLHD, v_BLHD, out_BLHD, p = ctx.saved_tensors + q_BLHD, k_BLHD, v_BLHD, out_BLHD = ctx.saved_tensors is_causal = ctx.is_causal if USE_TORCH_16BIT_BWD: - return attn_backward_16bit(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, p, is_causal) + return attn_backward_16bit(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_causal) else: return flash_attn_backward(q_BLHD, k_BLHD, v_BLHD, out_BLHD, grad_out_BLHD, is_causal) @@ -291,6 +292,6 @@ class SageAttention3Impl(AttentionImpl): value: torch.Tensor, attn_metadata: AttentionMetadata, ) -> torch.Tensor: - # output = sageattn_blackwell_with_16bit_bwd(query, key, value, is_causal=self.causal) - output = attn_forward_4bit_fwd_16bit_bwd(query, key, value) + output = sageattn_blackwell_with_16bit_bwd(query, key, value, is_causal=self.causal) + # output = attn_forward_4bit_fwd_16bit_bwd(query, key, value) return output diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index e0d7f3899..75140222e 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -928,91 +928,85 @@ class TrainingPipeline(LoRAPipeline, ABC): num_sp_groups = world_group.world_size // self.sp_group.world_size # Process each validation prompt for each validation step - for num_inference_steps in validation_steps: - logger.info("rank: %s: num_inference_steps: %s", - self.global_rank, - num_inference_steps, - local_main_process_only=False) - step_videos: list[np.ndarray] = [] - step_captions: list[str] = [] + # for num_inference_steps in validation_steps: + # logger.info("rank: %s: num_inference_steps: %s", + # self.global_rank, + # num_inference_steps, + # local_main_process_only=False) + # step_videos: list[np.ndarray] = [] + # step_captions: list[str] = [] - for validation_batch in validation_dataloader: - batch = self._prepare_validation_batch(sampling_param, - training_args, - validation_batch, - num_inference_steps) - logger.info("rank: %s: rank_in_sp_group: %s, batch.prompt: %s", - self.global_rank, - self.rank_in_sp_group, - batch.prompt, - local_main_process_only=False) + # for validation_batch in validation_dataloader: + # batch = self._prepare_validation_batch(sampling_param, + # training_args, + # validation_batch, + # num_inference_steps) + # logger.info("rank: %s: rank_in_sp_group: %s, batch.prompt: %s", + # self.global_rank, + # self.rank_in_sp_group, + # batch.prompt, + # local_main_process_only=False) - assert batch.prompt is not None and isinstance( - batch.prompt, str) - step_captions.append(batch.prompt) + # assert batch.prompt is not None and isinstance( + # batch.prompt, str) + # step_captions.append(batch.prompt) - # Run validation inference - output_batch = self.validation_pipeline.forward( - batch, training_args) - samples = output_batch.output + # # Run validation inference + # output_batch = self.validation_pipeline.forward( + # batch, training_args) + # samples = output_batch.output - if self.rank_in_sp_group != 0: - continue + # if self.rank_in_sp_group != 0: + # continue - # Process outputs - video = rearrange(samples, "b c t h w -> t b c h w") - frames = [] - for x in video: - x = torchvision.utils.make_grid(x, nrow=6) - x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) - frames.append((x * 255).numpy().astype(np.uint8)) - step_videos.append(frames) + # # Process outputs + # video = rearrange(samples, "b c t h w -> t b c h w") + # frames = [] + # for x in video: + # x = torchvision.utils.make_grid(x, nrow=6) + # x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) + # frames.append((x * 255).numpy().astype(np.uint8)) + # step_videos.append(frames) - # Only sp_group leaders (rank_in_sp_group == 0) need to send their - # results to global rank 0 - if self.rank_in_sp_group == 0: - if self.global_rank == 0: - # Global rank 0 collects results from all sp_group leaders - all_videos = step_videos # Start with own results - all_captions = step_captions + # # Only sp_group leaders (rank_in_sp_group == 0) need to send their + # # results to global rank 0 + # if self.rank_in_sp_group == 0: + # if self.global_rank == 0: + # # Global rank 0 collects results from all sp_group leaders + # all_videos = step_videos # Start with own results + # all_captions = step_captions - # Receive from other sp_group leaders - for sp_group_idx in range(1, num_sp_groups): - src_rank = sp_group_idx * self.sp_world_size # Global rank of other sp_group leaders - recv_videos = world_group.recv_object(src=src_rank) - recv_captions = world_group.recv_object(src=src_rank) - all_videos.extend(recv_videos) - all_captions.extend(recv_captions) + # # Receive from other sp_group leaders + # for sp_group_idx in range(1, num_sp_groups): + # src_rank = sp_group_idx * self.sp_world_size # Global rank of other sp_group leaders + # recv_videos = world_group.recv_object(src=src_rank) + # recv_captions = world_group.recv_object(src=src_rank) + # all_videos.extend(recv_videos) + # all_captions.extend(recv_captions) - video_filenames = [] - for i, (video, caption) in enumerate( - zip(all_videos, all_captions, strict=True)): - os.makedirs(training_args.output_dir, exist_ok=True) - filename = os.path.join( - training_args.output_dir, - f"validation_step_{global_step}_inference_steps_{num_inference_steps}_video_{i}.mp4" - ) - imageio.mimsave(filename, video, fps=sampling_param.fps) - video_filenames.append(filename) + # video_filenames = [] + # for i, (video, caption) in enumerate( + # zip(all_videos, all_captions, strict=True)): + # os.makedirs(training_args.output_dir, exist_ok=True) + # filename = os.path.join( + # training_args.output_dir, + # f"validation_step_{global_step}_inference_steps_{num_inference_steps}_video_{i}.mp4" + # ) + # imageio.mimsave(filename, video, fps=sampling_param.fps) + # video_filenames.append(filename) - artifacts = [] - for filename, caption in zip(video_filenames, - all_captions, - strict=True): - video_artifact = self.tracker.video(filename, - caption=caption) - if video_artifact is not None: - artifacts.append(video_artifact) - if artifacts: - logs = { - f"validation_videos_{num_inference_steps}_steps": - artifacts - } - self.tracker.log_artifacts(logs, global_step) - else: - # Other sp_group leaders send their results to global rank 0 - world_group.send_object(step_videos, dst=0) - world_group.send_object(step_captions, dst=0) + # logs = { + # f"validation_videos_{num_inference_steps}_steps": [ + # wandb.Video(filename, caption=caption) + # for filename, caption in zip( + # video_filenames, all_captions, strict=True) + # ] + # } + # wandb.log(logs, step=global_step) + # else: + # # Other sp_group leaders send their results to global rank 0 + # world_group.send_object(step_videos, dst=0) + # world_group.send_object(step_captions, dst=0) # Re-enable gradients for training training_args.inference_mode = False