update
This commit is contained in:
Peiyuan Zhang
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user