from pipeline.streaming_training import StreamingTrainingPipeline from typing import List, Optional, Tuple import torch import torch.distributed as dist class StreamingSwitchTrainingPipeline(StreamingTrainingPipeline): def __init__(self, *args, **kwargs): _apr_enabled = kwargs.pop('apr_enabled', False) _apr_alpha_max = kwargs.pop('apr_alpha_max', 0.8) _apr_d_window = kwargs.pop('apr_d_window', None) _apr_blend_sink = kwargs.pop('apr_blend_sink', False) _global_sink = kwargs.pop('global_sink', False) super().__init__(*args, **kwargs) self.global_sink = _global_sink self.apr_enabled = _apr_enabled self.apr_alpha_max = float(max(0.0, min(0.8, _apr_alpha_max))) self.apr_d_window = _apr_d_window self.apr_blend_sink = _apr_blend_sink if not dist.is_initialized() or dist.get_rank() == 0: print(f'[StreamingSwitchTrainingPipeline] global_sink={self.global_sink} (config-driven, Opt 14 kwargs.pop fix active)') def generate_chunk_with_cache(self, noise: torch.Tensor, conditional_dict: dict, *, current_start_frame: int=0, requires_grad: bool=True, switch_frame_index: Optional[int]=None, switch_conditional_dict: Optional[dict]=None, switch_recache_frames: Optional[torch.Tensor]=None, return_sim_step: bool=False) -> Tuple[torch.Tensor, Optional[int], Optional[int]]: if switch_conditional_dict is None or switch_frame_index is None: return super().generate_chunk_with_cache(noise=noise, conditional_dict=conditional_dict, current_start_frame=current_start_frame, requires_grad=requires_grad, return_sim_step=return_sim_step) batch_size, chunk_frames, num_channels, height, width = noise.shape assert chunk_frames % self.num_frame_per_block == 0 num_blocks = chunk_frames // self.num_frame_per_block all_num_frames = [self.num_frame_per_block] * num_blocks output = torch.zeros_like(noise) num_denoising_steps = len(self.denoising_step_list) exit_flags = self.generate_and_sync_list(len(all_num_frames), num_denoising_steps, device=noise.device) if not requires_grad: start_gradient_frame_index = chunk_frames else: start_gradient_frame_index = switch_frame_index local_start_frame = 0 self.generator.model.local_attn_size = int(self.local_attn_size) self._set_all_modules_max_attention_size(int(self.local_attn_size)) using_second = False cond_in_use = conditional_dict for block_index, current_num_frames in enumerate(all_num_frames): if not using_second and local_start_frame >= switch_frame_index: self._recache_after_switch(output[:, :local_start_frame, ...], current_start_frame + local_start_frame, switch_conditional_dict, local_start_frame, switch_recache_frames) cond_in_use = switch_conditional_dict using_second = True noisy_input = noise[:, local_start_frame:local_start_frame + current_num_frames] for step_idx, current_timestep in enumerate(self.denoising_step_list): exit_flag = step_idx == exit_flags[0] if self.same_step_across_blocks else step_idx == exit_flags[block_index] timestep = torch.ones([batch_size, current_num_frames], device=noise.device, dtype=torch.int64) * current_timestep if not exit_flag: with torch.no_grad(): _, denoised_pred = self.generator(noisy_image_or_video=noisy_input, conditional_dict=cond_in_use, timestep=timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=(current_start_frame + local_start_frame) * self.frame_seq_length) if step_idx < len(self.denoising_step_list) - 1: next_timestep = self.denoising_step_list[step_idx + 1] noisy_input = self.scheduler.add_noise(denoised_pred.flatten(0, 1), torch.randn_like(denoised_pred.flatten(0, 1)), next_timestep * torch.ones([batch_size * current_num_frames], device=noise.device, dtype=torch.long)).unflatten(0, denoised_pred.shape[:2]) else: enable_grad = local_start_frame >= start_gradient_frame_index context_manager = torch.enable_grad() if enable_grad else torch.no_grad() with context_manager: _, denoised_pred = self.generator(noisy_image_or_video=noisy_input, conditional_dict=cond_in_use, timestep=timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=(current_start_frame + local_start_frame) * self.frame_seq_length) break output[:, local_start_frame:local_start_frame + current_num_frames] = denoised_pred context_timestep = torch.ones_like(timestep) * self.context_noise context_noisy = self.scheduler.add_noise(denoised_pred.flatten(0, 1), torch.randn_like(denoised_pred.flatten(0, 1)), context_timestep.flatten(0, 1)).unflatten(0, denoised_pred.shape[:2]) with torch.no_grad(): self.generator(noisy_image_or_video=context_noisy, conditional_dict=cond_in_use, timestep=context_timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=(current_start_frame + local_start_frame) * self.frame_seq_length) local_start_frame += current_num_frames if not self.same_step_across_blocks: denoised_timestep_from, denoised_timestep_to = (None, None) elif exit_flags[0] == len(self.denoising_step_list) - 1: denoised_timestep_to = 0 denoised_timestep_from = 1000 - torch.argmin((self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0]].cuda()).abs(), dim=0).item() else: denoised_timestep_to = 1000 - torch.argmin((self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0] + 1].cuda()).abs(), dim=0).item() denoised_timestep_from = 1000 - torch.argmin((self.scheduler.timesteps.cuda() - self.denoising_step_list[exit_flags[0]].cuda()).abs(), dim=0).item() if return_sim_step: return (output, denoised_timestep_from, denoised_timestep_to, exit_flags[0] + 1) return (output, denoised_timestep_from, denoised_timestep_to) def _recache_after_switch(self, output, current_start_frame, new_conditional_dict, local_start_frame=None, switch_recache_frames=None): sink_tok_size = None sink_backup = None sink_norm_before = None if self.global_sink and current_start_frame > 0: gen = self.generator from torch.distributed.fsdp import FullyShardedDataParallel as _FSDP if hasattr(gen, '_fsdp_wrapped_module'): w = gen._fsdp_wrapped_module m = w.model _inner_tmp = m._fsdp_wrapped_module if isinstance(m, _FSDP) else m if hasattr(_inner_tmp, 'base_model') and hasattr(_inner_tmp.base_model, 'model'): _inner_tmp = _inner_tmp.base_model.model elif hasattr(gen, 'model'): _inner_tmp = gen.model else: _inner_tmp = None if _inner_tmp is not None: sink_tok_size = _inner_tmp.blocks[0].self_attn.sink_size * self.frame_seq_length sink_backup = [{'k': self.kv_cache1[i]['k'][:, :sink_tok_size].clone(), 'v': self.kv_cache1[i]['v'][:, :sink_tok_size].clone()} for i in range(self.num_transformer_blocks)] sink_norm_before = torch.norm(self.kv_cache1[0]['k'][:, :sink_tok_size]).float().item() if not dist.is_initialized() or dist.get_rank() == 0: print(f'[Recache-Train] switch@frame={current_start_frame}, global_sink=True, sink_tok={sink_tok_size}, sink_norm(block0) BEFORE = {sink_norm_before:.4f}') elif not self.global_sink: if not dist.is_initialized() or dist.get_rank() == 0: print(f'[Recache-Train] switch@frame={current_start_frame}, global_sink=False (sink will be recached with new prompt)') apr_old_cache = None if self.apr_enabled and current_start_frame > 0: apr_old_cache = [{'k': self.kv_cache1[i]['k'].clone(), 'v': self.kv_cache1[i]['v'].clone()} for i in range(self.num_transformer_blocks)] if not self.global_sink: for block_idx in range(self.num_transformer_blocks): cache = self.kv_cache1[block_idx] cache['k'].zero_() cache['v'].zero_() for blk in self.crossattn_cache: blk['k'].zero_() blk['v'].zero_() blk['is_init'] = False gen = self.generator from torch.distributed.fsdp import FullyShardedDataParallel as _FSDP if hasattr(gen, '_fsdp_wrapped_module'): w = gen._fsdp_wrapped_module m = w.model inner = m._fsdp_wrapped_module if isinstance(m, _FSDP) else m if hasattr(inner, 'base_model') and hasattr(inner.base_model, 'model'): inner = inner.base_model.model elif hasattr(gen, 'model'): inner = gen.model else: inner = None if current_start_frame == 0: return if switch_recache_frames is not None: frames_to_recache = torch.cat([switch_recache_frames, output], dim=1)[:, -21:, ...] num_recache_frames = frames_to_recache.shape[1] elif local_start_frame is not None: num_recache_frames = min(local_start_frame, 21) frames_to_recache = output[:, -num_recache_frames:] else: num_recache_frames = min(current_start_frame, 21) frames_to_recache = output[:, -num_recache_frames:] batch_size, num_recache_frames, c, h, w = frames_to_recache.shape device = frames_to_recache.device block_mask = self.generator.model._prepare_blockwise_causal_attn_mask(device=device, num_frames=num_recache_frames, frame_seqlen=self.frame_seq_length, num_frame_per_block=self.num_frame_per_block, local_attn_size=21) context_timestep = torch.ones([batch_size, num_recache_frames], device=device, dtype=torch.int64) * self.context_noise self.generator.model.block_mask = block_mask _has_memory = inner is not None and (getattr(inner, 'query_memory_encoder', None) is not None or getattr(inner, 'sink_memory', None) is not None) if _has_memory: frame_seqlen = self.frame_seq_length sink_frames = inner.blocks[0].self_attn.sink_size max_attn = inner.blocks[0].self_attn.max_attention_size recent_window_frames = (max_attn - sink_frames * frame_seqlen) // frame_seqlen inner._ei_prev_window_start = max(sink_frames, current_start_frame - recent_window_frames) with torch.no_grad(): self.generator(noisy_image_or_video=frames_to_recache, conditional_dict=new_conditional_dict, timestep=context_timestep, kv_cache=self.kv_cache1, crossattn_cache=self.crossattn_cache, current_start=(current_start_frame - num_recache_frames) * self.frame_seq_length) if sink_backup is not None: for i in range(self.num_transformer_blocks): self.kv_cache1[i]['k'][:, :sink_tok_size] = sink_backup[i]['k'] self.kv_cache1[i]['v'][:, :sink_tok_size] = sink_backup[i]['v'] sink_norm_after = torch.norm(self.kv_cache1[0]['k'][:, :sink_tok_size]).float().item() if not dist.is_initialized() or dist.get_rank() == 0: delta = abs(sink_norm_after - sink_norm_before) status = 'PRESERVED ✓' if delta < 0.001 else f'MISMATCH ✗ (delta={delta:.4e})' print(f'[Recache-Train] global_sink=True: sink_norm AFTER = {sink_norm_after:.4f} [{status}]') try: import model.streaming_training as _st_mod setattr(_st_mod, '_last_recache_sink_delta', delta) except Exception: pass del sink_backup for blk in self.crossattn_cache: blk['k'].zero_() blk['v'].zero_() blk['is_init'] = False if self.apr_enabled and apr_old_cache is not None: self._apply_apr_blend(old_cache=apr_old_cache, recache_start_frame=current_start_frame - num_recache_frames, current_start_frame=current_start_frame, num_recache_frames=num_recache_frames, inner=inner) del apr_old_cache if inner is not None and getattr(inner, 'query_memory_encoder', None) is not None: enc = inner.query_memory_encoder if getattr(enc, 'memory_recache', False): enc.reset(batch_size=frames_to_recache.shape[0], device=frames_to_recache.device, dtype=torch.bfloat16) last_blk = self.num_transformer_blocks - 1 cache = self.kv_cache1[last_blk] frame_seqlen = self.frame_seq_length sink_tok = inner.blocks[0].self_attn.sink_size * frame_seqlen local_end = cache['local_end_index'].item() if local_end > sink_tok: recache_k = cache['k'][:, sink_tok:local_end].clone() recache_v = cache['v'][:, sink_tok:local_end].clone() sink_k = cache['k'][:, :sink_tok].clone() if sink_tok > 0 else None sink_v = cache['v'][:, :sink_tok].clone() if sink_tok > 0 else None enc.update(recache_k, recache_v, sink_k, sink_v) if inner is not None and getattr(inner, 'sink_memory', None) is not None: sm = inner.sink_memory sm.reset() sink_frames_count = inner.blocks[0].self_attn.sink_size sink_output = output[:, :sink_frames_count] if output.shape[1] >= sink_frames_count else None if sink_output is None: if switch_recache_frames is not None and switch_recache_frames.shape[1] >= sink_frames_count: sink_output = switch_recache_frames[:, :sink_frames_count] if sink_output is not None: device = sink_output.device batch_size_local = sink_output.shape[0] sink_ts = torch.ones([batch_size_local, sink_frames_count], device=device, dtype=torch.int64) * self.context_noise n_heads = inner.blocks[0].self_attn.num_heads head_dim = inner.blocks[0].self_attn.head_dim sink_cache_size = sink_frames_count * self.frame_seq_length temp_kv = [{'k': torch.zeros([batch_size_local, sink_cache_size, n_heads, head_dim], device=device, dtype=torch.bfloat16), 'v': torch.zeros([batch_size_local, sink_cache_size, n_heads, head_dim], device=device, dtype=torch.bfloat16), 'global_end_index': torch.zeros(1, dtype=torch.long, device=device), 'local_end_index': torch.zeros(1, dtype=torch.long, device=device)} for _ in range(self.num_transformer_blocks)] temp_crossattn = [{'k': torch.zeros_like(self.crossattn_cache[0]['k']), 'v': torch.zeros_like(self.crossattn_cache[0]['v']), 'is_init': False} for _ in range(self.num_transformer_blocks)] with torch.no_grad(): self.generator(noisy_image_or_video=sink_output, conditional_dict=new_conditional_dict, timestep=sink_ts, kv_cache=temp_kv, crossattn_cache=temp_crossattn, current_start=0) frame_seqlen_local = self.frame_seq_length sink_frames_local = inner.blocks[0].self_attn.sink_size max_attn_local = inner.blocks[0].self_attn.max_attention_size recent_win = (max_attn_local - sink_frames_local * frame_seqlen_local) // frame_seqlen_local inner._ei_prev_window_start = max(sink_frames_local, current_start_frame - recent_win) if not dist.is_initialized() or dist.get_rank() == 0: print(f'[SinkMem] reset + re-captured on prompt switch at frame {current_start_frame}') def _apply_apr_blend(self, *, old_cache, recache_start_frame, current_start_frame, num_recache_frames, inner): frame_seqlen = self.frame_seq_length sink_size = inner.blocks[0].self_attn.sink_size sink_tok = sink_size * frame_seqlen D = self.apr_d_window if self.apr_d_window is not None else num_recache_frames D = max(1, int(D)) alpha_max = self.apr_alpha_max n_frames = num_recache_frames with torch.no_grad(): for blk_idx in range(self.num_transformer_blocks): new_k = self.kv_cache1[blk_idx]['k'] new_v = self.kv_cache1[blk_idx]['v'] old_k = old_cache[blk_idx]['k'] old_v = old_cache[blk_idx]['v'] if self.apr_blend_sink and (not self.global_sink): for f in range(sink_size): d_t = current_start_frame - f alpha = min(alpha_max, max(0.0, 1.0 - d_t / D)) s0, s1 = (f * frame_seqlen, (f + 1) * frame_seqlen) new_k[:, s0:s1] = (1 - alpha) * old_k[:, s0:s1] + alpha * new_k[:, s0:s1] new_v[:, s0:s1] = (1 - alpha) * old_v[:, s0:s1] + alpha * new_v[:, s0:s1] base = sink_tok for f in range(n_frames): absolute_frame = recache_start_frame + f d_t = current_start_frame - absolute_frame alpha = min(alpha_max, max(0.0, 1.0 - d_t / D)) s0, s1 = (base + f * frame_seqlen, base + (f + 1) * frame_seqlen) new_k[:, s0:s1] = (1 - alpha) * old_k[:, s0:s1] + alpha * new_k[:, s0:s1] new_v[:, s0:s1] = (1 - alpha) * old_v[:, s0:s1] + alpha * new_v[:, s0:s1] if not dist.is_initialized() or dist.get_rank() == 0: print(f'[APR] switch@{current_start_frame}: blended {n_frames} frames (α_max={alpha_max}, D_window={D})')