Skip to content

[Refactor] Migrate MPT to standardized output tracing decorators #44071

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

Closed
ArivunidhiA wants to merge 1 commit into from
Closed

[Refactor] Migrate MPT to standardized output tracing decorators #44071

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
150 changes: 38 additions & 112 deletions src/transformers/models/mpt/modeling_mpt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,10 @@
TokenClassifierOutput,
)
from ...modeling_utils import PreTrainedModel
from ...utils import auto_docstring, logging
from ...processing_utils import Unpack
from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
from ...utils.generic import merge_with_config_defaults
from ...utils.output_capturing import capture_outputs
from .configuration_mpt import MptConfig


Expand Down Expand Up @@ -184,7 +187,6 @@ def forward(
attention_mask: torch.Tensor,
layer_past: Cache | None = None,
use_cache: bool = False,
output_attentions: bool = False,
cache_position: torch.Tensor | None = None,
):
# hidden_states: [batch_size, seq_length, hidden_size]
Expand All @@ -211,7 +213,7 @@ def forward(

# MLP.
output = self.ffn(layernorm_output, residual)
return output, attn_weights
return output


@auto_docstring
Expand All @@ -220,6 +222,10 @@ class MptPreTrainedModel(PreTrainedModel):
base_model_prefix = "transformer"
supports_gradient_checkpointing = True
_no_split_modules = ["MptBlock"]
_can_record_outputs = {
"hidden_states": MptBlock,
"attentions": MptAttention,
}


@auto_docstring
Expand Down Expand Up @@ -255,6 +261,8 @@ def build_mpt_alibi_tensor(self, num_heads, sequence_length, alibi_bias_max=8, d
def set_input_embeddings(self, new_embeddings: torch.Tensor):
self.wte = new_embeddings

@merge_with_config_defaults
@capture_outputs
@auto_docstring
def forward(
self,
Expand All @@ -263,12 +271,9 @@ def forward(
attention_mask: torch.Tensor | None = None,
inputs_embeds: torch.LongTensor | None = None,
use_cache: bool | None = None,
output_attentions: bool | None = None,
output_hidden_states: bool | None = None,
return_dict: bool | None = None,
cache_position: torch.Tensor | None = None,
**kwargs, # NOOP kwargs, for now
) -> tuple[torch.Tensor, ...] | BaseModelOutputWithPastAndCrossAttentions:
**kwargs: Unpack[TransformersKwargs],
) -> BaseModelOutputWithPastAndCrossAttentions:
r"""
input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
`input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()`
Expand All @@ -282,28 +287,8 @@ def forward(

[What are input IDs?](../glossary#input-ids)
"""
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
output_hidden_states = (
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
)
use_cache = use_cache if use_cache is not None else self.config.use_cache
return_dict = return_dict if return_dict is not None else self.config.use_return_dict

if input_ids is not None and inputs_embeds is not None:
raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
elif input_ids is not None:
batch_size, seq_length = input_ids.shape
elif inputs_embeds is not None:
batch_size, seq_length, _ = inputs_embeds.shape
else:
raise ValueError("You have to specify either input_ids or inputs_embeds")

if self.gradient_checkpointing and self.training:
if use_cache:
logger.warning_once(
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
)
use_cache = False
if (input_ids is None) ^ (inputs_embeds is not None):
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

if inputs_embeds is None:
inputs_embeds = self.wte(input_ids)
Expand All @@ -313,12 +298,10 @@ def forward(

hidden_states = inputs_embeds

all_self_attentions = () if output_attentions else None
all_hidden_states = () if output_hidden_states else None

# Compute alibi tensor: check build_alibi_tensor documentation
alibi = self.build_mpt_alibi_tensor(self.num_heads, self.config.max_seq_len, device=hidden_states.device)

seq_length = inputs_embeds.shape[1]
past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
if cache_position is None:
cache_position = torch.arange(
Expand All @@ -334,39 +317,21 @@ def forward(
).to(torch.bool)

for block in self.blocks:
if output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)

outputs = block(
hidden_states = block(
hidden_states,
layer_past=past_key_values,
attention_mask=causal_mask,
use_cache=use_cache,
output_attentions=output_attentions,
position_bias=alibi,
cache_position=cache_position,
)

hidden_states = outputs[0]
if output_attentions:
all_self_attentions = all_self_attentions + (outputs[1],)

# Add last hidden state
hidden_states = self.norm_f(hidden_states)

if output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)

if not return_dict:
return tuple(
v for v in [hidden_states, past_key_values, all_hidden_states, all_self_attentions] if v is not None
)

return BaseModelOutputWithPastAndCrossAttentions(
last_hidden_state=hidden_states,
past_key_values=past_key_values,
hidden_states=all_hidden_states,
attentions=all_self_attentions,
)


Expand All @@ -390,6 +355,7 @@ def __init__(self, config: MptConfig):
def set_output_embeddings(self, new_embeddings: torch.Tensor):
self.lm_head = new_embeddings

@can_return_tuple
@auto_docstring
def forward(
self,
Expand All @@ -399,13 +365,10 @@ def forward(
inputs_embeds: torch.Tensor | None = None,
labels: torch.Tensor | None = None,
use_cache: bool | None = None,
output_attentions: bool | None = None,
output_hidden_states: bool | None = None,
return_dict: bool | None = None,
cache_position: torch.Tensor | None = None,
logits_to_keep: int | torch.Tensor = 0,
**kwargs,
) -> tuple[torch.Tensor] | CausalLMOutputWithCrossAttentions:
) -> CausalLMOutputWithCrossAttentions:
r"""
input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
`input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()`
Expand All @@ -423,21 +386,17 @@ def forward(
`labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
"""
return_dict = return_dict if return_dict is not None else self.config.use_return_dict

transformer_outputs = self.transformer(
transformer_outputs: BaseModelOutputWithPastAndCrossAttentions = self.transformer(
input_ids,
past_key_values=past_key_values,
attention_mask=attention_mask,
inputs_embeds=inputs_embeds,
use_cache=use_cache,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
cache_position=cache_position,
**kwargs,
)

hidden_states = transformer_outputs[0]
hidden_states = transformer_outputs.last_hidden_state
# Only compute necessary logits, and do not upcast them to float if we are not computing the loss
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
logits = self.lm_head(hidden_states[:, slice_indices, :])
Expand All @@ -446,10 +405,6 @@ def forward(
if labels is not None:
loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)

if not return_dict:
output = (logits,) + transformer_outputs[1:]
return ((loss,) + output) if loss is not None else output

return CausalLMOutputWithCrossAttentions(
loss=loss,
logits=logits,
Expand Down Expand Up @@ -486,6 +441,7 @@ def __init__(self, config: MptConfig):
def set_output_embeddings(self, new_embeddings: torch.Tensor):
self.score = new_embeddings

@can_return_tuple
@auto_docstring
def forward(
self,
Expand All @@ -495,11 +451,8 @@ def forward(
inputs_embeds: torch.Tensor | None = None,
labels: torch.Tensor | None = None,
use_cache: bool | None = None,
output_attentions: bool | None = None,
output_hidden_states: bool | None = None,
return_dict: bool | None = None,
**kwargs,
) -> tuple[torch.Tensor] | SequenceClassifierOutputWithPast:
) -> SequenceClassifierOutputWithPast:
r"""
input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
`input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()`
Expand All @@ -517,20 +470,16 @@ def forward(
config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
"""
return_dict = return_dict if return_dict is not None else self.config.use_return_dict

transformer_outputs = self.transformer(
transformer_outputs: BaseModelOutputWithPastAndCrossAttentions = self.transformer(
input_ids,
past_key_values=past_key_values,
attention_mask=attention_mask,
inputs_embeds=inputs_embeds,
use_cache=use_cache,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
**kwargs,
)

hidden_states = transformer_outputs[0]
hidden_states = transformer_outputs.last_hidden_state
logits = self.score(hidden_states)

if input_ids is not None:
Expand Down Expand Up @@ -578,9 +527,6 @@ def forward(
elif self.config.problem_type == "multi_label_classification":
loss_fct = BCEWithLogitsLoss()
loss = loss_fct(pooled_logits, labels)
if not return_dict:
output = (pooled_logits,) + transformer_outputs[1:]
return ((loss,) + output) if loss is not None else output

return SequenceClassifierOutputWithPast(
loss=loss,
Expand Down Expand Up @@ -610,6 +556,7 @@ def __init__(self, config: MptConfig):
# Initialize weights and apply final processing
self.post_init()

@can_return_tuple
@auto_docstring
def forward(
self,
Expand All @@ -619,11 +566,8 @@ def forward(
inputs_embeds: torch.Tensor | None = None,
labels: torch.Tensor | None = None,
use_cache: bool | None = None,
output_attentions: bool | None = None,
output_hidden_states: bool | None = None,
return_dict: bool | None = None,
**deprecated_arguments,
) -> tuple[torch.Tensor] | TokenClassifierOutput:
**kwargs,
) -> TokenClassifierOutput:
r"""
input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
`input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()`
Expand All @@ -641,20 +585,16 @@ def forward(
config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
"""
return_dict = return_dict if return_dict is not None else self.config.use_return_dict

transformer_outputs = self.transformer(
transformer_outputs: BaseModelOutputWithPastAndCrossAttentions = self.transformer(
input_ids,
past_key_values=past_key_values,
attention_mask=attention_mask,
inputs_embeds=inputs_embeds,
use_cache=use_cache,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
**kwargs,
)

hidden_states = transformer_outputs[0]
hidden_states = transformer_outputs.last_hidden_state
hidden_states = self.dropout(hidden_states)
logits = self.classifier(hidden_states)

Expand All @@ -668,10 +608,6 @@ def forward(
logits.view(batch_size * seq_length, self.num_labels), labels.view(batch_size * seq_length)
)

if not return_dict:
output = (logits,) + transformer_outputs[2:]
return ((loss,) + output) if loss is not None else output

return TokenClassifierOutput(
loss=loss,
logits=logits,
Expand All @@ -690,6 +626,7 @@ def __init__(self, config):
# Initialize weights and apply final processing
self.post_init()

@can_return_tuple
@auto_docstring
def forward(
self,
Expand All @@ -698,11 +635,8 @@ def forward(
inputs_embeds: torch.FloatTensor | None = None,
start_positions: torch.LongTensor | None = None,
end_positions: torch.LongTensor | None = None,
output_attentions: bool | None = None,
output_hidden_states: bool | None = None,
return_dict: bool | None = None,
**kwargs,
) -> tuple | QuestionAnsweringModelOutput:
) -> QuestionAnsweringModelOutput:
r"""
input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
`input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()`
Expand All @@ -716,18 +650,14 @@ def forward(

[What are input IDs?](../glossary#input-ids)
"""
return_dict = return_dict if return_dict is not None else self.config.use_return_dict

outputs = self.transformer(
outputs: BaseModelOutputWithPastAndCrossAttentions = self.transformer(
input_ids,
attention_mask=attention_mask,
inputs_embeds=inputs_embeds,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
**kwargs,
)

sequence_output = outputs[0]
sequence_output = outputs.last_hidden_state

logits = self.qa_outputs(sequence_output)
start_logits, end_logits = logits.split(1, dim=-1)
Expand All @@ -751,10 +681,6 @@ def forward(
end_loss = loss_fct(end_logits, end_positions)
total_loss = (start_loss + end_loss) / 2

if not return_dict:
output = (start_logits, end_logits) + outputs[2:]
return ((total_loss,) + output) if total_loss is not None else output

return QuestionAnsweringModelOutput(
loss=total_loss,
start_logits=start_logits,
Expand Down