Problem: granite-r2 model produces inconsistent embeddings with low cosine similarity (~0.5) for identical inputs. Expected similarity should be ~1.0.
Create a Python script to get reference embeddings:
from transformers import AutoModel, AutoTokenizer
import torch
import numpy as np
model_name = "ibm-granite/granite-embedding-english-r2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)
# Test inputs
texts = ["Hello world", "Hello world"] # Same text twice
# Get embeddings
inputs = tokenizer(texts, return_tensors="pt", padding=True)
with torch.no_grad():
outputs = model(**inputs)
# CLS token pooling
embeddings = outputs.last_hidden_state[:, 0, :]
# Normalize if model expects it
embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
# Check similarity
from sklearn.metrics.pairwise import cosine_similarity
similarity = cosine_similarity(embeddings[0:1], embeddings[1:2])[0][0]
print(f"Reference similarity: {similarity}") # Should be ~1.0
# Save for comparison
np.save("/tmp/reference_embedding.npy", embeddings[0].numpy())# Generate same embedding twice
echo "Hello world" | OLLAMA_HOST=127.0.0.1:13000 ./ollama run granite-r2 > /tmp/ollama_out1.txt
echo "Hello world" | OLLAMA_HOST=127.0.0.1:13000 ./ollama run granite-r2 > /tmp/ollama_out2.txt
# Compare outputs
diff /tmp/ollama_out1.txt /tmp/ollama_out2.txt- Is the output different on every run (non-determinism bug)?
- Is the output consistent but just wrong compared to HuggingFace (conversion bug)?
Modify ml/backend/ggml/bert.cpp to dump activations:
// After each layer, dump the output
static void dump_tensor(const char* name, struct ggml_tensor* tensor, int layer_idx) {
if (getenv("OLLAMA_DEBUG_ACTIVATIONS") == nullptr) return;
char filename[256];
snprintf(filename, sizeof(filename), "/tmp/activations_layer_%d_%s.bin", layer_idx, name);
FILE* f = fopen(filename, "wb");
if (f) {
// Dump tensor data
fwrite(ggml_get_data(tensor), 1, ggml_nbytes(tensor), f);
fclose(f);
fprintf(stderr, "Dumped %s at layer %d: shape=[%lld, %lld], size=%zu bytes\n",
name, layer_idx, tensor->ne[0], tensor->ne[1], ggml_nbytes(tensor));
}
}
// In your forward pass, after each layer:
for (int il = 0; il < n_layer; il++) {
// ... layer computation ...
dump_tensor("after_attn", cur, il);
dump_tensor("after_ffn", cur, il);
dump_tensor("after_norm", cur, il);
}import torch
import numpy as np
from transformers import AutoModel, AutoTokenizer
model_name = "ibm-granite/granite-embedding-english-r2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)
text = "Hello world"
inputs = tokenizer([text], return_tensors="pt")
# Hook to capture layer outputs
activations = {}
def make_hook(name):
def hook(module, input, output):
activations[name] = output[0].detach().cpu().numpy()
return hook
# Register hooks on each layer
for i, layer in enumerate(model.encoder.layer):
layer.attention.output.register_forward_hook(make_hook(f"layer_{i}_after_attn"))
layer.output.register_forward_hook(make_hook(f"layer_{i}_after_ffn"))
with torch.no_grad():
outputs = model(**inputs)
final_embedding = outputs.last_hidden_state[:, 0, :].numpy()
# Save all activations
for name, tensor in activations.items():
np.save(f"/tmp/hf_{name}.npy", tensor)
np.save("/tmp/hf_final.npy", final_embedding)import numpy as np
for layer in range(22):
hf_attn = np.load(f"/tmp/hf_layer_{layer}_after_attn.npy")
ollama_attn = np.fromfile(f"/tmp/activations_layer_{layer}_after_attn.bin", dtype=np.float32)
# Reshape to match
ollama_attn = ollama_attn.reshape(hf_attn.shape)
# Compute difference
max_diff = np.abs(hf_attn - ollama_attn).max()
mean_diff = np.abs(hf_attn - ollama_attn).mean()
print(f"Layer {layer} attention - max diff: {max_diff:.6f}, mean diff: {mean_diff:.6f}")
if max_diff > 0.01: # Significant divergence
print(f" ⚠️ DIVERGENCE DETECTED AT LAYER {layer}")
break- Token embeddings conversion
- Position embeddings (ModernBERT uses RoPE with dual theta)
- Input normalization
- Full attention implementation vs local attention
- Missing attention bias issue (verify global layers truly have no bias)
- Attention mask handling
- Sliding window attention implementation
- Attention bias application
- Window size and stride
- Gated FFN (GeGLU) implementation
- Verify
ffn_gateandffn_upsplit was correct - Check activation function (GELU vs others)
- CLS token selection
- Output normalization
- Pooling strategy
# Create a tiny 2-layer ModernBERT for easier debugging
from transformers import ModernBertConfig, ModernBertModel
config = ModernBertConfig(
hidden_size=256,
num_hidden_layers=2,
num_attention_heads=4,
intermediate_size=512,
global_attn_every_n_layers=1, # Layer 0, 1 both global
local_attention=128
)
tiny_model = ModernBertModel(config)
tiny_model.save_pretrained("/tmp/tiny-modernbert")Then convert and test this tiny model - easier to debug with fewer layers.
If non-determinism (different output each run):
- Check dropout is disabled during inference
- Check for uninitialized memory
- Check RNG seed is set
If consistent but wrong (same output each run, but differs from HuggingFace):
- Proceed to Phase 3 (layer-by-layer comparison)
- This will pinpoint the problematic component
- Dropout during inference: Ensure dropout is disabled
- Uninitialized tensors: Check all tensors are properly initialized
- Attention mask: Verify padding mask is applied correctly
- Normalization: Check LayerNorm epsilon values match
- Activation functions: Verify GELU vs GeLU_NEW vs other variants
- Tensor layout: Ensure GGML tensor dimensions match expected layout
- Precision: Check float16 vs float32 conversions
- RoPE implementation: Dual theta parameters (base=10000, mscale=1.0)
- Bias tensor application: Global layers should not apply attention biases
- FFN gate split: Verify the mlp.Wi tensor was split correctly into ffn_gate and ffn_up