[Bugfix] Prepare KD Models when Saving

[kylesayrs]

Purpose

• Fix saving knowledge distillation model checkpoints
• Fix finetune transformer tests

Prerequisites

• [Bugfix] Fix saving offloaded state dict #172
• Implement Missing Modifier Methods #166

Changes

• Have the session mixin check if the model is a KD model, if so it should prepare all the wrappers to be saved
• A wrapper which has been prepared acts transparently w.r.t. the named_modules function. The test failures arose because get_state_dict_offloaded_model(model) uses named_modules rather than model.state_dict() to gather all the state dict tensors, and the KD wrapper did not account for this use case.

Testing

• tests/llmcompressor/transformers/finetune/test_oneshot_then_finetune.py now passes without shared tensor conflicts

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

[kylesayrs]

Ready to go, just need to merge prerequisites first

[dsikka]

Is this ready for review?
In draft mode

[kylesayrs]

@dsikka Yes, ready for review but in draft mode because it should not be merged until the prereqs are merged

The base branch was changed.