Repository: mosaicml/llm-foundry

Issue:
MPT models on the Hub not working with `transformers` main
Hi there! 

Currently with transformers main loading MPT models from the Hub fails because it tries to import some private method (such as `_expand_mask` ) that has been recently removed: https://github.com/huggingface/transformers/pull/27086

The simple loading script below should work

```python
from accelerate import init_empty_weights
from transformers import AutoModelForCausalLM, AutoConfig

model_id = "mosaicml/mpt-7b"
config = AutoConfig.from_pretrained(
    model_id, trust_remote_code=True
)
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(
        config, trust_remote_code=True
    )
```



Discussion:
Thanks for letting us know Younes, will look into this ASAP
Thanks @dakinggg !


Resolving patch:
diff --git a/llmfoundry/models/utils/hf_prefixlm_converter.py b/llmfoundry/models/utils/hf_prefixlm_converter.py
--- a/llmfoundry/models/utils/hf_prefixlm_converter.py
+++ b/llmfoundry/models/utils/hf_prefixlm_converter.py
@@ -10,31 +10,14 @@
 and treat the input prompt as the prefix in `generate`.
 """
 
-import math
-import warnings
 from types import MethodType
 from typing import Any, List, MutableMapping, Optional, Tuple, Union
 
 import torch
-from transformers.models.bloom.modeling_bloom import (
-    BaseModelOutputWithPastAndCrossAttentions, BloomForCausalLM, BloomModel,
-    CausalLMOutputWithCrossAttentions, CrossEntropyLoss)
-from transformers.models.bloom.modeling_bloom import \
-    _expand_mask as _expand_mask_bloom
-from transformers.models.bloom.modeling_bloom import \
-    _make_causal_mask as _make_causal_mask_bloom
-from transformers.models.bloom.modeling_bloom import logging
 from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
 from transformers.models.gpt_neo.modeling_gpt_neo import GPTNeoForCausalLM
 from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXForCausalLM
 from transformers.models.gptj.modeling_gptj import GPTJForCausalLM
-from transformers.models.opt.modeling_opt import OPTForCausalLM
-from transformers.models.opt.modeling_opt import \
-    _expand_mask as _expand_mask_opt
-from transformers.models.opt.modeling_opt import \
-    _make_causal_mask as _make_causal_mask_opt
-
-logger = logging.get_logger(__name__)
 
 _SUPPORTED_GPT_MODELS = (
     GPT2LMHeadModel,
@@ -223,583 +206,10 @@ def generate(self: CAUSAL_GPT_TYPES, *args: Any, **kwargs: Any):
     return model
 
 
-def _convert_bloom_causal_lm_to_prefix_lm(
-        model: BloomForCausalLM) -> BloomForCausalLM:
-    """Converts a BLOOM Causal LM to a Prefix LM.
-
-    Supported HuggingFace model classes:
-        - `BloomForCausalLM`
-
-    See `convert_hf_causal_lm_to_prefix_lm` for more details.
-    """
-    if hasattr(model, '_prefix_lm_converted'):
-        return model
-
-    assert isinstance(model, BloomForCausalLM)
-    assert model.config.add_cross_attention == False, 'Only supports BLOOM decoder-only models'
-
-    # Modified from transformers.models.bloom.modeling_bloom.BloomModel._prepare_attn_mask
-    # https://github.com/huggingface/transformers/blob/v4.25.1/src/transformers/models/bloom/modeling_bloom.py#L648
-    def _prepare_attn_mask(
-        self: BloomModel,
-        attention_mask: torch.Tensor,
-        bidirectional_mask: Optional[torch.Tensor],
-        input_shape: Tuple[int, int],
-        past_key_values_length: int,
-    ) -> torch.BoolTensor:
-        # create causal mask
-        # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]
-        combined_attention_mask = None
-        device = attention_mask.device
-        _, src_length = input_shape
-
-        if src_length > 1:
-            combined_attention_mask = _make_causal_mask_bloom(
-                input_shape,
-                device=device,
-                past_key_values_length=past_key_values_length)
-            # Make use of the batch-specific `bidirectional_mask` attribute set
-            # by the parent module in its (new) `forward` method wrapper
-            if bidirectional_mask is not None:
-                # The two masks should have the same size
-                assert attention_mask.shape == bidirectional_mask.shape
-
-                # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]
-                expanded_bidirectional_mask = _expand_mask_bloom(
-                    bidirectional_mask, tgt_length=src_length)
-                combined_attention_mask = torch.logical_and(
-                    combined_attention_mask, expanded_bidirectional_mask)
-
-        # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]
-        expanded_attn_mask = _expand_mask_bloom(attention_mask,
-                                                tgt_length=src_length)
-        combined_attention_mask = (expanded_attn_mask
-                                   if combined_attention_mask is None else
-                                   expanded_attn_mask | combined_attention_mask)
-
-        return combined_attention_mask
-
-    # Modified from transformers.models.bloom.modeling_bloom._prepare_alibi_transformer
-    # https://github.com/huggingface/transformers/blob/v4.25.1/src/transformers/models/bloom/modeling_bloom.py#L87
-    def _build_alibi_tensor(
-        self: BloomModel,
-        batch_size: int,
-        query_length: int,
-        key_length: int,
-        dtype: torch.dtype,
-        device: torch.device,
-    ) -> torch.Tensor:
-        num_heads = self.config.n_head
-
-        closest_power_of_2 = 2**math.floor(math.log2(num_heads))
-        base = torch.tensor(2**(-(2**-(math.log2(closest_power_of_2) - 3))),
-                            device=device,
-                            dtype=torch.float32)
-        powers = torch.arange(1,
-                              1 + closest_power_of_2,
-                              device=device,
-                              dtype=torch.int32)
-        slopes = torch.pow(base, powers)
-
-        if closest_power_of_2 != num_heads:
-            extra_base = torch.tensor(
-                2**(-(2**-(math.log2(2 * closest_power_of_2) - 3))),
-                device=device,
-                dtype=torch.float32)
-            num_remaining_heads = min(closest_power_of_2,
-                                      num_heads - closest_power_of_2)
-            extra_powers = torch.arange(1,
-                                        1 + 2 * num_remaining_heads,
-                                        2,
-                                        device=device,
-                                        dtype=torch.int32)
-            slopes = torch.cat(
-                [slopes, torch.pow(extra_base, extra_powers)], dim=0)
-
-        qa = torch.arange(query_length, device=device,
-                          dtype=torch.int32).view(-1, 1)
-        ka = torch.arange(key_length, device=device,
-                          dtype=torch.int32).view(1, -1)
-        diffs = qa - ka + key_length - query_length
-        diffs = -diffs.abs()
-        alibi = slopes.view(1, num_heads, 1, 1) * diffs.view(
-            1, 1, query_length, key_length)
-        alibi = alibi.expand(batch_size, -1, -1,
-                             -1).reshape(-1, query_length, key_length)
-        return alibi.to(dtype)
-
-    # Modified from transformers.models.bloom.modeling_bloom.BloomModel.forward
-    # Note: The modified code is surrounded with #### START/END #### comments
-    # and one new argument (`bidirectional_mask`) is added to the signature.
-    KeyValueT = Tuple[torch.Tensor, torch.Tensor]
-
-    def transformer_forward(
-        self: BloomModel,
-        input_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[Tuple[KeyValueT, ...]] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        bidirectional_mask: Optional[torch.Tensor] = None,
-        head_mask: Optional[torch.LongTensor] = None,
-        inputs_embeds: Optional[torch.LongTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        **deprecated_arguments: Any
-    ) -> Union[Tuple[torch.Tensor, ...],
-               BaseModelOutputWithPastAndCrossAttentions]:
-        if deprecated_arguments.pop('position_ids', False) is not False:
-            # `position_ids` could have been `torch.Tensor` or `None` so
-            # defaulting pop to `False` allows to detect if users were
-            # passing explicitly `None`
-            warnings.warn(
-                '`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. ' +\
-                'You can safely ignore passing `position_ids`.',
-                FutureWarning,
-            )
-        if len(deprecated_arguments) > 0:
-            raise ValueError(
-                f'Got unexpected arguments: {deprecated_arguments}')
-
-        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 past_key_values is None:
-            past_key_values = tuple([None] * len(self.h))  # type: ignore
-
-        # Prepare head mask if needed
-        # 1.0 in head_mask indicate we keep the head
-        # attention_probs has shape batch_size x num_heads x N x N
-        # head_mask has shape n_layer x batch x num_heads x N x N
-        head_mask = self.get_head_mask(head_mask, self.config.n_layer)
-
-        if inputs_embeds is None:
-            inputs_embeds = self.word_embeddings(input_ids)
-
-        hidden_states = self.word_embeddings_layernorm(inputs_embeds)
-
-        presents = () if use_cache else None
-        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
-        seq_length_with_past = seq_length
-        past_key_values_length = 0
-        if past_key_values[0] is not None:  # type: ignore
-            tmp = past_key_values[0][0]  # type: ignore
-            past_key_values_length = tmp.shape[2]  # type: ignore
-            seq_length_with_past = seq_length_with_past + past_key_values_length
-        if attention_mask is None:
-            attention_mask = torch.ones((batch_size, seq_length_with_past),
-                                        device=hidden_states.device)
-        else:
-            attention_mask = attention_mask.to(hidden_states.device)
-
-        ##### ALL NON-SIGNATURE MODIFICATIONS ARE CONTAINED TO THIS BLOCK [STARTS HERE] #####
-        alibi = self._build_alibi_tensor(
-            batch_size=batch_size,
-            query_length=seq_length,
-            key_length=seq_length_with_past,
-            dtype=hidden_states.dtype,
-            device=hidden_states.device,
-        )
-
-        causal_mask = self._prepare_attn_mask(
-            attention_mask,
-            bidirectional_mask,
-            input_shape=(batch_size, seq_length),
-            past_key_values_length=past_key_values_length,
-        )
-        ##### ALL NON-SIGNATURE MODIFICATIONS ARE CONTAINED TO THIS BLOCK [ENDS HERE] #####
-
-        for i, (block,
-                layer_past) in enumerate(zip(self.h,
-                                             past_key_values)):  # type: ignore
-
-            if output_hidden_states:
-                hst = (hidden_states,)
-                all_hidden_states = all_hidden_states + hst  # type: ignore
-
-            if self.gradient_checkpointing and self.training:
-
-                if use_cache:
-                    logger.warning(
-                        '`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...'
-                    )
-                    use_cache = False
-
-                def create_custom_forward(module: torch.nn.Module):
-
-                    def custom_forward(*inputs: Any):
-                        # None for past_key_value
-                        return module(*inputs,
-                                      use_cache=use_cache,
-                                      output_attentions=output_attentions)
-
-                    return custom_forward
-
-                outputs = torch.utils.checkpoint.checkpoint(  # type: ignore
-                    create_custom_forward(block),
-                    hidden_states,
-                    alibi,
-                    causal_mask,
-                    head_mask[i],  # type: ignore
-                )
-            else:
-                outputs = block(
-                    hidden_states,
-                    layer_past=layer_past,
-                    attention_mask=causal_mask,
-                    head_mask=head_mask[i],  # type: ignore
-                    use_cache=use_cache,
-                    output_attentions=output_attentions,
-                    alibi=alibi,
-                )
-
-            hidden_states = outputs[0]
-            if use_cache is True:
-                presents = presents + (outputs[1],)  # type: ignore
-
-            if output_attentions:
-                oa = (outputs[2 if use_cache else 1],)  # type: ignore
-                all_self_attentions = all_self_attentions + oa  # type: ignore
-
-        # Add last hidden state
-        hidden_states = self.ln_f(hidden_states)
-
-        if output_hidden_states:
-            hst = (hidden_states,)
-            all_hidden_states = all_hidden_states + hst  # type: ignore
-
-        if not return_dict:
-            return tuple(v for v in [
-                hidden_states, presents, all_hidden_states, all_self_attentions
-            ] if v is not None)
-
-        return BaseModelOutputWithPastAndCrossAttentions(
-            last_hidden_state=hidden_states,
-            past_key_values=presents,
-            hidden_states=all_hidden_states,
-            attentions=all_self_attentions,
-        )
-
-    # Make it so model.transformer has the new helper methods and new
-    # `forward` method
-    setattr(model.transformer, '_prepare_attn_mask',
-            MethodType(_prepare_attn_mask, model.transformer))
-    setattr(model.transformer, '_build_alibi_tensor',
-            MethodType(_build_alibi_tensor, model.transformer))
-    setattr(model.transformer, 'forward',
-            MethodType(transformer_forward, model.transformer))
-
-    # In order to actually use the new argument we've added to
-    # model.transformer, we need to update the parent module's `forward` to
-    # accept/pass the same new argument.
-    # We add 2 lines to handle that change.
-    # Both lines are tagged with "# WE'RE ADDING A NEW ARGUMENT!"
-    KeyValueT = Tuple[torch.Tensor, torch.Tensor]
-
-    def forward(
-        self: BloomForCausalLM,
-        input_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[Tuple[KeyValueT, ...]] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        # WE'RE ADDING A NEW ARGUMENT! (Change 1/2)
-        bidirectional_mask: Optional[torch.Tensor] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        inputs_embeds: Optional[torch.Tensor] = None,
-        labels: Optional[torch.Tensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        **deprecated_arguments: Any,
-    ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:
-        """Replacement forward method for BloomCausalLM."""
-        if deprecated_arguments.pop('position_ids', False) is not False:
-            # `position_ids` could have been `torch.Tensor` or `None` so
-            # defaulting pop to `False` allows to detect if users were passing
-            # explicitly `None`
-            warnings.warn(
-                '`position_ids` have no functionality in BLOOM and will be removed ' +\
-                'in v5.0.0. You can safely ignore passing `position_ids`.',
-                FutureWarning,
-            )
-        if len(deprecated_arguments) > 0:
-            raise ValueError(
-                f'Got unexpected arguments: {deprecated_arguments}')
-
-        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,
-            # WE'RE ADDING A NEW ARGUMENT! (Change 2/2)
-            bidirectional_mask=bidirectional_mask,
-            head_mask=head_mask,
-            inputs_embeds=inputs_embeds,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-        hidden_states = transformer_outputs[0]
-
-        lm_logits = self.lm_head(hidden_states)
-
-        loss = None
-        if labels is not None:
-            # Shift so that tokens < n predict n
-            shift_logits = lm_logits[..., :-1, :].contiguous()
-            shift_labels = labels[..., 1:].contiguous()
-            batch_size, seq_length, vocab_size = shift_logits.shape
-            # Flatten the tokens
-            loss_fct = CrossEntropyLoss()
-            loss = loss_fct(
-                shift_logits.view(batch_size * seq_length, vocab_size),
-                shift_labels.view(batch_size * seq_length))
-
-        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,
-            past_key_values=transformer_outputs.past_key_values,
-            hidden_states=transformer_outputs.hidden_states,
-            attentions=transformer_outputs.attentions,
-        )
-
-    # To handle generation, re-write `prepare_inputs_for_generation` to
-    # implement the bidirectional logic.
-    def prepare_inputs_for_generation(self: BloomForCausalLM,
-                                      input_ids: torch.LongTensor,
-                                      past: Optional[torch.Tensor] = None,
-                                      attention_mask: Optional[
-                                          torch.Tensor] = None,
-                                      **kwargs: Any) -> dict:
-        del kwargs  # unused
-        # only last token for input_ids if past is not None
-        if past:
-            input_ids = input_ids[:, -1].unsqueeze(-1)  # type: ignore
-            # We can turn off bidirectional masking after the prefix
-            # has been encoded into `past`
-            bidirectional_mask = None
-
-            # the cache may be in the standard format (e.g. in contrastive
-            # search), convert to bloom's format if needed
-            if past[0][0].shape[0] == input_ids.shape[0]:
-                past = self._convert_to_bloom_cache(past)
-
-        else:
-            # If we're here, `input_ids` contains the prefix. Encode it with
-            # bidirectional attention.
-            bidirectional_mask = torch.ones_like(input_ids)
-
-        return {
-            'input_ids': input_ids,
-            'past_key_values': past,
-            # "use_cache": kwargs.get("use_cache"),
-            # Requires this. TODO(Alex): Confirm this supports other decoding strategies.
-            'use_cache': True,
-            'attention_mask': attention_mask,
-            'bidirectional_mask': bidirectional_mask,
-        }
-
-    # Register the new `forward` and `prepare_inputs_for_generation` methods
-    # with the model
-    setattr(model, 'forward', MethodType(forward, model))
-    setattr(model, 'prepare_inputs_for_generation',
-            MethodType(prepare_inputs_for_generation, model))
-
-    # Finally, tag the model so that this conversion cannot happen again.
-    setattr(model, '_prefix_lm_converted', True)
-    return model
-
-
-def _convert_opt_causal_lm_to_prefix_lm(
-        model: OPTForCausalLM) -> OPTForCausalLM:
-    """Converts an OPT Causal LM to a Prefix LM.
-
-    Supported HuggingFace model classes:
-        - `OPTForCausalLM`
-
-    See `convert_hf_causal_lm_to_prefix_lm` for more details.
-    """
-    if hasattr(model, '_prefix_lm_converted'):
-        return model
-
-    assert isinstance(model, OPTForCausalLM)
-    assert model.config.add_cross_attention == False, 'Only supports OPT decoder-only models'
-
-    # Rename methods to allow:
-    #  - new `forward` to wrap original `forward`
-    #  - new `generate` to wrap original `generate`
-    setattr(model, '_original_forward', getattr(model, 'forward'))
-    setattr(model, '_original_generate', getattr(model, 'generate'))
-
-    model.model.decoder.bidirectional_mask = None
-
-    # Modified from transformers.models.bloom.modeling_opt.OPTDecoder._prepare_decoder_attn_mask
-    # https://github.com/huggingface/transformers/blob/v4.25.1/src/transformers/models/opt/modeling_opt.py#L532
-    def _prepare_decoder_attention_mask(self: torch.nn.Module,
-                                        attention_mask: Optional[torch.Tensor],
-                                        input_shape: Tuple[int, int],
-                                        inputs_embeds: Optional[torch.Tensor],
-                                        past_key_values_length: int):
-        # create causal mask
-        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
-        combined_attention_mask = None
-        if input_shape[-1] > 1:
-            assert inputs_embeds is not None
-            # 'g' indicates generation mode. Causal mask replaced with 0.
-            if self.bidirectional_mask == 'g':
-                bsz, src_length = input_shape
-                combined_attention_mask = torch.zeros(
-                    (bsz, 1, src_length, src_length + past_key_values_length),
-                    dtype=inputs_embeds.dtype,
-                    device=inputs_embeds.device)
-            else:
-                combined_attention_mask = _make_causal_mask_opt(
-                    input_shape,
-                    inputs_embeds.dtype,
-                    past_key_values_length=past_key_values_length).to(
-                        inputs_embeds.device)
-
-                # Make use of the batch-specific `bidirectional_mask` attribute
-                # set by the parent module in its (new) `forward` method wrapper
-                if self.bidirectional_mask is not None:
-                    assert attention_mask is not None
-                    # The two masks should have the same size
-                    assert attention_mask.shape == self.bidirectional_mask.shape
-
-                    # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]
-                    expanded_bidirectional_mask = _expand_mask_opt(
-                        self.bidirectional_mask,
-                        inputs_embeds.dtype,
-                        tgt_len=input_shape[-1]).to(inputs_embeds.device)
-                    combined_attention_mask = torch.maximum(
-                        expanded_bidirectional_mask, combined_attention_mask)
-
-        if attention_mask is not None:
-            assert inputs_embeds is not None
-            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
-            expanded_attn_mask = _expand_mask_opt(attention_mask,
-                                                  inputs_embeds.dtype,
-                                                  tgt_len=input_shape[-1]).to(
-                                                      inputs_embeds.device)
-            combined_attention_mask = (expanded_attn_mask
-                                       if combined_attention_mask is None else
-                                       expanded_attn_mask +
-                                       combined_attention_mask)
-
-        return combined_attention_mask
-
-    # Make it so model.model.decoder uses the above `_prepare_decoder_attn_mask`
-    # in place of the original method
-    setattr(model.model.decoder, '_prepare_decoder_attention_mask',
-            MethodType(_prepare_decoder_attention_mask, model.model.decoder))
-
-    def forward(
-        self: OPTForCausalLM,
-        input_ids: Optional[torch.LongTensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        bidirectional_mask: Optional[torch.ByteTensor] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        labels: Optional[torch.LongTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-    ):
-
-        def call_og_forward():
-            return self._original_forward(
-                input_ids=input_ids,
-                attention_mask=attention_mask,
-                head_mask=head_mask,
-                past_key_values=past_key_values,
-                inputs_embeds=inputs_embeds,
-                labels=labels,
-                use_cache=use_cache,
-                output_attentions=output_attentions,
-                output_hidden_states=output_hidden_states,
-                return_dict=return_dict,
-            )
-
-        if bidirectional_mask is None:
-            # This wrapper is a no-op if bidirectional masks are not supplied
-            return call_og_forward()
-
-        # Temporarily set `bidirectional_mask` in the child module
-        self.model.decoder.bidirectional_mask = bidirectional_mask
-
-        # Apply the original forward method (the model will use the mask that
-        # was just set)
-        try:
-            outputs = call_og_forward()
-        except:
-            self.model.decoder.bidirectional_mask = None
-            raise
-
-        # Reset the `bidirectional_mask` attribute to None
-        self.model.decoder.bidirectional_mask = None
-
-        # Return the outputs
-        return outputs
-
-    def generate(self: OPTForCausalLM, *args: tuple, **kwargs: Any):
-        """Wraps original generate to enable PrefixLM-style attention."""
-        # Flag the child module to use generation-style attention masking
-        self.model.decoder.bidirectional_mask = 'g'
-
-        # Collect outputs using the model's original forward method
-        try:
-            output = self._original_generate(*args, **kwargs)
-        except:
-            self.model.decoder.bidirectional_mask = None
-            raise
-
-        # Reset the `bidirectional_mask` attribute to None
-        self.model.decoder.bidirectional_mask = None
-
-        # Return the output
-        return output
-
-    # Replace `forward` and `generate` with the new wrappers
-    setattr(model, 'forward', MethodType(forward, model))
-    setattr(model, 'generate', MethodType(generate, model))
-
-    # Finally, tag the model so that this conversion cannot happen again.
-    setattr(model, '_prefix_lm_converted', True)
-    return model
-
-
-_SUPPORTED_HF_MODELS = _SUPPORTED_GPT_MODELS + (BloomForCausalLM,
-                                                OPTForCausalLM)
+_SUPPORTED_HF_MODELS = _SUPPORTED_GPT_MODELS
 
 CAUSAL_LM_TYPES = Union[GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM,
-                        GPTNeoXForCausalLM, BloomForCausalLM, OPTForCausalLM]
+                        GPTNeoXForCausalLM]
 
 
 def convert_hf_causal_lm_to_prefix_lm(
@@ -811,8 +221,6 @@ def convert_hf_causal_lm_to_prefix_lm(
         - `GPTNeoForCausalLM`
         - `GPTNeoXForCausalLM`
         - `GPTJForCausalLM`
-        - `BloomForCausalLM`
-        - `OPTForCausalLM`
 
     Conversion to a Prefix LM is done by modifying the `forward` method, and possibly also the
     `generate` method and/or select underlying methods depending on the model class.
@@ -862,13 +270,6 @@ def convert_hf_causal_lm_to_prefix_lm(
     """
     if isinstance(model, _SUPPORTED_GPT_MODELS):
         return _convert_gpt_causal_lm_to_prefix_lm(model)
-
-    elif isinstance(model, BloomForCausalLM):
-        return _convert_bloom_causal_lm_to_prefix_lm(model)
-
-    elif isinstance(model, OPTForCausalLM):
-        return _convert_opt_causal_lm_to_prefix_lm(model)
-
     else:
         raise TypeError(
             f'Cannot convert model to Prefix LM. ' +\


Test patch:
diff --git a/tests/test_model.py b/tests/test_model.py
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -25,8 +25,7 @@
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.models.bloom.modeling_bloom import build_alibi_tensor
 
-from llmfoundry import (COMPOSER_MODEL_REGISTRY, ComposerHFCausalLM,
-                        ComposerHFPrefixLM)
+from llmfoundry import COMPOSER_MODEL_REGISTRY, ComposerHFCausalLM
 from llmfoundry.models.hf.model_wrapper import HuggingFaceModelWithZLoss
 from llmfoundry.models.layers import NORM_CLASS_REGISTRY, build_alibi_bias
 from llmfoundry.models.layers.blocks import MPTBlock
@@ -438,11 +437,10 @@ def test_loss_fn():
                                     atol=1e-4), f'differed at step {i}'
 
 
-@pytest.mark.parametrize('prefixlm', [False, True])
-def test_opt_wrapping(prefixlm: bool):
+def test_opt_wrapping():
     conf = {
         'model': {
-            'name': 'hf_prefix_lm' if prefixlm else 'hf_causal_lm',
+            'name': 'hf_causal_lm',
             'pretrained_model_name_or_path': 'facebook/opt-125m',
             'pretrained': 'false'
         },
@@ -456,10 +454,7 @@ def test_opt_wrapping(prefixlm: bool):
     tokenizer = build_tokenizer(config.tokenizer.name,
                                 tokenizer_cfg.get('kwargs', {}))
 
-    if prefixlm:
-        model = ComposerHFPrefixLM(config.model, tokenizer)
-    else:
-        model = ComposerHFCausalLM(config.model, tokenizer)
+    model = ComposerHFCausalLM(config.model, tokenizer)
 
     # check that all the modules we except are blocked from FSDP wrapping
     assert not model.model.model._fsdp_wrap