Fix Process Group for tensors shared across processes

[mrshenli]

Ops on a Process Group (pg) instance will hit an error when input/output tensors are created on a different process, because, pg calls recordStream on CUDACachingAllocator which only knows tensors created within the same process.

The proposed solution is to add a suppressError arg (suggestions for better names?) to recordStream. See comments in code for arguments.

CC @pichuang1984

[facebook-github-bot]

@mrshenli has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[facebook-github-bot]

@mrshenli has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[facebook-github-bot]

@mrshenli has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[facebook-github-bot]

@mrshenli has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[facebook-github-bot]

@mrshenli has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[ezyang]

Out of curiosity, why didn't we test for this situation using cast_context on the DataPtr? It seems like a safer test to do.

[ezyang]

You seem to have modified nearly every call site to recordStream; I saw maybe only two missing sites:

torch/csrc/distributed/c10d/ddp.cpp
165:      c10::cuda::CUDACachingAllocator::recordStream(
209:    c10::cuda::CUDACachingAllocator::recordStream(

torch/csrc/autograd/generated/python_variable_methods.cpp
390:  c10::cuda::CUDACachingAllocator::recordStream(data, at::cuda::CUDAStream::unpack(((THCPStream*)arg)->cdata));

and I suspect that these should not error if the pointer is not found. Maybe we should do away with the option / flip the default?

[ezyang]

Wouldn't world_size = 2 be a much more idiomatic way to write this?

[mrshenli]

Yes, let me make the change.

[ezyang]

Generally speaking, it's better not to leak temporary files (which is what delete=False does, unless you explicitly delete it later). What is the reasoning for using delete=false?

[mrshenli]

I might be wrong, but if I don't add delete=False, it complains it cannot find the tmp file on delete. I was thinking maybe with context and tempfile both try to delete the file when exiting the context, as a result, one of them hits the error. But let me double check if that is the case, and will ad a comment if yes.

[ezyang]

I would be extremely surprised if that were the case. Here's a simple test:

macbook-pro-116:~ ezyang$ cat test.py
import tempfile

with tempfile.NamedTemporaryFile() as f:
    pass
macbook-pro-116:~ ezyang$ python test.py
macbook-pro-116:~ ezyang$ 

[mrshenli]

@ezyang thanks a lot for reviewing!

Out of curiosity, why didn't we test for this situation using cast_context on the DataPtr? It seems like a safer test to do.

I agree this will test the changes to CUDACachingAllocator, but this is not what users are experiencing. I could add that test, but still want to keep the multiprocessing test, as the goal of the test is to make sure users don't hit error anymore when they use PG to work on tensors shared across processes. If we don't have the multiprocessing test, we would not know whether we need additional changes other than modifying CUDACachingAllocator.

You seem to have modified nearly every call site to recordStream; I saw maybe only two missing sites:

1. For the recordStream in DDP, I don't see how users would get a gradient tensor from a different process. If they do, I would rather expose the error instead of silently suppress it, because it could mean we need a bigger change in DDP for certain cases.

2. For the python API, I have a similar concern as above. The recordStream is used in Scatter function, where the input should not be a shared tensor.

[ezyang]

This is a bit suboptimal, because the way this will look to the test runner is that the multiprocessing spawned subprocess died unceremoniously, and you have to look at the tea leaves to see, "Ah, it failed due to an assert." The way that test_multiprocessing goes about arranging this, is to have the parent process (aka the test runner) always responsible for doing the actual asserts (where you can do a normal self.assertEquals), and just have the child process pass back the tensors for the parent process to do checking on (or, if the child process must do the test, passing back a boolean saying if the result worked or not.)

What does the test suite output look like when this assert fails?

[ezyang]

Hmm, I guess because you are using multiprocessing.spawn, the exceptions will get propagated backwards

[mrshenli]

I agree. I don't like reinventing the wheel either. Will make the change.

What does the test suite output look like when this assert fails?

I have only tried error cases in local pytest runs, where the error messages are clear. Let me add a deliberate error to see what CI test shows.

[ezyang]

FWIW, test_multiprocessing.py does this by just having the test runner methods as honest to goodness top-level functions. This is just an informational comment, since a class method is just as good.

[ezyang]

Please think about if we should get rid of suppressError argument and just suppress errors always. All the other comments are just nits.

[ezyang]

I agree this will test the changes to CUDACachingAllocator, but this is not what users are experiencing. I could add that test, but still want to keep the multiprocessing test, as the goal of the test is to make sure users don't hit error anymore when they use PG to work on tensors shared across processes. If we don't have the multiprocessing test, we would not know whether we need additional changes other than modifying CUDACachingAllocator.

Sorry, miscommunication here. I'm referring to how, in the logic of a PR, you determine if a pointer was produced by the CUDA caching allocator by attempting to look it up in the block map; if it's not in the block map, it's not a CUDA caching allocator pointer. My point is that another way to test if it's produced by the CUDA caching allocator is by checking if the deleter for the DataPtr in question is for the CUDA caching allocator (you can perform this test using cast_context). I'm not referring to the test suite.

For the recordStream in DDP, I don't see how users would get a gradient tensor from a different process. If they do, I would rather expose the error instead of silently suppress it, because it could mean we need a bigger change in DDP for certain cases.

Remember, "different process" isn't the root cause of the problem; if any cuda tensor we allocate, comes from an allocation source that is not the caching allocator, recordStream will fail. So I could have very well defined a backward function that uses some external system (e.g., Numba) to perform the computation, and share back some memory that was allocated by them. This seems perfectly legitimate to me.

Morally, CUDA caching allocator's record stream annotations are to solve a problem that exists specifically because of how the caching allocator is implemented. Arguably, it is a generic interface for all CUDA allocators, but it just so happens that other allocators don't need to do anything.

[mrshenli]

@ezyang if I feel timid about a PR, can I force test suite to run all 73 tests?

[facebook-github-bot]

@mrshenli has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[facebook-github-bot]

@mrshenli has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[facebook-github-bot]

@mrshenli merged this pull request in 25d1496.