add revision pins and edit github actions (#1423)

pstjohn · web-flow · commit f2c2b146e0e1 · 2026-01-20T20:16:29.000Z
Signed-off-by: Peter St. John &lt;pstjohn@nvidia.com&gt;
diff --git a/.github/workflows/unit-tests-recipes.yml b/.github/workflows/unit-tests-recipes.yml
@@ -138,6 +138,7 @@ jobs:
       env:
         CI: true
         HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        HF_HOME: /cache/huggingface
     strategy:
       matrix:
         recipe: ${{ fromJson(needs.changed-dirs.outputs.dirs) }}
@@ -155,6 +156,15 @@ jobs:
           sparse-checkout: "${{ matrix.recipe.dir }}"
           sparse-checkout-cone-mode: false
 
+      - name: Cache Hugging Face models
+        uses: actions/cache@v4
+        with:
+          path: /cache/huggingface
+          key: ${{ runner.os }}-huggingface-${{ matrix.recipe.name }}
+          restore-keys: |
+            ${{ runner.os }}-huggingface-${{ matrix.recipe.name }}-
+            ${{ runner.os }}-huggingface-
+
       - name: Install dependencies
         working-directory: ${{ matrix.recipe.dir }}
         run: |
@@ -179,7 +189,6 @@ jobs:
           fi
           pytest -v .
 
-
   verify-recipe-tests:
     # This job checks the status of the unit-tests matrix and fails if any matrix job failed or was cancelled.
     # Use this job as the required check for PRs.
diff --git a/bionemo-recipes/models/amplify/src/amplify/export.py b/bionemo-recipes/models/amplify/src/amplify/export.py
@@ -56,11 +56,11 @@ def export_hf_checkpoint(tag: str, export_path: Path):
         tag: The tag of the checkpoint to export.
         export_path: The parent path to export the checkpoint to.
     """
-    model_hf = AutoModel.from_pretrained(f"chandar-lab/{tag}", trust_remote_code=True)
+    model_hf = AutoModel.from_pretrained(f"chandar-lab/{tag}", trust_remote_code=True, revision="d918a9e8")
     model_te = convert_amplify_hf_to_te(model_hf)
     model_te.save_pretrained(export_path / tag)
 
-    tokenizer = AutoTokenizer.from_pretrained(f"chandar-lab/{tag}")
+    tokenizer = AutoTokenizer.from_pretrained(f"chandar-lab/{tag}", revision="d918a9e8")
     tokenizer.save_pretrained(export_path / tag)
 
     # Patch the config
diff --git a/bionemo-recipes/models/amplify/tests/conftest.py b/bionemo-recipes/models/amplify/tests/conftest.py
@@ -30,12 +30,12 @@ def requires_fp8(func):
 
 @pytest.fixture
 def tokenizer():
-    return AutoTokenizer.from_pretrained("chandar-lab/AMPLIFY_120M")
+    return AutoTokenizer.from_pretrained("chandar-lab/AMPLIFY_120M", revision="d918a9e8")
 
 
 @pytest.fixture
 def config():
-    config = AutoConfig.from_pretrained("chandar-lab/AMPLIFY_120M", trust_remote_code=True)
+    config = AutoConfig.from_pretrained("chandar-lab/AMPLIFY_120M", trust_remote_code=True, revision="d918a9e8")
     config.dtype = torch.bfloat16
     return config
 
diff --git a/bionemo-recipes/models/amplify/tests/test_amplify_model.py b/bionemo-recipes/models/amplify/tests/test_amplify_model.py
@@ -68,7 +68,7 @@ def test_te_model_has_all_te_layers(config):
 
 
 def test_models_have_identical_outputs(input_data):
-    model_hf = amp_hf.AMPLIFY.from_pretrained("chandar-lab/AMPLIFY_120M")
+    model_hf = amp_hf.AMPLIFY.from_pretrained("chandar-lab/AMPLIFY_120M", revision="d918a9e8")
     model_te = convert_amplify_hf_to_te(model_hf)
     input_data = {k: v.to("cuda") for k, v in input_data.items()}
 
@@ -85,7 +85,7 @@ def test_models_have_identical_outputs(input_data):
 
 
 def test_converted_model_roundtrip(input_data, tmp_path):
-    model_hf = amp_hf.AMPLIFY.from_pretrained("chandar-lab/AMPLIFY_120M")
+    model_hf = amp_hf.AMPLIFY.from_pretrained("chandar-lab/AMPLIFY_120M", revision="d918a9e8")
     model_te = convert_amplify_hf_to_te(model_hf)
 
     model_te.save_pretrained(tmp_path / "AMPLIFY_120M")
@@ -108,7 +108,7 @@ def test_converted_model_roundtrip(input_data, tmp_path):
 
 
 def test_convert_state_dict():
-    model_hf = amp_hf.AMPLIFY.from_pretrained("chandar-lab/AMPLIFY_120M")
+    model_hf = amp_hf.AMPLIFY.from_pretrained("chandar-lab/AMPLIFY_120M", revision="d918a9e8")
     model_te = convert_amplify_hf_to_te(model_hf)
 
     from amplify.state_dict_convert import _pack_qkv_weight, _pad_bias, _pad_weights, mapping
@@ -171,7 +171,7 @@ def test_convert_state_dict():
 
 
 def test_hf_trained_model_loss(input_data):
-    model = amp_hf.AMPLIFY.from_pretrained("chandar-lab/AMPLIFY_120M")
+    model = amp_hf.AMPLIFY.from_pretrained("chandar-lab/AMPLIFY_120M", revision="d918a9e8")
     model.to("cuda", dtype=torch.bfloat16)
     input_data = {k: v.to("cuda") for k, v in input_data.items()}
     model.eval()
@@ -182,7 +182,7 @@ def test_hf_trained_model_loss(input_data):
 
 
 def test_te_trained_model_loss(input_data):
-    model_hf = amp_hf.AMPLIFY.from_pretrained("chandar-lab/AMPLIFY_120M")
+    model_hf = amp_hf.AMPLIFY.from_pretrained("chandar-lab/AMPLIFY_120M", revision="d918a9e8")
     model = convert_amplify_hf_to_te(model_hf)
     model.to("cuda", dtype=torch.bfloat16)
     input_data = {k: v.to("cuda") for k, v in input_data.items()}
@@ -194,7 +194,7 @@ def test_te_trained_model_loss(input_data):
 
 
 def test_hf_reinitialized_model_loss(input_data):
-    config = amp_hf.AMPLIFYConfig.from_pretrained("chandar-lab/AMPLIFY_120M")
+    config = amp_hf.AMPLIFYConfig.from_pretrained("chandar-lab/AMPLIFY_120M", revision="d918a9e8")
     model = amp_hf.AMPLIFY(config)
     model.to("cuda", dtype=torch.bfloat16)
     input_data = {k: v.to("cuda") for k, v in input_data.items()}
@@ -207,7 +207,7 @@ def test_hf_reinitialized_model_loss(input_data):
 
 
 def test_te_reinitialized_model_loss(input_data):
-    config = amp_te.AMPLIFYConfig.from_pretrained("chandar-lab/AMPLIFY_120M")
+    config = amp_te.AMPLIFYConfig.from_pretrained("chandar-lab/AMPLIFY_120M", revision="d918a9e8")
     model = amp_te.AMPLIFYForMaskedLM(config)
     model.to("cuda", dtype=torch.bfloat16)
     input_data = {k: v.to("cuda") for k, v in input_data.items()}
diff --git a/bionemo-recipes/models/amplify/tests/test_encoder_block.py b/bionemo-recipes/models/amplify/tests/test_encoder_block.py
@@ -56,7 +56,7 @@ def data(self) -> torch.Tensor:
 
 @pytest.fixture
 def config():
-    config = AutoConfig.from_pretrained("chandar-lab/AMPLIFY_120M", trust_remote_code=True)
+    config = AutoConfig.from_pretrained("chandar-lab/AMPLIFY_120M", trust_remote_code=True, revision="d918a9e8")
     config.dtype = torch.bfloat16
     return config
 
diff --git a/bionemo-recipes/models/amplify/tests/test_rotary_embeddings.py b/bionemo-recipes/models/amplify/tests/test_rotary_embeddings.py
@@ -29,7 +29,7 @@ def test_apply_rotary_pos_emb():
     key = torch.randn([2, 72, 10, 64], dtype=torch.bfloat16, generator=rng).to("cuda")
 
     # AMPLIFY HF Rope
-    hf_config = AutoConfig.from_pretrained("chandar-lab/AMPLIFY_120M", trust_remote_code=True)
+    hf_config = AutoConfig.from_pretrained("chandar-lab/AMPLIFY_120M", trust_remote_code=True, revision="d918a9e8")
 
     freqs_cis = precompute_freqs_cis(hf_config.hidden_size // hf_config.num_attention_heads, 72).to("cuda")
     q_post, k_post = apply_rotary_emb(query, key, freqs_cis)
diff --git a/bionemo-recipes/models/esm2/tests/conftest.py b/bionemo-recipes/models/esm2/tests/conftest.py
@@ -86,7 +86,7 @@ def input_data(tokenizer, tokenized_proteins):
 
 @pytest.fixture
 def te_model_checkpoint(tmp_path):
-    model_hf = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D")
+    model_hf = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D", revision="c731040f")
     model_te = convert_esm_hf_to_te(model_hf)
     model_te.save_pretrained(tmp_path / "te_model_checkpoint")
     return tmp_path / "te_model_checkpoint"
diff --git a/bionemo-recipes/models/esm2/tests/test_convert.py b/bionemo-recipes/models/esm2/tests/test_convert.py
@@ -22,7 +22,7 @@ def test_convert_te_to_hf_roundtrip():
     """Test that converting HF -> TE -> HF produces the same model."""
     from esm.convert import convert_esm_hf_to_te, convert_esm_te_to_hf
 
-    model_hf_original = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D")
+    model_hf_original = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D", revision="c731040f")
 
     model_te = convert_esm_hf_to_te(model_hf_original)
     model_hf_converted = convert_esm_te_to_hf(model_te)
@@ -42,7 +42,7 @@ def test_qkv_unpacking():
     """Test that QKV unpacking works correctly."""
     from esm.convert import convert_esm_hf_to_te, convert_esm_te_to_hf
 
-    model_hf = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D")
+    model_hf = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D", revision="c731040f")
     model_te = convert_esm_hf_to_te(model_hf)
     model_hf_converted = convert_esm_te_to_hf(model_te)
 
@@ -64,7 +64,7 @@ def test_config_conversion():
     """Test that config conversion works correctly."""
     from esm.convert import convert_esm_hf_to_te, convert_esm_te_to_hf
 
-    model_hf = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D")
+    model_hf = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D", revision="c731040f")
     model_te = convert_esm_hf_to_te(model_hf)
     model_hf_converted = convert_esm_te_to_hf(model_te)
 
@@ -97,7 +97,7 @@ def test_padding_unpadding_operations():
     """Test that padding and unpadding operations work correctly for embeddings and decoder weights."""
     from esm.convert import convert_esm_hf_to_te, convert_esm_te_to_hf
 
-    model_hf = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D")
+    model_hf = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D", revision="c731040f")
     model_te = convert_esm_hf_to_te(model_hf)
     model_hf_converted = convert_esm_te_to_hf(model_te)
 
@@ -146,7 +146,7 @@ def test_weight_initialization_matches_hf():
 
     set_seed(42)
 
-    config_hf = AutoConfig.from_pretrained("facebook/esm2_t6_8M_UR50D", vocab_size=64)
+    config_hf = AutoConfig.from_pretrained("facebook/esm2_t6_8M_UR50D", vocab_size=64, revision="c731040f")
     model_hf = EsmForMaskedLM(config_hf)
     model_te_converted = convert_esm_hf_to_te(model_hf)
 
diff --git a/bionemo-recipes/models/esm2/tests/test_cp_bshd.py b/bionemo-recipes/models/esm2/tests/test_cp_bshd.py
@@ -72,7 +72,7 @@ def get_te_model_checkpoint(tmp_path):
     Returns:
         The path to the saved model checkpoint.
     """
-    model_hf = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D")
+    model_hf = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D", revision="c731040f")
     model_te = convert_esm_hf_to_te(model_hf, attn_mask_type="no_mask", attn_input_format="bshd")
     model_te.save_pretrained(tmp_path / "te_model_checkpoint")
     return tmp_path / "te_model_checkpoint"
@@ -183,7 +183,7 @@ def test_context_parallel_equivalence_2process():
         model_ckpt = get_te_model_checkpoint(tmp_path)
 
         # Create tokenizer for real protein sequences
-        tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D")
+        tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D", revision="c731040f")
         input_data_bshd_padded_dp0 = get_dummy_data_bshd_with_padding_dp0(tokenizer=tokenizer)
 
         model = NVEsmForMaskedLM.from_pretrained(
diff --git a/bionemo-recipes/models/esm2/tests/test_cp_thd.py b/bionemo-recipes/models/esm2/tests/test_cp_thd.py
@@ -82,7 +82,7 @@ def get_te_model_checkpoint(tmp_path):
     Returns:
         The path to the saved model checkpoint.
     """
-    model_hf = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D")
+    model_hf = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D", revision="c731040f")
     model_te = convert_esm_hf_to_te(model_hf)
     model_te.save_pretrained(tmp_path / "te_model_checkpoint")
     return tmp_path / "te_model_checkpoint"
@@ -174,7 +174,7 @@ def test_context_parallel_equivalence_2process():
         model_ckpt = get_te_model_checkpoint(tmp_path)
 
         # Create tokenizer for real protein sequences
-        tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D")
+        tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D", revision="c731040f")
         input_data_thd_padded_dp0 = get_dummy_data_thd_with_padding_dp0(tokenizer)
 
         model = NVEsmForMaskedLM.from_pretrained(
diff --git a/bionemo-recipes/models/esm2/tests/test_distributed_fp8.py b/bionemo-recipes/models/esm2/tests/test_distributed_fp8.py
@@ -152,7 +152,7 @@ def is_main_process(self) -> bool:
     )
     device = f"cuda:{dist_config.local_rank}"
 
-    config = NVEsmConfig.from_pretrained("nvidia/esm2_t6_8M_UR50D", dtype=torch.bfloat16)
+    config = NVEsmConfig.from_pretrained("facebook/esm2_t6_8M_UR50D", dtype=torch.bfloat16, revision="c731040f")
     model = NVEsmForMaskedLM(config)
 
     if args.strategy is Strategy.FSDP2:
diff --git a/bionemo-recipes/models/esm2/tests/test_fast_tokenizer.py b/bionemo-recipes/models/esm2/tests/test_fast_tokenizer.py
@@ -17,13 +17,13 @@
 
 
 def test_tokenizer_vocab_equivalence():
-    original_tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D")
+    original_tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D", revision="c731040f")
     fast_tokenizer = AutoTokenizer.from_pretrained("esm_fast_tokenizer")
     assert original_tokenizer.get_vocab() == fast_tokenizer.get_vocab()
 
 
 def test_tokenizer_tokenization_equivalence(test_proteins):
-    original_tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D")
+    original_tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D", revision="c731040f")
     fast_tokenizer = AutoTokenizer.from_pretrained("esm_fast_tokenizer")
 
     original_output = original_tokenizer(test_proteins)
diff --git a/bionemo-recipes/models/esm2/tests/test_meta_device_init.py b/bionemo-recipes/models/esm2/tests/test_meta_device_init.py
@@ -162,7 +162,7 @@ def test_meta_init():
 
 
 def test_cuda_fp8_init(fp8_recipe):
-    config = NVEsmConfig(**AutoConfig.from_pretrained("facebook/esm2_t6_8M_UR50D").to_dict())
+    config = NVEsmConfig(**AutoConfig.from_pretrained("facebook/esm2_t6_8M_UR50D", revision="c731040f").to_dict())
 
     set_seed(42)
     with transformer_engine.pytorch.fp8_model_init(recipe=fp8_recipe):
@@ -174,7 +174,7 @@ def test_cuda_fp8_init(fp8_recipe):
 
 
 def test_meta_fp8_init(fp8_recipe):
-    config = NVEsmConfig(**AutoConfig.from_pretrained("facebook/esm2_t6_8M_UR50D").to_dict())
+    config = NVEsmConfig(**AutoConfig.from_pretrained("facebook/esm2_t6_8M_UR50D", revision="c731040f").to_dict())
 
     set_seed(42)
     with transformer_engine.pytorch.fp8_model_init(recipe=fp8_recipe), torch.device("meta"):
@@ -186,12 +186,12 @@ def test_meta_fp8_init(fp8_recipe):
     verify_model_parameters_initialized_correctly(model, should_be_fp8=True)
 
 
-def test_model_for_token_classification_init():
-    config = NVEsmConfig(**AutoConfig.from_pretrained("nvidia/esm2_t6_8M_UR50D", trust_remote_code=True).to_dict())
+def test_model_for_token_classification_init(te_model_checkpoint):
+    config = NVEsmConfig.from_pretrained(te_model_checkpoint, trust_remote_code=True)
 
     set_seed(42)
     model = NVEsmForTokenClassification.from_pretrained(
-        "nvidia/esm2_t6_8M_UR50D", config=config, dtype=torch.bfloat16, trust_remote_code=True
+        te_model_checkpoint, config=config, dtype=torch.bfloat16, trust_remote_code=True
     )
     model.to("cuda")
 
diff --git a/bionemo-recipes/models/esm2/tests/test_modeling_esm_te.py b/bionemo-recipes/models/esm2/tests/test_modeling_esm_te.py
@@ -23,7 +23,7 @@
 def test_esm_model_for_masked_lm(input_data):
     from esm.modeling_esm_te import NVEsmConfig, NVEsmForMaskedLM
 
-    config = NVEsmConfig(**AutoConfig.from_pretrained("facebook/esm2_t6_8M_UR50D").to_dict())
+    config = NVEsmConfig(**AutoConfig.from_pretrained("facebook/esm2_t6_8M_UR50D", revision="c731040f").to_dict())
     model = NVEsmForMaskedLM(config)
     model.to("cuda")
     input_data = {k: v.to("cuda") for k, v in input_data.items()}
@@ -36,7 +36,7 @@ def test_esm_model_for_masked_lm(input_data):
 def test_esm_model_has_all_te_layers(input_data):
     from esm.modeling_esm_te import NVEsmConfig, NVEsmForMaskedLM
 
-    config = NVEsmConfig(**AutoConfig.from_pretrained("facebook/esm2_t6_8M_UR50D").to_dict())
+    config = NVEsmConfig(**AutoConfig.from_pretrained("facebook/esm2_t6_8M_UR50D", revision="c731040f").to_dict())
     model = NVEsmForMaskedLM(config)
     for name, module in model.named_modules():
         assert not isinstance(module, nn.Linear), f"Vanilla linear layer found in {name}"
@@ -46,7 +46,7 @@ def test_esm_model_has_all_te_layers(input_data):
 def test_convert_state_dict(input_data):
     from esm.convert import _pack_qkv_bias, _pack_qkv_weight, _pad_bias, _pad_weights, convert_esm_hf_to_te, mapping
 
-    model_hf = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D")
+    model_hf = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D", revision="c731040f")
     model_te = convert_esm_hf_to_te(model_hf)
     model_hf.to("cuda")
     model_te.to("cuda")
@@ -166,7 +166,7 @@ def test_golden_values(input_data):
     from esm.convert import convert_esm_hf_to_te
 
     model_hf = AutoModelForMaskedLM.from_pretrained(
-        "facebook/esm2_t6_8M_UR50D", attn_implementation="flash_attention_2"
+        "facebook/esm2_t6_8M_UR50D", attn_implementation="flash_attention_2", revision="c731040f"
     )
     model_te = convert_esm_hf_to_te(model_hf)
     model_te.to(torch.bfloat16)
@@ -195,7 +195,7 @@ def test_converted_model_roundtrip(tmp_path, input_data):
     from esm.convert import convert_esm_hf_to_te
     from esm.modeling_esm_te import NVEsmConfig, NVEsmEncoder, NVEsmForMaskedLM, NVEsmLMHead, NVEsmModel
 
-    model_hf = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D")
+    model_hf = AutoModelForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D", revision="c731040f")
     model_te = convert_esm_hf_to_te(model_hf)
 
     model_te.save_pretrained(tmp_path / "esm2_t6_8M_UR50D_te")
diff --git a/bionemo-recipes/models/geneformer/src/geneformer/export.py b/bionemo-recipes/models/geneformer/src/geneformer/export.py
@@ -35,7 +35,7 @@ def export_hf_checkpoint(model_name: str, export_path: Path):
     """
     print(f"Loading Geneformer model: {model_name}")
 
-    model_hf = AutoModelForMaskedLM.from_pretrained("ctheodoris/Geneformer", subfolder=model_name)
+    model_hf = AutoModelForMaskedLM.from_pretrained("ctheodoris/Geneformer", subfolder=model_name, revision="f45a6c7d")
 
     print(f"Loaded HF model with {len(list(model_hf.parameters()))} parameters")
 
diff --git a/bionemo-recipes/models/geneformer/tests/test_checkpoints_modeling_bert.py b/bionemo-recipes/models/geneformer/tests/test_checkpoints_modeling_bert.py
@@ -24,10 +24,10 @@ def load_geneformer_model(model_name):
     """Helper function to load the correct Geneformer model variant."""
     if model_name == "Geneformer-V2-316M":
         # Default model (no subfolder needed)
-        return AutoModelForMaskedLM.from_pretrained("ctheodoris/Geneformer")
+        return AutoModelForMaskedLM.from_pretrained("ctheodoris/Geneformer", revision="f45a6c7d")
     else:
         # Use subfolder for specific model variants
-        return AutoModelForMaskedLM.from_pretrained("ctheodoris/Geneformer", subfolder=model_name)
+        return AutoModelForMaskedLM.from_pretrained("ctheodoris/Geneformer", subfolder=model_name, revision="f45a6c7d")
 
 
 # Model variants with detailed information
diff --git a/bionemo-recipes/models/llama3/tests/test_convert.py b/bionemo-recipes/models/llama3/tests/test_convert.py
diff --git a/bionemo-recipes/models/llama3/tests/test_lm_eval.py b/bionemo-recipes/models/llama3/tests/test_lm_eval.py
diff --git a/bionemo-recipes/models/llama3/tests/test_modeling_llama_te.py b/bionemo-recipes/models/llama3/tests/test_modeling_llama_te.py
