Performance gap between models loaded with `modeling_hf_nomic_bert.py` and `modeling_nomic_bert.py`

[HSILA]

Hello, and thank you for your amazing work! @zanussbaum

I have been using your repository extensively for domain-specific adaptations, but I've encountered an issue that I couldn't find an explanation for.

To clarify, I fine-tuned the model starting from nomic-ai/nomic-embed-text-v1-unsupervised using the contrastive_finetune.yaml configuration with the exact data you used for the Weakly-Supervised Contrastive Pretraining stage. My aim was simply to verify my fine-tuning process before proceeding further.

However, during evaluation on a custom MTEB-like retrieval task, I observed a notable performance gap between:

• The final model uploaded to Hugging Face Hub with convert_to_hf.py (loaded via modeling_hf_nomic_bert.py using SentenceTransformer and mteb.get_model).

• A model loaded using modeling_nomic_bert.py with the following snippet (inspired by your MTEB evaluation method here, and suppose I have written a wrapper for it inspired by nomic_models.py in mteb):

from contrastors import BiEncoder, BiEncoderConfig

config = BiEncoderConfig(
    model_name='src/ckpts/test/final_model',
    encoder=True,
    pooling="mean"
)
model = BiEncoder(config).to(torch.bfloat16)

To double-check, I encoded the same sentence using both methods and computed their cosine similarity. Surprisingly, the similarity was nearly zero.

After reviewing my training setups, I suspect this issue might relate to the differing sequence lengths (seq_len: 2048 in contrastive_pretrain.yaml vs. seq_len: 512 in contrastive_finetune.yaml). Could this discrepancy cause such divergence, perhaps due to rotary embeddings or flash attention? Or is there possibly another argument I missed in the configuration?

Notably, this issue does not happen when the same seq_len value is maintained across training stages (e.g., fine-tuning from the nomic unsupervised model using a seq_len of 2048).

Also, could you clarify the recommended approach for loading and benchmarking models? My experience suggests that loading via modeling_nomic_bert.py more accurately reflects the effects of in-domain training compared to modeling_hf_nomic_bert.py, which tends to show similar performance across most models, although the first one is in bf16 and the second one is in fp32.

For reference, I am using the version before your Mixture-of-Experts (MoE) commit.

Any suggestions or insights would be greatly appreciated.

Thank you!

[zanussbaum]

Thank you for opening the issue!

What is the command you're running? What are you trying to compare? What is the performance gap? Can you provide a minimal code snippet so I can reproduce?

[HSILA]

Hi again, thanks for getting back, and sorry for the late reply, I was busy experimenting.

I think I've pinpointed the issue; it seems to be related to GradCache. TL;DR: I tested the following fine-tuning scenarios:

1. Finetune from nomic-embed-text-v1-unsupervised with pairs
2. Finetune from nomic-embed-text-v1 with triplets
3. Finetune from nomic-embed-text-v1-unsupervised with triplets

In scenario 3, I noticed a significant performance gap when using GradCache during training and prefixes during inference (MTEB evaluation for some tasks). Interestingly, this gap appeared on the ArguAna retrieval task but not on the AmazonPolarityClassification task, which is quite odd. I used a custom model wrapper (BiEncoderWrapper) for this evaluation (inspired by nomic_models.py in MTEB)

For clarity, the tested model variant was first fine-tuned on nomic-embed-text-v1-unsupervised with domain-specific pairs, and then further fine-tuned on a triplet dataset (let's say the ones in finetune_triplets.yaml). The result was as follows:

Task	With Prefix	Without Prefix
ArguAna	0.23	0.44
AmazonPolarityClassification	0.88	0.82

The model is a variant first finetuned on nomic-embed-text-v1-unsupervised with my domain pairs, then finetuned on a triplet dataset (doesn't matter what, let's say one you used for the second stage)

I'm unsure why exactly GradCache and prefix usage interact like this. Also, I noticed that Nomic doesn't use GradCache in the second training phase (contrastive_finetune.yaml), is there a specific reason? I resorted to GradCache since the default batch size (256) exceeded my available compute resources (4 x A100 40GB GPUs).

More context for the future readers:

For training, I'm using a slightly modified fork of this repo, mainly enhancements for better logging, evaluation, and local support for StreamingShardDataset in text_text_dataset.py. Training was executed using:

torchrun --nproc-per-node=4 train.py --config=configs/train/chem_contrastive_finetune.yaml --dtype=bf16

You can check out my sample evaluation code here.

Additional notes:

• Is my model wrapper reliable? Yes, I verified it by comparing it directly with mteb.get_model('nomic-embed-text-v1') on the MTEB(eng, v1) benchmark, and it performed slightly better.

• About the issue title: Initially, I thought the performance difference stemmed from different loading methods (BiEncoderWrapper vs. SentenceTransformer via mteb.get_model). It turned out I had forgotten to pass prefixes.

• Interestingly, evaluating nomic-ai/nomic-embed-text-v1 on MTEB(eng, v1) with and without prefixes showed minimal overall differences (average scores of 0.58 vs. 0.57), but omitting prefixes yielded slightly better performance in Retrieval, STS, and Summarization tasks.