Fix Transformers 5 compatibility and make xFormers optional

#3
by y-korobko - opened

Fix compatibility with Transformers 5 and make xFormers an optional dependency.

Includes changes to model.py, rotary.py, README.md, and a new swiglu.py module.

Digital Research Unit - Arab Center For Research and Policy Studies org
•
edited 19 days ago

Thank you for the PR, this is really cool! I'll take a closer look soon.

Before we can merge, could you check a few things?

  • Make sure it works across Transformers 5.x, not just 5.20, and let us know the minimum version it supports.
  • Benchmark the speed before and after your changes.
  • Run the Muradif benchmark before and after your changes. Both the embeddings and the final scores should be identical.

Thanks for quick response, I don't have access to linux GPU right now, so I won't be able to use "original" version of your model (xformers is not published for macOS). But I can benchmark the model using monkey-patched version of SwiGLU, like I did in PR.

Hi, I did the benchmarking. I checked transformers 5.0.0 and 5.20.2 it worked with both. Results are identical, speed difference is negligible. I checked on CPU.

I benchmarked and compared results using code from here, ran bash run_cpu.sh. Used python3.14. Those are the results:

=== TIMING / QUALITY ===
run           TF          load (s)   encode (s)  forward (s)       sent/s   accuracy   correct
tf4_57_6      4.57.6         0.082      998.087      991.970       115.88   0.784069     30229
tf5_0_0       5.0.0          0.114     1007.734     1001.981       114.77   0.784069     30229
tf5_20_2      5.20.2         0.057      996.630      990.978       116.05   0.784069     30229

Speedup ratios relative to TF 4.57.6 (higher is faster):
  tf5_0_0: load 0.7191x; forward 0.9900x; end-to-end 0.9904x
  tf5_20_2: load 1.4460x; forward 1.0010x; end-to-end 1.0015x

=== tf4_57_6__vs__tf5_0_0 ===
Weights: exact=True, allclose=True, different_tensors=0, max_abs=0, keys_missing=0
Answers disagree: 0/38,554; scores outside tolerance: 0; embeddings exact: True; allclose: True; maximum embedding absolute difference: 0
Cosine score max abs deltas: positive=0; negative=0
Mean embedding vector cosines: anchor=1.0000000000, synonym=1.0000000000, irrelevant=1.0000000000

=== tf4_57_6__vs__tf5_20_2 ===
Weights: exact=True, allclose=True, different_tensors=0, max_abs=0, keys_missing=0
Answers disagree: 0/38,554; scores outside tolerance: 0; embeddings exact: True; allclose: True; maximum embedding absolute difference: 0
Cosine score max abs deltas: positive=0; negative=0
Mean embedding vector cosines: anchor=1.0000000000, synonym=1.0000000000, irrelevant=1.0000000000

=== tf5_0_0__vs__tf5_20_2 ===
Weights: exact=True, allclose=True, different_tensors=0, max_abs=0, keys_missing=0
Answers disagree: 0/38,554; scores outside tolerance: 0; embeddings exact: True; allclose: True; maximum embedding absolute difference: 0
Cosine score max abs deltas: positive=0; negative=0
Mean embedding vector cosines: anchor=1.0000000000, synonym=1.0000000000, irrelevant=1.0000000000

Wrote comparison_summary.json, 3 weight comparison JSONs, and 3 CSVs (one row per Muradif triplet per model pair).
Digital Research Unit - Arab Center For Research and Policy Studies org

Thanks a lot for this! I tested it on GPU (RTX 3090) against main: Muradif results are identical, speed is the same, and it works on every Transformers release from 5.0.0 to 5.20.2. The only issue was Transformers 4, where transformers.initialization doesn't exist, so I pushed a small fallback on top of your commits. Merged now. I'll also bring the same change to NeoAraBERT_MSA and NeoAraBERT_DA.

One note on scoring: with mean pooling over all tokens and every _ filled, the model gets 0.870753 (33,571 / 38,554) on Muradif, vs your 0.784069 (30,229) with the CLS token. Nearly all of that gain comes from mean pooling; filling every _ instead of only the first adds 6 correct rows. Your script with CLS gives exactly your number on GPU too, so your runs were correct. Thanks again!

Osama-Rakan-Al-Mraikhat changed pull request status to merged
Digital Research Unit - Arab Center For Research and Policy Studies org

I'm also curious how you came across NeoAraBERT. Are you using it for a research paper? And was the macOS issue (not being able to install xFormers) what led you to make this PR? @y-korobko

I use it in a RAG, but it's related to the research. I've opened the PR mainly to fix incorrect model initialization. In Transformers v5 non-persistent buffers, in your case self.freqs_cis, which are used in rotary embeddings, are initialized randomly. So error propagates in self attention and causes model to return NaNs.

It is also the case for NeoBert, check this PR
https://huggingface.co/chandar-lab/NeoBERT/discussions/14

Sign up or log in to comment