Upload folder using huggingface_hub

README.md

@@ -5,24 +5,112 @@ widget:
   - text: MQIFVKTLTGKTITLEVEPS<mask>TIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG
 ---
 
-## ESM-2 (TransformerEngine-optimized)
+> [!NOTE]
+> This model has been optimized using NVIDIA's [TransformerEngine](https://github.com/NVIDIA/TransformerEngine)
+> library. Slight numerical differences may be observed between the original model and the optimized
+> version. For instructions on how to install TransformerEngine, please refer to the
+> [official documentation](https://github.com/NVIDIA/TransformerEngine?tab=readme-ov-file#installation).
+
+# ESM-2 (TransformerEngine-Optimized) Overview
+
+## Description:
+
+ESM-2 is a state-of-the-art protein model trained on a masked language modelling objective. It predicts protein
+structures from amino acid sequences, leveraging a transformer-based architecture for accurate 3D modeling. It is
+suitable for fine-tuning on a wide range of tasks that take protein sequences as input.
 
 This version of the ESM-2 model is optimized with NVIDIA's
-[TransformerEngine](https://github.com/NVIDIA/TransformerEngine) library. It is based on the
-[original ESM-2 model](https://huggingface.co/facebook/esm2_t48_15B_UR50D) from Facebook Research,
-and (within numerical precision) has identical weights and outputs.
-
-ESM-2 is a state-of-the-art protein model trained on a masked language modelling objective. It is
-suitable for fine-tuning on a wide range of tasks that take protein sequences as input. For detailed
-information on the model architecture and training data, please refer to the [accompanying
-paper](https://www.biorxiv.org/content/10.1101/2022.07.20.500902v2). You may also be interested in
-some demo notebooks
-([PyTorch](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/protein_language_modeling.ipynb),
-[TensorFlow](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/protein_language_modeling-tf.ipynb))
-which demonstrate how to fine-tune ESM-2 models on your tasks of interest.
-
-Several ESM-2 checkpoints are available in the Hub with varying sizes. Larger sizes generally have
-somewhat better accuracy, but require much more memory and time to train:
+[TransformerEngine](https://github.com/NVIDIA/TransformerEngine) library. It is based on the original ESM-2 model from
+Facebook Research, and (within numerical precision) has identical weights and outputs.
+
+This model is ready for commercial/non-commercial use.
+
+## Third-Party Community Consideration
+
+This model is not owned or developed by NVIDIA. This model has been developed and built to a third-party's requirements
+for this application and use case; see link to Non-NVIDIA Model Card [ESM-2 Model
+Card](https://huggingface.co/facebook/esm2_t33_650M_UR50D).
+
+### License/Terms of Use:
+
+ESM-2 is licensed under the [MIT license](https://github.com/facebookresearch/esm/blob/main/LICENSE).
+
+### Deployment Geography:
+
+Global
+
+### Use Case:
+
+Protein structure prediction, specifically predicting 3D protein structures from amino acid sequences.
+
+### Release Date:
+
+Hugging Face 07/29/2025 via [https://huggingface.co/nvidia/esm2_t12_35M_UR50D](https://huggingface.co/nvidia/esm2_t12_35M_UR50D)
+
+## Reference(s):
+
+- [Evolutionary-scale prediction of atomic level protein structure with a language
+  model](https://www.biorxiv.org/content/10.1101/2022.07.20.500902v2) - detailed information on the model architecture
+  and training data, please refer to the accompanying [paper].
+- Demo notebooks
+  ([PyTorch](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/protein_language_modeling.ipynb),
+  [TensorFlow](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/protein_language_modeling-tf.ipynb))
+  which demonstrate how to fine-tune ESM-2 models on your tasks of interest.
+
+## Model Architecture:
+
+**Architecture Type:** Transformer
+**Network Architecture:** ESM-2
+
+**This model was developed based on:** [ESM-2](https://huggingface.co/facebook/esm2_t33_650M_UR50D) <br>
+**Number of model parameters:** 3.4 x 10^7
+
+## Input:
+
+**Input Type:** Text (Protein Sequences) <br>
+**Input Format:** String <br>
+**Input Parameters:** One-Dimensional (1D) <br>
+**Other Properties Related to Input:** Protein sequence represented as a string of canonical amino acids, of maximum
+length 1022. Longer sequences are automatically truncated to this length.
+
+## Output:
+
+**Output Type:** Embeddings (Amino acid and sequence-level) <br>
+**Output Format:** Vector <br>
+**Output Parameters:** One-Dimensional (1D) <br>
+**Other Properties Related to Output:** Numeric vector with floating-point values corresponding to an embedding for each
+amino acid in the input protein sequence. Maximum output length is 1022 embeddings - one embedding vector per amino
+acid.
+
+Our AI models are designed and/or optimized to run on NVIDIA GPU-accelerated systems. By leveraging NVIDIA’s hardware
+(e.g. GPU cores) and software frameworks (e.g., CUDA libraries), the model achieves faster training and inference times
+compared to CPU-only solutions.
+
+## Software Integration:
+
+**Runtime Engine(s):**
+
+- Hugging Face Transformers
+
+**Supported Hardware Microarchitecture Compatibility:**
+
+- NVIDIA Ampere
+- NVIDIA Blackwell
+- NVIDIA Hopper
+
+**[Preferred/Supported] Operating System(s):**
+
+- Linux
+
+The integration of foundation and fine-tuned models into AI systems requires additional testing using use-case-specific
+data to ensure safe and effective deployment. Following the V-model methodology, iterative testing and validation at
+both unit and system levels are essential to mitigate risks, meet technical and functional requirements, and ensure
+compliance with safety and ethical standards before deployment.
+
+## Model Version: This model features the following version/checkpoints:
+
+Several ESM-2 checkpoints are available with varying sizes. Larger sizes have better accuracy, but require more memory
+and time to train:
 
 | Checkpoint name                                                          | Num layers | Num parameters |
 | ------------------------------------------------------------------------ | ---------- | -------------- |
@@ -32,3 +120,117 @@ somewhat better accuracy, but require much more memory and time to train:
 | [esm2_t30_150M_UR50D](https://huggingface.co/nvidia/esm2_t30_150M_UR50D) | 30         | 150M           |
 | [esm2_t12_35M_UR50D](https://huggingface.co/nvidia/esm2_t12_35M_UR50D)   | 12         | 35M            |
 | [esm2_t6_8M_UR50D](https://huggingface.co/nvidia/esm2_t6_8M_UR50D)       | 6          | 8M             |
+
+## Training and Evaluation Datasets:
+
+## Training Datasets:
+
+**Link:** [UniRef90](https://www.uniprot.org/uniref?query=%28identity%3A0.9%29)
+
+**Data Modality:**
+
+- Text (Protein Sequences)
+
+**Text Training Data Size:**
+
+- 1 Billion to 10 Trillion Tokens
+
+**Data Collection Method:**
+
+- Human
+
+**Labeling Method:**
+
+- N/A
+
+**Properties (Quantity, Dataset Descriptions, Sensor(s)):** UniRef90 clusters are generated from the UniRef100 seed
+sequences with a 90% sequence identity threshold using the MMseqs2 algorithm. The seed sequences are the longest members
+of the UniRef100 cluster. However, the longest sequence is not always the most informative. There is often more
+biologically relevant information and annotation (name, function, cross-references) available on other cluster members.
+All the proteins in each cluster are ranked to facilitate the selection of a biologically relevant representative for
+the cluster.
+
+**Link:** [UniRef50](https://www.uniprot.org/uniref?query=%28identity%3A0.5%29)
+
+**Data Modality:**
+
+- Text (Protein Sequences)
+
+**Text Training Data Size:**
+
+- 1 Billion to 10 Trillion Tokens
+
+**Data Collection Method:**
+
+- Human
+
+**Labeling Method:**
+
+- N/A
+
+**Properties:** UniRef50 clusters are generated from the UniRef90 seed sequences with a 50% sequence identity threshold
+using the MMseqs2 algorithm. The seed sequences are the longest members of the UniRef90 cluster. However, the longest
+sequence is not always the most informative. There is often more biologically relevant information and annotation (name,
+function, cross-references) available on other cluster members. All the proteins in each cluster are ranked to
+facilitate the selection of a biologically relevant representative for the cluster.
+
+## Evaluation Datasets:
+
+**Link:** [Continuous Automated Model Evaluation (CAMEO)](https://pmc.ncbi.nlm.nih.gov/articles/PMC8673552/)
+
+**Benchmark Score:** 0.56
+
+**Data Collection Method:**
+
+- Human
+
+**Labeling Method:**
+
+- N/A
+
+**Properties:** The data is collected by taking sequences of protein structures that are about to be released weekly by
+the Protein Data Bank (PDB). These sequences are sent as "blind targets" to participating protein structure prediction
+servers, which then return their predictions.
+
+**Link:** [CASP14 (Critical Assessment of Methods of Protein Structure Prediction)](https://pubmed.ncbi.nlm.nih.gov/34533838/)
+
+**Benchmark Score:** 0.41
+
+**Data Collection Method:**
+
+- Human
+
+**Labeling Method:**
+
+- N/A
+
+**Properties:** The data for CASP14 targets is collected from protein structures that are newly solved by experimental
+structural biologists. The CASP organizers receive the amino acid sequences of these proteins before their full,
+three-dimensional structures are publicly released in the Protein Data Bank (PDB). They then provide these sequences to
+participating research groups and servers, who must submit their predicted structures within a specific time frame.
+
+## Inference:
+
+**Acceleration Engine:**
+
+- Hugging Face Transformers
+
+**Test Hardware:**
+
+- A100
+- H100
+- H200
+- GB200
+
+## Ethical Considerations:
+
+NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable
+development for a wide array of AI applications. When downloaded or used in accordance with our terms of service,
+developers should work with their internal model team to ensure this model meets requirements for the relevant industry
+and use case and addresses unforeseen product misuse.
+
+Users are responsible for ensuring the physical properties of model-generated molecules are appropriately evaluated and
+comply with applicable safety regulations and ethical standards.
+
+Please report model quality, risk, security vulnerabilities or NVIDIA AI Concerns
+[here](https://www.nvidia.com/en-us/support/submit-security-vulnerability/).

config.json

@@ -4,12 +4,14 @@
   ],
   "attention_probs_dropout_prob": 0.0,
   "attn_input_format": "bshd",
+  "attn_mask_type": "padding",
   "auto_map": {
     "AutoConfig": "esm_nv.NVEsmConfig",
     "AutoModel": "esm_nv.NVEsmModel",
     "AutoModelForMaskedLM": "esm_nv.NVEsmForMaskedLM"
   },
   "classifier_dropout": null,
+  "dtype": "float32",
   "emb_layer_norm_before": false,
   "encoder_activation": "gelu",
   "esmfold_config": null,
@@ -33,8 +35,7 @@
   "position_embedding_type": "rotary",
   "qkv_weight_interleaved": true,
   "token_dropout": true,
-  "torch_dtype": "float32",
-  "transformers_version": "4.55.4",
+  "transformers_version": "4.56.2",
   "use_cache": true,
   "vocab_list": null,
   "vocab_size": 33

esm_nv.py

@@ -60,6 +60,7 @@ class NVEsmConfig(EsmConfig):
         micro_batch_size: Optional[int] = None,
         max_seq_length: Optional[int] = None,
         padded_vocab_size: Optional[int] = 64,
+        attn_mask_type: str = "padding",
         **kwargs,
     ):
         """Initialize the NVEsmConfig with additional TE-related config options.
@@ -89,6 +90,7 @@ class NVEsmConfig(EsmConfig):
                 ensure same kernels are used for forward propogation and activation recompute phase.
             padded_vocab_size: The padded vocabulary size to support FP8. If not provided, defaults
                 to vocab_size. Must be greater than or equal to vocab_size.
+            attn_mask_type: The type of attention mask to use.
             **kwargs: Additional config options to pass to EsmConfig.
         """
         super().__init__(**kwargs)
@@ -99,6 +101,7 @@ class NVEsmConfig(EsmConfig):
         self.fuse_qkv_params = fuse_qkv_params
         self.micro_batch_size = micro_batch_size
         self.max_seq_length = max_seq_length
+        self.attn_mask_type = attn_mask_type
 
         # Set padded_vocab_size with default fallback to vocab_size
         self.padded_vocab_size = padded_vocab_size if padded_vocab_size is not None else self.vocab_size
@@ -133,14 +136,14 @@ class NVEsmEncoder(nn.Module):
                     qkv_weight_interleaved=config.qkv_weight_interleaved,
                     layer_number=i + 1,
                     layer_type="encoder",
-                    self_attn_mask_type="padding",
+                    self_attn_mask_type=config.attn_mask_type,
                     activation=config.encoder_activation,
                     attn_input_format=config.attn_input_format,
                     seq_length=config.max_seq_length,
                     micro_batch_size=config.micro_batch_size,
                     num_gqa_groups=config.num_attention_heads,
                     fuse_qkv_params=config.fuse_qkv_params,
-                    params_dtype=config.torch_dtype,
+                    params_dtype=config.dtype,
                     window_size=(-1, -1),
                 )
                 for i in range(config.num_hidden_layers)