MCPcopy Create free account
hub / github.com/tdrussell/diffusion-pipe / move_blocks

Method move_blocks

utils/offloading.py:147–158  ·  view source on GitHub ↗
(bidx_to_cpu, block_to_cpu, bidx_to_cuda, block_to_cuda)

Source from the content-addressed store, hash-verified

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]

Callers

nothing calls this directly

Calls 1

swap_weight_devicesMethod · 0.95

Tested by

no test coverage detected