Amplify Troubleshooting

Readme.md

Setup

Repo and Rust

# setup rust
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
apt-get update && apt-get install -y libssl-dev pkg-config  libasound2-dev libudev-dev libwayland-dev libxkbcommon-dev fontconfig libfontconfig1-dev git-lfs


git clone [email protected]:zachcp/ferritin.git
git clone [email protected]:zachcp/AMPLIFY.git
git clone https://huggingface.co/chandar-lab/AMPLIFY_120M


cd ferritin && cargo test

Amplify Pyenv

cd AMPLIFY
export REPO_DIR="/workspace/candletestAMPLIFY"
python3 -m venv env && \
source env/bin/activate && \
python3 -m pip install --upgrade pip && \
#python3 -m pip install --editable $REPO_DIR[dev]
python3 -m pip install --editable .[dev]

pip install transformers datasets xformers biopython ipython importlib safetensors jupyter
python3 -m pytest

DevNotes

# use a local copy of Amplify what we can modify
wget https://huggingface.co/chandar-lab/AMPLIFY_120M/resolve/main/model.safetensors?download=true

# cargo test
cargo test -p ferritin-featurizers test_amplify_round_trip -- --exact

# python test
# in the Amplify env above.
python3 run.py

# config_path = "./conf/model/120M.yaml"
# checkpoint_file = "/workspace/candletest/models/AMPLIFY120M_model.safetensors"
# model, tokenizer = amplify.AMPLIFY.load(checkpoint_file, config_path)
# predictor = amplify.inference.Predictor(model, tokenizer, device=device)
# sequence = "MSVVGIDLGFQSCYVAVARAGGIETIANEYSDRCTPACISFGPKNR"
# logits = predictor.logits(sequence)
# embedder = (model, tokenizer)
# sequence_embedding = predictor.embed(sequence)

amplify.py

# From https://stackoverflow.com/a/23689767
# From https://github.com/pytorch/pytorch/issues/97899
# From https://github.com/facebookresearch/llama/blob/main/llama/model.py

import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

import torch
from torch import nn
from torch.nn.functional import scaled_dot_product_attention
from xformers.ops import SwiGLU, memory_efficient_attention

from .rmsnorm import RMSNorm
from .rotary import precompute_freqs_cis, apply_rotary_emb

from transformers import PreTrainedModel, PretrainedConfig
from transformers.modeling_outputs import MaskedLMOutput


class DotDict(dict):
    """Dictionary that supports the dot notation to access attributes (similarly to HuggingFace)."""

    __getattr__ = dict.get
    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__


class AMPLIFYConfig(PretrainedConfig):
    model_type = "AMPLIFY"

    # All config parameters must have a default value.
    def __init__(
        self,
        hidden_size: int = 960,
        num_hidden_layers: int = 32,
        num_attention_heads: int = 15,
        intermediate_size: int = 3840,
        dropout_prob: float = 0,
        embedding_init_range: float = 0.02,
        decoder_init_range: float = 0.02,
        rms_norm: bool = True,
        norm_eps: float = 1e-05,
        hidden_act: str = "SwiGLU",
        layer_norm_after_embedding: bool = False,
        layer_norm_before_last_layer: bool = True,
        vocab_size: int = 27,
        ffn_bias: bool = False,
        att_bias: bool = False,
        pad_token_id: int = 0,
        max_length: int = 2048,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.dropout_prob = dropout_prob
        self.embedding_init_range = embedding_init_range
        self.decoder_init_range = decoder_init_range
        self.rms_norm = rms_norm
        self.norm_eps = norm_eps
        self.hidden_act = hidden_act
        self.layer_norm_after_embedding = layer_norm_after_embedding
        self.layer_norm_before_last_layer = layer_norm_before_last_layer
        self.vocab_size = vocab_size
        self.ffn_bias = ffn_bias
        self.att_bias = att_bias
        self.pad_token_id = pad_token_id
        self.max_length = max_length


class EncoderBlock(nn.Module):
    """Transformer encoder block."""

    def __init__(self, config: AMPLIFYConfig):
        """Initialize a EncoderBlock.

        Args:
            hidden_size (int): _description_
            num_attention_heads (int): _description_
            intermediate_size (int, optional): _description_. Defaults to 2048.
            dropout_prob (float, optional): _description_. Defaults to 0.1.
            activation (str, optional): _description_. Defaults to "relu".
            rms_norm (bool, optional): _description_. Defaults to True.
            norm_eps (float, optional): _description_. Defaults to 1e-5.
            pad_token_id (int, optional): _description_. Defaults to 0.
            max_length (int, optional): _description_. Defaults to 2048.
            ffn_bias (bool, optional): _description_. Defaults to False.
            att_bias (bool, optional): _description_. Defaults to False.
        """
        super().__init__()

        self.config = config
        self.d_head = config.hidden_size // config.num_attention_heads

        # Attention
        self.q = nn.Linear(in_features=config.hidden_size, out_features=config.hidden_size, bias=config.att_bias)
        self.k = nn.Linear(in_features=config.hidden_size, out_features=config.hidden_size, bias=config.att_bias)
        self.v = nn.Linear(in_features=config.hidden_size, out_features=config.hidden_size, bias=config.att_bias)
        self.wo = nn.Linear(in_features=config.hidden_size, out_features=config.hidden_size, bias=config.att_bias)
        self.resid_dropout = nn.Dropout(config.dropout_prob)

        # Feedforward network
        match config.hidden_act.lower():
            case "swiglu":
                # To keep the number of parameters and the amount of computation constant, we reduce the number of
                # hidden units by a factor of 2/3 (https://arxiv.org/pdf/2002.05202.pdf) and make it a multiple of 8 to
                # avoid RuntimeError due to misaligned operand
                multiple_of = 8
                intermediate_size = int(2 * config.intermediate_size / 3)
                intermediate_size = multiple_of * ((intermediate_size + multiple_of - 1) // multiple_of)
                self.ffn = SwiGLU(config.hidden_size, intermediate_size, config.hidden_size, bias=config.ffn_bias)
            case "relu":
                self.ffn = nn.Sequential(
                    nn.Linear(config.hidden_size, config.intermediate_size, bias=config.ffn_bias),
                    nn.ReLU(),
                    nn.Linear(config.intermediate_size, config.hidden_size, bias=config.ffn_bias),
                )
            case "gelu":
                self.ffn = nn.Sequential(
                    nn.Linear(config.hidden_size, config.intermediate_size, bias=config.ffn_bias),
                    nn.GELU(),
                    nn.Linear(config.intermediate_size, config.hidden_size, bias=config.ffn_bias),
                )

        self.attention_norm = (
            RMSNorm(config.hidden_size, config.norm_eps)
            if config.rms_norm
            else nn.LayerNorm(config.hidden_size, config.norm_eps)
        )
        self.ffn_norm = (
            RMSNorm(config.hidden_size, config.norm_eps)
            if config.rms_norm
            else nn.LayerNorm(config.hidden_size, config.norm_eps)
        )

        self.ffn_dropout = nn.Dropout(config.dropout_prob)

    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor, freqs_cis: torch.Tensor, output_attentions: bool):
        # attention mask is None
        print(f"EncoderBlock. x: {x.shape}, freqs_cis: {freqs_cis.shape}")
        attn, contact = self._att_block(self.attention_norm(x), attention_mask, freqs_cis, output_attentions)
        print(f"ATT_BLOCK_OUT. attn: {attn.shape},  contact: {contact}")
        x = x + attn
        ffn_norm = self.ffn_norm(x)
        print(f"FFN_NORM. ffn_norm: {ffn_norm.shape}")
        x = x + self._ff_block(ffn_norm)
        print(f"FFN_FINAL.: {x.shape}")
        return x, contact

    def _att_block(
        self, x: torch.Tensor, attention_mask: torch.Tensor, freqs_cis: torch.Tensor, output_attentions: bool
    ):
        batch_size, seq_len, _ = x.shape

        xq, xk, xv = self.q(x), self.k(x), self.v(x)
        print(f"EncoderBlock_ATT. xq: {xq.shape}, xk: {xk.shape}, xv: {xv.shape}")

        # Reshape for rotary embeddings
        xq = xq.view(batch_size, seq_len, self.config.num_attention_heads, self.d_head)
        xk = xk.view(batch_size, seq_len, self.config.num_attention_heads, self.d_head)
        xv = xv.view(batch_size, seq_len, self.config.num_attention_heads, self.d_head)
        print(f"EncoderBlock_ATT_reshaped. xq: {xq.shape}, xk: {xk.shape}, xv: {xv.shape}")

        xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
        print(f"EncoderBlock_after_rotary. xq: {xq.shape}, xk: {xk.shape}")


        # Compute the attention weight
        attn_weights = None
        if output_attentions:
            attn_weights = xq.permute(0, 2, 1, 3) @ xk.permute(0, 2, 3, 1) / (xq.size(-1) ** 0.5)
            if attention_mask is not None:
                attn_weights = attn_weights + attention_mask
            attn_weights = attn_weights.softmax(-1)

        print(f"attn_weights: {attn_weights}")

        # Compute the attention using xformers if the tensors are on GPU
        if x.is_cuda:
            # Input and output are of dimension (B, M, H, K) where B is the batch size, M the sequence length,
            # H the number of heads, and K the embeding size per head
            attn = memory_efficient_attention(
                query=xq,
                key=xk,
                value=xv,
                attn_bias=attention_mask,
                p=self.config.dropout_prob if self.training else 0,
            )
        else:
            # Input and output are of dimension (B, H, M, K)
            print(f"ATTN CALC IN: xq: {xq.shape}")

            xq_permute = xq.transpose(1, 2)
            print(f"ATTN CALC IN:  xq_permute:{xq_permute.shape}")

            attn = scaled_dot_product_attention(
                query=xq.transpose(1, 2),
                key=xk.transpose(1, 2),
                value=xv.transpose(1, 2),
                attn_mask=attention_mask,
                dropout_p=self.config.dropout_prob if self.training else 0,
            )
            print(f"ATTN CALC: {attn.shape}")
            attn = attn.transpose(1, 2)
            print(f"ATTN CALC TRANSPOSE: {attn.shape}")


        print(f"attention: {attn.shape}")
        att_view = attn.view(batch_size, seq_len, self.config.num_attention_heads * self.d_head)
        print(f"attention_view: {att_view.shape}")
        attn_scores = self.wo(att_view)
        print(f"attn_scores: {attn_scores.shape}")
        return (self.resid_dropout(attn_scores), attn_weights)

    def _ff_block(self, x: torch.Tensor):
        print(f"FFN_BLOCK: x:{x.shape}")
        ffn = self.ffn(x)
        print(f"FFN_BLOCK_FFN: x:{ffn.shape}")
        ffn_drop = self.ffn_dropout(ffn)
        print(f"FFN_BLOCK_DROP: x:{ffn_drop.shape}")
        return ffn_drop


class AMPLIFYPreTrainedModel(PreTrainedModel):
    config_class = AMPLIFYConfig

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            module.weight.data.uniform_(-self.config.decoder_init_range, self.config.decoder_init_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.uniform_(-self.config.embedding_init_range, self.config.embedding_init_range)


class AMPLIFY(AMPLIFYPreTrainedModel):
    """The main model class.

    Args:
       config (amplify.model.amplify.AMPLIFYConfig): model configuration, usually defined from the Hydra configuration.
    """

    def __init__(self, config: AMPLIFYConfig, **kwargs):
        super().__init__(config)

        self.config = config

        self.encoder = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)

        if config.layer_norm_after_embedding:
            self.layer_norm_1 = (
                RMSNorm(config.hidden_size, config.norm_eps)
                if config.rms_norm
                else nn.LayerNorm(config.hidden_size, config.norm_eps)
            )

        self.transformer_encoder = nn.ModuleList()
        for _ in range(config.num_hidden_layers):
            self.transformer_encoder.append(EncoderBlock(config))

        if config.layer_norm_before_last_layer:
            self.layer_norm_2 = (
                RMSNorm(config.hidden_size, config.norm_eps)
                if config.rms_norm
                else nn.LayerNorm(config.hidden_size, config.norm_eps)
            )

        self.decoder = nn.Linear(config.hidden_size, config.vocab_size)

        self.freqs_cis = precompute_freqs_cis(config.hidden_size // config.num_attention_heads, config.max_length)

        # Initialize weights and apply final processing
        self.post_init()

    def forward(self, input_ids, attention_mask=None, output_hidden_states=False, output_attentions=False, **kwargs):
        print("AMPLIFY Forward: Initialize.");
        # Initialize
        hidden_states, attentions = [], []

        # Expand and repeat: (Batch, Length) -> (Batch, Heads, Length, Length)
        if attention_mask is not None and not torch.all(attention_mask == 0):
            attention_mask = (
                attention_mask.unsqueeze(1)
                .unsqueeze(1)
                .repeat(1, self.config.num_attention_heads, attention_mask.size(-1), 1)
            )
        else:
            attention_mask = None

        # RoPE
        self.freqs_cis = self.freqs_cis.to(input_ids.device, non_blocking=True)
        print(f"Freqs CIS: {self.freqs_cis.shape}")
        freqs_cis = self.freqs_cis[: input_ids.shape[1]]
        print(f"Freqs CIS sliced: {freqs_cis.shape}")

        # Embedding
        x = self.encoder(input_ids)
        print(f"Encoded: {x.shape}")
        if self.config.layer_norm_after_embedding:
            x = self.layer_norm_1(x)

        # Transformer encoder
        for i,layer in enumerate(self.transformer_encoder):
            x, attn = layer(x, attention_mask, freqs_cis, output_attentions)
            print(f"Transformer Layer {i}; X dim :{x.shape}; attn: {attn}")
            if output_hidden_states:
                hidden_states.append(x)
            if output_attentions:
                attentions.append(attn)

        # Classification head with layer norm
        logits = self.decoder(self.layer_norm_2(x) if self.config.layer_norm_before_last_layer else x)

        # Return logits or the output of the last hidden layer
        return MaskedLMOutput(logits=logits, hidden_states=hidden_states, attentions=attentions)

rotary.py

import torch
import importlib
from transformers import AutoModel
from transformers import AutoTokenizer
import transformers.models.auto.modeling_auto


torch.set_default_device('cpu')

model = AutoModel.from_pretrained(
    "/workspace/candletest/AMPLIFY_120M",
    force_reload=True,
    trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(
    "/workspace/candletest/AMPLIFY_120M",
    force_reload=True,
    trust_remote_code=True)

# Move the model to GPU (required due to Flash Attention)
# model = model.to("cuda")
model = model.to("cpu")


sequence = "MSVVGIDLGFQSCYVAVARAGGIETIANEYSDRCTPACISFGPKNR"

input = tokenizer.encode(sequence, return_tensors="pt")
# print("Input: ", input)

#input = input.to("cuda")
input = input.to("cpu")

output = model(input)
# print("Output: ", output)
predictions = output.logits.argmax(dim=-1)  # Get the most likely token IDs
# print("Output shape: ", predictions.shape)
predictions = predictions.cpu().tolist()  # Convert to list for the tokenizer
print(tokenizer.decode(predictions[0]).replace(" ",""))  # Decode the first