Merge pull request #3151 from bghira/feature/musubi-h2d-hybrid

Stream checkpointed Musubi blocks through a packed H2D ring

Author: bghira

simpletuner/helpers/models/auraflow/transformer.py

@@ -806,8 +806,6 @@ def _to_pos(idx):
 
         # MMDiT blocks.
         for index_block, block in enumerate(self.joint_transformer_blocks):
-            if musubi_offload_active and musubi_manager.is_managed_block(global_idx):
-                musubi_manager.stream_in(block, hidden_states.device)
             # TREAD: START a route?
             if use_routing and route_ptr < len(routes) and global_idx == routes[route_ptr]["start_layer_idx"]:
                 mask_ratio = routes[route_ptr]["selection_ratio"]
@@ -827,7 +825,7 @@ def _to_pos(idx):
                     hidden_states = hidden_states + block_controlnet_hidden_states[index_block // interval_control]
                 continue
 
-            if (
+            checkpoint_this_block = (
                 self.training
                 and torch.is_grad_enabled()
                 and should_checkpoint_block(
@@ -836,7 +834,11 @@ def _to_pos(idx):
                     self.gradient_checkpointing_interval,
                     self.gradient_checkpointing_segment_stride,
                 )
-            ):
+            )
+            if musubi_offload_active and musubi_manager.is_managed_block(global_idx):
+                musubi_manager.stream_in(block, hidden_states.device, checkpointed=checkpoint_this_block)
+
+            if checkpoint_this_block:
 
                 if self.gradient_checkpointing_backend.startswith("unsloth"):
                     from simpletuner.helpers.training.offloaded_gradient_checkpointer import offloaded_checkpoint
@@ -916,9 +918,6 @@ def custom_forward(hidden_states, encoder_hidden_states, temb, context_temb, che
             for index_block, block in enumerate(self.single_transformer_blocks):
                 actual_index = len(self.joint_transformer_blocks) + index_block
 
-                if musubi_offload_active and musubi_manager.is_managed_block(global_idx):
-                    musubi_manager.stream_in(block, combined_hidden_states.device)
-
                 # TREAD: START a route?
                 if use_routing and route_ptr < len(routes) and global_idx == routes[route_ptr]["start_layer_idx"]:
                     mask_ratio = routes[route_ptr]["selection_ratio"]
@@ -959,7 +958,7 @@ def custom_forward(hidden_states, encoder_hidden_states, temb, context_temb, che
                         )
                     continue
 
-                if (
+                checkpoint_this_block = (
                     self.training
                     and torch.is_grad_enabled()
                     and should_checkpoint_block(
@@ -968,7 +967,15 @@ def custom_forward(hidden_states, encoder_hidden_states, temb, context_temb, che
                         self.gradient_checkpointing_interval,
                         self.gradient_checkpointing_segment_stride,
                     )
-                ):
+                )
+                if musubi_offload_active and musubi_manager.is_managed_block(global_idx):
+                    musubi_manager.stream_in(
+                        block,
+                        combined_hidden_states.device,
+                        checkpointed=checkpoint_this_block,
+                    )
+
+                if checkpoint_this_block:
 
                     if self.gradient_checkpointing_backend.startswith("unsloth"):
                         from simpletuner.helpers.training.offloaded_gradient_checkpointer import offloaded_checkpoint

simpletuner/helpers/models/chroma/transformer.py

@@ -964,8 +964,6 @@ def _to_pos(idx):
             grounding_objs = self.position_net(**grounding_kwargs)
 
         for index_block, block in enumerate(self.transformer_blocks):
-            if musubi_offload_active and musubi_manager.is_managed_block(global_idx):
-                musubi_manager.stream_in(block, hidden_states.device)
             actual_index = global_idx
             img_offset = 3 * len(self.single_transformer_blocks)
             txt_offset = img_offset + 6 * len(self.transformer_blocks)
@@ -1035,6 +1033,13 @@ def _to_pos(idx):
                 )
             )
 
+            if musubi_offload_active and musubi_manager.is_managed_block(global_idx):
+                musubi_manager.stream_in(
+                    block,
+                    hidden_states.device,
+                    checkpointed=use_checkpoint and not self.gradient_checkpointing_backend.endswith("-ffn"),
+                )
+
             if use_checkpoint:
                 if self.gradient_checkpointing_backend.startswith("unsloth"):
                     from simpletuner.helpers.training.offloaded_gradient_checkpointer import offloaded_checkpoint
@@ -1130,8 +1135,6 @@ def run_checkpointed_block(
         txt_len = encoder_hidden_states.shape[1]
 
         for index_block, block in enumerate(self.single_transformer_blocks):
-            if musubi_offload_active and musubi_manager.is_managed_block(global_idx):
-                musubi_manager.stream_in(block, hidden_states.device)
             actual_index = global_idx
             start_idx = 3 * index_block
             if pooled_temb.ndim == 4:
@@ -1237,6 +1240,13 @@ def run_checkpointed_block(
                 )
             )
 
+            if musubi_offload_active and musubi_manager.is_managed_block(global_idx):
+                musubi_manager.stream_in(
+                    block,
+                    hidden_states.device,
+                    checkpointed=use_checkpoint and not self.gradient_checkpointing_backend.endswith("-ffn"),
+                )
+
             if use_checkpoint:
                 if self.gradient_checkpointing_backend.startswith("unsloth"):
                     from simpletuner.helpers.training.offloaded_gradient_checkpointer import offloaded_checkpoint

simpletuner/helpers/models/cosmos/transformer.py

@@ -895,9 +895,7 @@ def forward(
                         )
                         hidden_states = router.start_route(hidden_states, mask_info)
                         break
-            if musubi_offload_active and musubi_manager.is_managed_block(bid):
-                musubi_manager.stream_in(block, hidden_states.device)
-            if (
+            checkpoint_this_block = (
                 grad_enabled
                 and self.gradient_checkpointing
                 and should_checkpoint_block(
@@ -906,7 +904,10 @@ def forward(
                     self.gradient_checkpointing_interval,
                     self.gradient_checkpointing_segment_stride,
                 )
-            ):
+            )
+            if musubi_offload_active and musubi_manager.is_managed_block(bid):
+                musubi_manager.stream_in(block, hidden_states.device, checkpointed=checkpoint_this_block)
+            if checkpoint_this_block:
                 hidden_states = self._gradient_checkpointing_func(
                     block,
                     hidden_states,

simpletuner/helpers/models/cosmos3/transformer.py

@@ -707,7 +707,7 @@ def _should_gradient_checkpoint_layer(self, layer_idx: int) -> bool:
     @staticmethod
     def _stream_in_for_checkpoint_recompute(musubi_manager, layer_idx: int, layer: nn.Module, device: torch.device) -> None:
         if musubi_manager is not None and musubi_manager.is_managed_block(layer_idx):
-            musubi_manager.stream_in(layer, device)
+            musubi_manager.stream_in(layer, device, checkpointed=True)
 
     # -------------------------------------------------------------------------
     # Pure-tensor packing/unpacking helpers (no layer state).
@@ -1057,9 +1057,14 @@ def forward(
             gen_seq = hidden_states
             rotary_gen = (cos[und_len:], sin[und_len:])
             for layer_idx, decoder_layer in enumerate(self.layers):
+                checkpoint_this_block = self._should_gradient_checkpoint_layer(layer_idx)
                 if musubi_offload_active and musubi_manager.is_managed_block(layer_idx):
-                    musubi_manager.stream_in(decoder_layer, gen_seq.device)
-                if self._should_gradient_checkpoint_layer(layer_idx):
+                    musubi_manager.stream_in(
+                        decoder_layer,
+                        gen_seq.device,
+                        checkpointed=checkpoint_this_block,
+                    )
+                if checkpoint_this_block:
 
                     def checkpointed_gen_only(
                         x,
@@ -1088,9 +1093,14 @@ def checkpointed_gen_only(
             rotary_emb = (cos[:und_len], sin[:und_len], cos[und_len:], sin[und_len:])
             vision_gen_indexes = vision_mse_loss_indexes - und_len
             for layer_idx, decoder_layer in enumerate(self.layers):
+                checkpoint_this_block = self._should_gradient_checkpoint_layer(layer_idx)
                 if musubi_offload_active and musubi_manager.is_managed_block(layer_idx):
-                    musubi_manager.stream_in(decoder_layer, gen_seq.device)
-                if self._should_gradient_checkpoint_layer(layer_idx):
+                    musubi_manager.stream_in(
+                        decoder_layer,
+                        gen_seq.device,
+                        checkpointed=checkpoint_this_block,
+                    )
+                if checkpoint_this_block:
 
                     def checkpointed_layer(
                         und,

simpletuner/helpers/models/ernie/transformer.py

@@ -260,8 +260,19 @@ def apply_layer(layer_idx, layer, batch_first_hidden_states, rotary_emb, attn_ma
             musubi_offload_active = musubi_manager.activate(combined_blocks, hidden_states.device, grad_enabled)
 
         for layer_idx, layer in enumerate(self.layers):
+            checkpoint_this_block = (
+                layer_idx not in skip_set
+                and grad_enabled
+                and self.gradient_checkpointing
+                and should_checkpoint_block(
+                    layer_idx,
+                    True,
+                    self.gradient_checkpointing_interval,
+                    getattr(self, "gradient_checkpointing_segment_stride", None),
+                )
+            )
             if musubi_offload_active and musubi_manager.is_managed_block(layer_idx):
-                musubi_manager.stream_in(layer, hidden_states.device)
+                musubi_manager.stream_in(layer, hidden_states.device, checkpointed=checkpoint_this_block)
             if use_routing and route_ptr < len(routes) and layer_idx == routes[route_ptr]["start_layer_idx"]:
                 keep_mask = torch.zeros(
                     (batch_size, hidden_states.shape[1]),

simpletuner/helpers/models/flux/transformer.py

@@ -1208,8 +1208,18 @@ def run_double_block(
                     capture_idx += len(segment_blocks)
                     continue
 
+            checkpoint_this_block = (
+                self.training
+                and self.gradient_checkpointing
+                and not run_gap_eagerly
+                and (self.gradient_checkpointing_interval is None or index_block % self.gradient_checkpointing_interval == 0)
+            )
             if musubi_offload_active and musubi_manager.is_managed_block(global_idx):
-                musubi_manager.stream_in(block, hidden_states.device)
+                musubi_manager.stream_in(
+                    block,
+                    hidden_states.device,
+                    checkpointed=checkpoint_this_block and not self.gradient_checkpointing_backend.endswith("-ffn"),
+                )
             # TREAD: START a route?
             if use_routing and route_ptr < len(routes) and global_idx == routes[route_ptr]["start_layer_idx"]:
                 mask_ratio = routes[route_ptr]["selection_ratio"]
@@ -1237,12 +1247,7 @@ def run_double_block(
 
                 # concatenate text + image rope
                 current_rope = tuple(torch.cat([tr, ir], dim=1) for tr, ir in zip(text_rope_b, img_rope_r))
-            if (
-                self.training
-                and self.gradient_checkpointing
-                and not run_gap_eagerly
-                and (self.gradient_checkpointing_interval is None or index_block % self.gradient_checkpointing_interval == 0)
-            ):
+            if checkpoint_this_block:
                 checkpoint_ffn = self.gradient_checkpointing_backend.endswith("-ffn")
 
                 def create_custom_forward(module):
@@ -1365,8 +1370,18 @@ def run_single_block(_relative_index, segment_block, segment_hidden_states):
                     capture_idx += len(segment_blocks)
                     continue
 
+            checkpoint_this_block = (
+                self.training
+                and self.gradient_checkpointing
+                and not run_gap_eagerly
+                and (self.gradient_checkpointing_interval is None or index_block % self.gradient_checkpointing_interval == 0)
+            )
             if musubi_offload_active and musubi_manager.is_managed_block(global_idx):
-                musubi_manager.stream_in(block, hidden_states.device)
+                musubi_manager.stream_in(
+                    block,
+                    hidden_states.device,
+                    checkpointed=checkpoint_this_block and not self.gradient_checkpointing_backend.endswith("-ffn"),
+                )
             # TREAD: START? (operate on *image* tokens only)
             if use_routing and route_ptr < len(routes) and global_idx == routes[route_ptr]["start_layer_idx"]:
                 mask_ratio = routes[route_ptr]["selection_ratio"]
@@ -1402,12 +1417,7 @@ def run_single_block(_relative_index, segment_block, segment_hidden_states):
                 )
                 current_rope = tuple(torch.cat([tr, ir], dim=1) for tr, ir in zip(text_rope_b, img_rope_r))
 
-            if (
-                self.training
-                and self.gradient_checkpointing
-                and not run_gap_eagerly
-                and (self.gradient_checkpointing_interval is None or index_block % self.gradient_checkpointing_interval == 0)
-            ):
+            if checkpoint_this_block:
                 checkpoint_ffn = self.gradient_checkpointing_backend.endswith("-ffn")
 
                 def create_custom_forward(module):

simpletuner/helpers/models/flux2/transformer.py

@@ -1315,9 +1315,14 @@ def run_double_block(_idx, block, encoder_hidden_states, hidden_states):
                 break
 
             global_layer_idx = index_block
+            checkpoint_this_block = grad_enabled and self.gradient_checkpointing
 
             if musubi_offload_active and musubi_manager.is_managed_block(global_layer_idx):
-                musubi_manager.stream_in(block, hidden_states.device)
+                musubi_manager.stream_in(
+                    block,
+                    hidden_states.device,
+                    checkpointed=checkpoint_this_block,
+                )
 
             # Check for TREAD routing
             if self._tread_router is not None and self.training:
@@ -1356,7 +1361,7 @@ def run_double_block(_idx, block, encoder_hidden_states, hidden_states):
                 torch.cat([current_pe_txt[1], current_pe_img[1]], dim=0),
             )
 
-            if torch.is_grad_enabled() and self.gradient_checkpointing:
+            if checkpoint_this_block:
 
                 def create_custom_forward(module):
                     def custom_forward(*inputs):
@@ -1442,9 +1447,14 @@ def run_single_block(_idx, block, hidden_states):
                 break
 
             global_layer_idx = num_double + index_block
+            checkpoint_this_block = grad_enabled and self.gradient_checkpointing
 
             if musubi_offload_active and musubi_manager.is_managed_block(global_layer_idx):
-                musubi_manager.stream_in(block, hidden_states.device)
+                musubi_manager.stream_in(
+                    block,
+                    hidden_states.device,
+                    checkpointed=checkpoint_this_block,
+                )
 
             # Check for TREAD routing
             if self._tread_router is not None and self.training:
@@ -1490,7 +1500,7 @@ def run_single_block(_idx, block, hidden_states):
                     tread_routing_info = None
                     current_concat_pe = concat_rotary_emb
 
-            if torch.is_grad_enabled() and self.gradient_checkpointing:
+            if checkpoint_this_block:
 
                 def create_custom_forward(module):
                     def custom_forward(*inputs):

simpletuner/helpers/models/hidream/transformer.py

@@ -1666,8 +1666,18 @@ def custom_forward(pe_embedder):
         # Process through double stream blocks
         capture_idx = 0
         for bid, block in enumerate(self.double_stream_blocks):
+            checkpoint_this_block = self.training and should_checkpoint_block(
+                bid,
+                self.gradient_checkpointing,
+                self.gradient_checkpointing_interval,
+                self.gradient_checkpointing_segment_stride,
+            )
             if musubi_offload_active and musubi_manager.is_managed_block(block_id):
-                musubi_manager.stream_in(block, hidden_states.device)
+                musubi_manager.stream_in(
+                    block,
+                    hidden_states.device,
+                    checkpointed=checkpoint_this_block,
+                )
             # TREAD routing for this layer
             if use_routing:
                 # Check if this layer should use routing
@@ -1698,12 +1708,7 @@ def custom_forward(pe_embedder):
             )
 
             # Process through the block with optional gradient checkpointing
-            if self.training and should_checkpoint_block(
-                bid,
-                self.gradient_checkpointing,
-                self.gradient_checkpointing_interval,
-                self.gradient_checkpointing_segment_stride,
-            ):
+            if checkpoint_this_block:
 
                 def create_custom_forward(module, return_dict=None):
                     def custom_forward(*inputs):
@@ -1797,8 +1802,18 @@ def custom_forward(*inputs):
 
         # 7. Process through single stream blocks
         for bid, block in enumerate(self.single_stream_blocks):
+            checkpoint_this_block = self.training and should_checkpoint_block(
+                len(self.double_stream_blocks) + bid,
+                self.gradient_checkpointing,
+                self.gradient_checkpointing_interval,
+                self.gradient_checkpointing_segment_stride,
+            )
             if musubi_offload_active and musubi_manager.is_managed_block(block_id):
-                musubi_manager.stream_in(block, hidden_states.device)
+                musubi_manager.stream_in(
+                    block,
+                    hidden_states.device,
+                    checkpointed=checkpoint_this_block,
+                )
             # TREAD routing for single stream layers
             if use_routing:
                 # Check if this layer should use routing
@@ -1835,12 +1850,7 @@ def custom_forward(*inputs):
             hidden_states = torch.cat([hidden_states, cur_llama_embedding], dim=1)
 
             # Process through the block with optional gradient checkpointing
-            if self.training and should_checkpoint_block(
-                len(self.double_stream_blocks) + bid,
-                self.gradient_checkpointing,
-                self.gradient_checkpointing_interval,
-                self.gradient_checkpointing_segment_stride,
-            ):
+            if checkpoint_this_block:
 
                 def create_custom_forward(module, return_dict=None):
                     def custom_forward(*inputs):