File size: 6,404 Bytes
eae424a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | //! Binding a HuggingFace DINOv3 checkpoint to the encoder graph's
//! parameters.
//!
//! Three shape conventions differ between the checkpoint and the graph,
//! and all three are handled here rather than in the graph:
//!
//! * `nn.Linear` stores `[out_features, in_features]`; meganeura's
//! `matmul` wants `[in, out]`. Every projection is transposed.
//! * The patch embedding is a 4D `[out, 3, k, k]` conv weight; it is
//! flattened to `[out, patch_dim]` and transposed.
//! * `cls_token` and `register_tokens` are separate parameters; the graph
//! declares one `prefix_tokens` matrix, assembled here.
use meganeura::Session;
use meganeura::data::safetensors::SafeTensorsModel;
use crate::dinov3::Config;
use crate::preprocess::conv_weight_to_matmul;
type Error = Box<dyn std::error::Error>;
/// Checkpoints appear both with and without a `model.` prefix on the
/// encoder layers, depending on whether they were saved from
/// `DINOv3ViTModel` or from the bare backbone. Resolve whichever this one
/// uses once, up front, instead of guessing per tensor.
fn layer_prefix(model: &SafeTensorsModel) -> Result<&'static str, Error> {
let info = model.tensor_info();
for candidate in ["layer.", "model.layer."] {
if info.keys().any(|k| k.starts_with(candidate)) {
return Ok(candidate);
}
}
Err(format!(
"no encoder layers found; checkpoint has {} tensors, e.g. {:?}",
info.len(),
info.keys().take(5).collect::<Vec<_>>()
)
.into())
}
/// Read a tensor under an optional `model.` prefix.
fn embedding_tensor(model: &SafeTensorsModel, name: &str) -> Result<Vec<f32>, Error> {
if model.tensor_info().contains_key(name) {
model.tensor_f32_auto(name)
} else {
model.tensor_f32_auto(&format!("model.{name}"))
}
}
/// Upload every parameter the encoder graph declares.
///
/// Fails loudly on a missing or mis-shaped tensor: a silently skipped
/// weight would leave a zero buffer and produce plausible-looking
/// garbage.
pub fn load_encoder(
session: &mut Session,
model: &SafeTensorsModel,
config: &Config,
) -> Result<(), Error> {
let hidden = config.hidden_size;
let prefix = layer_prefix(model)?;
log::info!("checkpoint layer prefix: {prefix:?}");
// --- Patch embedding: [out, 3, k, k] -> [patch_dim, out] ---
let conv = embedding_tensor(model, "embeddings.patch_embeddings.weight")?;
let expected = hidden * config.patch_dim();
if conv.len() != expected {
return Err(format!(
"patch embedding has {} values, expected {expected} \
(is the checkpoint's patch_size or hidden_size different?)",
conv.len()
)
.into());
}
session.set_parameter(
"embeddings.patch_embeddings.weight",
&conv_weight_to_matmul(&conv, hidden, config.patch_dim()),
);
session.set_parameter(
"embeddings.patch_embeddings.bias",
&embedding_tensor(model, "embeddings.patch_embeddings.bias")?,
);
// --- Prefix tokens: CLS first, then the registers ---
//
// Order matters and is not arbitrary: the reference builds the
// sequence as cat([cls, registers, patches]), and RoPE identifies
// prefix tokens purely by their position at the front.
let cls = embedding_tensor(model, "embeddings.cls_token")?;
let registers = embedding_tensor(model, "embeddings.register_tokens")?;
if cls.len() != hidden {
return Err(format!("cls_token has {} values, expected {hidden}", cls.len()).into());
}
if registers.len() != config.num_register_tokens * hidden {
return Err(format!(
"register_tokens has {} values, expected {} * {hidden}",
registers.len(),
config.num_register_tokens
)
.into());
}
let mut prefix_tokens = Vec::with_capacity(config.num_prefix_tokens() * hidden);
prefix_tokens.extend_from_slice(&cls);
prefix_tokens.extend_from_slice(®isters);
session.set_parameter("prefix_tokens", &prefix_tokens);
// --- Encoder layers ---
for i in 0..config.num_hidden_layers {
let src = format!("{prefix}{i}");
let dst = format!("layer.{i}");
for norm in ["norm1", "norm2"] {
for part in ["weight", "bias"] {
session.set_parameter(
&format!("{dst}.{norm}.{part}"),
&model.tensor_f32_auto(&format!("{src}.{norm}.{part}"))?,
);
}
}
// K has no bias — `key_bias: false` in the config.
for (proj, has_bias) in [
("q_proj", true),
("k_proj", false),
("v_proj", true),
("o_proj", true),
] {
session.set_parameter(
&format!("{dst}.attention.{proj}.weight"),
&model.tensor_f32_auto_transposed(&format!("{src}.attention.{proj}.weight"))?,
);
if has_bias {
session.set_parameter(
&format!("{dst}.attention.{proj}.bias"),
&model.tensor_f32_auto(&format!("{src}.attention.{proj}.bias"))?,
);
}
}
for proj in ["up_proj", "down_proj"] {
session.set_parameter(
&format!("{dst}.mlp.{proj}.weight"),
&model.tensor_f32_auto_transposed(&format!("{src}.mlp.{proj}.weight"))?,
);
session.set_parameter(
&format!("{dst}.mlp.{proj}.bias"),
&model.tensor_f32_auto(&format!("{src}.mlp.{proj}.bias"))?,
);
}
for ls in ["layer_scale1", "layer_scale2"] {
session.set_parameter(
&format!("{dst}.{ls}.lambda1"),
&model.tensor_f32_auto(&format!("{src}.{ls}.lambda1"))?,
);
}
}
// --- Final norm ---
for part in ["weight", "bias"] {
session.set_parameter(
&format!("norm.{part}"),
&embedding_tensor(model, &format!("norm.{part}"))?,
);
}
log::info!(
"loaded DINOv3 weights: {} layers, hidden {hidden}",
config.num_hidden_layers
);
Ok(())
}
|