fix(bridge): rotate Cohere and Cohere2 with their own RoPE convention

[Mudassiruddin7]

Description

CohereForCausalLM and Cohere2ForCausalLM do not match Hugging Face through the bridge, because the bridge rotates Q and K with the wrong RoPE convention. Cohere's CohereRotaryEmbedding returns cos/sin with every frequency repeated twice (torch.repeat_interleave(freqs, 2, dim=-1), so [c0, c0, c1, c1, ...]), and Cohere's apply_rotary_pos_emb uses a rotate_half that pairs adjacent elements. PositionEmbeddingsAttentionBridge used the Llama formula for Cohere. The adapter comment says the rotation is "delegated to HF", but the generic attention bridge reconstructs it.

rotary_adjacent_pairs (GLM/ERNIE) does not fix it either: those models get the usual half-duplicated cos/sin and the helper interleaves them itself, so with Cohere's already interleaved cos/sin it would interleave twice. I tried that first and it left the attention output wrong.

So this adds a separate config flag, rotary_interleaved_cos_sin (default False, declared next to rotary_adjacent_pairs), and a helper _apply_rotary_pos_emb_interleaved_cos_sin that rotates adjacent pairs with cos/sin as given. The Cohere adapter sets the flag, and Cohere2ArchitectureAdapter inherits it. Other architectures are unchanged.

I found it with a parity battery over the supported causal-LM architectures (a tiny random-weight Hugging Face model per architecture wrapped with build_bridge_from_module, bridge against HF logits). For Cohere the q/k/v projections, cos and sin matched and the rotated q did not, so the rotation is the only difference. The project's own record agrees: supported_models.json has hyper-accel/tiny-random-cohere below threshold with forward_pass_logits failing at max_diff=0.009886.

Bridge against an untouched copy of the Hugging Face model (fp32, CPU):

	dev	this PR
hyper-accel/tiny-random-cohere, real checkpoint	max diff 1.07e-02, relative 7.4%, argmax differs	max diff 7.8e-08, relative 5.4e-07, argmax equal
trl-internal-testing/tiny-CohereForCausalLM, real checkpoint	2.2e-09	2.2e-09 (logit scale is 0.004, so it cannot show the difference)
tiny random CohereForCausalLM, wide weights	1.6e-01	1.2e-07
tiny random Cohere2ForCausalLM, wide weights	3.4e-01	2.2e-07

With the fix, greedy generation of the tiny random Cohere model, with and without the KV cache, is token for token the same as the untouched Hugging Face model. I did not run verify_models and I could not try the gated Cohere checkpoints.

One thing I learned on the way, that you may know already: build_bridge_from_module replaces the submodules of the model it is given (its state_dict keys get _original_component), so calling the Hugging Face model afterwards runs the bridge's code. A reference has to be computed before bridging, or on a deepcopy, which is what the new tests do.

Fixes # (no issue, found by the parity check above)

Type of change

• Bug fix (non-breaking change which fixes an issue)

Tests

In test_cohere_adapter.py: the adapter sets rotary_interleaved_cos_sin, the default is False (on a fresh config, the class-scoped fixture is mutated by the adapter), and TestCohereParity builds a tiny random CohereForCausalLM and compares the bridge with an untouched copy using assert_tiny_parity. test_cohere2_adapter.py has the same for Cohere2ForCausalLM. Four of the new tests fail on dev (the two config tests and both parity tests), and all pass with this change. The two files pass, 83 tests.

All of tests/unit (Windows, Python 3.12, torch 2.14.1 CPU, transformers 5.19.0, with timm and torchvision installed): 7273 passed, 3 failed. The same 3 fail on unmodified dev: test_model_structure_doc.py::test_every_qualified_hook_name_across_docs_exists (a UnicodeDecodeError reading a file with the Windows default codec) and two in test_sparse_probing.py (test_default_tolerance_accepts_large_scale_activations and test_stop_reason_distinguishes_converged_from_capped_fits), which I did not look into. black, isort, pycln and mypy are clean on the changed files.

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes (3 fail, identical on unmodified dev, see above)
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.