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"
