Layer-Wise Distillation (#1272)

* Initial Commit with Alex's Work

* Update `student_names` -> `student_layer_names`
Update `teacher_names` -> `teacher_layer_names`

* Intermediate commit

* Styling

* Reorg initialize

* More cleanups

* Update docstring

* Moving finalize logic to update

* Tests passing a bit

* Fixing lifecycle tests

* Changing projection to dict

* Cleanup

* Adding quantization hooks test

* Add failing test for optimizer serialization

* Monkey patching  optimizer state_dict method

* Apply suggestions from code review

Co-authored-by: Konstantin Gulin <[email protected]>

* Update src/sparseml/pytorch/sparsification/distillation/modifier_per_layer.py

* Adding missing docstrings

* Respond to review on modifier/optimizer state_dict

* Add a test for modifier load before forward pass

* Updating comments

* Fix failing test

* Add more asserts based on @bfineran 's  comments

* * Rename `_DISTILL_PARAM_GROUP_KEY` -> `DISTILL_PARAM_GROUP_KEY`
* Add to `DISTILL_PARAM_GROUP_KEY` to `__all__`

* Move state dict patching to a helper function

* Quality

Co-authored-by: Corey Lowman <[email protected]>
Co-authored-by: corey-nm <[email protected]>
Co-authored-by: Konstantin Gulin <[email protected]>

Author: 3 people

src/sparseml/pytorch/sparsification/distillation/__init__.py

@@ -15,3 +15,4 @@
 # limitations under the License.
 
 from .modifier_distillation import *
+from .modifier_per_layer import *

src/sparseml/pytorch/sparsification/distillation/modifier_distillation_base.py

@@ -332,6 +332,7 @@ def loss_update(
             teacher_outputs=teacher_outputs,
             student_labels=student_labels,
             teacher_labels=teacher_labels,
+            optimizer=optimizer,
         )
 
         total_loss = self.compute_total_loss(loss, distillation_loss)