[RFC] Early fusion multimodal models

[RdoubleA]

TODO: fix tests

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.

❌ 5 New Failures, 2 Cancelled Jobs

As of commit 024bfc7 with merge base d3039da ():

NEW FAILURES - The following jobs have failed:

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

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

• GPU tests / gpu_test (3.10, stable) (gh)
  ##[error]The operation was canceled.
• GPU tests / gpu_test (3.9, stable) (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

[acisseJZhong]

nit: EarlyFusionModel

[acisseJZhong]

Thanks for the RFC, you made it very clear what's difference between early fusion and late fusion!

About the design choice, I personally prefer Option 2 for the same reason as you mentioned. I think it's fine to "polluting" the decoder model forward a bit with some optional arguments for each modality. We might need something like

h = self.tok_embeddings(tokens) 
if speech: 
   h[speech_mask] += speech_encoder(input) 
if image: 
   h[image_mask] += image_encoder(input)

[acisseJZhong]

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

Why is this a concern? we only set decoder.tok_embeddings to identity within EarlyFusionModule? If outside EarlyFusionModule, we just use the normal decoder right?

[RdoubleA]

right, but if I passed in the decoder into EarlyFusionModule and modified a layer here, even if I use it separately the embedding layer will have been modified

[acisseJZhong]

curious why would we need this?

[RdoubleA]

Not sure yet if we need this... once we have a better idea of the full architecture let's chat with @pbontrager to see which components need to be fusion modules

[ebsmothers]

11th hour comment on the open design question: in my mind there are nonzero UX costs to either approach. If we patch the decoder embeddings to nn.Identity we introduce additional indirection that is pretty trivial but also pretty non-obvious (I claim any state dict hook is non-obvious when first debugging the inevitable key mismatch error until you find the actual code pointer). On the plus side, we fully contain the blast radius to multimodal model code, and text-only users do not have to worry about it. Conversely, I know we don't want to just add a bunch of random arguments to TransformerDecoder forward, especially ones that are very specific to multimodal models.

Personally I really don't like state dict hooks for the reason I described above. As soon as something (inevitably) goes wrong, it will take a lot more debugging and head-banging-against-the-wall before the user realizes that things are being swapped out under the hood. So perhaps it's no surprise, but I vote for the simple and dumb thing: just add an extra parameter to TransformerDecoder forward. I know that may be controversial, but I like doing the obvious thing, and I like to think our users would appreciate that as well.