Amgen's AMPLIFY Port

[ynashed]

Porting and reproducing AMPLIFY protein language mode to BioNeMo paper Hugging Face

• Added a new sub-package under sub-packages/bionemo-amplify following ESM2 sub-package
• Added pretrain script under scripts/protein/amplify
• Tokenizer class from pretrained HF tokenizer "chandar-lab/AMPLIFY_350M"
• Model config and model class very similar to ESM2 with a few changes to accommodate RMSNorm, swiglu activations, and noop LMHead
• Added a new BiobertSpecOption.amplify_bert_layer_with_transformer_engine_spec

---

Summary

High level summary of changes. Try to keep this as short and informative as possible: less is more.

Details

Describe your changes. You can be more detailed and descriptive here.

Usage

python scripts/protein/amplify/amplify_pretrain.py  --hf-dataset-name chandar-lab/UR100P

Testing

Tests were copied from ESM2 but are still WIP

(NOTE: also delete this checklist as you fill-out this PR template)

Most of the changes to files with extensions *.py, *.yaml, *.yml, Dockerfile* or requirements.txt DO REQUIRE both pytest- and jet- CI stages.

• Did you review the Before your PR is "Ready for review" section before asking for review?
• Did you make sure your changes have tests? Did you test your changes locally?
• Can you add the SKIP_CI label to your PR?
• Can you add the PYTEST_NOT_REQUIRED label to your PR?
• Can you add the JET_NOT_REQUIRED label to your PR?