Akoumparouli/fix get params for weight decay optimization
State: closed | Pull request: yes
Labels: name: NLP

# What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

**Collection**: [Note which collection this PR will affect]

# Changelog 
- Add specific line by line info of high level changes in this PR.

# Usage
* You can potentially add a usage example below

```python
# Add a code snippet demonstrating how to use this 
```

# Jenkins CI
To run Jenkins, a NeMo User with write access must comment `jenkins` on the PR.

# Before your PR is "Ready for review"
**Pre checks**:
- [ ] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [ ] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Related to # (issue)


jenkins

@akoumpa is this still in draft phase?

Thanks @akoumpa. This change **does** fix the mlperf backward-compatibility, although I'm worried it might introduce another one for different scenarios, i.e. for models that previously had no params in "no_weight_decay_params" bucket? Did you test this case maybe?

jenkins

@ericharper will remove draft once jenkins passes.
@mikolajblaz I assumed that if a bucket does not have any parameters it should not be part of optimizer's param groups, I'm not sure what the use-case would be here since it's empty.

> @ericharper will remove draft once jenkins passes. @mikolajblaz I assumed that if a bucket does not have any parameters it should not be part of optimizer's param groups, I'm not sure what the use-case would be here since it's empty.

I totally agree, just not sure what are the consequences of omitting that group altogether in terms of backwards compatibility

@mikolajblaz I think it's fine, the code will yield a group if it has parameters, consumer code does not appear to make explicit assumptions on the number of returned groups and from the PyTorch documentation that also seems to be the right way to handle it. The problem would be to yield groups with empty parameters, rather than the current case. If you have an explicit example in mind happy to look into it. Apex also [seems ok](https://github.com/NVIDIA/apex/blob/141bbf1cf362d4ca4d94f4284393e91dda5105a5/apex/contrib/optimizers/distributed_fused_adam.py#L523) with Iterable[dict]

> @mikolajblaz I think it's fine, the code will yield a group if it has parameters, consumer code does not appear to make explicit assumptions on the number of returned groups and from the PyTorch documentation that also seems to be the right way to handle it. The problem would be to yield groups with empty parameters, rather than the current case. If you have an explicit example in mind happy to look into it. Apex also [seems ok](https://github.com/NVIDIA/apex/blob/141bbf1cf362d4ca4d94f4284393e91dda5105a5/apex/contrib/optimizers/distributed_fused_adam.py#L523) with Iterable[dict]

ok, thanks for checking
I would approve the PR but I don't think I can