diff --git a/simpletuner/helpers/models/kandinsky5_video/transformer_kandinsky5.py b/simpletuner/helpers/models/kandinsky5_video/transformer_kandinsky5.py index 82770db97..094fef576 100644 --- a/simpletuner/helpers/models/kandinsky5_video/transformer_kandinsky5.py +++ b/simpletuner/helpers/models/kandinsky5_video/transformer_kandinsky5.py @@ -39,7 +39,9 @@ validate_flowmap_deltatime_type, ) from simpletuner.helpers.musubi_block_swap import MusubiBlockSwapManager +from simpletuner.helpers.training.gradient_checkpointing_interval import should_checkpoint_block from simpletuner.helpers.training.grounding.gligen_layers import apply_grounding_fuser +from simpletuner.helpers.training.offloaded_gradient_checkpointer import activation_offload_context from simpletuner.helpers.training.qk_clip_logging import publish_attention_max_logits from simpletuner.helpers.training.tread import TREADRouter @@ -407,7 +409,15 @@ def __init__(self): if not hasattr(F, "scaled_dot_product_attention"): raise ImportError(f"{self.__class__.__name__} requires PyTorch 2.0. Please upgrade your pytorch version.") - def __call__(self, attn, hidden_states, encoder_hidden_states=None, rotary_emb=None, sparse_params=None): + def __call__( + self, + attn, + hidden_states, + encoder_hidden_states=None, + rotary_emb=None, + sparse_params=None, + offload_attention: bool = False, + ): def _describe_tensor(name: str, tensor: Optional[Tensor]) -> str: if tensor is None: return f"{name}=None" @@ -474,13 +484,14 @@ def apply_rotary(x, rope): getattr(attn, "to_query", None) and attn.to_query.weight, getattr(attn, "to_key", None) and attn.to_key.weight, ) - attn_output = dispatch_attention_fn( - query, - key, - value, - attn_mask=attn_mask, - backend=self._attention_backend, - ) + with activation_offload_context(offload_attention, label=f"{attn.__class__.__qualname__}:attention"): + attn_output = dispatch_attention_fn( + query, + key, + value, + attn_mask=attn_mask, + backend=self._attention_backend, + ) except Exception: logger.error( "dispatch_attention_fn failed (backend=%s): %s; %s; %s; hidden_states=%s; encoder_hidden_states=%s; attn_mask=%s; %s", @@ -529,6 +540,7 @@ def forward( encoder_hidden_states: Optional[torch.Tensor] = None, sparse_params: Optional[torch.Tensor] = None, rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + offload_attention: bool = False, **kwargs, ) -> torch.Tensor: attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys()) @@ -546,6 +558,7 @@ def forward( encoder_hidden_states=encoder_hidden_states, sparse_params=sparse_params, rotary_emb=rotary_emb, + offload_attention=offload_attention, **kwargs, ) @@ -607,7 +620,7 @@ def __init__(self, model_dim, time_dim, ff_dim, head_dim): self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) self.feed_forward = Kandinsky5FeedForward(model_dim, ff_dim) - def forward(self, x, time_embed, rope): + def forward(self, x, time_embed, rope, offload_attention: bool = False): self_attn_params, ff_params = torch.chunk( _broadcast_modulation_params(self.text_modulation(time_embed), x), 2, dim=-1 ) @@ -627,7 +640,7 @@ def forward(self, x, time_embed, rope): ] ) raise RuntimeError(debug_msg) from err - out = self.self_attention(out, rotary_emb=rope) + out = self.self_attention(out, rotary_emb=rope, offload_attention=offload_attention) x = (x.float() + gate.float() * out.float()).type_as(x) shift, scale, gate = torch.chunk(ff_params, 3, dim=-1) @@ -652,7 +665,7 @@ def __init__(self, model_dim, time_dim, ff_dim, head_dim): self.feed_forward_norm = nn.LayerNorm(model_dim, elementwise_affine=False) self.feed_forward = Kandinsky5FeedForward(model_dim, ff_dim) - def forward(self, visual_embed, text_embed, time_embed, rope, sparse_params): + def forward(self, visual_embed, text_embed, time_embed, rope, sparse_params, offload_attention: bool = False): self_attn_params, cross_attn_params, ff_params = torch.chunk( _broadcast_modulation_params(self.visual_modulation(time_embed), visual_embed), 3, dim=-1 ) @@ -661,14 +674,23 @@ def forward(self, visual_embed, text_embed, time_embed, rope, sparse_params): visual_out = (self.self_attention_norm(visual_embed.float()) * (scale.float() + 1.0) + shift.float()).type_as( visual_embed ) - visual_out = self.self_attention(visual_out, rotary_emb=rope, sparse_params=sparse_params) + visual_out = self.self_attention( + visual_out, + rotary_emb=rope, + sparse_params=sparse_params, + offload_attention=offload_attention, + ) visual_embed = (visual_embed.float() + gate.float() * visual_out.float()).type_as(visual_embed) shift, scale, gate = torch.chunk(cross_attn_params, 3, dim=-1) visual_out = (self.cross_attention_norm(visual_embed.float()) * (scale.float() + 1.0) + shift.float()).type_as( visual_embed ) - visual_out = self.cross_attention(visual_out, encoder_hidden_states=text_embed) + visual_out = self.cross_attention( + visual_out, + encoder_hidden_states=text_embed, + offload_attention=offload_attention, + ) visual_embed = (visual_embed.float() + gate.float() * visual_out.float()).type_as(visual_embed) shift, scale, gate = torch.chunk(ff_params, 3, dim=-1) @@ -701,6 +723,7 @@ class Kandinsky5Transformer3DModel( "Kandinsky5TransformerDecoderBlock", ] _supports_gradient_checkpointing = True + _supports_attention_activation_offload = True _cp_plan = { "": { "hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False), @@ -784,6 +807,9 @@ def __init__( self.out_layer = Kandinsky5OutLayer(model_dim, time_dim, out_visual_dim, patch_size) self.gradient_checkpointing = False self.gradient_checkpointing_backend = "torch" + self.gradient_checkpointing_offload_attention = False + self.gradient_checkpointing_interval = None + self.gradient_checkpointing_segment_stride = None self._musubi_block_swap = MusubiBlockSwapManager.build( depth=num_text_blocks + num_visual_blocks, blocks_to_swap=musubi_blocks_to_swap, @@ -798,6 +824,15 @@ def enable_flowmap_time_conditioning(self, gate_value: float = 0.25, deltatime_t def set_gradient_checkpointing_backend(self, backend: str): self.gradient_checkpointing_backend = backend + def set_gradient_checkpointing_offload_attention(self, enabled: bool): + self.gradient_checkpointing_offload_attention = bool(enabled) + + def set_gradient_checkpointing_interval(self, interval: int): + self.gradient_checkpointing_interval = interval + + def set_gradient_checkpointing_segment_stride(self, segment_stride: int | None): + self.gradient_checkpointing_segment_stride = segment_stride + def set_router(self, router: TREADRouter, routes: List[Dict[str, Any]]): """Attach a TREAD router and route definitions.""" self._tread_router = router @@ -891,20 +926,38 @@ def forward( f"Tokenwise timestep count {visual_time_embed.shape[1]} does not match visual token count {expected_tokens}." ) - for text_transformer_block in self.text_transformer_blocks: - if torch.is_grad_enabled() and self.gradient_checkpointing: - if self.gradient_checkpointing_backend == "unsloth": + for text_layer_idx, text_transformer_block in enumerate(self.text_transformer_blocks): + if torch.is_grad_enabled() and should_checkpoint_block( + text_layer_idx, + self.gradient_checkpointing, + self.gradient_checkpointing_interval, + self.gradient_checkpointing_segment_stride, + ): + if self.gradient_checkpointing_backend.startswith("unsloth"): from simpletuner.helpers.training.offloaded_gradient_checkpointer import offloaded_checkpoint checkpoint_fn = offloaded_checkpoint else: checkpoint_fn = torch.utils.checkpoint.checkpoint + def _checkpointed_text_block(text_embed, text_time_embed, text_rope, block=text_transformer_block): + return block( + text_embed, + text_time_embed, + text_rope, + offload_attention=self.gradient_checkpointing_offload_attention, + ) + text_embed = checkpoint_fn( - text_transformer_block, text_embed, text_time_embed, text_rope, use_reentrant=False + _checkpointed_text_block, text_embed, text_time_embed, text_rope, use_reentrant=False ) else: - text_embed = text_transformer_block(text_embed, text_time_embed, text_rope) + text_embed = text_transformer_block( + text_embed, + text_time_embed, + text_rope, + offload_attention=self.gradient_checkpointing_offload_attention, + ) visual_rope = self.visual_rope_embeddings(visual_shape, visual_rope_pos, scale_factor) to_fractal = sparse_params["to_fractal"] if sparse_params is not None else False @@ -982,16 +1035,38 @@ def _to_pos(idx: int) -> int: current_rope = self._route_rope(visual_rope, tread_mask_info, keep_len=visual_embed.size(1)) routing_now = True - if torch.is_grad_enabled() and self.gradient_checkpointing: - if self.gradient_checkpointing_backend == "unsloth": + if torch.is_grad_enabled() and should_checkpoint_block( + global_idx, + self.gradient_checkpointing, + self.gradient_checkpointing_interval, + self.gradient_checkpointing_segment_stride, + ): + if self.gradient_checkpointing_backend.startswith("unsloth"): from simpletuner.helpers.training.offloaded_gradient_checkpointer import offloaded_checkpoint checkpoint_fn = offloaded_checkpoint else: checkpoint_fn = torch.utils.checkpoint.checkpoint + def _checkpointed_visual_block( + visual_embed, + text_embed, + visual_time_embed, + current_rope, + sparse_params, + block=visual_transformer_block, + ): + return block( + visual_embed, + text_embed, + visual_time_embed, + current_rope, + sparse_params, + offload_attention=self.gradient_checkpointing_offload_attention, + ) + visual_embed = checkpoint_fn( - visual_transformer_block, + _checkpointed_visual_block, visual_embed, text_embed, visual_time_embed, @@ -1001,7 +1076,12 @@ def _to_pos(idx: int) -> int: ) else: visual_embed = visual_transformer_block( - visual_embed, text_embed, visual_time_embed, current_rope, sparse_params + visual_embed, + text_embed, + visual_time_embed, + current_rope, + sparse_params, + offload_attention=self.gradient_checkpointing_offload_attention, ) if grounding_objs is not None and hasattr(visual_transformer_block, "fuser"): diff --git a/simpletuner/helpers/training/default_settings/safety_check.py b/simpletuner/helpers/training/default_settings/safety_check.py index fc879e5a4..9eb1eb887 100644 --- a/simpletuner/helpers/training/default_settings/safety_check.py +++ b/simpletuner/helpers/training/default_settings/safety_check.py @@ -156,6 +156,8 @@ def safety_check(args, accelerator): "flux2", "hidream", "ideogram", + "kandinsky5_image", + "kandinsky5_video", ] gradient_checkpointing_segment_stride_supported_models = [ "ace_step", @@ -170,12 +172,18 @@ def safety_check(args, accelerator): "hidream", "hunyuanvideo", "ideogram", + "kandinsky5_image", + "kandinsky5_video", ] attention_activation_offload_supported_models = [ "chroma", "flux", "flux2", "hunyuanvideo", + "kandinsky5-image", + "kandinsky5-video", + "kandinsky5_image", + "kandinsky5_video", ] if getattr(args, "gradient_checkpointing_offload_attention", False): if args.model_family.lower() not in attention_activation_offload_supported_models: diff --git a/tests/test_segmented_checkpointing_model_support.py b/tests/test_segmented_checkpointing_model_support.py index bf5658ba9..2cdd785c8 100644 --- a/tests/test_segmented_checkpointing_model_support.py +++ b/tests/test_segmented_checkpointing_model_support.py @@ -356,3 +356,19 @@ def test_checkpointing_controls(self): ffn=False, attention_offload=False, ) + + +class Kandinsky5SegmentedCheckpointingSupportTests(unittest.TestCase): + def test_checkpointing_controls(self): + from simpletuner.helpers.models.kandinsky5_video.transformer_kandinsky5 import Kandinsky5Transformer3DModel + + assert_checkpointing_controls( + self, + Kandinsky5Transformer3DModel, + backend=True, + interval=True, + stride=True, + checkpoint_attention_offload=True, + ffn=False, + attention_offload=True, + )