[ZeRo] Parameter group support in constructor

[rohan-varma]

🚀 The feature, motivation and pitch

Motivated by: https://discuss.pytorch.org/t/adamw-zeroredundancyoptimizer-weight-decay-dictionary/141516

User might create optimizer such as:

optimizer = optim.AdamW(
                [
                    {"params": gain_or_bias_params, "weight_decay": 0.},
                    {"params": rest_params, "weight_decay": args.wd},
                ],
                lr=args.lr,
                betas=(args.beta1, args.beta2),
                eps=args.eps,
            )

where the first list is a dict specifying param groups. However,

optimizer = ZeroRedundancyOptimizer(
                [
                    {"params": gain_or_bias_params, "weight_decay": 0.},
                    {"params": rest_params, "weight_decay": args.wd},
                ],
                optim.AdamW,
                lr=args.lr,
                betas=(args.beta1, args.beta2),
                eps=args.eps,
            )

does not work as expected. A workaround is the following:

optimizer.add_param_group({"params": gain_or_bias_params, "weight_decay": 0.})

but for a better dev experience we should support passing them directly in the constructor.

Alternatives

No response

Additional context

No response

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang

[awgu]

Fixing this issue will fix this other one as well: #59973

[awgu]

Here is a rough outline of what to do to fix this issue. I can take on this fix if/when I get access to a multi-GPU machine if no one objects. This serves mostly as a self-documentation in case I do end up taking this on.

Testing Changes

• Modify the single-rank test test_constructor() so that iterables of dict do not raise a TypeError (see here).
• Add a multi-rank test that exercises the iterable of dict code path for the constructor. The test should be multi-rank (i.e. should be a method in TestZeroRedundancyOptimizerDistributed) to ensure that parameters are being partitioned non-trivially.

Core Changes

The constraint to keep in mind is that (1) add_param_group() exists as a method in the Optimizer class, (2) the Optimizer constructor relies on that method, and (3) ZeroRedundancyOptimizer.add_param_group() has additional semantics on top of its parent's version. Hence, the current implementation uses a self.initialized() flag to determine when to follow ZeRO's additional semantics.

Approach 1
After re-examining the code, I am inclined to believe that there are no issues if we simply modify _verify_and_init_params() to permit the argument params to be an iterable of dicts. The additional semantics of ZeRO's add_param_group() on top of Optimizer's add_param_group() are as follows:

• Clear caches and repartition the parameters
• Add the new parameter group to the local optimizer if needed
• Build the parameter buckets if needed

Now, at the time of calling the parent Optimizer constructor (with params as an iterable of dicts), (1) all cached data structures are not even defined, (2) the initial _partition_parameters() call has not happened yet, and (3) the local optimizer has not been initialized yet.

1. N/A
2. _partition_parameters() should work with multiple parameter groups. See here.
3. _init_local_optimizer() should work with multiple parameter groups. See here.
   As such, there should be no issue to simply remove the check against params being an iterable of dict. If anyone recalls the original reason why we could not support this or sees any issue, feel free to let me know. (Although in all likelihood, someone else or I can simply try this.)

Approach 2
If Approach 1 is infeasible for some reason, the other approach I have in mind follows the existing code path, except restricting itself to only the first parameter group in the iterable of dict parameter groups and then calls the existing ZeroRedundancyOptimizer.add_param_group() for each remaining parameter group at the end of the constructor.

• Pros: This effectively builds the existing workaround (see here) into the implementation itself. The code change should not be too complex.
• Cons: This re-partitions the parameters across workers N - 1 times, where N is the number of parameter groups and where each re-partition requires data copying when parameters_as_bucket_view=True. This may be a non-issue since we do not expect N to be more than 2 or 3 typically and since the cost is only incurred once up front.
• Implementation Sketch:
  • Modify this code block in _verify_and_init_params() to additionally accommodate the argument params being an iterable of dict.
  • One option is to have _verify_and_init_params() return a bool indicating if params contains Tensors or dicts. Then, given that bool in the constructor, after self.initialized is set to True (see here), call add_param_group() on the remaining parameter groups as needed.
  • Of course, a proper fix should rename variables and functions and update documentation appropriately.

[rohan-varma]

Thanks for the details @awgu! I'll assign you for now, feel free to unassign if you don't end up working on it