[Gradient Compression] Allow PowerSGD to run vallina allreduce for the first K iterations

[wayi1]

Stack from ghstack:

• [Gradient Compression] Refactor default_hooks.py and powerSGD_hook.py by creating a util function that make a vanilla allreduce future #51094 [Gradient Compression] Refactor default_hooks.py and powerSGD_hook.py by creating a util function that make a vanilla allreduce future
• [Gradient Compression] Explicitly specify the dtype of the error tensor #50985 [Gradient Compression] Explicitly specify the dtype of the error tensor
• [Gradient Compression] Simplify the implementation of warm-start #50981 [Gradient Compression] Simplify the implementation of warm-start
• [Gradient Compression] Typo fixes in PowerSGD #50974 [Gradient Compression] Typo fixes in PowerSGD
• [Gradient Compression] Allow PowerSGD to run vallina allreduce for the first K iterations #50973 [Gradient Compression] Allow PowerSGD to run vanilla allreduce for the first K iterations

This can extend the original PowerSGD method to a hybrid approach: vanilla allreduce + PowerSGD. This can help further improve the accuracy, at the cost of a lower speedup.

Also add more comments on the fields in PowerSGDState.

Original PR issue: Investigate Applying PowerSGD to Communication Hook for Gradient Compression #47202

Differential Revision: D26031478

[facebook-github-bot]

💊 CI failures summary and remediations

As of commit 61f32b6 (more details on the Dr. CI page):

---

• 1/1 failures possibly* introduced in this PR
  • 1/1 non-CircleCI failure(s)

---

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions to the (internal) Dr. CI Users group.

[rohan-varma]

Awesome, excited to see the results of this! One interesting area to explore would be hook composition, i.e. the user can specify whether to run vanilla allreduce or fp16 or even some other algorithm for the first t iterations. Ideally, it would be great if we can just express this by composing the different hooks as building blocks, maybe something like:

def main_hook(state, iter, args):
    hook_for_iter = dispatch_table[iter, args] # dispatch table that determines which hook to run based on args
   return hook_for_iter(state, args)

This is admittedly overkill for the current use case, but may be an interesting area to explore if we realize we are composing more hooks.

[rohan-varma]

Also, can we add an appropriate unittest for this change?

[wayi1]

Awesome, excited to see the results of this! One interesting area to explore would be hook composition, i.e. the user can specify whether to run vanilla allreduce or fp16 or even some other algorithm for the first t iterations. Ideally, it would be great if we can just express this by composing the different hooks as building blocks, maybe something like:

def main_hook(state, iter, args):
    hook_for_iter = dispatch_table[iter, args] # dispatch table that determines which hook to run based on args
   return hook_for_iter(state, args)

This is admittedly overkill for the current use case, but may be an interesting area to explore if we realize we are composing more hooks.

That will be one direction of the future work. Since we don't have many comm hooks available, it's hard to see how it will work out. I would prioritize exploring other compression approaches over this. A few comments about the more generic combination:

1. I don't plan to consider fp16 as a compression hook in this setup. Instead, I will consider it as a wrapper, so it can be combined with both vanilla allreduce and PowerSGD, as well as any other future gradient compression. FP16 is just a type casting, it can be applied in conjunction with any other comm hook, rather than being mutually exclusive.

Therefore, if we do not consider fp16 as a valid choice, then so far we only have vanilla allreduce and PowerSGD two valid options.

2. I am not sure how useful if we support switching back and forth different compression schemes, especially for biased compressor like PowerSGD that typically requires error compensation. PowerSGD needs to cache some error tensors and warm-up tensors. Switching to another compression scheme will discard these cached states, so there will be some extra accuracy loss due to the discontinuation.

[wayi1]

Also, can we add an appropriate unittest for this change?

Since by default, start_powerSGD_iter is 10 instead of 0, so the first 10 iterations will just run vanilla allreduce. It's already covered by the unit test.