[RFC] Early fusion multimodal models

[RdoubleA]

TODO: fix tests, update DeepFusion to support multiple encoders, add docstrings

Context

This is a focused RFC based on @pbontrager 's excellent original RFC on multimodal fusion models #1283. Since the RFC, we have already landed Deep Fusion model components. This PR discusses and implements the EarlyFusionModel component, along with testing and some lint updates.

Early fusion is simply a decoder with 1 or more extra encoders that merges their outputs with the token embeddings tables. The challenge lies in how we merge the embeddings and pass it into the decoder.

Design

There is one design consideration I am seeking feedback on, and that is the EarlyFusionModel's usage of self.decoder.tok_embeddings. It accesses the decoder's token embedding table outside of the decoder forward because we need to merge the image encoder and any other modality encoder's output embeddings with the text embeddings (in this case just concatenate in sequence dimension):

embeds = self.tok_embeddings(tokens)
bsz, seq_len, embed_dim = embeds.shape
for encoder, inp in (encoder_input or {}).items():
    encoder_embeds = self.encoders[encoder](**inp)
    encoder_mask = (tokens == self.encoder_tokens[encoder]).expand(bsz, seq_len, embed_dim)
    embeds[encoder_mask] = encoder_embeds
    
output = self.decoder(embeds, mask, input_pos)
return output

Now, instead of token ids, we are passing in the merged embeddings directly into the decoder. But since we already used the text-only tok_embeddings from the decoder, we need to skip it when passing in the merged embeddings for the final decoder output. There are two ways we can do this.

State dict surgery

In the current code changes and suggested by the original RFC, we can manually set self.decoder.tok_embeddings = nn.Identity() so that it becomes a no-op when you forward pass with merged embeddings.

• This will require additional state dict hooks to make sure checkpoint saving and loading is still maintained despite the module change
• If a user wants to use the decoder outside of the EarlyFusionModule in the same script, they will need to restore the original tok_embeddings module from nn.Identity

Additional input_embeds kwarg

We could add a new keyword argument in TransformerDecoder forward for input embeddings. If this is passed in, we automatically skip the token embeddings:

h = self.tok_embeddings(tokens) if input_embeds is None else input_embeds

This way we don't need any state dict hooks or decoder modifications. However, we are polluting the decoder model forward with more arguments.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 3 New Failures, 4 Cancelled Jobs

As of commit e37a3e1 with merge base 74139c9 ():

NEW FAILURES - The following jobs have failed:

• GPU tests / gpu_test (3.10, stable) (gh)
  tests/torchtune/modules/model_fusion/test_fusion_models.py::TestDeepFusionModel::test_set_trainable_params
• Lint / lint (3.10) (gh)
  ##[error]Process completed with exit code 1.
• Unit Test / unit_tests (3.9) (gh)
  tests/torchtune/modules/model_fusion/test_fusion_models.py::TestDeepFusionModel::test_set_trainable_params

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• GPU tests / gpu_test (3.11, stable) (gh)
  ##[error]The operation was canceled.
• GPU tests / gpu_test (3.9, stable) (gh)
  ##[error]The operation was canceled.
• Unit Test / unit_tests (3.10) (gh)
  ##[error]The operation was canceled.
• Unit Test / unit_tests (3.11) (gh)
  ##[error]The operation was canceled.

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

[pbontrager]

Thanks for putting this up Rafi! I left some comments on the implementation, but I'll leave the state dict discussion to others as we've already chatted on this.

[pbontrager]

nit: I think it would be nice if we allowed all of the encoder params to be a list/dict or single value input to make single encoder builders look much cleaner. Then we can package them as an iterable in the init.

[pbontrager]

An error should be thrown if the different input dicts don't have the same keys

[pbontrager]

This is missing the logic and the parameter to make fusion modules trainable/untrainable

[pbontrager]

You can't do this because the encoder tokens won't be in the tok_embeddings table. You need to first filter those out as in here https://www.internalfb.com/intern/paste/P1666298928/

[RdoubleA]

Yeah let's talk about this offline, because from the reference code I was using the encoder tokens are part of the embedding table