[ARCH]Add SiameseNorm and DepthAttention architecture

README.md

@@ -1,80 +1,31 @@
-## Notes
+# SiameseNorm and DepthAttention for OLMo 3
 
-💡 To make your Proposal easier to validate and accept, provide implementation code that is **reproducible**, **runnable**, and **easy to use**, together with **clear and complete documentation**.
+This branch implements SiameseNorm and DepthAttention on top of the OLMo 3 architecture.
 
-💡 The model architecture code must be converted to a Hugging Face Transformers-compatible format and placed in the `archs/` directory. **Only changes within `archs/` will be merged**.
+## Repository Layout
 
-💡 **All reproduction code must be placed in the `reproduce/` directory**, which you may also use as your working directory during development.
+| Path | Description |
+|---|---|
+| [`archs/SiameseNorm-DepthAttention`](archs/SiameseNorm-DepthAttention) | Hugging Face Transformers-compatible implementation for inference and generation |
+| [`reproduce/Megatron-LM`](reproduce/Megatron-LM) | Ascend/MindSpeed training, checkpoint conversion, inference, and evaluation pipeline |
 
-# Architecture: `<Architecture Name>`
+The reproduction pipeline covers:
 
----
+- Stage 1: 8K pretraining
+- Stage 2: 8K mid-training
+- Stage 3: 65K long-context training
+- Stage 4: Think SFT and Instruct SFT
+- Checkpoint save, resume, transition, and Hugging Face conversion
+- Native inference, PPL validation, and OLMES evaluation
 
-## 1. Basic Information
+Detailed setup and execution commands are available in the
+[reproduction runbook](reproduce/Megatron-LM/docs/RUNBOOK.md).
 
-| Item                 |                          Details                          |
-| -------------------- | :-------------------------------------------------------: |
-| Architecture Name    |                   `<Architecture Name>`                   |
-| Parent ARCH-PROP ID  | [Issue \#N](https://github.com/InternLM/archspace/issues) |
-| Current ARCH-PROP ID | [Issue \#N](https://github.com/InternLM/archspace/issues) |
+## Released Checkpoints and Logs
 
-## 2. Reproducing the Experiments
+- [Hugging Face checkpoints](https://huggingface.co/ArchSpace-Collection/SiameseNorm-DepthAttention)
+- [Weights & Biases report](https://wandb.ai/archspace/SiameseNormDepthAttention/reports/Siamese-Norm-and-Depth-Attention-in-OLMo-3-1B--VmlldzoxNzYwMzAwNw)
 
-### 2.1 Environment Setup
-
-> Specify the required hardware and software environment, and provide complete installation instructions. Pin key dependency versions to ensure the environment can be reproduced reliably.
-
-### 2.2 Data Preparation
-
-> Architecture experiments should generally use the same data as the baseline. If the data or data-processing pipeline differs from the baseline, describe the data source and the complete preparation process here.
-
-### 2.3 Training Pipeline
-
-> Provide all training scripts, configuration files, and commands required to reproduce the training process. The commands should run without requiring modifications to the source code.
-
-### 2.4 Evaluation Pipeline
-
-> Provide all evaluation scripts, configuration files, and commands required to reproduce the reported results. Clearly specify the evaluation metrics and expected outputs.
-
-### 2.5 Model Weights and Experiment Logs
-
-> 1. Use [Weights & Biases](https://wandb.ai/site/) to record training logs.
-> 2. After completing the validation experiments, convert the model architecture code to a [Hugging Face Transformers-compatible format](https://huggingface.co/docs/transformers/v5.14.0/en/main_classes/model).## Notes
->    💡 To make your Proposal easier to validate and accept, provide implementation code that is **reproducible**, **runnable**, and **easy to use**, together with **clear and complete documentation**.
->    💡 The model architecture code must be converted to a Hugging Face Transformers-compatible format and placed in the `archs/` directory. **Only changes within `archs/` will be merged**.
->    💡 **All reproduction code must be placed in the `reproduce/` directory**, which you may also use as your working directory during development.
-
-# Architecture: `<Architecture Name>`
-
----
-
-## 1. Basic Information
-
-| Item                 |                          Details                          |
-| -------------------- | :-------------------------------------------------------: |
-| Architecture Name    |                   `<Architecture Name>`                   |
-| Parent ARCH-PROP ID  | [Issue \#N](https://github.com/InternLM/archspace/issues) |
-| Current ARCH-PROP ID | [Issue \#N](https://github.com/InternLM/archspace/issues) |
-
-## 2. Reproducing the Experiments
-
-### 2.1 Environment Setup
-
-> Specify the required hardware and software environment, and provide complete installation instructions. Pin key dependency versions to ensure the environment can be reproduced reliably.
-
-### 2.2 Data Preparation
-
-> Architecture experiments should generally use the same data as the baseline. If the data or data-processing pipeline differs from the baseline, describe the data source and the complete preparation process here.
-
-### 2.3 Training Pipeline
-
-> Provide all training scripts, configuration files, and commands required to reproduce the training process. The commands should run without requiring modifications to the source code.
-
-### 2.4 Evaluation Pipeline
-
-> Provide all evaluation scripts, configuration files, and commands required to reproduce the reported results. Clearly specify the evaluation metrics and expected outputs.
-
-### 2.5 Model Weights and Experiment Logs
-
-> 1. Use [Weights & Biases](https://wandb.ai/site/) to record training logs.
-> 2. After completing the validation experiments, convert the model architecture code to a [Hugging Face Transformers-compatible format](https://huggingface.co/docs/transformers/v5.14.0/en/main_classes/model).
+The released checkpoints currently cover the complete four-stage OLMo 3 1B
+SiameseNorm–DepthAttention pipeline. Matched baseline training and additional
+model-scale validation are ongoing.

archs/SiameseNorm-DepthAttention/README.md

@@ -0,0 +1,85 @@
+---
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+  - olmo3
+  - custom-code
+  - siamese-norm
+  - depth-attention
+  - sliding-window-attention
+---
+
+# OLMo 3 with SiameseNorm and DepthAttention
+
+This repository contains the 1B checkpoints from the four-stage OLMo 3
+training pipeline with SiameseNorm and DepthAttention.
+
+## Checkpoints
+
+| Checkpoint | Hub subfolder | Context length |
+|---|---|---:|
+| Stage 1 pretraining | `olmo3/1b/stage1` | 8,192 |
+| Stage 2 mid-training | `olmo3/1b/stage2` | 8,192 |
+| Stage 3 long-context training | `olmo3/1b/stage3` | 65,536 |
+| Stage 4 Think SFT | `olmo3/1b/stage4/think` | 65,536 |
+| Stage 4 Instruct SFT | `olmo3/1b/stage4/instruct` | 65,536 |
+
+Stage 3 and Stage 4 apply YaRN only to Full-attention layers. Sliding-window
+attention layers retain the original RoPE and a 4,096-token window.
+
+## Loading
+
+Select one checkpoint through `subfolder`. SDPA is the recommended and
+release-validated BF16 inference backend:
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+repo_id = "ArchSpace-Collection/SiameseNorm-DepthAttention"
+subfolder = "olmo3/1b/stage4/instruct"
+
+tokenizer = AutoTokenizer.from_pretrained(
+    repo_id,
+    subfolder=subfolder,
+    trust_remote_code=True,
+    fix_mistral_regex=False,
+)
+model = AutoModelForCausalLM.from_pretrained(
+    repo_id,
+    subfolder=subfolder,
+    trust_remote_code=True,
+    dtype=torch.bfloat16,
+    attn_implementation="sdpa",
+)
+```
+
+`fix_mistral_regex=False` is intentional and preserves the tokenizer behavior
+used during training.
+
+The eager backend can also load these checkpoints. For BF16 generation, SDPA
+is recommended because eager cache partitioning can introduce small rounding
+differences when the leading logits are nearly tied; in that narrow case,
+greedy generation can select a different token. This is a numerical
+backend/cache-partition effect, not a checkpoint conversion or weight-integrity
+problem.
+
+The repository root contains the shared custom modeling code required by
+Transformers remote-code loading. Each checkpoint subfolder also contains a
+self-contained copy of its configuration, tokenizer, modeling code, and
+weights.
+
+## Architecture
+
+- 16 transformer layers
+- hidden size 2,048
+- intermediate size 8,192
+- 16 query heads and 16 key/value heads
+- 128-dimensional attention heads
+- 3:1 sliding-window/full-attention pattern
+- 4,096-token sliding window
+- reordered RMSNorm, SiameseNorm, and DepthAttention
+
+The Hugging Face implementation is intended for inference and generation.
+Exact continuation of the native distributed training objective should use
+the accompanying MindSpeed/Megatron training pipeline.

archs/SiameseNorm-DepthAttention/__init__.py

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

archs/SiameseNorm-DepthAttention/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"]