diff options
Diffstat (limited to 'CLAUDE.md')
| -rw-r--r-- | CLAUDE.md | 1381 |
1 files changed, 1381 insertions, 0 deletions
diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..1d83dec --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,1381 @@ +# CLAUDE.md — DAGFormer Project Specification + +**Read this file in full before writing any code.** This document is the +single source of truth for all design decisions. If something contradicts +the README or any other file, this document wins. + +--- + +## 1. What Is This Project? + +**DAGFormer** trains a small neural network (the "structure predictor") to +predict, for each token, the optimal wiring diagram (a DAG) of a frozen +1B-parameter language model (OLMo2-1B). The predicted DAG controls which +attention heads talk to which other heads. The entire system is trained +end-to-end with language modeling loss — no labeled topology data needed. + +### Why? + +Standard transformers have a fixed sequential computation graph: layer 0 → +layer 1 → ... → layer 15. Every input sees the same wiring. We showed via +expensive oracle search that **context-dependent topologies** can reduce +next-token prediction loss (NLL) from 2.58 to 0.12 (median across 50 +evaluation windows, 100% of windows improved). The oracle search is too +slow to use at scale (500 gradient steps per window), so we need a learned +predictor that produces the topology in a single forward pass. + +### What is NOT this project? + +- This is NOT the oracle search codebase (that exists separately) +- This is NOT a Mixture-of-Experts project (despite the repo history) +- This does NOT modify the OLMo2-1B weights in Phase 1 +- This does NOT implement Phase 2 (joint training) yet — only the + infrastructure to support it later + +--- + +## 2. Architecture — Exact Specification + +### 2.1 The Computation Graph of OLMo2-1B + +OLMo2-1B (HuggingFace ID: `allenai/OLMo-2-0425-1B`) has: + +- **16 transformer layers**, each with **16 attention heads** +- This gives **256 "nodes"** total: node `(l, h)` = layer `l`, head `h` +- We flatten to a single index: `node_id = l * 16 + h` (0-indexed) + +**Standard forward pass** — each layer does: +```python +# Input: residual (shared across all heads in this layer) +normed = RMSNorm(residual) +attn_out = self_attn(normed) # all 16 heads compute in parallel, + # outputs concatenated and projected by o_proj +residual = residual + attn_out # attention residual connection +normed2 = RMSNorm(residual) +mlp_out = MLP(normed2) +residual = residual + mlp_out # MLP residual connection +``` + +The **residual stream** at the start of layer `l` is therefore: +``` +residual_l = embedding + Σ_{l'<l} (attn_output[l'] + mlp_output[l']) + = embedding + Σ_{l'<l} (Σ_h head_output[l',h] + mlp_output[l']) +``` +where `head_output[l',h]` is head h's individual contribution to the +attention output (its slice of `o_proj`, see §2.2). ALL heads in layer l +see the SAME `residual_l` as input — there is no per-head differentiation +in standard transformers. + +### 2.2 The Adjacency Matrix A + +We introduce a **256×256 adjacency matrix A** that controls information +routing between attention heads across layers. + +``` +A[i][j] ∈ [0, 1] where i = source node, j = target node + i = l_i * 16 + h_i, j = l_j * 16 + h_j +``` + +#### Mask: Block-Upper-Triangular (NOT element-upper-triangular) + +**CRITICAL**: The mask is based on LAYER indices, not node indices. + +```python +# CORRECT: block-upper-triangular based on layer +mask[i, j] = 1 if layer(j) > layer(i) # i.e. j//16 > i//16 +mask[i, j] = 0 if layer(j) <= layer(i) # same layer or backward + +# WRONG: do NOT use torch.triu() — it would allow same-layer connections +# e.g. triu would set mask[0,15]=1, but both are in layer 0 +``` + +Heads in the same layer execute in parallel and cannot see each other's +outputs. Only cross-layer forward connections are meaningful. + +#### Connection Count + +| Type | Definition | Count | Role | +|------|-----------|-------|------| +| **Adjacent-layer** | `layer(j) = layer(i) + 1`, all head pairs | 15 × 16 × 16 = **3,840** | These exist in standard transformer (via shared residual). When gated to 1, behavior matches baseline. | +| **Skip** | `layer(j) > layer(i) + 1`, all head pairs | 105 × 16 × 16 = **26,880** | These do NOT exist in standard transformer. They are additional direct routes that bypass intermediate layers. | +| **Total** | All entries where `layer(j) > layer(i)` | **30,720** | | + +For logging and analysis, label connections as "adjacent" or "skip", but +the forward pass treats all 30,720 entries identically. + +> **Note on "31K" count**: The oracle search reported "256 sequential + +> 30,720 hyperconnection ≈ 31K" using a different parameterization +> (separate per-head activity gates + routing gates). In our unified +> 256×256 matrix, there are exactly 30,720 free entries. Both represent +> the same underlying structure. + +#### What "head output" means (resolves shape ambiguity) + +In HuggingFace OLMo2, the `o_proj` layer concatenates all 16 heads and +projects back to model_dim: + +```python +# Inside self_attn: +# Each head computes: attn_weights @ V → [batch, seq, head_dim] (head_dim = 128) +# All heads concatenated: [batch, seq, 16 * 128] = [batch, seq, 2048] +# Then: o_proj([batch, seq, 2048]) → [batch, seq, 2048] +``` + +The `o_proj` weight matrix `W_o ∈ R^{2048 × 2048}` can be viewed as 16 +column blocks: `W_o = [W_o^0 | W_o^1 | ... | W_o^15]`, each `W_o^h ∈ +R^{2048 × 128}`. The full attention output is: + +``` +attn_output = Σ_h W_o^h @ head_value_h = Σ_h head_output[l, h] +``` + +So **`head_output[l, h]` is `W_o^h @ head_value_h`, shape `[batch, seq, +model_dim]` (= 2048)**. This is each head's contribution to the attention +output in MODEL DIM space. Heads can be summed directly because they're +already in the same 2048-dim space. + +#### Per-Head Input Assembly (the core modification) + +In DAGFormer, each head j = (l_j, h_j) has its **own input**, assembled +from three sources: + +```python +input_j = embedding # (1) always present + + Σ_{l' < l_j} mlp_output[l'] # (2) all prior MLP outputs, NOT gated + + Σ_{i: layer(i) < l_j} A[i,j] * head_output[i] # (3) gated head outputs +``` + +**Source (1) — Token embedding**: The output of the token embedding layer +(before any transformer layer). Always included for every head. This is +NOT part of A. Think of it as the "seed" of the computation. + +**Source (2) — MLP outputs**: Each layer's MLP computes on a shared +aggregated state (see below) and its output is added to ALL downstream +head inputs equally. MLP outputs are **never gated by A**. This keeps +the MLP's contribution identical to the standard forward pass. + +**Source (3) — Gated head outputs**: The only part controlled by A. +Each entry A[i,j] scales how much of head i's output reaches head j's +input. Different heads h within the same layer l_j can receive different +weighted combinations of prior head outputs — this is what makes the +256×256 per-head routing meaningful. + +#### Baseline Reproduction Proof (resolves the "double-counting" concern) + +When A[i,j] = 1 for ALL 30,720 valid entries: + +``` +input_j (where j is in layer l) + = embedding + Σ_{l'<l} mlp_output[l'] + Σ_{i: layer(i)<l} 1.0 * head_output[i] + = embedding + Σ_{l'<l} mlp_output[l'] + Σ_{l'<l} Σ_h head_output[l',h] + = embedding + Σ_{l'<l} (mlp_output[l'] + attn_output[l']) + = residual_l ← the standard residual stream! +``` + +All heads in the same layer l receive the SAME input (because A=1 gives +the same weighted sum for all target heads). This equals the standard +`residual_l`. Therefore **A=1 exactly reproduces vanilla OLMo2-1B**. ✓ + +There is NO double-counting because each head output appears exactly once +in the sum. Head (0,5)'s output is included via the gated sum (source 3), +and it flows to layer 8 only via that direct A[0*16+5, 8*16+h] entry, NOT +also through intermediate layers. The intermediate layers' head outputs +are separate entries in the sum. + +#### MLP Execution (not gated, shared input) + +After all 16 heads in layer l compute, the MLP runs: + +```python +# MLP input = standard aggregation (same as baseline when A=1) +mlp_input_l = embedding + + Σ_{l'<l} mlp_output[l'] + + Σ_{l'<=l} attn_output[l'] # includes current layer's heads +# where attn_output[l'] = Σ_h head_output[l', h] (UNGATED sum of this layer's heads) + +# IMPORTANT: MLP always sees the UNGATED sum of its own layer's head outputs. +# The gating by A only affects what OTHER layers' heads receive. +# This is necessary for baseline reproduction. + +mlp_output[l] = MLP_l(RMSNorm(mlp_input_l)) +``` + +The MLP input includes the **ungated** sum of the current layer's head +outputs. Gating only affects cross-layer routing. This design ensures +that A does not interfere with the MLP's computation within its own layer. + +#### Layer 0 Special Case + +Layer 0's 16 heads have no prior head outputs (no nodes with layer < 0). +Their input is simply: +```python +input_j (j in layer 0) = embedding # no prior heads, no prior MLPs +``` +This requires no special handling — the formula naturally produces +`embedding` when there are no prior sources. + +### 2.3 The Structure Predictor + +The predictor takes the current context and outputs A. Architecture: + +``` +Input: raw text → Qwen tokenizer → qwen_ids [batch, qwen_seq_len] + (qwen_seq_len may differ from OLMo's seq_len — that's fine, + because mean pooling collapses to a single vector per sequence) + │ + ▼ + ┌─────────────────────────────────┐ + │ Qwen-3-Embedding-0.6B │ + │ HF ID: "Qwen/Qwen3-Embedding-0.6B" │ + │ │ + │ This is a TEXT EMBEDDING MODEL, │ + │ NOT a generative LLM. It takes │ + │ text and produces a single │ + │ fixed-size vector per sequence. │ + │ │ + │ FROZEN — no gradients ever │ + │ Uses its OWN tokenizer │ + │ (separate from OLMo's) │ + │ │ + │ Input: raw text → Qwen tokenizer│ + │ → Qwen forward → last_hidden │ + │ Pooling: mean over seq_len dim │ + │ → embedding e ∈ R^d │ + │ (d = model.config.hidden_size │ + │ — do NOT hardcode, query at │ + │ runtime from the model) │ + │ │ + │ e is a SINGLE VECTOR per │ + │ sequence — it summarizes the │ + │ entire context. The predictor │ + │ MLP then maps e → A matrix. │ + │ Qwen has nothing to do with │ + │ OLMo's vocabulary or tokens. │ + └─────────────────────────────────┘ + │ + ▼ + ┌─────────────────────────────────┐ + │ MLP Decoder (TRAINABLE) │ + │ │ + │ Linear(d, hidden_dim) │ + │ → GELU │ + │ → Linear(hidden_dim, hidden_dim)│ + │ → GELU │ + │ → two heads: │ + │ Linear(hidden_dim, 256 * r) │ → reshape to U ∈ R^{256×r} + │ Linear(hidden_dim, 256 * r) │ → reshape to V ∈ R^{256×r} + │ │ + │ hidden_dim = 1024 (default) │ + │ r = rank hyperparameter │ + │ ablate: r ∈ {8, 16, 32, 64} │ + │ default: r = 32 │ + └─────────────────────────────────┘ + │ + ▼ + ┌─────────────────────────────────┐ + │ Low-Rank Logits │ + │ │ + │ Z = U @ V.transpose(-1, -2) │ + │ Z ∈ R^{batch × 256 × 256} │ + │ │ + │ This is the logit matrix before │ + │ any gating or masking. │ + └─────────────────────────────────┘ + │ + ▼ + ┌─────────────────────────────────┐ + │ Block-Upper-Triangular Mask │ + │ │ + │ mask[i,j] = 1 if j//16 > i//16 │ + │ (layer(j) > layer(i)) │ + │ │ + │ ⚠ Do NOT use torch.triu() — │ + │ it allows same-layer connections│ + │ │ + │ Z_masked = Z * mask + (-1e9) * (1 - mask) │ + │ (force invalid positions to │ + │ -inf so sigmoid → 0) │ + └─────────────────────────────────┘ + │ + ▼ + ┌─────────────────────────────────┐ + │ Gumbel-Sigmoid (3 modes) │ + │ │ + │ MODE 1 — TRAINING: │ + │ G ~ Logistic(0, 1) │ + │ (= log(U) - log(1-U), │ + │ U ~ Uniform(0,1)) │ + │ A = σ((Z_masked + G) / τ) │ + │ Gumbel noise G adds stochastic│ + │ exploration. Gradients flow │ + │ through σ naturally (this is │ + │ continuous relaxation, NOT │ + │ STE — no straight-through │ + │ estimator needed here). │ + │ │ + │ MODE 2 — EVAL SOFT: │ + │ A = σ(Z_masked / τ) │ + │ NO Gumbel noise. Deterministic│ + │ soft gates. Values in (0,1). │ + │ Used for eval/nll_soft metric.│ + │ │ + │ MODE 3 — EVAL HARD: │ + │ A = (Z_masked > 0).float() │ + │ NO Gumbel noise. Binary 0/1. │ + │ Threshold at logit=0 (= prob │ + │ 0.5). Used for eval/nll_hard │ + │ and for final inference. │ + │ │ + │ τ = temperature, annealed │ + │ during training (see §3) │ + └─────────────────────────────────┘ + +**WHY Gumbel-Sigmoid — the gradient problem and its solution:** + +At inference we want binary A ∈ {0, 1} (discrete routing decisions). +But discrete decisions have zero gradient — you can't backprop through +`(Z > 0).float()`. The predictor MLP would receive no learning signal. + +Gumbel-Sigmoid is a **continuous relaxation** (also called the "concrete +distribution" / "surrogate gradient" technique): + +``` +Training: A = σ((Z + G) / τ) ← continuous, differentiable + ∂A/∂Z = A(1-A)/τ ← well-defined gradient + backprop: ∂Loss/∂Z = ∂Loss/∂A · ∂A/∂Z + +Inference: A = (Z > 0).float() ← discrete, but no gradient needed +``` + +The complete gradient chain during training: +``` +Loss (NLL from OLMo) + → ∂Loss/∂logits (OLMo's output layer) + → ∂logits/∂head_inputs (through OLMo's frozen layers — computed + but NOT used to update OLMo weights) + → ∂head_inputs/∂A (the gate multiplication: input_j = Σ A[i,j] * head_out[i]) + → ∂A/∂Z (sigmoid derivative: A(1-A)/τ — THIS is the surrogate) + → ∂Z/∂(U,V) (low-rank matmul: Z = UV^T) + → ∂(U,V)/∂MLP_params (the trainable predictor MLP) + → optimizer updates MLP_params +``` + +Key points: +- OLMo is frozen but its forward computation is ON the gradient tape + (it's a differentiable function of A). Gradients flow THROUGH OLMo + back to A, even though OLMo's own parameters don't get updated. +- The sigmoid σ acts as a differentiable surrogate for the discrete + threshold. As τ → 0, σ((Z+G)/τ) approaches a step function, but + during training τ is always > 0 so gradients are always nonzero. +- Gumbel noise G provides stochastic exploration: even if Z is small, + the noise can push A above or below 0.5, letting the model discover + which connections matter. +- NO straight-through estimator (STE) is needed. The sigmoid itself + provides the gradient. STE would be needed if we hard-thresholded + during training, but we don't. + │ + ▼ + ┌─────────────────────────────────┐ + │ Cascading Gate │ + │ │ + │ Purpose: if node j has no │ + │ incoming edges, it has no info │ + │ to propagate, so kill outgoing. │ + │ │ + │ ONE-PASS computation (not │ + │ sequential by layer): │ + │ │ + │ 1. Compute ALL incoming sums: │ + │ inc_j = Σ_i A[i, j] ∀j │ + │ 2. Compute ALL gates: │ + │ g_j = σ(k * inc_j) ∀j │ + │ 3. Apply ALL gates at once: │ + │ A[j, :] *= g_j ∀j │ + │ │ + │ This is a single vectorized op, │ + │ NOT a layer-by-layer cascade. │ + │ All incoming sums use the │ + │ ORIGINAL A values (before any │ + │ gates are applied). │ + │ │ + │ k = 5.0 (fixed scalar) │ + │ k can be made learnable later. │ + │ This is fully differentiable. │ + │ │ + │ EVAL HARD MODE: After hard- │ + │ thresholding A to binary 0/1, │ + │ the cascading gate is also │ + │ hard: g_j = 1 if inc_j > 0, │ + │ else g_j = 0. (Because σ(5*0) │ + │ = 0.5 would be wrong for │ + │ binary gates — a disconnected │ + │ node should be fully killed, │ + │ not half-alive.) │ + │ │ + │ SOFT MODE NOTE: In training and │ + │ eval-soft mode, the cascading │ + │ gate uses σ(k * inc_j) with the│ + │ ORIGINAL (pre-gate) A values. │ + │ If all incoming gates are small │ + │ (e.g. 0.01 each), inc_j can │ + │ still be > 0, giving g_j > 0.5.│ + │ This is INTENTIONAL: in soft │ + │ mode, "weakly connected" is a │ + │ valid state (the gradient can │ + │ still flow). The cascading gate │ + │ is a soft penalty, not a hard │ + │ kill, during training. │ + └─────────────────────────────────┘ + │ + ▼ + Output: A ∈ [0,1]^{batch × 256 × 256}, block-upper-triangular + (30,720 free entries, rest forced to 0) +``` + +### 2.4 End-to-End Pipeline + +``` +raw text ──→ [Qwen tokenizer] ──→ qwen_ids ──→ [Qwen encoder] ──→ e ──→ [Predictor MLP] ──→ A + │ │ + │ ▼ + └──→ [OLMo tokenizer] ──→ olmo_ids ──→ [OLMo2-1B modified forward with A] ──→ logits ──→ NLL + │ + ∇ backprop to predictor MLP +``` + +**Tokenization**: Qwen and OLMo have DIFFERENT tokenizers and vocabularies. +The dataloader produces raw text strings. Each model tokenizes independently: +```python +# In the dataloader or pipeline: +raw_text = batch["text"] # list of strings +qwen_ids = qwen_tokenizer(raw_text, ...)["input_ids"] # Qwen's token IDs +olmo_ids = olmo_tokenizer(raw_text, ...)["input_ids"] # OLMo's token IDs +# These are DIFFERENT tensors with DIFFERENT lengths and vocabularies. +# Qwen produces a pooled embedding (one vector per sequence). +# OLMo produces per-token logits for NLL computation. +``` + +- Qwen and OLMo are FROZEN (Phase 1). Only the MLP decoder trains. +- The same **raw text** goes to both Qwen and OLMo, but they use + **separate tokenizers**. Qwen tokenizes independently to produce an + embedding; OLMo tokenizes independently for language modeling. + Token IDs are NOT shared between the two models. +- Qwen's output (a single pooled vector per sequence) goes to the + predictor MLP → A matrix. OLMo uses A in its modified forward pass. +- Loss = NLL (from OLMo) + λ · mean(A) (sparsity regularization) + +--- + +## 3. Training — Exact Specification + +### 3.1 Phase 1: Learn Topology (IMPLEMENT THIS) + +| What | Frozen/Trainable | +|------|-----------------| +| OLMo2-1B (`allenai/OLMo-2-0425-1B`) | ❄ FROZEN, no grad | +| Qwen-3-Embedding (`Qwen/Qwen3-Embedding-0.6B`) | ❄ FROZEN, no grad | +| Structure Predictor (MLP decoder) | 🔥 TRAINABLE | + +| Hyperparameter | Value | Notes | +|----------------|-------|-------| +| Dataset | Dolma v1.7 (`allenai/dolma`, `name="v1_7"`, streamed) | Specify version explicitly | +| Token budget | 5–10B | Configurable | +| Sequence length | 1024 | OLMo token count, matches oracle search | +| Batch size | 32 (start) | Reduce if OOM | +| Learning rate | 3e-4 | For predictor MLP only | +| Optimizer | AdamW | β1=0.9, β2=0.999, weight_decay=0.01 | +| LR schedule | Cosine decay to 0 | Standard | +| τ (temperature) init | 5.0 | | +| τ final | 0.2 | | +| τ schedule | Cosine annealing | τ(t) = τ_f + 0.5(τ_i - τ_f)(1 + cos(πt/T)) | +| Sparsity λ max | 0.01 | | +| Sparsity ramp | Linear 0 → λ_max over first 20% of steps | | +| Hardware | 4× A40 (48GB each) | Use DDP | + +**Loss function:** +```python +total_loss = nll_loss + lambda_t * A.mean() +``` +where `lambda_t` ramps linearly from 0 to `lambda_max` over the first 20% +of training steps. + +**τ schedule (exact formula):** +```python +tau_t = tau_final + 0.5 * (tau_init - tau_final) * (1 + math.cos(math.pi * step / total_steps)) +``` + +### 3.1.1 Data Processing Pipeline + +**Dolma version**: Use `allenai/dolma` with `name="v1_7"`. If v1_7 is not +available on HuggingFace, fall back to whatever default version loads. +Verify by printing the dataset info at startup and logging to wandb. + +**Sequence packing** (critical for training efficiency): + +Dolma contains documents of varying lengths. Do NOT pad short documents +or discard them. Instead, **pack multiple documents into fixed-length +sequences**: + +```python +# Pseudocode for sequence packing: +buffer = [] +for doc in dolma_stream: + olmo_tokens = olmo_tokenizer(doc["text"], add_special_tokens=False)["input_ids"] + buffer.extend(olmo_tokens) + buffer.append(olmo_tokenizer.eos_token_id) # document separator + + while len(buffer) >= seq_len + 1: # +1 for labels (next-token prediction) + chunk = buffer[:seq_len + 1] + buffer = buffer[seq_len + 1:] + yield { + "olmo_ids": chunk[:seq_len], # input + "olmo_labels": chunk[1:seq_len+1], # shifted target + "raw_text": olmo_tokenizer.decode(chunk[:seq_len]) # for Qwen + } +``` + +This means: +- No padding ever. Every token contributes to NLL. +- A single training sequence may contain parts of multiple documents, + separated by EOS tokens. The causal mask handles this correctly (each + token can only attend to prior tokens, so cross-document leakage is + minimal and standard practice). +- `raw_text` is decoded from OLMo tokens to feed to Qwen. This is a + slight approximation (Qwen re-tokenizes the decoded text), but Qwen + only needs to understand the gist of the context to produce a good + embedding — exact token alignment doesn't matter. +- **Qwen sees the entire packed sequence** (which may contain multiple + document fragments separated by EOS). This is intentional: the + predictor should condition on exactly what OLMo will process. Qwen's + mean-pooling produces a single summary vector of the full 1024-token + window, which is the right granularity — one A matrix per window. +- Qwen's `seq_len` will differ from OLMo's 1024 (different tokenizer + granularity), but this is fine because Qwen's output is mean-pooled + to a single vector regardless of input length. + +**Qwen input format**: Qwen3-Embedding-0.6B is an embedding model. +**Decision**: Use raw text directly (no prefix). Rationale: our input is +a packed sequence of document fragments, not a search query or passage +in the retrieval sense. Prefixes like "query:" or "passage:" are designed +for retrieval tasks and would be semantically misleading here. The Qwen +encoder just needs a general-purpose text representation. + +Log `qwen_input_prefix: ""` to wandb config for reproducibility. If +future experiments show that a prefix helps, it can be changed via the +`qwen_input_prefix` config field. + +### 3.2 Evaluation Data + +Use a **fixed held-out subset of Dolma** for eval. Implementation: + +```python +# At startup, skip the first N examples to get to a held-out region +eval_dataset = load_dataset("allenai/dolma", name="v1_7", split="train", streaming=True) +eval_dataset = eval_dataset.skip(1_000_000) # skip 1M examples +eval_dataset = eval_dataset.take(1_000) # use next 1K as eval set + +# Pack these into fixed-length sequences using the same packing logic +# as training (§3.1.1), then cache in memory at startup +eval_batches = list(eval_dataloader) # ~1K sequences, fits in RAM +``` + +**Eval caching**: The `skip(1_000_000)` on a streaming dataset is slow +(~minutes). To avoid this cost on every restart, **cache the packed eval +sequences to disk** after the first run: +```python +eval_cache_path = os.path.join(config.save_dir, "eval_cache.pt") +if os.path.exists(eval_cache_path): + eval_batches = torch.load(eval_cache_path) +else: + eval_batches = list(build_eval_dataloader(...)) + torch.save(eval_batches, eval_cache_path) +``` + +**Multi-GPU eval**: Run eval on **rank 0 only**. Other ranks skip eval +and wait at a barrier. This avoids redundant computation and ensures +consistent eval metrics (no need to reduce across GPUs). + +Eval runs in **two modes** (both deterministic, no Gumbel noise): +- **Soft**: A = σ(Z / τ) at current temperature. Reports `eval/nll_soft`. +- **Hard**: A = (Z > 0).float(), binary 0/1. Reports `eval/nll_hard`. +Also report `eval/nll_baseline` with A = all-ones (should be constant). + +### 3.3 Multi-GPU Strategy + +**Standard DDP — each GPU holds a full replica of all three models.** + +Memory budget per GPU (fp16/bf16): +``` +OLMo2-1B parameters: ~2.0 GB +Qwen-0.6B parameters: ~1.2 GB +Predictor MLP: ~0.05 GB +Optimizer states: ~0.1 GB (only predictor params) +───────────────────────────── +Model total: ~3.4 GB + +Activations (seq_len=1024, batch=8): + OLMo per-head forward: ~12-16 GB (per-head inputs means 16 separate + attention computations per layer instead of 1 batched MHA — this + increases intermediate activation storage significantly vs standard) + Qwen forward: ~3 GB + Per-head A storage: ~0.5 GB (256×256 × batch × float32) +───────────────────────────── +Activations total: ~16-20 GB + +Grand total: ~20-24 GB per GPU (with batch=8) +Headroom on 48GB A40: ~24-28 GB +``` + +This fits but is tighter than a standard OLMo forward pass. If OOM, +reduce batch_size to 4 first (halves activation memory). If still OOM, +use **gradient accumulation** to maintain effective batch size: +```python +# Example: effective_batch=8, micro_batch=2, accumulation_steps=4 +accumulation_steps = config.batch_size // config.micro_batch_size +for micro_step in range(accumulation_steps): + loss = forward(micro_batch) / accumulation_steps + loss.backward() +optimizer.step() +optimizer.zero_grad() +``` +Add `micro_batch_size` to config (default: same as `batch_size`, i.e. no +accumulation). The per-head computation path creates ~2-3x more +intermediates than standard batched MHA because we cannot fuse across +heads when each has a different input. + +**Gradient checkpointing**: Since OLMo is frozen (no gradients computed +for its parameters), we do NOT need to store OLMo's intermediate +activations for backprop through OLMo's weights. However, we DO need +gradients to flow through OLMo's forward pass back to A (and then to +the predictor). This means OLMo's forward activations are needed for +the chain rule through A's gate multiplications, but NOT for OLMo's +own parameter gradients. + +If memory is still tight, apply `torch.utils.checkpoint` to the OLMo +layer loop — recompute each layer's forward pass during backward instead +of storing all intermediates. This trades compute for memory. +Gradient checkpointing is optional; try without it first. + +**Storing head_outputs**: The forward pass accumulates `head_outputs` for +all 256 nodes (each `[batch, seq, 2048]`). At fp16 with batch=8, +seq=1024: 256 × 8 × 1024 × 2048 × 2 bytes ≈ 8.6 GB. This is the +dominant memory cost. Gradient checkpointing can reduce this by +recomputing head outputs layer-by-layer during backward. + +Use `DistributedDataParallel` wrapping ONLY on the predictor MLP (the +only trainable component). The frozen Qwen and OLMo models do not need +DDP wrapping — just load them on each GPU. + +```python +# DDP setup +predictor = StructurePredictor(config).to(device) +predictor = DDP(predictor, device_ids=[local_rank]) + +# Frozen models — no DDP needed +olmo = AutoModelForCausalLM.from_pretrained(...).to(device).eval() +qwen = AutoModel.from_pretrained(...).to(device).eval() + +# Gradient disabled for frozen models +for p in olmo.parameters(): p.requires_grad_(False) +for p in qwen.parameters(): p.requires_grad_(False) +``` + +Data parallelism: Dolma is a streaming `IterableDataset` — do NOT use +`DistributedSampler` (which requires map-style datasets). Instead, shard +the stream manually: +```python +dataset = load_dataset("allenai/dolma", name="v1_7", split="train", streaming=True) +dataset = dataset.shard(num_shards=world_size, index=rank) +``` +Each GPU processes a disjoint shard. Gradients are synchronized by DDP. + +### 3.4 o_proj Bias and Weight Layout + +OLMo2-1B uses **no bias** in its linear layers (`bias=False` for q_proj, +k_proj, v_proj, o_proj, and MLP layers). This is standard for modern LLMs. + +**Verify this at runtime:** +```python +assert not model.layers[0].self_attn.o_proj.bias, \ + "Expected no bias in o_proj — update per-head splitting if bias exists" +``` + +If a future model version adds bias, the per-head split must also split +the bias: `bias_h = bias[h * head_dim : (h+1) * head_dim]` for input +projections, or `bias / num_heads` for o_proj (since it's additive). +But for OLMo2-1B, this is not needed. + +**o_proj weight layout** for per-head splitting: +```python +# o_proj.weight: [model_dim, model_dim] = [2048, 2048] +# This maps concatenated head outputs → model_dim +# +# Split by INPUT dimension (head_dim chunks): +# o_proj_h.weight = o_proj.weight[:, h*head_dim : (h+1)*head_dim] +# Shape: [2048, 128] +# +# head_output[l, h] = attn_values_h @ o_proj_h.weight.T +# Shape: [batch, seq, 128] @ [128, 2048] → [batch, seq, 2048] +``` + +### 3.5 Phase 2: Joint CPT (DO NOT IMPLEMENT — FUTURE WORK) + +In Phase 2, OLMo would be unfrozen and co-trained with the predictor using +differential learning rates (OLMo: 3e-5, predictor: 1e-4). The training +loop should be DESIGNED to support this (parameter groups with different +LRs) but the actual unfreezing logic should NOT be implemented yet. + +--- + +## 4. OLMo2-1B Modification — Implementation Guide + +This is the hardest part. The goal: intercept OLMo2-1B's forward pass to +apply the adjacency matrix A without forking the model code. + +### 4.1 Strategy: Hook-Based or Subclass + +**Option A (preferred): Monkey-patch the attention forward.** +- Load the model normally via HuggingFace +- Replace each layer's attention module's forward method with a wrapped + version that accepts and applies A +- Pro: no code duplication, stays in sync with HF updates +- Con: fragile if HF changes internal API + +**Option B: Subclass the model.** +- Subclass `OlmoForCausalLM` and override the relevant methods +- Pro: cleaner, more explicit +- Con: more code to maintain + +Choose whichever is cleaner after inspecting the actual OLMo2 source. +The key requirement is: **when A is not provided or is all-ones, the +output must be IDENTICAL to vanilla OLMo2-1B.** + +### 4.2 What Exactly to Modify + +**Reference pseudocode — this is the exact semantics to implement:** + +```python +def dagformer_forward(model, olmo_ids, A, input_norm="none"): + """ + Args: + model: OLMo2-1B loaded from HuggingFace + olmo_ids: [batch, seq_len] — tokenized by OLMo's tokenizer + A: [batch, 256, 256] — block-upper-triangular gate matrix + (produced by the predictor from Qwen embeddings — but this + function doesn't know about Qwen, it just takes A as input) + input_norm: normalization method for assembled head inputs + Returns: + logits: [batch, seq_len, vocab_size] + """ + batch, seq = olmo_ids.shape + model_dim = 2048 + num_layers = 16 + num_heads = 16 + + # Token embedding (always included, not gated) + embedding = model.embed_tokens(olmo_ids) # [batch, seq, 2048] + + # Storage for per-head outputs (in model_dim space) + head_outputs = {} # (layer, head) -> [batch, seq, model_dim] + mlp_outputs = {} # layer -> [batch, seq, model_dim] + + for l in range(num_layers): + layer_module = model.layers[l] + + # === ASSEMBLE PER-HEAD INPUTS === + # Base: embedding + all prior MLP outputs (shared, ungated) + base_l = embedding.clone() + for prev_l in range(l): + base_l = base_l + mlp_outputs[prev_l] + + # Per-head: add gated head outputs from earlier layers + # This is the KEY difference from standard transformer: + # each head h in layer l gets its OWN input. + per_head_inputs = [] + for h in range(num_heads): + j = l * num_heads + h + assembled = base_l.clone() # [batch, seq, model_dim] + + # Gated sum of all prior head outputs + gated_sum = torch.zeros_like(base_l) + for src_l in range(l): + for src_h in range(num_heads): + i = src_l * num_heads + src_h + gate = A[:, i, j] # [batch] + gated_sum = gated_sum + gate[:, None, None] * head_outputs[(src_l, src_h)] + + # Apply normalization ONLY to gated_sum, then add to base (see §6.1) + assembled = base_l + apply_norm(gated_sum, method=input_norm) + per_head_inputs.append(assembled) + + # === RUN ATTENTION WITH PER-HEAD INPUTS === + # This requires splitting the attention computation: + # 1. Each head h gets its own Q from per_head_inputs[h] + # 2. Each head h gets its own K, V from per_head_inputs[h] + # 3. Each head computes attention independently + # 4. Each head's output is projected by its slice of o_proj + # + # In practice: need to split q_proj, k_proj, v_proj into per-head + # projections, run attention per-head, then apply per-head o_proj. + # + # Efficient batch implementation: + # Stack per_head_inputs → [batch, num_heads, seq, model_dim] + # Apply qkv projection in batch across heads + # Run attention, apply o_proj per-head + stacked = torch.stack(per_head_inputs, dim=1) # [batch, 16, seq, 2048] + normed = layer_module.input_layernorm(stacked) # RMSNorm per-head + # NOTE: RMSNorm(2048) operates on the last dimension. When applied to + # [batch, 16, seq, 2048], it normalizes each head's input independently. + # When A=1, all 16 heads have identical inputs, so RMSNorm produces + # identical outputs — matching the standard single-RMSNorm behavior. + # No numerical divergence occurs in the baseline case. + + for h in range(num_heads): + q = q_proj_head(normed[:, h], head=h) # [batch, seq, head_dim] + k = k_proj_head(normed[:, h], head=h) + v = v_proj_head(normed[:, h], head=h) + attn_out_h = attention(q, k, v) # [batch, seq, head_dim] + head_outputs[(l, h)] = o_proj_head(attn_out_h, head=h) # [batch, seq, model_dim] + + # === RUN MLP (ungated, standard) === + # MLP input = base_l + ungated sum of THIS layer's head outputs + attn_output_l = sum(head_outputs[(l, h)] for h in range(num_heads)) + mlp_input = base_l + attn_output_l # standard residual + mlp_outputs[l] = layer_module.mlp(layer_module.post_attention_layernorm(mlp_input)) + + # Final: assemble output state = embedding + all heads + all MLPs + # IMPORTANT: final output uses UNGATED sum of ALL head outputs. + # A only controls intermediate routing (which heads feed into which). + # The final output is NOT gated — every head's output contributes + # equally to the final state. This is a deliberate design choice: + # (1) it matches the standard residual stream when A=1, + # (2) it means A controls *how* heads compute (via their inputs), + # not *whether* their outputs are used in the final prediction. + final_state = embedding + for l in range(num_layers): + final_state = final_state + mlp_outputs[l] + for h in range(num_heads): + final_state = final_state + head_outputs[(l, h)] + + logits = model.lm_head(model.norm(final_state)) + return logits +``` + +⚠ **This pseudocode is for SEMANTIC CLARITY. The actual implementation +MUST use batched operations for performance:** + +```python +# Efficient version for the gated sum: +# Stack all prior head outputs into a tensor +prior_outputs = torch.stack([head_outputs[(l', h')] + for l' in range(l) for h' in range(16)], dim=1) +# prior_outputs: [batch, l*16, seq, model_dim] + +# Slice A for connections into this layer +A_slice = A[:, :l*16, l*16:(l+1)*16] # [batch, l*16, 16] + +# Batched gated sum: einsum or matmul +# per_head_gated[h] = sum_i A_slice[:, i, h] * prior_outputs[:, i] +per_head_gated = torch.einsum('bih,bisd->bhsd', A_slice, prior_outputs) +# per_head_gated: [batch, 16, seq, model_dim] + +# Apply normalization ONLY to per_head_gated, then add base (consistent with §6.1) +per_head_gated = apply_norm(per_head_gated, method=input_norm) +assembled = base_l.unsqueeze(1) + per_head_gated # [batch, 16, seq, model_dim] +``` + +**Splitting Q/K/V projections per-head**: OLMo2's `q_proj`, `k_proj`, +`v_proj` are single linear layers projecting from `model_dim` to +`num_heads * head_dim`. To apply different inputs per head: + +```python +# Option A: reshape the weight matrix and apply per-head +W_q = layer.self_attn.q_proj.weight # [2048, 2048] +W_q_heads = W_q.view(num_heads, head_dim, model_dim) # [16, 128, 2048] +# For each head h: q_h = assembled[:, h] @ W_q_heads[h].T + +# Option B: run the full projection on each head's input separately +# Less efficient but simpler + +# Option C (RECOMMENDED): run full projection on stacked inputs +# assembled: [batch, 16, seq, 2048] +# Reshape to [batch*16, seq, 2048], run q_proj, reshape back +``` + +**RoPE (Rotary Position Embeddings)**: OLMo2 applies RoPE to Q and K +AFTER projection, BEFORE the attention dot product. This is critical for +per-head computation: + +```python +# Standard OLMo2 attention (simplified): +q = q_proj(hidden_states) # [batch, seq, num_heads * head_dim] +k = k_proj(hidden_states) # [batch, seq, num_kv_heads * head_dim] +q, k = apply_rotary_emb(q, k, cos, sin, position_ids) +# Then attention(q, k, v) + +# DAGFormer per-head version: +# For each head h with its own input: +q_h = q_proj_head(assembled_h, head=h) # [batch, seq, head_dim] +k_h = k_proj_head(assembled_h, head=h) # [batch, seq, head_dim] +q_h, k_h = apply_rotary_emb(q_h, k_h, cos, sin, position_ids) # SAME cos/sin/position_ids for all heads +v_h = v_proj_head(assembled_h, head=h) # no RoPE on V +attn_out_h = attention(q_h, k_h, v_h) +``` + +The cos/sin cache and position_ids are **shared across all heads** (they +depend on sequence position, not on head identity). Extract them once from +the model's rotary embedding module at the start of each layer, then reuse +for all 16 heads. If the implementation fails to apply RoPE, the baseline +reproduction sanity check WILL fail because position information is lost. + +**OLMo2-1B uses standard MHA (NOT GQA)**: OLMo2-1B has +`num_attention_heads = 16` and `num_key_value_heads = 16` (every Q head +has its own K and V head). This means per-head splitting is straightforward +— no KV sharing complications. **Verify at runtime:** +```python +config = model.config +assert config.num_attention_heads == 16 +assert config.num_key_value_heads == 16, \ + f"Expected MHA (16 KV heads), got {config.num_key_value_heads} — GQA detected, update splitting logic" +``` +If a future model uses GQA, the per-head splitting must account for shared +KV heads. But OLMo2-1B does not require this. + +**Causal attention mask**: The standard causal mask within each head's +self-attention (preventing token t from attending to tokens t+1, t+2, ...) +is UNCHANGED. DAGFormer's adjacency matrix A controls **cross-layer** +routing (which heads' outputs feed into which heads' inputs). The causal +mask controls **within-sequence** attention (which token positions can +attend to which). These are orthogonal — A operates at the head/layer +level, causal mask operates at the token position level. Use OLMo's +existing causal mask implementation as-is. + +### 4.3 Sanity Checks (MUST PASS before proceeding) + +1. **Baseline reproduction**: Set A[i,j]=1 for all 30,720 valid cross-layer + entries, `input_norm: "none"`. NLL must match vanilla OLMo2-1B within + 0.01 nats. This works because A=1 makes every head's input equal to the + standard residual stream (see §2.2 proof). If this fails, the gate + injection or per-head splitting logic is wrong. +2. **A = all-zeros** (all 30,720 entries = 0): Every head sees only + embedding + MLP outputs, no cross-layer attention routing. NLL should + be significantly higher than baseline. +3. **A = random in [0,1]**: NLL should be between the all-ones and all-zeros + cases. +4. **Gradient check**: Create A as a leaf tensor with `requires_grad=True`, + compute NLL, call `.backward()`, verify `A.grad` is not None and has + nonzero entries at all 30,720 valid positions. +5. **Normalization smoke test**: For each `input_norm` in {gate_mean, + rms_post, ln_post, rms_pre}, run forward with A=all-ones. NLL will NOT + match baseline (normalization changes the scale), but must be finite + (no NaN/Inf). This confirms the norm implementations don't crash. +6. **Per-head input divergence**: Set A to a matrix where different heads + in the same layer have different gate values. Verify that the per-head + inputs are actually different (not collapsed to the same tensor). + This confirms the per-head routing works. + +--- + +## 5. Directory Structure + +``` +dagformer/ +├── CLAUDE.md # THIS FILE — the complete spec +├── README.md # Brief public description +├── pyproject.toml # Dependencies +│ +├── configs/ +│ ├── sanity_check.yaml # 1K steps, verify baseline NLL reproduction +│ ├── ablation_rank.yaml # r ∈ {8, 16, 32, 64} +│ ├── ablation_tau.yaml # τ_init/τ_final sweep +│ ├── ablation_lambda.yaml # sparsity coefficient sweep +│ └── phase1_full.yaml # Full Phase 1 run +│ +├── src/ +│ ├── __init__.py +│ ├── model/ +│ │ ├── __init__.py +│ │ ├── olmo_graph.py # Modified OLMo2 forward with A injection +│ │ ├── predictor.py # Qwen encoder + MLP decoder + Gumbel + cascade +│ │ └── pipeline.py # Combines predictor + OLMo into single forward +│ ├── data/ +│ │ ├── __init__.py +│ │ └── dolma.py # Streaming dataloader: produces raw text, +│ │ # tokenizes with BOTH Qwen and OLMo tokenizers, +│ │ # returns {qwen_ids, olmo_ids, labels} +│ ├── training/ +│ │ ├── __init__.py +│ │ ├── trainer.py # Training loop (pure PyTorch + DDP) +│ │ ├── schedulers.py # τ annealing, λ ramp, LR schedule +│ │ └── checkpointing.py # Save/load predictor + optimizer + schedule state +│ └── utils/ +│ ├── __init__.py +│ ├── logging.py # Wandb integration +│ └── topology.py # A matrix analysis utilities +│ +├── scripts/ +│ ├── train.py # Entry: python scripts/train.py --config configs/X.yaml +│ ├── eval.py # Evaluate NLL with/without predictor +│ ├── sanity_check.py # Verify A=1 reproduces baseline +│ └── visualize_topology.py # Plot A matrices and gate distributions +│ +└── tests/ + ├── test_olmo_graph.py # Baseline reproduction test + ├── test_predictor.py # Shape and gradient tests + └── test_gumbel.py # Gumbel-Sigmoid correctness +``` + +**File responsibilities (be precise):** + +- `olmo_graph.py`: ONLY handles injecting A into OLMo's forward. Does NOT + know about the predictor, Qwen, or training. Exports a function or class + that takes `(model, olmo_ids, A)` and returns logits. + +- `predictor.py`: ONLY the structure predictor. Takes raw text (or + pre-tokenized Qwen IDs), returns A. Contains Qwen loading, Qwen + tokenizer, MLP, Gumbel-Sigmoid, cascading gate. Does NOT know about + OLMo or training. + +- `pipeline.py`: Glue. Takes raw text (or pre-tokenized batch dict with + both qwen_ids and olmo_ids), calls predictor to get A, calls modified + OLMo forward with A, returns loss. This is what the trainer calls. + +- `trainer.py`: Pure training loop. Loads config, builds pipeline, runs + forward/backward/step. Handles DDP, logging, checkpointing. No model + logic here. + +--- + +## 6. Implementation Order (FOLLOW THIS EXACTLY) + +### Step 0: Project Scaffolding +- Create directory structure above +- Write `pyproject.toml`: + ``` + dependencies: torch>=2.2, transformers>=4.40, datasets, wandb, pyyaml, einops + ``` +- Verify model loading: + ```python + from transformers import AutoModelForCausalLM, AutoModel, AutoTokenizer + olmo = AutoModelForCausalLM.from_pretrained("allenai/OLMo-2-0425-1B") + qwen = AutoModel.from_pretrained("Qwen/Qwen3-Embedding-0.6B") + ``` +- Verify Dolma streaming: + ```python + from datasets import load_dataset + ds = load_dataset("allenai/dolma", name="v1_7", split="train", streaming=True) + print(next(iter(ds))) # verify a document loads + ``` + +### Step 1: `olmo_graph.py` — Modified OLMo Forward +**This is the foundation. Get it right.** + +1. Load OLMo2-1B, inspect its architecture: + - Find the attention layer class name + - Find where head outputs are computed and merged + - Find the residual connection logic +2. Implement adjacency injection +3. Run sanity checks (§4.3) — ALL FOUR must pass +4. Do not proceed to Step 2 until Step 1 passes all checks + +### Step 2: `predictor.py` — Structure Predictor +1. Implement Qwen encoder wrapper (frozen, mean-pooled) +2. Implement MLP decoder with low-rank heads +3. Implement Gumbel-Sigmoid with 3 modes: (a) training: noise + τ, + (b) eval soft: σ(Z/τ) no noise, (c) eval hard: (Z>0).float() no noise +4. Implement block-upper-triangular mask (based on layer index, NOT torch.triu) +5. Implement cascading gate +6. Test: output shape [batch, 256, 256], values in [0,1], block-upper-tri +7. Test: `loss = A.sum(); loss.backward()` produces gradients in MLP params + +### Step 3: `pipeline.py` — End-to-End +1. Wire predictor output into OLMo modified forward +2. Verify full gradient chain: NLL.backward() updates predictor MLP +3. Profile memory on single GPU (must fit seq_len=1024, batch=1 on 48GB) + +### Step 4: Training Infrastructure +1. YAML config → Python dataclass (not raw dicts) +2. `schedulers.py`: τ annealing, λ ramp, LR cosine decay +3. `trainer.py`: training loop +4. `logging.py`: Wandb metrics (see §7) +5. `checkpointing.py`: save/load + +### Step 5: Sanity Check Training Run +- Config: `sanity_check.yaml` — 1K steps, batch=4, high τ=5.0 +- Verify: loss decreases over 1K steps +- Verify: A is not collapsing (not all-ones, not all-zeros) +- Verify: gradient norms are reasonable +- **STOP HERE if loss doesn't decrease.** Debug before proceeding. + +### Step 6: Ablations (later) +- Rank r, τ schedule, sparsity λ, cascading gate on/off +- **Input normalization** (see §6.1 below) + +### 6.1 Input Normalization Ablation (IMPORTANT) + +When head j receives gated inputs from multiple source heads across +different layers, the magnitudes of these representations can differ +significantly. Layer 0 outputs and layer 14 outputs live at different +scales. The choice of how to normalize the aggregated input to each +head is a critical design decision. + +**The problem** — referring to the per-head assembly from §2.2: +```python +gated_sum = Σ_{i: layer(i) < l} A[i,j] * head_output[i] +# If 50 sources have A[i,j] > 0, gated_sum has much larger magnitude +# than any single head_output. Scale depends on sparsity pattern of A. + +assembled = base_l + gated_sum # base_l = embedding + prior MLPs +# The gated_sum component can dwarf or be dwarfed by base_l. +``` + +**Normalization is applied ONLY to the gated_sum, before adding to base_l:** +```python +assembled = base_l + normalize(gated_sum, method=config.input_norm) +``` +This way, base_l (which is the standard, ungated component) is preserved +as-is, and only the novel gated routing is normalized. + +**Ablation candidates (implement all, sweep in configs):** + +| ID | Method | Formula | Learnable params | Rationale | +|----|--------|---------|-----------------|-----------| +| `none` | Raw weighted sum | `gated_sum` as-is | 0 | Baseline. When A=1 this reproduces vanilla OLMo. | +| `gate_mean` | Divide by gate sum | `gated_sum / (Σ_i A[i,j] + ε)` | 0 | Normalizes for fan-in. ε=1e-8. | +| `rms_post` | RMSNorm after sum | `RMSNorm(gated_sum)` | 2048 (one gain vector) | One shared `nn.RMSNorm(model_dim)` instance, applied to each head's gated_sum. | +| `ln_post` | LayerNorm after sum | `LayerNorm(gated_sum)` | 2×2048 (gain + bias) | One shared `nn.LayerNorm(model_dim)` instance. Affine params are trainable and counted as predictor params. | +| `rms_pre` | RMSNorm each source before sum | `Σ_i A[i,j] * RMSNorm_i(head_output[i])` | 256 × 2048 = 524,288 | One `nn.RMSNorm(model_dim)` **per source node** (256 total). Each head gets its own learnable gain, allowing fine-grained per-head scale correction before mixing. | + +All learnable norm params (if any) are part of the predictor's parameter +group and trained with the same LR and optimizer as the MLP decoder. + +**Default:** Start with `none` (to verify baseline reproduction), then +switch to `gate_mean` for actual training. If that doesn't work, try +`rms_post`. + +**Implementation:** The normalization method is a config string +(`input_norm: "gate_mean"`). The `olmo_graph.py` code dispatches on this +config. All five methods must be implemented. + +**Config example for ablation:** +```yaml +# In configs/ablation_norm.yaml +sweep: + input_norm: ["none", "gate_mean", "rms_post", "ln_post", "rms_pre"] +``` + +--- + +## 7. Logging & Monitoring + +Log ALL of these to Wandb at every training step: + +| Metric | Formula / Source | Why | +|--------|-----------------|-----| +| `train/nll` | Cross-entropy loss | Primary objective | +| `train/sparsity_loss` | λ_t · mean(A) | Regularization term | +| `train/total_loss` | nll + sparsity_loss | What optimizer sees | +| `eval/nll_soft` | NLL with **deterministic** soft gates: A = σ(Z / τ), NO Gumbel noise | Smooth relaxation perf | +| `eval/nll_hard` | NLL with hard gates: A = (Z > 0).float(), NO Gumbel noise | Inference-mode perf | +| `eval/nll_baseline` | NLL with A = all-ones | Should be constant | +| `topology/mean_A` | mean(A) | Overall gate activation | +| `topology/seq_gate_frac` | Fraction of sequential gates > 0.5 | | +| `topology/hyp_gate_frac` | Fraction of hyperconnection gates > 0.5 | | +| `topology/jaccard_var` | Variance of pairwise Jaccard across batch | Context-dependency | +| `schedule/tau` | Current temperature | | +| `schedule/lambda` | Current sparsity coefficient | | +| `grad/predictor_norm` | Total L2 norm of predictor gradients | | + +**Collapse alarm:** If `topology/mean_A` < 0.01 or > 0.99 for 100 +consecutive steps, log a WARNING. The predictor has degenerated. + +--- + +## 8. Config Format + +Use YAML. Example (`configs/sanity_check.yaml`): + +```yaml +# Model +olmo_model_id: "allenai/OLMo-2-0425-1B" +qwen_model_id: "Qwen/Qwen3-Embedding-0.6B" + +# Predictor +predictor_hidden_dim: 1024 +predictor_rank: 32 +cascading_gate_k: 5.0 +input_norm: "none" # one of: none, gate_mean, rms_post, ln_post, rms_pre + # use "none" for sanity check, "gate_mean" for training + +# Data +dataset: "allenai/dolma" +dataset_name: "v1_7" # Dolma version / subset +seq_len: 1024 # OLMo token count per packed sequence +batch_size: 4 +micro_batch_size: 4 # per-GPU micro batch; if < batch_size, use gradient accumulation +qwen_input_prefix: "" # use raw text directly (see §3.1.1) + +# Eval +eval_skip: 1000000 # skip this many examples to reach held-out region +eval_size: 1000 # number of eval sequences (cached in memory) + +# Training +total_steps: 1000 +lr: 3e-4 +weight_decay: 0.01 +optimizer: "adamw" # only adamw supported + +# Schedules +tau_init: 5.0 +tau_final: 0.2 +tau_schedule: "cosine" +lambda_max: 0.0 # no sparsity for sanity check +lambda_warmup_frac: 0.2 + +# Logging +wandb_project: "dagformer" +wandb_run_name: "sanity-check" +log_every: 10 +eval_every: 100 + +# Checkpointing +save_every: 500 +save_dir: "checkpoints/" + +# Hardware +num_gpus: 1 +``` + +Parse into a `@dataclass` with validation. Crash on unknown keys. + +--- + +## 9. Key Invariants (ALWAYS enforce) + +1. **Baseline reproduction**: A=1 (all 30,720 entries) with `input_norm: + "none"` → NLL matches vanilla OLMo within 0.01. This validates the + entire gate injection and per-head splitting logic. Test BEFORE and + AFTER any architectural change. + +2. **DAG constraint**: A is block-upper-triangular based on LAYER indices. + A[i,j] = 0 whenever `j//16 <= i//16`. Enforced by multiplicative mask, + never by loss or gradient clipping. Do NOT use `torch.triu()`. + +3. **Gradient flow**: After every forward-backward, assert that all + predictor parameters have non-None, non-zero gradients. + +4. **Memory budget**: Must fit on 4×A40 for seq_len=1024. If OOM, reduce + batch size. Do NOT change the architecture to fix memory. + +5. **Frozen models stay frozen**: OLMo and Qwen must have + `requires_grad=False` on ALL parameters. Verify this at startup. + In Phase 1, the only trainable parameters are the MLP decoder. + +6. **Deterministic eval**: Eval uses NO Gumbel noise, ever. Two eval modes: + (a) Soft: A = σ(Z/τ), continuous [0,1]. (b) Hard: A = (Z>0).float(), + binary {0,1}, with hard cascading gate (g_j = 1 if inc_j>0, else 0). + Always report both `eval/nll_soft` and `eval/nll_hard`. + +--- + +## 10. Oracle Search Reference Numbers + +These are the results from the completed oracle search. Use them to +validate that the learned predictor is heading in the right direction. + +| Metric | Value | +|--------|-------| +| Windows evaluated | 50 | +| Window size | 1024 tokens | +| Improvement rate | 100% (all 50 improved) | +| Baseline NLL (median) | 2.58 | +| Oracle NLL (median) | 0.12 | +| NLL delta (median) | +2.38 | +| Oracle sequential gates ON | ~91% | +| Oracle hyperconnection gates ON | ~70% | +| Oracle search steps per window | 500 | +| Oracle search method | Surrogate gradient (STE) | + +The learned predictor does NOT need to match oracle performance. The +**decision gate** for Phase 1 success is: predictor NLL ≤ dense baseline +NLL (i.e., the predictor must not make things WORSE). + +--- + +## 11. Dependencies (exact) + +```toml +[project] +name = "dagformer" +version = "0.1.0" +requires-python = ">=3.10" +dependencies = [ + "torch>=2.2", + "transformers>=4.40", + "datasets", + "wandb", + "pyyaml", + "einops", +] +``` + +**NOT used (by design):** +- No HuggingFace Accelerate +- No PyTorch Lightning +- No DeepSpeed +- No custom CUDA kernels + +Multi-GPU via `torch.nn.parallel.DistributedDataParallel` only. + +--- + +## 12. What NOT to Do + +- **Do NOT implement Phase 2** (joint OLMo training). Design the code to + support it (param groups, differential LR) but do not implement unfreezing. +- **Do NOT implement a diffusion-based predictor.** The MLP decoder is + the current design. Diffusion is future work. +- **Do NOT write custom CUDA kernels.** Use dense matmuls with masking. +- **Do NOT support other base models.** Hardcode OLMo2-1B for now. +- **Do NOT use Accelerate or Lightning.** Pure PyTorch. +- **Do NOT run hyperparameter search.** Manual ablations only. +- **Do NOT fork OLMo2's source code.** Load from HF and modify via + hooks, monkey-patching, or subclassing. +- **Do NOT use `nn.DataParallel`.** Use `DistributedDataParallel` only. + +--- + +## 13. Code Style + +- Type hints on all function signatures +- Docstrings on all public functions and classes +- Config as `@dataclass`, not raw dicts +- `assert` for shape checks in every forward pass (e.g., + `assert A.shape == (batch, 256, 256)`) +- No silent failures — crash loudly with informative messages +- Prefer explicit loops over clever one-liners when clarity matters +- One class per file is fine; don't over-split +- Use `einops.rearrange` for complex tensor reshaping (clearer than + chains of `.view().permute().contiguous()`) + +--- + +## 14. Quick Reference — Model IDs and Shapes + +| Thing | Value | +|-------|-------| +| OLMo model ID | `allenai/OLMo-2-0425-1B` | +| Qwen model ID | `Qwen/Qwen3-Embedding-0.6B` | +| OLMo layers | 16 | +| OLMo heads per layer | 16 | +| Total nodes | 256 | +| A matrix shape | `[batch, 256, 256]` | +| A constraint | Block-upper-triangular: `A[i,j]=0` when `j//16 <= i//16` | +| A free entries | 30,720 (cross-layer only) | +| Predictor rank r | 32 (default) | +| Predictor hidden dim | 1024 (default) | +| Sequence length | 1024 | +| Gumbel-Sigmoid τ range | 5.0 → 0.2 | +| Cascading gate k | 5.0 | +| Input normalization | `none` (sanity check), `gate_mean` (training default), ablate all 5 | +| Sparsity λ range | 0 → 0.01 | |