[Transform] Serialize with tied weights

[kylesayrs]

Purpose

• Support saving models with transforms attached

Prerequisites

• [Transform] apply_transform_config #348

Semi-Prerequisites

These changes are required in order to support saving models with offloaded Transforms. Models without offloading do not require these changes.

• [Core] [Offloading] Enable saving offloaded models with multiple shared tensor groups huggingface/transformers#39263
• [Core] [Offloading] Fix saving offloaded submodules huggingface/transformers#39280

Changes

• Implement _update_tied_weights which
  • Updates the _dynamic_tied_weights_keys attribute of the transform modules. This property is read by transformers during saving and tells transformers to delete duplicates of the tied weights before saving.
  • Sets the meta tensors of shared weights to be identical, so they can be recognized and deleted by transformers

Testing

• Add serialization tests

The base branch was changed.

[brian-dellabetta]

Looks like your style checks have different behavior than what's on main. I've seen this elsewhere, not sure what's causing it. Maybe our version pin on flake8>=3.8.3 is too loose and later versions have different behavior?

[brian-dellabetta]

This and #391 are both ready to merge right? We can send over to team for review if so

[kylesayrs]

@brian-dellabetta This is ready to merge. Afaict the only thing left from the head branch is applying transformers in higher granularity, which is still unknown whether this actually has an effect on accuracy

[brian-dellabetta]

LGTM!

[brian-dellabetta]

👍