Add L2NormHook and use it in megatron.py

[danielkorzekwa]

What does this PR do?

• Add L2NormHook and use it in megatron.py
• Using L2NormHook removes code duplication between _DynamicSelfAttention and _DynamicMLP

This is the first step towards reusing activation scores logic across Minitron and Puzzle. Next steps:

• complete redesign of megatron.py - move other activation hooks logic to hooks.py
• then combined those hooks.py with a similar hooks.py functoriality in puzzle (modelopt/torch/_compress/activation_scoring/activation_hooks/hooks.py)

Questions:

• why in the code before and after this redesign we store temp variables in two ways _register_temp_attribute and self.hook_handle)?

self._register_temp_attribute("_activation_hook", activation_hook)
# TODO: confusion: why hook_handle is removed manually in export() and not using _register_temp_attribute?
self.hook_handle = self.linear_fc2.register_forward_hook(activation_hook)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 74.37%. Comparing base (f10be0d) to head (5aaaf1a).
⚠️ Report is 2 commits behind head on feature/compress.

Additional details and impacted files

@@                Coverage Diff                @@
##           feature/compress     #599   +/-   ##
=================================================
  Coverage             74.37%   74.37%           
=================================================
  Files                   182      182           
  Lines                 18219    18219           
=================================================
  Hits                  13550    13550           
  Misses                 4669     4669           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[kevalmorabia97]

even if we register hook_handle as temp attribute, we still need to call hook_handle.remove() to remove the hook so there's no change. Temp attribute will be remove from model i.e. self.hook_handle reference will be dropped but that still doesnt remove the actuall pytorch hook added to the forward

[danielkorzekwa]

I understand now.