Selective mode saving

[AAnoosheh]

What does this PR do?

Type of change: ? Bug fix

Overview: Filter out KD state from ModelOpt state list when saving. This allows for applying the KD mode after a modelopt checkpoint restore without it complaining that it was already applied previously.

Usage

# Add a code snippet demonstrating how to use this

Testing

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes/No
• Did you write any new necessary tests?: Yes/No
• Did you add or update any necessary documentation?: Yes/No
• Did you update Changelog?: Yes/No

Additional Information

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 74.46%. Comparing base (38550b0) to head (d4883ed).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #596   +/-   ##
=======================================
  Coverage   74.45%   74.46%           
=======================================
  Files         182      182           
  Lines       18250    18255    +5     
=======================================
+ Hits        13588    13593    +5     
  Misses       4662     4662           

☔ 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.

[realAsma]

Approving the PR to unblock merging for the release.
Attaching the teacher model design was not a good choice looking back - it is creating lot of headaches. We should just deal with distillation with the loss function or trainer.

We dont need to necessarily need to remove the mtd.convert - but stop maintaining it and use a trainer/loss function based distillation support -

QAT + distillation is a critical piece going forward - It would be great to simplify things - If any design choices turned out problematic - we dont have to keep patching it - we could redirect our energy for a better design.

Cc @kevalmorabia97 @jenchen13 @ChenhanYu