Support layers with a mix of frozen and learnable parameters

[karthikprasad]

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Docs change / refactoring / dependency upgrade

Motivation and Context / Related issue

#436

How Has This Been Tested (if it applies)

[WIP] unit tests

In addition to fixing the issue, this also results in massive memory savings. See CUDA memory usage before and after this change in https://colab.research.google.com/drive/1jYpZ2uVz2UXCoj2v0Vf30uqWO2oAl3de?usp=sharing (5.7 GB --> 27MB)

Checklist

• The documentation is up-to-date with the changes I made.
• I have read the CONTRIBUTING document and completed the CLA (see CONTRIBUTING).
• All tests passed, and additional code has been covered with new tests.

[facebook-github-bot]

@karthikprasad has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@karthikprasad has updated the pull request. You must reimport the pull request before landing.

[facebook-github-bot]

@karthikprasad has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[ffuuugor]

(5.7 GB --> 27MB)

I now bestow upon you the proud title of "The Chief Metric Gaming Officer", that's a beautiful experiment design.

But in all seriousness, great change, thanks a lot!
Having to do that in every grad sampler is not ideal, but I can't see any other option either

[facebook-github-bot]

@karthikprasad has updated the pull request. You must reimport the pull request before landing.

[facebook-github-bot]

@karthikprasad has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.