Add Distributed Model Parallel Tests for models_test.py #363
Conversation
swong3-sc
requested review from
kmontemayor2-sc,
mkolodner-sc,
nshah-sc,
svij-sc,
xgao4-sc,
yliu2-sc and
zfan3-sc
as code owners
| """ | ||
| Test that DMP-wrapped LightGCN produces the same output as non-wrapped model. Note: We only test with a single process for unit test. | ||
| """ | ||
| from torchrec.distributed.model_parallel import DistributedModelParallel as DMP |
There was a problem hiding this comment.
nit. import at the top of the file?
| if not dist.is_initialized(): | ||
| dist.init_process_group( | ||
| backend="gloo", | ||
| init_method="tcp://localhost:29500", | ||
| rank=0, | ||
| world_size=1, # Single process for unit test | ||
| ) |
There was a problem hiding this comment.
lets cleanup the process group after every test? Like we do here
There was a problem hiding this comment.
This way we can get rid of the try/catch here.
There was a problem hiding this comment.
Thanks, I added a tear down method
There was a problem hiding this comment.
Can we follow the pattern here to test against world size > 1?
There was a problem hiding this comment.
Added a test to do so. I think we should have some discussion when you get back about the nature of this test, ie. use CPU vs CUDA. I went ahead and did a world size of 2, but with CPU, so we weren't really testing the sharding here, just that it works with a larger world size.
There was a problem hiding this comment.
This PR is not urgent, however, so we can discuss late.
| # When using DMP, EmbeddingBagCollection returns Awaitable that needs to be resolved | ||
| if isinstance(embeddings_0, Awaitable): | ||
| embeddings_0 = embeddings_0.wait() |
There was a problem hiding this comment.
It does seem unfortunate/surprising that the rest of our code and/or pyg code doesn't support this type of tensor.
Can we add a TODO to look into this?
There was a problem hiding this comment.
I will add a TODO. This seems like an expected result of TorchRec sharding, as introduced by TorchRec itself.
| if not dist.is_initialized(): | ||
| dist.init_process_group( | ||
| backend="gloo", | ||
| init_method="tcp://localhost:29500", |
There was a problem hiding this comment.
let's also use get_process_group_init_method 1 so we can always have a free port?
There was a problem hiding this comment.
Changed to this.
2 participants
Scope of work done
models_test.pyjust to make sure we can use DMP to wrap a model. One test for forward, one for gradient flow.Where is the documentation for this feature?: N/A
Did you add automated tests or write a test plan?
Yes, added two unit tests.
Updated Changelog.md? NO
Ready for code review?: YES