[Deps] Upgrade to transformers 4.56.x (#587)

* [Deps] Upgrade to transformers 4.56.x

* remove triton conv1d backend for mamba2

* skip mom test

Author: zhiyuan1i

.github/workflows/reusable-ci-tests.yml

@@ -137,13 +137,13 @@ jobs:
               NIGHTLY_URL="https://download.pytorch.org/whl/nightly/${{ inputs.pytorch_cuda_version }}"
               echo "Using nightly index URL: $NIGHTLY_URL"
               $CONDA_BIN_PATH/pip install -U torch pytorch-triton --index-url $NIGHTLY_URL
-              $CONDA_BIN_PATH/pip install -U numpy packaging psutil ninja einops datasets
+              $CONDA_BIN_PATH/pip install -U numpy packaging psutil ninja einops datasets transformers
               $CONDA_BIN_PATH/pip install --no-deps .
             else
               STABLE_URL="https://download.pytorch.org/whl/${{ inputs.pytorch_cuda_version }}"
               echo "Using stable index URL: $STABLE_URL"
               $CONDA_BIN_PATH/pip install -U torch~=${{ inputs.pytorch_version }} triton --index-url $STABLE_URL
-              $CONDA_BIN_PATH/pip install -U numpy packaging psutil ninja einops datasets
+              $CONDA_BIN_PATH/pip install -U numpy packaging psutil ninja einops datasets transformers
               $CONDA_BIN_PATH/pip install .
               if [[ "${{ inputs.runner }}" == nvidia-h100* ]]; then
                 echo "Installing causal-conv1d for H100"
@@ -156,7 +156,7 @@ jobs:
             XPU_URL="https://download.pytorch.org/whl/xpu"
             echo "Using XPU index URL: $XPU_URL"
             $CONDA_BIN_PATH/pip install -U torch~=${{ inputs.pytorch_version }} pytorch-triton-xpu --index-url $XPU_URL
-            $CONDA_BIN_PATH/pip install -U numpy packaging psutil ninja einops datasets
+            $CONDA_BIN_PATH/pip install -U numpy packaging psutil ninja einops datasets transformers
             $CONDA_BIN_PATH/pip install .
           else
             echo "::error::Unsupported GPU type: ${{ inputs.gpu_type }}"
@@ -319,7 +319,7 @@ jobs:
               NIGHTLY_URL="https://download.pytorch.org/whl/nightly/${{ inputs.pytorch_cuda_version }}"
               echo "Using nightly index URL: $NIGHTLY_URL"
               $CONDA_BIN_PATH/pip install -U torch pytorch-triton --index-url $NIGHTLY_URL
-              $CONDA_BIN_PATH/pip install -U numpy packaging psutil ninja einops datasets
+              $CONDA_BIN_PATH/pip install -U numpy packaging psutil ninja einops datasets transformers
               $CONDA_BIN_PATH/pip install --no-deps .
             else
               STABLE_URL="https://download.pytorch.org/whl/${{ inputs.pytorch_cuda_version }}"

fla/layers/mamba2.py

@@ -211,6 +211,10 @@ def __init__(
             from fla.modules.convolution import causal_conv1d_update as causal_conv1d_update_triton
             self.causal_conv1d_fn = causal_conv1d_triton
             self.causal_conv1d_update = causal_conv1d_update_triton
+            logger.warning(
+                "Mamba2 does not recommend using Triton's conv1d backend, "
+                "as it is untested and may contain bugs."
+            )
         else:
             self.causal_conv1d_fn = causal_conv1d_fn
             self.causal_conv1d_update = causal_conv1d_update

fla/models/abc/modeling_abc.py

@@ -8,7 +8,6 @@
 
 import torch
 import torch.nn as nn
-from transformers.generation import GenerationMixin
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import logging
@@ -17,7 +16,7 @@
 from fla.layers.abc import ABCAttention
 from fla.layers.attn import Attention
 from fla.models.abc.configuration_abc import ABCConfig
-from fla.models.utils import Cache
+from fla.models.utils import Cache, FLAGenerationMixin
 from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss
 from fla.modules import GatedMLP as ABCMLP
 from fla.modules import RMSNorm
@@ -259,7 +258,7 @@ def forward(
         )
 
 
-class ABCForCausalLM(ABCPreTrainedModel, GenerationMixin):
+class ABCForCausalLM(ABCPreTrainedModel, FLAGenerationMixin):
 
     _tied_weights_keys = ["lm_head.weight"]
 
@@ -306,40 +305,6 @@ def generate(self, *args, **kwargs):
             else:
                 raise exception
 
-    @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-    def prepare_inputs_for_generation(
-        self,
-        input_ids: torch.LongTensor = None,
-        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        use_cache: bool = True,
-        logits_to_keep: Optional[int] = None,
-        **kwargs
-    ):
-        # only last token for `inputs_ids` if the `past_key_values` is not empty.
-        if past_key_values is not None and len(past_key_values) > 0:
-            input_ids = input_ids[:, -1:]
-        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
-        if inputs_embeds is not None and len(past_key_values) == 0:
-            model_inputs = {'inputs_embeds': inputs_embeds}
-        else:
-            # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
-            # recompiles graphs as the stride of the inputs is a guard.
-            # Ref: https://github.com/huggingface/transformers/pull/29114
-            # TODO: use `next_tokens` directly instead.
-            model_inputs = {'input_ids': input_ids.contiguous()}
-
-        if logits_to_keep is not None:
-            model_inputs['logits_to_keep'] = logits_to_keep
-
-        model_inputs.update({
-            'past_key_values': past_key_values,
-            'use_cache': use_cache,
-            'attention_mask': attention_mask,
-        })
-        return model_inputs
-
     @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
     def forward(
         self,

fla/models/bitnet/modeling_bitnet.py

@@ -8,15 +8,14 @@
 
 import torch
 import torch.nn as nn
-from transformers.generation import GenerationMixin
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import logging
 from transformers.utils.deprecation import deprecate_kwarg
 
 from fla.layers.bitattn import BitAttention
 from fla.models.bitnet.configuration_bitnet import BitNetConfig
-from fla.models.utils import Cache
+from fla.models.utils import Cache, FLAGenerationMixin
 from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm
 from fla.modules.activations import swiglu
 from fla.modules.fused_bitlinear import FusedBitLinear
@@ -296,7 +295,7 @@ def forward(
         )
 
 
-class BitNetForCausalLM(BitNetPreTrainedModel, GenerationMixin):
+class BitNetForCausalLM(BitNetPreTrainedModel, FLAGenerationMixin):
 
     _tied_weights_keys = ["lm_head.weight"]
 
@@ -328,40 +327,6 @@ def set_decoder(self, decoder):
     def get_decoder(self):
         return self.model
 
-    @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-    def prepare_inputs_for_generation(
-        self,
-        input_ids: torch.LongTensor = None,
-        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        use_cache: bool = True,
-        logits_to_keep: Optional[int] = None,
-        **kwargs
-    ):
-        # only last token for `inputs_ids` if the `past_key_values` is not empty.
-        if past_key_values is not None and len(past_key_values) > 0:
-            input_ids = input_ids[:, -1:]
-        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
-        if inputs_embeds is not None and len(past_key_values) == 0:
-            model_inputs = {'inputs_embeds': inputs_embeds}
-        else:
-            # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
-            # recompiles graphs as the stride of the inputs is a guard.
-            # Ref: https://github.com/huggingface/transformers/pull/29114
-            # TODO: use `next_tokens` directly instead.
-            model_inputs = {'input_ids': input_ids.contiguous()}
-
-        if logits_to_keep is not None:
-            model_inputs['logits_to_keep'] = logits_to_keep
-
-        model_inputs.update({
-            'past_key_values': past_key_values,
-            'use_cache': use_cache,
-            'attention_mask': attention_mask,
-        })
-        return model_inputs
-
     @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
     def forward(
         self,

fla/models/comba/modeling_comba.py

@@ -8,7 +8,6 @@
 
 import torch
 import torch.nn as nn
-from transformers.generation import GenerationMixin
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import logging
@@ -17,7 +16,7 @@
 from fla.layers.attn import Attention
 from fla.layers.comba import Comba
 from fla.models.comba.configuration_comba import CombaConfig
-from fla.models.utils import Cache
+from fla.models.utils import Cache, FLAGenerationMixin
 from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss
 from fla.modules import GatedMLP as CombaMLP
 from fla.modules import RMSNorm
@@ -266,7 +265,7 @@ def forward(
         )
 
 
-class CombaForCausalLM(CombaPreTrainedModel, GenerationMixin):
+class CombaForCausalLM(CombaPreTrainedModel, FLAGenerationMixin):
 
     _tied_weights_keys = ["lm_head.weight"]
 
@@ -313,40 +312,6 @@ def generate(self, *args, **kwargs):
             else:
                 raise exception
 
-    @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-    def prepare_inputs_for_generation(
-        self,
-        input_ids: torch.LongTensor = None,
-        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        use_cache: bool = True,
-        logits_to_keep: Optional[int] = None,
-        **kwargs
-    ):
-        # only last token for `inputs_ids` if the `past_key_values` is not empty.
-        if past_key_values is not None and len(past_key_values) > 0:
-            input_ids = input_ids[:, -1:]
-        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
-        if inputs_embeds is not None and len(past_key_values) == 0:
-            model_inputs = {'inputs_embeds': inputs_embeds}
-        else:
-            # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
-            # recompiles graphs as the stride of the inputs is a guard.
-            # Ref: https://github.com/huggingface/transformers/pull/29114
-            # TODO: use `next_tokens` directly instead.
-            model_inputs = {'input_ids': input_ids.contiguous()}
-
-        if logits_to_keep is not None:
-            model_inputs['logits_to_keep'] = logits_to_keep
-
-        model_inputs.update({
-            'past_key_values': past_key_values,
-            'use_cache': use_cache,
-            'attention_mask': attention_mask,
-        })
-        return model_inputs
-
     @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
     def forward(
         self,

fla/models/delta_net/modeling_delta_net.py

@@ -8,7 +8,6 @@
 
 import torch
 import torch.nn as nn
-from transformers.generation import GenerationMixin
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import logging
@@ -17,7 +16,7 @@
 from fla.layers.attn import Attention
 from fla.layers.delta_net import DeltaNet
 from fla.models.delta_net.configuration_delta_net import DeltaNetConfig
-from fla.models.utils import Cache
+from fla.models.utils import Cache, FLAGenerationMixin
 from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss
 from fla.modules import GatedMLP as DeltaNetMLP
 from fla.modules import RMSNorm
@@ -256,7 +255,7 @@ def forward(
         )
 
 
-class DeltaNetForCausalLM(DeltaNetPreTrainedModel, GenerationMixin):
+class DeltaNetForCausalLM(DeltaNetPreTrainedModel, FLAGenerationMixin):
 
     _tied_weights_keys = ["lm_head.weight"]
 
@@ -303,40 +302,6 @@ def generate(self, *args, **kwargs):
             else:
                 raise exception
 
-    @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-    def prepare_inputs_for_generation(
-        self,
-        input_ids: torch.LongTensor = None,
-        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        use_cache: bool = True,
-        logits_to_keep: Optional[int] = None,
-        **kwargs
-    ):
-        # only last token for `inputs_ids` if the `past_key_values` is not empty.
-        if past_key_values is not None and len(past_key_values) > 0:
-            input_ids = input_ids[:, -1:]
-        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
-        if inputs_embeds is not None and len(past_key_values) == 0:
-            model_inputs = {'inputs_embeds': inputs_embeds}
-        else:
-            # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
-            # recompiles graphs as the stride of the inputs is a guard.
-            # Ref: https://github.com/huggingface/transformers/pull/29114
-            # TODO: use `next_tokens` directly instead.
-            model_inputs = {'input_ids': input_ids.contiguous()}
-
-        if logits_to_keep is not None:
-            model_inputs['logits_to_keep'] = logits_to_keep
-
-        model_inputs.update({
-            'past_key_values': past_key_values,
-            'use_cache': use_cache,
-            'attention_mask': attention_mask,
-        })
-        return model_inputs
-
     @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
     def forward(
         self,

fla/models/forgetting_transformer/modeling_forgetting_transformer.py

@@ -8,15 +8,14 @@
 
 import torch
 import torch.nn as nn
-from transformers.generation import GenerationMixin
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import logging
 from transformers.utils.deprecation import deprecate_kwarg
 
 from fla.layers.forgetting_attn import ForgettingAttention
 from fla.models.forgetting_transformer.configuration_forgetting_transformer import ForgettingTransformerConfig
-from fla.models.utils import Cache
+from fla.models.utils import Cache, FLAGenerationMixin
 from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss
 from fla.modules import GatedMLP as ForgettingTransformerMLP
 from fla.modules import RMSNorm
@@ -260,7 +259,7 @@ def forward(
         )
 
 
-class ForgettingTransformerForCausalLM(ForgettingTransformerPreTrainedModel, GenerationMixin):
+class ForgettingTransformerForCausalLM(ForgettingTransformerPreTrainedModel, FLAGenerationMixin):
 
     _tied_weights_keys = ["lm_head.weight"]
 
@@ -292,40 +291,6 @@ def set_decoder(self, decoder):
     def get_decoder(self):
         return self.model
 
-    @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-    def prepare_inputs_for_generation(
-        self,
-        input_ids: torch.LongTensor = None,
-        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        use_cache: bool = True,
-        logits_to_keep: Optional[int] = None,
-        **kwargs
-    ):
-        # only last token for `inputs_ids` if the `past_key_values` is not empty.
-        if past_key_values is not None and len(past_key_values) > 0:
-            input_ids = input_ids[:, -1:]
-        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
-        if inputs_embeds is not None and len(past_key_values) == 0:
-            model_inputs = {'inputs_embeds': inputs_embeds}
-        else:
-            # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
-            # recompiles graphs as the stride of the inputs is a guard.
-            # Ref: https://github.com/huggingface/transformers/pull/29114
-            # TODO: use `next_tokens` directly instead.
-            model_inputs = {'input_ids': input_ids.contiguous()}
-
-        if logits_to_keep is not None:
-            model_inputs['logits_to_keep'] = logits_to_keep
-
-        model_inputs.update({
-            'past_key_values': past_key_values,
-            'use_cache': use_cache,
-            'attention_mask': attention_mask,
-        })
-        return model_inputs
-
     @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
     def forward(
         self,