Fixing memory leak in Joint QKV Attention Bridge

[jlarson4]

Description

#1043 introduced a flaw in the deepcopy initially added in #960.

Before #1043, deepcopy(blocks_template) was cheap because the template's submodules didn't hold bound methods referencing the adapter. Once split_qkv_matrix became a bound adapter method stored on the attention submodule, the deepcopy started transitively copying the entire adapter (and everything it references) once per attention layer.

This affected all architectures that used the JointQKVAttentionBridge class: bloom, neox, pythia, and qwen.

The solution was to override __deepcopy__ for the JointQKVAttentionBridge, to ensure that each copy shares a reference to the original split_qkv_matrix & config, rather than creating their own unique copies in memory

Type of change

Please delete options that are not relevant.

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility