Add files using upload-large-folder tool

olmo3/1b/stage1/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dea5562210667dc24ecec56159f05b9f8f2245aa83d64bf31e95dd02fff39695
+size 2969458432

olmo3/1b/stage2/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:226d1174f58fccc681a916163eabdfb7c5c9c15990fa4f7e4d3b88d721f9f849
+size 2969458432

olmo3/1b/stage3/__init__.py

@@ -0,0 +1,2 @@
+from .configuration_olmo3_siamese_depth import Olmo3SiameseDepthConfig
+from .modeling_olmo3_siamese_depth import Olmo3SiameseDepthForCausalLM, Olmo3SiameseDepthModel

olmo3/1b/stage3/config.json

@@ -0,0 +1,76 @@
+{
+  "_name_or_path": "",
+  "architectures": [
+    "Olmo3SiameseDepthForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_olmo3_siamese_depth.Olmo3SiameseDepthConfig",
+    "AutoModel": "modeling_olmo3_siamese_depth.Olmo3SiameseDepthModel",
+    "AutoModelForCausalLM": "modeling_olmo3_siamese_depth.Olmo3SiameseDepthForCausalLM"
+  },
+  "bos_token_id": null,
+  "depth_attention_recent_window": 0,
+  "depth_attention_stride": 8,
+  "eos_token_id": 100257,
+  "hidden_act": "silu",
+  "hidden_dropout": 0.0,
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 8192,
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 65536,
+  "max_sequence_length": 65536,
+  "mlp_bias": false,
+  "model_type": "olmo3_siamese_depth",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 16,
+  "olmo3_weight_decay": true,
+  "pad_token_id": 100277,
+  "padded_vocab_size": 100352,
+  "qk_norm": true,
+  "qk_norm_mode": "full_projection",
+  "rms_norm_eps": 1e-06,
+  "rope_full_precision": true,
+  "rope_scaling": {
+    "beta_fast": 32.0,
+    "beta_slow": 1.0,
+    "factor": 8.0,
+    "original_max_position_embeddings": 8192,
+    "rope_type": "yarn"
+  },
+  "rope_theta": 500000.0,
+  "scale_embeds": false,
+  "scale_output_layer_init": false,
+  "siamese_norm_variant": "hybrid_pre",
+  "sliding_window": 4096,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.57.6",
+  "true_vocab_size": 100278,
+  "truncated_normal_factor": 3.0,
+  "use_cache": true,
+  "use_depth_attention": true,
+  "use_siamese_norm": true,
+  "vocab_size": 100278,
+  "vocab_z_loss_coeff": 1e-05
+}

olmo3/1b/stage3/configuration_olmo3_siamese_depth.py

@@ -0,0 +1,78 @@
+"""Hugging Face configuration for the OLMo 3 Siamese-Norm/Depth-Attention model."""
+
+from __future__ import annotations
+
+from transformers.models.olmo3.configuration_olmo3 import Olmo3Config
+
+
+class Olmo3SiameseDepthConfig(Olmo3Config):
+    """OLMo 3 plus the checkpoint-compatible Siamese/Depth extensions."""
+
+    model_type = "olmo3_siamese_depth"
+
+    def __init__(
+        self,
+        *args,
+        vocab_size: int | None = None,
+        true_vocab_size: int = 100278,
+        padded_vocab_size: int = 100352,
+        qk_norm_mode: str = "full_projection",
+        use_siamese_norm: bool = True,
+        siamese_norm_variant: str = "hybrid_pre",
+        use_depth_attention: bool = True,
+        depth_attention_stride: int = 8,
+        depth_attention_recent_window: int = 0,
+        rope_full_precision: bool = True,
+        **kwargs,
+    ):
+        # ``PretrainedConfig.to_diff_dict()`` constructs a no-argument instance
+        # of this class.  Make that default instance internally consistent
+        # instead of inheriting OLMo3's unrelated 50304-token default.
+        if vocab_size is None:
+            vocab_size = true_vocab_size
+        super().__init__(*args, vocab_size=vocab_size, **kwargs)
+        self.true_vocab_size = int(true_vocab_size)
+        self.padded_vocab_size = int(padded_vocab_size)
+        self.qk_norm = True
+        self.qk_norm_mode = qk_norm_mode
+        self.use_siamese_norm = bool(use_siamese_norm)
+        self.siamese_norm_variant = siamese_norm_variant
+        self.use_depth_attention = bool(use_depth_attention)
+        self.depth_attention_stride = int(depth_attention_stride)
+        self.depth_attention_recent_window = int(depth_attention_recent_window)
+        self.rope_full_precision = bool(rope_full_precision)
+        self._validate_siamese_depth()
+
+    def _validate_siamese_depth(self) -> None:
+        if self.vocab_size != self.true_vocab_size:
+            raise ValueError(
+                "HF vocab_size must equal true_vocab_size after padded-row removal; "
+                f"got {self.vocab_size} and {self.true_vocab_size}."
+            )
+        if self.padded_vocab_size < self.true_vocab_size:
+            raise ValueError("padded_vocab_size cannot be smaller than true_vocab_size.")
+        if self.qk_norm_mode != "full_projection":
+            raise ValueError("qk_norm_mode must be 'full_projection'.")
+        if not self.use_siamese_norm or self.siamese_norm_variant != "hybrid_pre":
+            raise ValueError("This remote model requires Hybrid-Pre Siamese Norm.")
+        if not self.use_depth_attention:
+            raise ValueError("This remote model requires Depth Attention.")
+        if self.depth_attention_stride < 1:
+            raise ValueError("depth_attention_stride must be positive.")
+        if self.depth_attention_recent_window < 0:
+            raise ValueError("depth_attention_recent_window must be non-negative.")
+        if not self.rope_full_precision:
+            raise ValueError("OLMo 3 requires FP32 Q/K RoPE.")
+        if self.hidden_size % self.num_attention_heads:
+            raise ValueError("hidden_size must be divisible by num_attention_heads.")
+        if self.num_attention_heads % self.num_key_value_heads:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads.")
+        expected_layer_types = [
+            "sliding_attention" if (index + 1) % 4 else "full_attention"
+            for index in range(self.num_hidden_layers)
+        ]
+        if list(self.layer_types) != expected_layer_types:
+            raise ValueError("OLMo 3 requires the repeating SWA,SWA,SWA,Full pattern.")
+
+
+__all__ = ["Olmo3SiameseDepthConfig"]

olmo3/1b/stage3/generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": null,
+  "eos_token_id": 100257,
+  "pad_token_id": 100277,
+  "transformers_version": "4.57.6"
+}

olmo3/1b/stage3/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d380f9b5ff781d7c10089db0559215ad2da91443100d33e8e6fd677ff219622
+size 2969458432

olmo3/1b/stage3/modeling_olmo3_siamese_depth.py

@@ -0,0 +1,548 @@
+"""Transformers remote code for OLMo 3 + Siamese Norm + Depth Attention.
+
+The Full/SWA masks, Full-only YaRN rotary embeddings, and bounded sliding KV
+cache are delegated to the official Transformers OLMo 3 implementation.  This
+file adds only the two checkpoint-persistent mechanisms used during training:
+Hybrid-Pre Siamese Norm and recursive same-position Depth Attention.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import torch
+import torch.nn as nn
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.models.olmo3.modeling_olmo3 import (
+    Olmo3MLP,
+    Olmo3PreTrainedModel,
+    Olmo3RMSNorm,
+    Olmo3RotaryEmbedding,
+    apply_rotary_pos_emb,
+    eager_attention_forward,
+)
+
+from .configuration_olmo3_siamese_depth import Olmo3SiameseDepthConfig
+
+
+@dataclass
+class _DepthSource:
+    layer_index: int
+    key: torch.Tensor
+    value: torch.Tensor
+
+
+class _DepthContext:
+    def __init__(self, config: Olmo3SiameseDepthConfig):
+        self.num_layers = int(config.num_hidden_layers)
+        self.stride = int(config.depth_attention_stride)
+        self.recent_window = int(config.depth_attention_recent_window)
+        self.records: dict[int, _DepthSource] = {}
+
+    def sources_for(self, layer_index: int) -> tuple[_DepthSource, ...]:
+        recent_start = max(0, layer_index - self.recent_window)
+        return tuple(
+            self.records[index]
+            for index in sorted(self.records)
+            if index < layer_index
+            and (index % self.stride == 0 or index >= recent_start)
+        )
+
+    def record(self, layer_index: int, key: torch.Tensor, value: torch.Tensor) -> None:
+        if layer_index in self.records:
+            raise RuntimeError(f"Depth Attention layer {layer_index} was recorded twice.")
+        self.records[layer_index] = _DepthSource(layer_index, key, value)
+        next_recent_start = layer_index + 1 - self.recent_window
+        stale = [
+            index
+            for index in self.records
+            if index % self.stride != 0 and index < next_recent_start
+        ]
+        for index in stale:
+            del self.records[index]
+
+
+def _depth_attention_mix(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    sources: tuple[_DepthSource, ...],
+) -> torch.Tensor:
+    """Mix current V with recursively mixed V from selected earlier layers.
+
+    Tensors use Transformers' ``[batch, heads, sequence, head_dim]`` layout.
+    """
+
+    if not sources:
+        return value
+    query_heads = query.shape[1]
+    kv_heads = key.shape[1]
+    if query_heads % kv_heads:
+        raise ValueError("Depth Attention requires Q heads divisible by KV heads.")
+    group_size = query_heads // kv_heads
+    grouped_query = (
+        query
+        if group_size == 1
+        else query.reshape(
+            query.shape[0],
+            kv_heads,
+            group_size,
+            query.shape[2],
+            query.shape[3],
+        ).mean(dim=2)
+    )
+    scale = 1.0 / math.sqrt(query.shape[-1])
+    scores = [
+        (grouped_query * source.key).sum(dim=-1, keepdim=True) * scale
+        for source in sources
+    ]
+    scores.append((grouped_query * key).sum(dim=-1, keepdim=True) * scale)
+    weights = torch.softmax(torch.cat(scores, dim=-1), dim=-1, dtype=torch.float32)
+    weights = weights.to(dtype=value.dtype)
+    mixed = weights[..., -1:].mul(value)
+    for source_index, source in enumerate(sources):
+        mixed = mixed + weights[..., source_index : source_index + 1] * source.value
+    return mixed
+
+
+class Olmo3SiameseDepthAttention(nn.Module):
+    def __init__(self, config: Olmo3SiameseDepthConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = int(layer_idx)
+        self.head_dim = getattr(
+            config, "head_dim", config.hidden_size // config.num_attention_heads
+        )
+        self.num_key_value_groups = (
+            config.num_attention_heads // config.num_key_value_heads
+        )
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = float(config.attention_dropout)
+        self.is_causal = True
+        self.q_proj = nn.Linear(
+            config.hidden_size,
+            config.num_attention_heads * self.head_dim,
+            bias=False,
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=False,
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=False,
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim,
+            config.hidden_size,
+            bias=False,
+        )
+        # These norms cover the complete Q and K projections, not one head.
+        self.q_norm = Olmo3RMSNorm(
+            config.num_attention_heads * self.head_dim, config.rms_norm_eps
+        )
+        self.k_norm = Olmo3RMSNorm(
+            config.num_key_value_heads * self.head_dim, config.rms_norm_eps
+        )
+        self.attention_type = config.layer_types[self.layer_idx]
+        self.sliding_window = (
+            int(config.sliding_window)
+            if self.attention_type == "sliding_attention"
+            else None
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        depth_context: _DepthContext,
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        output_attentions: bool = False,
+        **kwargs,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        batch_size, sequence_length, _ = hidden_states.shape
+        query = self.q_norm(self.q_proj(hidden_states)).view(
+            batch_size, sequence_length, self.config.num_attention_heads, self.head_dim
+        ).transpose(1, 2)
+        key = self.k_norm(self.k_proj(hidden_states)).view(
+            batch_size, sequence_length, self.config.num_key_value_heads, self.head_dim
+        ).transpose(1, 2)
+        value = self.v_proj(hidden_states).view(
+            batch_size, sequence_length, self.config.num_key_value_heads, self.head_dim
+        ).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query, key = apply_rotary_pos_emb(query.float(), key.float(), cos.float(), sin.float())
+        query = query.to(dtype=hidden_states.dtype)
+        key = key.to(dtype=hidden_states.dtype)
+
+        mixed_value = _depth_attention_mix(
+            query,
+            key,
+            value,
+            depth_context.sources_for(self.layer_idx),
+        )
+        depth_context.record(self.layer_idx, key, mixed_value)
+
+        attention_key = key
+        attention_value = mixed_value
+        if past_key_values is not None:
+            cache_kwargs = {
+                "sin": sin,
+                "cos": cos,
+                "cache_position": cache_position,
+            }
+            attention_key, attention_value = past_key_values.update(
+                key, mixed_value, self.layer_idx, cache_kwargs
+            )
+
+        attention_interface = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[
+                self.config._attn_implementation
+            ]
+        attention_output, attention_weights = attention_interface(
+            self,
+            query,
+            attention_key,
+            attention_value,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+        attention_output = attention_output.reshape(
+            batch_size, sequence_length, -1
+        ).contiguous()
+        attention_output = self.o_proj(attention_output)
+        return attention_output, attention_weights if output_attentions else None
+
+
+class Olmo3SiameseDepthDecoderLayer(nn.Module):
+    def __init__(self, config: Olmo3SiameseDepthConfig, layer_idx: int):
+        super().__init__()
+        self.layer_idx = int(layer_idx)
+        self.self_attn = Olmo3SiameseDepthAttention(config, layer_idx)
+        self.mlp = Olmo3MLP(config)
+        self.pre_mlp_layernorm = Olmo3RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.siamese_attn_pre_norm = Olmo3RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.siamese_mlp_pre_norm = Olmo3RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.siamese_mlp_input_norm = Olmo3RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.siamese_hybrid_attn_scale = nn.Parameter(
+            torch.ones(config.hidden_size)
+        )
+        self.post_residual_scale = 1.0 / math.sqrt(2.0 * (self.layer_idx + 1))
+
+    def forward(
+        self,
+        post_stream: torch.Tensor,
+        pre_stream: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        depth_context: _DepthContext,
+        past_key_values: Optional[Cache],
+        cache_position: Optional[torch.LongTensor],
+        output_attentions: bool,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        pre_normalized = self.siamese_attn_pre_norm(pre_stream)
+        attention_input = torch.addcmul(
+            pre_normalized,
+            post_stream,
+            self.siamese_hybrid_attn_scale.to(dtype=post_stream.dtype),
+        )
+        attention_output, attention_weights = self.self_attn(
+            hidden_states=attention_input,
+            position_embeddings=position_embeddings,
+            attention_mask=attention_mask,
+            depth_context=depth_context,
+            past_key_values=past_key_values,
+            cache_position=cache_position,
+            output_attentions=output_attentions,
+            **kwargs,
+        )
+        post_stream = torch.add(
+            post_stream, attention_output, alpha=self.post_residual_scale
+        )
+        pre_stream = pre_stream + attention_output
+
+        # The native Hybrid-Pre block advances its post stream to the
+        # pre-MLP-normalized value before both the MLP input construction and
+        # the MLP residual update.  This is intentionally not a conventional
+        # pre-norm residual that adds the branch back to the unnormalized
+        # stream.
+        post_stream = self.pre_mlp_layernorm(post_stream)
+        pre_normalized = self.siamese_mlp_pre_norm(pre_stream)
+        mlp_input = self.siamese_mlp_input_norm(
+            post_stream + pre_normalized
+        )
+        mlp_output = self.mlp(mlp_input)
+        post_stream = torch.add(
+            post_stream, mlp_output, alpha=self.post_residual_scale
+        )
+        pre_stream = pre_stream + mlp_output
+        return post_stream, pre_stream, attention_weights
+
+
+class Olmo3SiameseDepthPreTrainedModel(Olmo3PreTrainedModel):
+    config_class = Olmo3SiameseDepthConfig
+    base_model_prefix = "model"
+    _no_split_modules = ["Olmo3SiameseDepthDecoderLayer"]
+    # Whole-layer replay is invalid for the mutable cross-layer Depth context.
+    # The native trainer supports only its dedicated selective-MLP checkpoint
+    # path; do not advertise Transformers' generic layer checkpointing.
+    supports_gradient_checkpointing = False
+
+
+class Olmo3SiameseDepthModel(Olmo3SiameseDepthPreTrainedModel):
+    def __init__(self, config: Olmo3SiameseDepthConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [
+                Olmo3SiameseDepthDecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.siamese_post_final_layernorm = Olmo3RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.siamese_pre_final_layernorm = Olmo3RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.norm = Olmo3RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.rotary_embs = nn.ModuleDict(
+            {
+                "sliding_attention": Olmo3RotaryEmbedding(
+                    config=config, rope_type="default"
+                ),
+                "full_attention": Olmo3RotaryEmbedding(config=config),
+            }
+        )
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[tuple, BaseModelOutputWithPast]:
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Specify exactly one of input_ids or inputs_embeds.")
+        output_attentions = (
+            self.config.output_attentions
+            if output_attentions is None
+            else output_attentions
+        )
+        output_hidden_states = (
+            self.config.output_hidden_states
+            if output_hidden_states is None
+            else output_hidden_states
+        )
+        return_dict = self.config.use_return_dict if return_dict is None else return_dict
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+        if cache_position is None:
+            past_seen_tokens = (
+                past_key_values.get_seq_length()
+                if past_key_values is not None
+                else 0
+            )
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        if not isinstance(attention_mask, dict):
+            mask_kwargs = {
+                "config": self.config,
+                "input_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+                "position_ids": position_ids,
+            }
+            causal_masks = {
+                "full_attention": create_causal_mask(**mask_kwargs),
+                "sliding_attention": create_sliding_window_causal_mask(
+                    **mask_kwargs
+                ),
+            }
+        else:
+            causal_masks = attention_mask
+
+        post_stream = inputs_embeds
+        # The optimized training path intentionally aliases both initial streams.
+        pre_stream = inputs_embeds
+        depth_context = _DepthContext(self.config)
+        position_embeddings = {
+            kind: rotary(post_stream, position_ids)
+            for kind, rotary in self.rotary_embs.items()
+        }
+        hidden_history = () if output_hidden_states else None
+        attention_history = () if output_attentions else None
+        for layer in self.layers:
+            if output_hidden_states:
+                hidden_history += (post_stream,)
+            attention_type = layer.self_attn.attention_type
+            post_stream, pre_stream, attention_weights = layer(
+                post_stream=post_stream,
+                pre_stream=pre_stream,
+                attention_mask=causal_masks[attention_type],
+                position_embeddings=position_embeddings[attention_type],
+                depth_context=depth_context,
+                past_key_values=past_key_values,
+                cache_position=cache_position,
+                output_attentions=output_attentions,
+                **kwargs,
+            )
+            if output_attentions:
+                attention_history += (attention_weights,)
+
+        hidden_states = (
+            self.siamese_post_final_layernorm(post_stream)
+            + self.siamese_pre_final_layernorm(pre_stream)
+        )
+        hidden_states = self.norm(hidden_states)
+        if output_hidden_states:
+            hidden_history += (hidden_states,)
+        if not return_dict:
+            values = (hidden_states, past_key_values, hidden_history, attention_history)
+            return tuple(value for value in values if value is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+            hidden_states=hidden_history,
+            attentions=attention_history,
+        )
+
+
+class Olmo3SiameseDepthForCausalLM(
+    Olmo3SiameseDepthPreTrainedModel, GenerationMixin
+):
+    _tied_weights_keys = []
+
+    def __init__(self, config: Olmo3SiameseDepthConfig):
+        super().__init__(config)
+        self.model = Olmo3SiameseDepthModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[tuple, CausalLMOutputWithPast]:
+        return_dict = self.config.use_return_dict if return_dict is None else return_dict
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            **kwargs,
+        )
+        hidden_states = outputs.last_hidden_state
+        indices = (
+            slice(-logits_to_keep, None)
+            if isinstance(logits_to_keep, int)
+            else logits_to_keep
+        )
+        logits = self.lm_head(hidden_states[:, indices, :])
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                vocab_size=self.config.vocab_size,
+                **kwargs,
+            )
+        if not return_dict:
+            result = (logits, outputs.past_key_values)
+            return ((loss,) + result) if loss is not None else result
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = [
+    "Olmo3SiameseDepthForCausalLM",
+    "Olmo3SiameseDepthModel",
+    "Olmo3SiameseDepthPreTrainedModel",
+]

olmo3/1b/stage3/special_tokens_map.json

@@ -0,0 +1,11 @@
+{
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|pad|>",
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

olmo3/1b/stage3/tokenizer_config.json

@@ -0,0 +1,190 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "100256": {
+      "content": "<|extra_id_0|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100257": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100258": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100259": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100260": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100261": {
+      "content": "|||PHONE_NUMBER|||",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100262": {
+      "content": "|||EMAIL_ADDRESS|||",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100263": {
+      "content": "|||IP_ADDRESS|||",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100264": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100265": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100266": {
+      "content": "<|extra_id_1|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100267": {
+      "content": "<|extra_id_2|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100268": {
+      "content": "<|extra_id_3|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100269": {
+      "content": "<|extra_id_4|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100270": {
+      "content": "<|extra_id_5|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100271": {
+      "content": "<|extra_id_6|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100272": {
+      "content": "<|extra_id_7|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100273": {
+      "content": "<|extra_id_8|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100274": {
+      "content": "<|extra_id_9|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100275": {
+      "content": "<|extra_id_10|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100276": {
+      "content": "<|endofprompt|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100277": {
+      "content": "<|pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "fix_mistral_regex": false,
+  "model_max_length": 65536,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}

olmo3/1b/stage4/instruct/PUBLIC_FILES.SHA256

@@ -0,0 +1,13 @@
+6c77801731dc5611cff8478138ef3fd1d8bdd37e2b992aa034d5da66fec61eca  README.md
+e2144e7990128799f6eaf6430ea0a472dbb17e076bdb08b327887e59f84657be  __init__.py
+a3d6a7ad3fde8e26fa9953d80ac30682f7d56ac1d03f3428910a4239a4e23bef  chat_template.jinja
+77c545f18298aa5960d58d0c83eb752b8559ebac016f5755a3aef50b226588c8  config.json
+d66ccce18efda42443a91536bf907e7bcb6e2d0acd0c8dc70d69eb15f277d634  configuration_olmo3_siamese_depth.py
+dbd4aa5dc6f7f1bdc0c181d9e5582f6571e1d2a22878dfcd98149324b9f5c4a2  generation_config.json
+b6fe424e334903f7fb84d3a106d9730455f4744b9fe3c21ee136d97a00e72502  merges.txt
+f0b8b78613e1b49a5f6164acefec39e06f7f938176ff23663f703802799e2497  model.safetensors
+12329a1f6cb1732bf100bda2ee2dd62761f6485aefb0a977879eb6d74f764aa6  modeling_olmo3_siamese_depth.py
+78afb564e81264029b25f9caf24bda2521d5bdaeff5cd3fdbc01d3da2e8ce2f2  special_tokens_map.json
+7738a25c46a6043a0e2e605138851baa450269ce9d800af6ca9bf6192e45510c  tokenizer.json
+5f959b9c0cff767a63f22a0569e0554cce28214fb49e41479ba326f44169c5ec  tokenizer_config.json
+d8eb44d0d51f5e39c9b14825ca8bd10a2500f99c438ca81d4462f0263acd3b53  vocab.json

olmo3/1b/stage4/instruct/README.md

@@ -0,0 +1,6 @@
+# Stage 4 Instruct SFT checkpoint
+
+This subfolder contains a complete Hugging Face checkpoint with a context
+length of 65536 tokens. Load it from the parent model repository using
+`trust_remote_code=True`, `subfolder="olmo3/1b/stage4/instruct"`, and
+`attn_implementation="sdpa"`.

olmo3/1b/stage4/instruct/__init__.py

@@ -0,0 +1,2 @@
+from .configuration_olmo3_siamese_depth import Olmo3SiameseDepthConfig
+from .modeling_olmo3_siamese_depth import Olmo3SiameseDepthForCausalLM, Olmo3SiameseDepthModel

olmo3/1b/stage4/instruct/chat_template.jinja

@@ -0,0 +1,16 @@
+{%- set has_system = messages|selectattr('role', 'equalto', 'system')|list|length > 0 -%}{%- if not has_system -%}{{- '<|im_start|>system
+You are a helpful function-calling AI assistant. ' -}}{%- if tools is none or (tools | length) == 0 -%}{{- 'You do not currently have access to any functions. <functions></functions><|im_end|>
+' -}}{%- else -%}{{- 'You are provided with function signatures within <functions></functions> XML tags. You may call one or more functions to assist with the user query. Output any function calls within <function_calls></function_calls> XML tags. Do not make assumptions about what values to plug into functions.' -}}{{- '<functions>' -}}{{- tools | tojson -}}{{- '</functions><|im_end|>
+' -}}{%- endif -%}{%- endif -%}{%- for message in messages -%}{%- if message['role'] == 'system' -%}{{- '<|im_start|>system
+' + message['content'] -}}{%- if tools is not none -%}{{- '<functions>' -}}{{- tools | tojson -}}{{- '</functions>' -}}{%- elif message.get('functions', none) is not none -%}{{- ' <functions>' + message['functions'] + '</functions>' -}}{%- endif -%}{{- '<|im_end|>
+' -}}{%- elif message['role'] == 'user' -%}{{- '<|im_start|>user
+' + message['content'] + '<|im_end|>
+' -}}{%- elif message['role'] == 'assistant' -%}{{- '<|im_start|>assistant
+' -}}{%- if message.get('content', none) is not none -%}{{- message['content'] -}}{%- endif -%}{%- if message.get('function_calls', none) is not none and (message['function_calls'] | length) > 0 -%}{{- '<function_calls>' + message['function_calls'] + '</function_calls>' -}}{%- elif message.get('tool_calls', none) is not none and (message['tool_calls'] | length) > 0 -%}{{- '<function_calls>' -}}{%- for tool_call in message['tool_calls'] %}{%- if tool_call is mapping and tool_call.get('function', none) is not none %}{%- set args = tool_call['function']['arguments'] -%}{%- set ns = namespace(arguments_list=[]) -%}{%- for key, value in args.items() -%}{%- set ns.arguments_list = ns.arguments_list + [key ~ '=' ~ (value | tojson)] -%}{%- endfor -%}{%- set arguments = ns.arguments_list | join(', ') -%}{{- tool_call['function']['name'] + '(' + arguments + ')' -}}{%- if not loop.last -%}{{ '
+' }}{%- endif -%}{% else %}{{- tool_call -}}{%- endif %}{%- endfor %}{{- '</function_calls>' -}}{%- endif -%}{%- if not loop.last -%}{{- '<|im_end|>' + '
+' -}}{%- else -%}{{- eos_token -}}{%- endif -%}{%- elif message['role'] == 'environment' -%}{{- '<|im_start|>environment
+' + message['content'] + '<|im_end|>
+' -}}{%- elif message['role'] == 'tool' -%}{{- '<|im_start|>environment
+' + message['content'] + '<|im_end|>
+' -}}{%- endif -%}{%- if loop.last and add_generation_prompt -%}{{- '<|im_start|>assistant
+' -}}{%- endif -%}{%- endfor -%}

olmo3/1b/stage4/instruct/config.json

@@ -0,0 +1,76 @@
+{
+  "_name_or_path": "",
+  "architectures": [
+    "Olmo3SiameseDepthForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_olmo3_siamese_depth.Olmo3SiameseDepthConfig",
+    "AutoModel": "modeling_olmo3_siamese_depth.Olmo3SiameseDepthModel",
+    "AutoModelForCausalLM": "modeling_olmo3_siamese_depth.Olmo3SiameseDepthForCausalLM"
+  },
+  "bos_token_id": null,
+  "depth_attention_recent_window": 0,
+  "depth_attention_stride": 8,
+  "eos_token_id": 100257,
+  "hidden_act": "silu",
+  "hidden_dropout": 0.0,
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 8192,
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 65536,
+  "max_sequence_length": 65536,
+  "mlp_bias": false,
+  "model_type": "olmo3_siamese_depth",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 16,
+  "olmo3_weight_decay": true,
+  "pad_token_id": 100277,
+  "padded_vocab_size": 100352,
+  "qk_norm": true,
+  "qk_norm_mode": "full_projection",
+  "rms_norm_eps": 1e-06,
+  "rope_full_precision": true,
+  "rope_scaling": {
+    "beta_fast": 32.0,
+    "beta_slow": 1.0,
+    "factor": 8.0,
+    "original_max_position_embeddings": 8192,
+    "rope_type": "yarn"
+  },
+  "rope_theta": 500000.0,
+  "scale_embeds": false,
+  "scale_output_layer_init": false,
+  "siamese_norm_variant": "hybrid_pre",
+  "sliding_window": 4096,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.57.6",
+  "true_vocab_size": 100278,
+  "truncated_normal_factor": 3.0,
+  "use_cache": true,
+  "use_depth_attention": true,
+  "use_siamese_norm": true,
+  "vocab_size": 100278,
+  "vocab_z_loss_coeff": 1e-05
+}

olmo3/1b/stage4/instruct/configuration_olmo3_siamese_depth.py

@@ -0,0 +1,78 @@
+"""Hugging Face configuration for the OLMo 3 Siamese-Norm/Depth-Attention model."""
+
+from __future__ import annotations
+
+from transformers.models.olmo3.configuration_olmo3 import Olmo3Config
+
+
+class Olmo3SiameseDepthConfig(Olmo3Config):
+    """OLMo 3 plus the checkpoint-compatible Siamese/Depth extensions."""
+
+    model_type = "olmo3_siamese_depth"
+
+    def __init__(
+        self,
+        *args,
+        vocab_size: int | None = None,
+        true_vocab_size: int = 100278,
+        padded_vocab_size: int = 100352,
+        qk_norm_mode: str = "full_projection",
+        use_siamese_norm: bool = True,
+        siamese_norm_variant: str = "hybrid_pre",
+        use_depth_attention: bool = True,
+        depth_attention_stride: int = 8,
+        depth_attention_recent_window: int = 0,
+        rope_full_precision: bool = True,
+        **kwargs,
+    ):
+        # ``PretrainedConfig.to_diff_dict()`` constructs a no-argument instance
+        # of this class.  Make that default instance internally consistent
+        # instead of inheriting OLMo3's unrelated 50304-token default.
+        if vocab_size is None:
+            vocab_size = true_vocab_size
+        super().__init__(*args, vocab_size=vocab_size, **kwargs)
+        self.true_vocab_size = int(true_vocab_size)
+        self.padded_vocab_size = int(padded_vocab_size)
+        self.qk_norm = True
+        self.qk_norm_mode = qk_norm_mode
+        self.use_siamese_norm = bool(use_siamese_norm)
+        self.siamese_norm_variant = siamese_norm_variant
+        self.use_depth_attention = bool(use_depth_attention)
+        self.depth_attention_stride = int(depth_attention_stride)
+        self.depth_attention_recent_window = int(depth_attention_recent_window)
+        self.rope_full_precision = bool(rope_full_precision)
+        self._validate_siamese_depth()
+
+    def _validate_siamese_depth(self) -> None:
+        if self.vocab_size != self.true_vocab_size:
+            raise ValueError(
+                "HF vocab_size must equal true_vocab_size after padded-row removal; "
+                f"got {self.vocab_size} and {self.true_vocab_size}."
+            )
+        if self.padded_vocab_size < self.true_vocab_size:
+            raise ValueError("padded_vocab_size cannot be smaller than true_vocab_size.")
+        if self.qk_norm_mode != "full_projection":
+            raise ValueError("qk_norm_mode must be 'full_projection'.")
+        if not self.use_siamese_norm or self.siamese_norm_variant != "hybrid_pre":
+            raise ValueError("This remote model requires Hybrid-Pre Siamese Norm.")
+        if not self.use_depth_attention:
+            raise ValueError("This remote model requires Depth Attention.")
+        if self.depth_attention_stride < 1:
+            raise ValueError("depth_attention_stride must be positive.")
+        if self.depth_attention_recent_window < 0:
+            raise ValueError("depth_attention_recent_window must be non-negative.")
+        if not self.rope_full_precision:
+            raise ValueError("OLMo 3 requires FP32 Q/K RoPE.")
+        if self.hidden_size % self.num_attention_heads:
+            raise ValueError("hidden_size must be divisible by num_attention_heads.")
+        if self.num_attention_heads % self.num_key_value_heads:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads.")
+        expected_layer_types = [
+            "sliding_attention" if (index + 1) % 4 else "full_attention"
+            for index in range(self.num_hidden_layers)
+        ]
+        if list(self.layer_types) != expected_layer_types:
+            raise ValueError("OLMo 3 requires the repeating SWA,SWA,SWA,Full pattern.")
+
+
+__all__ = ["Olmo3SiameseDepthConfig"]

olmo3/1b/stage4/instruct/generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": null,
+  "eos_token_id": 100257,
+  "pad_token_id": 100277,
+  "transformers_version": "4.57.6"
+}

olmo3/1b/stage4/instruct/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f0b8b78613e1b49a5f6164acefec39e06f7f938176ff23663f703802799e2497
+size 2969458432

olmo3/1b/stage4/instruct/modeling_olmo3_siamese_depth.py

@@ -0,0 +1,548 @@
+"""Transformers remote code for OLMo 3 + Siamese Norm + Depth Attention.
+
+The Full/SWA masks, Full-only YaRN rotary embeddings, and bounded sliding KV
+cache are delegated to the official Transformers OLMo 3 implementation.  This
+file adds only the two checkpoint-persistent mechanisms used during training:
+Hybrid-Pre Siamese Norm and recursive same-position Depth Attention.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import torch
+import torch.nn as nn
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.models.olmo3.modeling_olmo3 import (
+    Olmo3MLP,
+    Olmo3PreTrainedModel,
+    Olmo3RMSNorm,
+    Olmo3RotaryEmbedding,
+    apply_rotary_pos_emb,
+    eager_attention_forward,
+)
+
+from .configuration_olmo3_siamese_depth import Olmo3SiameseDepthConfig
+
+
+@dataclass
+class _DepthSource:
+    layer_index: int
+    key: torch.Tensor
+    value: torch.Tensor
+
+
+class _DepthContext:
+    def __init__(self, config: Olmo3SiameseDepthConfig):
+        self.num_layers = int(config.num_hidden_layers)
+        self.stride = int(config.depth_attention_stride)
+        self.recent_window = int(config.depth_attention_recent_window)
+        self.records: dict[int, _DepthSource] = {}
+
+    def sources_for(self, layer_index: int) -> tuple[_DepthSource, ...]:
+        recent_start = max(0, layer_index - self.recent_window)
+        return tuple(
+            self.records[index]
+            for index in sorted(self.records)
+            if index < layer_index
+            and (index % self.stride == 0 or index >= recent_start)
+        )
+
+    def record(self, layer_index: int, key: torch.Tensor, value: torch.Tensor) -> None:
+        if layer_index in self.records:
+            raise RuntimeError(f"Depth Attention layer {layer_index} was recorded twice.")
+        self.records[layer_index] = _DepthSource(layer_index, key, value)
+        next_recent_start = layer_index + 1 - self.recent_window
+        stale = [
+            index
+            for index in self.records
+            if index % self.stride != 0 and index < next_recent_start
+        ]
+        for index in stale:
+            del self.records[index]
+
+
+def _depth_attention_mix(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    sources: tuple[_DepthSource, ...],
+) -> torch.Tensor:
+    """Mix current V with recursively mixed V from selected earlier layers.
+
+    Tensors use Transformers' ``[batch, heads, sequence, head_dim]`` layout.
+    """
+
+    if not sources:
+        return value
+    query_heads = query.shape[1]
+    kv_heads = key.shape[1]
+    if query_heads % kv_heads:
+        raise ValueError("Depth Attention requires Q heads divisible by KV heads.")
+    group_size = query_heads // kv_heads
+    grouped_query = (
+        query
+        if group_size == 1
+        else query.reshape(
+            query.shape[0],
+            kv_heads,
+            group_size,
+            query.shape[2],
+            query.shape[3],
+        ).mean(dim=2)
+    )
+    scale = 1.0 / math.sqrt(query.shape[-1])
+    scores = [
+        (grouped_query * source.key).sum(dim=-1, keepdim=True) * scale
+        for source in sources
+    ]
+    scores.append((grouped_query * key).sum(dim=-1, keepdim=True) * scale)
+    weights = torch.softmax(torch.cat(scores, dim=-1), dim=-1, dtype=torch.float32)
+    weights = weights.to(dtype=value.dtype)
+    mixed = weights[..., -1:].mul(value)
+    for source_index, source in enumerate(sources):
+        mixed = mixed + weights[..., source_index : source_index + 1] * source.value
+    return mixed
+
+
+class Olmo3SiameseDepthAttention(nn.Module):
+    def __init__(self, config: Olmo3SiameseDepthConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = int(layer_idx)
+        self.head_dim = getattr(
+            config, "head_dim", config.hidden_size // config.num_attention_heads
+        )
+        self.num_key_value_groups = (
+            config.num_attention_heads // config.num_key_value_heads
+        )
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = float(config.attention_dropout)
+        self.is_causal = True
+        self.q_proj = nn.Linear(
+            config.hidden_size,
+            config.num_attention_heads * self.head_dim,
+            bias=False,
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=False,
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=False,
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim,
+            config.hidden_size,
+            bias=False,
+        )
+        # These norms cover the complete Q and K projections, not one head.
+        self.q_norm = Olmo3RMSNorm(
+            config.num_attention_heads * self.head_dim, config.rms_norm_eps
+        )
+        self.k_norm = Olmo3RMSNorm(
+            config.num_key_value_heads * self.head_dim, config.rms_norm_eps
+        )
+        self.attention_type = config.layer_types[self.layer_idx]
+        self.sliding_window = (
+            int(config.sliding_window)
+            if self.attention_type == "sliding_attention"
+            else None
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        depth_context: _DepthContext,
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        output_attentions: bool = False,
+        **kwargs,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        batch_size, sequence_length, _ = hidden_states.shape
+        query = self.q_norm(self.q_proj(hidden_states)).view(
+            batch_size, sequence_length, self.config.num_attention_heads, self.head_dim
+        ).transpose(1, 2)
+        key = self.k_norm(self.k_proj(hidden_states)).view(
+            batch_size, sequence_length, self.config.num_key_value_heads, self.head_dim
+        ).transpose(1, 2)
+        value = self.v_proj(hidden_states).view(
+            batch_size, sequence_length, self.config.num_key_value_heads, self.head_dim
+        ).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query, key = apply_rotary_pos_emb(query.float(), key.float(), cos.float(), sin.float())
+        query = query.to(dtype=hidden_states.dtype)
+        key = key.to(dtype=hidden_states.dtype)
+
+        mixed_value = _depth_attention_mix(
+            query,
+            key,
+            value,
+            depth_context.sources_for(self.layer_idx),
+        )
+        depth_context.record(self.layer_idx, key, mixed_value)
+
+        attention_key = key
+        attention_value = mixed_value
+        if past_key_values is not None:
+            cache_kwargs = {
+                "sin": sin,
+                "cos": cos,
+                "cache_position": cache_position,
+            }
+            attention_key, attention_value = past_key_values.update(
+                key, mixed_value, self.layer_idx, cache_kwargs
+            )
+
+        attention_interface = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[
+                self.config._attn_implementation
+            ]
+        attention_output, attention_weights = attention_interface(
+            self,
+            query,
+            attention_key,
+            attention_value,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+        attention_output = attention_output.reshape(
+            batch_size, sequence_length, -1
+        ).contiguous()
+        attention_output = self.o_proj(attention_output)
+        return attention_output, attention_weights if output_attentions else None
+
+
+class Olmo3SiameseDepthDecoderLayer(nn.Module):
+    def __init__(self, config: Olmo3SiameseDepthConfig, layer_idx: int):
+        super().__init__()
+        self.layer_idx = int(layer_idx)
+        self.self_attn = Olmo3SiameseDepthAttention(config, layer_idx)
+        self.mlp = Olmo3MLP(config)
+        self.pre_mlp_layernorm = Olmo3RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.siamese_attn_pre_norm = Olmo3RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.siamese_mlp_pre_norm = Olmo3RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.siamese_mlp_input_norm = Olmo3RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.siamese_hybrid_attn_scale = nn.Parameter(
+            torch.ones(config.hidden_size)
+        )
+        self.post_residual_scale = 1.0 / math.sqrt(2.0 * (self.layer_idx + 1))
+
+    def forward(
+        self,
+        post_stream: torch.Tensor,
+        pre_stream: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        depth_context: _DepthContext,
+        past_key_values: Optional[Cache],
+        cache_position: Optional[torch.LongTensor],
+        output_attentions: bool,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        pre_normalized = self.siamese_attn_pre_norm(pre_stream)
+        attention_input = torch.addcmul(
+            pre_normalized,
+            post_stream,
+            self.siamese_hybrid_attn_scale.to(dtype=post_stream.dtype),
+        )
+        attention_output, attention_weights = self.self_attn(
+            hidden_states=attention_input,
+            position_embeddings=position_embeddings,
+            attention_mask=attention_mask,
+            depth_context=depth_context,
+            past_key_values=past_key_values,
+            cache_position=cache_position,
+            output_attentions=output_attentions,
+            **kwargs,
+        )
+        post_stream = torch.add(
+            post_stream, attention_output, alpha=self.post_residual_scale
+        )
+        pre_stream = pre_stream + attention_output
+
+        # The native Hybrid-Pre block advances its post stream to the
+        # pre-MLP-normalized value before both the MLP input construction and
+        # the MLP residual update.  This is intentionally not a conventional
+        # pre-norm residual that adds the branch back to the unnormalized
+        # stream.
+        post_stream = self.pre_mlp_layernorm(post_stream)
+        pre_normalized = self.siamese_mlp_pre_norm(pre_stream)
+        mlp_input = self.siamese_mlp_input_norm(
+            post_stream + pre_normalized
+        )
+        mlp_output = self.mlp(mlp_input)
+        post_stream = torch.add(
+            post_stream, mlp_output, alpha=self.post_residual_scale
+        )
+        pre_stream = pre_stream + mlp_output
+        return post_stream, pre_stream, attention_weights
+
+
+class Olmo3SiameseDepthPreTrainedModel(Olmo3PreTrainedModel):
+    config_class = Olmo3SiameseDepthConfig
+    base_model_prefix = "model"
+    _no_split_modules = ["Olmo3SiameseDepthDecoderLayer"]
+    # Whole-layer replay is invalid for the mutable cross-layer Depth context.
+    # The native trainer supports only its dedicated selective-MLP checkpoint
+    # path; do not advertise Transformers' generic layer checkpointing.
+    supports_gradient_checkpointing = False
+
+
+class Olmo3SiameseDepthModel(Olmo3SiameseDepthPreTrainedModel):
+    def __init__(self, config: Olmo3SiameseDepthConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [
+                Olmo3SiameseDepthDecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.siamese_post_final_layernorm = Olmo3RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.siamese_pre_final_layernorm = Olmo3RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.norm = Olmo3RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.rotary_embs = nn.ModuleDict(
+            {
+                "sliding_attention": Olmo3RotaryEmbedding(
+                    config=config, rope_type="default"
+                ),
+                "full_attention": Olmo3RotaryEmbedding(config=config),
+            }
+        )
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[tuple, BaseModelOutputWithPast]:
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Specify exactly one of input_ids or inputs_embeds.")
+        output_attentions = (
+            self.config.output_attentions
+            if output_attentions is None
+            else output_attentions
+        )
+        output_hidden_states = (
+            self.config.output_hidden_states
+            if output_hidden_states is None
+            else output_hidden_states
+        )
+        return_dict = self.config.use_return_dict if return_dict is None else return_dict
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+        if cache_position is None:
+            past_seen_tokens = (
+                past_key_values.get_seq_length()
+                if past_key_values is not None
+                else 0
+            )
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        if not isinstance(attention_mask, dict):
+            mask_kwargs = {
+                "config": self.config,
+                "input_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+                "position_ids": position_ids,
+            }
+            causal_masks = {
+                "full_attention": create_causal_mask(**mask_kwargs),
+                "sliding_attention": create_sliding_window_causal_mask(
+                    **mask_kwargs
+                ),
+            }
+        else:
+            causal_masks = attention_mask
+
+        post_stream = inputs_embeds
+        # The optimized training path intentionally aliases both initial streams.
+        pre_stream = inputs_embeds
+        depth_context = _DepthContext(self.config)
+        position_embeddings = {
+            kind: rotary(post_stream, position_ids)
+            for kind, rotary in self.rotary_embs.items()
+        }
+        hidden_history = () if output_hidden_states else None
+        attention_history = () if output_attentions else None
+        for layer in self.layers:
+            if output_hidden_states:
+                hidden_history += (post_stream,)
+            attention_type = layer.self_attn.attention_type
+            post_stream, pre_stream, attention_weights = layer(
+                post_stream=post_stream,
+                pre_stream=pre_stream,
+                attention_mask=causal_masks[attention_type],
+                position_embeddings=position_embeddings[attention_type],
+                depth_context=depth_context,
+                past_key_values=past_key_values,
+                cache_position=cache_position,
+                output_attentions=output_attentions,
+                **kwargs,
+            )
+            if output_attentions:
+                attention_history += (attention_weights,)
+
+        hidden_states = (
+            self.siamese_post_final_layernorm(post_stream)
+            + self.siamese_pre_final_layernorm(pre_stream)
+        )
+        hidden_states = self.norm(hidden_states)
+        if output_hidden_states:
+            hidden_history += (hidden_states,)
+        if not return_dict:
+            values = (hidden_states, past_key_values, hidden_history, attention_history)
+            return tuple(value for value in values if value is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+            hidden_states=hidden_history,
+            attentions=attention_history,
+        )
+
+
+class Olmo3SiameseDepthForCausalLM(
+    Olmo3SiameseDepthPreTrainedModel, GenerationMixin
+):
+    _tied_weights_keys = []
+
+    def __init__(self, config: Olmo3SiameseDepthConfig):
+        super().__init__(config)
+        self.model = Olmo3SiameseDepthModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[tuple, CausalLMOutputWithPast]:
+        return_dict = self.config.use_return_dict if return_dict is None else return_dict
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            **kwargs,
+        )
+        hidden_states = outputs.last_hidden_state
+        indices = (
+            slice(-logits_to_keep, None)
+            if isinstance(logits_to_keep, int)
+            else logits_to_keep
+        )
+        logits = self.lm_head(hidden_states[:, indices, :])
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                vocab_size=self.config.vocab_size,
+                **kwargs,
+            )
+        if not return_dict:
+            result = (logits, outputs.past_key_values)
+            return ((loss,) + result) if loss is not None else result
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = [
+    "Olmo3SiameseDepthForCausalLM",
+    "Olmo3SiameseDepthModel",
+    "Olmo3SiameseDepthPreTrainedModel",
+]

olmo3/1b/stage4/instruct/special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|pad|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

olmo3/1b/stage4/instruct/tokenizer_config.json

@@ -0,0 +1,191 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "100256": {
+      "content": "<|extra_id_0|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100257": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100258": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100259": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100260": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100261": {
+      "content": "|||PHONE_NUMBER|||",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100262": {
+      "content": "|||EMAIL_ADDRESS|||",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100263": {
+      "content": "|||IP_ADDRESS|||",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100264": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100265": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100266": {
+      "content": "<functions>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100267": {
+      "content": "</functions>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100268": {
+      "content": "<function_calls>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100269": {
+      "content": "</function_calls>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100270": {
+      "content": "<|extra_id_1|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100271": {
+      "content": "<|extra_id_2|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100272": {
+      "content": "<|extra_id_3|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100273": {
+      "content": "<|extra_id_4|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100274": {
+      "content": "<|extra_id_5|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100275": {
+      "content": "<|extra_id_6|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100276": {
+      "content": "<|endofprompt|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100277": {
+      "content": "<|pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "chat_template": "{%- set has_system = messages|selectattr('role', 'equalto', 'system')|list|length > 0 -%}{%- if not has_system -%}{{- '<|im_start|>system\nYou are a helpful function-calling AI assistant. ' -}}{%- if tools is none or (tools | length) == 0 -%}{{- 'You do not currently have access to any functions. <functions></functions><|im_end|>\n' -}}{%- else -%}{{- 'You are provided with function signatures within <functions></functions> XML tags. You may call one or more functions to assist with the user query. Output any function calls within <function_calls></function_calls> XML tags. Do not make assumptions about what values to plug into functions.' -}}{{- '<functions>' -}}{{- tools | tojson -}}{{- '</functions><|im_end|>\n' -}}{%- endif -%}{%- endif -%}{%- for message in messages -%}{%- if message['role'] == 'system' -%}{{- '<|im_start|>system\n' + message['content'] -}}{%- if tools is not none -%}{{- '<functions>' -}}{{- tools | tojson -}}{{- '</functions>' -}}{%- elif message.get('functions', none) is not none -%}{{- ' <functions>' + message['functions'] + '</functions>' -}}{%- endif -%}{{- '<|im_end|>\n' -}}{%- elif message['role'] == 'user' -%}{{- '<|im_start|>user\n' + message['content'] + '<|im_end|>\n' -}}{%- elif message['role'] == 'assistant' -%}{{- '<|im_start|>assistant\n' -}}{%- if message.get('content', none) is not none -%}{{- message['content'] -}}{%- endif -%}{%- if message.get('function_calls', none) is not none and (message['function_calls'] | length) > 0 -%}{{- '<function_calls>' + message['function_calls'] + '</function_calls>' -}}{%- elif message.get('tool_calls', none) is not none and (message['tool_calls'] | length) > 0 -%}{{- '<function_calls>' -}}{%- for tool_call in message['tool_calls'] %}{%- if tool_call is mapping and tool_call.get('function', none) is not none %}{%- set args = tool_call['function']['arguments'] -%}{%- set ns = namespace(arguments_list=[]) -%}{%- for key, value in args.items() -%}{%- set ns.arguments_list = ns.arguments_list + [key ~ '=' ~ (value | tojson)] -%}{%- endfor -%}{%- set arguments = ns.arguments_list | join(', ') -%}{{- tool_call['function']['name'] + '(' + arguments + ')' -}}{%- if not loop.last -%}{{ '\n' }}{%- endif -%}{% else %}{{- tool_call -}}{%- endif %}{%- endfor %}{{- '</function_calls>' -}}{%- endif -%}{%- if not loop.last -%}{{- '<|im_end|>' + '\n' -}}{%- else -%}{{- eos_token -}}{%- endif -%}{%- elif message['role'] == 'environment' -%}{{- '<|im_start|>environment\n' + message['content'] + '<|im_end|>\n' -}}{%- elif message['role'] == 'tool' -%}{{- '<|im_start|>environment\n' + message['content'] + '<|im_end|>\n' -}}{%- endif -%}{%- if loop.last and add_generation_prompt -%}{{- '<|im_start|>assistant\n' -}}{%- endif -%}{%- endfor -%}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "fix_mistral_regex": false,
+  "model_max_length": 65536,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}

olmo3/1b/stage4/think/PUBLIC_FILES.SHA256

@@ -0,0 +1,13 @@
+7ad9734423206b9dc3ddc76b320d5fe29f142245ef0df47981d69bdebe148556  README.md
+e2144e7990128799f6eaf6430ea0a472dbb17e076bdb08b327887e59f84657be  __init__.py
+a3d6a7ad3fde8e26fa9953d80ac30682f7d56ac1d03f3428910a4239a4e23bef  chat_template.jinja
+77c545f18298aa5960d58d0c83eb752b8559ebac016f5755a3aef50b226588c8  config.json
+d66ccce18efda42443a91536bf907e7bcb6e2d0acd0c8dc70d69eb15f277d634  configuration_olmo3_siamese_depth.py
+dbd4aa5dc6f7f1bdc0c181d9e5582f6571e1d2a22878dfcd98149324b9f5c4a2  generation_config.json
+b6fe424e334903f7fb84d3a106d9730455f4744b9fe3c21ee136d97a00e72502  merges.txt
+29c4567bcbcaf74acfca88df34c23fd23604f1ad66b933ff92bf72991c01ed58  model.safetensors
+12329a1f6cb1732bf100bda2ee2dd62761f6485aefb0a977879eb6d74f764aa6  modeling_olmo3_siamese_depth.py
+78afb564e81264029b25f9caf24bda2521d5bdaeff5cd3fdbc01d3da2e8ce2f2  special_tokens_map.json
+7738a25c46a6043a0e2e605138851baa450269ce9d800af6ca9bf6192e45510c  tokenizer.json
+5f959b9c0cff767a63f22a0569e0554cce28214fb49e41479ba326f44169c5ec  tokenizer_config.json
+d8eb44d0d51f5e39c9b14825ca8bd10a2500f99c438ca81d4462f0263acd3b53  vocab.json

olmo3/1b/stage4/think/README.md

@@ -0,0 +1,6 @@
+# Stage 4 Think SFT checkpoint
+
+This subfolder contains a complete Hugging Face checkpoint with a context
+length of 65536 tokens. Load it from the parent model repository using
+`trust_remote_code=True`, `subfolder="olmo3/1b/stage4/think"`, and
+`attn_implementation="sdpa"`.

olmo3/1b/stage4/think/__init__.py

@@ -0,0 +1,2 @@
+from .configuration_olmo3_siamese_depth import Olmo3SiameseDepthConfig
+from .modeling_olmo3_siamese_depth import Olmo3SiameseDepthForCausalLM, Olmo3SiameseDepthModel

olmo3/1b/stage4/think/chat_template.jinja

@@ -0,0 +1,16 @@
+{%- set has_system = messages|selectattr('role', 'equalto', 'system')|list|length > 0 -%}{%- if not has_system -%}{{- '<|im_start|>system
+You are a helpful function-calling AI assistant. ' -}}{%- if tools is none or (tools | length) == 0 -%}{{- 'You do not currently have access to any functions. <functions></functions><|im_end|>
+' -}}{%- else -%}{{- 'You are provided with function signatures within <functions></functions> XML tags. You may call one or more functions to assist with the user query. Output any function calls within <function_calls></function_calls> XML tags. Do not make assumptions about what values to plug into functions.' -}}{{- '<functions>' -}}{{- tools | tojson -}}{{- '</functions><|im_end|>
+' -}}{%- endif -%}{%- endif -%}{%- for message in messages -%}{%- if message['role'] == 'system' -%}{{- '<|im_start|>system
+' + message['content'] -}}{%- if tools is not none -%}{{- '<functions>' -}}{{- tools | tojson -}}{{- '</functions>' -}}{%- elif message.get('functions', none) is not none -%}{{- ' <functions>' + message['functions'] + '</functions>' -}}{%- endif -%}{{- '<|im_end|>
+' -}}{%- elif message['role'] == 'user' -%}{{- '<|im_start|>user
+' + message['content'] + '<|im_end|>
+' -}}{%- elif message['role'] == 'assistant' -%}{{- '<|im_start|>assistant
+' -}}{%- if message.get('content', none) is not none -%}{{- message['content'] -}}{%- endif -%}{%- if message.get('function_calls', none) is not none and (message['function_calls'] | length) > 0 -%}{{- '<function_calls>' + message['function_calls'] + '</function_calls>' -}}{%- elif message.get('tool_calls', none) is not none and (message['tool_calls'] | length) > 0 -%}{{- '<function_calls>' -}}{%- for tool_call in message['tool_calls'] %}{%- if tool_call is mapping and tool_call.get('function', none) is not none %}{%- set args = tool_call['function']['arguments'] -%}{%- set ns = namespace(arguments_list=[]) -%}{%- for key, value in args.items() -%}{%- set ns.arguments_list = ns.arguments_list + [key ~ '=' ~ (value | tojson)] -%}{%- endfor -%}{%- set arguments = ns.arguments_list | join(', ') -%}{{- tool_call['function']['name'] + '(' + arguments + ')' -}}{%- if not loop.last -%}{{ '
+' }}{%- endif -%}{% else %}{{- tool_call -}}{%- endif %}{%- endfor %}{{- '</function_calls>' -}}{%- endif -%}{%- if not loop.last -%}{{- '<|im_end|>' + '
+' -}}{%- else -%}{{- eos_token -}}{%- endif -%}{%- elif message['role'] == 'environment' -%}{{- '<|im_start|>environment
+' + message['content'] + '<|im_end|>
+' -}}{%- elif message['role'] == 'tool' -%}{{- '<|im_start|>environment
+' + message['content'] + '<|im_end|>
+' -}}{%- endif -%}{%- if loop.last and add_generation_prompt -%}{{- '<|im_start|>assistant
+' -}}{%- endif -%}{%- endfor -%}

olmo3/1b/stage4/think/config.json

@@ -0,0 +1,76 @@
+{
+  "_name_or_path": "",
+  "architectures": [
+    "Olmo3SiameseDepthForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_olmo3_siamese_depth.Olmo3SiameseDepthConfig",
+    "AutoModel": "modeling_olmo3_siamese_depth.Olmo3SiameseDepthModel",
+    "AutoModelForCausalLM": "modeling_olmo3_siamese_depth.Olmo3SiameseDepthForCausalLM"
+  },
+  "bos_token_id": null,
+  "depth_attention_recent_window": 0,
+  "depth_attention_stride": 8,
+  "eos_token_id": 100257,
+  "hidden_act": "silu",
+  "hidden_dropout": 0.0,
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 8192,
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 65536,
+  "max_sequence_length": 65536,
+  "mlp_bias": false,
+  "model_type": "olmo3_siamese_depth",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 16,
+  "num_key_value_heads": 16,
+  "olmo3_weight_decay": true,
+  "pad_token_id": 100277,
+  "padded_vocab_size": 100352,
+  "qk_norm": true,
+  "qk_norm_mode": "full_projection",
+  "rms_norm_eps": 1e-06,
+  "rope_full_precision": true,
+  "rope_scaling": {
+    "beta_fast": 32.0,
+    "beta_slow": 1.0,
+    "factor": 8.0,
+    "original_max_position_embeddings": 8192,
+    "rope_type": "yarn"
+  },
+  "rope_theta": 500000.0,
+  "scale_embeds": false,
+  "scale_output_layer_init": false,
+  "siamese_norm_variant": "hybrid_pre",
+  "sliding_window": 4096,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.57.6",
+  "true_vocab_size": 100278,
+  "truncated_normal_factor": 3.0,
+  "use_cache": true,
+  "use_depth_attention": true,
+  "use_siamese_norm": true,
+  "vocab_size": 100278,
+  "vocab_z_loss_coeff": 1e-05
+}

olmo3/1b/stage4/think/configuration_olmo3_siamese_depth.py

@@ -0,0 +1,78 @@
+"""Hugging Face configuration for the OLMo 3 Siamese-Norm/Depth-Attention model."""
+
+from __future__ import annotations
+
+from transformers.models.olmo3.configuration_olmo3 import Olmo3Config
+
+
+class Olmo3SiameseDepthConfig(Olmo3Config):
+    """OLMo 3 plus the checkpoint-compatible Siamese/Depth extensions."""
+
+    model_type = "olmo3_siamese_depth"
+
+    def __init__(
+        self,
+        *args,
+        vocab_size: int | None = None,
+        true_vocab_size: int = 100278,
+        padded_vocab_size: int = 100352,
+        qk_norm_mode: str = "full_projection",
+        use_siamese_norm: bool = True,
+        siamese_norm_variant: str = "hybrid_pre",
+        use_depth_attention: bool = True,
+        depth_attention_stride: int = 8,
+        depth_attention_recent_window: int = 0,
+        rope_full_precision: bool = True,
+        **kwargs,
+    ):
+        # ``PretrainedConfig.to_diff_dict()`` constructs a no-argument instance
+        # of this class.  Make that default instance internally consistent
+        # instead of inheriting OLMo3's unrelated 50304-token default.
+        if vocab_size is None:
+            vocab_size = true_vocab_size
+        super().__init__(*args, vocab_size=vocab_size, **kwargs)
+        self.true_vocab_size = int(true_vocab_size)
+        self.padded_vocab_size = int(padded_vocab_size)
+        self.qk_norm = True
+        self.qk_norm_mode = qk_norm_mode
+        self.use_siamese_norm = bool(use_siamese_norm)
+        self.siamese_norm_variant = siamese_norm_variant
+        self.use_depth_attention = bool(use_depth_attention)
+        self.depth_attention_stride = int(depth_attention_stride)
+        self.depth_attention_recent_window = int(depth_attention_recent_window)
+        self.rope_full_precision = bool(rope_full_precision)
+        self._validate_siamese_depth()
+
+    def _validate_siamese_depth(self) -> None:
+        if self.vocab_size != self.true_vocab_size:
+            raise ValueError(
+                "HF vocab_size must equal true_vocab_size after padded-row removal; "
+                f"got {self.vocab_size} and {self.true_vocab_size}."
+            )
+        if self.padded_vocab_size < self.true_vocab_size:
+            raise ValueError("padded_vocab_size cannot be smaller than true_vocab_size.")
+        if self.qk_norm_mode != "full_projection":
+            raise ValueError("qk_norm_mode must be 'full_projection'.")
+        if not self.use_siamese_norm or self.siamese_norm_variant != "hybrid_pre":
+            raise ValueError("This remote model requires Hybrid-Pre Siamese Norm.")
+        if not self.use_depth_attention:
+            raise ValueError("This remote model requires Depth Attention.")
+        if self.depth_attention_stride < 1:
+            raise ValueError("depth_attention_stride must be positive.")
+        if self.depth_attention_recent_window < 0:
+            raise ValueError("depth_attention_recent_window must be non-negative.")
+        if not self.rope_full_precision:
+            raise ValueError("OLMo 3 requires FP32 Q/K RoPE.")
+        if self.hidden_size % self.num_attention_heads:
+            raise ValueError("hidden_size must be divisible by num_attention_heads.")
+        if self.num_attention_heads % self.num_key_value_heads:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads.")
+        expected_layer_types = [
+            "sliding_attention" if (index + 1) % 4 else "full_attention"
+            for index in range(self.num_hidden_layers)
+        ]
+        if list(self.layer_types) != expected_layer_types:
+            raise ValueError("OLMo 3 requires the repeating SWA,SWA,SWA,Full pattern.")
+
+
+__all__ = ["Olmo3SiameseDepthConfig"]

olmo3/1b/stage4/think/generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": null,
+  "eos_token_id": 100257,
+  "pad_token_id": 100277,
+  "transformers_version": "4.57.6"
+}

olmo3/1b/stage4/think/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:29c4567bcbcaf74acfca88df34c23fd23604f1ad66b933ff92bf72991c01ed58
+size 2969458432

olmo3/1b/stage4/think/modeling_olmo3_siamese_depth.py

@@ -0,0 +1,548 @@
+"""Transformers remote code for OLMo 3 + Siamese Norm + Depth Attention.
+
+The Full/SWA masks, Full-only YaRN rotary embeddings, and bounded sliding KV
+cache are delegated to the official Transformers OLMo 3 implementation.  This
+file adds only the two checkpoint-persistent mechanisms used during training:
+Hybrid-Pre Siamese Norm and recursive same-position Depth Attention.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import torch
+import torch.nn as nn
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+from transformers.models.olmo3.modeling_olmo3 import (
+    Olmo3MLP,
+    Olmo3PreTrainedModel,
+    Olmo3RMSNorm,
+    Olmo3RotaryEmbedding,
+    apply_rotary_pos_emb,
+    eager_attention_forward,
+)
+
+from .configuration_olmo3_siamese_depth import Olmo3SiameseDepthConfig
+
+
+@dataclass
+class _DepthSource:
+    layer_index: int
+    key: torch.Tensor
+    value: torch.Tensor
+
+
+class _DepthContext:
+    def __init__(self, config: Olmo3SiameseDepthConfig):
+        self.num_layers = int(config.num_hidden_layers)
+        self.stride = int(config.depth_attention_stride)
+        self.recent_window = int(config.depth_attention_recent_window)
+        self.records: dict[int, _DepthSource] = {}
+
+    def sources_for(self, layer_index: int) -> tuple[_DepthSource, ...]:
+        recent_start = max(0, layer_index - self.recent_window)
+        return tuple(
+            self.records[index]
+            for index in sorted(self.records)
+            if index < layer_index
+            and (index % self.stride == 0 or index >= recent_start)
+        )
+
+    def record(self, layer_index: int, key: torch.Tensor, value: torch.Tensor) -> None:
+        if layer_index in self.records:
+            raise RuntimeError(f"Depth Attention layer {layer_index} was recorded twice.")
+        self.records[layer_index] = _DepthSource(layer_index, key, value)
+        next_recent_start = layer_index + 1 - self.recent_window
+        stale = [
+            index
+            for index in self.records
+            if index % self.stride != 0 and index < next_recent_start
+        ]
+        for index in stale:
+            del self.records[index]
+
+
+def _depth_attention_mix(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    sources: tuple[_DepthSource, ...],
+) -> torch.Tensor:
+    """Mix current V with recursively mixed V from selected earlier layers.
+
+    Tensors use Transformers' ``[batch, heads, sequence, head_dim]`` layout.
+    """
+
+    if not sources:
+        return value
+    query_heads = query.shape[1]
+    kv_heads = key.shape[1]
+    if query_heads % kv_heads:
+        raise ValueError("Depth Attention requires Q heads divisible by KV heads.")
+    group_size = query_heads // kv_heads
+    grouped_query = (
+        query
+        if group_size == 1
+        else query.reshape(
+            query.shape[0],
+            kv_heads,
+            group_size,
+            query.shape[2],
+            query.shape[3],
+        ).mean(dim=2)
+    )
+    scale = 1.0 / math.sqrt(query.shape[-1])
+    scores = [
+        (grouped_query * source.key).sum(dim=-1, keepdim=True) * scale
+        for source in sources
+    ]
+    scores.append((grouped_query * key).sum(dim=-1, keepdim=True) * scale)
+    weights = torch.softmax(torch.cat(scores, dim=-1), dim=-1, dtype=torch.float32)
+    weights = weights.to(dtype=value.dtype)
+    mixed = weights[..., -1:].mul(value)
+    for source_index, source in enumerate(sources):
+        mixed = mixed + weights[..., source_index : source_index + 1] * source.value
+    return mixed
+
+
+class Olmo3SiameseDepthAttention(nn.Module):
+    def __init__(self, config: Olmo3SiameseDepthConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = int(layer_idx)
+        self.head_dim = getattr(
+            config, "head_dim", config.hidden_size // config.num_attention_heads
+        )
+        self.num_key_value_groups = (
+            config.num_attention_heads // config.num_key_value_heads
+        )
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = float(config.attention_dropout)
+        self.is_causal = True
+        self.q_proj = nn.Linear(
+            config.hidden_size,
+            config.num_attention_heads * self.head_dim,
+            bias=False,
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=False,
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=False,
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim,
+            config.hidden_size,
+            bias=False,
+        )
+        # These norms cover the complete Q and K projections, not one head.
+        self.q_norm = Olmo3RMSNorm(
+            config.num_attention_heads * self.head_dim, config.rms_norm_eps
+        )
+        self.k_norm = Olmo3RMSNorm(
+            config.num_key_value_heads * self.head_dim, config.rms_norm_eps
+        )
+        self.attention_type = config.layer_types[self.layer_idx]
+        self.sliding_window = (
+            int(config.sliding_window)
+            if self.attention_type == "sliding_attention"
+            else None
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        depth_context: _DepthContext,
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        output_attentions: bool = False,
+        **kwargs,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        batch_size, sequence_length, _ = hidden_states.shape
+        query = self.q_norm(self.q_proj(hidden_states)).view(
+            batch_size, sequence_length, self.config.num_attention_heads, self.head_dim
+        ).transpose(1, 2)
+        key = self.k_norm(self.k_proj(hidden_states)).view(
+            batch_size, sequence_length, self.config.num_key_value_heads, self.head_dim
+        ).transpose(1, 2)
+        value = self.v_proj(hidden_states).view(
+            batch_size, sequence_length, self.config.num_key_value_heads, self.head_dim
+        ).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query, key = apply_rotary_pos_emb(query.float(), key.float(), cos.float(), sin.float())
+        query = query.to(dtype=hidden_states.dtype)
+        key = key.to(dtype=hidden_states.dtype)
+
+        mixed_value = _depth_attention_mix(
+            query,
+            key,
+            value,
+            depth_context.sources_for(self.layer_idx),
+        )
+        depth_context.record(self.layer_idx, key, mixed_value)
+
+        attention_key = key
+        attention_value = mixed_value
+        if past_key_values is not None:
+            cache_kwargs = {
+                "sin": sin,
+                "cos": cos,
+                "cache_position": cache_position,
+            }
+            attention_key, attention_value = past_key_values.update(
+                key, mixed_value, self.layer_idx, cache_kwargs
+            )
+
+        attention_interface = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[
+                self.config._attn_implementation
+            ]
+        attention_output, attention_weights = attention_interface(
+            self,
+            query,
+            attention_key,
+            attention_value,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            **kwargs,
+        )
+        attention_output = attention_output.reshape(
+            batch_size, sequence_length, -1
+        ).contiguous()
+        attention_output = self.o_proj(attention_output)
+        return attention_output, attention_weights if output_attentions else None
+
+
+class Olmo3SiameseDepthDecoderLayer(nn.Module):
+    def __init__(self, config: Olmo3SiameseDepthConfig, layer_idx: int):
+        super().__init__()
+        self.layer_idx = int(layer_idx)
+        self.self_attn = Olmo3SiameseDepthAttention(config, layer_idx)
+        self.mlp = Olmo3MLP(config)
+        self.pre_mlp_layernorm = Olmo3RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.siamese_attn_pre_norm = Olmo3RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.siamese_mlp_pre_norm = Olmo3RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.siamese_mlp_input_norm = Olmo3RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.siamese_hybrid_attn_scale = nn.Parameter(
+            torch.ones(config.hidden_size)
+        )
+        self.post_residual_scale = 1.0 / math.sqrt(2.0 * (self.layer_idx + 1))
+
+    def forward(
+        self,
+        post_stream: torch.Tensor,
+        pre_stream: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        depth_context: _DepthContext,
+        past_key_values: Optional[Cache],
+        cache_position: Optional[torch.LongTensor],
+        output_attentions: bool,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        pre_normalized = self.siamese_attn_pre_norm(pre_stream)
+        attention_input = torch.addcmul(
+            pre_normalized,
+            post_stream,
+            self.siamese_hybrid_attn_scale.to(dtype=post_stream.dtype),
+        )
+        attention_output, attention_weights = self.self_attn(
+            hidden_states=attention_input,
+            position_embeddings=position_embeddings,
+            attention_mask=attention_mask,
+            depth_context=depth_context,
+            past_key_values=past_key_values,
+            cache_position=cache_position,
+            output_attentions=output_attentions,
+            **kwargs,
+        )
+        post_stream = torch.add(
+            post_stream, attention_output, alpha=self.post_residual_scale
+        )
+        pre_stream = pre_stream + attention_output
+
+        # The native Hybrid-Pre block advances its post stream to the
+        # pre-MLP-normalized value before both the MLP input construction and
+        # the MLP residual update.  This is intentionally not a conventional
+        # pre-norm residual that adds the branch back to the unnormalized
+        # stream.
+        post_stream = self.pre_mlp_layernorm(post_stream)
+        pre_normalized = self.siamese_mlp_pre_norm(pre_stream)
+        mlp_input = self.siamese_mlp_input_norm(
+            post_stream + pre_normalized
+        )
+        mlp_output = self.mlp(mlp_input)
+        post_stream = torch.add(
+            post_stream, mlp_output, alpha=self.post_residual_scale
+        )
+        pre_stream = pre_stream + mlp_output
+        return post_stream, pre_stream, attention_weights
+
+
+class Olmo3SiameseDepthPreTrainedModel(Olmo3PreTrainedModel):
+    config_class = Olmo3SiameseDepthConfig
+    base_model_prefix = "model"
+    _no_split_modules = ["Olmo3SiameseDepthDecoderLayer"]
+    # Whole-layer replay is invalid for the mutable cross-layer Depth context.
+    # The native trainer supports only its dedicated selective-MLP checkpoint
+    # path; do not advertise Transformers' generic layer checkpointing.
+    supports_gradient_checkpointing = False
+
+
+class Olmo3SiameseDepthModel(Olmo3SiameseDepthPreTrainedModel):
+    def __init__(self, config: Olmo3SiameseDepthConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [
+                Olmo3SiameseDepthDecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.siamese_post_final_layernorm = Olmo3RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.siamese_pre_final_layernorm = Olmo3RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.norm = Olmo3RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.rotary_embs = nn.ModuleDict(
+            {
+                "sliding_attention": Olmo3RotaryEmbedding(
+                    config=config, rope_type="default"
+                ),
+                "full_attention": Olmo3RotaryEmbedding(config=config),
+            }
+        )
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[tuple, BaseModelOutputWithPast]:
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Specify exactly one of input_ids or inputs_embeds.")
+        output_attentions = (
+            self.config.output_attentions
+            if output_attentions is None
+            else output_attentions
+        )
+        output_hidden_states = (
+            self.config.output_hidden_states
+            if output_hidden_states is None
+            else output_hidden_states
+        )
+        return_dict = self.config.use_return_dict if return_dict is None else return_dict
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+        if cache_position is None:
+            past_seen_tokens = (
+                past_key_values.get_seq_length()
+                if past_key_values is not None
+                else 0
+            )
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        if not isinstance(attention_mask, dict):
+            mask_kwargs = {
+                "config": self.config,
+                "input_embeds": inputs_embeds,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+                "position_ids": position_ids,
+            }
+            causal_masks = {
+                "full_attention": create_causal_mask(**mask_kwargs),
+                "sliding_attention": create_sliding_window_causal_mask(
+                    **mask_kwargs
+                ),
+            }
+        else:
+            causal_masks = attention_mask
+
+        post_stream = inputs_embeds
+        # The optimized training path intentionally aliases both initial streams.
+        pre_stream = inputs_embeds
+        depth_context = _DepthContext(self.config)
+        position_embeddings = {
+            kind: rotary(post_stream, position_ids)
+            for kind, rotary in self.rotary_embs.items()
+        }
+        hidden_history = () if output_hidden_states else None
+        attention_history = () if output_attentions else None
+        for layer in self.layers:
+            if output_hidden_states:
+                hidden_history += (post_stream,)
+            attention_type = layer.self_attn.attention_type
+            post_stream, pre_stream, attention_weights = layer(
+                post_stream=post_stream,
+                pre_stream=pre_stream,
+                attention_mask=causal_masks[attention_type],
+                position_embeddings=position_embeddings[attention_type],
+                depth_context=depth_context,
+                past_key_values=past_key_values,
+                cache_position=cache_position,
+                output_attentions=output_attentions,
+                **kwargs,
+            )
+            if output_attentions:
+                attention_history += (attention_weights,)
+
+        hidden_states = (
+            self.siamese_post_final_layernorm(post_stream)
+            + self.siamese_pre_final_layernorm(pre_stream)
+        )
+        hidden_states = self.norm(hidden_states)
+        if output_hidden_states:
+            hidden_history += (hidden_states,)
+        if not return_dict:
+            values = (hidden_states, past_key_values, hidden_history, attention_history)
+            return tuple(value for value in values if value is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+            hidden_states=hidden_history,
+            attentions=attention_history,
+        )
+
+
+class Olmo3SiameseDepthForCausalLM(
+    Olmo3SiameseDepthPreTrainedModel, GenerationMixin
+):
+    _tied_weights_keys = []
+
+    def __init__(self, config: Olmo3SiameseDepthConfig):
+        super().__init__(config)
+        self.model = Olmo3SiameseDepthModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[tuple, CausalLMOutputWithPast]:
+        return_dict = self.config.use_return_dict if return_dict is None else return_dict
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            **kwargs,
+        )
+        hidden_states = outputs.last_hidden_state
+        indices = (
+            slice(-logits_to_keep, None)
+            if isinstance(logits_to_keep, int)
+            else logits_to_keep
+        )
+        logits = self.lm_head(hidden_states[:, indices, :])
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                vocab_size=self.config.vocab_size,
+                **kwargs,
+            )
+        if not return_dict:
+            result = (logits, outputs.past_key_values)
+            return ((loss,) + result) if loss is not None else result
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+__all__ = [
+    "Olmo3SiameseDepthForCausalLM",
+    "Olmo3SiameseDepthModel",
+    "Olmo3SiameseDepthPreTrainedModel",
+]

olmo3/1b/stage4/think/special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|pad|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}