Skip to content

Refactoring falcon model to match standardized output collection interface #43995

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

Open
akeemlh wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Refactoring falcon model to match standardized output collection interface #43995

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
108 changes: 21 additions & 87 deletions src/transformers/models/falcon/modeling_falcon.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,11 +41,15 @@
dynamic_rope_update,
)
from ...modeling_utils import PreTrainedModel
from ...processing_utils import Unpack
from ...utils import (
TransformersKwargs,
auto_docstring,
can_return_tuple,
logging,
)
from ...utils.generic import maybe_autocast
from ...utils.output_capturing import capture_outputs
from .configuration_falcon import FalconConfig


Expand Down Expand Up @@ -321,9 +325,9 @@ def forward(
position_ids: torch.LongTensor | None = None,
layer_past: Cache | None = None,
use_cache: bool = False,
output_attentions: bool = False,
cache_position: torch.LongTensor | None = None,
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
**kwargs: Unpack[TransformersKwargs],
):
fused_qkv = self.query_key_value(hidden_states) # [batch_size, seq_length, 3 x hidden_size]
num_kv_heads = self.num_heads if self.new_decoder_architecture else self.num_kv_heads
Expand All @@ -349,7 +353,7 @@ def forward(
kv_length = key_layer.shape[-2]

if alibi is None:
if self.config._attn_implementation == "sdpa" and not output_attentions:
if self.config._attn_implementation == "sdpa":
# We dispatch to SDPA's Flash Attention or Efficient kernels via this if statement instead of an
# inline conditional assignment to support both torch.compile's `dynamic=True` and `fullgraph=True`
# The query_length > 1 is necessary to match with a bidirectional attention mask we do not have
Expand Down Expand Up @@ -381,7 +385,7 @@ def forward(
return attn_output, attention_scores

else:
if self.config._attn_implementation == "sdpa" and not output_attentions:
if self.config._attn_implementation == "sdpa":
# We dispatch to SDPA's Flash Attention or Efficient kernels via this if statement instead of an
# inline conditional assignment to support both torch.compile's `dynamic=True` and `fullgraph=True`
is_causal = self.is_causal and attention_mask is None and query_length > 1
Expand Down Expand Up @@ -453,7 +457,6 @@ def forward(
position_ids: torch.LongTensor | None = None,
layer_past: Cache | None = None,
use_cache: bool = False,
output_attentions: bool = False,
cache_position: torch.LongTensor | None = None,
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
):
Expand Down Expand Up @@ -528,9 +531,6 @@ def forward(
attn_weights = attn_output.reshape(batch_size, query_length, self.num_heads * self.head_dim)
attn_output = self.dense(attn_weights)

if not output_attentions:
attn_weights = None

return attn_output, attn_weights


Expand Down Expand Up @@ -591,10 +591,9 @@ def forward(
position_ids: torch.LongTensor | None = None,
layer_past: Cache | tuple[torch.Tensor, torch.Tensor] | None = None,
use_cache: bool = False,
output_attentions: bool = False,
cache_position: torch.LongTensor | None = None,
position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
**kwargs,
**kwargs: Unpack[TransformersKwargs],
):
residual = hidden_states

Expand All @@ -612,7 +611,6 @@ def forward(
position_ids=position_ids,
alibi=alibi,
use_cache=use_cache,
output_attentions=output_attentions,
cache_position=cache_position,
position_embeddings=position_embeddings,
)
Expand Down Expand Up @@ -641,7 +639,7 @@ def forward(

output = dropout_add(mlp_output, residual, self.config.hidden_dropout, training=self.training)

return output, attn_weights
return output


@auto_docstring
Expand All @@ -653,6 +651,7 @@ class FalconPreTrainedModel(PreTrainedModel):
_supports_flash_attn = True
_supports_sdpa = True
_can_compile_fullgraph = True
_can_record_outputs = {"hidden_states": FalconDecoderLayer, "attentions": FalconAttention}

@torch.no_grad()
def _init_weights(self, module: nn.Module):
Expand Down Expand Up @@ -704,6 +703,7 @@ def get_input_embeddings(self):
def set_input_embeddings(self, new_embeddings: torch.Tensor):
self.word_embeddings = new_embeddings

@capture_outputs
@auto_docstring
def forward(
self,
Expand All @@ -713,9 +713,6 @@ def forward(
position_ids: torch.LongTensor | 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.LongTensor | None = None,
**kwargs,
) -> tuple[torch.Tensor, ...] | BaseModelOutputWithPastAndCrossAttentions:
Expand All @@ -732,12 +729,7 @@ 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 None) ^ (inputs_embeds is not None):
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
Expand Down Expand Up @@ -806,45 +798,25 @@ def forward(
hidden_states = inputs_embeds
position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)

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

for i, block in enumerate(self.h):
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,
position_ids=position_ids,
use_cache=use_cache,
output_attentions=output_attentions,
alibi=alibi,
cache_position=cache_position,
position_embeddings=position_embeddings,
**kwargs,
)

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

# Add last hidden state
hidden_states = self.ln_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 @@ -867,6 +839,7 @@ def __init__(self, config: FalconConfig):
def set_output_embeddings(self, new_embeddings: torch.Tensor):
self.lm_head = new_embeddings

@can_return_tuple
@auto_docstring
def forward(
self,
Expand All @@ -877,9 +850,6 @@ 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.LongTensor | None = None,
logits_to_keep: int | torch.Tensor = 0,
**kwargs,
Expand All @@ -902,19 +872,15 @@ def forward(
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(
input_ids,
past_key_values=past_key_values,
attention_mask=attention_mask,
position_ids=position_ids,
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]

Expand All @@ -930,10 +896,6 @@ def forward(
**kwargs,
)

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

return CausalLMOutputWithCrossAttentions(
loss=loss,
logits=lm_logits,
Expand Down Expand Up @@ -967,6 +929,7 @@ def __init__(self, config: FalconConfig):
# Initialize weights and apply final processing
self.post_init()

@can_return_tuple
@auto_docstring
def forward(
self,
Expand All @@ -976,9 +939,6 @@ 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:
r"""
Expand All @@ -999,17 +959,13 @@ def forward(
`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(
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]
Expand Down Expand Up @@ -1060,9 +1016,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 @@ -1092,6 +1045,7 @@ def __init__(self, config: FalconConfig):
# Initialize weights and apply final processing
self.post_init()

@can_return_tuple
@auto_docstring
def forward(
self,
Expand All @@ -1101,9 +1055,6 @@ 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] | TokenClassifierOutput:
r"""
Expand All @@ -1124,17 +1075,13 @@ def forward(
`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(
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]
Expand All @@ -1149,10 +1096,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 @@ -1171,6 +1114,7 @@ def __init__(self, config):
# Initialize weights and apply final processing
self.post_init()

@can_return_tuple
@auto_docstring
def forward(
self,
Expand All @@ -1179,9 +1123,6 @@ 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:
r"""
Expand All @@ -1197,15 +1138,12 @@ 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(
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]
Expand All @@ -1232,10 +1170,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