Proper prefix handling in EarlyFusion sd hooks

[ebsmothers]

Thanks to @mkrainin for pointing this out and making the initial fix.

Our EarlyFusion state dict hooks don't work when an EarlyFusion module is wrapped as part of a larger model (as is the case with e.g. DDP). This is because state dict keys in the hooks have global key names, so we can't just use e.g. del state_dict["decoder.tok_embeddings.weight"] since decoder may not be a top-level module.

The fix is to use the prefix arg to state dict hooks, which gets recursively prepended while traversing the module tree (ref).

Test plan

Added a unit test wrapping EarlyFusion into an nn.Sequential so that the state dict keys have integer prefixes. This plus the existing test of state dict hooks pass after these changes:

pytest tests/torchtune/modules/model_fusion/test_early_fusion.py
...
================== 12 passed in 0.18s =============

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/2291

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit d470857 with merge base 890deab ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 66.64%. Comparing base (97e857f) to head (d470857).
Report is 10 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##            main    #2291       +/-   ##
==========================================
+ Coverage   8.97%   66.64%   +57.66%     
==========================================
  Files        305      353       +48     
  Lines      18166    20651     +2485     
==========================================
+ Hits        1631    13762    +12131     
+ Misses     16535     6889     -9646     

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

[RdoubleA]

Good find with the prefix argument, do we need to update DeepFusion as well?

[RdoubleA]

prefix will include the period?

[ebsmothers]

Yep, see here

[ebsmothers]

@RdoubleA DeepFusion should be good. State dict hooks are not defined at the model level there, but at least individual layers seem to already be using the prefix argument in a similar fashion. E.g.