[FlaxBart] make sure no grads are computed an bias (#16345)

* [FlaxBart] make sure no grads are computed an bias

* correct all other seq2seq models

src/transformers/models/bart/modeling_flax_bart.py

@@ -1292,7 +1292,7 @@ class FlaxBartForConditionalGenerationModule(nn.Module):
         else:
             lm_logits = self.lm_head(hidden_states)
 
-        lm_logits += self.final_logits_bias.astype(self.dtype)
+        lm_logits += jax.lax.stop_gradient(self.final_logits_bias.astype(self.dtype))
 
         if not return_dict:
             output = (lm_logits,) + outputs[1:]

src/transformers/models/blenderbot/modeling_flax_blenderbot.py

@@ -1270,7 +1270,7 @@ class FlaxBlenderbotForConditionalGenerationModule(nn.Module):
         else:
             lm_logits = self.lm_head(hidden_states)
 
-        lm_logits += self.final_logits_bias.astype(self.dtype)
+        lm_logits += jax.lax.stop_gradient(self.final_logits_bias.astype(self.dtype))
 
         if not return_dict:
             output = (lm_logits,) + outputs[1:]

src/transformers/models/blenderbot_small/modeling_flax_blenderbot_small.py

@@ -1267,7 +1267,7 @@ class FlaxBlenderbotSmallForConditionalGenerationModule(nn.Module):
         else:
             lm_logits = self.lm_head(hidden_states)
 
-        lm_logits += self.final_logits_bias.astype(self.dtype)
+        lm_logits += jax.lax.stop_gradient(self.final_logits_bias.astype(self.dtype))
 
         if not return_dict:
             output = (lm_logits,) + outputs[1:]

src/transformers/models/mbart/modeling_flax_mbart.py

@@ -1329,7 +1329,7 @@ class FlaxMBartForConditionalGenerationModule(nn.Module):
         else:
             lm_logits = self.lm_head(hidden_states)
 
-        lm_logits += self.final_logits_bias.astype(self.dtype)
+        lm_logits += jax.lax.stop_gradient(self.final_logits_bias.astype(self.dtype))
 
         if not return_dict:
             output = (lm_logits,) + outputs[1:]

src/transformers/models/pegasus/modeling_flax_pegasus.py

@@ -1280,7 +1280,7 @@ class FlaxPegasusForConditionalGenerationModule(nn.Module):
         else:
             lm_logits = self.lm_head(hidden_states)
 
-        lm_logits += self.final_logits_bias.astype(self.dtype)
+        lm_logits += jax.lax.stop_gradient(self.final_logits_bias.astype(self.dtype))
 
         if not return_dict:
             output = (lm_logits,) + outputs[1:]