(
self,
block_type: str,
blocks: list[nn.Module],
num_blocks: int,
blocks_to_swap: int,
supports_backward: bool,
device: torch.device,
reentrant_activation_checkpointing: bool,
debug: bool = False,
)
| 187 | """ |
| 188 | |
| 189 | def __init__( |
| 190 | self, |
| 191 | block_type: str, |
| 192 | blocks: list[nn.Module], |
| 193 | num_blocks: int, |
| 194 | blocks_to_swap: int, |
| 195 | supports_backward: bool, |
| 196 | device: torch.device, |
| 197 | reentrant_activation_checkpointing: bool, |
| 198 | debug: bool = False, |
| 199 | ): |
| 200 | super().__init__(block_type, blocks, num_blocks, blocks_to_swap, device, debug) |
| 201 | |
| 202 | self.supports_backward = supports_backward |
| 203 | self.forward_only = not supports_backward # forward only offloading: can be changed to True for inference |
| 204 | self.reentrant_activation_checkpointing = reentrant_activation_checkpointing |
| 205 | |
| 206 | if self.supports_backward: |
| 207 | # register backward hooks |
| 208 | self.remove_handles = [] |
| 209 | for i, block in enumerate(blocks): |
| 210 | hook = self.create_backward_hook(i) |
| 211 | if hook is not None: |
| 212 | handle = block.register_full_backward_hook(hook) |
| 213 | self.remove_handles.append(handle) |
| 214 | |
| 215 | def disable_block_swap(self): |
| 216 | self.blocks_to_swap_tmp = self.blocks_to_swap |
nothing calls this directly
no test coverage detected