this is not needed anymore with modern pytorch, it was there initially because this functionality was not present, I think that you could remove it altogether

+1

I looked into replacing _broadcast_object() with torch.distributed.broadcast_object_list() as suggested by @mrshenli. However, broadcast_object_list() yields a significant slowdown in my preliminary testing. Precisely, running test_collect_shards() on the AI AWS cluster with a world size of 4 regresses from ~6 seconds to ~18 seconds (where the affected function is consolidate_state_dict()).

I am unable to resolve the issue, but this is what I found: There are two (possibly-related) points of slowdown in broadcast_object_list():

1. The implicit torch.cuda.synchronize() resulting from moving data from object_tensor_sizes from GPU to CPU via .item() takes ~4 seconds for one of the broadcasts on ranks 1 and 2 (while generally it takes ~1 millisecond). This confuses me since _broadcast_object() follows the exact same data movement pattern, yet it does not suffer from a slow synchronization.
2. Unlike _broadcast_object(), broadcast_object_list() pickles/unpickles the object via _object_to_tensor()/_tensor_to_object(), respectively. _tensor_to_object() consistently takes ~4 seconds.
   I am unable to see any connection between the two, but the similarity between their latencies is curious. Together, these two points cause consolidate_state_dict() to take >4 seconds when using broadcast_object_list(), while it only takes a few milliseconds when using _broadcast_object().

Any help investigating this is greatly appreciated!

NB: Without pickling, the tensor allocations are 939, 939, and 1131 bytes, while with pickling, the allocations are 571, 591, and 755 bytes.

I think I found the issue, and fixing this will require changing code in distributed_c10d.py. Here is my hypothesis:

Currently, broadcast_object_list() takes care to specially handle the NCCL backend case, making sure to move object_sizes_tensor and object_tensor to device memory before broadcasting. Upon receiving the broadcasted tensor, the code calls _tensor_to_object(), which loads the tensor to an object. The existing implementation uses a Pickler and Unpickler to save/load, but I have verified that the slowdown behavior is replicated using torch.save() and torch.load(), which itself uses pickling. For the remainder, I will assume a torch.save()/torch.load()-based implementation of _object_to_tensor() and _tensor_to_object().

The crux is the default behavior of torch.load() (and presumably Unpickler.load()), which deserializes (on CPU) and loads contained tensors to the device where they were saved. Suppose that the world size is 4. Then, in the received local_state_dict on rank j from the broadcast, the contained tensors are loaded to cuda:i where i != j. The time bottleneck appears to be exactly in loading these received objects to device, taking > 3.5 seconds.

The bandaid fix I have right now is changing:

• _object_to_tensor(obj) to _object_to_tensor(obj, pickle=True)
• _tensor_to_object(tensor, tensor_size) to _tensor_to_object(tensor, tensor_size, pickle=True, device=None)

If pickle == True, then the behavior is as-is to preserve backward compatibility. Otherwise, torch.save(obj, f) and torch.load(io.BytesIO(buf), map_location=device) are used instead of the _pickler and _unpickler, respectively. Then, broadcast_object_list() passes in device=current_device to _tensor_to_object() if using NCCL backend. I have confirmed that this change removes the slowdown.

However, this solution is only a temporary fix, and I think a more thorough examination is needed. I would need to see if a solution is feasible while still using the Pickler/Unpickler (though at first glance this is not promising since the Unpickler.load() provides no options).

I do wonder how much this behavior has affected the speed of other use cases. It seems to be pretty common that the received tensor from the broadcast is from a different device -- in fact, that should be the typical case?

Discussed with @andwgu offline on this. Since there are gaps in existing broadcast_object_list(), we are going to keep the current _broadcast_object method. Later, if bandwidth allows, we can work on a followup PR to improve broadcast_object_list() and use it here.

this was for compatibility with pytorch 1.5 (this comes from fairscale, which is compatible with older pytorch, current-2 and at the time this came out it was 1.5). this is not needed anymore indeed. There were a couple of update PRs for this code which never got reviewed, it's really good that you can spend some time in that

We have both _update_trainable and _init_local_optimizer called here?

I think following your other comment and removing _update_trainable() altogether in favor of only _build_param_buckets() (which checks parameters_as_bucket_view) might make this more clear. Then, the constructor will call _init_local_optimizer() to initialize the optimizer and call _build_param_buckets() to build the buckets, as the function names suggest.

I think we can get rid of this method entirely and just call _build_param_buckets and inside _build_param_buckets have the check for self.parameters_as_bucket_view.

I'm guessing this is referring to NCCL? If so we should mention this here and if we don't already have a gh issue, we should mention the lack of gather support in NCCL and add a code pointer for this to that issue.

from the top of my head that was with older Gloo (compatibility matrix was pytorch 1.5/1.6/1.7 & NCCL/Gloo & CPU/GPU with gloo), maybe that this is worth revisiting ?

We shouldn't have notes about test methods in the docs of public methods. Also, shouldn't this method be private? Do we actually expect users to call this method? If not, it should be private.

I'm wondering if we need this interaction between state_dict and consolidate_state_dict? Every time we call .step() can we update a variable needs_consolidation=True. Then when state_dict is called, we first consolidate the state dicts if needs_consolidateion=True, cache the result and set needs_consolidation=False. This way the user only has to worry about calling state_dict.

No, that would not work because of the collective needs. The issue here is that users might call .state_dict() on one rank only (in practice they often will), which means that you cannot consolidate at that point with any of the collective communication primitives. Something like RPC would alleviate that, but that's a bigger change, without this you need the collective consolidation before you can checkout (and doing this for every step is super slow, of course).

this can be replaced by the following one?

+1

Please feel free to ignore if you already have a convenient way to inspect generated docs. I usually use this gist to serve html files locally.

(This is prior to this PR) IIUC, a Tensor will return its id by default when used as hash key. Is it necessary to call id(p) here?

if this is unnecessary, we can also revert the var name to _param_to_index_cache`

(seconding that personally, but soft preference)

I examined the hashing for both Tensor and Parameter and neither seem to use id() by default.

>>> model = torch.nn.Linear(1, 1)
>>> t = model.weight.data  # t: torch.Tensor
>>> d = dict()
>>> d[t] = 1
>>> d[id(t)]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
KeyError: 139809367343568
>>> p = model.weight  # p: torch.nn.parameter.Parameter
>>> d = dict()
>>> d[p] = 1
>>> d[id(p)]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
KeyError: 139809367327664

Let me know if I tested this incorrectly.

Thanks for checking this. My mistake.

Hold on, I am not sure the test is correct. We are passing a Tensor: int pair, but checking an int: int pair? Isn't that suppose to fail? Below is Tensor'e hash function:

Ah, I misunderstood. Yes, I believe we can switch back to _param_to_index and make it of type dict[torch.Tensor, int]. Then, instead of the key being id(p), it can directly be p.

nit: other parts of this file and this code base usually break long arg list into one arg per line, e.g.:

Could you please post a screenshot of the rendered the docs? Thanks!

I changed it to be like this because any other text on the same line as * state or * param_groups appears in the grey box, but anything in the next line appears underneath in the white. The old way looked strange with differs between optimizer classes. being indented and in the white, while * state - a dict holding current optimization state. Its content was in the grey box above.

add an error message for this assert?

has -> have?

Very minor, but I believe the current way is correct since it should be trainability [...] has changed rather than trainability [...] have changed (trainability being singular not plural)
:)

Oh I see, my bad, thanks!

I would remove that method altogether, I don't think that there's a use case for it actually (we don't expose a way to reload local state dicts)

could be worth it adding an assert with a nice error message ? not super likely to be triggered, but in case that happens it would help debugging a fair bit

The assumptions in the NOTE: ...s should be verified at the very beginning of the constructor (line 171). I was intending the notes for documentation purposes (for a user reading the source or for a developer planning to modify the implementation).

Do you think we should do a second pass of verifications each time this function _build_param_buckets() is called?

sorry, a bit late to react, I meant checking for this assumption explicitly (at least the first time the buckets are built ?) and nicely error out if any issue, but maybe that it was already there indeed. Non blocking, just a brain dump

shall we use one arg per line format to stay consistent with other APIs in this file and other files?

(no need to change this PR, just a heads-up) there is another recursive_to in DDP. It might be helpful to consolidate these two into one common utility function.

(no need to change this PR, just FYI) if you would like to add a note block, this needs to be .. note::. But since this is a private API and the docs won't render, the current one also works.