(self, block_idx_to_cpu, block_idx_to_cuda)
| 144 | swap_weight_devices_no_cuda(self.device, block_to_cpu, block_to_cuda) |
| 145 | |
| 146 | def _submit_move_blocks(self, block_idx_to_cpu, block_idx_to_cuda): |
| 147 | def move_blocks(bidx_to_cpu, block_to_cpu, bidx_to_cuda, block_to_cuda): |
| 148 | if self.debug: |
| 149 | start_time = time.perf_counter() |
| 150 | print( |
| 151 | f"[{self.block_type}] Move block {bidx_to_cpu} to CPU and block {bidx_to_cuda} to {'CUDA' if self.cuda_available else 'device'}" |
| 152 | ) |
| 153 | |
| 154 | self.swap_weight_devices(block_to_cpu, block_to_cuda) |
| 155 | |
| 156 | if self.debug: |
| 157 | print(f"[{self.block_type}] Moved blocks {bidx_to_cpu} and {bidx_to_cuda} in {time.perf_counter()-start_time:.2f}s") |
| 158 | return bidx_to_cpu, bidx_to_cuda # , event |
| 159 | |
| 160 | block_to_cpu = self.blocks[block_idx_to_cpu] |
| 161 | block_to_cuda = self.blocks[block_idx_to_cuda] |
| 162 | |
| 163 | self.futures[block_idx_to_cuda] = self.thread_pool.submit( |
| 164 | move_blocks, block_idx_to_cpu, block_to_cpu, block_idx_to_cuda, block_to_cuda |
| 165 | ) |
| 166 | |
| 167 | def _wait_blocks_move(self, block_idx): |
| 168 | if block_idx not in self.futures: |
no outgoing calls
no test coverage detected