Release verified completed 5B weights, tokenizer and inference loader

README.md

@@ -17,7 +17,7 @@ trained from scratch on the frozen FineWeb-Edu corpus. It is an inference releas
 the original model tensors are preserved exactly, without optimizer or RNG state.
 
 **Visibility:** public baseline.
-**Parameters:** 373,867,520. **Training context:** 2,048 tokens.
+**Parameters:** 341,099,520. **Training context:** 2,048 tokens.
 The 5B endpoint uses the original 15B-token learning-rate schedule, stopped at 5B.
 See `research_config.json` for the actual architecture and training history policy.
 

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f43364962b0905335d7b61cedcdb75c235539a230911d7157042fefd9f454901
-size 1495492664
+oid sha256:36f524026cfb6054d8ca0226159e78a1e4062715a869241287a593de6670d251
+size 1364420568

provenance.json

@@ -11,8 +11,8 @@
   "model_id": "dense",
   "native_training_context": 2048,
   "new_gpu_evaluation": false,
-  "parameter_count": 373867520,
-  "release_git_revision": "ee32e5d51c2202ca07513745b75107b8fba1a821",
+  "parameter_count": 341099520,
+  "release_git_revision": "9e0034ce9a476545474af9cadc49e2fdf163ada9",
   "source_artifact_sha256": "36984a2c331fa28c0d730e8e8f5d887edf6b65b42f711daee079a447d20ee52f",
   "source_checkpoint_sha256": "010f0d620a9ada040d0e8e35934ec846c45e9df69684759455f68a951ad030dc",
   "source_repository": "https://github.com/Nicholas0228/tail_drop_attn",

release_manifest.json

@@ -6,7 +6,7 @@
     },
     "README.md": {
       "bytes": 2196,
-      "sha256": "d0761fd2b72a9cdb2ce710666ebe10cfdfae1b2af0058a4a79872ea42888552e"
+      "sha256": "17ef029fea91e9f1ea8c02103f123a50f23c2319540bd3d2c7f30b94055e73dc"
     },
     "THIRD_PARTY.md": {
       "bytes": 1629,
@@ -21,12 +21,12 @@
       "sha256": "e2389cabfefc5c39a630ec110cb2f27701c3d79b807454b6b4acf1d40955f7a4"
     },
     "model.safetensors": {
-      "bytes": 1495492664,
-      "sha256": "f43364962b0905335d7b61cedcdb75c235539a230911d7157042fefd9f454901"
+      "bytes": 1364420568,
+      "sha256": "36f524026cfb6054d8ca0226159e78a1e4062715a869241287a593de6670d251"
     },
     "provenance.json": {
       "bytes": 1100,
-      "sha256": "80a8af9e0cc1ae83c8600bf3b943272e08be479029dcc8bff975aa0f69c73402"
+      "sha256": "ab12f01ef61bfa3b3e76d7279698570a037e98815515dc243a86115981d03e95"
     },
     "research_code/pretraining_5b/__init__.py": {
       "bytes": 59,
@@ -80,6 +80,14 @@
       "bytes": 2352,
       "sha256": "16f5e391510f832d14ca312754b46268fff890d225d357462e37e651455a72af"
     },
+    "research_code/pretraining_suffix_min/__init__.py": {
+      "bytes": 67,
+      "sha256": "833cc39de2615b8f92ae34dea999d46926516ce79a9e6840732f927dd17bd7ca"
+    },
+    "research_code/pretraining_suffix_min/sampling.py": {
+      "bytes": 3253,
+      "sha256": "a3aa97878a9e75429da8bc99d9806e532174d8c3c0014acca56d0479ad5b4457"
+    },
     "research_config.json": {
       "bytes": 206,
       "sha256": "57806f0e1869206599088fa814bddddca1121d10f38c4ec4cea643cf90968a70"
@@ -93,8 +101,8 @@
       "sha256": "c2391c2c7dbedaaad4bac2b854ececd3a668bf546cc9c8d665876bcc0796dc66"
     },
     "tied_weights.json": {
-      "bytes": 3,
-      "sha256": "ca3d163bab055381827226140568f3bef7eaac187cebd76878e0b63e9e442356"
+      "bytes": 44,
+      "sha256": "6621f33b8064e92a055e9d6a323805631e93a884410e4a81294277f6a8b91095"
     },
     "tokenizer.json": {
       "bytes": 3505608,

research_code/pretraining_suffix_min/__init__.py

@@ -0,0 +1 @@
+cat: pretraining_suffix_min/__init__.py: No such file or directory

research_code/pretraining_suffix_min/sampling.py

@@ -0,0 +1,83 @@
+"""Minimum-window suffix budgets; pure contract helpers do not import Torch.
+
+This changes the distribution of retained lengths, not attention visibility or
+inference caps. A sampled K always includes the current position within K.
+"""
+from __future__ import annotations
+
+import math
+
+
+def validate_minimum(minimum):
+    if type(minimum) is not int or minimum < 1 or minimum & (minimum - 1):
+        raise ValueError('Minimum retained history must be a positive power of two')
+
+
+def retained_candidates(available, minimum):
+    """Unique ascending candidates, including full available history once."""
+    validate_minimum(minimum)
+    if type(available) is not int or available < 1:
+        raise ValueError('Available history must be a positive integer')
+    if available < minimum:
+        return (available,)
+    values = []
+    value = minimum
+    while value <= available:
+        values.append(value)
+        value *= 2
+    if values[-1] != available:
+        values.append(available)
+    return tuple(values)
+
+
+def retained_from_uniform(available, minimum, uniform):
+    """Scalar reference used to verify equal-probability sampling intervals."""
+    if not math.isfinite(uniform) or not 0 <= uniform < 1:
+        raise ValueError('Uniform variate must lie in [0, 1)')
+    choices = retained_candidates(available, minimum)
+    return choices[int(uniform * len(choices))]
+
+
+def minimum_budget_policy(canonical_policy, minimum):
+    """Bind the new distribution while preserving legacy min1 serialization."""
+    validate_minimum(minimum)
+    if minimum == 1:
+        return canonical_policy
+    if canonical_policy is None:
+        raise ValueError('Minimum-window training requires an explicit canonical budget policy')
+    return {
+        **canonical_policy,
+        'retained_length_distribution': {
+            'schema': 1,
+            'name': 'uniform_unique_dyadic_minimum_plus_full',
+            'minimum_history': minimum,
+            'short_history': 'dense',
+            'full_history_included_once': True,
+            'scope': 'independent_sequence_query_shared_heads_and_middle_layers',
+        },
+    }
+
+
+def sample_suffix_lengths(batch, length, *, minimum, generator, device):
+    """Consume one original-shaped uniform draw; dense short queries still draw.
+
+    CanonicalBudgetPlan calls this separately for each original 16x2048 draw,
+    preserving per-example quantiles and generator advancement on world2/4.
+    The original min1 trainer continues calling its unchanged sample_prefixes.
+    """
+    import torch
+
+    validate_minimum(minimum)
+    if batch < 1 or length < 1:
+        raise ValueError('Sampler dimensions must be positive')
+    n = torch.arange(1, length + 1, device=device, dtype=torch.int64)
+    powers = torch.tensor(
+        [minimum << p for p in range(max(1, length // minimum).bit_length())],
+        device=device,
+    )
+    power_count = torch.bucketize(n, powers, right=True)
+    extra_full = ((n & (n - 1)) != 0).long()
+    count = torch.where(n < minimum, 1, power_count + extra_full)
+    choice = (torch.rand((batch, length), device=device, generator=generator) * count).long()
+    retained = torch.where(choice < power_count, minimum << choice, n)
+    return torch.where(n < minimum, n, retained)

tied_weights.json

@@ -1 +1,3 @@
-{}
+{
+  "lm_head.weight": "embeddings.weight"
+}