dragonscale 0.1.0__tar.gz

dragonscale-0.1.0/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2026 Shih-Ying Yeh (A.K.A. KohakuBlueLeaf)
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

dragonscale-0.1.0/PKG-INFO

@@ -0,0 +1,463 @@
+Metadata-Version: 2.4
+Name: dragonscale
+Version: 0.1.0
+Summary: Keeps weight sizes from drifting during training: finds the scales a model's loss cannot see and normalizes the weights behind them.
+Author-email: "Shih-Ying Yeh(KohakuBlueLeaf)" <apolloyeh0123@gmail.com>
+License: Apache-2.0
+Project-URL: Homepage, https://github.com/KohakuBlueleaf/DragonScale
+Keywords: pytorch,weight-normalization,training-stability,triton
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.9
+Provides-Extra: triton
+Requires-Dist: triton>=3.7; extra == "triton"
+Provides-Extra: tilelang
+Requires-Dist: tilelang; extra == "tilelang"
+Provides-Extra: hf
+Requires-Dist: transformers; extra == "hf"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: black; extra == "dev"
+Requires-Dist: matplotlib; extra == "dev"
+Requires-Dist: kohaku-engine; extra == "dev"
+Provides-Extra: measure
+Requires-Dist: transformers; extra == "measure"
+Requires-Dist: matplotlib; extra == "measure"
+Requires-Dist: kohaku-engine; extra == "measure"
+Dynamic: license-file
+
+# DragonScale
+
+![DragonScale: a pre-norm transformer's free scales drift where the loss cannot see them; EDM2 and nGPT fix every scale by construction; DragonScale finds the free scales in the model you have and gives each one owner, a gain the loss moves. Right: JiT-B/16's stream peak, stock 1,694 against arch 15, and a byte-level LM's loss against learning rate](docs/assets/teaser.png)
+
+A PyTorch library that keeps weight and activation scales from drifting during training.
+
+Some scales in a network can change without changing what the loss sees: in a transformer MLP you
+can multiply `up` by 2 and divide `down` by 2 and get the same output. The loss gives no gradient
+along such a direction, but every optimizer step still moves the weights a little along it, and
+those moves add up. Over training the residual stream, the attention logits and the weight norms
+grow, until fp16 overflows or training diverges. DragonScale finds these directions in any model,
+gives the rows of the weights behind them a fixed length, and keeps one learned gain per row in
+their place. Scales can still grow, but only through gains the loss sees and moves.
+
+**Contents:** [Highlights](#highlights) · [Install](#install) · [Quick start](#quick-start) ·
+[How it works](#how-it-works) · [Why scales drift](#why-scales-drift) · [Analysis](#analysis) ·
+[Patching options](#patching-options) · [Optimizer](#optimizer) · [Checkpoints](#checkpoints) ·
+[Kernels and speed](#kernels-and-speed) · [Measured](#measured) · [Repository](#repository)
+
+## Highlights
+
+- **Finds the free scales of any model** in one forward pass: `ds.analysis.free_scales` follows the
+  scale of every parameter through every op and reports which directions the loss cannot see,
+  which of them move activations, and what the architecture could change.
+- **Patches the model you already have, same output at step 0.** Linear, Conv, ConvTranspose,
+  Embedding and GPT-2's `Conv1D` get unit rows and a gain; tied weights share one gain; `bake`
+  returns stock modules.
+- **The stream stays where the loss can see it.** JiT-B/16 at lr 1e-3: stock's stream peaks at
+  1,694, `arch` at 15. Byte-level LM: stock trains up to lr 5e-3, `arch` up to 1.0, at a lower
+  best loss (0.930 against 0.947).
+- **Costs little.** Triton kernels at 85–96% of memory bandwidth on large weights (V100, RTX 2080
+  Ti, 3090, 4090), against 54–64% for `torch.compile` and 21–25% for eager; a grouped launch for
+  many small layers; under `torch.compile` a patched model runs within 3–6% of stock at the same
+  memory.
+
+## Install
+
+```bash
+uv pip install -e .                  # extras: triton, hf, measure, dev
+```
+
+Needs Python 3.11+ and torch 2.9+.
+
+## Quick start
+
+```python
+import torch
+import dragonscale as ds
+
+model = MyModel()
+
+# Which scales does the loss not see, which of them are risky, what should the model change?
+report = ds.analysis.free_scales(model, example_input)
+print(report.describe())
+
+# Unit rows and a gain on every Linear / Conv / Embedding; the output is unchanged.
+session = ds.apply(model, ds.analysis.suggest(model, example_input, level="matrices"))
+
+optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)   # gains are ordinary parameters
+# ... train as usual ...
+
+session.bake()      # back to stock modules with the trained weights, same output
+```
+
+`ds.apply` works on a model you already have: an existing checkpoint is converted with the same
+output at step 0. Training from scratch works the same way.
+
+## How it works
+
+```mermaid
+flowchart LR
+    model["your model"] --> report["free_scales(model, x)<br/>free directions, dangerous ones,<br/>architecture notes"]
+    report --> plan["suggest(model, x, level)<br/>report / minimal / matrices / absorb"]
+    plan --> apply["ds.apply(model, plan)<br/>unit rows + gain per layer<br/>same output"]
+    apply --> train["train as usual<br/>any optimizer"]
+    train --> out["session.bake()<br/>stock modules"]
+    train --> native["session.canonicalize()<br/>ds.nn modules"]
+```
+
+A patched layer computes with unit rows times a gain; the architecture notes are what a weight
+patch cannot do, and belong in the model code. The `arch` models in [Measured](#measured) follow
+them: every input enters through a norm without affine, the blocks write the residual stream with
+unit rows and gains, and the output layer reads the stream directly with unit rows and no gain.
+There is no final norm: a norm in front of the output hides the stream's scale from the loss and
+makes it free. In JiT the conditioning and in-context tokens are attended to and never written,
+so they keep the scale they entered with.
+
+## Why scales drift
+
+Give every parameter tensor $\theta_i$ a log-scale $s_i$ and look at the model as a function of
+$\theta_i e^{s_i}$. A **free direction** is a vector $s$ along which the output does not change:
+
+```math
+f(\theta_1 e^{s_1}, \dots, \theta_n e^{s_n}) = f(\theta_1, \dots, \theta_n).
+```
+
+Examples: `up` $\times c$ with `down` $\times 1/c$; `q` $\times c$ with `k` $\times 1/c$; a norm's
+gain $\times c$ with every layer that reads it $\times 1/c$; and any weight followed by a norm on
+its own.
+
+The loss is constant along a free direction, so its gradient has no component there. For a weight
+$w$ whose scale is free, $\langle \nabla_w L, w\rangle = 0$, and one step $w' = w - \eta g$ gives
+
+```math
+\lVert w' \rVert^2 = \lVert w \rVert^2 + \eta^2 \lVert g \rVert^2 \ \ge\ \lVert w \rVert^2 .
+```
+
+The norm can only grow. With Adam every entry moves about $\eta$ per step whatever its gradient,
+so a row of $n$ entries grows like a random walk, and a matrix row moves about $\sqrt{n}$ times
+faster than a single gain. A free direction is **dangerous** when it moves activations and has a
+matrix in it: that is where the drift is fast and the network sees it.
+
+## What DragonScale changes
+
+A patched layer uses unit rows times a gain in place of its weight:
+
+```math
+W_{\text{eff}} = \mathrm{diag}(g)\, r\, \frac{W}{\lVert W \rVert_{\text{row}}}
+```
+
+with $r$ the row radius (1 for linears and convs, $\sqrt{\text{dim}}$ for embeddings, so their rows
+have unit RMS). The raw row length $\lVert W \rVert_{\text{row}}$ no longer reaches the output, so
+the scale that drifted now lives in $g$: one number per row, moved by a gradient that is aligned
+with the loss instead of by noise. The gradient reaching $W$ is tangent to its rows,
+
+```math
+\frac{\partial L}{\partial W} = \frac{g\, r}{\lVert W \rVert}\,\bigl(I - \hat w \hat w^\top\bigr)\, \frac{\partial L}{\partial W_{\text{eff}}},
+```
+
+so the raw rows still grow, slowly, and their growth shrinks the effective step of the direction:
+a per-row $1/\sqrt{t}$ schedule.
+
+## Analysis
+
+`ds.analysis.free_scales(model, *inputs)` runs one forward pass, follows the scale of every
+parameter through every op, and solves for the free directions. For a one-block pre-norm LM:
+
+```text
+7 free direction(s), 7 dangerous, 6 constraints over 13 groups
+  dangerous: +1*embed.weight +1*pos.weight +1*blocks.0.out.weight +1*blocks.0.down.weight
+  dangerous: +1*blocks.0.norm1.weight -2*blocks.0.k.weight -1*blocks.0.out.weight
+  dangerous: +1*blocks.0.q.weight -1*blocks.0.k.weight
+  dangerous: +1*blocks.0.v.weight -1*blocks.0.out.weight
+  dangerous: +1*blocks.0.norm2.weight -1*blocks.0.gate.weight -1*blocks.0.down.weight
+  dangerous: +1*blocks.0.up.weight -1*blocks.0.down.weight
+  dangerous: +1*norm.weight -1*head.weight
+architecture:
+  attention calls [0]: q and k trade scale freely; normalize q and k (QK-norm) with a gain on one side only
+  embed.weight, pos.weight, blocks.0.out.weight, blocks.0.down.weight scale the residual stream together; give the model input a fixed scale (an embedding: mode 'unit')
+```
+
+`gate` and `down` do not trade freely: SiLU pins the scale of `gate`.
+
+`ds.analysis.suggest(model, *inputs, level=...)` turns the report into a plan:
+
+| level | plan |
+| --- | --- |
+| `report` | nothing; read the report and its architecture notes |
+| `minimal` | unit rows and a gain only on the layers in a dangerous direction |
+| `matrices` | unit rows and a gain on every Linear / Conv / Embedding |
+| `absorb` | each norm's affine is moved into the layers that read it and frozen; each layer whose row scale can move into the next layer (`v` into `out`, `up` into `down`) gets unit rows without a gain; the rest get unit rows and a gain |
+
+Every level keeps the model's output. `absorb` plans carry the moves; run them before applying:
+
+```python
+plan = ds.analysis.suggest(model, x, level="absorb")
+session = ds.apply(model, ds.analysis.absorb_plan(model, plan, x))
+```
+
+`ds.analysis.absorb(model, source, *inputs, into=None)` makes one move by hand; it finds the
+consumers when `into` is None, refuses a move that is not free, and checks that the output did not
+change. `ds.analysis.scale_jacobian` gives the spectrum of the same scale map numerically.
+
+The architecture notes are what a weight patch cannot do; they belong in the model code.
+
+## Patching options
+
+Plans are data: `ds.make_plan(model, rules)` matches `ds.Rule(match, **options)` against module
+paths or types, first match wins. `ds.apply` accepts a plan, a rule list or a preset name
+(`unit_gain`, `unit`, `linear_unit_gain`, `conv_unit_gain`, `parametrize_unit_gain`).
+
+| option | values | meaning |
+| --- | --- | --- |
+| `mode` | `unit_gain`, `unit`, `plain` | unit rows times a gain, unit rows alone, untouched |
+| `gain` | `channel`, `scalar` | one gain per row, or one for the layer |
+| `init` | `preserve`, `ones`, `reset`, `absorb` | gain = row length (same output), gain = 1, re-initialize then preserve, rows already moved by `absorb_plan` |
+| `install` | `forward`, `parametrize`, `bake` | replace the module's forward, a torch parametrization, write the unit-row weight once |
+| `method` | `normalize`, `rescale` | divide by the row length in every forward, or use the raw rows and let the optimizer rescale them after each step |
+
+Supported layers: Linear, Conv1d/2d/3d, ConvTranspose1d/2d/3d, Embedding and transformers'
+`Conv1D` (GPT-2); `ds.layer_types.register_layer_type` adds more. Tied weights (an LM head shared
+with the embedding) are patched as one layer with one gain.
+
+`ds.nn` has the same layers as modules (`UnitRowLinear`, `UnitRowConv2d`, `UnitRowEmbedding`, ...)
+for new models, and `ds.functional.unit_rows(w, spec, gain)` does it for a bare weight.
+
+## Optimizer
+
+`method="normalize"` needs nothing from the optimizer. `method="rescale"` needs the gradient
+decorator, which works with any torch optimizer through its step hooks:
+
+```python
+decorator = ds.optim.attach(optimizer, session)
+```
+
+Before each step it removes the part of each gradient that would change a row's length; after the
+step it removes the same part from the momentum (Adam's `exp_avg`, SGD's and Muon's
+`momentum_buffer`) and rescales each row back to its radius. Second moments are left alone.
+
+## Checkpoints
+
+| call | result |
+| --- | --- |
+| `session.state_dict("raw")` | raw weights and gains, under the module's own keys plus `<path>.gain` |
+| `session.state_dict("baked")` | the unit-row weights under the stock keys; loads into the stock model |
+| `session.bake()` | the model becomes stock modules with the unit-row weights, same output |
+| `session.canonicalize()` | the patched modules become `ds.nn` modules in place (`native` state dict) |
+| `session.remove()` | everything undone; the weights keep what training wrote |
+
+`load_state_dict(state, format)` reads all three formats. Patched models work with
+`copy.deepcopy`, `torch.save` of the whole model (`forward` install) and FSDP2.
+
+`ds.monitor.Monitor(session)` reads each patched layer's output size, gains and raw row lengths
+during training.
+
+## Kernels and speed
+
+The unit-row op is memory-bound: it reads a weight and writes one of the same size, so its speed
+is its effective bandwidth, the bytes it must move over its time. The Triton kernels read a row
+once when it fits 64 KB. On an 11008 x 4096 weight, forward + backward, they reach 85–96% of the
+card's memory bandwidth, against 54–64% for `torch.compile` and 21–25% for eager. Every fp32
+result matches fp64 to about 1e-7.
+
+![Unit-row forward + backward effective bandwidth per backend on three weight shapes, on a V100, RTX 2080 Ti, 3090 and 4090](docs/assets/kernels.png)
+
+A unit-row Linear against a plain `nn.Linear`, forward + backward with the matmul, fp32 weights
+under fp16 autocast: 96% of the plain layer's speed on a 4096 -> 11008 layer and 99% on 768 ->
+2048, with Triton or `torch.compile`; eager reaches 87% and 96%.
+
+![Unit-row Linear speed as a share of the plain Linear's, per backend, on three layer shapes, on an RTX 2080 Ti, 3090 and 4090](docs/assets/layer_speed.png)
+
+On a 256 -> 768 layer the unit-row launches are a visible share of a 1 ms step where the matmul
+finishes faster than the host launches (73% on the 3090), and hidden where it does not (98% on
+the 2080 Ti). Small layers are bound by launch time, about 0.06 ms per call on an RTX 4090. In eager mode a
+forward pre-hook computes the weights of every eligible layer in one autograd node and a few
+launches (`ds.apply(..., group=True)`, the default); under `torch.compile` the op is traced as
+plain torch ops and fused. Whole training step of the byte-level LM (8 blocks, dim 256, batch
+64 x 256, fp16 AMP) on an RTX 4090:
+
+| preset | eager, per layer | eager, grouped | `torch.compile` |
+| --- | --- | --- | --- |
+| stock | 18.8 ms | 18.8 ms | 15.6 ms |
+| matrices | 30.5 ms | 21.2 ms | 16.6 ms (+6%) |
+| arch | 33.3 ms | 25.5 ms | 16.0 ms (+3%) |
+
+`arch` in eager mode pays for its own extra norms and casts (830 kernel launches against 681);
+compile fuses them. Peak memory under compile: 1.44 GB stock, 1.43 GB matrices, 1.55 GB arch.
+
+| backend | used when | ops |
+| --- | --- | --- |
+| `triton` | CUDA, sm70 and up | every row op, grouped launch |
+| `compile` | CUDA | `torch.compile` of the eager ops |
+| `eager` | always | the reference, fp32 accumulation, fp64 kept |
+
+`ds.apply(model, plan, backend=...)` pins one; by default the fastest available is used.
+Benchmarks: `kogine run measure/ops/speed.py` (the op), `kogine run measure/ops/overhead.py`
+(a layer against the plain layer).
+
+## Measured
+
+These runs check one thing: whether the scales stay put and the loss trains normally. They report
+held-out loss and hidden-state statistics only. All runs are in [`measure/`](measure) and run
+through [kogine](https://github.com/KohakuBlueleaf/KohakuEngine); every number comes from one
+config file.
+
+| preset | what it is |
+| --- | --- |
+| `stock` | the model as written |
+| `matrices` | `suggest(level="matrices")` on the stock model |
+| `arch` | `suggest(level="absorb")` plus the architecture notes done in the model code (see [How it works](#how-it-works)) |
+
+A run **failed** when it diverged or its best held-out loss stayed above 2.0.
+
+### Byte-level LM
+
+A pre-norm transformer (dim 256, 8 blocks, SwiGLU) on the bytes of the Python standard library,
+AdamW, constant learning rate, no weight decay, 5k steps unless stated.
+
+**A: learning rate x precision.**
+
+![Byte-level LM, best held-out loss and peak residual stream against learning rate, fp16 and fp32, per preset](docs/assets/a_lr_precision.png)
+
+| preset | precision | best loss (lr) | largest lr that trains |
+| --- | --- | --- | --- |
+| stock | fp16 | 0.947 (2e-3) | 5e-3 |
+| stock | fp32 | 0.957 (2e-3) | 5e-3 |
+| matrices | fp16 | 0.954 (5e-3) | 5e-2 |
+| matrices | fp32 | 0.955 (5e-3) | 5e-2 |
+| arch | fp16 | 0.930 (5e-3) | 1.0 |
+| arch | fp32 | 0.936 (0.1) | 1.0 |
+
+`arch` trains at every learning rate up to 1.0; past 0.2 its best loss rises (1.05 at 0.5, 1.00
+at 1.0 in fp16). fp32 fails where fp16 does: stock at lr 1e-2 and above ends near loss 2.4 with a stream of 1e6
+to 1e12 in fp32, and diverges in fp16. The precision is not the cause.
+
+**B: longer runs, 16k steps** at lr 2e-3 and 1e-2, fp16.
+
+![Byte-level LM over 16k steps: peak stream, largest raw row length and largest writer gain against step](docs/assets/b_long.png)
+
+| preset | lr 2e-3: best loss, peak stream | lr 1e-2: best loss, peak stream |
+| --- | --- | --- |
+| stock | 0.957, 1,438 | failed at step 1k |
+| matrices | 0.958, 147 | 0.962, 770 |
+| arch | 0.939, 184 | 0.939, 263 |
+
+The corpus is 10.7 MB; past a few thousand steps every preset memorizes it and the held-out loss
+climbs back to 1.1–1.2.
+
+**C and E: depth 4 / 8 / 16 / 32 and width 128 / 256 / 512 / 1024**, peak stream:
+
+<p align="center">
+  <img src="docs/assets/depth.png" width="49%" alt="Byte-level LM, peak stream against depth per preset">
+  <img src="docs/assets/width.png" width="49%" alt="Byte-level LM, peak stream against width per preset">
+</p>
+
+C, depth (best held-out loss):
+
+| preset | lr | depth 4 | depth 8 | depth 16 | depth 32 |
+| --- | --- | --- | --- | --- | --- |
+| stock | 2e-3 | 0.995 | 0.956 | 0.941 | 0.936 |
+| stock | 1e-2 | failed | failed | failed | failed |
+| stock | 5e-2 | failed | failed | failed | failed |
+| matrices | 2e-3 | 0.994 | 0.954 | 0.939 | 0.935 |
+| matrices | 1e-2 | 1.006 | 0.960 | 0.939 | 0.932 |
+| matrices | 5e-2 | 1.110 | 1.147 | 0.991 | failed |
+| arch | 2e-3 | 0.946 | 0.936 | 0.932 | 0.927 |
+| arch | 1e-2 | 0.938 | 0.937 | 0.925 | 0.925 |
+| arch | 5e-2 | 0.943 | 0.953 | 0.938 | 0.925 |
+
+**D: weight decay on stock** (best held-out loss, peak stream).
+
+| weight decay | lr 1e-3 | lr 3e-3 | lr 1e-2 | lr 3e-2 | lr 0.1 |
+| --- | --- | --- | --- | --- | --- |
+| 0.01 | 0.975, 136 | 0.947, 791 | failed | failed | failed |
+| 0.1 | 0.965, 52 | 0.942, 156 | 1.038, 6,790 | 1.299, 25,300 | failed |
+
+Weight decay 0.1 keeps stock training at lr 1e-2 and 3e-2, with the stream still at 7k–25k.
+
+E, width (best held-out loss):
+
+| preset | lr | width 128 | width 256 | width 512 | width 1024 |
+| --- | --- | --- | --- | --- | --- |
+| stock | 2e-3 | 0.990 | 0.955 | 0.948 | failed |
+| stock | 1e-2 | 1.035 | failed | failed | failed |
+| stock | 5e-2 | failed | failed | failed | failed |
+| matrices | 2e-3 | 0.987 | 0.958 | 0.940 | 0.982 |
+| matrices | 1e-2 | 0.977 | 0.967 | 1.121 | 1.187 |
+| matrices | 5e-2 | 0.970 | 1.047 | failed | failed |
+| arch | 2e-3 | 0.970 | 0.938 | 0.932 | 0.953 |
+| arch | 1e-2 | 0.951 | 0.942 | 0.940 | 0.944 |
+| arch | 5e-2 | 0.952 | 0.942 | 0.935 | 0.946 |
+
+Stock's stream grows with width (207, 516, 1,812 at lr 2e-3) and width 1024 diverges at the
+smallest learning rate. `arch` trains at every depth and width at all three learning rates.
+
+**Loss against stream size**, every run of A–E: the frontier per preset (left) and each setup's
+learning-rate sweep as one line (right).
+
+<p align="center">
+  <img src="docs/assets/pareto.png" width="49%" alt="Best held-out loss against peak stream for every byte-level LM run, Pareto frontier per preset">
+  <img src="docs/assets/pareto_setups.png" width="49%" alt="Best held-out loss against peak stream, one line per setup">
+</p>
+
+### JiT-B/16 on Flowers-102
+
+Pixel-space diffusion with the JiT recipe (x-prediction, velocity loss) at 256 px, batch 64, the
+time and class as conditioning tokens, fp16.
+
+![JiT-B/16 at lr 1e-3: image-token stream RMS through depth, and its peak over 2k steps, for stock, arch with a final norm and arch](docs/assets/final_norm.png)
+
+| model, lr 1e-3, 2k steps | stream RMS, block 0 -> 11 | peak | held-out loss |
+| --- | --- | --- | --- |
+| stock | 49 -> 150 | 1,694 | 0.419 |
+| arch with a final norm | 19 -> 93 | 1,239 | 0.425 |
+| arch | 1.4 -> 0.8 | 15 | 0.444 |
+
+A final norm in front of the output makes the stream's scale free, and it grows in the patched
+model as in stock. Without it, and with the context tokens read-only, the stream stays near 1.
+
+![JiT-B/16 over 10k steps per preset and learning rate: held-out velocity loss and peak stream](docs/assets/jit.png)
+
+The loss is compared from step 2k: the first few hundred steps' losses are outliers.
+
+| preset | lr | best held-out loss, steps 2k-10k | peak stream |
+| --- | --- | --- | --- |
+| stock | 2e-4 | 0.264 | 779 |
+| stock | 1e-3 | 0.219 | 4,506 |
+| stock | 3e-3 | 0.247 | 23,804 |
+| matrices | 2e-4 | 0.260 | 710 |
+| matrices | 1e-3 | 0.216 | 1,294 |
+| matrices | 3e-3 | 0.217 | 2,055 |
+| arch | 2e-4 | 0.234 | 29 |
+| arch | 1e-3 | 0.175 | 28 |
+| arch | 3e-3 | 0.183 | 86 |
+
+Over 10k steps stock's stream peaks at 779 to 23,804, rising with the learning rate; `arch` stays
+at 28 to 86, and its held-out loss is the lowest at every learning rate.
+
+Reproduce, for example:
+
+```bash
+kogine run measure/bytelm/train.py --set BUILD_CORPUS=true
+kogine workflow parallel measure/bytelm/train.py --config measure/bytelm/configs/a_stock.py --workers 2
+kogine run measure/plots.py
+```
+
+## Repository
+
+| path | contents |
+| --- | --- |
+| `src/dragonscale/` | the library: analysis, plans, patching, the grouped launch, the gradient decorator, unit-row modules |
+| `src/dragonscale_kernels/` | the row ops per backend: eager reference, `torch.compile`, Triton |
+| `src/dragonscale_common/` | helpers shared by both, without torch: logger, name registry, dotted-path import, problem lists |
+| `measure/` | experiments, op benchmarks and figures, one folder each, run with kogine |
+| `tests/` | patch / remove / bake / canonicalize, the grouped launch, and the analysis on hand-checked models |
+
+## License
+
+Apache-2.0.