def run(self, pipeline_record: DistributedPipelineRecord) -> None: """Runs pipeline parallelism. It modifies the given batches in place.""" m = len(pipeline_record.batches) self.stream = current_stream(self.device) for chunk in range(m): with record_function("feed"): pipeline_record.wait_for(chunk) pipeline_record.fence(chunk) self.compute(pipeline_record, chunk) with use_stream(self.stream): pipeline_record.forward_results(chunk)
def run_pipeline(self, pipeline_record_rref: rpc.RRef) -> Optional[Tensor]: """Processes a min-batch on this partition. If this is the last partition (pipeline_record has no consumer), concatenates results of processing all chunks and returns the result as the output of the model on the whole mini-batch. """ pipeline_record = pipeline_record_rref.local_value() self.run(pipeline_record) if not pipeline_record.consumers: result = microbatch.gather(pipeline_record.batches) assert len(result) == 1 result = result[0] s0 = current_stream(result.device) if is_cuda(s0): # TODO. Investigate why this is needed and remove it if possible. as_cuda(s0).synchronize() return result return None
def test_copy_wait_cuda_cuda(cuda_sleep): prev_stream = current_stream(torch.device("cuda")) next_stream = new_stream(torch.device("cuda")) _test_copy_wait(prev_stream, next_stream, cuda_sleep)
def test_use_stream_cuda(self): stream = new_stream(torch.device("cuda")) with use_stream(stream): assert current_stream(torch.device("cuda")) == stream
def test_current_stream_cuda(self): stream = current_stream(torch.device("cuda")) assert isinstance(stream, torch.cuda.Stream) assert stream == torch.cuda.current_stream()
def test_current_stream_cpu(self): stream = current_stream(torch.device("cpu")) assert stream is CPUStream
def test_wait_stream_cuda_cuda(self, cuda_sleep): source = current_stream(torch.device("cuda")) target = new_stream(torch.device("cuda")) self._test_wait_stream(source, target, cuda_sleep)
def test_get_device_cuda(self): stream = current_stream(torch.device("cuda")) assert get_device(stream).type == "cuda"