diff --git a/src/transformers/models/dbrx/modeling_dbrx.py b/src/transformers/models/dbrx/modeling_dbrx.py index f94d35ba510f..0db15aa71499 100644 --- a/src/transformers/models/dbrx/modeling_dbrx.py +++ b/src/transformers/models/dbrx/modeling_dbrx.py @@ -697,7 +697,6 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_router_logits: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -725,10 +724,6 @@ def forward( "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ``` """ - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.transformer( input_ids=input_ids, @@ -737,8 +732,8 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -752,7 +747,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/dbrx/modular_dbrx.py b/src/transformers/models/dbrx/modular_dbrx.py index 9e2d11334489..9df4262dd305 100644 --- a/src/transformers/models/dbrx/modular_dbrx.py +++ b/src/transformers/models/dbrx/modular_dbrx.py @@ -485,7 +485,6 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_router_logits: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -513,10 +512,6 @@ def forward( "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ``` """ - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.transformer( input_ids=input_ids, @@ -525,8 +520,8 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -540,7 +535,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/doge/modeling_doge.py b/src/transformers/models/doge/modeling_doge.py index cc32bf7a0e49..9e805ba354c4 100644 --- a/src/transformers/models/doge/modeling_doge.py +++ b/src/transformers/models/doge/modeling_doge.py @@ -755,7 +755,6 @@ def forward( use_cache: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, - output_router_logits: bool | None = None, **kwargs: Unpack[TransformersKwargs], ) -> MoeCausalLMOutputWithPast: r""" @@ -780,10 +779,6 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -793,6 +788,7 @@ def forward( inputs_embeds=inputs_embeds, use_cache=use_cache, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -806,7 +802,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if self.config.is_moe and outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/doge/modular_doge.py b/src/transformers/models/doge/modular_doge.py index 02a86ad31579..2f2371fb277a 100644 --- a/src/transformers/models/doge/modular_doge.py +++ b/src/transformers/models/doge/modular_doge.py @@ -678,7 +678,6 @@ def forward( use_cache: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, - output_router_logits: bool | None = None, **kwargs: Unpack[TransformersKwargs], ) -> MoeCausalLMOutputWithPast: r""" @@ -703,10 +702,6 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -716,6 +711,7 @@ def forward( inputs_embeds=inputs_embeds, use_cache=use_cache, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -729,7 +725,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if self.config.is_moe and outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py b/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py index b3e545776e02..0962d5824a0c 100644 --- a/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py +++ b/src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py @@ -692,7 +692,6 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_router_logits: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -703,11 +702,6 @@ def forward( config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. """ - - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -716,8 +710,8 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -731,7 +725,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/ernie4_5_vl_moe/modeling_ernie4_5_vl_moe.py b/src/transformers/models/ernie4_5_vl_moe/modeling_ernie4_5_vl_moe.py index a0198139bc02..3c23fbee2370 100644 --- a/src/transformers/models/ernie4_5_vl_moe/modeling_ernie4_5_vl_moe.py +++ b/src/transformers/models/ernie4_5_vl_moe/modeling_ernie4_5_vl_moe.py @@ -1594,7 +1594,6 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_router_logits: bool | None = None, pixel_values: torch.Tensor | None = None, pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, @@ -1620,10 +1619,6 @@ def forward( rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): The rope index difference between sequence length and multimodal rope. """ - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.text_config.output_router_logits - ) - outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, @@ -1633,7 +1628,6 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, return_dict=True, pixel_values=pixel_values, pixel_values_videos=pixel_values_videos, @@ -1654,7 +1648,7 @@ def forward( loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py b/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py index 68a181eccfa3..6d727c45b4ee 100644 --- a/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py +++ b/src/transformers/models/ernie4_5_vl_moe/modular_ernie4_5_vl_moe.py @@ -1400,7 +1400,6 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_router_logits: bool | None = None, pixel_values: torch.Tensor | None = None, pixel_values_videos: torch.FloatTensor | None = None, image_grid_thw: torch.LongTensor | None = None, @@ -1426,10 +1425,6 @@ def forward( rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): The rope index difference between sequence length and multimodal rope. """ - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.text_config.output_router_logits - ) - outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, @@ -1439,7 +1434,6 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, return_dict=True, pixel_values=pixel_values, pixel_values_videos=pixel_values_videos, @@ -1460,7 +1454,7 @@ def forward( loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/flex_olmo/modeling_flex_olmo.py b/src/transformers/models/flex_olmo/modeling_flex_olmo.py index aa6ce9f5723a..95962033e193 100644 --- a/src/transformers/models/flex_olmo/modeling_flex_olmo.py +++ b/src/transformers/models/flex_olmo/modeling_flex_olmo.py @@ -634,7 +634,6 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_router_logits: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -662,11 +661,6 @@ def forward( 'Hey, are you conscious? Can you talk to me?\nI’m not sure if you’re conscious of this, but I’m' ``` """ - - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -675,8 +669,8 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -690,7 +684,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py index 52640e6fff74..db726b1adafe 100644 --- a/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/configuration_glm4v_moe.py @@ -96,6 +96,9 @@ class Glm4vMoeTextConfig(PreTrainedConfig): Beginning of stream token id. router_aux_loss_coef (`float`, *optional*, defaults to 0.0001): The aux loss factor for the loss. + output_router_logits (`bool`, *optional*, defaults to `False`): + Whether or not the router logits should be returned by the model. Enabling this will also + allow the model to output the auxiliary loss. ```python >>> from transformers import Glm4vMoeTextModel, Glm4vMoeConfig @@ -161,6 +164,7 @@ def __init__( eos_token_id: int | None = None, bos_token_id: int | None = None, router_aux_loss_coef: float | None = 0.0001, + output_router_logits: bool | None = False, **kwargs, ): self.pad_token_id = pad_token_id @@ -194,6 +198,7 @@ def __init__( self.first_k_dense_replace = first_k_dense_replace self.norm_topk_prob = norm_topk_prob self.router_aux_loss_coef = router_aux_loss_coef + self.output_router_logits = output_router_logits super().__init__(ignore_keys_at_rope_validation={"mrope_section"}, **kwargs) diff --git a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py index 4821bf829cd1..ced87a4d5b80 100644 --- a/src/transformers/models/glm4v_moe/modular_glm4v_moe.py +++ b/src/transformers/models/glm4v_moe/modular_glm4v_moe.py @@ -131,6 +131,9 @@ class Glm4vMoeTextConfig(Glm4MoeConfig): Beginning of stream token id. router_aux_loss_coef (`float`, *optional*, defaults to 0.0001): The aux loss factor for the loss. + output_router_logits (`bool`, *optional*, defaults to `False`): + Whether or not the router logits should be returned by the model. Enabling this will also + allow the model to output the auxiliary loss. ```python >>> from transformers import Glm4vMoeTextModel, Glm4vMoeConfig @@ -193,6 +196,7 @@ def __init__( eos_token_id: int | None = None, bos_token_id: int | None = None, router_aux_loss_coef: float | None = 0.0001, + output_router_logits: bool | None = False, **kwargs, ): self.pad_token_id = pad_token_id @@ -226,6 +230,7 @@ def __init__( self.first_k_dense_replace = first_k_dense_replace self.norm_topk_prob = norm_topk_prob self.router_aux_loss_coef = router_aux_loss_coef + self.output_router_logits = output_router_logits PreTrainedConfig.__init__(self, ignore_keys_at_rope_validation={"mrope_section"}, **kwargs) diff --git a/src/transformers/models/gpt_oss/configuration_gpt_oss.py b/src/transformers/models/gpt_oss/configuration_gpt_oss.py index a93481fc7d26..e6e33a2940d8 100644 --- a/src/transformers/models/gpt_oss/configuration_gpt_oss.py +++ b/src/transformers/models/gpt_oss/configuration_gpt_oss.py @@ -71,7 +71,7 @@ def __init__( }, attention_dropout: float | None = 0.0, num_experts_per_tok: int | None = 4, - router_aux_loss_coef: float | None = 0.9, + router_aux_loss_coef: float | None = 0.001, output_router_logits: bool | None = False, use_cache: bool | None = True, layer_types: list[str] | None = None, diff --git a/src/transformers/models/gpt_oss/modeling_gpt_oss.py b/src/transformers/models/gpt_oss/modeling_gpt_oss.py index 1b0b8b8ca394..2423b1863054 100644 --- a/src/transformers/models/gpt_oss/modeling_gpt_oss.py +++ b/src/transformers/models/gpt_oss/modeling_gpt_oss.py @@ -625,7 +625,6 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_router_logits: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -652,11 +651,6 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -665,8 +659,8 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -680,7 +674,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/granitemoe/modeling_granitemoe.py b/src/transformers/models/granitemoe/modeling_granitemoe.py index e6d98911f362..0fcbf6f6fdeb 100644 --- a/src/transformers/models/granitemoe/modeling_granitemoe.py +++ b/src/transformers/models/granitemoe/modeling_granitemoe.py @@ -663,7 +663,6 @@ def forward( past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, - output_router_logits: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs, @@ -690,9 +689,6 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs = self.model( input_ids=input_ids, @@ -701,6 +697,7 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -721,7 +718,7 @@ def forward( ) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/granitemoe/modular_granitemoe.py b/src/transformers/models/granitemoe/modular_granitemoe.py index 88c50171096e..96a66da6e986 100644 --- a/src/transformers/models/granitemoe/modular_granitemoe.py +++ b/src/transformers/models/granitemoe/modular_granitemoe.py @@ -245,7 +245,6 @@ def forward( past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, - output_router_logits: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs, @@ -272,9 +271,6 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs = self.model( input_ids=input_ids, @@ -283,6 +279,7 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -303,7 +300,7 @@ def forward( ) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py index 2e1625742cce..8d6583c0745e 100644 --- a/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py +++ b/src/transformers/models/granitemoehybrid/modeling_granitemoehybrid.py @@ -1477,7 +1477,6 @@ def forward( past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, - output_router_logits: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs, @@ -1504,9 +1503,6 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs = self.model( input_ids=input_ids, @@ -1515,6 +1511,7 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -1535,7 +1532,7 @@ def forward( ) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py index 91f6a4ed5158..8989151089c6 100644 --- a/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py +++ b/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py @@ -732,7 +732,6 @@ def forward( past_key_values: Cache | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, - output_router_logits: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs, @@ -759,9 +758,6 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs = self.model( input_ids=input_ids, @@ -770,6 +766,7 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -790,7 +787,7 @@ def forward( ) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/jamba/modeling_jamba.py b/src/transformers/models/jamba/modeling_jamba.py index c3a72cd916fd..d0c230f32574 100755 --- a/src/transformers/models/jamba/modeling_jamba.py +++ b/src/transformers/models/jamba/modeling_jamba.py @@ -1024,11 +1024,6 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -1037,8 +1032,8 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -1052,7 +1047,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/minimax/modeling_minimax.py b/src/transformers/models/minimax/modeling_minimax.py index 028068a78a77..e554d3801b7a 100644 --- a/src/transformers/models/minimax/modeling_minimax.py +++ b/src/transformers/models/minimax/modeling_minimax.py @@ -827,7 +827,6 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_router_logits: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -854,11 +853,6 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -867,8 +861,8 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -882,7 +876,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/minimax_m2/modeling_minimax_m2.py b/src/transformers/models/minimax_m2/modeling_minimax_m2.py index e4e276294324..1cf3691f7b0d 100644 --- a/src/transformers/models/minimax_m2/modeling_minimax_m2.py +++ b/src/transformers/models/minimax_m2/modeling_minimax_m2.py @@ -625,7 +625,6 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_router_logits: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -652,11 +651,6 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -665,8 +659,8 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -680,7 +674,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/mixtral/modeling_mixtral.py b/src/transformers/models/mixtral/modeling_mixtral.py index 859300d4bb69..c5baeda16054 100644 --- a/src/transformers/models/mixtral/modeling_mixtral.py +++ b/src/transformers/models/mixtral/modeling_mixtral.py @@ -617,7 +617,6 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_router_logits: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -644,11 +643,6 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -657,8 +651,8 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -672,7 +666,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/mixtral/modular_mixtral.py b/src/transformers/models/mixtral/modular_mixtral.py index 34996ad8a404..40036c2eee69 100644 --- a/src/transformers/models/mixtral/modular_mixtral.py +++ b/src/transformers/models/mixtral/modular_mixtral.py @@ -360,7 +360,6 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_router_logits: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -387,11 +386,6 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -400,8 +394,8 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -415,7 +409,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/nllb_moe/modeling_nllb_moe.py b/src/transformers/models/nllb_moe/modeling_nllb_moe.py index 313b574518f4..35a6aadd07c0 100644 --- a/src/transformers/models/nllb_moe/modeling_nllb_moe.py +++ b/src/transformers/models/nllb_moe/modeling_nllb_moe.py @@ -1085,13 +1085,9 @@ def forward( decoder_inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_router_logits: bool | None = None, cache_position: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor] | Seq2SeqMoEOutput: - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) if labels is not None: if decoder_input_ids is None: decoder_input_ids = shift_tokens_right( @@ -1108,8 +1104,8 @@ def forward( inputs_embeds=inputs_embeds, decoder_inputs_embeds=decoder_inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, + return_dict=True, **kwargs, ) lm_logits = self.lm_head(outputs[0]) @@ -1122,7 +1118,7 @@ def forward( loss_fct = CrossEntropyLoss(ignore_index=-100) # todo check in the config if router loss enables - if output_router_logits: + if outputs.encoder_router_logits is not None: encoder_router_logits = outputs.encoder_router_logits decoder_router_logits = outputs.decoder_router_logits encoder_aux_loss = load_balancing_loss_func( @@ -1134,7 +1130,7 @@ def forward( loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1)) - if output_router_logits and labels is not None: + if outputs.encoder_router_logits is not None and labels is not None: aux_loss = self.router_aux_loss_coef * (encoder_aux_loss + decoder_aux_loss) loss = loss + aux_loss diff --git a/src/transformers/models/olmoe/modeling_olmoe.py b/src/transformers/models/olmoe/modeling_olmoe.py index ff367f1ab4f7..5ea5e04598fd 100644 --- a/src/transformers/models/olmoe/modeling_olmoe.py +++ b/src/transformers/models/olmoe/modeling_olmoe.py @@ -641,7 +641,6 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_router_logits: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -669,11 +668,6 @@ def forward( 'Hey, are you conscious? Can you talk to me?\nI’m not sure if you’re conscious of this, but I’m' ``` """ - - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -682,8 +676,8 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -697,7 +691,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/phimoe/modeling_phimoe.py b/src/transformers/models/phimoe/modeling_phimoe.py index 8d88f3b041c5..78ff29a075ca 100644 --- a/src/transformers/models/phimoe/modeling_phimoe.py +++ b/src/transformers/models/phimoe/modeling_phimoe.py @@ -809,7 +809,6 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_router_logits: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -836,11 +835,6 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -849,8 +843,8 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -864,7 +858,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py index d53d9ae6ff32..74c8145cfcc3 100644 --- a/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py +++ b/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py @@ -655,7 +655,6 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_router_logits: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -682,11 +681,6 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -695,8 +689,8 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -710,7 +704,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py index d21e12ebebb5..ad7ba0c4a543 100644 --- a/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py +++ b/src/transformers/models/qwen3_5_moe/modeling_qwen3_5_moe.py @@ -1972,11 +1972,6 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -1985,8 +1980,8 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -2000,7 +1995,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py index b07006597a66..3130c39fcced 100644 --- a/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/modeling_qwen3_moe.py @@ -646,7 +646,6 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_router_logits: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -673,11 +672,6 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -686,8 +680,8 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -701,7 +695,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/qwen3_moe/modular_qwen3_moe.py b/src/transformers/models/qwen3_moe/modular_qwen3_moe.py index 4dfe23ebfd8a..ba57386a7802 100644 --- a/src/transformers/models/qwen3_moe/modular_qwen3_moe.py +++ b/src/transformers/models/qwen3_moe/modular_qwen3_moe.py @@ -109,7 +109,6 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, - output_router_logits: bool | None = None, cache_position: torch.LongTensor | None = None, logits_to_keep: int | torch.Tensor = 0, **kwargs: Unpack[TransformersKwargs], @@ -136,11 +135,6 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -149,8 +143,8 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -164,7 +158,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/qwen3_next/modeling_qwen3_next.py b/src/transformers/models/qwen3_next/modeling_qwen3_next.py index 2190ec03f83b..11a6481d4989 100644 --- a/src/transformers/models/qwen3_next/modeling_qwen3_next.py +++ b/src/transformers/models/qwen3_next/modeling_qwen3_next.py @@ -1230,11 +1230,6 @@ def forward( >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." ```""" - - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs: MoeModelOutputWithPast = self.model( input_ids=input_ids, @@ -1243,8 +1238,8 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, + return_dict=True, **kwargs, ) @@ -1258,7 +1253,7 @@ def forward( loss = self.loss_function(logits, labels, self.vocab_size, **kwargs) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py index 13b083a2ab6b..3dcd5be33a68 100644 --- a/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/modeling_qwen3_omni_moe.py @@ -2068,7 +2068,6 @@ def forward( rope_deltas=None, labels=None, use_cache=None, - output_router_logits: bool | None = None, use_audio_in_video=None, cache_position=None, video_second_per_grid=None, @@ -2129,10 +2128,6 @@ def forward( >>> response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] ```""" - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.text_config.output_router_logits - ) - if inputs_embeds is None: # 1. Extract the input embeddings inputs_embeds = self.get_input_embeddings()(input_ids) @@ -2232,10 +2227,10 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, deepstack_visual_embeds=visual_embeds_multiscale, visual_pos_masks=visual_pos_masks, + return_dict=True, **kwargs, ) @@ -2249,7 +2244,7 @@ def forward( ) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py index 1a964b0701d9..30f5b058846f 100644 --- a/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py +++ b/src/transformers/models/qwen3_omni_moe/modular_qwen3_omni_moe.py @@ -1544,16 +1544,11 @@ def forward( rope_deltas=None, labels=None, use_cache=None, - output_router_logits: bool | None = None, use_audio_in_video=None, cache_position=None, video_second_per_grid=None, **kwargs, ) -> tuple | Qwen3OmniMoeThinkerCausalLMOutputWithPast: - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.text_config.output_router_logits - ) - if inputs_embeds is None: # 1. Extract the input embeddings inputs_embeds = self.get_input_embeddings()(input_ids) @@ -1653,10 +1648,10 @@ def forward( past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, - output_router_logits=output_router_logits, cache_position=cache_position, deepstack_visual_embeds=visual_embeds_multiscale, visual_pos_masks=visual_pos_masks, + return_dict=True, **kwargs, ) @@ -1670,7 +1665,7 @@ def forward( ) aux_loss = None - if output_router_logits: + if outputs.router_logits is not None: aux_loss = load_balancing_loss_func( outputs.router_logits, self.num_experts, diff --git a/src/transformers/models/switch_transformers/configuration_switch_transformers.py b/src/transformers/models/switch_transformers/configuration_switch_transformers.py index 909049c32d5d..3ea101e753c5 100644 --- a/src/transformers/models/switch_transformers/configuration_switch_transformers.py +++ b/src/transformers/models/switch_transformers/configuration_switch_transformers.py @@ -91,6 +91,9 @@ class SwitchTransformersConfig(PreTrainedConfig): Whether to output router probabilities to compute router auxiliary loss. use_cache (`bool`, *optional*, defaults to `True`): Whether or not the model should return the last key/values attentions (not used by all models). + output_router_logits (`bool`, *optional*, defaults to `False`): + Whether or not the router logits should be returned by the model. Enabling this will also + allow the model to output the auxiliary loss. """ model_type = "switch_transformers" @@ -131,6 +134,7 @@ def __init__( tie_word_embeddings=True, is_decoder=False, add_cross_attention=False, + output_router_logits=False, **kwargs, ): self.is_decoder = is_decoder @@ -181,6 +185,7 @@ def __init__( self.router_z_loss_coef = router_z_loss_coef self.router_aux_loss_coef = router_aux_loss_coef + self.output_router_logits = output_router_logits self.dense_act_fn = dense_act_fn self.pad_token_id = pad_token_id self.bos_token_id = bos_token_id diff --git a/src/transformers/models/switch_transformers/modeling_switch_transformers.py b/src/transformers/models/switch_transformers/modeling_switch_transformers.py index d0949ff6d389..cf1696815903 100644 --- a/src/transformers/models/switch_transformers/modeling_switch_transformers.py +++ b/src/transformers/models/switch_transformers/modeling_switch_transformers.py @@ -966,7 +966,6 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, decoder_inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, - output_router_logits: bool | None = False, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor] | Seq2SeqMoEOutput: @@ -975,7 +974,6 @@ def forward( input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, - output_router_logits=output_router_logits, **kwargs, ) @@ -994,7 +992,6 @@ def forward( encoder_hidden_states=hidden_states, encoder_attention_mask=attention_mask, cache_position=cache_position, - output_router_logits=output_router_logits, **kwargs, ) @@ -1013,7 +1010,7 @@ def forward( decoder_z_loss = None decoder_aux_loss = None - if output_router_logits: + if encoder_outputs[-1] is not None: # Compute the router loss (z_loss + auxiliary loss) for each router in the encoder and decoder if self.encoder.config.encoder_sparse_step > 1: encoder_router_logits, encoder_expert_indexes = self._unpack_router_logits(encoder_outputs[-1]) @@ -1039,7 +1036,7 @@ def forward( labels = labels.to(lm_logits.device) loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1)) - if output_router_logits: + if encoder_outputs[-1] is not None: z_loss = self.router_z_loss_coef * (encoder_z_loss + decoder_z_loss) aux_loss = self.router_aux_loss_coef * (encoder_aux_loss + decoder_aux_loss) loss = loss + z_loss + aux_loss diff --git a/src/transformers/models/switch_transformers/modular_switch_transformers.py b/src/transformers/models/switch_transformers/modular_switch_transformers.py index dcfe061da793..a433daf38d0c 100644 --- a/src/transformers/models/switch_transformers/modular_switch_transformers.py +++ b/src/transformers/models/switch_transformers/modular_switch_transformers.py @@ -665,7 +665,6 @@ def forward( inputs_embeds: torch.FloatTensor | None = None, decoder_inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, - output_router_logits: bool | None = False, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.FloatTensor] | Seq2SeqMoEOutput: @@ -674,7 +673,6 @@ def forward( input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, - output_router_logits=output_router_logits, **kwargs, ) @@ -693,7 +691,6 @@ def forward( encoder_hidden_states=hidden_states, encoder_attention_mask=attention_mask, cache_position=cache_position, - output_router_logits=output_router_logits, **kwargs, ) @@ -712,7 +709,7 @@ def forward( decoder_z_loss = None decoder_aux_loss = None - if output_router_logits: + if encoder_outputs[-1] is not None: # Compute the router loss (z_loss + auxiliary loss) for each router in the encoder and decoder if self.encoder.config.encoder_sparse_step > 1: encoder_router_logits, encoder_expert_indexes = self._unpack_router_logits(encoder_outputs[-1]) @@ -738,7 +735,7 @@ def forward( labels = labels.to(lm_logits.device) loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1)) - if output_router_logits: + if encoder_outputs[-1] is not None: z_loss = self.router_z_loss_coef * (encoder_z_loss + decoder_z_loss) aux_loss = self.router_aux_loss_coef * (encoder_aux_loss + decoder_aux_loss) loss = loss + z_loss + aux_loss diff --git a/src/transformers/utils/generic.py b/src/transformers/utils/generic.py index 51f30090f752..66769dc3884a 100644 --- a/src/transformers/utils/generic.py +++ b/src/transformers/utils/generic.py @@ -906,6 +906,16 @@ def wrapper(self, *args, **kwargs): self.config.is_causal = is_causal kwargs["is_causal"] = is_causal + # Maybe temporarily enable collecting router logits for aux loss during training + original_output_router_logits = None + if ( + self.training + and (router_aux_loss_coef := getattr(self.config, "router_aux_loss_coef", None)) + and router_aux_loss_coef != 0 + ): + original_output_router_logits = self.config.output_router_logits + self.config.output_router_logits = True + # Call the original forward with the updated kwargs/config try: if kwargs.get("debug_io", False): @@ -923,6 +933,9 @@ def wrapper(self, *args, **kwargs): else: del self.config.is_causal + if original_output_router_logits is not None: + self.config.output_router_logits = original_output_router_logits + return output return wrapper diff --git a/utils/check_config_attributes.py b/utils/check_config_attributes.py index a44cc0d05471..3023a8c903f3 100644 --- a/utils/check_config_attributes.py +++ b/utils/check_config_attributes.py @@ -40,7 +40,6 @@ "BambaConfig": ["attn_layer_indices"], "Dots1Config": ["max_window_layers"], "JambaConfig": ["attn_layer_offset", "attn_layer_period", "expert_layer_offset", "expert_layer_period"], - "JetMoeConfig": ["output_router_logits"], "Phi3Config": ["embd_pdrop"], "EncodecConfig": ["overlap"], "XcodecConfig": ["sample_rate", "audio_channels"], @@ -172,6 +171,8 @@ "vision_feature_layer", "vision_feature_select_strategy", "vision_aspect_ratio", + # moe attributes that are used via merge_with_config_defaults and/or capture hooks + "output_router_logits", )