This map also lives in DistributedOptimizer. Shall we consolidate these two into a single source of truth?

This looks fine for now. But let's add a followup PR to automatically do the translation for users, similar to how DistributedOptimizer handles this.

nit (can be done in follow PRs): curious, callbacks on buckets will fire in order, so this can be simplified to a list of GradBucket objects?

For a given rank, we only need to store references to the GradBuckets that the rank uses to update. As such, even though the buckets fire in order, the "keys" for this data structure are only a subset of all indices.

If the total number of buckets is not large, using a list and either storing all buckets or storing Nones in unused positions may be faster than storing a dict, but I used a dict for now.

Ah, I see, makes sense.

nit (can be done in followup PRs): looks like each Future in this dict is just read once, so this cache probably won't save on Python-CPP context switch cost. Maybe this can be deleted, and directly use bucket.get_future()?

In general, I would try to avoid the state vars unless that would give us either perf gain or considerable code simplification. As more vars might make it harder to understand and maintain.

I may be misunderstanding, but the reason I stored this dict is to have a reference to all relevant Futures when running the hook for the last bucket to be able to call wait() on them. bucket is of type GradBucket, which does not offer a get_future() method. Rather, c10d.ProcessGroup.work objects provide a get_future() method. However, at the moment, this seems a bit backward. The hook passed into hook_then_zero_step() (most typically allreduce_hook()) returns a Future, where that Future is exactly gotten via dist.all_reduce().get_future(). Either way, I do not see how to get around storing at least something per bucket to update that gives access to the bucket's corresponding Future to give the rank access to the Future when executing the hook for the last bucket.

curious, why this is device_to_params instead of device_to_buckets?

Good point -- I will change it. It was from before I created the class _DDPBucket and instead had a List[torch.Tensor] as the inner-most element in the dict value.

(can be done in followup PRs) This conceptually overlaps with DDP's Gradient Bucket. We probably want to let DDP expose its GradBuckets instead of creating new classes?

I think that it may be possible to get rid of this class altogether with how DDP's GradBucket is exposed currently. I can investigate this in a follow-up PR.

This is still the allreduce future, right? Meaning that waiting on this future (by DDP backward()) won't guarantee that the ZeRO broadcast is done? We probably need to return torch.futures.collect_all(all_zero_broadcast_futures) for the last bucket instead.

Among these reviews, this is the most important one, as it might have impact on correctness.

Resolved offline.

@andwgu Do you mind describing how this was resolved?

https://github.com/andwgu/pytorch/blob/ab54c52fa54dac017ff6c9eadae49a97aa117432/torch/distributed/algorithms/ddp_comm_hooks/ddp_zero_hook.py#L175
That line should provide the necessary synchronization already. (The wait() on the returned future in finalize_backward() will not guarantee that the ZeRO broadcasts are done, but these additional wait()s will.)

Also, as a side note, the torch.futures.collect_all(all_zero_broadcast_futures) would change the hook semantics since it would then return a future for the broadcasts instead of a future for the bucket all-reduce.

state (object): ... ?

Fixed.

Actually, I was reading this: https://stackoverflow.com/questions/39817081/typing-any-vs-object

Even though the type annotation for state in register_comm_hook() in distributed.py is object, I wonder if that should be Any instead. If I am understanding correctly, if you type annotate state with object, then to appease a type checker, every operation involving state must work for any general object type. In other words, annotating with object is actually restrictive rather than all-encompassing.

curious, is this the conventional way to document function signature? Curious, why this does not match Callable[[Any, dist.GradBucket], torch.futures.Future]?

Good catch. I accidentally used the type signature style from the language Standard ML.

:meth:step -> :meth:step` ?

Fixed.

print a warning if parameters_as_bucket_view=True?

Added a warning using logging.warning() in that case.

add some explanation that ZeroRedundancyOptimizer is fully initialized after DDP buckets are rebuilt?

Added.

assert overlap_with_ddp == False?

Added.

IIUC, the recommended new API for this is param.set_(). Let's do that change in a followup PR, as there might be subtle differences.

https://pytorch.org/docs/stable/generated/torch.Tensor.set_.html#torch.Tensor.set_

One issue with using set_() is that we must set requires_grad=False and later reset it to True since otherwise we get the error saying that a leaf Variable that requires grad is being used in an in-place operation.

More importantly, I am unable to get it to manage the strides correctly. It seems like if you specify offset, then you must also specify size and stride as arguments. The method uses that stride when accessing the shared underlying storage. However, we want param to preserve its non-flattened stride, while sharing the storage with the flattened bucket. This suggests the need for additional view manipulation (which I have not figured out how to get correct). Hence, as far as I can see, I do not think that using set_() will be any simpler for this copying.

Do we need to error out if the backend is not NCCL? (as gloo might hang today)

Added check.

this is a bit different from the above hook, which also checks assert overlap_info.status == _OverlapStatus.UNINITIALIZED

This means there is no optimizer step in the first 2 iterations and gradients are discarded? Do we need to mention this in the doc, or do you plan to change this in followup PRs?

I added warnings to mention this in both hook functions and ZeroRedundancyOptimizer.

The motivation for this is so that we can avoid initializing a non-functional local optimizer and then re-initializing a functional optimizer soon thereafter. Even if we did that, we would have to transfer optimizer state correctly.
We do not know until DDP buckets are built which parameters a given rank is assigned to update. I have not yet come up with a clean solution that allows not throwing away the first two inputs.

Small note: The first three inputs are unused if static_graph=True for the DDP model.

(can be done in followup PRs)

We do not know until DDP buckets are built which parameters a given rank is assigned to update. I have not yet come up with a clean solution that allows not throwing away the first two inputs.

Since ZeRO's overlap states knows whether the buckets are rebuilt or not, is it possible that, in the first 2-3 iterations, ZeroRedundancyOptimizer.step() still performs comp + comm the same way as no-overlap cases? And only after the buckets are rebuilt, ZeroRedundancyOptimizer.step() becomes no-op?

Yes, we can do that. Originally, I was averse to the idea of re-initializing the local optimizer, but upon further thought, providing an equivalent functionality to the user is definitely more important.

What we can do is assume parameter_as_bucket_view=False behavior until buckets are rebuilt. This means that the parameters will be partitioned according to the normal sorted-greedy algorithm and a local optimizer will be initialized as normal at construction time. Then, when buckets are rebuilt, the parameter partitioning will change and the local optimizer will be reinitialized. I can do this in a follow-up PR.

I made these uniform across the file to use two colons. I checked the render, and the warning still appears as desired.

This will not actually be rendered since the method is private, but just in case, I think the docstring render expects a newline before any bulleted list.

Similar thing here.

This rename is to prevent implicit shadowing of the method's argument params_per_rank.

If not, this should error out, as it might hang with Gloo?

NVM, I saw, there are TODOs to cover CPU tests as well.

Curious, why do we need a barrier here?

Good point. Removed.

sort?

Sorted.

Do we have to explicitly import distributed_c10d? If the goal was to use Backend.NCCL, will dist.Backend.NCCL work?

https://pytorch.org/docs/master/distributed.html#torch.distributed.Backend

Changed to dist.get_backend() and dist.Backend.NCCL.

Let's cite issue #62300 here

Added.

(can be done in followup PRs)

We do not know until DDP buckets are built which parameters a given rank is assigned to update. I have not yet come up with a clean solution that allows not throwing away the first two inputs.

Since ZeRO's overlap states knows whether the buckets are rebuilt or not, is it possible that, in the first 2-3 iterations, ZeroRedundancyOptimizer.step() still performs comp + comm the same way as no-overlap cases? And only after the buckets are rebuilt, ZeroRedundancyOptimizer.step() becomes no-op?

Shall we add a check here to error out if this assumption was violated? (maybe by remembering the last bucket index)

This is a good idea. Currently, I am a bit uncertain on how to do this. I will address it in a follow-up PR.

Remembering the last bucket index might be circular since the way we currently tell if it is the last one is by its index. If the buckets fire in order, then we should be guaranteed that the CPU part of each hook executes in order right? In that case, we can reuse the bucket_indices_seen data structures used in the other hook and make sure that the current bucket index is 1 more than the previous index seen (clearing it after each iteration).

Let's add a short comment to explain why this cannot be chained as a callback to the returned future.

Added:

# NOTE: This should not be chained as a callback to the last bucket's
# all-reduce future since that would add synchronization that delays
# all optimizer computation to wait for that last all-reduce

nit (can be done in followup PR, and this is just my personal preference, please feel free to ignore): The _OverlapInfo type is not owned by this hook function and lives a bit far away from this file. It might be better to wrap its state manipulations into member method in _OverlapInfo, instead of directly manipulating them here and above.

I think this definitely makes sense. I will try to refactor this in a follow-up PR.

Ah, I see, makes sense.