Fix TransformerQA distributed dead-lock bug - #193
Conversation
| if dist.get_backend() == "nccl": | ||
| device = torch.cuda.current_device() | ||
| else: | ||
| device = torch.device("cpu") | ||
| _min_node_batch_size = torch.tensor(min_node_batch_size, dtype=torch.int, device=device) | ||
| dist.all_reduce(_min_node_batch_size, op=dist.ReduceOp.MIN) | ||
| min_node_batch_size = _min_node_batch_size.item() |
There was a problem hiding this comment.
This is becoming a really common pattern. It would be to have a helper function for this called dist_reduce or something. Then we could just do:
| if dist.get_backend() == "nccl": | |
| device = torch.cuda.current_device() | |
| else: | |
| device = torch.device("cpu") | |
| _min_node_batch_size = torch.tensor(min_node_batch_size, dtype=torch.int, device=device) | |
| dist.all_reduce(_min_node_batch_size, op=dist.ReduceOp.MIN) | |
| min_node_batch_size = _min_node_batch_size.item() | |
| min_node_batch_size = dist_reduce(min_node_batch_size, dist.ReduceOp.MIN, dtype=torch.int) |
|
I don't think there's an easy way to write a test for this. |
dirkgr
left a comment
There was a problem hiding this comment.
There is no way to make this work correctly with uneven batch sizes, so that it doesn't throw anything away?
|
There might be a way. I'll investigate some more tomorrow. |
|
@dirkgr updated. Now can handle different batch sizes. |
| cast(List[str], best_span_string) | ||
| cast(List[List[str]], answer_strings) |
There was a problem hiding this comment.
I never knew you could do that.
dirkgr
left a comment
There was a problem hiding this comment.
This is great! If it comes up with the other metrics, we should use this as a template for how to fix the others.
| ) | ||
| def __call__( | ||
| self, | ||
| best_span_string: Union[str, List[str]], |
There was a problem hiding this comment.
Should be called best_span_strings now?
|
|
||
| ### Fixed | ||
|
|
||
| - Fixed the potentially for a dead-lock when training the `TransformerQA` model on multiple GPUs |
Closes allenai/allennlp#4886