(self, block_idx: int)
| 283 | self._wait_blocks_move(block_idx) |
| 284 | |
| 285 | def submit_move_blocks_forward(self, block_idx: int): |
| 286 | # check if blocks_to_swap is enabled |
| 287 | if self.blocks_to_swap is None or self.blocks_to_swap == 0: |
| 288 | return |
| 289 | |
| 290 | if self.reentrant_activation_checkpointing and torch.is_grad_enabled(): |
| 291 | # Second forward pass, don't do block swapping |
| 292 | return |
| 293 | |
| 294 | # if supports_backward and backward is enabled, we swap blocks more than blocks_to_swap in backward pass |
| 295 | if not self.forward_only and block_idx >= self.blocks_to_swap: |
| 296 | return |
| 297 | |
| 298 | block_idx_to_cpu = block_idx |
| 299 | block_idx_to_cuda = self.num_blocks - self.blocks_to_swap + block_idx |
| 300 | block_idx_to_cuda = block_idx_to_cuda % self.num_blocks # this works for forward-only offloading |
| 301 | self._submit_move_blocks(block_idx_to_cpu, block_idx_to_cuda) |
no test coverage detected