-
Notifications
You must be signed in to change notification settings - Fork 411
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[RFC] Early fusion multimodal models #1904
base: main
Are you sure you want to change the base?
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/1904
Note: Links to docs will display an error until the docs builds have been completed. ❌ 3 New Failures, 4 Cancelled JobsAs of commit e37a3e1 with merge base 74139c9 (): NEW FAILURES - The following jobs have failed:
CANCELLED JOBS - The following jobs were cancelled. Please retry:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
def __init__( | ||
self, | ||
decoder: TransformerDecoder, | ||
encoders: nn.ModuleDict, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
encoders: nn.ModuleDict, | ||
encoder_tokens: Dict[str, int], | ||
decoder_trainable: bool, | ||
encoders_trainable: Dict[str, bool], |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
An error should be thrown if the different input dicts don't have the same keys
if decoder_trainable: | ||
trainable_params |= { | ||
f"decoder.{n}" for n, p in self.decoder.named_parameters() | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is missing the logic and the parameter to make fusion modules trainable/untrainable
been expanded to the number of tokens encoded for the given media. For example, if an image is tiled/patched | ||
and tokenized to 100 tokens, we assume the text sequence already has 100 "image" tokens as placeholders. | ||
""" | ||
embeds = self.tok_embeddings(tokens) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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/
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah let's talk about this offline, because from the reference code I was using the encoder tokens are part of the embedding table
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):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.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:This way we don't need any state dict hooks or decoder modifications. However, we are polluting the decoder model forward with more arguments.