Skip to content

[InstructBlip] Add accelerate support for instructblip #24488

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.

Sign up for GitHub

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

Merged
younesbelkada merged 5 commits into from
Jun 26, 2023
Merged

[InstructBlip] Add accelerate support for instructblip #24488

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7ee03df
add accelerate support for instructblip
younesbelkada Jun 26, 2023
838db9b
add `_keep_in_fp32_modules`
younesbelkada Jun 26, 2023
88ccd18
dynamically adapt `_no_split_modules`
younesbelkada Jun 26, 2023
802b536
better fix
younesbelkada Jun 26, 2023
1f4fd81
same logic for `_keep_in_fp32_modules`
younesbelkada Jun 26, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion src/transformers/models/instructblip/modeling_instructblip.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -281,6 +281,8 @@ class InstructBlipPreTrainedModel(PreTrainedModel):
r"language_model.decoder.embed_tokens.weight",
r"language_model.lm_head.weight",
]
_no_split_modules = ["InstructBlipAttention", "InstructBlipQFormerMultiHeadAttention"]
_keep_in_fp32_modules = []

# Copied from transformers.models.blip_2.modeling_blip_2.Blip2PreTrainedModel._init_weights with Blip2->InstructBlip
def _init_weights(self, module):
Expand Down Expand Up @@ -1264,11 +1266,18 @@ def __init__(self, config: InstructBlipConfig):
self.qformer = InstructBlipQFormerModel(config.qformer_config)

self.language_projection = nn.Linear(config.qformer_config.hidden_size, config.text_config.hidden_size)

if config.use_decoder_only_language_model:
language_model = AutoModelForCausalLM.from_config(config.text_config)
else:
language_model = AutoModelForSeq2SeqLM.from_config(config.text_config)

if language_model._no_split_modules is not None:
self._no_split_modules.extend(language_model._no_split_modules)

if language_model._keep_in_fp32_modules is not None:
self._keep_in_fp32_modules.extend(language_model._keep_in_fp32_modules)

self.language_model = language_model

# Initialize weights and apply final processing
Expand Down Expand Up @@ -1422,7 +1431,7 @@ def forward(

if attention_mask is None:
attention_mask = torch.ones_like(input_ids)
attention_mask = torch.cat([language_model_attention_mask, attention_mask], dim=1)
attention_mask = torch.cat([language_model_attention_mask.to(attention_mask.device), attention_mask], dim=1)

if self.config.use_decoder_only_language_model:
outputs = self.language_model(
Expand Down