diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..dcfe4ecf6c93447c9a580a0d19cc17ca8bf28042 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +paper/*.pdf filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore index 7a8b945e825fed28d5040f1a6d6533b7d715a7e7..e0055f576e602a27d84dc7116f45a3ac0c4702bb 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,12 @@ __pycache__/ *.py[cod] *.pdf +*.aux +*.blg +*.fdb_latexmk +*.fls +*.log +*.out *.npy .venv/ *.egg-info/ diff --git a/paper/dropout_decay_regime_specific_pressure_law.bbl b/paper/dropout_decay_regime_specific_pressure_law.bbl new file mode 100644 index 0000000000000000000000000000000000000000..ce2aea89a73846e285baca8107f6255753d3787a --- /dev/null +++ b/paper/dropout_decay_regime_specific_pressure_law.bbl @@ -0,0 +1,73 @@ +\begin{thebibliography}{8} +\providecommand{\natexlab}[1]{#1} +\providecommand{\url}[1]{\texttt{#1}} +\expandafter\ifx\csname urlstyle\endcsname\relax + \providecommand{\doi}[1]{doi: #1}\else + \providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi + +\bibitem[Eldan and Li(2023)]{eldan2023tinystories} +Ronen Eldan and Yuanzhi Li. +\newblock Tinystories: How small can language models be and still speak + coherent english? +\newblock \emph{arXiv preprint arXiv:2305.07759}, 2023. +\newblock URL \url{https://arxiv.org/abs/2305.07759}. + +\bibitem[Gokaslan et~al.(2019)Gokaslan, Cohen, Pavlick, and + Tellex]{gokaslan2019openwebtext} +Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex. +\newblock Openwebtext corpus. +\newblock \url{https://zenodo.org/records/3834942}, 2019. +\newblock Open reproduction of the WebText-style corpus. + +\bibitem[Hoffmann et~al.(2022)Hoffmann, Borgeaud, Mensch, Buchatskaya, Cai, + Rutherford, Casas, Hendricks, Welbl, Clark, Hennigan, Noland, Millican, + van~den Driessche, Damoc, Guy, Osindero, Simonyan, Elsen, Rae, Vinyals, and + Sifre]{hoffmann2022training} +Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor + Cai, Eliza Rutherford, Diego de~Las Casas, Lisa~Anne Hendricks, Johannes + Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van~den + Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich + Elsen, Jack~W. Rae, Oriol Vinyals, and Laurent Sifre. +\newblock Training compute-optimal large language models. +\newblock \emph{arXiv preprint arXiv:2203.15556}, 2022. +\newblock URL \url{https://arxiv.org/abs/2203.15556}. + +\bibitem[Kaplan et~al.(2020)Kaplan, McCandlish, Henighan, Brown, Chess, Child, + Gray, Radford, Wu, and Amodei]{kaplan2020scaling} +Jared Kaplan, Sam McCandlish, Tom Henighan, Tom~B. Brown, Benjamin Chess, Rewon + Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. +\newblock Scaling laws for neural language models. +\newblock \emph{arXiv preprint arXiv:2001.08361}, 2020. +\newblock URL \url{https://arxiv.org/abs/2001.08361}. + +\bibitem[Karpathy(2025)]{karpathy_nanochat} +Andrej Karpathy. +\newblock nanochat. +\newblock \url{https://github.com/karpathy/nanochat}, 2025. +\newblock MIT-licensed repository, accessed 2026-05-29. + +\bibitem[Merity et~al.(2017)Merity, Xiong, Bradbury, and + Socher]{merity2017pointer} +Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. +\newblock Pointer sentinel mixture models. +\newblock In \emph{International Conference on Learning Representations}, 2017. +\newblock URL \url{https://arxiv.org/abs/1609.07843}. + +\bibitem[Srivastava et~al.(2014)Srivastava, Hinton, Krizhevsky, Sutskever, and + Salakhutdinov]{srivastava2014dropout} +Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan + Salakhutdinov. +\newblock Dropout: A simple way to prevent neural networks from overfitting. +\newblock \emph{Journal of Machine Learning Research}, 15\penalty0 + (56):\penalty0 1929--1958, 2014. +\newblock URL \url{https://www.jmlr.org/papers/v15/srivastava14a.html}. + +\bibitem[Vaswani et~al.(2017)Vaswani, Shazeer, Parmar, Uszkoreit, Jones, Gomez, + Kaiser, and Polosukhin]{vaswani2017attention} +Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, + Aidan~N. Gomez, Lukasz Kaiser, and Illia Polosukhin. +\newblock Attention is all you need. +\newblock In \emph{Advances in Neural Information Processing Systems}, 2017. +\newblock URL \url{https://arxiv.org/abs/1706.03762}. + +\end{thebibliography} diff --git a/paper/dropout_decay_regime_specific_pressure_law.pdf b/paper/dropout_decay_regime_specific_pressure_law.pdf new file mode 100644 index 0000000000000000000000000000000000000000..17781ca4d54fff2804b7bf1dfc721a4fe37c36b5 --- /dev/null +++ b/paper/dropout_decay_regime_specific_pressure_law.pdf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:46ba08475ba2566bc3733eed5403db061a3ce4145ed6bd79965796719a396bc8 +size 267426 diff --git a/paper/dropout_decay_regime_specific_pressure_law.tex b/paper/dropout_decay_regime_specific_pressure_law.tex new file mode 100644 index 0000000000000000000000000000000000000000..1c63a35c44b084dffa9ebb41b94a50248eb0ae02 --- /dev/null +++ b/paper/dropout_decay_regime_specific_pressure_law.tex @@ -0,0 +1,427 @@ +\documentclass[11pt]{article} + +\usepackage[margin=1in]{geometry} +\usepackage{amsmath} +\usepackage{amssymb} +\usepackage{booktabs} +\usepackage{graphicx} +\usepackage{microtype} +\usepackage{tikz} +\usepackage{pgfplots} +\usepgfplotslibrary{groupplots} +\pgfplotsset{compat=1.18} +\usepackage[numbers,sort&compress]{natbib} +\usepackage[colorlinks=true,linkcolor=blue,citecolor=blue,urlcolor=blue]{hyperref} + +\title{Regime-Specific Pressure Laws for Dropout Decay in Streaming Language Model Training} +\author{Anonymous Authors} +\date{June 2026} + +\newcommand{\clamp}{\operatorname{clamp}} +\newcommand{\params}{P} +\newcommand{\unique}{U_t} +\newcommand{\sampled}{C_t} +\newcommand{\dropout}{p_t} + +\begin{document} +\maketitle + +\begin{abstract} +Dropout is usually selected as a fixed hyperparameter before training, but +streaming language-model training changes the regularization problem over time: +early stages expose few unique tokens to a comparatively large model, while +later stages reveal more data and reduce the need for strong stochastic +regularization. We study this setting with nanochat-derived causal +Transformers trained under an expanding-prefix protocol on three text regimes: +OpenWebText10K, TinyStories, and WikiText-103. Static dropout sweeps show a +repeatable pattern: for a fixed model, the validation-loss curve over dropout +is typically unimodal or plateau-shaped, and the optimum falls as unique data +increases; for a fixed prefix, the optimum rises with model capacity. We turn +these per-cell optima into a regime-specific pressure law, +\[ +\dropout=\clamp(p_{\min},p_{\max}, + A\log_{10}(\params/\unique) ++B\log_{10}(\sampled/\unique) ++D\log_{10}(\params/\unique)\log_{10}(\sampled/\unique)+C_0), +\] +where $\params$ is parameter count, $\unique$ is currently available unique +tokens, and $\sampled$ is cumulative sampled training tokens. The fitted +interaction law explains static optimum cells with mean absolute errors of +0.0158, 0.0180, and 0.0162 dropout on OpenWebText10K, TinyStories, and +WikiText-103 respectively. In locked streaming validation, fitted or +formula-derived decay schedules beat the best fixed-dropout baseline on final +validation loss in all three regimes at 4M-token validation scale and again in +larger 51.84M-parameter, 8M-token extrapolations. The large extrapolation tier +wins all paired seeds across all three regimes (9/9 paired wins), but stage-wise +analysis shows the schedules are final-loss optimized rather than best at every +intermediate prefix. The results support a practical hypothesis: dropout decay +should be calibrated as a regime-specific pressure law, not chosen as a +universal constant or an arbitrary high-to-low schedule. +\end{abstract} + +\section{Introduction} + +Dropout was introduced as a simple and effective regularizer for neural +networks \citep{srivastava2014dropout}. In language model practice, however, +dropout is often treated as a fixed scalar selected before a run. That choice +is mismatched to streaming settings. When a model begins training on a small +prefix of a stream, its parameter count can be large relative to the number of +unique tokens currently available; later, the same model receives more unique +data and can safely expose more functional capacity. A fixed dropout rate must +therefore compromise between early overfitting risk and late underfitting risk. + +The motivating hypothesis of this project was intentionally ambitious: a +high-capacity model could start with strong regularization on small initial +data and then decay dropout as streaming data accumulates, allowing capacity to +come online without catastrophic overfitting. The experiments refined that +hypothesis. Extremely high initial dropout is not automatically useful; early +tests with dropout near 0.8 underfit badly. The stronger and better-supported +claim is that useful dropout is a function of training pressure: model capacity +relative to the current data prefix, and cumulative optimizer sampling relative +to that prefix. + +This paper makes four empirical contributions. + +\begin{enumerate} + \item We define calibration cells for streaming dropout: one model, one + unique-token prefix, one sampled-token budget, and a sweep over fixed + dropout rates. + \item We show that calibration cells often have a parabolic or plateau-like + validation-loss curve over dropout, making the local optimum a usable + supervised target. + \item We fit a regime-specific interaction pressure law that maps + $(\params,\unique,\sampled)$ to a dropout rate. + \item We validate the resulting schedules against broad fixed-dropout + baselines across three text regimes and a larger-model/larger-stream + extrapolation tier. +\end{enumerate} + +The central result is not that one coefficient vector works everywhere. It does +not. The central result is that the pressure-law structure is useful, while the +coefficients are regime-specific. + +\section{Related Work} + +\paragraph{Dropout and regularization.} +Dropout applies multiplicative Bernoulli noise during training and can be +understood as an approximate ensemble over subnetworks +\citep{srivastava2014dropout}. With inverted dropout, an activation $h_i$ is +transformed as $h'_i=h_i m_i/(1-p)$, $m_i\sim\operatorname{Bernoulli}(1-p)$. +The expectation is preserved, but the variance grows as $h_i^2p/(1-p)$. This +captures the key tradeoff in our experiments: low dropout can overfit repeated +small prefixes, but high dropout can remove too much usable capacity. + +\paragraph{Transformer language models and scaling.} +Our model family is a decoder-only Transformer following the attention-based +architecture of \citet{vaswani2017attention}. Scaling-law work shows that loss +varies systematically with model size, data size, and compute +\citep{kaplan2020scaling,hoffmann2022training}. Our study is narrower: it does +not propose a global loss scaling law. Instead, it asks whether a +regularization hyperparameter can be mapped to local pressure variables during +streaming. + +\paragraph{Text regimes.} +We evaluate on three text regimes. TinyStories is a synthetic simple-English +story corpus designed for small language-model studies \citep{eldan2023tinystories}. +WikiText-103 is a standard word-level language-modeling benchmark derived from +Wikipedia articles \citep{merity2017pointer}. OpenWebText is an open +WebText-style corpus \citep{gokaslan2019openwebtext}. We use cached/tokenized +subsets and therefore report results as regime-level evidence rather than +claims about full-corpus state of the art. + +\section{Experimental Setting} + +\subsection{Model and implementation} + +All experiments use a standalone research codebase derived from Andrej +Karpathy's MIT-licensed nanochat repository \citep{karpathy_nanochat}. Chat +interfaces and deployment scripts are removed; the retained components are the +core BPE tokenizer and causal Transformer template. Derived source files retain +the original copyright and permission notice. + +The architecture family keeps the standard Transformer feed-forward expansion +ratio of four. The central model sizes are L8-H8-D256 (8.39M parameters), +L12-H8-D320 (17.37M), L16-H8-D384 (31.46M or 32.71M depending on +vocabulary/run metadata), and L20-H8-D448 (51.84M). The tokenizer vocabulary is +4096 in the reported runs. Unless noted otherwise, training uses AdamW, +learning rate $3\times 10^{-4}$, weight decay 0.1, gradient clipping 1.0, batch +size 16, and block size 128. All Torch training runs are MPS-only; the runner +refuses CUDA, CPU execution, and MPS fallback. + +\subsection{Streaming protocol} + +We simulate streaming by revealing larger token prefixes over time. A locked +streaming run keeps the same model and optimizer state across stages. At stage +$t$, minibatches are sampled only from the currently revealed prefix +$\unique$. The two main protocols are: + +\begin{itemize} + \item 4M validation: OpenWebText10K and WikiText-103 use 250k, 500k, 1M, + 2M, and 4M prefixes. TinyStories uses 500k, 1M, 2M, and 4M prefixes. These + runs use five random seeds. + \item 8M extrapolation: L20-H8-D448 with an 8M final prefix. These runs use + three seeds because they are substantially more expensive. +\end{itemize} + +Each run compares a dropout schedule against a grid of static-dropout +baselines. The best static baseline is selected after training from the same +run family, so the decay schedule competes against a strong fixed-rate oracle +within the tested grid. + +\section{From Calibration Cells to a Pressure Law} + +\subsection{Calibration cell definition} + +A calibration cell is one static dropout sweep at a fixed tuple: +\[ +(\text{model architecture}, \params, \unique, \sampled, \text{validation split}). +\] +Within that cell, we train/evaluate several fixed dropout values and obtain a +curve +\[ +p \mapsto L_{\mathrm{val}}(p \mid \params,\unique,\sampled). +\] +The target dropout $p^\star$ is the best observed grid point or, when the best +point is bracketed by worse neighbors, a local quadratic interpolation of the +minimum. Boundary optima are retained but downweighted in coefficient fitting +because the true optimum may lie outside the tested grid. + +Figure~\ref{fig:cell-curves} shows representative cells. The curves are not +perfect mathematical parabolas, but they show the practical pattern needed for +calibration: validation loss improves as dropout moves away from zero, reaches +a local minimum or plateau, and then rises when dropout becomes too large. + +\begin{figure}[t] +\centering +\input{figures/calibration_cell_curves} +\caption{Representative static dropout calibration cells. The y-axis is +validation loss above the best tested dropout in that cell, so each curve is +normalized to show shape rather than absolute corpus difficulty.} +\label{fig:cell-curves} +\end{figure} + +\subsection{Pressure variables} + +The formula uses two dimensionless pressure variables: +\[ +x_t=\log_{10}(\params/\unique), \qquad +y_t=\log_{10}(\sampled/\unique). +\] +The first term measures capacity pressure. Increasing parameters at fixed +prefix increases memorization capacity and tends to increase useful dropout. +Increasing unique tokens at fixed model reduces this pressure. The second term +measures repeated-sampling pressure. If the optimizer consumes many sampled +tokens while the unique prefix remains small, it repeatedly revisits the same +data and can benefit from stronger regularization. + +The fitted interaction family is +\begin{equation} +\dropout=\clamp(p_{\min},p_{\max}, + A x_t+B y_t+D x_t y_t+C_0). +\label{eq:interaction-law} +\end{equation} +For schedule generation we use $p_{\min}=0.02$ and $p_{\max}=0.65$ unless a +run explicitly specifies a different clamp. Static sweeps may still include +0.0 as an ablation. + +\subsection{Regime-specific coefficients} + +We define a regime as the combination of architecture family, tokenizer, corpus +family, optimizer protocol, dropout placement, stream protocol, and evaluation +distribution. Model size and prefix size are not regime changes; they are +inputs to the formula. Corpus family, tokenizer, or optimizer changes can shift +the coefficients. + +Table~\ref{tab:coefficients} gives the fitted interaction coefficients. Each +regime uses 15 or 16 calibration cells. The fitted mean absolute error is about +0.016--0.018 dropout, meaning the coefficient law typically predicts the +calibrated static optimum to within roughly two hundredths of dropout. + +\begin{table}[t] +\centering +\caption{Regime-specific interaction coefficient fits. L-model and L-prefix +are leave-model-out and leave-prefix-out mean absolute errors.} +\label{tab:coefficients} +\resizebox{\textwidth}{!}{\input{tables/coefficient_table}} +\end{table} + +Figure~\ref{fig:static-optima} shows the fitted target dropout values from the +calibration cells. Across regimes, useful dropout generally decreases as the +unique prefix grows. At a fixed prefix, larger models tend to require larger +dropout. The absolute surface differs by corpus regime, which is why the paper +does not propose universal coefficients. + +\begin{figure}[t] +\centering +\input{figures/static_optima_by_regime} +\caption{Calibration targets by regime. At fixed model size, target dropout +tends to fall as more unique tokens are available; at fixed prefix, larger +models tend to have higher target dropout. WikiText-103 reaches the dropout +floor earlier than the other regimes under its calibration protocol.} +\label{fig:static-optima} +\end{figure} + +\section{Streaming Results} + +\subsection{Main validation} + +Table~\ref{tab:streaming-validation} summarizes the main streaming results. +The ``Gain'' column is best static final validation loss minus decay final +validation loss, so positive values favor decay. The best static baseline is +chosen from the tested fixed-dropout grid in the same run family. + +\begin{table}[t] +\centering +\caption{Streaming validation against best fixed-dropout baselines.} +\label{tab:streaming-validation} +\resizebox{\textwidth}{!}{\input{tables/streaming_validation_table}} +\end{table} + +OpenWebText10K provides the cleanest five-seed coefficient-derived validation: +the interaction schedule improves final validation loss by 0.0473 over the +best static baseline and wins all five paired seeds. TinyStories shows a +smaller but consistent final-loss margin: 0.0133 over the best static baseline +with five paired wins. WikiText-103 is more nuanced. The earlier +formula-derived L12 schedule wins the five-seed 4M run, but the later +static-cell-fitted interaction schedule decays too quickly for L12 and is only +near the best static baseline. This negative result is important: a good static +cell fit does not by itself guarantee the best streaming schedule unless the +schedule objective is aligned with final-loss optimization. + +\subsection{L20 8M extrapolation} + +The strongest comparable evidence is the L20 extrapolation tier. Here the +model is fixed to L20-H8-D448 with 51.84M parameters, and each regime uses +its fitted interaction coefficients to produce a frozen schedule before +training. No coefficients are refit on the L20 results. Across OpenWebText10K, +TinyStories, and WikiText-103, the decay schedule beats the best static +baseline on final validation loss and wins all paired seeds (3/3 in each +regime, 9/9 total). + +Figure~\ref{fig:final-improvements} visualizes the final-loss gains. The +largest L20 gain is OpenWebText10K (0.0363), followed by WikiText-103 (0.0184) +and TinyStories (0.0100). The TinyStories gain is smaller because a static +dropout near 0.20 is already a strong fixed compromise for that regime. + +\begin{figure}[t] +\centering +\input{figures/final_improvements} +\caption{Final validation-loss improvement over the best fixed-dropout +baseline. Positive values favor the decay schedule.} +\label{fig:final-improvements} +\end{figure} + +\subsection{Stage-wise behavior} + +Final-loss improvement does not imply per-prefix dominance. Table~\ref{tab:l20-stage} +and Figure~\ref{fig:l20-stage-deltas} compare the L20 schedule against the best +static baseline at each stage. OpenWebText10K starts worse than static 0.30 at +250k but becomes better from 1M onward. TinyStories is nearly tied early, +slightly worse at 1M/2M, then better at 4M/8M. WikiText-103 is worse through +the middle of the stream and wins only at the final 8M prefix. + +\begin{table}[t] +\centering +\caption{L20 stage-wise comparison. Delta is decay validation loss minus the +best static validation loss at the same prefix; negative favors decay.} +\label{tab:l20-stage} +\resizebox{\textwidth}{!}{\input{tables/l20_stage_table}} +\end{table} + +\begin{figure}[t] +\centering +\input{figures/l20_stage_deltas} +\caption{L20 stage-wise deltas versus the best static dropout at each prefix. +The schedules are final-loss schedules, not greedy per-prefix optima.} +\label{fig:l20-stage-deltas} +\end{figure} + +\section{Patterns and Interpretation} + +\paragraph{Moderate high dropout, not massive high dropout.} +The earliest experiments tested very large dropout and failed. A static 0.8 +dropout stream underfit badly, and a high-dropout decay also lost to moderate +static dropout. This falsified the naive version of the hypothesis. The +successful schedules start high only when the pressure law predicts high +dropout, and high means roughly 0.25--0.45 in the L20 extrapolations, not 0.8. + +\paragraph{Dropout optima move with pressure.} +The calibration cells support the core mechanism. For the same model, target +dropout declines as the unique prefix expands. For the same prefix, target +dropout generally increases with parameter count. This is consistent with the +view that dropout is compensating for capacity pressure and repeated sampling. + +\paragraph{The interaction term matters.} +The interaction term $xy$ captures the fact that model/data pressure and +sampled-token pressure are not independent. For example, repeated sampling is +more dangerous when a large model sees a small prefix than when a smaller model +sees the same prefix. In WikiText-103, the interaction form reduces calibration +MAE from 0.0449 for the first-order base formula to 0.0162. Similar interaction +fits are used for OpenWebText10K and TinyStories. + +\paragraph{Regime specificity is real.} +The coefficients differ substantially across regimes. TinyStories needs +different signs and offsets than OpenWebText10K, and WikiText-103 reaches low +dropout optima earlier under its calibration protocol. This is not a weakness +of the pressure-law structure; it is evidence that the corpus/tokenizer/eval +distribution defines the coefficient regime. + +\paragraph{Final-loss schedules can sacrifice early local optimality.} +The L20 stage-wise analysis shows that a schedule can lose to a fixed static +dropout at early or middle prefixes and still win final loss. This is expected +if the goal is to train one model continuously through a growing stream, not to +restart and optimize each prefix independently. Static baselines can be locally +optimal at one prefix but cannot adapt when the stream changes. + +\section{Limitations} + +The experiments are intentionally small compared with frontier language-model +training. The largest model has 51.84M parameters and the largest stream prefix +has 8M tokens. All training uses Apple MPS, short sequence length, and a 4096 +vocabulary. We therefore claim a reproducible small-scale phenomenon, not a +confirmed law at billion-parameter scale. + +The static coefficient fit optimizes dropout targets extracted from static +sweeps, while the streaming objective is final validation loss after continued +training. WikiText-103 shows that these objectives can diverge: the fitted +interaction schedule is excellent as a static optimum model but decays too +quickly for L12 streaming. Future work should fit coefficients directly against +streaming objective functions or include trajectory-aware loss terms. + +The L20 extrapolation tier uses three seeds for cost reasons. This is balanced +across all three regimes and supported by five-seed smaller-stream validation, +but larger seed counts would narrow confidence intervals. We also do not claim +that the tested fixed-dropout grids contain every possible fixed optimum. + +\section{Reproducibility} + +The repository is standalone: cached corpora/tokenizers, saved metrics, +coefficient files, scripts, and reports are stored in the project. The main +training entry point is \texttt{scripts/run\_experiments.py}; coefficient +fitting uses \texttt{scripts/fit\_dropout\_coefficients.py}; this manuscript's +figures and tables are generated by +\texttt{scripts/generate\_paper\_assets.py}. Existing result artifacts are read +from \texttt{runs/}. The experiment runner enforces MPS-only Torch execution. + +All generated tables in this paper are derived from saved metrics JSONL files +and coefficient JSON files. The paper-generation summary is written to +\texttt{paper/paper\_results\_summary.json}. + +\section{Conclusion} + +The evidence supports a practical, falsifiable hypothesis: in streaming +language-model training, dropout should decay according to a regime-specific +pressure law over model size, currently available unique tokens, and cumulative +sampled tokens. The hypothesis is stronger than ``larger dropout always helps'' +and weaker than ``one universal formula works everywhere.'' Static calibration +cells reveal local dropout optima; interaction pressure laws fit those optima +with low error; and the resulting schedules improve final validation loss over +the best tested fixed-dropout baselines across three regimes, including a +larger-model 8M-token extrapolation tier. The next scientific step is to test +the same fitting procedure at larger scale and to optimize coefficients against +streaming final-loss objectives directly. + +\bibliographystyle{plainnat} +\bibliography{references} + +\end{document} diff --git a/paper/figures/calibration_cell_curves.tex b/paper/figures/calibration_cell_curves.tex new file mode 100644 index 0000000000000000000000000000000000000000..1134c00c54dd46d743965cd78cb4a8b11bf15c87 --- /dev/null +++ b/paper/figures/calibration_cell_curves.tex @@ -0,0 +1,51 @@ +\begin{tikzpicture} +\begin{groupplot}[ + group style={group size=3 by 1, horizontal sep=1.0cm}, + width=0.32\textwidth, + height=0.26\textwidth, + legend style={font=\scriptsize}, + tick label style={font=\scriptsize}, + label style={font=\scriptsize}, + title style={font=\scriptsize}, +] + +\nextgroupplot[ + title={OpenWebText10K, L12, 1M}, + xlabel={Dropout}, + ylabel={Validation loss above best}, + ymin=0, + grid=both, +] +\addplot+[thick, mark=*] coordinates {(0,0.64089) (0.02,0.41177) (0.05,0.21561) (0.08,0.10277) (0.1,0.06467) (0.14,0.02993) (0.2,0.00000) (0.3,0.03437) (0.4,0.09347) (0.5,0.21260) (0.6,0.37234) (0.7,0.57061) (0.8,0.88258) (0.9,2.06141)}; +\addlegendentry{static sweep} +\addplot+[only marks, mark=star, mark size=3pt] coordinates {(0.2,0)}; +\addlegendentry{target $p^\star=0.20$} + + +\nextgroupplot[ + title={TinyStories, L12, 1M}, + xlabel={Dropout}, + ylabel={Validation loss above best}, + ymin=0, + grid=both, +] +\addplot+[thick, mark=*] coordinates {(0,0.40226) (0.04,0.13035) (0.1,0.02502) (0.2,0.00000) (0.35,0.09066) (0.5,0.22441)}; +\addlegendentry{static sweep} +\addplot+[only marks, mark=star, mark size=3pt] coordinates {(0.2,0)}; +\addlegendentry{target $p^\star=0.20$} + + +\nextgroupplot[ + title={WikiText-103, L12, 250k}, + xlabel={Dropout}, + ylabel={Validation loss above best}, + ymin=0, + grid=both, +] +\addplot+[thick, mark=*] coordinates {(0,0.17165) (0.02,0.10888) (0.04,0.07189) (0.06,0.04641) (0.08,0.02649) (0.1,0.00897) (0.14,0.00363) (0.18,0.00000) (0.2,0.01021) (0.26,0.03365) (0.3,0.04794) (0.4,0.13084)}; +\addlegendentry{static sweep} +\addplot+[only marks, mark=star, mark size=3pt] coordinates {(0.18,0)}; +\addlegendentry{target $p^\star=0.18$} + +\end{groupplot} +\end{tikzpicture} diff --git a/paper/figures/final_improvements.tex b/paper/figures/final_improvements.tex new file mode 100644 index 0000000000000000000000000000000000000000..36fd6df7355400d5ab74d4c9713c6152edfdcfee --- /dev/null +++ b/paper/figures/final_improvements.tex @@ -0,0 +1,23 @@ + +\begin{tikzpicture} +\begin{axis}[ + ybar, + bar width=8pt, + width=0.72\textwidth, + height=0.34\textwidth, + ymin=0, + ylabel={Best static final loss $-$ decay final loss}, + symbolic x coords={OWT10K,TinyStories,WikiText}, + xtick=data, + xtick={OWT10K,TinyStories,WikiText}, + grid=both, + legend style={font=\scriptsize, at={(0.5,1.04)}, anchor=south, legend columns=2}, + tick label style={font=\scriptsize}, + label style={font=\scriptsize}, +] +\addplot coordinates {(OWT10K,0.04732) (TinyStories,0.01326) (WikiText,0.02966)}; +\addlegendentry{4M validation} +\addplot coordinates {(OWT10K,0.03631) (TinyStories,0.01000) (WikiText,0.01838)}; +\addlegendentry{L20 8M extrapolation} +\end{axis} +\end{tikzpicture} diff --git a/paper/figures/l20_stage_deltas.tex b/paper/figures/l20_stage_deltas.tex new file mode 100644 index 0000000000000000000000000000000000000000..835cecda3cc05e63c8d53d1a7dcc259726b9f857 --- /dev/null +++ b/paper/figures/l20_stage_deltas.tex @@ -0,0 +1,47 @@ +\begin{tikzpicture} +\begin{groupplot}[ + group style={group size=3 by 1, horizontal sep=1.0cm}, + width=0.32\textwidth, + height=0.27\textwidth, + tick label style={font=\scriptsize}, + label style={font=\scriptsize}, + title style={font=\scriptsize}, +] + +\nextgroupplot[ + title={OpenWebText10K}, + xlabel={Prefix}, + ylabel={Decay $-$ best static}, + xtick={0,1,2,3,4,5}, + xticklabels={250k,500k,1M,2M,4M,8M}, + grid=both, +] +\addplot+[thick, mark=*] coordinates {(0,0.06300) (1,0.00103) (2,-0.01477) (3,-0.02370) (4,-0.02720) (5,-0.03631)}; +\addplot+[black, dashed, mark=none] coordinates {(0,0.00000) (5,0.00000)}; + + +\nextgroupplot[ + title={TinyStories}, + xlabel={Prefix}, + ylabel={Decay $-$ best static}, + xtick={0,1,2,3,4}, + xticklabels={500k,1M,2M,4M,8M}, + grid=both, +] +\addplot+[thick, mark=*] coordinates {(0,0.00050) (1,0.00729) (2,0.00347) (3,-0.00570) (4,-0.01000)}; +\addplot+[black, dashed, mark=none] coordinates {(0,0.00000) (4,0.00000)}; + + +\nextgroupplot[ + title={WikiText-103}, + xlabel={Prefix}, + ylabel={Decay $-$ best static}, + xtick={0,1,2,3,4,5}, + xticklabels={250k,500k,1M,2M,4M,8M}, + grid=both, +] +\addplot+[thick, mark=*] coordinates {(0,0.00803) (1,0.01591) (2,0.03177) (3,0.03542) (4,0.00371) (5,-0.01838)}; +\addplot+[black, dashed, mark=none] coordinates {(0,0.00000) (5,0.00000)}; + +\end{groupplot} +\end{tikzpicture} diff --git a/paper/figures/static_optima_by_regime.tex b/paper/figures/static_optima_by_regime.tex new file mode 100644 index 0000000000000000000000000000000000000000..436e21dc907f0ddc09232a17ddd16241c239b4e7 --- /dev/null +++ b/paper/figures/static_optima_by_regime.tex @@ -0,0 +1,66 @@ +\begin{tikzpicture} +\begin{groupplot}[ + group style={group size=3 by 1, horizontal sep=1.0cm}, + width=0.32\textwidth, + height=0.28\textwidth, + legend style={font=\tiny}, + tick label style={font=\scriptsize}, + label style={font=\scriptsize}, + title style={font=\scriptsize}, +] + +\nextgroupplot[ + title={OpenWebText10K}, + xlabel={Unique prefix $U$}, + ylabel={Fitted target dropout}, + xtick={0,1,2,3,4}, + xticklabels={250k,500k,1M,2M,4M}, + ymin=0, + ymax=0.65, + grid=both, +] +\addplot+[thick, mark=*] coordinates {(0,0.5062) (1,0.3559) (2,0.2174) (3,0.1364) (4,0.0224)}; +\addlegendentry{L12\_H8\_D320} +\addplot+[thick, mark=*] coordinates {(0,0.6039) (1,0.4475) (2,0.2837) (3,0.1388) (4,0.0272)}; +\addlegendentry{L16\_H8\_D384} +\addplot+[thick, mark=*] coordinates {(0,0.3837) (1,0.2401) (2,0.1463) (3,0.0691) (4,0.0000)}; +\addlegendentry{L8\_H8\_D256} + + +\nextgroupplot[ + title={TinyStories}, + xlabel={Unique prefix $U$}, + ylabel={Fitted target dropout}, + xtick={0,1,2,3,4}, + xticklabels={250k,500k,1M,2M,4M}, + ymin=0, + ymax=0.65, + grid=both, +] +\addplot+[thick, mark=*] coordinates {(0,0.4575) (1,0.2947) (2,0.1866) (3,0.0703) (4,0.0303)}; +\addlegendentry{L12\_H8\_D320} +\addplot+[thick, mark=*] coordinates {(0,0.5699) (2,0.2450) (4,0.0468)}; +\addlegendentry{L16\_H8\_D384} +\addplot+[thick, mark=*] coordinates {(0,0.3491) (2,0.0977) (4,0.0000)}; +\addlegendentry{L8\_H8\_D256} + + +\nextgroupplot[ + title={WikiText-103}, + xlabel={Unique prefix $U$}, + ylabel={Fitted target dropout}, + xtick={0,1,2,3,4}, + xticklabels={250k,500k,1M,2M,4M}, + ymin=0, + ymax=0.65, + grid=both, +] +\addplot+[thick, mark=*] coordinates {(0,0.1645) (1,0.0299) (2,0.0000) (3,0.0000) (4,0.0000)}; +\addlegendentry{L12\_H8\_D320} +\addplot+[thick, mark=*] coordinates {(0,0.2459) (1,0.0822) (2,0.0225) (3,0.0000) (4,0.0000)}; +\addlegendentry{L16\_H8\_D384} +\addplot+[thick, mark=*] coordinates {(0,0.0965) (1,0.0114) (2,0.0000) (3,0.0000) (4,0.0000)}; +\addlegendentry{L8\_H8\_D256} + +\end{groupplot} +\end{tikzpicture} diff --git a/paper/paper_results_summary.json b/paper/paper_results_summary.json new file mode 100644 index 0000000000000000000000000000000000000000..b497e6e09c8dda1e1392286c65fe9baacf406c16 --- /dev/null +++ b/paper/paper_results_summary.json @@ -0,0 +1,520 @@ +{ + "openweb": { + "coefficients": { + "regime": "OpenWebText10K", + "cells": 15, + "A": 0.08189560629877818, + "B": -0.023948107975524876, + "D": 0.1536984646521253, + "C0": -0.05559686586455526, + "mae": 0.01579101988615387, + "rmse": 0.018845563252758567, + "leave_model_mae": 0.024771721924833498, + "leave_prefix_mae": 0.030983393290384315, + "formula": "p = +0.081896 * log10(P / U) -0.023948 * log10(C / U) +0.153698 * log10(P / U) * log10(C / U) -0.055597" + }, + "small_stream": { + "n": 5, + "final_stage": 4, + "final_prefix": 4000000, + "condition_final_mean": 4.3981304407119755, + "condition_final_std": 0.009545784836147743, + "best_static_condition": "static_dropout_0.14", + "best_static_mean": 4.44545366615057, + "best_static_std": 0.012017216742245517, + "improvement": 0.04732322543859446, + "paired_deltas": [ + -0.03942108899354935, + -0.05825117975473404, + -0.032753005623817444, + -0.052590467035770416, + -0.053600385785102844 + ], + "paired_wins": 5, + "paired_mean_delta": -0.04732322543859482, + "paired_std_delta": 0.010738969078098962, + "stage_deltas": [ + { + "stage": 0, + "prefix": 250000, + "condition_dropout": 0.385, + "condition_mean": 5.4946602284908295, + "condition_std": 0.01093302726132647, + "best_static_condition": "static_dropout_0.3", + "best_static_mean": 5.448330155014991, + "best_static_std": 0.0138285316736341, + "delta": 0.04633007347583806 + }, + { + "stage": 1, + "prefix": 500000, + "condition_dropout": 0.319, + "condition_mean": 5.071460470557213, + "condition_std": 0.01179360076463939, + "best_static_condition": "static_dropout_0.3", + "best_static_mean": 5.066737298667431, + "best_static_std": 0.017273470277214743, + "delta": 0.004723171889781774 + }, + { + "stage": 2, + "prefix": 1000000, + "condition_dropout": 0.227, + "condition_mean": 4.781069016456604, + "condition_std": 0.008428247627355752, + "best_static_condition": "static_dropout_0.3", + "best_static_mean": 4.79825523942709, + "best_static_std": 0.01441949497608529, + "delta": -0.01718622297048622 + }, + { + "stage": 3, + "prefix": 2000000, + "condition_dropout": 0.139, + "condition_mean": 4.558990895748138, + "condition_std": 0.014177173687439483, + "best_static_condition": "static_dropout_0.3", + "best_static_mean": 4.603498187661171, + "best_static_std": 0.014129740963263297, + "delta": -0.04450729191303271 + }, + { + "stage": 4, + "prefix": 4000000, + "condition_dropout": 0.066, + "condition_mean": 4.3981304407119755, + "condition_std": 0.009545784836147743, + "best_static_condition": "static_dropout_0.14", + "best_static_mean": 4.44545366615057, + "best_static_std": 0.012017216742245517, + "delta": -0.04732322543859446 + } + ] + }, + "l20_stream": { + "n": 3, + "final_stage": 5, + "final_prefix": 8000000, + "condition_final_mean": 4.179318442940712, + "condition_final_std": 0.01344880214961143, + "best_static_condition": "static_dropout_0.2", + "best_static_mean": 4.215624410659075, + "best_static_std": 0.014242200766109402, + "improvement": 0.03630596771836281, + "paired_deltas": [ + -0.03935118764638901, + -0.025513865053653717, + -0.03717709332704544 + ], + "paired_wins": 3, + "paired_mean_delta": -0.03401404867569605, + "paired_std_delta": 0.007441203715372809, + "stage_deltas": [ + { + "stage": 0, + "prefix": 250000, + "condition_dropout": 0.433, + "condition_mean": 5.460076779127121, + "condition_std": 0.0012588964293079841, + "best_static_condition": "static_dropout_0.3", + "best_static_mean": 5.39707263559103, + "best_static_std": 0.014451385161445679, + "delta": 0.06300414353609085 + }, + { + "stage": 1, + "prefix": 500000, + "condition_dropout": 0.367, + "condition_mean": 5.017915904521942, + "condition_std": 0.005764488447180914, + "best_static_condition": "static_dropout_0.3", + "best_static_mean": 5.01688560595115, + "best_static_std": 0.012061853156578773, + "delta": 0.0010302985707921763 + }, + { + "stage": 2, + "prefix": 1000000, + "condition_dropout": 0.27, + "condition_mean": 4.7070309941967325, + "condition_std": 0.015739562982191573, + "best_static_condition": "static_dropout_0.3", + "best_static_mean": 4.72180004666249, + "best_static_std": 0.02152489745362506, + "delta": -0.014769052465757326 + }, + { + "stage": 3, + "prefix": 2000000, + "condition_dropout": 0.17600000000000002, + "condition_mean": 4.487937837839127, + "condition_std": 0.009026555295310244, + "best_static_condition": "static_dropout_0.2", + "best_static_mean": 4.511637444297473, + "best_static_std": 0.016469397224544248, + "delta": -0.023699606458346345 + }, + { + "stage": 4, + "prefix": 4000000, + "condition_dropout": 0.09500000000000001, + "condition_mean": 4.311577834188938, + "condition_std": 0.015007670903308238, + "best_static_condition": "static_dropout_0.2", + "best_static_mean": 4.338780450324218, + "best_static_std": 0.01698295581756926, + "delta": -0.027202616135279634 + }, + { + "stage": 5, + "prefix": 8000000, + "condition_dropout": 0.033, + "condition_mean": 4.179318442940712, + "condition_std": 0.01344880214961143, + "best_static_condition": "static_dropout_0.2", + "best_static_mean": 4.215624410659075, + "best_static_std": 0.014242200766109402, + "delta": -0.03630596771836281 + } + ] + } + }, + "tinystories": { + "coefficients": { + "regime": "TinyStories", + "cells": 16, + "A": -0.08926115772929914, + "B": -0.1297539607354037, + "D": 0.2550686141985067, + "C0": 0.08152490344304414, + "mae": 0.01799638717496229, + "rmse": 0.023120178816187012, + "leave_model_mae": 0.029523601005481504, + "leave_prefix_mae": 0.03253257579449062, + "formula": "p = -0.089261 * log10(P / U) -0.129754 * log10(C / U) +0.255069 * log10(P / U) * log10(C / U) +0.081525" + }, + "small_stream": { + "n": 5, + "final_stage": 3, + "final_prefix": 4000000, + "condition_final_mean": 2.5311126589775084, + "condition_final_std": 0.021347338722582713, + "best_static_condition": "static_dropout_0.08", + "best_static_mean": 2.5443769969046115, + "best_static_std": 0.021138839821398296, + "improvement": 0.013264337927103131, + "paired_deltas": [ + -0.0005062669515609741, + -0.021110694855451584, + -0.009244173765182495, + -0.016586389392614365, + -0.014097623527050018 + ], + "paired_wins": 5, + "paired_mean_delta": -0.012309029698371887, + "paired_std_delta": 0.00786865194496459, + "stage_deltas": [ + { + "stage": 0, + "prefix": 500000, + "condition_dropout": 0.184, + "condition_mean": 3.232573334872723, + "condition_std": 0.012322249630223228, + "best_static_condition": "static_dropout_0.12", + "best_static_mean": 3.2225750528275965, + "best_static_std": 0.014316453554089874, + "delta": 0.009998282045126317 + }, + { + "stage": 1, + "prefix": 1000000, + "condition_dropout": 0.141, + "condition_mean": 2.8907706700265408, + "condition_std": 0.00267614743609897, + "best_static_condition": "static_dropout_0.12", + "best_static_mean": 2.893001724779606, + "best_static_std": 0.012108793042410648, + "delta": -0.002231054753065287 + }, + { + "stage": 2, + "prefix": 2000000, + "condition_dropout": 0.084, + "condition_mean": 2.6690264880657195, + "condition_std": 0.020694473868817478, + "best_static_condition": "static_dropout_0.12", + "best_static_mean": 2.6794611223042013, + "best_static_std": 0.016337678176899628, + "delta": -0.010434634238481788 + }, + { + "stage": 3, + "prefix": 4000000, + "condition_dropout": 0.045, + "condition_mean": 2.5311126589775084, + "condition_std": 0.021347338722582713, + "best_static_condition": "static_dropout_0.08", + "best_static_mean": 2.5443769969046115, + "best_static_std": 0.021138839821398296, + "delta": -0.013264337927103131 + } + ] + }, + "l20_stream": { + "n": 3, + "final_stage": 4, + "final_prefix": 8000000, + "condition_final_mean": 2.208572393283248, + "condition_final_std": 0.004634657584019386, + "best_static_condition": "static_dropout_0.2", + "best_static_mean": 2.218574500953158, + "best_static_std": 0.003241222983607362, + "improvement": 0.010002107669909943, + "paired_deltas": [ + -0.00869341753423214, + -0.006095590069890022, + -0.013002954423427582 + ], + "paired_wins": 3, + "paired_mean_delta": -0.00926398734251658, + "paired_std_delta": 0.0034888512447114515, + "stage_deltas": [ + { + "stage": 0, + "prefix": 500000, + "condition_dropout": 0.253, + "condition_mean": 3.030243801573912, + "condition_std": 0.006391037258560546, + "best_static_condition": "static_dropout_0.2", + "best_static_mean": 3.02974146977067, + "best_static_std": 0.006865445597812618, + "delta": 0.0005023318032422175 + }, + { + "stage": 1, + "prefix": 1000000, + "condition_dropout": 0.209, + "condition_mean": 2.7195058229068914, + "condition_std": 0.01655123058742941, + "best_static_condition": "static_dropout_0.26", + "best_static_mean": 2.712211973965168, + "best_static_std": 0.01060853933383755, + "delta": 0.007293848941723358 + }, + { + "stage": 2, + "prefix": 2000000, + "condition_dropout": 0.137, + "condition_mean": 2.476810533553362, + "condition_std": 0.006082638165661795, + "best_static_condition": "static_dropout_0.2", + "best_static_mean": 2.4733379955093064, + "best_static_std": 0.009827865290864221, + "delta": 0.003472538044055451 + }, + { + "stage": 3, + "prefix": 4000000, + "condition_dropout": 0.077, + "condition_mean": 2.3110335220893226, + "condition_std": 0.015477651713059048, + "best_static_condition": "static_dropout_0.2", + "best_static_mean": 2.3167371824383736, + "best_static_std": 0.008877543510132386, + "delta": -0.00570366034905101 + }, + { + "stage": 4, + "prefix": 8000000, + "condition_dropout": 0.041, + "condition_mean": 2.208572393283248, + "condition_std": 0.004634657584019386, + "best_static_condition": "static_dropout_0.2", + "best_static_mean": 2.218574500953158, + "best_static_std": 0.003241222983607362, + "delta": -0.010002107669909943 + } + ] + } + }, + "wikitext": { + "coefficients": { + "regime": "WikiText-103", + "cells": 15, + "A": -0.002674916387321836, + "B": -0.2836210480242947, + "D": 0.2686063066585002, + "C0": -0.01899454503425616, + "mae": 0.016244092648220347, + "rmse": 0.019131667586788335, + "leave_model_mae": 0.030438322968748356, + "leave_prefix_mae": 0.04435212839791333, + "formula": "p = -0.002675 * log10(P / U) -0.283621 * log10(C / U) +0.268606 * log10(P / U) * log10(C / U) -0.018995" + }, + "small_stream": { + "n": 5, + "final_stage": 4, + "final_prefix": 4000000, + "condition_final_mean": 4.080797865241766, + "condition_final_std": 0.019458062276034267, + "best_static_condition": "static_dropout_0.1", + "best_static_mean": 4.1104619488120075, + "best_static_std": 0.018806640195540816, + "improvement": 0.02966408357024175, + "paired_deltas": [ + -0.018418047577142715, + -0.018113188445568085, + -0.027215413749217987, + -0.02507050335407257, + -0.037036627531051636 + ], + "paired_wins": 5, + "paired_mean_delta": -0.0251707561314106, + "paired_std_delta": 0.007752385341386418, + "stage_deltas": [ + { + "stage": 0, + "prefix": 250000, + "condition_dropout": 0.3, + "condition_mean": 5.21483291387558, + "condition_std": 0.018060268226475495, + "best_static_condition": "static_dropout_0.18", + "best_static_mean": 5.16162400841713, + "best_static_std": 0.014987630048664583, + "delta": 0.05320890545845014 + }, + { + "stage": 1, + "prefix": 500000, + "condition_dropout": 0.26, + "condition_mean": 4.80534844994545, + "condition_std": 0.027829930328897787, + "best_static_condition": "static_dropout_0.2", + "best_static_mean": 4.787301687896251, + "best_static_std": 0.023601504941398347, + "delta": 0.01804676204919886 + }, + { + "stage": 2, + "prefix": 1000000, + "condition_dropout": 0.18, + "condition_mean": 4.49377838075161, + "condition_std": 0.014674108647494532, + "best_static_condition": "static_dropout_0.14", + "best_static_mean": 4.500055414438248, + "best_static_std": 0.014819407090300617, + "delta": -0.006277033686638234 + }, + { + "stage": 3, + "prefix": 2000000, + "condition_dropout": 0.09, + "condition_mean": 4.260725708305836, + "condition_std": 0.018098739272622868, + "best_static_condition": "static_dropout_0.14", + "best_static_mean": 4.281563547253609, + "best_static_std": 0.02138858846966249, + "delta": -0.020837838947772802 + }, + { + "stage": 4, + "prefix": 4000000, + "condition_dropout": 0.02, + "condition_mean": 4.080797865241766, + "condition_std": 0.019458062276034267, + "best_static_condition": "static_dropout_0.1", + "best_static_mean": 4.1104619488120075, + "best_static_std": 0.018806640195540816, + "delta": -0.02966408357024175 + } + ] + }, + "l20_stream": { + "n": 3, + "final_stage": 5, + "final_prefix": 8000000, + "condition_final_mean": 3.957240684578816, + "condition_final_std": 0.008487626829204101, + "best_static_condition": "static_dropout_0.14", + "best_static_mean": 3.9756230898201466, + "best_static_std": 0.008723729300606414, + "improvement": 0.018382405241330613, + "paired_deltas": [ + -0.016518089920282364, + -0.00015844032168388367, + -0.027258697897195816 + ], + "paired_wins": 3, + "paired_mean_delta": -0.014645076046387354, + "paired_std_delta": 0.013646872385042126, + "stage_deltas": [ + { + "stage": 0, + "prefix": 250000, + "condition_dropout": 0.284, + "condition_mean": 5.2075243567427, + "condition_std": 0.013993230772716539, + "best_static_condition": "static_dropout_0.26", + "best_static_mean": 5.1994956483443575, + "best_static_std": 0.013634536236337126, + "delta": 0.008028708398342133 + }, + { + "stage": 1, + "prefix": 500000, + "condition_dropout": 0.211, + "condition_mean": 4.816951217750709, + "condition_std": 0.0272929978517565, + "best_static_condition": "static_dropout_0.3", + "best_static_mean": 4.8010413100322085, + "best_static_std": 0.023497612891923463, + "delta": 0.015909907718500094 + }, + { + "stage": 2, + "prefix": 1000000, + "condition_dropout": 0.116, + "condition_mean": 4.538959090908368, + "condition_std": 0.018682984974041847, + "best_static_condition": "static_dropout_0.26", + "best_static_mean": 4.507193277279536, + "best_static_std": 0.022063926876841285, + "delta": 0.03176581362883191 + }, + { + "stage": 3, + "prefix": 2000000, + "condition_dropout": 0.03599999999999999, + "condition_mean": 4.310809387514989, + "condition_std": 0.02431455246903313, + "best_static_condition": "static_dropout_0.26", + "best_static_mean": 4.275388439496358, + "best_static_std": 0.021064935851257058, + "delta": 0.03542094801863094 + }, + { + "stage": 4, + "prefix": 4000000, + "condition_dropout": 0.02, + "condition_mean": 4.1124994705120725, + "condition_std": 0.024692308091551114, + "best_static_condition": "static_dropout_0.18", + "best_static_mean": 4.108785220732291, + "best_static_std": 0.016546231012667934, + "delta": 0.003714249779781298 + }, + { + "stage": 5, + "prefix": 8000000, + "condition_dropout": 0.02, + "condition_mean": 3.957240684578816, + "condition_std": 0.008487626829204101, + "best_static_condition": "static_dropout_0.14", + "best_static_mean": 3.9756230898201466, + "best_static_std": 0.008723729300606414, + "delta": -0.018382405241330613 + } + ] + } + } +} \ No newline at end of file diff --git a/paper/references.bib b/paper/references.bib index 5837e4783444c1c3d85b024974c66738e03da1a8..a98393500825c212ba459fc96fd3d991f689b7c9 100644 --- a/paper/references.bib +++ b/paper/references.bib @@ -33,6 +33,14 @@ url = {https://arxiv.org/abs/2203.15556} } +@misc{gokaslan2019openwebtext, + title = {OpenWebText Corpus}, + author = {Gokaslan, Aaron and Cohen, Vanya and Pavlick, Ellie and Tellex, Stefanie}, + year = {2019}, + howpublished = {\url{https://zenodo.org/records/3834942}}, + note = {Open reproduction of the WebText-style corpus} +} + @inproceedings{merity2017pointer, title = {Pointer Sentinel Mixture Models}, author = {Merity, Stephen and Xiong, Caiming and Bradbury, James and Socher, Richard}, diff --git a/paper/tables/coefficient_table.tex b/paper/tables/coefficient_table.tex new file mode 100644 index 0000000000000000000000000000000000000000..40a72bd3788b6aa1457ff6844c14f2a6471e0fdf --- /dev/null +++ b/paper/tables/coefficient_table.tex @@ -0,0 +1,9 @@ +\begin{tabular}{lrrrrrrrr} +\toprule +Regime & Cells & A & B & D & C0 & MAE & L-model MAE & L-prefix MAE \\ +\midrule +OpenWebText10K & 15 & 0.0819 & -0.0239 & 0.1537 & -0.0556 & 0.0158 & 0.0248 & 0.0310 \\ +TinyStories & 16 & -0.0893 & -0.1298 & 0.2551 & 0.0815 & 0.0180 & 0.0295 & 0.0325 \\ +WikiText-103 & 15 & -0.0027 & -0.2836 & 0.2686 & -0.0190 & 0.0162 & 0.0304 & 0.0444 \\ +\bottomrule +\end{tabular} diff --git a/paper/tables/l20_stage_table.tex b/paper/tables/l20_stage_table.tex new file mode 100644 index 0000000000000000000000000000000000000000..3e081f438af116a37ecd1a56b9766b06867b7e80 --- /dev/null +++ b/paper/tables/l20_stage_table.tex @@ -0,0 +1,23 @@ +\begin{tabular}{lrrrrr} +\toprule +Regime & Prefix & Decay dropout & Decay val & Best static val & Delta \\ +\midrule +OpenWebText10K & 250k & 0.433 & 5.4601 & 5.3971 & +0.0630 \\ +OpenWebText10K & 500k & 0.367 & 5.0179 & 5.0169 & +0.0010 \\ +OpenWebText10K & 1M & 0.270 & 4.7070 & 4.7218 & -0.0148 \\ +OpenWebText10K & 2M & 0.176 & 4.4879 & 4.5116 & -0.0237 \\ +OpenWebText10K & 4M & 0.095 & 4.3116 & 4.3388 & -0.0272 \\ +OpenWebText10K & 8M & 0.033 & 4.1793 & 4.2156 & -0.0363 \\ +TinyStories & 500k & 0.253 & 3.0302 & 3.0297 & +0.0005 \\ +TinyStories & 1M & 0.209 & 2.7195 & 2.7122 & +0.0073 \\ +TinyStories & 2M & 0.137 & 2.4768 & 2.4733 & +0.0035 \\ +TinyStories & 4M & 0.077 & 2.3110 & 2.3167 & -0.0057 \\ +TinyStories & 8M & 0.041 & 2.2086 & 2.2186 & -0.0100 \\ +WikiText-103 & 250k & 0.284 & 5.2075 & 5.1995 & +0.0080 \\ +WikiText-103 & 500k & 0.211 & 4.8170 & 4.8010 & +0.0159 \\ +WikiText-103 & 1M & 0.116 & 4.5390 & 4.5072 & +0.0318 \\ +WikiText-103 & 2M & 0.036 & 4.3108 & 4.2754 & +0.0354 \\ +WikiText-103 & 4M & 0.020 & 4.1125 & 4.1088 & +0.0037 \\ +WikiText-103 & 8M & 0.020 & 3.9572 & 3.9756 & -0.0184 \\ +\bottomrule +\end{tabular} diff --git a/paper/tables/streaming_validation_table.tex b/paper/tables/streaming_validation_table.tex new file mode 100644 index 0000000000000000000000000000000000000000..6e6455addf8f8701cb1a732ee8c62af89b8974aa --- /dev/null +++ b/paper/tables/streaming_validation_table.tex @@ -0,0 +1,12 @@ +\begin{tabular}{llrrrrr} +\toprule +Regime & Protocol & Decay final & Best static final & Gain & Seeds & Paired wins \\ +\midrule +OpenWebText10K & 4M / 5 seeds & 4.3981 $\pm$ 0.0095 & 4.4455 $\pm$ 0.0120 & 0.0473 & 5 & 5/5 \\ +OpenWebText10K & 8M / 3 seeds & 4.1793 $\pm$ 0.0134 & 4.2156 $\pm$ 0.0142 & 0.0363 & 3 & 3/3 \\ +TinyStories & 4M / 5 seeds & 2.5311 $\pm$ 0.0213 & 2.5444 $\pm$ 0.0211 & 0.0133 & 5 & 5/5 \\ +TinyStories & 8M / 3 seeds & 2.2086 $\pm$ 0.0046 & 2.2186 $\pm$ 0.0032 & 0.0100 & 3 & 3/3 \\ +WikiText-103 & 4M / 5 seeds & 4.0808 $\pm$ 0.0195 & 4.1105 $\pm$ 0.0188 & 0.0297 & 5 & 5/5 \\ +WikiText-103 & 8M / 3 seeds & 3.9572 $\pm$ 0.0085 & 3.9756 $\pm$ 0.0087 & 0.0184 & 3 & 3/3 \\ +\bottomrule +\end{tabular} diff --git a/runs/smoke_compat/locked_stream/20260605-224834/RESULT_SUMMARY.md b/runs/smoke_compat/locked_stream/20260605-224834/RESULT_SUMMARY.md new file mode 100644 index 0000000000000000000000000000000000000000..df576ac5c1aafa16f645ba6df04bb31407330971 --- /dev/null +++ b/runs/smoke_compat/locked_stream/20260605-224834/RESULT_SUMMARY.md @@ -0,0 +1,29 @@ +# Locked Streaming Dropout Summary + +Run directory: `runs/smoke_compat/locked_stream/20260605-224834` + +Model: `smoke` causal Transformer, 134,144 parameters, 1 layers, 1 heads, 16 embedding dim. +Training per stage: 2 steps. Sampled tokens are cumulative in each stage row. Seeds present: 123. + +## Condition Ranking + +| Condition | Kind | Final dropout | Mean trajectory val loss | Final val loss | Final gap | Dropout path | +|---|---|---:|---:|---:|---:|---| +| `static_dropout_0.05` | static | 0.05 | 8.3189 | 8.3194 | 0.0069 | 0.05 -> 0.05 | +| `smoke_decay` | anchor_decay | 0.05 | 8.3189 | 8.3194 | 0.0014 | 0.10 -> 0.05 | + +## Stage Trajectory + +### Stage 0: 512 Prefix Tokens + +| Condition | Dropout | Mean val loss | Mean train loss | Mean gap | N | +|---|---:|---:|---:|---:|---:| +| `static_dropout_0.05` | 0.05 | 8.3184 | 8.3184 | -0.0001 | 1 | +| `smoke_decay` | 0.10 | 8.3185 | 8.3181 | 0.0003 | 1 | + +### Stage 1: 768 Prefix Tokens + +| Condition | Dropout | Mean val loss | Mean train loss | Mean gap | N | +|---|---:|---:|---:|---:|---:| +| `smoke_decay` | 0.05 | 8.3194 | 8.3180 | 0.0014 | 1 | +| `static_dropout_0.05` | 0.05 | 8.3194 | 8.3125 | 0.0069 | 1 | diff --git a/runs/smoke_compat/locked_stream/20260605-224834/config.json b/runs/smoke_compat/locked_stream/20260605-224834/config.json new file mode 100644 index 0000000000000000000000000000000000000000..939e79c8a29414d33332d0dee46753ab3c5707b0 --- /dev/null +++ b/runs/smoke_compat/locked_stream/20260605-224834/config.json @@ -0,0 +1,103 @@ +{ + "args": { + "mode": "locked_stream", + "corpus": null, + "corpus_glob": null, + "text_column": "text", + "use_cached_data": true, + "output_dir": "runs/smoke_compat", + "resume_from": null, + "cache_dir": ".cache/dropout_decay", + "models": [ + "smoke=1x1x16" + ], + "seeds": [ + 123 + ], + "token_limits": [ + 5000000 + ], + "stream_token_caps": [ + 512, + 768 + ], + "val_tokens": 128, + "allow_short_corpus": true, + "force_retokenize": false, + "vocab_size": 4096, + "tokenizer_train_chars": 10000000, + "block_size": 16, + "batch_size": 2, + "steps": 2000, + "stage_steps": 2, + "dropout_rates": [ + 0.05 + ], + "decays": [], + "anchor_decays": [ + { + "name": "smoke_decay", + "kind": "anchor_decay", + "initial": 0.1, + "final": 0.05, + "schedule": "log_prefix_anchor", + "decay_tokens": null, + "anchors": [ + [ + 512, + 0.1 + ], + [ + 768, + 0.05 + ] + ] + } + ], + "decay_tokens": null, + "eval_batches": 1, + "train_eval_batches": 1, + "trace_eval_batches": 1, + "eval_every": 0, + "log_every": 1, + "lr": 0.0003, + "weight_decay": 0.1, + "grad_clip": 1.0, + "plateau_delta": 0.01, + "target_min_dropout": 0.1, + "min_nonzero_margin": 0.01, + "min_high_dropout_margin": 0.03, + "screen_early_stop": false, + "screen_prune_patience": 3, + "screen_prune_min_delta": 0.01 + }, + "mode": "locked_stream", + "seeds": [ + 123 + ], + "models": [ + { + "model_name": "smoke", + "n_layer": 1, + "n_head": 1, + "n_embd": 16 + } + ], + "device": "mps", + "torch": "2.12.0", + "python": "3.11.15 (main, Mar 3 2026, 00:52:57) [Clang 21.0.0 (clang-2100.0.123.102)]", + "mps_available": true, + "attribution": "Derived from Andrej Karpathy's nanochat project (https://github.com/karpathy/nanochat), MIT License, Copyright (c) 2025 Andrej Karpathy.", + "tokenizer_path": ".cache/dropout_decay/tokenizer-v4096.json", + "encoded_path": ".cache/dropout_decay/tokens-v4096-uint16.npy", + "train_tokens": 5500842, + "val_tokens": 128, + "effective_token_limits": [ + 5000000 + ], + "effective_stream_token_caps": [ + 512, + 768 + ], + "resume_from": null +} \ No newline at end of file diff --git a/runs/smoke_compat/locked_stream/20260605-224834/metrics.jsonl b/runs/smoke_compat/locked_stream/20260605-224834/metrics.jsonl new file mode 100644 index 0000000000000000000000000000000000000000..6baa03432d640e354c5f5fbee35707f76c1aa4f0 --- /dev/null +++ b/runs/smoke_compat/locked_stream/20260605-224834/metrics.jsonl @@ -0,0 +1,4 @@ +{"condition": "smoke_decay", "condition_kind": "anchor_decay", "dropout_active_final": 0.1, "dropout_final": 0.05, "dropout_initial": 0.1, "dropout_schedule": "log_prefix_anchor", "elapsed_sec": 0.08690309524536133, "eval_loss": 8.318451881408691, "generalization_gap": 0.00031757354736328125, "model_config": {"block_size": 16, "dropout": 0.1, "n_embd": 16, "n_head": 1, "n_layer": 1, "vocab_size": 4096}, "model_name": "smoke", "n_embd": 16, "n_head": 1, "n_layer": 1, "parameters": 134144, "run_mode": "locked_stream", "seed": 123, "stage": 0, "steps": 2, "token_limit": 512, "tokens_seen": 64, "train_eval_loss": 8.318134307861328, "train_loss_last": 8.317201614379883, "val_eval_loss": 8.318451881408691} +{"condition": "smoke_decay", "condition_kind": "anchor_decay", "dropout_active_final": 0.05, "dropout_final": 0.05, "dropout_initial": 0.1, "dropout_schedule": "log_prefix_anchor", "elapsed_sec": 0.018578290939331055, "eval_loss": 8.319382667541504, "generalization_gap": 0.0013904571533203125, "model_config": {"block_size": 16, "dropout": 0.1, "n_embd": 16, "n_head": 1, "n_layer": 1, "vocab_size": 4096}, "model_name": "smoke", "n_embd": 16, "n_head": 1, "n_layer": 1, "parameters": 134144, "run_mode": "locked_stream", "seed": 123, "stage": 1, "steps": 2, "token_limit": 768, "tokens_seen": 128, "train_eval_loss": 8.317992210388184, "train_loss_last": 8.3181734085083, "val_eval_loss": 8.319382667541504} +{"condition": "static_dropout_0.05", "condition_kind": "static", "dropout_active_final": 0.05, "dropout_final": 0.05, "dropout_initial": 0.05, "dropout_schedule": "constant", "elapsed_sec": 0.019870996475219727, "eval_loss": 8.318363189697266, "generalization_gap": -8.20159912109375e-05, "model_config": {"block_size": 16, "dropout": 0.05, "n_embd": 16, "n_head": 1, "n_layer": 1, "vocab_size": 4096}, "model_name": "smoke", "n_embd": 16, "n_head": 1, "n_layer": 1, "parameters": 134144, "run_mode": "locked_stream", "seed": 123, "stage": 0, "steps": 2, "token_limit": 512, "tokens_seen": 64, "train_eval_loss": 8.318445205688477, "train_loss_last": 8.317625999450684, "val_eval_loss": 8.318363189697266} +{"condition": "static_dropout_0.05", "condition_kind": "static", "dropout_active_final": 0.05, "dropout_final": 0.05, "dropout_initial": 0.05, "dropout_schedule": "constant", "elapsed_sec": 0.017475128173828125, "eval_loss": 8.319437026977539, "generalization_gap": 0.0068874359130859375, "model_config": {"block_size": 16, "dropout": 0.05, "n_embd": 16, "n_head": 1, "n_layer": 1, "vocab_size": 4096}, "model_name": "smoke", "n_embd": 16, "n_head": 1, "n_layer": 1, "parameters": 134144, "run_mode": "locked_stream", "seed": 123, "stage": 1, "steps": 2, "token_limit": 768, "tokens_seen": 128, "train_eval_loss": 8.312549591064453, "train_loss_last": 8.316361427307129, "val_eval_loss": 8.319437026977539} diff --git a/runs/smoke_compat/locked_stream/20260605-224834/summary.csv b/runs/smoke_compat/locked_stream/20260605-224834/summary.csv new file mode 100644 index 0000000000000000000000000000000000000000..9ebe04f17e446e8b7b5b969ef229217aca0b2bd9 --- /dev/null +++ b/runs/smoke_compat/locked_stream/20260605-224834/summary.csv @@ -0,0 +1,5 @@ +run_mode,condition,condition_kind,stage,token_limit,model_name,n_layer,n_head,n_embd,parameters,dropout_initial,dropout_final,dropout_schedule,n,mean_train_eval_loss,std_train_eval_loss,mean_val_eval_loss,std_val_eval_loss,mean_generalization_gap,std_generalization_gap +locked_stream,smoke_decay,anchor_decay,0,512,smoke,1,1,16,134144,0.1,0.05,log_prefix_anchor,1,8.318134307861328,0.0,8.318451881408691,0.0,0.00031757354736328125,0.0 +locked_stream,static_dropout_0.05,static,0,512,smoke,1,1,16,134144,0.05,0.05,constant,1,8.318445205688477,0.0,8.318363189697266,0.0,-8.20159912109375e-05,0.0 +locked_stream,smoke_decay,anchor_decay,1,768,smoke,1,1,16,134144,0.1,0.05,log_prefix_anchor,1,8.317992210388184,0.0,8.319382667541504,0.0,0.0013904571533203125,0.0 +locked_stream,static_dropout_0.05,static,1,768,smoke,1,1,16,134144,0.05,0.05,constant,1,8.312549591064453,0.0,8.319437026977539,0.0,0.0068874359130859375,0.0 diff --git a/runs/smoke_compat/locked_stream/20260605-224834/summary.json b/runs/smoke_compat/locked_stream/20260605-224834/summary.json new file mode 100644 index 0000000000000000000000000000000000000000..db3b3af2a08972ad7a434c8f242faff93a0e5675 --- /dev/null +++ b/runs/smoke_compat/locked_stream/20260605-224834/summary.json @@ -0,0 +1,90 @@ +[ + { + "run_mode": "locked_stream", + "condition": "smoke_decay", + "condition_kind": "anchor_decay", + "stage": 0, + "token_limit": 512, + "model_name": "smoke", + "n_layer": 1, + "n_head": 1, + "n_embd": 16, + "parameters": 134144, + "dropout_initial": 0.1, + "dropout_final": 0.05, + "dropout_schedule": "log_prefix_anchor", + "n": 1, + "mean_train_eval_loss": 8.318134307861328, + "std_train_eval_loss": 0.0, + "mean_val_eval_loss": 8.318451881408691, + "std_val_eval_loss": 0.0, + "mean_generalization_gap": 0.00031757354736328125, + "std_generalization_gap": 0.0 + }, + { + "run_mode": "locked_stream", + "condition": "static_dropout_0.05", + "condition_kind": "static", + "stage": 0, + "token_limit": 512, + "model_name": "smoke", + "n_layer": 1, + "n_head": 1, + "n_embd": 16, + "parameters": 134144, + "dropout_initial": 0.05, + "dropout_final": 0.05, + "dropout_schedule": "constant", + "n": 1, + "mean_train_eval_loss": 8.318445205688477, + "std_train_eval_loss": 0.0, + "mean_val_eval_loss": 8.318363189697266, + "std_val_eval_loss": 0.0, + "mean_generalization_gap": -8.20159912109375e-05, + "std_generalization_gap": 0.0 + }, + { + "run_mode": "locked_stream", + "condition": "smoke_decay", + "condition_kind": "anchor_decay", + "stage": 1, + "token_limit": 768, + "model_name": "smoke", + "n_layer": 1, + "n_head": 1, + "n_embd": 16, + "parameters": 134144, + "dropout_initial": 0.1, + "dropout_final": 0.05, + "dropout_schedule": "log_prefix_anchor", + "n": 1, + "mean_train_eval_loss": 8.317992210388184, + "std_train_eval_loss": 0.0, + "mean_val_eval_loss": 8.319382667541504, + "std_val_eval_loss": 0.0, + "mean_generalization_gap": 0.0013904571533203125, + "std_generalization_gap": 0.0 + }, + { + "run_mode": "locked_stream", + "condition": "static_dropout_0.05", + "condition_kind": "static", + "stage": 1, + "token_limit": 768, + "model_name": "smoke", + "n_layer": 1, + "n_head": 1, + "n_embd": 16, + "parameters": 134144, + "dropout_initial": 0.05, + "dropout_final": 0.05, + "dropout_schedule": "constant", + "n": 1, + "mean_train_eval_loss": 8.312549591064453, + "std_train_eval_loss": 0.0, + "mean_val_eval_loss": 8.319437026977539, + "std_val_eval_loss": 0.0, + "mean_generalization_gap": 0.0068874359130859375, + "std_generalization_gap": 0.0 + } +] \ No newline at end of file diff --git a/runs/smoke_compat/locked_stream/20260605-224834/trace.jsonl b/runs/smoke_compat/locked_stream/20260605-224834/trace.jsonl new file mode 100644 index 0000000000000000000000000000000000000000..7023fe9a49b453cd95666e06b407b9e9df3b8513 --- /dev/null +++ b/runs/smoke_compat/locked_stream/20260605-224834/trace.jsonl @@ -0,0 +1,8 @@ +{"condition": "smoke_decay", "dropout": 0.1, "event": "train_step", "model_name": "smoke", "run_mode": "locked_stream", "seed": 123, "stage": 0, "step": 1, "steps": 2, "token_limit": 512, "tokens_seen": 32, "train_batch_loss": 8.31665325164795} +{"condition": "smoke_decay", "dropout": 0.1, "event": "train_step", "model_name": "smoke", "run_mode": "locked_stream", "seed": 123, "stage": 0, "step": 2, "steps": 2, "token_limit": 512, "tokens_seen": 64, "train_batch_loss": 8.317201614379883} +{"condition": "smoke_decay", "dropout": 0.05, "event": "train_step", "model_name": "smoke", "run_mode": "locked_stream", "seed": 123, "stage": 1, "step": 1, "steps": 2, "token_limit": 768, "tokens_seen": 96, "train_batch_loss": 8.316765785217285} +{"condition": "smoke_decay", "dropout": 0.05, "event": "train_step", "model_name": "smoke", "run_mode": "locked_stream", "seed": 123, "stage": 1, "step": 2, "steps": 2, "token_limit": 768, "tokens_seen": 128, "train_batch_loss": 8.3181734085083} +{"condition": "static_dropout_0.05", "dropout": 0.05, "event": "train_step", "model_name": "smoke", "run_mode": "locked_stream", "seed": 123, "stage": 0, "step": 1, "steps": 2, "token_limit": 512, "tokens_seen": 32, "train_batch_loss": 8.317569732666016} +{"condition": "static_dropout_0.05", "dropout": 0.05, "event": "train_step", "model_name": "smoke", "run_mode": "locked_stream", "seed": 123, "stage": 0, "step": 2, "steps": 2, "token_limit": 512, "tokens_seen": 64, "train_batch_loss": 8.317625999450684} +{"condition": "static_dropout_0.05", "dropout": 0.05, "event": "train_step", "model_name": "smoke", "run_mode": "locked_stream", "seed": 123, "stage": 1, "step": 1, "steps": 2, "token_limit": 768, "tokens_seen": 96, "train_batch_loss": 8.318181037902832} +{"condition": "static_dropout_0.05", "dropout": 0.05, "event": "train_step", "model_name": "smoke", "run_mode": "locked_stream", "seed": 123, "stage": 1, "step": 2, "steps": 2, "token_limit": 768, "tokens_seen": 128, "train_batch_loss": 8.316361427307129} diff --git a/runs/smoke_compat/screen_static/20260605-224821/RESULT_SUMMARY.md b/runs/smoke_compat/screen_static/20260605-224821/RESULT_SUMMARY.md new file mode 100644 index 0000000000000000000000000000000000000000..a497a6bbb1574ab12ae45a1cb23e2b9cad8c5e1f --- /dev/null +++ b/runs/smoke_compat/screen_static/20260605-224821/RESULT_SUMMARY.md @@ -0,0 +1,24 @@ +# Static Dropout Screen Summary + +Run directory: `runs/smoke_compat/screen_static/20260605-224821` + +## Models + +| Model | Params | Layers | Heads | Embedding | Block | Vocab | Seeds | +|---|---:|---:|---:|---:|---:|---:|---| +| `smoke` | 134,144 | 1 | 1 | 16 | 16 | 4096 | 123 | + +## Best Dropout By Model And Prefix + +| Model | Prefix tokens | Effective epochs | Best dropout | Mean val loss | Val std | Mean train loss | Mean gap | Plateau/bracket note | +|---|---:|---:|---:|---:|---:|---:|---:|---| +| `smoke` | 512 | 0.12 | 0.10 | 8.3185 | 0.0000 | 8.3181 | 0.0003 | not bracketed; best at top of tested grid | + +## Model `smoke` + +### Prefix 512 Tokens (0.12 Effective Epochs) + +| Dropout | N | Mean val loss | Val std | Mean train loss | Train std | Mean gap | Gap std | Sampled tokens | Params | +|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| 0.00 | 1 | 8.3186 | 0.0000 | 8.3134 | 0.0000 | 0.0052 | 0.0000 | 64 | 134,144 | +| 0.10 | 1 | 8.3185 | 0.0000 | 8.3181 | 0.0000 | 0.0003 | 0.0000 | 64 | 134,144 | diff --git a/runs/smoke_compat/screen_static/20260605-224821/config.json b/runs/smoke_compat/screen_static/20260605-224821/config.json new file mode 100644 index 0000000000000000000000000000000000000000..f3e3639b1855f08d10cb99e5e7d5680b419c8090 --- /dev/null +++ b/runs/smoke_compat/screen_static/20260605-224821/config.json @@ -0,0 +1,89 @@ +{ + "args": { + "mode": "screen_static", + "corpus": null, + "corpus_glob": null, + "text_column": "text", + "use_cached_data": true, + "output_dir": "runs/smoke_compat", + "resume_from": null, + "cache_dir": ".cache/dropout_decay", + "models": [ + "smoke=1x1x16" + ], + "seeds": [ + 123 + ], + "token_limits": [ + 512 + ], + "stream_token_caps": [ + 5000000, + 10000000, + 20000000, + 40000000 + ], + "val_tokens": 128, + "allow_short_corpus": true, + "force_retokenize": false, + "vocab_size": 4096, + "tokenizer_train_chars": 10000000, + "block_size": 16, + "batch_size": 2, + "steps": 2, + "stage_steps": 1000, + "dropout_rates": [ + 0.0, + 0.1 + ], + "decays": [], + "anchor_decays": [], + "decay_tokens": null, + "eval_batches": 1, + "train_eval_batches": 1, + "trace_eval_batches": 1, + "eval_every": 0, + "log_every": 1, + "lr": 0.0003, + "weight_decay": 0.1, + "grad_clip": 1.0, + "plateau_delta": 0.01, + "target_min_dropout": 0.1, + "min_nonzero_margin": 0.01, + "min_high_dropout_margin": 0.03, + "screen_early_stop": false, + "screen_prune_patience": 3, + "screen_prune_min_delta": 0.01 + }, + "mode": "screen_static", + "seeds": [ + 123 + ], + "models": [ + { + "model_name": "smoke", + "n_layer": 1, + "n_head": 1, + "n_embd": 16 + } + ], + "device": "mps", + "torch": "2.12.0", + "python": "3.11.15 (main, Mar 3 2026, 00:52:57) [Clang 21.0.0 (clang-2100.0.123.102)]", + "mps_available": true, + "attribution": "Derived from Andrej Karpathy's nanochat project (https://github.com/karpathy/nanochat), MIT License, Copyright (c) 2025 Andrej Karpathy.", + "tokenizer_path": ".cache/dropout_decay/tokenizer-v4096.json", + "encoded_path": ".cache/dropout_decay/tokens-v4096-uint16.npy", + "train_tokens": 5500842, + "val_tokens": 128, + "effective_token_limits": [ + 512 + ], + "effective_stream_token_caps": [ + 5000000, + 5500842, + 5500842, + 5500842 + ], + "resume_from": null +} \ No newline at end of file diff --git a/runs/smoke_compat/screen_static/20260605-224821/dropout_curves.svg b/runs/smoke_compat/screen_static/20260605-224821/dropout_curves.svg new file mode 100644 index 0000000000000000000000000000000000000000..b652d7bebab5727d67d5ff0e76afc377b6ad55fe --- /dev/null +++ b/runs/smoke_compat/screen_static/20260605-224821/dropout_curves.svg @@ -0,0 +1,24 @@ + + + +Static dropout law: validation loss vs dropout +Each panel uses its own y-scale. Points are one-seed means unless N > 1. +512 prefix tokens +smoke (0.1M) + + + +8.34 +8.30 +0 +0.9 + + + +best p=0.10 +loss=8.318 + \ No newline at end of file diff --git a/runs/smoke_compat/screen_static/20260605-224821/metrics.jsonl b/runs/smoke_compat/screen_static/20260605-224821/metrics.jsonl new file mode 100644 index 0000000000000000000000000000000000000000..86629346de1f870b4fd7322741bb40c7565726ba --- /dev/null +++ b/runs/smoke_compat/screen_static/20260605-224821/metrics.jsonl @@ -0,0 +1,2 @@ +{"condition": "static_dropout_0", "condition_kind": "static", "dropout_active_final": 0.0, "dropout_final": 0.0, "dropout_initial": 0.0, "dropout_schedule": "constant", "elapsed_sec": 0.5898959636688232, "eval_loss": 8.318593978881836, "generalization_gap": 0.005218505859375, "model_config": {"block_size": 16, "dropout": 0.0, "n_embd": 16, "n_head": 1, "n_layer": 1, "vocab_size": 4096}, "model_name": "smoke", "n_embd": 16, "n_head": 1, "n_layer": 1, "parameters": 134144, "run_mode": "screen_static", "seed": 123, "stage": null, "steps": 2, "token_limit": 512, "tokens_seen": 64, "train_eval_loss": 8.313375473022461, "train_loss_last": 8.315808296203613, "val_eval_loss": 8.318593978881836} +{"condition": "static_dropout_0.1", "condition_kind": "static", "dropout_active_final": 0.1, "dropout_final": 0.1, "dropout_initial": 0.1, "dropout_schedule": "constant", "elapsed_sec": 0.10111308097839355, "eval_loss": 8.318451881408691, "generalization_gap": 0.00031757354736328125, "model_config": {"block_size": 16, "dropout": 0.1, "n_embd": 16, "n_head": 1, "n_layer": 1, "vocab_size": 4096}, "model_name": "smoke", "n_embd": 16, "n_head": 1, "n_layer": 1, "parameters": 134144, "run_mode": "screen_static", "seed": 123, "stage": null, "steps": 2, "token_limit": 512, "tokens_seen": 64, "train_eval_loss": 8.318134307861328, "train_loss_last": 8.317201614379883, "val_eval_loss": 8.318451881408691} diff --git a/runs/smoke_compat/screen_static/20260605-224821/model_selection.csv b/runs/smoke_compat/screen_static/20260605-224821/model_selection.csv new file mode 100644 index 0000000000000000000000000000000000000000..97d3fec873057fd229516d72e6f2c68c5468d7c9 --- /dev/null +++ b/runs/smoke_compat/screen_static/20260605-224821/model_selection.csv @@ -0,0 +1,2 @@ +run_mode,token_limit,model_name,n_layer,n_head,n_embd,parameters,n,best_dropout,best_val_loss,best_val_std,plateau_start_dropout,plateau_end_dropout,plateau_delta,zero_dropout_val_loss,zero_minus_best,best_nonzero_dropout,best_nonzero_val_loss,zero_minus_best_nonzero,max_dropout,max_dropout_val_loss,max_dropout_minus_best,has_nonzero_optimum,meets_target_dropout,curve_json +screen_static,512,smoke,1,1,16,134144,1,0.1,8.318451881408691,0.0,0.0,0.1,0.01,8.318593978881836,0.00014209747314453125,0.1,8.318451881408691,0.00014209747314453125,0.1,8.318451881408691,0.0,False,True,"[{""dropout"": 0.0, ""mean_generalization_gap"": 0.005218505859375, ""mean_train_loss"": 8.313375473022461, ""mean_val_loss"": 8.318593978881836, ""n"": 1, ""std_val_loss"": 0.0}, {""dropout"": 0.1, ""mean_generalization_gap"": 0.00031757354736328125, ""mean_train_loss"": 8.318134307861328, ""mean_val_loss"": 8.318451881408691, ""n"": 1, ""std_val_loss"": 0.0}]" diff --git a/runs/smoke_compat/screen_static/20260605-224821/model_selection.json b/runs/smoke_compat/screen_static/20260605-224821/model_selection.json new file mode 100644 index 0000000000000000000000000000000000000000..2dba651c4b684d7702105fe796e9fa36c20eb9d5 --- /dev/null +++ b/runs/smoke_compat/screen_static/20260605-224821/model_selection.json @@ -0,0 +1,29 @@ +[ + { + "run_mode": "screen_static", + "token_limit": 512, + "model_name": "smoke", + "n_layer": 1, + "n_head": 1, + "n_embd": 16, + "parameters": 134144, + "n": 1, + "best_dropout": 0.1, + "best_val_loss": 8.318451881408691, + "best_val_std": 0.0, + "plateau_start_dropout": 0.0, + "plateau_end_dropout": 0.1, + "plateau_delta": 0.01, + "zero_dropout_val_loss": 8.318593978881836, + "zero_minus_best": 0.00014209747314453125, + "best_nonzero_dropout": 0.1, + "best_nonzero_val_loss": 8.318451881408691, + "zero_minus_best_nonzero": 0.00014209747314453125, + "max_dropout": 0.1, + "max_dropout_val_loss": 8.318451881408691, + "max_dropout_minus_best": 0.0, + "has_nonzero_optimum": false, + "meets_target_dropout": true, + "curve_json": "[{\"dropout\": 0.0, \"mean_generalization_gap\": 0.005218505859375, \"mean_train_loss\": 8.313375473022461, \"mean_val_loss\": 8.318593978881836, \"n\": 1, \"std_val_loss\": 0.0}, {\"dropout\": 0.1, \"mean_generalization_gap\": 0.00031757354736328125, \"mean_train_loss\": 8.318134307861328, \"mean_val_loss\": 8.318451881408691, \"n\": 1, \"std_val_loss\": 0.0}]" + } +] \ No newline at end of file diff --git a/runs/smoke_compat/screen_static/20260605-224821/summary.csv b/runs/smoke_compat/screen_static/20260605-224821/summary.csv new file mode 100644 index 0000000000000000000000000000000000000000..bf3da06d813e2660587741eb50b985adab3e2589 --- /dev/null +++ b/runs/smoke_compat/screen_static/20260605-224821/summary.csv @@ -0,0 +1,3 @@ +run_mode,condition,condition_kind,stage,token_limit,model_name,n_layer,n_head,n_embd,parameters,dropout_initial,dropout_final,dropout_schedule,n,mean_train_eval_loss,std_train_eval_loss,mean_val_eval_loss,std_val_eval_loss,mean_generalization_gap,std_generalization_gap +screen_static,static_dropout_0,static,,512,smoke,1,1,16,134144,0.0,0.0,constant,1,8.313375473022461,0.0,8.318593978881836,0.0,0.005218505859375,0.0 +screen_static,static_dropout_0.1,static,,512,smoke,1,1,16,134144,0.1,0.1,constant,1,8.318134307861328,0.0,8.318451881408691,0.0,0.00031757354736328125,0.0 diff --git a/runs/smoke_compat/screen_static/20260605-224821/summary.json b/runs/smoke_compat/screen_static/20260605-224821/summary.json new file mode 100644 index 0000000000000000000000000000000000000000..a978e321d704893898d7d306c93f4acd3aee8403 --- /dev/null +++ b/runs/smoke_compat/screen_static/20260605-224821/summary.json @@ -0,0 +1,46 @@ +[ + { + "run_mode": "screen_static", + "condition": "static_dropout_0", + "condition_kind": "static", + "stage": null, + "token_limit": 512, + "model_name": "smoke", + "n_layer": 1, + "n_head": 1, + "n_embd": 16, + "parameters": 134144, + "dropout_initial": 0.0, + "dropout_final": 0.0, + "dropout_schedule": "constant", + "n": 1, + "mean_train_eval_loss": 8.313375473022461, + "std_train_eval_loss": 0.0, + "mean_val_eval_loss": 8.318593978881836, + "std_val_eval_loss": 0.0, + "mean_generalization_gap": 0.005218505859375, + "std_generalization_gap": 0.0 + }, + { + "run_mode": "screen_static", + "condition": "static_dropout_0.1", + "condition_kind": "static", + "stage": null, + "token_limit": 512, + "model_name": "smoke", + "n_layer": 1, + "n_head": 1, + "n_embd": 16, + "parameters": 134144, + "dropout_initial": 0.1, + "dropout_final": 0.1, + "dropout_schedule": "constant", + "n": 1, + "mean_train_eval_loss": 8.318134307861328, + "std_train_eval_loss": 0.0, + "mean_val_eval_loss": 8.318451881408691, + "std_val_eval_loss": 0.0, + "mean_generalization_gap": 0.00031757354736328125, + "std_generalization_gap": 0.0 + } +] \ No newline at end of file diff --git a/runs/smoke_compat/screen_static/20260605-224821/trace.jsonl b/runs/smoke_compat/screen_static/20260605-224821/trace.jsonl new file mode 100644 index 0000000000000000000000000000000000000000..5dff175f18c745825ed4742887da375a95e90080 --- /dev/null +++ b/runs/smoke_compat/screen_static/20260605-224821/trace.jsonl @@ -0,0 +1,4 @@ +{"condition": "static_dropout_0", "dropout": 0.0, "event": "train_step", "model_name": "smoke", "run_mode": "screen_static", "seed": 123, "stage": null, "step": 1, "steps": 2, "token_limit": 512, "tokens_seen": 32, "train_batch_loss": 8.318573951721191} +{"condition": "static_dropout_0", "dropout": 0.0, "event": "train_step", "model_name": "smoke", "run_mode": "screen_static", "seed": 123, "stage": null, "step": 2, "steps": 2, "token_limit": 512, "tokens_seen": 64, "train_batch_loss": 8.315808296203613} +{"condition": "static_dropout_0.1", "dropout": 0.1, "event": "train_step", "model_name": "smoke", "run_mode": "screen_static", "seed": 123, "stage": null, "step": 1, "steps": 2, "token_limit": 512, "tokens_seen": 32, "train_batch_loss": 8.31665325164795} +{"condition": "static_dropout_0.1", "dropout": 0.1, "event": "train_step", "model_name": "smoke", "run_mode": "screen_static", "seed": 123, "stage": null, "step": 2, "steps": 2, "token_limit": 512, "tokens_seen": 64, "train_batch_loss": 8.317201614379883} diff --git a/scripts/generate_paper_assets.py b/scripts/generate_paper_assets.py new file mode 100644 index 0000000000000000000000000000000000000000..0aed462742c3e7aa66ac12942ed553ff7167ae24 --- /dev/null +++ b/scripts/generate_paper_assets.py @@ -0,0 +1,579 @@ +#!/usr/bin/env python3 +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from __future__ import annotations + +import csv +import json +import math +from pathlib import Path +import statistics + + +ROOT = Path(__file__).resolve().parents[1] +PAPER_DIR = ROOT / "paper" +FIG_DIR = PAPER_DIR / "figures" +TABLE_DIR = PAPER_DIR / "tables" + + +REGIMES = { + "openweb": { + "label": "OpenWebText10K", + "short": "OWT10K", + "coeff": ROOT + / "runs/coefficient_calibration/cross_regime_backtest/openwebtext10k_main_interaction/coefficients.json", + "cells": ROOT + / "runs/coefficient_calibration/cross_regime_backtest/openwebtext10k_main_interaction/calibration_cells.csv", + "curve_run": ROOT / "runs/screen_static/20260525-133008", + "curve_model": "L12_H8_D320", + "curve_prefix": 1_000_000, + "curve_title": "OpenWebText10K, L12, 1M", + "small_stream": ROOT + / "runs/openwebtext10k_l16_updated_formula_clean_5seed/locked_stream/20260530-174525", + "small_condition": "openwebtext10k_interaction", + "small_model": "L16_H8_D384", + "small_protocol": "4M / 5 seeds", + "l20_stream": ROOT + / "runs/openwebtext10k_l20_8m_extrapolation_3seed/locked_stream/20260601-191150", + "l20_condition": "openwebtext10k_l20_interaction_8m", + }, + "tinystories": { + "label": "TinyStories", + "short": "TinyStories", + "coeff": ROOT + / "runs/coefficient_calibration/tinystories_combined_plus_all_holdouts_interaction/coefficients.json", + "cells": ROOT + / "runs/coefficient_calibration/tinystories_combined_plus_all_holdouts_interaction/calibration_cells.csv", + "curve_run": ROOT + / "runs/regime_calibration_tinystories_coarse/screen_static/20260529-161726", + "curve_model": "L12_H8_D320", + "curve_prefix": 1_000_000, + "curve_title": "TinyStories, L12, 1M", + "small_stream": [ + ROOT + / "runs/streaming_tinystories_interaction_schedule_l12/locked_stream/20260530-053831", + ROOT + / "runs/streaming_tinystories_multiseed_validation_l12/locked_stream/20260530-111523", + ROOT + / "runs/streaming_tinystories_multiseed_validation_l12/locked_stream/20260530-141335", + ], + "small_condition": "interaction", + "small_model": "L12_H8_D320", + "small_protocol": "4M / 5 seeds", + "l20_stream": ROOT + / "runs/tinystories_l20_8m_extrapolation_3seed/locked_stream/20260602-112249", + "l20_condition": "tinystories_l20_interaction_8m", + }, + "wikitext": { + "label": "WikiText-103", + "short": "WikiText", + "coeff": ROOT / "runs/coefficient_calibration/wikitext103_interaction/coefficients.json", + "cells": ROOT / "runs/coefficient_calibration/wikitext103_interaction/calibration_cells.csv", + "curve_run": ROOT + / "runs/wikitext103_interaction_static_screen/screen_static/20260603-153314", + "curve_model": "L12_H8_D320", + "curve_prefix": 250_000, + "curve_title": "WikiText-103, L12, 250k", + "small_stream": ROOT + / "runs/wikitext103_l12_streaming_validation_5seed/locked_stream/20260531-093525", + "small_condition": "wikitext103_formula_l12", + "small_model": "L12_H8_D320", + "small_protocol": "4M / 5 seeds", + "l20_stream": ROOT + / "runs/wikitext103_l20_8m_extrapolation_3seed/locked_stream/20260603-231952", + "l20_condition": "wikitext103_l20_interaction_8m", + }, +} + + +def read_json(path: Path) -> dict | list: + return json.loads(path.read_text(encoding="utf-8")) + + +def read_jsonl(path: Path) -> list[dict]: + rows = [] + for line in path.read_text(encoding="utf-8").splitlines(): + if line.strip(): + rows.append(json.loads(line)) + return rows + + +def read_csv(path: Path) -> list[dict]: + with path.open(newline="", encoding="utf-8") as handle: + return list(csv.DictReader(handle)) + + +def mean_std(values: list[float]) -> tuple[float, float]: + if not values: + return float("nan"), float("nan") + mean = statistics.fmean(values) + std = statistics.stdev(values) if len(values) > 1 else 0.0 + return mean, std + + +def fmt(x: float, digits: int = 4) -> str: + return f"{x:.{digits}f}" + + +def tex_escape(value: object) -> str: + text = str(value) + return ( + text.replace("\\", "\\textbackslash{}") + .replace("_", "\\_") + .replace("&", "\\&") + .replace("%", "\\%") + .replace("#", "\\#") + ) + + +def token_label(value: int) -> str: + if value >= 1_000_000: + whole = value / 1_000_000 + return f"{whole:g}M" + return f"{value // 1000}k" + + +def metric_stage_summary(run_dir: Path | list[Path], condition_name: str) -> dict: + run_dirs = run_dir if isinstance(run_dir, list) else [run_dir] + rows = [] + for path in run_dirs: + rows.extend(read_jsonl(path / "metrics.jsonl")) + stages = sorted({int(row["stage"]) for row in rows if row.get("stage") is not None}) + final_stage = max(stages) + by_condition_stage: dict[tuple[str, int], list[dict]] = {} + for row in rows: + if row.get("stage") is None: + continue + key = (row["condition"], int(row["stage"])) + by_condition_stage.setdefault(key, []).append(row) + + condition_rows = by_condition_stage[(condition_name, final_stage)] + final_vals = [float(row["val_eval_loss"]) for row in condition_rows] + condition_final_mean, condition_final_std = mean_std(final_vals) + + final_static_by_condition: dict[str, list[dict]] = {} + for row in rows: + if int(row.get("stage", -1)) != final_stage: + continue + if row["condition_kind"] == "static": + final_static_by_condition.setdefault(row["condition"], []).append(row) + static_means = { + condition: statistics.fmean(float(row["val_eval_loss"]) for row in group) + for condition, group in final_static_by_condition.items() + } + best_static_condition = min(static_means, key=static_means.get) + best_static_rows = final_static_by_condition[best_static_condition] + best_static_vals = [float(row["val_eval_loss"]) for row in best_static_rows] + best_static_mean, best_static_std = mean_std(best_static_vals) + + final_static_by_seed: dict[int, tuple[str, float]] = {} + for row in rows: + if int(row.get("stage", -1)) != final_stage or row["condition_kind"] != "static": + continue + seed = int(row["seed"]) + value = float(row["val_eval_loss"]) + if seed not in final_static_by_seed or value < final_static_by_seed[seed][1]: + final_static_by_seed[seed] = (row["condition"], value) + + paired_deltas = [] + for row in condition_rows: + seed = int(row["seed"]) + paired_deltas.append(float(row["val_eval_loss"]) - final_static_by_seed[seed][1]) + paired_mean, paired_std = mean_std(paired_deltas) + + stage_deltas = [] + for stage in stages: + condition_group = by_condition_stage[(condition_name, stage)] + condition_mean = statistics.fmean( + float(row["val_eval_loss"]) for row in condition_group + ) + static_candidates = [] + for (condition, candidate_stage), group in by_condition_stage.items(): + if candidate_stage != stage: + continue + if group[0]["condition_kind"] != "static": + continue + static_candidates.append( + ( + condition, + statistics.fmean(float(row["val_eval_loss"]) for row in group), + statistics.stdev(float(row["val_eval_loss"]) for row in group) + if len(group) > 1 + else 0.0, + ) + ) + best_static_stage = min(static_candidates, key=lambda item: item[1]) + first = condition_group[0] + stage_deltas.append( + { + "stage": stage, + "prefix": int(first["token_limit"]), + "condition_dropout": statistics.fmean( + float(row["dropout_active_final"]) for row in condition_group + ), + "condition_mean": condition_mean, + "condition_std": statistics.stdev( + float(row["val_eval_loss"]) for row in condition_group + ) + if len(condition_group) > 1 + else 0.0, + "best_static_condition": best_static_stage[0], + "best_static_mean": best_static_stage[1], + "best_static_std": best_static_stage[2], + "delta": condition_mean - best_static_stage[1], + } + ) + + return { + "n": len(condition_rows), + "final_stage": final_stage, + "final_prefix": int(condition_rows[0]["token_limit"]), + "condition_final_mean": condition_final_mean, + "condition_final_std": condition_final_std, + "best_static_condition": best_static_condition, + "best_static_mean": best_static_mean, + "best_static_std": best_static_std, + "improvement": best_static_mean - condition_final_mean, + "paired_deltas": paired_deltas, + "paired_wins": sum(1 for delta in paired_deltas if delta < 0), + "paired_mean_delta": paired_mean, + "paired_std_delta": paired_std, + "stage_deltas": stage_deltas, + } + + +def static_curve_points(run_dir: Path, model: str, prefix: int) -> list[tuple[float, float]]: + rows = read_json(run_dir / "summary.json") + curve = [] + for row in rows: + if row["run_mode"] != "screen_static": + continue + if row["condition_kind"] != "static": + continue + if row["model_name"] != model: + continue + if int(row["token_limit"]) != prefix: + continue + curve.append((float(row["dropout_initial"]), float(row["mean_val_eval_loss"]))) + return sorted(curve) + + +def calibration_points(cells_path: Path) -> dict[str, list[tuple[int, float]]]: + rows = read_csv(cells_path) + grouped: dict[str, dict[int, float]] = {} + for row in rows: + sampled = int(float(row["sampled_tokens"])) + unique = int(float(row["unique_tokens"])) + source = row["source"] + if "sample_pressure_low" in source: + continue + if "heldout_model" in source: + continue + model = row["model_name"] + target = max(0.0, float(row["target_dropout"])) + grouped.setdefault(model, {})[unique] = target + return { + model: sorted(points.items()) + for model, points in sorted(grouped.items()) + if len(points) >= 2 + } + + +def coeff_summary(regime_key: str, paths: dict) -> dict: + coeff = read_json(paths["coeff"]) + return { + "regime": paths["label"], + "cells": int(coeff["metrics"]["n"]), + "A": float(coeff["coefficients"]["A"]), + "B": float(coeff["coefficients"]["B"]), + "D": float(coeff["coefficients"]["D"]), + "C0": float(coeff["coefficients"]["C0"]), + "mae": float(coeff["metrics"]["mae"]), + "rmse": float(coeff["metrics"]["rmse"]), + "leave_model_mae": float(coeff["cv"]["leave_model"]["mae"]), + "leave_prefix_mae": float(coeff["cv"]["leave_prefix"]["mae"]), + "formula": coeff["formula"], + } + + +def tikz_coordinates(points: list[tuple[float, float]], digits: int = 4) -> str: + return " ".join(f"({x:.4g},{y:.{digits}f})" for x, y in points) + + +def write_calibration_curve_figure() -> None: + axes = [] + for key, paths in REGIMES.items(): + raw_curve = static_curve_points( + paths["curve_run"], paths["curve_model"], paths["curve_prefix"] + ) + minimum = min(loss for _, loss in raw_curve) + points = [(dropout, loss - minimum) for dropout, loss in raw_curve] + best = min(raw_curve, key=lambda item: item[1]) + axes.append( + rf""" +\nextgroupplot[ + title={{{paths["curve_title"]}}}, + xlabel={{Dropout}}, + ylabel={{Validation loss above best}}, + ymin=0, + grid=both, +] +\addplot+[thick, mark=*] coordinates {{{tikz_coordinates(points, digits=5)}}}; +\addlegendentry{{static sweep}} +\addplot+[only marks, mark=star, mark size=3pt] coordinates {{({best[0]:.4g},0)}}; +\addlegendentry{{target $p^\star={best[0]:.2f}$}} +""" + ) + content = ( + "\\begin{tikzpicture}\n" + "\\begin{groupplot}[\n" + " group style={group size=3 by 1, horizontal sep=1.0cm},\n" + " width=0.32\\textwidth,\n" + " height=0.26\\textwidth,\n" + " legend style={font=\\scriptsize},\n" + " tick label style={font=\\scriptsize},\n" + " label style={font=\\scriptsize},\n" + " title style={font=\\scriptsize},\n" + "]\n" + + "\n".join(axes) + + "\n\\end{groupplot}\n\\end{tikzpicture}\n" + ) + (FIG_DIR / "calibration_cell_curves.tex").write_text(content, encoding="utf-8") + + +def write_static_optima_figure() -> None: + axes = [] + for key, paths in REGIMES.items(): + grouped = calibration_points(paths["cells"]) + labels = sorted({prefix for points in grouped.values() for prefix, _ in points}) + label_to_x = {prefix: idx for idx, prefix in enumerate(labels)} + plots = [] + for model, points in grouped.items(): + coords = [ + (label_to_x[prefix], target) + for prefix, target in points + if prefix in label_to_x + ] + plots.append( + "\\addplot+[thick, mark=*] coordinates {" + + tikz_coordinates(coords) + + "};\n" + + f"\\addlegendentry{{{tex_escape(model)}}}" + ) + xticks = ",".join(str(i) for i in range(len(labels))) + xticklabels = ",".join(token_label(prefix) for prefix in labels) + axes.append( + rf""" +\nextgroupplot[ + title={{{paths["label"]}}}, + xlabel={{Unique prefix $U$}}, + ylabel={{Fitted target dropout}}, + xtick={{{xticks}}}, + xticklabels={{{xticklabels}}}, + ymin=0, + ymax=0.65, + grid=both, +] +{chr(10).join(plots)} +""" + ) + content = ( + "\\begin{tikzpicture}\n" + "\\begin{groupplot}[\n" + " group style={group size=3 by 1, horizontal sep=1.0cm},\n" + " width=0.32\\textwidth,\n" + " height=0.28\\textwidth,\n" + " legend style={font=\\tiny},\n" + " tick label style={font=\\scriptsize},\n" + " label style={font=\\scriptsize},\n" + " title style={font=\\scriptsize},\n" + "]\n" + + "\n".join(axes) + + "\n\\end{groupplot}\n\\end{tikzpicture}\n" + ) + (FIG_DIR / "static_optima_by_regime.tex").write_text(content, encoding="utf-8") + + +def write_l20_stage_delta_figure(summaries: dict[str, dict]) -> None: + axes = [] + for key, paths in REGIMES.items(): + stage_rows = summaries[key]["l20"]["stage_deltas"] + labels = [token_label(row["prefix"]) for row in stage_rows] + coords = [(idx, row["delta"]) for idx, row in enumerate(stage_rows)] + zero_coords = [(0, 0.0), (len(stage_rows) - 1, 0.0)] + xticks = ",".join(str(i) for i in range(len(labels))) + xticklabels = ",".join(labels) + axes.append( + rf""" +\nextgroupplot[ + title={{{paths["label"]}}}, + xlabel={{Prefix}}, + ylabel={{Decay $-$ best static}}, + xtick={{{xticks}}}, + xticklabels={{{xticklabels}}}, + grid=both, +] +\addplot+[thick, mark=*] coordinates {{{tikz_coordinates(coords, digits=5)}}}; +\addplot+[black, dashed, mark=none] coordinates {{{tikz_coordinates(zero_coords, digits=5)}}}; +""" + ) + content = ( + "\\begin{tikzpicture}\n" + "\\begin{groupplot}[\n" + " group style={group size=3 by 1, horizontal sep=1.0cm},\n" + " width=0.32\\textwidth,\n" + " height=0.27\\textwidth,\n" + " tick label style={font=\\scriptsize},\n" + " label style={font=\\scriptsize},\n" + " title style={font=\\scriptsize},\n" + "]\n" + + "\n".join(axes) + + "\n\\end{groupplot}\n\\end{tikzpicture}\n" + ) + (FIG_DIR / "l20_stage_deltas.tex").write_text(content, encoding="utf-8") + + +def write_final_improvement_figure(summaries: dict[str, dict]) -> None: + labels = [] + coords_small = [] + coords_l20 = [] + for idx, (key, paths) in enumerate(REGIMES.items()): + labels.append(paths["short"]) + coords_small.append((idx, summaries[key]["small"]["improvement"])) + coords_l20.append((idx, summaries[key]["l20"]["improvement"])) + xticks = ",".join(str(i) for i in range(len(labels))) + xticklabels = ",".join(labels) + content = rf""" +\begin{{tikzpicture}} +\begin{{axis}}[ + ybar, + bar width=8pt, + width=0.72\textwidth, + height=0.34\textwidth, + ymin=0, + ylabel={{Best static final loss $-$ decay final loss}}, + symbolic x coords={{{xticklabels}}}, + xtick=data, + xtick={{{xticklabels}}}, + grid=both, + legend style={{font=\scriptsize, at={{(0.5,1.04)}}, anchor=south, legend columns=2}}, + tick label style={{font=\scriptsize}}, + label style={{font=\scriptsize}}, +] +\addplot coordinates {{{" ".join(f"({label},{y:.5f})" for label, (_, y) in zip(labels, coords_small))}}}; +\addlegendentry{{4M validation}} +\addplot coordinates {{{" ".join(f"({label},{y:.5f})" for label, (_, y) in zip(labels, coords_l20))}}}; +\addlegendentry{{L20 8M extrapolation}} +\end{{axis}} +\end{{tikzpicture}} +""" + (FIG_DIR / "final_improvements.tex").write_text(content, encoding="utf-8") + + +def write_tables(summaries: dict[str, dict], coeffs: dict[str, dict]) -> None: + coeff_lines = [ + "\\begin{tabular}{lrrrrrrrr}", + "\\toprule", + "Regime & Cells & A & B & D & C0 & MAE & L-model MAE & L-prefix MAE \\\\", + "\\midrule", + ] + for key in REGIMES: + row = coeffs[key] + coeff_lines.append( + f"{tex_escape(row['regime'])} & {row['cells']} & " + f"{row['A']:.4f} & {row['B']:.4f} & {row['D']:.4f} & {row['C0']:.4f} & " + f"{row['mae']:.4f} & {row['leave_model_mae']:.4f} & {row['leave_prefix_mae']:.4f} \\\\" + ) + coeff_lines.extend(["\\bottomrule", "\\end{tabular}"]) + (TABLE_DIR / "coefficient_table.tex").write_text( + "\n".join(coeff_lines) + "\n", encoding="utf-8" + ) + + val_lines = [ + "\\begin{tabular}{llrrrrr}", + "\\toprule", + "Regime & Protocol & Decay final & Best static final & Gain & Seeds & Paired wins \\\\", + "\\midrule", + ] + for key, paths in REGIMES.items(): + for protocol_key, protocol_label in [("small", paths["small_protocol"]), ("l20", "8M / 3 seeds")]: + row = summaries[key][protocol_key] + val_lines.append( + f"{tex_escape(paths['label'])} & {tex_escape(protocol_label)} & " + f"{row['condition_final_mean']:.4f} $\\pm$ {row['condition_final_std']:.4f} & " + f"{row['best_static_mean']:.4f} $\\pm$ {row['best_static_std']:.4f} & " + f"{row['improvement']:.4f} & {row['n']} & {row['paired_wins']}/{row['n']} \\\\" + ) + val_lines.extend(["\\bottomrule", "\\end{tabular}"]) + (TABLE_DIR / "streaming_validation_table.tex").write_text( + "\n".join(val_lines) + "\n", encoding="utf-8" + ) + + l20_lines = [ + "\\begin{tabular}{lrrrrr}", + "\\toprule", + "Regime & Prefix & Decay dropout & Decay val & Best static val & Delta \\\\", + "\\midrule", + ] + for key, paths in REGIMES.items(): + for row in summaries[key]["l20"]["stage_deltas"]: + l20_lines.append( + f"{tex_escape(paths['label'])} & {token_label(row['prefix'])} & " + f"{row['condition_dropout']:.3f} & {row['condition_mean']:.4f} & " + f"{row['best_static_mean']:.4f} & {row['delta']:+.4f} \\\\" + ) + l20_lines.extend(["\\bottomrule", "\\end{tabular}"]) + (TABLE_DIR / "l20_stage_table.tex").write_text( + "\n".join(l20_lines) + "\n", encoding="utf-8" + ) + + summary = { + key: { + "coefficients": coeffs[key], + "small_stream": summaries[key]["small"], + "l20_stream": summaries[key]["l20"], + } + for key in REGIMES + } + (PAPER_DIR / "paper_results_summary.json").write_text( + json.dumps(summary, indent=2), encoding="utf-8" + ) + + +def main() -> None: + FIG_DIR.mkdir(parents=True, exist_ok=True) + TABLE_DIR.mkdir(parents=True, exist_ok=True) + summaries = {} + coeffs = {} + for key, paths in REGIMES.items(): + coeffs[key] = coeff_summary(key, paths) + summaries[key] = { + "small": metric_stage_summary(paths["small_stream"], paths["small_condition"]), + "l20": metric_stage_summary(paths["l20_stream"], paths["l20_condition"]), + } + write_calibration_curve_figure() + write_static_optima_figure() + write_l20_stage_delta_figure(summaries) + write_final_improvement_figure(summaries) + write_tables(summaries, coeffs) + print(json.dumps({"figures": str(FIG_DIR), "tables": str(TABLE_DIR)}, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/src/dropout_decay/data/__init__.py b/src/dropout_decay/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ba3c26d9439afb1fb2d80d3e64f96bd13b6e9fed --- /dev/null +++ b/src/dropout_decay/data/__init__.py @@ -0,0 +1,45 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from dropout_decay.datasets import ( + CachedTokenizer, + TokenSplits, + encode_corpus, + iter_documents, + load_cached_splits, + preview_documents, + required_token_count, + resolve_paths, + train_or_load_tokenizer, +) +from dropout_decay.specs import DropoutCondition, ModelSpec, anchor_dropout + +__all__ = [ + "CachedTokenizer", + "DropoutCondition", + "ModelSpec", + "TokenSplits", + "anchor_dropout", + "encode_corpus", + "iter_documents", + "load_cached_splits", + "preview_documents", + "required_token_count", + "resolve_paths", + "train_or_load_tokenizer", +] diff --git a/src/dropout_decay/datasets/__init__.py b/src/dropout_decay/datasets/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..27693b0ff04d6249770c1e0cf5b7744b8ba52a37 --- /dev/null +++ b/src/dropout_decay/datasets/__init__.py @@ -0,0 +1,41 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from dropout_decay.datasets.corpus import ( + CachedTokenizer, + TokenSplits, + encode_corpus, + iter_documents, + load_cached_splits, + preview_documents, + required_token_count, + resolve_paths, + train_or_load_tokenizer, +) + +__all__ = [ + "CachedTokenizer", + "TokenSplits", + "encode_corpus", + "iter_documents", + "load_cached_splits", + "preview_documents", + "required_token_count", + "resolve_paths", + "train_or_load_tokenizer", +] diff --git a/src/dropout_decay/data.py b/src/dropout_decay/datasets/corpus.py similarity index 99% rename from src/dropout_decay/data.py rename to src/dropout_decay/datasets/corpus.py index 34e26173c27672db817f604e91b18d432b8542db..4fc81d114f769c05017e506b4137a828f6ebcf52 100644 --- a/src/dropout_decay/data.py +++ b/src/dropout_decay/datasets/corpus.py @@ -143,7 +143,7 @@ def train_or_load_tokenizer( text_column: str, force_retrain: bool = False, ) -> object: - from dropout_decay.tokenizer import BpeTokenizer + from dropout_decay.tokenization import BpeTokenizer tokenizer_path = output_dir / f"tokenizer-v{vocab_size}.json" if tokenizer_path.exists() and not force_retrain: diff --git a/src/dropout_decay/experiment.py b/src/dropout_decay/experiment.py index 6c8c16667d9b08623097e9062bef5474b1e2fb6c..ece9fc489af3924b55c39b3566a4a83553194a7a 100644 --- a/src/dropout_decay/experiment.py +++ b/src/dropout_decay/experiment.py @@ -16,1694 +16,13 @@ The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. """ -from __future__ import annotations +from dropout_decay.experiments.runner import build_parser, main, run -import argparse -import csv -from collections import defaultdict -from dataclasses import dataclass -from datetime import datetime -import json -import math -import os -from pathlib import Path -import random -import statistics -import sys -import time -from typing import Callable - -import numpy as np -import torch - -from dropout_decay.data import ( - encode_corpus, - load_cached_splits, - resolve_paths, - train_or_load_tokenizer, -) -from dropout_decay.license import NANOCHAT_ATTRIBUTION -from dropout_decay.model import DropoutGPT, GPTConfig -from dropout_decay.scheduler import DropoutDecayConfig, DropoutDecayScheduler - - -DEFAULT_DROPOUT_RATES = [0.0, 0.02, 0.05, 0.08, 0.10, 0.14, 0.20, 0.30, 0.50] -SUMMARY_FIELDS = [ - "run_mode", - "condition", - "condition_kind", - "stage", - "token_limit", - "model_name", - "n_layer", - "n_head", - "n_embd", - "parameters", - "dropout_initial", - "dropout_final", - "dropout_schedule", - "n", - "mean_train_eval_loss", - "std_train_eval_loss", - "mean_val_eval_loss", - "std_val_eval_loss", - "mean_generalization_gap", - "std_generalization_gap", +__all__ = [ + "build_parser", + "main", + "run", ] -SELECTION_FIELDS = [ - "run_mode", - "token_limit", - "model_name", - "n_layer", - "n_head", - "n_embd", - "parameters", - "n", - "best_dropout", - "best_val_loss", - "best_val_std", - "plateau_start_dropout", - "plateau_end_dropout", - "plateau_delta", - "zero_dropout_val_loss", - "zero_minus_best", - "best_nonzero_dropout", - "best_nonzero_val_loss", - "zero_minus_best_nonzero", - "max_dropout", - "max_dropout_val_loss", - "max_dropout_minus_best", - "has_nonzero_optimum", - "meets_target_dropout", - "curve_json", -] - - -@dataclass(frozen=True) -class ModelSpec: - name: str - n_layer: int - n_head: int - n_embd: int - - def config(self, vocab_size: int, block_size: int, dropout: float) -> GPTConfig: - return GPTConfig( - block_size=block_size, - vocab_size=vocab_size, - n_layer=self.n_layer, - n_head=self.n_head, - n_embd=self.n_embd, - dropout=dropout, - ) - - def to_dict(self) -> dict[str, int | str]: - return { - "model_name": self.name, - "n_layer": self.n_layer, - "n_head": self.n_head, - "n_embd": self.n_embd, - } - - -@dataclass(frozen=True) -class DropoutCondition: - name: str - kind: str - initial: float - final: float - schedule: str = "constant" - decay_tokens: int | None = None - anchors: tuple[tuple[int, float], ...] = () - - def to_dict(self) -> dict: - return { - "name": self.name, - "kind": self.kind, - "initial": self.initial, - "final": self.final, - "schedule": self.schedule, - "decay_tokens": self.decay_tokens, - "anchors": [list(anchor) for anchor in self.anchors], - } - - def make_fn( - self, - fallback_decay_tokens: int, - unique_tokens: int | None = None, - ) -> Callable[[int], float]: - if self.kind == "static": - return lambda _tokens_seen, p=self.initial: p - if self.kind == "anchor_decay": - if unique_tokens is None: - raise ValueError("anchor_decay conditions require unique_tokens") - p = anchor_dropout(unique_tokens, self.anchors) - return lambda _tokens_seen, p=p: p - scheduler = DropoutDecayScheduler( - DropoutDecayConfig( - initial_dropout=self.initial, - final_dropout=self.final, - decay_tokens=self.decay_tokens or fallback_decay_tokens, - schedule=self.schedule, - ) - ) - return scheduler.value - - -def anchor_dropout(unique_tokens: int, anchors: tuple[tuple[int, float], ...]) -> float: - if not anchors: - raise ValueError("anchor dropout schedule requires at least one anchor") - ordered = sorted(anchors) - if unique_tokens <= ordered[0][0]: - return ordered[0][1] - if unique_tokens >= ordered[-1][0]: - return ordered[-1][1] - log_unique = math.log(unique_tokens) - for (left_tokens, left_dropout), (right_tokens, right_dropout) in zip( - ordered, ordered[1:] - ): - if left_tokens <= unique_tokens <= right_tokens: - left_log = math.log(left_tokens) - right_log = math.log(right_tokens) - mix = (log_unique - left_log) / (right_log - left_log) - return left_dropout + mix * (right_dropout - left_dropout) - return ordered[-1][1] - - -def assert_mps_only() -> torch.device: - if os.environ.get("PYTORCH_ENABLE_MPS_FALLBACK") == "1": - raise SystemExit( - "Refusing to run: PYTORCH_ENABLE_MPS_FALLBACK=1 could execute CPU fallbacks." - ) - if not torch.backends.mps.is_built(): - raise SystemExit("PyTorch was not built with MPS. Stopping.") - if not torch.backends.mps.is_available(): - raise SystemExit("MPS is not available. Stopping before any Torch experiment work.") - device = torch.device("mps") - torch.set_default_device(device) - return device - - -def clean_name(value: str) -> str: - cleaned = "".join( - c if c.isalnum() or c in {"-", "_"} else "_" for c in value.strip() - ) - return cleaned.strip("_") or "model" - - -def rate_label(rate: float) -> str: - label = f"{rate:.3f}".rstrip("0").rstrip(".") - return label if label else "0" - - -def parse_model_spec(raw: str) -> ModelSpec: - if "=" in raw: - name, dims = raw.split("=", 1) - name = clean_name(name) - else: - dims = raw - name = "" - parts = dims.lower().replace(",", "x").split("x") - if len(parts) != 3: - raise argparse.ArgumentTypeError( - "model specs must look like 8x8x256 or name=8x8x256" - ) - try: - n_layer, n_head, n_embd = [int(part) for part in parts] - except ValueError as exc: - raise argparse.ArgumentTypeError("model dimensions must be integers") from exc - if n_layer <= 0 or n_head <= 0 or n_embd <= 0: - raise argparse.ArgumentTypeError("model dimensions must be positive") - if n_embd % n_head != 0: - raise argparse.ArgumentTypeError("n_embd must be divisible by n_head") - if (n_embd // n_head) % 2 != 0: - raise argparse.ArgumentTypeError("n_embd / n_head must be even for rotary attention") - return ModelSpec( - name or f"L{n_layer}_H{n_head}_D{n_embd}", - n_layer, - n_head, - n_embd, - ) - - -def parse_decay_spec(raw: str) -> DropoutCondition: - parts = raw.split(":") - if len(parts) not in {3, 4, 5}: - raise argparse.ArgumentTypeError( - "decay specs must look like " - "name:initial:final[:cosine|smoothstep|linear[:decay_tokens]]" - ) - name = clean_name(parts[0]) - try: - initial = float(parts[1]) - final = float(parts[2]) - decay_tokens = int(parts[4]) if len(parts) == 5 else None - except ValueError as exc: - raise argparse.ArgumentTypeError( - "decay dropout values and decay_tokens must be numeric" - ) from exc - schedule = parts[3] if len(parts) >= 4 else "cosine" - if schedule not in {"cosine", "smoothstep", "linear"}: - raise argparse.ArgumentTypeError( - "decay schedule must be cosine, smoothstep, or linear" - ) - return DropoutCondition( - name=name, - kind="decay", - initial=initial, - final=final, - schedule=schedule, - decay_tokens=decay_tokens, - ) - - -def parse_anchor_decay_spec(raw: str) -> DropoutCondition: - if ":" not in raw: - raise argparse.ArgumentTypeError( - "anchor decay specs must look like name:250000=0.60,500000=0.40" - ) - name, raw_anchors = raw.split(":", 1) - anchors: list[tuple[int, float]] = [] - for piece in raw_anchors.split(","): - if "=" not in piece: - raise argparse.ArgumentTypeError( - "anchor decay anchors must look like token_count=dropout" - ) - raw_tokens, raw_dropout = piece.split("=", 1) - try: - token_count = int(raw_tokens) - dropout = float(raw_dropout) - except ValueError as exc: - raise argparse.ArgumentTypeError( - "anchor token counts must be integers and dropout values numeric" - ) from exc - if token_count <= 0: - raise argparse.ArgumentTypeError("anchor token counts must be positive") - if not 0.0 <= dropout < 1.0: - raise argparse.ArgumentTypeError("anchor dropout must satisfy 0 <= p < 1") - anchors.append((token_count, dropout)) - anchors = sorted(set(anchors)) - if len(anchors) < 2: - raise argparse.ArgumentTypeError("provide at least two dropout anchors") - for left, right in zip(anchors, anchors[1:]): - if right[1] > left[1]: - raise argparse.ArgumentTypeError( - "anchor dropout values must be non-increasing as token counts grow" - ) - return DropoutCondition( - name=clean_name(name), - kind="anchor_decay", - initial=anchors[0][1], - final=anchors[-1][1], - schedule="log_prefix_anchor", - anchors=tuple(anchors), - ) - - -def default_seeds(mode: str, seeds: list[int] | None) -> list[int]: - if seeds: - return seeds - return [1] if mode == "screen_static" else [1, 2, 3] - - -def set_seed(seed: int) -> None: - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - - -def batch_seed(seed: int, model: ModelSpec, dropout_code: int, stage: int | None) -> int: - return ( - seed * 1_000_003 - + model.n_layer * 100_003 - + model.n_head * 10_007 - + model.n_embd * 101 - + dropout_code * 37 - + (stage or 0) * 997 - ) - - -def make_batch( - tokens: np.ndarray, - token_limit: int, - batch_size: int, - block_size: int, - rng: np.random.Generator, - device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor]: - limit = min(token_limit, len(tokens)) - max_start = limit - block_size - 1 - if max_start <= 0: - raise ValueError("token_limit is too small for the requested block_size") - starts = rng.integers(0, max_start, size=batch_size) - x_np = np.stack([tokens[start : start + block_size] for start in starts]).astype( - np.int64 - ) - y_np = np.stack( - [tokens[start + 1 : start + 1 + block_size] for start in starts] - ).astype(np.int64) - return ( - torch.tensor(x_np, dtype=torch.long, device=device), - torch.tensor(y_np, dtype=torch.long, device=device), - ) - - -@torch.no_grad() -def estimate_loss( - model: DropoutGPT, - tokens: np.ndarray, - token_limit: int, - batches: int, - args: argparse.Namespace, - device: torch.device, - rng_seed: int, -) -> float: - if batches <= 0: - return float("nan") - model.eval() - rng = np.random.default_rng(rng_seed) - losses: list[float] = [] - for _ in range(batches): - x, y = make_batch(tokens, token_limit, args.batch_size, args.block_size, rng, device) - _, loss = model(x, y) - losses.append(float(loss.item())) - model.train() - return float(statistics.fmean(losses)) - - -def write_jsonl_row(handle, row: dict) -> None: - handle.write(json.dumps(row, sort_keys=True) + "\n") - handle.flush() - - -def metric_key(row: dict) -> tuple: - return ( - row["run_mode"], - row["condition"], - row["condition_kind"], - row.get("stage"), - int(row["token_limit"]), - row["model_name"], - int(row["n_layer"]), - int(row["n_head"]), - int(row["n_embd"]), - int(row["seed"]), - float(row["dropout_initial"]), - float(row["dropout_final"]), - row["dropout_schedule"], - ) - - -def planned_metric_key( - *, - mode: str, - condition: DropoutCondition, - model_spec: ModelSpec, - seed: int, - token_limit: int, - stage: int | None = None, -) -> tuple: - return ( - mode, - condition.name, - condition.kind, - stage, - int(token_limit), - model_spec.name, - int(model_spec.n_layer), - int(model_spec.n_head), - int(model_spec.n_embd), - int(seed), - float(condition.initial), - float(condition.final), - condition.schedule, - ) - - -def load_metrics(path: Path) -> list[dict]: - if not path.exists(): - return [] - rows = [] - for line in path.read_text(encoding="utf-8").splitlines(): - if line.strip(): - rows.append(json.loads(line)) - return rows - - -def train_segment( - *, - run_mode: str, - condition: DropoutCondition, - model_spec: ModelSpec, - config: GPTConfig, - train_tokens: np.ndarray, - val_tokens: np.ndarray, - token_limit: int, - steps: int, - seed: int, - args: argparse.Namespace, - device: torch.device, - dropout_fn: Callable[[int], float], - metrics_file, - trace_file, - stage: int | None = None, - model: DropoutGPT | None = None, - optimizer: torch.optim.Optimizer | None = None, - tokens_seen_start: int = 0, -) -> tuple[DropoutGPT, torch.optim.Optimizer, int, dict]: - if model is None: - set_seed(seed) - model = DropoutGPT(config).to(device) - optimizer = torch.optim.AdamW( - model.parameters(), - lr=args.lr, - betas=(0.9, 0.95), - weight_decay=args.weight_decay, - ) - else: - torch.manual_seed(seed + 10_000 + (stage or 0)) - if optimizer is None: - raise ValueError("optimizer is required when reusing a model") - - model.train() - dropout_code = int(round(condition.initial * 10_000)) - rng = np.random.default_rng(batch_seed(seed, model_spec, dropout_code, stage)) - tokens_seen = tokens_seen_start - last_loss = float("nan") - active_dropout = condition.initial - t0 = time.time() - - for step in range(1, steps + 1): - active_dropout = dropout_fn(tokens_seen) - model.set_dropout(active_dropout) - x, y = make_batch(train_tokens, token_limit, args.batch_size, args.block_size, rng, device) - _, loss = model(x, y) - optimizer.zero_grad(set_to_none=True) - loss.backward() - if args.grad_clip > 0: - torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) - optimizer.step() - - tokens_seen += args.batch_size * args.block_size - last_loss = float(loss.item()) - if args.log_every > 0 and step % args.log_every == 0: - write_jsonl_row( - trace_file, - { - "event": "train_step", - "run_mode": run_mode, - "condition": condition.name, - "model_name": model_spec.name, - "seed": seed, - "stage": stage, - "step": step, - "steps": steps, - "token_limit": int(token_limit), - "tokens_seen": int(tokens_seen), - "dropout": float(active_dropout), - "train_batch_loss": last_loss, - }, - ) - if args.eval_every > 0 and step % args.eval_every == 0: - train_eval = estimate_loss( - model, - train_tokens, - token_limit, - args.trace_eval_batches, - args, - device, - rng_seed=seed + 20_000 + step, - ) - val_eval = estimate_loss( - model, - val_tokens, - len(val_tokens), - args.trace_eval_batches, - args, - device, - rng_seed=seed + 30_000 + step, - ) - write_jsonl_row( - trace_file, - { - "event": "eval_step", - "run_mode": run_mode, - "condition": condition.name, - "model_name": model_spec.name, - "seed": seed, - "stage": stage, - "step": step, - "steps": steps, - "token_limit": int(token_limit), - "tokens_seen": int(tokens_seen), - "dropout": float(active_dropout), - "train_eval_loss": train_eval, - "val_eval_loss": val_eval, - "generalization_gap": val_eval - train_eval, - }, - ) - - train_eval = estimate_loss( - model, - train_tokens, - token_limit, - args.train_eval_batches, - args, - device, - rng_seed=seed + 40_000 + (stage or 0), - ) - val_eval = estimate_loss( - model, - val_tokens, - len(val_tokens), - args.eval_batches, - args, - device, - rng_seed=seed + 50_000 + (stage or 0), - ) - row = { - "run_mode": run_mode, - "condition": condition.name, - "condition_kind": condition.kind, - "seed": seed, - "stage": stage, - "token_limit": int(token_limit), - "steps": int(steps), - "tokens_seen": int(tokens_seen), - "dropout_initial": float(condition.initial), - "dropout_final": float(condition.final), - "dropout_schedule": condition.schedule, - "dropout_active_final": float(active_dropout), - "train_loss_last": last_loss, - "train_eval_loss": train_eval, - "val_eval_loss": val_eval, - "eval_loss": val_eval, - "generalization_gap": val_eval - train_eval, - "elapsed_sec": time.time() - t0, - "parameters": model.num_parameters(), - "model_config": config.to_dict(), - **model_spec.to_dict(), - } - write_jsonl_row(metrics_file, row) - return model, optimizer, tokens_seen, row - - -def summarize(rows: list[dict]) -> list[dict]: - groups: dict[tuple, list[dict]] = defaultdict(list) - for row in rows: - key = ( - row["run_mode"], - row["condition"], - row["condition_kind"], - row["stage"], - row["token_limit"], - row["model_name"], - row["n_layer"], - row["n_head"], - row["n_embd"], - row["parameters"], - row["dropout_initial"], - row["dropout_final"], - row["dropout_schedule"], - ) - groups[key].append(row) - - summary: list[dict] = [] - for group_rows in groups.values(): - first = group_rows[0] - item = {field: first[field] for field in SUMMARY_FIELDS if field in first} - item["n"] = len(group_rows) - for source, mean_key, std_key in [ - ("train_eval_loss", "mean_train_eval_loss", "std_train_eval_loss"), - ("val_eval_loss", "mean_val_eval_loss", "std_val_eval_loss"), - ("generalization_gap", "mean_generalization_gap", "std_generalization_gap"), - ]: - values = [float(row[source]) for row in group_rows] - item[mean_key] = statistics.fmean(values) - item[std_key] = statistics.stdev(values) if len(values) > 1 else 0.0 - summary.append(item) - return sorted( - summary, - key=lambda row: ( - row["run_mode"], - row["model_name"], - row["token_limit"], - row["condition"], - row["stage"] or -1, - ), - ) - - -def write_csv(path: Path, rows: list[dict], fieldnames: list[str]) -> None: - with path.open("w", newline="", encoding="utf-8") as handle: - writer = csv.DictWriter(handle, fieldnames=fieldnames, extrasaction="ignore") - writer.writeheader() - writer.writerows(rows) - - -def format_duration(seconds: float) -> str: - seconds = max(0, int(seconds)) - hours, remainder = divmod(seconds, 3600) - minutes, secs = divmod(remainder, 60) - if hours: - return f"{hours}h{minutes:02d}m" - if minutes: - return f"{minutes}m{secs:02d}s" - return f"{secs}s" - - -class ProgressMeter: - def __init__(self, total: int): - self.total = max(0, total) - self.done = 0 - self.started_at = time.time() - - def mark_done(self, row: dict) -> None: - self.done += 1 - elapsed = time.time() - self.started_at - mean = elapsed / self.done if self.done else 0.0 - remaining = max(0, self.total - self.done) - eta = remaining * mean - stage = row.get("stage") - stage_text = f"stage={stage} " if stage is not None else "" - print( - "progress " - f"{self.done}/{self.total} " - f"eta={format_duration(eta)} " - f"mode={row['run_mode']} " - f"model={row['model_name']} " - f"params={int(row['parameters']):,} " - f"prefix={int(row['token_limit']):,} " - f"{stage_text}" - f"seed={row['seed']} " - f"condition={row['condition']} " - f"dropout={float(row['dropout_active_final']):.3f} " - f"val={row['val_eval_loss']:.4f} " - f"train={row['train_eval_loss']:.4f} " - f"gap={row['generalization_gap']:.4f} " - f"elapsed={format_duration(float(row['elapsed_sec']))}", - flush=True, - ) - - -def build_model_selection(summary: list[dict], args: argparse.Namespace) -> list[dict]: - groups: dict[tuple, list[dict]] = defaultdict(list) - for row in summary: - if row["condition_kind"] == "static" and row["stage"] is None: - groups[(row["run_mode"], row["token_limit"], row["model_name"])].append(row) - - selection: list[dict] = [] - for (run_mode, token_limit, model_name), rows in groups.items(): - curve = sorted(rows, key=lambda row: row["dropout_initial"]) - best = min(curve, key=lambda row: row["mean_val_eval_loss"]) - plateau = [ - row for row in curve - if row["mean_val_eval_loss"] <= best["mean_val_eval_loss"] + args.plateau_delta - ] - zero = next((row for row in curve if row["dropout_initial"] == 0.0), None) - nonzero = [row for row in curve if row["dropout_initial"] > 0.0] - best_nonzero = min(nonzero, key=lambda row: row["mean_val_eval_loss"]) if nonzero else None - max_dropout = max(curve, key=lambda row: row["dropout_initial"]) - - zero_loss = zero["mean_val_eval_loss"] if zero else None - zero_minus_best = zero_loss - best["mean_val_eval_loss"] if zero_loss is not None else None - zero_minus_best_nonzero = ( - zero_loss - best_nonzero["mean_val_eval_loss"] - if zero_loss is not None and best_nonzero is not None - else None - ) - max_minus_best = max_dropout["mean_val_eval_loss"] - best["mean_val_eval_loss"] - has_nonzero_optimum = ( - best["dropout_initial"] > 0.0 - and zero_minus_best is not None - and zero_minus_best >= args.min_nonzero_margin - and max_minus_best >= args.min_high_dropout_margin - ) - curve_json = json.dumps( - [ - { - "dropout": row["dropout_initial"], - "mean_val_loss": row["mean_val_eval_loss"], - "std_val_loss": row["std_val_eval_loss"], - "mean_train_loss": row["mean_train_eval_loss"], - "mean_generalization_gap": row["mean_generalization_gap"], - "n": row["n"], - } - for row in curve - ], - sort_keys=True, - ) - selection.append( - { - "run_mode": run_mode, - "token_limit": token_limit, - "model_name": model_name, - "n_layer": best["n_layer"], - "n_head": best["n_head"], - "n_embd": best["n_embd"], - "parameters": best["parameters"], - "n": best["n"], - "best_dropout": best["dropout_initial"], - "best_val_loss": best["mean_val_eval_loss"], - "best_val_std": best["std_val_eval_loss"], - "plateau_start_dropout": min(row["dropout_initial"] for row in plateau), - "plateau_end_dropout": max(row["dropout_initial"] for row in plateau), - "plateau_delta": args.plateau_delta, - "zero_dropout_val_loss": zero_loss, - "zero_minus_best": zero_minus_best, - "best_nonzero_dropout": ( - best_nonzero["dropout_initial"] if best_nonzero else None - ), - "best_nonzero_val_loss": ( - best_nonzero["mean_val_eval_loss"] if best_nonzero else None - ), - "zero_minus_best_nonzero": zero_minus_best_nonzero, - "max_dropout": max_dropout["dropout_initial"], - "max_dropout_val_loss": max_dropout["mean_val_eval_loss"], - "max_dropout_minus_best": max_minus_best, - "has_nonzero_optimum": has_nonzero_optimum, - "meets_target_dropout": best["dropout_initial"] >= args.target_min_dropout, - "curve_json": curve_json, - } - ) - return sorted( - selection, - key=lambda row: ( - not row["has_nonzero_optimum"], - not row["meets_target_dropout"], - row["best_val_loss"], - ), - ) - - -def write_screen_markdown_summary(output_dir: Path, rows: list[dict]) -> None: - if not rows: - return - static_rows = [ - row - for row in rows - if row["run_mode"] == "screen_static" and row["condition_kind"] == "static" - ] - if not static_rows: - return - - by_model_prefix_rate: dict[tuple[str, int, float], list[dict]] = defaultdict(list) - for row in rows: - if row["run_mode"] == "screen_static" and row["condition_kind"] == "static": - by_model_prefix_rate[ - ( - row["model_name"], - int(row["token_limit"]), - float(row["dropout_initial"]), - ) - ].append(row) - - aggregates: list[dict] = [] - for (model_name, prefix, dropout), group_rows in by_model_prefix_rate.items(): - first = group_rows[0] - val_losses = [float(row["val_eval_loss"]) for row in group_rows] - train_losses = [float(row["train_eval_loss"]) for row in group_rows] - gaps = [float(row["generalization_gap"]) for row in group_rows] - aggregates.append( - { - "model_name": model_name, - "token_limit": prefix, - "dropout_initial": dropout, - "n": len(group_rows), - "mean_val_eval_loss": statistics.fmean(val_losses), - "std_val_eval_loss": statistics.stdev(val_losses) - if len(val_losses) > 1 - else 0.0, - "mean_train_eval_loss": statistics.fmean(train_losses), - "std_train_eval_loss": statistics.stdev(train_losses) - if len(train_losses) > 1 - else 0.0, - "mean_generalization_gap": statistics.fmean(gaps), - "std_generalization_gap": statistics.stdev(gaps) - if len(gaps) > 1 - else 0.0, - "parameters": int(first["parameters"]), - "n_layer": int(first["n_layer"]), - "n_head": int(first["n_head"]), - "n_embd": int(first["n_embd"]), - "block_size": int(first["model_config"]["block_size"]), - "vocab_size": int(first["model_config"]["vocab_size"]), - "tokens_seen": int(first["tokens_seen"]), - "seeds": sorted({int(row["seed"]) for row in group_rows}), - } - ) - - by_model: dict[str, list[dict]] = defaultdict(list) - for row in aggregates: - by_model[row["model_name"]].append(row) - - model_rows = [] - for model_name, model_group in by_model.items(): - first = model_group[0] - seeds = sorted({seed for row in model_group for seed in row["seeds"]}) - model_rows.append( - { - "model_name": model_name, - "parameters": first["parameters"], - "n_layer": first["n_layer"], - "n_head": first["n_head"], - "n_embd": first["n_embd"], - "block_size": first["block_size"], - "vocab_size": first["vocab_size"], - "seeds": seeds, - } - ) - - lines = [ - "# Static Dropout Screen Summary", - "", - f"Run directory: `{output_dir}`", - "", - "## Models", - "", - "| Model | Params | Layers | Heads | Embedding | Block | Vocab | Seeds |", - "|---|---:|---:|---:|---:|---:|---:|---|", - ] - for model in sorted(model_rows, key=lambda item: item["parameters"]): - lines.append( - "| " - f"`{model['model_name']}` | {model['parameters']:,} | " - f"{model['n_layer']} | {model['n_head']} | {model['n_embd']} | " - f"{model['block_size']} | {model['vocab_size']} | " - f"{', '.join(str(seed) for seed in model['seeds'])} |" - ) - lines.extend( - [ - "", - "## Best Dropout By Model And Prefix", - "", - "| Model | Prefix tokens | Effective epochs | Best dropout | Mean val loss | Val std | Mean train loss | Mean gap | Plateau/bracket note |", - "|---|---:|---:|---:|---:|---:|---:|---:|---|", - ] - ) - for model_name, model_group in sorted(by_model.items()): - by_prefix: dict[int, list[dict]] = defaultdict(list) - for row in model_group: - by_prefix[int(row["token_limit"])].append(row) - for prefix, prefix_rows in sorted(by_prefix.items()): - best = min(prefix_rows, key=lambda row: row["mean_val_eval_loss"]) - rates = [float(row["dropout_initial"]) for row in prefix_rows] - eff_epochs = float(best["tokens_seen"]) / prefix - if best["dropout_initial"] == max(rates): - note = "not bracketed; best at top of tested grid" - elif best["dropout_initial"] == min(rates): - note = "not bracketed; best at bottom of tested grid" - else: - note = "bracketed by tested grid" - lines.append( - "| " - f"`{model_name}` | {prefix:,} | {eff_epochs:.2f} | " - f"{best['dropout_initial']:.2f} | " - f"{best['mean_val_eval_loss']:.4f} | " - f"{best['std_val_eval_loss']:.4f} | " - f"{best['mean_train_eval_loss']:.4f} | " - f"{best['mean_generalization_gap']:.4f} | {note} |" - ) - - for model_name, model_group in sorted(by_model.items()): - by_prefix = defaultdict(list) - for row in model_group: - by_prefix[int(row["token_limit"])].append(row) - lines.extend( - [ - "", - f"## Model `{model_name}`", - ] - ) - for prefix, prefix_rows in sorted(by_prefix.items()): - eff_epochs = float(prefix_rows[0]["tokens_seen"]) / prefix - lines.extend( - [ - "", - f"### Prefix {prefix:,} Tokens ({eff_epochs:.2f} Effective Epochs)", - "", - "| Dropout | N | Mean val loss | Val std | Mean train loss | Train std | Mean gap | Gap std | Sampled tokens | Params |", - "|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|", - ] - ) - for row in sorted(prefix_rows, key=lambda item: item["dropout_initial"]): - lines.append( - "| " - f"{row['dropout_initial']:.2f} | {row['n']} | " - f"{row['mean_val_eval_loss']:.4f} | " - f"{row['std_val_eval_loss']:.4f} | " - f"{row['mean_train_eval_loss']:.4f} | " - f"{row['std_train_eval_loss']:.4f} | " - f"{row['mean_generalization_gap']:.4f} | " - f"{row['std_generalization_gap']:.4f} | " - f"{int(row['tokens_seen']):,} | {int(row['parameters']):,} |" - ) - - output = "\n".join(lines) + "\n" - (output_dir / "RESULT_SUMMARY.md").write_text(output, encoding="utf-8") - - -def svg_escape(value: object) -> str: - return ( - str(value) - .replace("&", "&") - .replace("<", "<") - .replace(">", ">") - .replace('"', """) - ) - - -def write_dropout_curve_svg(output_dir: Path, summary: list[dict]) -> None: - rows = [ - row - for row in summary - if row["run_mode"] == "screen_static" and row["condition_kind"] == "static" - ] - if not rows: - return - - grouped: dict[tuple[str, int], list[dict]] = defaultdict(list) - model_params: dict[str, int] = {} - for row in rows: - model_name = row["model_name"] - grouped[(model_name, int(row["token_limit"]))].append(row) - model_params[model_name] = int(row["parameters"]) - - models = sorted(model_params, key=lambda name: model_params[name]) - prefixes = sorted({int(row["token_limit"]) for row in rows}) - panel_w, panel_h = 230, 170 - margin_l, margin_t = 58, 34 - plot_w, plot_h = 142, 94 - gap_x, gap_y = 18, 38 - width = margin_l + len(prefixes) * panel_w + gap_x - height = 70 + len(models) * (panel_h + gap_y) - colors = ["#1f77b4", "#d62728", "#2ca02c", "#9467bd", "#ff7f0e"] - - parts = [ - f'', - "", - '', - 'Static dropout law: validation loss vs dropout', - 'Each panel uses its own y-scale. Points are one-seed means unless N > 1.', - ] - - for col, prefix in enumerate(prefixes): - x = margin_l + col * panel_w + plot_w / 2 - parts.append( - f'{prefix:,} prefix tokens' - ) - - for row_idx, model_name in enumerate(models): - row_y = 92 + row_idx * (panel_h + gap_y) - parts.append( - f'' - f'{svg_escape(model_name)} ({model_params[model_name] / 1_000_000:.1f}M)' - ) - for col, prefix in enumerate(prefixes): - panel_x = margin_l + col * panel_w - panel_y = row_y - curve = sorted( - grouped.get((model_name, prefix), []), - key=lambda item: float(item["dropout_initial"]), - ) - if not curve: - continue - losses = [float(item["mean_val_eval_loss"]) for item in curve] - min_loss, max_loss = min(losses), max(losses) - pad = max(0.02, (max_loss - min_loss) * 0.08) - y_min, y_max = min_loss - pad, max_loss + pad - best = min(curve, key=lambda item: float(item["mean_val_eval_loss"])) - - def px(dropout: float) -> float: - return panel_x + (dropout / 0.9) * plot_w - - def py(loss: float) -> float: - scale = (loss - y_min) / (y_max - y_min) - return panel_y + plot_h - scale * plot_h - - parts.extend( - [ - f'', - f'', - f'', - f'{y_max:.2f}', - f'{y_min:.2f}', - f'0', - f'0.9', - ] - ) - points = " ".join( - f"{px(float(item['dropout_initial'])):.1f},{py(float(item['mean_val_eval_loss'])):.1f}" - for item in curve - ) - color = colors[row_idx % len(colors)] - parts.append(f'') - for item in curve: - dropout = float(item["dropout_initial"]) - loss = float(item["mean_val_eval_loss"]) - radius = 4 if item is best else 2.7 - fill = "#111827" if item is best else "#ffffff" - parts.append( - f'' - ) - parts.append( - f'' - f'best p={float(best["dropout_initial"]):.2f}' - ) - parts.append( - f'' - f'loss={float(best["mean_val_eval_loss"]):.3f}' - ) - - parts.append("") - (output_dir / "dropout_curves.svg").write_text("\n".join(parts), encoding="utf-8") - - -def write_stream_markdown_summary(output_dir: Path, rows: list[dict]) -> None: - stream_rows = [row for row in rows if row["run_mode"] == "locked_stream"] - if not stream_rows: - return - - by_condition_stage: dict[tuple[str, int], list[dict]] = defaultdict(list) - by_condition: dict[str, list[dict]] = defaultdict(list) - for row in stream_rows: - condition = row["condition"] - by_condition_stage[(condition, int(row["stage"]))].append(row) - by_condition[condition].append(row) - - first = stream_rows[0] - seeds = sorted({int(row["seed"]) for row in stream_rows}) - conditions = sorted( - by_condition, - key=lambda name: ( - by_condition[name][0]["condition_kind"] != "anchor_decay", - by_condition[name][0]["dropout_initial"], - name, - ), - ) - stages = sorted({int(row["stage"]) for row in stream_rows}) - - lines = [ - "# Locked Streaming Dropout Summary", - "", - f"Run directory: `{output_dir}`", - "", - ( - f"Model: `{first['model_name']}` causal Transformer, " - f"{int(first['parameters']):,} parameters, {first['n_layer']} layers, " - f"{first['n_head']} heads, {first['n_embd']} embedding dim." - ), - ( - f"Training per stage: {first['steps']:,} steps. " - "Sampled tokens are cumulative in each stage row. " - f"Seeds present: {', '.join(str(seed) for seed in seeds)}." - ), - "", - "## Condition Ranking", - "", - "| Condition | Kind | Final dropout | Mean trajectory val loss | Final val loss | Final gap | Dropout path |", - "|---|---|---:|---:|---:|---:|---|", - ] - - ranking = [] - for condition in conditions: - stage_items = [] - for stage in stages: - group = by_condition_stage.get((condition, stage), []) - if not group: - continue - stage_items.append( - { - "stage": stage, - "token_limit": int(group[0]["token_limit"]), - "mean_val": statistics.fmean( - float(row["val_eval_loss"]) for row in group - ), - "mean_gap": statistics.fmean( - float(row["generalization_gap"]) for row in group - ), - "mean_dropout": statistics.fmean( - float(row["dropout_active_final"]) for row in group - ), - "kind": group[0]["condition_kind"], - } - ) - if not stage_items: - continue - final = max(stage_items, key=lambda item: item["stage"]) - ranking.append( - { - "condition": condition, - "kind": stage_items[0]["kind"], - "trajectory_val": statistics.fmean(item["mean_val"] for item in stage_items), - "final_val": final["mean_val"], - "final_gap": final["mean_gap"], - "final_dropout": final["mean_dropout"], - "dropout_path": " -> ".join( - f"{item['mean_dropout']:.2f}" for item in stage_items - ), - } - ) - - for item in sorted(ranking, key=lambda row: row["trajectory_val"]): - lines.append( - "| " - f"`{item['condition']}` | {item['kind']} | " - f"{item['final_dropout']:.2f} | {item['trajectory_val']:.4f} | " - f"{item['final_val']:.4f} | {item['final_gap']:.4f} | " - f"{item['dropout_path']} |" - ) - - lines.extend(["", "## Stage Trajectory", ""]) - for stage in stages: - stage_groups = { - condition: by_condition_stage[(condition, stage)] - for condition in conditions - if (condition, stage) in by_condition_stage - } - if not stage_groups: - continue - prefix = int(next(iter(stage_groups.values()))[0]["token_limit"]) - lines.extend( - [ - f"### Stage {stage}: {prefix:,} Prefix Tokens", - "", - "| Condition | Dropout | Mean val loss | Mean train loss | Mean gap | N |", - "|---|---:|---:|---:|---:|---:|", - ] - ) - for condition, group in sorted( - stage_groups.items(), - key=lambda item: statistics.fmean( - float(row["val_eval_loss"]) for row in item[1] - ), - ): - val = statistics.fmean(float(row["val_eval_loss"]) for row in group) - train = statistics.fmean(float(row["train_eval_loss"]) for row in group) - gap = statistics.fmean(float(row["generalization_gap"]) for row in group) - dropout = statistics.fmean( - float(row["dropout_active_final"]) for row in group - ) - lines.append( - "| " - f"`{condition}` | {dropout:.2f} | {val:.4f} | " - f"{train:.4f} | {gap:.4f} | {len(group)} |" - ) - lines.append("") - - (output_dir / "RESULT_SUMMARY.md").write_text( - "\n".join(lines).rstrip() + "\n", - encoding="utf-8", - ) - - -def static_conditions(dropout_rates: list[float]) -> list[DropoutCondition]: - return [ - DropoutCondition( - name=f"static_dropout_{rate_label(rate)}", - kind="static", - initial=rate, - final=rate, - ) - for rate in dropout_rates - ] - - -def run_fixed_static_sweep( - *, - args: argparse.Namespace, - model_specs: list[ModelSpec], - seeds: list[int], - train_tokens: np.ndarray, - val_tokens: np.ndarray, - tokenizer_vocab_size: int, - token_limits: list[int], - device: torch.device, - metrics_file, - trace_file, - completed_keys: set[tuple] | None = None, -) -> list[dict]: - rows: list[dict] = [] - completed_keys = completed_keys or set() - conditions = static_conditions(sorted(set(args.dropout_rates))) - planned = 0 - for token_limit in token_limits: - for model_spec in model_specs: - for condition in conditions: - for seed in seeds: - key = planned_metric_key( - mode=args.mode, - condition=condition, - model_spec=model_spec, - seed=seed, - token_limit=token_limit, - ) - if key not in completed_keys: - planned += 1 - progress = ProgressMeter(planned) - for token_limit in token_limits: - for model_spec in model_specs: - best_val_loss = float("inf") - worse_streak = 0 - for condition in conditions: - condition_rows: list[dict] = [] - for seed in seeds: - key = planned_metric_key( - mode=args.mode, - condition=condition, - model_spec=model_spec, - seed=seed, - token_limit=token_limit, - ) - if key in completed_keys: - write_jsonl_row( - trace_file, - { - "event": "skipped_completed_condition", - "run_mode": args.mode, - "condition": condition.name, - "model_name": model_spec.name, - "seed": seed, - "stage": None, - "token_limit": int(token_limit), - }, - ) - continue - config = model_spec.config( - tokenizer_vocab_size, - args.block_size, - condition.initial, - ) - model, optimizer, _, row = train_segment( - run_mode=args.mode, - condition=condition, - model_spec=model_spec, - config=config, - train_tokens=train_tokens, - val_tokens=val_tokens, - token_limit=token_limit, - steps=args.steps, - seed=seed, - args=args, - device=device, - dropout_fn=condition.make_fn( - args.steps * args.batch_size * args.block_size - ), - metrics_file=metrics_file, - trace_file=trace_file, - ) - rows.append(row) - condition_rows.append(row) - completed_keys.add(metric_key(row)) - progress.mark_done(row) - del model, optimizer - torch.mps.empty_cache() - - if not condition_rows: - continue - mean_val_loss = statistics.fmean( - float(row["val_eval_loss"]) for row in condition_rows - ) - if mean_val_loss < best_val_loss - args.screen_prune_min_delta: - best_val_loss = mean_val_loss - worse_streak = 0 - elif mean_val_loss > best_val_loss + args.screen_prune_min_delta: - worse_streak += 1 - - if ( - args.mode == "screen_static" - and args.screen_early_stop - and worse_streak >= args.screen_prune_patience - and condition.initial >= args.target_min_dropout - ): - write_jsonl_row( - trace_file, - { - "event": "screen_pruned_model", - "run_mode": args.mode, - "model_name": model_spec.name, - "token_limit": int(token_limit), - "best_val_loss": best_val_loss, - "pruned_after_dropout": condition.initial, - "worse_streak": worse_streak, - "remaining_dropouts": [ - rate - for rate in args.dropout_rates - if rate > condition.initial - ], - }, - ) - break - return rows - - -def run_locked_stream( - *, - args: argparse.Namespace, - model_specs: list[ModelSpec], - seeds: list[int], - train_tokens: np.ndarray, - val_tokens: np.ndarray, - tokenizer_vocab_size: int, - stream_caps: list[int], - device: torch.device, - metrics_file, - trace_file, - completed_keys: set[tuple] | None = None, -) -> list[dict]: - rows: list[dict] = [] - completed_keys = completed_keys or set() - conditions = args.anchor_decays + static_conditions(args.dropout_rates) + args.decays - fallback_decay_tokens = ( - args.decay_tokens - or args.stage_steps * args.batch_size * args.block_size * len(stream_caps) - ) - planned = 0 - for model_spec in model_specs: - for condition in conditions: - for seed in seeds: - for stage, token_limit in enumerate(stream_caps): - key = planned_metric_key( - mode=args.mode, - condition=condition, - model_spec=model_spec, - seed=seed, - token_limit=token_limit, - stage=stage, - ) - if key not in completed_keys: - planned += 1 - progress = ProgressMeter(planned) - for model_spec in model_specs: - for condition in conditions: - for seed in seeds: - model = None - optimizer = None - tokens_seen = 0 - for stage, token_limit in enumerate(stream_caps): - key = planned_metric_key( - mode=args.mode, - condition=condition, - model_spec=model_spec, - seed=seed, - token_limit=token_limit, - stage=stage, - ) - if key in completed_keys: - write_jsonl_row( - trace_file, - { - "event": "skipped_completed_condition", - "run_mode": args.mode, - "condition": condition.name, - "model_name": model_spec.name, - "seed": seed, - "stage": stage, - "token_limit": int(token_limit), - }, - ) - continue - config = model_spec.config( - tokenizer_vocab_size, - args.block_size, - condition.initial, - ) - model, optimizer, tokens_seen, row = train_segment( - run_mode=args.mode, - condition=condition, - model_spec=model_spec, - config=config, - train_tokens=train_tokens, - val_tokens=val_tokens, - token_limit=token_limit, - steps=args.stage_steps, - seed=seed, - args=args, - device=device, - dropout_fn=condition.make_fn( - fallback_decay_tokens, - unique_tokens=token_limit, - ), - metrics_file=metrics_file, - trace_file=trace_file, - stage=stage, - model=model, - optimizer=optimizer, - tokens_seen_start=tokens_seen, - ) - rows.append(row) - completed_keys.add(metric_key(row)) - progress.mark_done(row) - del model, optimizer - torch.mps.empty_cache() - return rows - - -def prepare_data(args: argparse.Namespace, output_dir: Path, required_train_tokens: int): - cache_dir = Path(args.cache_dir) if args.cache_dir else output_dir / "cache" - cache_dir.mkdir(parents=True, exist_ok=True) - if args.use_cached_data: - if args.force_retokenize: - raise ValueError("--use-cached-data cannot be combined with --force-retokenize") - return load_cached_splits( - cache_dir=cache_dir, - vocab_size=args.vocab_size, - max_required_train_tokens=required_train_tokens, - val_tokens=args.val_tokens, - allow_short_corpus=args.allow_short_corpus, - ) - - paths = resolve_paths(args.corpus, args.corpus_glob) - tokenizer = train_or_load_tokenizer( - paths=paths, - output_dir=cache_dir, - vocab_size=args.vocab_size, - tokenizer_train_chars=args.tokenizer_train_chars, - text_column=args.text_column, - force_retrain=args.force_retokenize, - ) - splits = encode_corpus( - paths=paths, - tokenizer=tokenizer, - output_dir=cache_dir, - max_required_train_tokens=required_train_tokens, - val_tokens=args.val_tokens, - text_column=args.text_column, - allow_short_corpus=args.allow_short_corpus, - force_reencode=args.force_retokenize, - ) - return tokenizer, splits - - -def run(args: argparse.Namespace) -> Path: - device = assert_mps_only() - seeds = default_seeds(args.mode, args.seeds) - model_specs = [parse_model_spec(spec) for spec in args.models] - if args.mode != "locked_stream" and (args.decays or args.anchor_decays): - raise ValueError("--decays and --anchor-decays are only used with --mode locked_stream") - if args.resume_from and args.mode == "locked_stream": - raise ValueError("--resume-from currently supports fixed static sweeps only") - - if args.resume_from: - output_dir = Path(args.resume_from) - if not output_dir.exists(): - raise FileNotFoundError(f"resume directory does not exist: {output_dir}") - else: - run_id = datetime.now().strftime("%Y%m%d-%H%M%S") - output_dir = Path(args.output_dir) / args.mode / run_id - output_dir.mkdir(parents=True, exist_ok=True) - required_train_tokens = max( - args.stream_token_caps if args.mode == "locked_stream" else args.token_limits - ) - tokenizer, splits = prepare_data(args, output_dir, required_train_tokens) - token_limits = [min(limit, len(splits.train)) for limit in args.token_limits] - stream_caps = [min(limit, len(splits.train)) for limit in args.stream_token_caps] - - args_payload = vars(args).copy() - args_payload["decays"] = [condition.to_dict() for condition in args.decays] - args_payload["anchor_decays"] = [ - condition.to_dict() for condition in args.anchor_decays - ] - config_payload = { - "args": args_payload, - "mode": args.mode, - "seeds": seeds, - "models": [model.to_dict() for model in model_specs], - "device": str(device), - "torch": torch.__version__, - "python": sys.version, - "mps_available": torch.backends.mps.is_available(), - "attribution": NANOCHAT_ATTRIBUTION, - "tokenizer_path": str(splits.tokenizer_path), - "encoded_path": str(splits.encoded_path), - "train_tokens": int(len(splits.train)), - "val_tokens": int(len(splits.val)), - "effective_token_limits": [int(limit) for limit in token_limits], - "effective_stream_token_caps": [int(limit) for limit in stream_caps], - "resume_from": str(args.resume_from) if args.resume_from else None, - } - config_name = "config.resume.json" if args.resume_from else "config.json" - (output_dir / config_name).write_text( - json.dumps(config_payload, indent=2), - encoding="utf-8", - ) - - metrics_path = output_dir / "metrics.jsonl" - trace_path = output_dir / "trace.jsonl" - existing_rows = load_metrics(metrics_path) if args.resume_from else [] - completed_keys = {metric_key(row) for row in existing_rows} - with ( - metrics_path.open("a" if args.resume_from else "w", encoding="utf-8") as metrics_file, - trace_path.open("a" if args.resume_from else "w", encoding="utf-8") as trace_file, - ): - if args.mode in {"screen_static", "confirm_static"}: - new_rows = run_fixed_static_sweep( - args=args, - model_specs=model_specs, - seeds=seeds, - train_tokens=splits.train, - val_tokens=splits.val, - tokenizer_vocab_size=tokenizer.vocab_size, - token_limits=token_limits, - device=device, - metrics_file=metrics_file, - trace_file=trace_file, - completed_keys=completed_keys, - ) - else: - new_rows = run_locked_stream( - args=args, - model_specs=model_specs, - seeds=seeds, - train_tokens=splits.train, - val_tokens=splits.val, - tokenizer_vocab_size=tokenizer.vocab_size, - stream_caps=stream_caps, - device=device, - metrics_file=metrics_file, - trace_file=trace_file, - ) - rows = existing_rows + new_rows - - summary = summarize(rows) - (output_dir / "summary.json").write_text( - json.dumps(summary, indent=2), - encoding="utf-8", - ) - write_csv(output_dir / "summary.csv", summary, SUMMARY_FIELDS) - if args.mode in {"screen_static", "confirm_static"}: - selection = build_model_selection(summary, args) - (output_dir / "model_selection.json").write_text( - json.dumps(selection, indent=2), - encoding="utf-8", - ) - write_csv(output_dir / "model_selection.csv", selection, SELECTION_FIELDS) - write_screen_markdown_summary(output_dir, rows) - write_dropout_curve_svg(output_dir, summary) - elif args.mode == "locked_stream": - write_stream_markdown_summary(output_dir, rows) - - print( - json.dumps( - { - "output_dir": str(output_dir), - "new_rows": len(new_rows), - "total_metric_rows": len(rows), - "summary_rows": len(summary), - }, - indent=2, - ) - ) - return output_dir - - -def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser( - description="MPS-only dropout/model selection experiments" - ) - parser.add_argument( - "--mode", - choices=["screen_static", "confirm_static", "locked_stream"], - default="screen_static", - ) - parser.add_argument("--corpus", default=None, help="Text or parquet corpus path") - parser.add_argument("--corpus-glob", default=None, help="Glob of text/parquet corpus paths") - parser.add_argument("--text-column", default="text", help="Parquet text column") - parser.add_argument( - "--use-cached-data", - action="store_true", - help=( - "Load tokenizer-v{vocab}.json and tokens-v{vocab}-*.npy from --cache-dir " - "instead of requiring the original text/parquet corpus." - ), - ) - parser.add_argument("--output-dir", default="runs") - parser.add_argument( - "--resume-from", - default=None, - help="Existing fixed-static run directory; completed metric rows are skipped", - ) - parser.add_argument("--cache-dir", default=".cache/dropout_decay") - parser.add_argument( - "--models", - nargs="+", - default=["8x8x256"], - help="Model specs like 8x8x256 or name=8x8x256", - ) - parser.add_argument("--seeds", nargs="+", type=int, default=None) - parser.add_argument("--token-limits", nargs="+", type=int, default=[5_000_000]) - parser.add_argument( - "--stream-token-caps", - nargs="+", - type=int, - default=[5_000_000, 10_000_000, 20_000_000, 40_000_000], - ) - parser.add_argument("--val-tokens", type=int, default=500_000) - parser.add_argument("--allow-short-corpus", action="store_true") - parser.add_argument("--force-retokenize", action="store_true") - parser.add_argument("--vocab-size", type=int, default=4096) - parser.add_argument("--tokenizer-train-chars", type=int, default=10_000_000) - parser.add_argument("--block-size", type=int, default=128) - parser.add_argument("--batch-size", type=int, default=16) - parser.add_argument("--steps", type=int, default=2000) - parser.add_argument("--stage-steps", type=int, default=1000) - parser.add_argument("--dropout-rates", nargs="*", type=float, default=DEFAULT_DROPOUT_RATES) - parser.add_argument("--decays", nargs="*", type=parse_decay_spec, default=[]) - parser.add_argument( - "--anchor-decays", - nargs="*", - type=parse_anchor_decay_spec, - default=[], - help=( - "Prefix-token anchor schedules like " - "fit:250000=0.60,500000=0.40,1000000=0.30" - ), - ) - parser.add_argument("--decay-tokens", type=int, default=None) - parser.add_argument("--eval-batches", type=int, default=64) - parser.add_argument("--train-eval-batches", type=int, default=32) - parser.add_argument("--trace-eval-batches", type=int, default=8) - parser.add_argument("--eval-every", type=int, default=0) - parser.add_argument("--log-every", type=int, default=100) - parser.add_argument("--lr", type=float, default=3e-4) - parser.add_argument("--weight-decay", type=float, default=0.1) - parser.add_argument("--grad-clip", type=float, default=1.0) - parser.add_argument("--plateau-delta", type=float, default=0.01) - parser.add_argument("--target-min-dropout", type=float, default=0.10) - parser.add_argument("--min-nonzero-margin", type=float, default=0.01) - parser.add_argument("--min-high-dropout-margin", type=float, default=0.03) - parser.add_argument("--screen-early-stop", action="store_true") - parser.add_argument("--screen-prune-patience", type=int, default=3) - parser.add_argument("--screen-prune-min-delta", type=float, default=0.01) - return parser - - -def main() -> None: - args = build_parser().parse_args() - run(args) if __name__ == "__main__": diff --git a/src/dropout_decay/experiments/__init__.py b/src/dropout_decay/experiments/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1d8b21f8e3c57f4ea8e439cbf22f42da30da8a20 --- /dev/null +++ b/src/dropout_decay/experiments/__init__.py @@ -0,0 +1,25 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from dropout_decay.experiments.runner import build_parser, main, run + +__all__ = [ + "build_parser", + "main", + "run", +] diff --git a/src/dropout_decay/experiments/artifacts.py b/src/dropout_decay/experiments/artifacts.py new file mode 100644 index 0000000000000000000000000000000000000000..9583cca2ac658064608e22db6c50057b34176df4 --- /dev/null +++ b/src/dropout_decay/experiments/artifacts.py @@ -0,0 +1,281 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from __future__ import annotations + +import argparse +import csv +from collections import defaultdict +import json +from pathlib import Path +import statistics + +from dropout_decay.specs import DropoutCondition, ModelSpec + + +SUMMARY_FIELDS = [ + "run_mode", + "condition", + "condition_kind", + "stage", + "token_limit", + "model_name", + "n_layer", + "n_head", + "n_embd", + "parameters", + "dropout_initial", + "dropout_final", + "dropout_schedule", + "n", + "mean_train_eval_loss", + "std_train_eval_loss", + "mean_val_eval_loss", + "std_val_eval_loss", + "mean_generalization_gap", + "std_generalization_gap", +] +SELECTION_FIELDS = [ + "run_mode", + "token_limit", + "model_name", + "n_layer", + "n_head", + "n_embd", + "parameters", + "n", + "best_dropout", + "best_val_loss", + "best_val_std", + "plateau_start_dropout", + "plateau_end_dropout", + "plateau_delta", + "zero_dropout_val_loss", + "zero_minus_best", + "best_nonzero_dropout", + "best_nonzero_val_loss", + "zero_minus_best_nonzero", + "max_dropout", + "max_dropout_val_loss", + "max_dropout_minus_best", + "has_nonzero_optimum", + "meets_target_dropout", + "curve_json", +] + + +def write_jsonl_row(handle, row: dict) -> None: + handle.write(json.dumps(row, sort_keys=True) + "\n") + handle.flush() + + +def metric_key(row: dict) -> tuple: + return ( + row["run_mode"], + row["condition"], + row["condition_kind"], + row.get("stage"), + int(row["token_limit"]), + row["model_name"], + int(row["n_layer"]), + int(row["n_head"]), + int(row["n_embd"]), + int(row["seed"]), + float(row["dropout_initial"]), + float(row["dropout_final"]), + row["dropout_schedule"], + ) + + +def planned_metric_key( + *, + mode: str, + condition: DropoutCondition, + model_spec: ModelSpec, + seed: int, + token_limit: int, + stage: int | None = None, +) -> tuple: + return ( + mode, + condition.name, + condition.kind, + stage, + int(token_limit), + model_spec.name, + int(model_spec.n_layer), + int(model_spec.n_head), + int(model_spec.n_embd), + int(seed), + float(condition.initial), + float(condition.final), + condition.schedule, + ) + + +def load_metrics(path: Path) -> list[dict]: + if not path.exists(): + return [] + rows = [] + for line in path.read_text(encoding="utf-8").splitlines(): + if line.strip(): + rows.append(json.loads(line)) + return rows + + +def summarize(rows: list[dict]) -> list[dict]: + groups: dict[tuple, list[dict]] = defaultdict(list) + for row in rows: + key = ( + row["run_mode"], + row["condition"], + row["condition_kind"], + row["stage"], + row["token_limit"], + row["model_name"], + row["n_layer"], + row["n_head"], + row["n_embd"], + row["parameters"], + row["dropout_initial"], + row["dropout_final"], + row["dropout_schedule"], + ) + groups[key].append(row) + + summary: list[dict] = [] + for group_rows in groups.values(): + first = group_rows[0] + item = {field: first[field] for field in SUMMARY_FIELDS if field in first} + item["n"] = len(group_rows) + for source, mean_key, std_key in [ + ("train_eval_loss", "mean_train_eval_loss", "std_train_eval_loss"), + ("val_eval_loss", "mean_val_eval_loss", "std_val_eval_loss"), + ("generalization_gap", "mean_generalization_gap", "std_generalization_gap"), + ]: + values = [float(row[source]) for row in group_rows] + item[mean_key] = statistics.fmean(values) + item[std_key] = statistics.stdev(values) if len(values) > 1 else 0.0 + summary.append(item) + return sorted( + summary, + key=lambda row: ( + row["run_mode"], + row["model_name"], + row["token_limit"], + row["condition"], + row["stage"] or -1, + ), + ) + + +def write_csv(path: Path, rows: list[dict], fieldnames: list[str]) -> None: + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fieldnames, extrasaction="ignore") + writer.writeheader() + writer.writerows(rows) + + +def build_model_selection(summary: list[dict], args: argparse.Namespace) -> list[dict]: + groups: dict[tuple, list[dict]] = defaultdict(list) + for row in summary: + if row["condition_kind"] == "static" and row["stage"] is None: + groups[(row["run_mode"], row["token_limit"], row["model_name"])].append(row) + + selection: list[dict] = [] + for (run_mode, token_limit, model_name), rows in groups.items(): + curve = sorted(rows, key=lambda row: row["dropout_initial"]) + best = min(curve, key=lambda row: row["mean_val_eval_loss"]) + plateau = [ + row for row in curve + if row["mean_val_eval_loss"] <= best["mean_val_eval_loss"] + args.plateau_delta + ] + zero = next((row for row in curve if row["dropout_initial"] == 0.0), None) + nonzero = [row for row in curve if row["dropout_initial"] > 0.0] + best_nonzero = min(nonzero, key=lambda row: row["mean_val_eval_loss"]) if nonzero else None + max_dropout = max(curve, key=lambda row: row["dropout_initial"]) + + zero_loss = zero["mean_val_eval_loss"] if zero else None + zero_minus_best = zero_loss - best["mean_val_eval_loss"] if zero_loss is not None else None + zero_minus_best_nonzero = ( + zero_loss - best_nonzero["mean_val_eval_loss"] + if zero_loss is not None and best_nonzero is not None + else None + ) + max_minus_best = max_dropout["mean_val_eval_loss"] - best["mean_val_eval_loss"] + has_nonzero_optimum = ( + best["dropout_initial"] > 0.0 + and zero_minus_best is not None + and zero_minus_best >= args.min_nonzero_margin + and max_minus_best >= args.min_high_dropout_margin + ) + curve_json = json.dumps( + [ + { + "dropout": row["dropout_initial"], + "mean_val_loss": row["mean_val_eval_loss"], + "std_val_loss": row["std_val_eval_loss"], + "mean_train_loss": row["mean_train_eval_loss"], + "mean_generalization_gap": row["mean_generalization_gap"], + "n": row["n"], + } + for row in curve + ], + sort_keys=True, + ) + selection.append( + { + "run_mode": run_mode, + "token_limit": token_limit, + "model_name": model_name, + "n_layer": best["n_layer"], + "n_head": best["n_head"], + "n_embd": best["n_embd"], + "parameters": best["parameters"], + "n": best["n"], + "best_dropout": best["dropout_initial"], + "best_val_loss": best["mean_val_eval_loss"], + "best_val_std": best["std_val_eval_loss"], + "plateau_start_dropout": min(row["dropout_initial"] for row in plateau), + "plateau_end_dropout": max(row["dropout_initial"] for row in plateau), + "plateau_delta": args.plateau_delta, + "zero_dropout_val_loss": zero_loss, + "zero_minus_best": zero_minus_best, + "best_nonzero_dropout": ( + best_nonzero["dropout_initial"] if best_nonzero else None + ), + "best_nonzero_val_loss": ( + best_nonzero["mean_val_eval_loss"] if best_nonzero else None + ), + "zero_minus_best_nonzero": zero_minus_best_nonzero, + "max_dropout": max_dropout["dropout_initial"], + "max_dropout_val_loss": max_dropout["mean_val_eval_loss"], + "max_dropout_minus_best": max_minus_best, + "has_nonzero_optimum": has_nonzero_optimum, + "meets_target_dropout": best["dropout_initial"] >= args.target_min_dropout, + "curve_json": curve_json, + } + ) + return sorted( + selection, + key=lambda row: ( + not row["has_nonzero_optimum"], + not row["meets_target_dropout"], + row["best_val_loss"], + ), + ) diff --git a/src/dropout_decay/experiments/device.py b/src/dropout_decay/experiments/device.py new file mode 100644 index 0000000000000000000000000000000000000000..6d4d894525cebb35d9afe23aba74f036f49f7d51 --- /dev/null +++ b/src/dropout_decay/experiments/device.py @@ -0,0 +1,37 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from __future__ import annotations + +import os + +import torch + + +def assert_mps_only() -> torch.device: + if os.environ.get("PYTORCH_ENABLE_MPS_FALLBACK") == "1": + raise SystemExit( + "Refusing to run: PYTORCH_ENABLE_MPS_FALLBACK=1 could execute CPU fallbacks." + ) + if not torch.backends.mps.is_built(): + raise SystemExit("PyTorch was not built with MPS. Stopping.") + if not torch.backends.mps.is_available(): + raise SystemExit("MPS is not available. Stopping before any Torch experiment work.") + device = torch.device("mps") + torch.set_default_device(device) + return device diff --git a/src/dropout_decay/experiments/parsing.py b/src/dropout_decay/experiments/parsing.py new file mode 100644 index 0000000000000000000000000000000000000000..1779e1143d5dea18d33b9816002df55fa4085181 --- /dev/null +++ b/src/dropout_decay/experiments/parsing.py @@ -0,0 +1,148 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from __future__ import annotations + +import argparse + +from dropout_decay.specs import DropoutCondition, ModelSpec + + +DEFAULT_DROPOUT_RATES = [0.0, 0.02, 0.05, 0.08, 0.10, 0.14, 0.20, 0.30, 0.50] + + +def clean_name(value: str) -> str: + cleaned = "".join( + c if c.isalnum() or c in {"-", "_"} else "_" for c in value.strip() + ) + return cleaned.strip("_") or "model" + + +def rate_label(rate: float) -> str: + label = f"{rate:.3f}".rstrip("0").rstrip(".") + return label if label else "0" + + +def parse_model_spec(raw: str) -> ModelSpec: + if "=" in raw: + name, dims = raw.split("=", 1) + name = clean_name(name) + else: + dims = raw + name = "" + parts = dims.lower().replace(",", "x").split("x") + if len(parts) != 3: + raise argparse.ArgumentTypeError( + "model specs must look like 8x8x256 or name=8x8x256" + ) + try: + n_layer, n_head, n_embd = [int(part) for part in parts] + except ValueError as exc: + raise argparse.ArgumentTypeError("model dimensions must be integers") from exc + if n_layer <= 0 or n_head <= 0 or n_embd <= 0: + raise argparse.ArgumentTypeError("model dimensions must be positive") + if n_embd % n_head != 0: + raise argparse.ArgumentTypeError("n_embd must be divisible by n_head") + if (n_embd // n_head) % 2 != 0: + raise argparse.ArgumentTypeError("n_embd / n_head must be even for rotary attention") + return ModelSpec( + name or f"L{n_layer}_H{n_head}_D{n_embd}", + n_layer, + n_head, + n_embd, + ) + + +def parse_decay_spec(raw: str) -> DropoutCondition: + parts = raw.split(":") + if len(parts) not in {3, 4, 5}: + raise argparse.ArgumentTypeError( + "decay specs must look like " + "name:initial:final[:cosine|smoothstep|linear[:decay_tokens]]" + ) + name = clean_name(parts[0]) + try: + initial = float(parts[1]) + final = float(parts[2]) + decay_tokens = int(parts[4]) if len(parts) == 5 else None + except ValueError as exc: + raise argparse.ArgumentTypeError( + "decay dropout values and decay_tokens must be numeric" + ) from exc + schedule = parts[3] if len(parts) >= 4 else "cosine" + if schedule not in {"cosine", "smoothstep", "linear"}: + raise argparse.ArgumentTypeError( + "decay schedule must be cosine, smoothstep, or linear" + ) + return DropoutCondition( + name=name, + kind="decay", + initial=initial, + final=final, + schedule=schedule, + decay_tokens=decay_tokens, + ) + + +def parse_anchor_decay_spec(raw: str) -> DropoutCondition: + if ":" not in raw: + raise argparse.ArgumentTypeError( + "anchor decay specs must look like name:250000=0.60,500000=0.40" + ) + name, raw_anchors = raw.split(":", 1) + anchors: list[tuple[int, float]] = [] + for piece in raw_anchors.split(","): + if "=" not in piece: + raise argparse.ArgumentTypeError( + "anchor decay anchors must look like token_count=dropout" + ) + raw_tokens, raw_dropout = piece.split("=", 1) + try: + token_count = int(raw_tokens) + dropout = float(raw_dropout) + except ValueError as exc: + raise argparse.ArgumentTypeError( + "anchor token counts must be integers and dropout values numeric" + ) from exc + if token_count <= 0: + raise argparse.ArgumentTypeError("anchor token counts must be positive") + if not 0.0 <= dropout < 1.0: + raise argparse.ArgumentTypeError("anchor dropout must satisfy 0 <= p < 1") + anchors.append((token_count, dropout)) + anchors = sorted(set(anchors)) + if len(anchors) < 2: + raise argparse.ArgumentTypeError("provide at least two dropout anchors") + for left, right in zip(anchors, anchors[1:]): + if right[1] > left[1]: + raise argparse.ArgumentTypeError( + "anchor dropout values must be non-increasing as token counts grow" + ) + return DropoutCondition( + name=clean_name(name), + kind="anchor_decay", + initial=anchors[0][1], + final=anchors[-1][1], + schedule="log_prefix_anchor", + anchors=tuple(anchors), + ) + + +def default_seeds(mode: str, seeds: list[int] | None) -> list[int]: + if seeds: + return seeds + return [1] if mode == "screen_static" else [1, 2, 3] diff --git a/src/dropout_decay/experiments/progress.py b/src/dropout_decay/experiments/progress.py new file mode 100644 index 0000000000000000000000000000000000000000..67c5df60f27033a67a22c53f24aec1634044b99c --- /dev/null +++ b/src/dropout_decay/experiments/progress.py @@ -0,0 +1,66 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from __future__ import annotations + +import time + + +def format_duration(seconds: float) -> str: + seconds = max(0, int(seconds)) + hours, remainder = divmod(seconds, 3600) + minutes, secs = divmod(remainder, 60) + if hours: + return f"{hours}h{minutes:02d}m" + if minutes: + return f"{minutes}m{secs:02d}s" + return f"{secs}s" + + +class ProgressMeter: + def __init__(self, total: int): + self.total = max(0, total) + self.done = 0 + self.started_at = time.time() + + def mark_done(self, row: dict) -> None: + self.done += 1 + elapsed = time.time() - self.started_at + mean = elapsed / self.done if self.done else 0.0 + remaining = max(0, self.total - self.done) + eta = remaining * mean + stage = row.get("stage") + stage_text = f"stage={stage} " if stage is not None else "" + print( + "progress " + f"{self.done}/{self.total} " + f"eta={format_duration(eta)} " + f"mode={row['run_mode']} " + f"model={row['model_name']} " + f"params={int(row['parameters']):,} " + f"prefix={int(row['token_limit']):,} " + f"{stage_text}" + f"seed={row['seed']} " + f"condition={row['condition']} " + f"dropout={float(row['dropout_active_final']):.3f} " + f"val={row['val_eval_loss']:.4f} " + f"train={row['train_eval_loss']:.4f} " + f"gap={row['generalization_gap']:.4f} " + f"elapsed={format_duration(float(row['elapsed_sec']))}", + flush=True, + ) diff --git a/src/dropout_decay/experiments/reports.py b/src/dropout_decay/experiments/reports.py new file mode 100644 index 0000000000000000000000000000000000000000..4040567b38a494722ef3cb54fd56c6b5787a8f97 --- /dev/null +++ b/src/dropout_decay/experiments/reports.py @@ -0,0 +1,448 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from __future__ import annotations + +from collections import defaultdict +from pathlib import Path +import statistics + + +def write_screen_markdown_summary(output_dir: Path, rows: list[dict]) -> None: + if not rows: + return + static_rows = [ + row + for row in rows + if row["run_mode"] == "screen_static" and row["condition_kind"] == "static" + ] + if not static_rows: + return + + by_model_prefix_rate: dict[tuple[str, int, float], list[dict]] = defaultdict(list) + for row in rows: + if row["run_mode"] == "screen_static" and row["condition_kind"] == "static": + by_model_prefix_rate[ + ( + row["model_name"], + int(row["token_limit"]), + float(row["dropout_initial"]), + ) + ].append(row) + + aggregates: list[dict] = [] + for (model_name, prefix, dropout), group_rows in by_model_prefix_rate.items(): + first = group_rows[0] + val_losses = [float(row["val_eval_loss"]) for row in group_rows] + train_losses = [float(row["train_eval_loss"]) for row in group_rows] + gaps = [float(row["generalization_gap"]) for row in group_rows] + aggregates.append( + { + "model_name": model_name, + "token_limit": prefix, + "dropout_initial": dropout, + "n": len(group_rows), + "mean_val_eval_loss": statistics.fmean(val_losses), + "std_val_eval_loss": statistics.stdev(val_losses) + if len(val_losses) > 1 + else 0.0, + "mean_train_eval_loss": statistics.fmean(train_losses), + "std_train_eval_loss": statistics.stdev(train_losses) + if len(train_losses) > 1 + else 0.0, + "mean_generalization_gap": statistics.fmean(gaps), + "std_generalization_gap": statistics.stdev(gaps) + if len(gaps) > 1 + else 0.0, + "parameters": int(first["parameters"]), + "n_layer": int(first["n_layer"]), + "n_head": int(first["n_head"]), + "n_embd": int(first["n_embd"]), + "block_size": int(first["model_config"]["block_size"]), + "vocab_size": int(first["model_config"]["vocab_size"]), + "tokens_seen": int(first["tokens_seen"]), + "seeds": sorted({int(row["seed"]) for row in group_rows}), + } + ) + + by_model: dict[str, list[dict]] = defaultdict(list) + for row in aggregates: + by_model[row["model_name"]].append(row) + + model_rows = [] + for model_name, model_group in by_model.items(): + first = model_group[0] + seeds = sorted({seed for row in model_group for seed in row["seeds"]}) + model_rows.append( + { + "model_name": model_name, + "parameters": first["parameters"], + "n_layer": first["n_layer"], + "n_head": first["n_head"], + "n_embd": first["n_embd"], + "block_size": first["block_size"], + "vocab_size": first["vocab_size"], + "seeds": seeds, + } + ) + + lines = [ + "# Static Dropout Screen Summary", + "", + f"Run directory: `{output_dir}`", + "", + "## Models", + "", + "| Model | Params | Layers | Heads | Embedding | Block | Vocab | Seeds |", + "|---|---:|---:|---:|---:|---:|---:|---|", + ] + for model in sorted(model_rows, key=lambda item: item["parameters"]): + lines.append( + "| " + f"`{model['model_name']}` | {model['parameters']:,} | " + f"{model['n_layer']} | {model['n_head']} | {model['n_embd']} | " + f"{model['block_size']} | {model['vocab_size']} | " + f"{', '.join(str(seed) for seed in model['seeds'])} |" + ) + lines.extend( + [ + "", + "## Best Dropout By Model And Prefix", + "", + "| Model | Prefix tokens | Effective epochs | Best dropout | Mean val loss | Val std | Mean train loss | Mean gap | Plateau/bracket note |", + "|---|---:|---:|---:|---:|---:|---:|---:|---|", + ] + ) + for model_name, model_group in sorted(by_model.items()): + by_prefix: dict[int, list[dict]] = defaultdict(list) + for row in model_group: + by_prefix[int(row["token_limit"])].append(row) + for prefix, prefix_rows in sorted(by_prefix.items()): + best = min(prefix_rows, key=lambda row: row["mean_val_eval_loss"]) + rates = [float(row["dropout_initial"]) for row in prefix_rows] + eff_epochs = float(best["tokens_seen"]) / prefix + if best["dropout_initial"] == max(rates): + note = "not bracketed; best at top of tested grid" + elif best["dropout_initial"] == min(rates): + note = "not bracketed; best at bottom of tested grid" + else: + note = "bracketed by tested grid" + lines.append( + "| " + f"`{model_name}` | {prefix:,} | {eff_epochs:.2f} | " + f"{best['dropout_initial']:.2f} | " + f"{best['mean_val_eval_loss']:.4f} | " + f"{best['std_val_eval_loss']:.4f} | " + f"{best['mean_train_eval_loss']:.4f} | " + f"{best['mean_generalization_gap']:.4f} | {note} |" + ) + + for model_name, model_group in sorted(by_model.items()): + by_prefix = defaultdict(list) + for row in model_group: + by_prefix[int(row["token_limit"])].append(row) + lines.extend( + [ + "", + f"## Model `{model_name}`", + ] + ) + for prefix, prefix_rows in sorted(by_prefix.items()): + eff_epochs = float(prefix_rows[0]["tokens_seen"]) / prefix + lines.extend( + [ + "", + f"### Prefix {prefix:,} Tokens ({eff_epochs:.2f} Effective Epochs)", + "", + "| Dropout | N | Mean val loss | Val std | Mean train loss | Train std | Mean gap | Gap std | Sampled tokens | Params |", + "|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|", + ] + ) + for row in sorted(prefix_rows, key=lambda item: item["dropout_initial"]): + lines.append( + "| " + f"{row['dropout_initial']:.2f} | {row['n']} | " + f"{row['mean_val_eval_loss']:.4f} | " + f"{row['std_val_eval_loss']:.4f} | " + f"{row['mean_train_eval_loss']:.4f} | " + f"{row['std_train_eval_loss']:.4f} | " + f"{row['mean_generalization_gap']:.4f} | " + f"{row['std_generalization_gap']:.4f} | " + f"{int(row['tokens_seen']):,} | {int(row['parameters']):,} |" + ) + + output = "\n".join(lines) + "\n" + (output_dir / "RESULT_SUMMARY.md").write_text(output, encoding="utf-8") + + +def svg_escape(value: object) -> str: + return ( + str(value) + .replace("&", "&") + .replace("<", "<") + .replace(">", ">") + .replace('"', """) + ) + + +def write_dropout_curve_svg(output_dir: Path, summary: list[dict]) -> None: + rows = [ + row + for row in summary + if row["run_mode"] == "screen_static" and row["condition_kind"] == "static" + ] + if not rows: + return + + grouped: dict[tuple[str, int], list[dict]] = defaultdict(list) + model_params: dict[str, int] = {} + for row in rows: + model_name = row["model_name"] + grouped[(model_name, int(row["token_limit"]))].append(row) + model_params[model_name] = int(row["parameters"]) + + models = sorted(model_params, key=lambda name: model_params[name]) + prefixes = sorted({int(row["token_limit"]) for row in rows}) + panel_w, panel_h = 230, 170 + margin_l, margin_t = 58, 34 + plot_w, plot_h = 142, 94 + gap_x, gap_y = 18, 38 + width = margin_l + len(prefixes) * panel_w + gap_x + height = 70 + len(models) * (panel_h + gap_y) + colors = ["#1f77b4", "#d62728", "#2ca02c", "#9467bd", "#ff7f0e"] + + parts = [ + f'', + "", + '', + 'Static dropout law: validation loss vs dropout', + 'Each panel uses its own y-scale. Points are one-seed means unless N > 1.', + ] + + for col, prefix in enumerate(prefixes): + x = margin_l + col * panel_w + plot_w / 2 + parts.append( + f'{prefix:,} prefix tokens' + ) + + for row_idx, model_name in enumerate(models): + row_y = 92 + row_idx * (panel_h + gap_y) + parts.append( + f'' + f'{svg_escape(model_name)} ({model_params[model_name] / 1_000_000:.1f}M)' + ) + for col, prefix in enumerate(prefixes): + panel_x = margin_l + col * panel_w + panel_y = row_y + curve = sorted( + grouped.get((model_name, prefix), []), + key=lambda item: float(item["dropout_initial"]), + ) + if not curve: + continue + losses = [float(item["mean_val_eval_loss"]) for item in curve] + min_loss, max_loss = min(losses), max(losses) + pad = max(0.02, (max_loss - min_loss) * 0.08) + y_min, y_max = min_loss - pad, max_loss + pad + best = min(curve, key=lambda item: float(item["mean_val_eval_loss"])) + + def px(dropout: float) -> float: + return panel_x + (dropout / 0.9) * plot_w + + def py(loss: float) -> float: + scale = (loss - y_min) / (y_max - y_min) + return panel_y + plot_h - scale * plot_h + + parts.extend( + [ + f'', + f'', + f'', + f'{y_max:.2f}', + f'{y_min:.2f}', + f'0', + f'0.9', + ] + ) + points = " ".join( + f"{px(float(item['dropout_initial'])):.1f},{py(float(item['mean_val_eval_loss'])):.1f}" + for item in curve + ) + color = colors[row_idx % len(colors)] + parts.append(f'') + for item in curve: + dropout = float(item["dropout_initial"]) + loss = float(item["mean_val_eval_loss"]) + radius = 4 if item is best else 2.7 + fill = "#111827" if item is best else "#ffffff" + parts.append( + f'' + ) + parts.append( + f'' + f'best p={float(best["dropout_initial"]):.2f}' + ) + parts.append( + f'' + f'loss={float(best["mean_val_eval_loss"]):.3f}' + ) + + parts.append("") + (output_dir / "dropout_curves.svg").write_text("\n".join(parts), encoding="utf-8") + + +def write_stream_markdown_summary(output_dir: Path, rows: list[dict]) -> None: + stream_rows = [row for row in rows if row["run_mode"] == "locked_stream"] + if not stream_rows: + return + + by_condition_stage: dict[tuple[str, int], list[dict]] = defaultdict(list) + by_condition: dict[str, list[dict]] = defaultdict(list) + for row in stream_rows: + condition = row["condition"] + by_condition_stage[(condition, int(row["stage"]))].append(row) + by_condition[condition].append(row) + + first = stream_rows[0] + seeds = sorted({int(row["seed"]) for row in stream_rows}) + conditions = sorted( + by_condition, + key=lambda name: ( + by_condition[name][0]["condition_kind"] != "anchor_decay", + by_condition[name][0]["dropout_initial"], + name, + ), + ) + stages = sorted({int(row["stage"]) for row in stream_rows}) + + lines = [ + "# Locked Streaming Dropout Summary", + "", + f"Run directory: `{output_dir}`", + "", + ( + f"Model: `{first['model_name']}` causal Transformer, " + f"{int(first['parameters']):,} parameters, {first['n_layer']} layers, " + f"{first['n_head']} heads, {first['n_embd']} embedding dim." + ), + ( + f"Training per stage: {first['steps']:,} steps. " + "Sampled tokens are cumulative in each stage row. " + f"Seeds present: {', '.join(str(seed) for seed in seeds)}." + ), + "", + "## Condition Ranking", + "", + "| Condition | Kind | Final dropout | Mean trajectory val loss | Final val loss | Final gap | Dropout path |", + "|---|---|---:|---:|---:|---:|---|", + ] + + ranking = [] + for condition in conditions: + stage_items = [] + for stage in stages: + group = by_condition_stage.get((condition, stage), []) + if not group: + continue + stage_items.append( + { + "stage": stage, + "token_limit": int(group[0]["token_limit"]), + "mean_val": statistics.fmean( + float(row["val_eval_loss"]) for row in group + ), + "mean_gap": statistics.fmean( + float(row["generalization_gap"]) for row in group + ), + "mean_dropout": statistics.fmean( + float(row["dropout_active_final"]) for row in group + ), + "kind": group[0]["condition_kind"], + } + ) + if not stage_items: + continue + final = max(stage_items, key=lambda item: item["stage"]) + ranking.append( + { + "condition": condition, + "kind": stage_items[0]["kind"], + "trajectory_val": statistics.fmean(item["mean_val"] for item in stage_items), + "final_val": final["mean_val"], + "final_gap": final["mean_gap"], + "final_dropout": final["mean_dropout"], + "dropout_path": " -> ".join( + f"{item['mean_dropout']:.2f}" for item in stage_items + ), + } + ) + + for item in sorted(ranking, key=lambda row: row["trajectory_val"]): + lines.append( + "| " + f"`{item['condition']}` | {item['kind']} | " + f"{item['final_dropout']:.2f} | {item['trajectory_val']:.4f} | " + f"{item['final_val']:.4f} | {item['final_gap']:.4f} | " + f"{item['dropout_path']} |" + ) + + lines.extend(["", "## Stage Trajectory", ""]) + for stage in stages: + stage_groups = { + condition: by_condition_stage[(condition, stage)] + for condition in conditions + if (condition, stage) in by_condition_stage + } + if not stage_groups: + continue + prefix = int(next(iter(stage_groups.values()))[0]["token_limit"]) + lines.extend( + [ + f"### Stage {stage}: {prefix:,} Prefix Tokens", + "", + "| Condition | Dropout | Mean val loss | Mean train loss | Mean gap | N |", + "|---|---:|---:|---:|---:|---:|", + ] + ) + for condition, group in sorted( + stage_groups.items(), + key=lambda item: statistics.fmean( + float(row["val_eval_loss"]) for row in item[1] + ), + ): + val = statistics.fmean(float(row["val_eval_loss"]) for row in group) + train = statistics.fmean(float(row["train_eval_loss"]) for row in group) + gap = statistics.fmean(float(row["generalization_gap"]) for row in group) + dropout = statistics.fmean( + float(row["dropout_active_final"]) for row in group + ) + lines.append( + "| " + f"`{condition}` | {dropout:.2f} | {val:.4f} | " + f"{train:.4f} | {gap:.4f} | {len(group)} |" + ) + lines.append("") + + (output_dir / "RESULT_SUMMARY.md").write_text( + "\n".join(lines).rstrip() + "\n", + encoding="utf-8", + ) diff --git a/src/dropout_decay/experiments/runner.py b/src/dropout_decay/experiments/runner.py new file mode 100644 index 0000000000000000000000000000000000000000..e85d4638ef72e29fe0715f9aadb13ce2cc37308a --- /dev/null +++ b/src/dropout_decay/experiments/runner.py @@ -0,0 +1,555 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from __future__ import annotations + +import argparse +from datetime import datetime +import json +from pathlib import Path +import statistics +import sys + +import numpy as np +import torch + +from dropout_decay.datasets import ( + encode_corpus, + load_cached_splits, + resolve_paths, + train_or_load_tokenizer, +) +from dropout_decay.experiments.artifacts import ( + SELECTION_FIELDS, + SUMMARY_FIELDS, + build_model_selection, + load_metrics, + metric_key, + planned_metric_key, + summarize, + write_csv, + write_jsonl_row, +) +from dropout_decay.experiments.device import assert_mps_only +from dropout_decay.experiments.parsing import ( + DEFAULT_DROPOUT_RATES, + default_seeds, + parse_anchor_decay_spec, + parse_decay_spec, + parse_model_spec, + rate_label, +) +from dropout_decay.experiments.progress import ProgressMeter +from dropout_decay.experiments.reports import ( + write_dropout_curve_svg, + write_screen_markdown_summary, + write_stream_markdown_summary, +) +from dropout_decay.experiments.training import train_segment +from dropout_decay.license import NANOCHAT_ATTRIBUTION +from dropout_decay.specs import DropoutCondition, ModelSpec + + +def static_conditions(dropout_rates: list[float]) -> list[DropoutCondition]: + return [ + DropoutCondition( + name=f"static_dropout_{rate_label(rate)}", + kind="static", + initial=rate, + final=rate, + ) + for rate in dropout_rates + ] + + +def run_fixed_static_sweep( + *, + args: argparse.Namespace, + model_specs: list[ModelSpec], + seeds: list[int], + train_tokens: np.ndarray, + val_tokens: np.ndarray, + tokenizer_vocab_size: int, + token_limits: list[int], + device: torch.device, + metrics_file, + trace_file, + completed_keys: set[tuple] | None = None, +) -> list[dict]: + rows: list[dict] = [] + completed_keys = completed_keys or set() + conditions = static_conditions(sorted(set(args.dropout_rates))) + planned = 0 + for token_limit in token_limits: + for model_spec in model_specs: + for condition in conditions: + for seed in seeds: + key = planned_metric_key( + mode=args.mode, + condition=condition, + model_spec=model_spec, + seed=seed, + token_limit=token_limit, + ) + if key not in completed_keys: + planned += 1 + progress = ProgressMeter(planned) + for token_limit in token_limits: + for model_spec in model_specs: + best_val_loss = float("inf") + worse_streak = 0 + for condition in conditions: + condition_rows: list[dict] = [] + for seed in seeds: + key = planned_metric_key( + mode=args.mode, + condition=condition, + model_spec=model_spec, + seed=seed, + token_limit=token_limit, + ) + if key in completed_keys: + write_jsonl_row( + trace_file, + { + "event": "skipped_completed_condition", + "run_mode": args.mode, + "condition": condition.name, + "model_name": model_spec.name, + "seed": seed, + "stage": None, + "token_limit": int(token_limit), + }, + ) + continue + config = model_spec.config( + tokenizer_vocab_size, + args.block_size, + condition.initial, + ) + model, optimizer, _, row = train_segment( + run_mode=args.mode, + condition=condition, + model_spec=model_spec, + config=config, + train_tokens=train_tokens, + val_tokens=val_tokens, + token_limit=token_limit, + steps=args.steps, + seed=seed, + args=args, + device=device, + dropout_fn=condition.make_fn( + args.steps * args.batch_size * args.block_size + ), + metrics_file=metrics_file, + trace_file=trace_file, + ) + rows.append(row) + condition_rows.append(row) + completed_keys.add(metric_key(row)) + progress.mark_done(row) + del model, optimizer + torch.mps.empty_cache() + + if not condition_rows: + continue + mean_val_loss = statistics.fmean( + float(row["val_eval_loss"]) for row in condition_rows + ) + if mean_val_loss < best_val_loss - args.screen_prune_min_delta: + best_val_loss = mean_val_loss + worse_streak = 0 + elif mean_val_loss > best_val_loss + args.screen_prune_min_delta: + worse_streak += 1 + + if ( + args.mode == "screen_static" + and args.screen_early_stop + and worse_streak >= args.screen_prune_patience + and condition.initial >= args.target_min_dropout + ): + write_jsonl_row( + trace_file, + { + "event": "screen_pruned_model", + "run_mode": args.mode, + "model_name": model_spec.name, + "token_limit": int(token_limit), + "best_val_loss": best_val_loss, + "pruned_after_dropout": condition.initial, + "worse_streak": worse_streak, + "remaining_dropouts": [ + rate + for rate in args.dropout_rates + if rate > condition.initial + ], + }, + ) + break + return rows + + +def run_locked_stream( + *, + args: argparse.Namespace, + model_specs: list[ModelSpec], + seeds: list[int], + train_tokens: np.ndarray, + val_tokens: np.ndarray, + tokenizer_vocab_size: int, + stream_caps: list[int], + device: torch.device, + metrics_file, + trace_file, + completed_keys: set[tuple] | None = None, +) -> list[dict]: + rows: list[dict] = [] + completed_keys = completed_keys or set() + conditions = args.anchor_decays + static_conditions(args.dropout_rates) + args.decays + fallback_decay_tokens = ( + args.decay_tokens + or args.stage_steps * args.batch_size * args.block_size * len(stream_caps) + ) + planned = 0 + for model_spec in model_specs: + for condition in conditions: + for seed in seeds: + for stage, token_limit in enumerate(stream_caps): + key = planned_metric_key( + mode=args.mode, + condition=condition, + model_spec=model_spec, + seed=seed, + token_limit=token_limit, + stage=stage, + ) + if key not in completed_keys: + planned += 1 + progress = ProgressMeter(planned) + for model_spec in model_specs: + for condition in conditions: + for seed in seeds: + model = None + optimizer = None + tokens_seen = 0 + for stage, token_limit in enumerate(stream_caps): + key = planned_metric_key( + mode=args.mode, + condition=condition, + model_spec=model_spec, + seed=seed, + token_limit=token_limit, + stage=stage, + ) + if key in completed_keys: + write_jsonl_row( + trace_file, + { + "event": "skipped_completed_condition", + "run_mode": args.mode, + "condition": condition.name, + "model_name": model_spec.name, + "seed": seed, + "stage": stage, + "token_limit": int(token_limit), + }, + ) + continue + config = model_spec.config( + tokenizer_vocab_size, + args.block_size, + condition.initial, + ) + model, optimizer, tokens_seen, row = train_segment( + run_mode=args.mode, + condition=condition, + model_spec=model_spec, + config=config, + train_tokens=train_tokens, + val_tokens=val_tokens, + token_limit=token_limit, + steps=args.stage_steps, + seed=seed, + args=args, + device=device, + dropout_fn=condition.make_fn( + fallback_decay_tokens, + unique_tokens=token_limit, + ), + metrics_file=metrics_file, + trace_file=trace_file, + stage=stage, + model=model, + optimizer=optimizer, + tokens_seen_start=tokens_seen, + ) + rows.append(row) + completed_keys.add(metric_key(row)) + progress.mark_done(row) + del model, optimizer + torch.mps.empty_cache() + return rows + + +def prepare_data(args: argparse.Namespace, output_dir: Path, required_train_tokens: int): + cache_dir = Path(args.cache_dir) if args.cache_dir else output_dir / "cache" + cache_dir.mkdir(parents=True, exist_ok=True) + if args.use_cached_data: + if args.force_retokenize: + raise ValueError("--use-cached-data cannot be combined with --force-retokenize") + return load_cached_splits( + cache_dir=cache_dir, + vocab_size=args.vocab_size, + max_required_train_tokens=required_train_tokens, + val_tokens=args.val_tokens, + allow_short_corpus=args.allow_short_corpus, + ) + + paths = resolve_paths(args.corpus, args.corpus_glob) + tokenizer = train_or_load_tokenizer( + paths=paths, + output_dir=cache_dir, + vocab_size=args.vocab_size, + tokenizer_train_chars=args.tokenizer_train_chars, + text_column=args.text_column, + force_retrain=args.force_retokenize, + ) + splits = encode_corpus( + paths=paths, + tokenizer=tokenizer, + output_dir=cache_dir, + max_required_train_tokens=required_train_tokens, + val_tokens=args.val_tokens, + text_column=args.text_column, + allow_short_corpus=args.allow_short_corpus, + force_reencode=args.force_retokenize, + ) + return tokenizer, splits + + +def run(args: argparse.Namespace) -> Path: + device = assert_mps_only() + seeds = default_seeds(args.mode, args.seeds) + model_specs = [parse_model_spec(spec) for spec in args.models] + if args.mode != "locked_stream" and (args.decays or args.anchor_decays): + raise ValueError("--decays and --anchor-decays are only used with --mode locked_stream") + if args.resume_from and args.mode == "locked_stream": + raise ValueError("--resume-from currently supports fixed static sweeps only") + + if args.resume_from: + output_dir = Path(args.resume_from) + if not output_dir.exists(): + raise FileNotFoundError(f"resume directory does not exist: {output_dir}") + else: + run_id = datetime.now().strftime("%Y%m%d-%H%M%S") + output_dir = Path(args.output_dir) / args.mode / run_id + output_dir.mkdir(parents=True, exist_ok=True) + required_train_tokens = max( + args.stream_token_caps if args.mode == "locked_stream" else args.token_limits + ) + tokenizer, splits = prepare_data(args, output_dir, required_train_tokens) + token_limits = [min(limit, len(splits.train)) for limit in args.token_limits] + stream_caps = [min(limit, len(splits.train)) for limit in args.stream_token_caps] + + args_payload = vars(args).copy() + args_payload["decays"] = [condition.to_dict() for condition in args.decays] + args_payload["anchor_decays"] = [ + condition.to_dict() for condition in args.anchor_decays + ] + config_payload = { + "args": args_payload, + "mode": args.mode, + "seeds": seeds, + "models": [model.to_dict() for model in model_specs], + "device": str(device), + "torch": torch.__version__, + "python": sys.version, + "mps_available": torch.backends.mps.is_available(), + "attribution": NANOCHAT_ATTRIBUTION, + "tokenizer_path": str(splits.tokenizer_path), + "encoded_path": str(splits.encoded_path), + "train_tokens": int(len(splits.train)), + "val_tokens": int(len(splits.val)), + "effective_token_limits": [int(limit) for limit in token_limits], + "effective_stream_token_caps": [int(limit) for limit in stream_caps], + "resume_from": str(args.resume_from) if args.resume_from else None, + } + config_name = "config.resume.json" if args.resume_from else "config.json" + (output_dir / config_name).write_text( + json.dumps(config_payload, indent=2), + encoding="utf-8", + ) + + metrics_path = output_dir / "metrics.jsonl" + trace_path = output_dir / "trace.jsonl" + existing_rows = load_metrics(metrics_path) if args.resume_from else [] + completed_keys = {metric_key(row) for row in existing_rows} + with ( + metrics_path.open("a" if args.resume_from else "w", encoding="utf-8") as metrics_file, + trace_path.open("a" if args.resume_from else "w", encoding="utf-8") as trace_file, + ): + if args.mode in {"screen_static", "confirm_static"}: + new_rows = run_fixed_static_sweep( + args=args, + model_specs=model_specs, + seeds=seeds, + train_tokens=splits.train, + val_tokens=splits.val, + tokenizer_vocab_size=tokenizer.vocab_size, + token_limits=token_limits, + device=device, + metrics_file=metrics_file, + trace_file=trace_file, + completed_keys=completed_keys, + ) + else: + new_rows = run_locked_stream( + args=args, + model_specs=model_specs, + seeds=seeds, + train_tokens=splits.train, + val_tokens=splits.val, + tokenizer_vocab_size=tokenizer.vocab_size, + stream_caps=stream_caps, + device=device, + metrics_file=metrics_file, + trace_file=trace_file, + ) + rows = existing_rows + new_rows + + summary = summarize(rows) + (output_dir / "summary.json").write_text( + json.dumps(summary, indent=2), + encoding="utf-8", + ) + write_csv(output_dir / "summary.csv", summary, SUMMARY_FIELDS) + if args.mode in {"screen_static", "confirm_static"}: + selection = build_model_selection(summary, args) + (output_dir / "model_selection.json").write_text( + json.dumps(selection, indent=2), + encoding="utf-8", + ) + write_csv(output_dir / "model_selection.csv", selection, SELECTION_FIELDS) + write_screen_markdown_summary(output_dir, rows) + write_dropout_curve_svg(output_dir, summary) + elif args.mode == "locked_stream": + write_stream_markdown_summary(output_dir, rows) + + print( + json.dumps( + { + "output_dir": str(output_dir), + "new_rows": len(new_rows), + "total_metric_rows": len(rows), + "summary_rows": len(summary), + }, + indent=2, + ) + ) + return output_dir + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="MPS-only dropout/model selection experiments" + ) + parser.add_argument( + "--mode", + choices=["screen_static", "confirm_static", "locked_stream"], + default="screen_static", + ) + parser.add_argument("--corpus", default=None, help="Text or parquet corpus path") + parser.add_argument("--corpus-glob", default=None, help="Glob of text/parquet corpus paths") + parser.add_argument("--text-column", default="text", help="Parquet text column") + parser.add_argument( + "--use-cached-data", + action="store_true", + help=( + "Load tokenizer-v{vocab}.json and tokens-v{vocab}-*.npy from --cache-dir " + "instead of requiring the original text/parquet corpus." + ), + ) + parser.add_argument("--output-dir", default="runs") + parser.add_argument( + "--resume-from", + default=None, + help="Existing fixed-static run directory; completed metric rows are skipped", + ) + parser.add_argument("--cache-dir", default=".cache/dropout_decay") + parser.add_argument( + "--models", + nargs="+", + default=["8x8x256"], + help="Model specs like 8x8x256 or name=8x8x256", + ) + parser.add_argument("--seeds", nargs="+", type=int, default=None) + parser.add_argument("--token-limits", nargs="+", type=int, default=[5_000_000]) + parser.add_argument( + "--stream-token-caps", + nargs="+", + type=int, + default=[5_000_000, 10_000_000, 20_000_000, 40_000_000], + ) + parser.add_argument("--val-tokens", type=int, default=500_000) + parser.add_argument("--allow-short-corpus", action="store_true") + parser.add_argument("--force-retokenize", action="store_true") + parser.add_argument("--vocab-size", type=int, default=4096) + parser.add_argument("--tokenizer-train-chars", type=int, default=10_000_000) + parser.add_argument("--block-size", type=int, default=128) + parser.add_argument("--batch-size", type=int, default=16) + parser.add_argument("--steps", type=int, default=2000) + parser.add_argument("--stage-steps", type=int, default=1000) + parser.add_argument("--dropout-rates", nargs="*", type=float, default=DEFAULT_DROPOUT_RATES) + parser.add_argument("--decays", nargs="*", type=parse_decay_spec, default=[]) + parser.add_argument( + "--anchor-decays", + nargs="*", + type=parse_anchor_decay_spec, + default=[], + help=( + "Prefix-token anchor schedules like " + "fit:250000=0.60,500000=0.40,1000000=0.30" + ), + ) + parser.add_argument("--decay-tokens", type=int, default=None) + parser.add_argument("--eval-batches", type=int, default=64) + parser.add_argument("--train-eval-batches", type=int, default=32) + parser.add_argument("--trace-eval-batches", type=int, default=8) + parser.add_argument("--eval-every", type=int, default=0) + parser.add_argument("--log-every", type=int, default=100) + parser.add_argument("--lr", type=float, default=3e-4) + parser.add_argument("--weight-decay", type=float, default=0.1) + parser.add_argument("--grad-clip", type=float, default=1.0) + parser.add_argument("--plateau-delta", type=float, default=0.01) + parser.add_argument("--target-min-dropout", type=float, default=0.10) + parser.add_argument("--min-nonzero-margin", type=float, default=0.01) + parser.add_argument("--min-high-dropout-margin", type=float, default=0.03) + parser.add_argument("--screen-early-stop", action="store_true") + parser.add_argument("--screen-prune-patience", type=int, default=3) + parser.add_argument("--screen-prune-min-delta", type=float, default=0.01) + return parser + + +def main() -> None: + args = build_parser().parse_args() + run(args) + + +if __name__ == "__main__": + main() diff --git a/src/dropout_decay/experiments/training.py b/src/dropout_decay/experiments/training.py new file mode 100644 index 0000000000000000000000000000000000000000..744a6bc2db54230d94bad8896aa3e16818494369 --- /dev/null +++ b/src/dropout_decay/experiments/training.py @@ -0,0 +1,261 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from __future__ import annotations + +import argparse +import random +import statistics +import time +from typing import Callable + +import numpy as np +import torch + +from dropout_decay.experiments.artifacts import write_jsonl_row +from dropout_decay.models import DropoutGPT, GPTConfig +from dropout_decay.specs import DropoutCondition, ModelSpec + + +def set_seed(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + + +def batch_seed( + seed: int, model: ModelSpec, dropout_code: int, stage: int | None +) -> int: + return ( + seed * 1_000_003 + + model.n_layer * 100_003 + + model.n_head * 10_007 + + model.n_embd * 101 + + dropout_code * 37 + + (stage or 0) * 997 + ) + + +def make_batch( + tokens: np.ndarray, + token_limit: int, + batch_size: int, + block_size: int, + rng: np.random.Generator, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + limit = min(token_limit, len(tokens)) + max_start = limit - block_size - 1 + if max_start <= 0: + raise ValueError("token_limit is too small for the requested block_size") + starts = rng.integers(0, max_start, size=batch_size) + x_np = np.stack([tokens[start : start + block_size] for start in starts]).astype( + np.int64 + ) + y_np = np.stack( + [tokens[start + 1 : start + 1 + block_size] for start in starts] + ).astype(np.int64) + return ( + torch.tensor(x_np, dtype=torch.long, device=device), + torch.tensor(y_np, dtype=torch.long, device=device), + ) + + +@torch.no_grad() +def estimate_loss( + model: DropoutGPT, + tokens: np.ndarray, + token_limit: int, + batches: int, + args: argparse.Namespace, + device: torch.device, + rng_seed: int, +) -> float: + if batches <= 0: + return float("nan") + model.eval() + rng = np.random.default_rng(rng_seed) + losses: list[float] = [] + for _ in range(batches): + x, y = make_batch( + tokens, token_limit, args.batch_size, args.block_size, rng, device + ) + _, loss = model(x, y) + losses.append(float(loss.item())) + model.train() + return float(statistics.fmean(losses)) + + +def train_segment( + *, + run_mode: str, + condition: DropoutCondition, + model_spec: ModelSpec, + config: GPTConfig, + train_tokens: np.ndarray, + val_tokens: np.ndarray, + token_limit: int, + steps: int, + seed: int, + args: argparse.Namespace, + device: torch.device, + dropout_fn: Callable[[int], float], + metrics_file, + trace_file, + stage: int | None = None, + model: DropoutGPT | None = None, + optimizer: torch.optim.Optimizer | None = None, + tokens_seen_start: int = 0, +) -> tuple[DropoutGPT, torch.optim.Optimizer, int, dict]: + if model is None: + set_seed(seed) + model = DropoutGPT(config).to(device) + optimizer = torch.optim.AdamW( + model.parameters(), + lr=args.lr, + betas=(0.9, 0.95), + weight_decay=args.weight_decay, + ) + else: + torch.manual_seed(seed + 10_000 + (stage or 0)) + if optimizer is None: + raise ValueError("optimizer is required when reusing a model") + + model.train() + dropout_code = int(round(condition.initial * 10_000)) + rng = np.random.default_rng(batch_seed(seed, model_spec, dropout_code, stage)) + tokens_seen = tokens_seen_start + last_loss = float("nan") + active_dropout = condition.initial + t0 = time.time() + + for step in range(1, steps + 1): + active_dropout = dropout_fn(tokens_seen) + model.set_dropout(active_dropout) + x, y = make_batch( + train_tokens, token_limit, args.batch_size, args.block_size, rng, device + ) + _, loss = model(x, y) + optimizer.zero_grad(set_to_none=True) + loss.backward() + if args.grad_clip > 0: + torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) + optimizer.step() + + tokens_seen += args.batch_size * args.block_size + last_loss = float(loss.item()) + if args.log_every > 0 and step % args.log_every == 0: + write_jsonl_row( + trace_file, + { + "event": "train_step", + "run_mode": run_mode, + "condition": condition.name, + "model_name": model_spec.name, + "seed": seed, + "stage": stage, + "step": step, + "steps": steps, + "token_limit": int(token_limit), + "tokens_seen": int(tokens_seen), + "dropout": float(active_dropout), + "train_batch_loss": last_loss, + }, + ) + if args.eval_every > 0 and step % args.eval_every == 0: + train_eval = estimate_loss( + model, + train_tokens, + token_limit, + args.trace_eval_batches, + args, + device, + rng_seed=seed + 20_000 + step, + ) + val_eval = estimate_loss( + model, + val_tokens, + len(val_tokens), + args.trace_eval_batches, + args, + device, + rng_seed=seed + 30_000 + step, + ) + write_jsonl_row( + trace_file, + { + "event": "eval_step", + "run_mode": run_mode, + "condition": condition.name, + "model_name": model_spec.name, + "seed": seed, + "stage": stage, + "step": step, + "steps": steps, + "token_limit": int(token_limit), + "tokens_seen": int(tokens_seen), + "dropout": float(active_dropout), + "train_eval_loss": train_eval, + "val_eval_loss": val_eval, + "generalization_gap": val_eval - train_eval, + }, + ) + + train_eval = estimate_loss( + model, + train_tokens, + token_limit, + args.train_eval_batches, + args, + device, + rng_seed=seed + 40_000 + (stage or 0), + ) + val_eval = estimate_loss( + model, + val_tokens, + len(val_tokens), + args.eval_batches, + args, + device, + rng_seed=seed + 50_000 + (stage or 0), + ) + row = { + "run_mode": run_mode, + "condition": condition.name, + "condition_kind": condition.kind, + "seed": seed, + "stage": stage, + "token_limit": int(token_limit), + "steps": int(steps), + "tokens_seen": int(tokens_seen), + "dropout_initial": float(condition.initial), + "dropout_final": float(condition.final), + "dropout_schedule": condition.schedule, + "dropout_active_final": float(active_dropout), + "train_loss_last": last_loss, + "train_eval_loss": train_eval, + "val_eval_loss": val_eval, + "eval_loss": val_eval, + "generalization_gap": val_eval - train_eval, + "elapsed_sec": time.time() - t0, + "parameters": model.num_parameters(), + "model_config": config.to_dict(), + **model_spec.to_dict(), + } + write_jsonl_row(metrics_file, row) + return model, optimizer, tokens_seen, row diff --git a/src/dropout_decay/model.py b/src/dropout_decay/model.py index 4608d0c3511efe393395934c3d23a61be6d20094..e5aa9e3056190c4114acd62feb2b52bfcce17a9d 100644 --- a/src/dropout_decay/model.py +++ b/src/dropout_decay/model.py @@ -16,212 +16,24 @@ The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. """ -from __future__ import annotations - -from dataclasses import asdict, dataclass -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F - - -@dataclass(frozen=True) -class GPTConfig: - block_size: int - vocab_size: int - n_layer: int - n_head: int - n_embd: int - dropout: float - - def to_dict(self) -> dict[str, int | float]: - return asdict(self) - - -def rms_norm(x: torch.Tensor) -> torch.Tensor: - return F.rms_norm(x, (x.size(-1),)) - - -class Linear(nn.Linear): - """Bias-free linear layer matching the nanochat architectural style.""" - - def __init__(self, in_features: int, out_features: int): - super().__init__(in_features, out_features, bias=False) - - -def apply_rotary_emb( - x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor -) -> torch.Tensor: - half = x.shape[-1] // 2 - x1, x2 = x[..., :half], x[..., half:] - return torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1) - - -class CausalSelfAttention(nn.Module): - def __init__(self, config: GPTConfig): - super().__init__() - if config.n_embd % config.n_head != 0: - raise ValueError("n_embd must be divisible by n_head") - self.n_head = config.n_head - self.head_dim = config.n_embd // config.n_head - if self.head_dim % 2 != 0: - raise ValueError("rotary attention requires an even head dimension") - self.dropout_p = config.dropout - self.c_q = Linear(config.n_embd, config.n_embd) - self.c_k = Linear(config.n_embd, config.n_embd) - self.c_v = Linear(config.n_embd, config.n_embd) - self.c_proj = Linear(config.n_embd, config.n_embd) - self.resid_dropout = nn.Dropout(config.dropout) - mask = torch.ones(config.block_size, config.block_size, dtype=torch.bool).tril() - self.register_buffer( - "causal_mask", - mask.view(1, 1, config.block_size, config.block_size), - persistent=False, - ) - - def set_dropout(self, p: float) -> None: - self.dropout_p = p - self.resid_dropout.p = p - - def forward( - self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor - ) -> torch.Tensor: - batch, seq_len, channels = x.shape - q = self.c_q(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2) - k = self.c_k(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2) - v = self.c_v(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2) - - q = rms_norm(apply_rotary_emb(q, cos, sin)) * 1.2 - k = rms_norm(apply_rotary_emb(k, cos, sin)) * 1.2 - - att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) - mask = self.causal_mask[:, :, :seq_len, :seq_len] - att = att.masked_fill(~mask, float("-inf")) - att = F.softmax(att.float(), dim=-1).to(dtype=x.dtype) - att = F.dropout(att, p=self.dropout_p, training=self.training) - y = att @ v - y = y.transpose(1, 2).contiguous().view(batch, seq_len, channels) - return self.resid_dropout(self.c_proj(y)) - - -class MLP(nn.Module): - def __init__(self, config: GPTConfig): - super().__init__() - self.c_fc = Linear(config.n_embd, 4 * config.n_embd) - self.c_proj = Linear(4 * config.n_embd, config.n_embd) - self.dropout = nn.Dropout(config.dropout) - - def set_dropout(self, p: float) -> None: - self.dropout.p = p - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.c_fc(x) - x = F.relu(x).square() - x = self.c_proj(x) - return self.dropout(x) - - -class Block(nn.Module): - def __init__(self, config: GPTConfig): - super().__init__() - self.attn = CausalSelfAttention(config) - self.mlp = MLP(config) - - def set_dropout(self, p: float) -> None: - self.attn.set_dropout(p) - self.mlp.set_dropout(p) - - def forward( - self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor - ) -> torch.Tensor: - x = x + self.attn(rms_norm(x), cos, sin) - x = x + self.mlp(rms_norm(x)) - return x - - -class DropoutGPT(nn.Module): - """Minimal nanochat-style causal Transformer with dynamic dropout control.""" - - def __init__(self, config: GPTConfig): - super().__init__() - self.config = config - self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd) - self.embedding_dropout = nn.Dropout(config.dropout) - self.blocks = nn.ModuleList(Block(config) for _ in range(config.n_layer)) - self.lm_head = Linear(config.n_embd, config.vocab_size) - self.register_buffer( - "cos", - torch.empty(1, 1, config.block_size, config.n_embd // config.n_head // 2), - persistent=False, - ) - self.register_buffer( - "sin", - torch.empty(1, 1, config.block_size, config.n_embd // config.n_head // 2), - persistent=False, - ) - self._init_weights() - self._init_rotary() - - @torch.no_grad() - def _init_weights(self) -> None: - nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.8) - nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.001) - scale = (3.0 ** 0.5) * (self.config.n_embd ** -0.5) - for block in self.blocks: - nn.init.uniform_(block.attn.c_q.weight, -scale, scale) - nn.init.uniform_(block.attn.c_k.weight, -scale, scale) - nn.init.uniform_(block.attn.c_v.weight, -scale, scale) - nn.init.zeros_(block.attn.c_proj.weight) - nn.init.uniform_(block.mlp.c_fc.weight, -0.4 * scale, 0.4 * scale) - nn.init.zeros_(block.mlp.c_proj.weight) - - @torch.no_grad() - def _init_rotary(self) -> None: - head_dim = self.config.n_embd // self.config.n_head - channel_range = torch.arange( - 0, - head_dim, - 2, - dtype=torch.float32, - device=self.cos.device, - ) - inv_freq = 1.0 / (100000 ** (channel_range / head_dim)) - positions = torch.arange( - self.config.block_size, - dtype=torch.float32, - device=self.cos.device, - ) - freqs = torch.outer(positions, inv_freq) - self.cos.copy_(freqs.cos().view(1, 1, self.config.block_size, head_dim // 2)) - self.sin.copy_(freqs.sin().view(1, 1, self.config.block_size, head_dim // 2)) - - def set_dropout(self, p: float) -> None: - if not 0.0 <= p < 1.0: - raise ValueError("dropout must satisfy 0 <= p < 1") - self.embedding_dropout.p = p - for block in self.blocks: - block.set_dropout(p) - - def num_parameters(self) -> int: - return sum(p.numel() for p in self.parameters()) - - def forward( - self, idx: torch.Tensor, targets: torch.Tensor | None = None - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - _, seq_len = idx.shape - if seq_len > self.config.block_size: - raise ValueError("sequence length exceeds block_size") - cos = self.cos[:, :, :seq_len, :] - sin = self.sin[:, :, :seq_len, :] - x = self.embedding_dropout(rms_norm(self.token_embedding(idx))) - for block in self.blocks: - x = block(x, cos, sin) - logits = self.lm_head(rms_norm(x)).float() - if targets is None: - return logits - loss = F.cross_entropy( - logits.reshape(-1, logits.size(-1)), - targets.reshape(-1), - ) - return logits, loss +from dropout_decay.models.gpt import ( + Block, + CausalSelfAttention, + DropoutGPT, + GPTConfig, + Linear, + MLP, + apply_rotary_emb, + rms_norm, +) + +__all__ = [ + "Block", + "CausalSelfAttention", + "DropoutGPT", + "GPTConfig", + "Linear", + "MLP", + "apply_rotary_emb", + "rms_norm", +] diff --git a/src/dropout_decay/models/__init__.py b/src/dropout_decay/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e5aa9e3056190c4114acd62feb2b52bfcce17a9d --- /dev/null +++ b/src/dropout_decay/models/__init__.py @@ -0,0 +1,39 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from dropout_decay.models.gpt import ( + Block, + CausalSelfAttention, + DropoutGPT, + GPTConfig, + Linear, + MLP, + apply_rotary_emb, + rms_norm, +) + +__all__ = [ + "Block", + "CausalSelfAttention", + "DropoutGPT", + "GPTConfig", + "Linear", + "MLP", + "apply_rotary_emb", + "rms_norm", +] diff --git a/src/dropout_decay/models/gpt.py b/src/dropout_decay/models/gpt.py new file mode 100644 index 0000000000000000000000000000000000000000..4608d0c3511efe393395934c3d23a61be6d20094 --- /dev/null +++ b/src/dropout_decay/models/gpt.py @@ -0,0 +1,227 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +@dataclass(frozen=True) +class GPTConfig: + block_size: int + vocab_size: int + n_layer: int + n_head: int + n_embd: int + dropout: float + + def to_dict(self) -> dict[str, int | float]: + return asdict(self) + + +def rms_norm(x: torch.Tensor) -> torch.Tensor: + return F.rms_norm(x, (x.size(-1),)) + + +class Linear(nn.Linear): + """Bias-free linear layer matching the nanochat architectural style.""" + + def __init__(self, in_features: int, out_features: int): + super().__init__(in_features, out_features, bias=False) + + +def apply_rotary_emb( + x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> torch.Tensor: + half = x.shape[-1] // 2 + x1, x2 = x[..., :half], x[..., half:] + return torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1) + + +class CausalSelfAttention(nn.Module): + def __init__(self, config: GPTConfig): + super().__init__() + if config.n_embd % config.n_head != 0: + raise ValueError("n_embd must be divisible by n_head") + self.n_head = config.n_head + self.head_dim = config.n_embd // config.n_head + if self.head_dim % 2 != 0: + raise ValueError("rotary attention requires an even head dimension") + self.dropout_p = config.dropout + self.c_q = Linear(config.n_embd, config.n_embd) + self.c_k = Linear(config.n_embd, config.n_embd) + self.c_v = Linear(config.n_embd, config.n_embd) + self.c_proj = Linear(config.n_embd, config.n_embd) + self.resid_dropout = nn.Dropout(config.dropout) + mask = torch.ones(config.block_size, config.block_size, dtype=torch.bool).tril() + self.register_buffer( + "causal_mask", + mask.view(1, 1, config.block_size, config.block_size), + persistent=False, + ) + + def set_dropout(self, p: float) -> None: + self.dropout_p = p + self.resid_dropout.p = p + + def forward( + self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + ) -> torch.Tensor: + batch, seq_len, channels = x.shape + q = self.c_q(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2) + k = self.c_k(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2) + v = self.c_v(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2) + + q = rms_norm(apply_rotary_emb(q, cos, sin)) * 1.2 + k = rms_norm(apply_rotary_emb(k, cos, sin)) * 1.2 + + att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) + mask = self.causal_mask[:, :, :seq_len, :seq_len] + att = att.masked_fill(~mask, float("-inf")) + att = F.softmax(att.float(), dim=-1).to(dtype=x.dtype) + att = F.dropout(att, p=self.dropout_p, training=self.training) + y = att @ v + y = y.transpose(1, 2).contiguous().view(batch, seq_len, channels) + return self.resid_dropout(self.c_proj(y)) + + +class MLP(nn.Module): + def __init__(self, config: GPTConfig): + super().__init__() + self.c_fc = Linear(config.n_embd, 4 * config.n_embd) + self.c_proj = Linear(4 * config.n_embd, config.n_embd) + self.dropout = nn.Dropout(config.dropout) + + def set_dropout(self, p: float) -> None: + self.dropout.p = p + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.c_fc(x) + x = F.relu(x).square() + x = self.c_proj(x) + return self.dropout(x) + + +class Block(nn.Module): + def __init__(self, config: GPTConfig): + super().__init__() + self.attn = CausalSelfAttention(config) + self.mlp = MLP(config) + + def set_dropout(self, p: float) -> None: + self.attn.set_dropout(p) + self.mlp.set_dropout(p) + + def forward( + self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + ) -> torch.Tensor: + x = x + self.attn(rms_norm(x), cos, sin) + x = x + self.mlp(rms_norm(x)) + return x + + +class DropoutGPT(nn.Module): + """Minimal nanochat-style causal Transformer with dynamic dropout control.""" + + def __init__(self, config: GPTConfig): + super().__init__() + self.config = config + self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd) + self.embedding_dropout = nn.Dropout(config.dropout) + self.blocks = nn.ModuleList(Block(config) for _ in range(config.n_layer)) + self.lm_head = Linear(config.n_embd, config.vocab_size) + self.register_buffer( + "cos", + torch.empty(1, 1, config.block_size, config.n_embd // config.n_head // 2), + persistent=False, + ) + self.register_buffer( + "sin", + torch.empty(1, 1, config.block_size, config.n_embd // config.n_head // 2), + persistent=False, + ) + self._init_weights() + self._init_rotary() + + @torch.no_grad() + def _init_weights(self) -> None: + nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.8) + nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.001) + scale = (3.0 ** 0.5) * (self.config.n_embd ** -0.5) + for block in self.blocks: + nn.init.uniform_(block.attn.c_q.weight, -scale, scale) + nn.init.uniform_(block.attn.c_k.weight, -scale, scale) + nn.init.uniform_(block.attn.c_v.weight, -scale, scale) + nn.init.zeros_(block.attn.c_proj.weight) + nn.init.uniform_(block.mlp.c_fc.weight, -0.4 * scale, 0.4 * scale) + nn.init.zeros_(block.mlp.c_proj.weight) + + @torch.no_grad() + def _init_rotary(self) -> None: + head_dim = self.config.n_embd // self.config.n_head + channel_range = torch.arange( + 0, + head_dim, + 2, + dtype=torch.float32, + device=self.cos.device, + ) + inv_freq = 1.0 / (100000 ** (channel_range / head_dim)) + positions = torch.arange( + self.config.block_size, + dtype=torch.float32, + device=self.cos.device, + ) + freqs = torch.outer(positions, inv_freq) + self.cos.copy_(freqs.cos().view(1, 1, self.config.block_size, head_dim // 2)) + self.sin.copy_(freqs.sin().view(1, 1, self.config.block_size, head_dim // 2)) + + def set_dropout(self, p: float) -> None: + if not 0.0 <= p < 1.0: + raise ValueError("dropout must satisfy 0 <= p < 1") + self.embedding_dropout.p = p + for block in self.blocks: + block.set_dropout(p) + + def num_parameters(self) -> int: + return sum(p.numel() for p in self.parameters()) + + def forward( + self, idx: torch.Tensor, targets: torch.Tensor | None = None + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + _, seq_len = idx.shape + if seq_len > self.config.block_size: + raise ValueError("sequence length exceeds block_size") + cos = self.cos[:, :, :seq_len, :] + sin = self.sin[:, :, :seq_len, :] + x = self.embedding_dropout(rms_norm(self.token_embedding(idx))) + for block in self.blocks: + x = block(x, cos, sin) + logits = self.lm_head(rms_norm(x)).float() + if targets is None: + return logits + loss = F.cross_entropy( + logits.reshape(-1, logits.size(-1)), + targets.reshape(-1), + ) + return logits, loss diff --git a/src/dropout_decay/scheduler.py b/src/dropout_decay/scheduler.py index 536ab4ab13e7ade069ad192e400f38d7aa4e2ece..99b07cf00ea188a8b95975f7ec11d8cdf03b2421 100644 --- a/src/dropout_decay/scheduler.py +++ b/src/dropout_decay/scheduler.py @@ -16,41 +16,9 @@ The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. """ -from __future__ import annotations +from dropout_decay.schedules.dropout import DropoutDecayConfig, DropoutDecayScheduler -from dataclasses import dataclass -import math - - -@dataclass(frozen=True) -class DropoutDecayConfig: - initial_dropout: float = 0.80 - final_dropout: float = 0.10 - decay_tokens: int = 40_000_000 - schedule: str = "cosine" - - -class DropoutDecayScheduler: - """Smoothly maps total processed/observed tokens to an active dropout rate.""" - - def __init__(self, config: DropoutDecayConfig): - if not 0.0 <= config.final_dropout <= config.initial_dropout < 1.0: - raise ValueError("dropout values must satisfy 0 <= final <= initial < 1") - if config.decay_tokens <= 0: - raise ValueError("decay_tokens must be positive") - if config.schedule not in {"cosine", "smoothstep", "linear"}: - raise ValueError("schedule must be cosine, smoothstep, or linear") - self.config = config - - def value(self, tokens_seen: int) -> float: - progress = min(max(tokens_seen / self.config.decay_tokens, 0.0), 1.0) - if self.config.schedule == "cosine": - mix = 0.5 * (1.0 + math.cos(math.pi * progress)) - elif self.config.schedule == "smoothstep": - smooth = progress * progress * (3.0 - 2.0 * progress) - mix = 1.0 - smooth - else: - mix = 1.0 - progress - return self.config.final_dropout + ( - self.config.initial_dropout - self.config.final_dropout - ) * mix +__all__ = [ + "DropoutDecayConfig", + "DropoutDecayScheduler", +] diff --git a/src/dropout_decay/schedules/__init__.py b/src/dropout_decay/schedules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..99b07cf00ea188a8b95975f7ec11d8cdf03b2421 --- /dev/null +++ b/src/dropout_decay/schedules/__init__.py @@ -0,0 +1,24 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from dropout_decay.schedules.dropout import DropoutDecayConfig, DropoutDecayScheduler + +__all__ = [ + "DropoutDecayConfig", + "DropoutDecayScheduler", +] diff --git a/src/dropout_decay/schedules/dropout.py b/src/dropout_decay/schedules/dropout.py new file mode 100644 index 0000000000000000000000000000000000000000..536ab4ab13e7ade069ad192e400f38d7aa4e2ece --- /dev/null +++ b/src/dropout_decay/schedules/dropout.py @@ -0,0 +1,56 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import math + + +@dataclass(frozen=True) +class DropoutDecayConfig: + initial_dropout: float = 0.80 + final_dropout: float = 0.10 + decay_tokens: int = 40_000_000 + schedule: str = "cosine" + + +class DropoutDecayScheduler: + """Smoothly maps total processed/observed tokens to an active dropout rate.""" + + def __init__(self, config: DropoutDecayConfig): + if not 0.0 <= config.final_dropout <= config.initial_dropout < 1.0: + raise ValueError("dropout values must satisfy 0 <= final <= initial < 1") + if config.decay_tokens <= 0: + raise ValueError("decay_tokens must be positive") + if config.schedule not in {"cosine", "smoothstep", "linear"}: + raise ValueError("schedule must be cosine, smoothstep, or linear") + self.config = config + + def value(self, tokens_seen: int) -> float: + progress = min(max(tokens_seen / self.config.decay_tokens, 0.0), 1.0) + if self.config.schedule == "cosine": + mix = 0.5 * (1.0 + math.cos(math.pi * progress)) + elif self.config.schedule == "smoothstep": + smooth = progress * progress * (3.0 - 2.0 * progress) + mix = 1.0 - smooth + else: + mix = 1.0 - progress + return self.config.final_dropout + ( + self.config.initial_dropout - self.config.final_dropout + ) * mix diff --git a/src/dropout_decay/specs/__init__.py b/src/dropout_decay/specs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7c4f955e59c19e42b7dcd718a32a974c618e6d3f --- /dev/null +++ b/src/dropout_decay/specs/__init__.py @@ -0,0 +1,26 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from dropout_decay.specs.dropout_condition import DropoutCondition, anchor_dropout +from dropout_decay.specs.model_spec import ModelSpec + +__all__ = [ + "DropoutCondition", + "ModelSpec", + "anchor_dropout", +] diff --git a/src/dropout_decay/specs/dropout_condition.py b/src/dropout_decay/specs/dropout_condition.py new file mode 100644 index 0000000000000000000000000000000000000000..6ad104cd94f9afbe6524daa533961b75f4b6fb0d --- /dev/null +++ b/src/dropout_decay/specs/dropout_condition.py @@ -0,0 +1,89 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import math +from typing import Callable + +from dropout_decay.schedules import DropoutDecayConfig, DropoutDecayScheduler + + +@dataclass(frozen=True) +class DropoutCondition: + name: str + kind: str + initial: float + final: float + schedule: str = "constant" + decay_tokens: int | None = None + anchors: tuple[tuple[int, float], ...] = () + + def to_dict(self) -> dict: + return { + "name": self.name, + "kind": self.kind, + "initial": self.initial, + "final": self.final, + "schedule": self.schedule, + "decay_tokens": self.decay_tokens, + "anchors": [list(anchor) for anchor in self.anchors], + } + + def make_fn( + self, + fallback_decay_tokens: int, + unique_tokens: int | None = None, + ) -> Callable[[int], float]: + if self.kind == "static": + return lambda _tokens_seen, p=self.initial: p + if self.kind == "anchor_decay": + if unique_tokens is None: + raise ValueError("anchor_decay conditions require unique_tokens") + p = anchor_dropout(unique_tokens, self.anchors) + return lambda _tokens_seen, p=p: p + scheduler = DropoutDecayScheduler( + DropoutDecayConfig( + initial_dropout=self.initial, + final_dropout=self.final, + decay_tokens=self.decay_tokens or fallback_decay_tokens, + schedule=self.schedule, + ) + ) + return scheduler.value + + +def anchor_dropout(unique_tokens: int, anchors: tuple[tuple[int, float], ...]) -> float: + if not anchors: + raise ValueError("anchor dropout schedule requires at least one anchor") + ordered = sorted(anchors) + if unique_tokens <= ordered[0][0]: + return ordered[0][1] + if unique_tokens >= ordered[-1][0]: + return ordered[-1][1] + log_unique = math.log(unique_tokens) + for (left_tokens, left_dropout), (right_tokens, right_dropout) in zip( + ordered, ordered[1:] + ): + if left_tokens <= unique_tokens <= right_tokens: + left_log = math.log(left_tokens) + right_log = math.log(right_tokens) + mix = (log_unique - left_log) / (right_log - left_log) + return left_dropout + mix * (right_dropout - left_dropout) + return ordered[-1][1] diff --git a/src/dropout_decay/specs/model_spec.py b/src/dropout_decay/specs/model_spec.py new file mode 100644 index 0000000000000000000000000000000000000000..e0ca6e2750721971f1d0ffa6e6bbd65ca2b36f18 --- /dev/null +++ b/src/dropout_decay/specs/model_spec.py @@ -0,0 +1,49 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from dropout_decay.models import GPTConfig + + +@dataclass(frozen=True) +class ModelSpec: + name: str + n_layer: int + n_head: int + n_embd: int + + def config(self, vocab_size: int, block_size: int, dropout: float) -> GPTConfig: + return GPTConfig( + block_size=block_size, + vocab_size=vocab_size, + n_layer=self.n_layer, + n_head=self.n_head, + n_embd=self.n_embd, + dropout=dropout, + ) + + def to_dict(self) -> dict[str, int | str]: + return { + "model_name": self.name, + "n_layer": self.n_layer, + "n_head": self.n_head, + "n_embd": self.n_embd, + } diff --git a/src/dropout_decay/tokenization/__init__.py b/src/dropout_decay/tokenization/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..97f35117a9adc9befdccd47892379f8d77f57320 --- /dev/null +++ b/src/dropout_decay/tokenization/__init__.py @@ -0,0 +1,25 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from dropout_decay.tokenization.bpe import BpeTokenizer, SPECIAL_TOKENS, SPLIT_PATTERN + +__all__ = [ + "BpeTokenizer", + "SPECIAL_TOKENS", + "SPLIT_PATTERN", +] diff --git a/src/dropout_decay/tokenization/bpe.py b/src/dropout_decay/tokenization/bpe.py new file mode 100644 index 0000000000000000000000000000000000000000..874803dee204888846dbb806124085efd4c6a136 --- /dev/null +++ b/src/dropout_decay/tokenization/bpe.py @@ -0,0 +1,105 @@ +""" +Derived from Andrej Karpathy's nanochat project. + +MIT License + +Copyright (c) 2025 Andrej Karpathy + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Iterable + +from tokenizers import Regex +from tokenizers import Tokenizer as HFTokenizer +from tokenizers import decoders, pre_tokenizers +from tokenizers.models import BPE +from tokenizers.trainers import BpeTrainer + +SPECIAL_TOKENS = [ + "<|bos|>", + "<|user_start|>", + "<|user_end|>", + "<|assistant_start|>", + "<|assistant_end|>", + "<|python_start|>", + "<|python_end|>", + "<|output_start|>", + "<|output_end|>", +] + +SPLIT_PATTERN = ( + r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,2}|""" + r""" ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+""" +) + + +class BpeTokenizer: + """Minimal HuggingFace BPE wrapper following nanochat's GPT-4-style splitter.""" + + def __init__(self, tokenizer: HFTokenizer): + self.tokenizer = tokenizer + + @classmethod + def train_from_iterator( + cls, text_iterator: Iterable[str], vocab_size: int + ) -> "BpeTokenizer": + tokenizer = HFTokenizer(BPE(byte_fallback=True, unk_token=None, fuse_unk=False)) + tokenizer.normalizer = None + tokenizer.pre_tokenizer = pre_tokenizers.Sequence( + [ + pre_tokenizers.Split( + pattern=Regex(SPLIT_PATTERN), + behavior="isolated", + invert=False, + ), + pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False), + ] + ) + tokenizer.decoder = decoders.ByteLevel() + trainer = BpeTrainer( + vocab_size=vocab_size, + show_progress=True, + min_frequency=0, + initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), + special_tokens=SPECIAL_TOKENS, + ) + tokenizer.train_from_iterator(text_iterator, trainer) + return cls(tokenizer) + + @classmethod + def from_file(cls, path: str | Path) -> "BpeTokenizer": + return cls(HFTokenizer.from_file(str(path))) + + def save(self, path: str | Path) -> None: + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + self.tokenizer.save(str(path)) + + @property + def vocab_size(self) -> int: + return self.tokenizer.get_vocab_size() + + @property + def bos_id(self) -> int: + bos = self.tokenizer.token_to_id("<|bos|>") + if bos is None: + raise ValueError("tokenizer is missing <|bos|>") + return bos + + def encode(self, text: str, prepend_bos: bool = False) -> list[int]: + ids = self.tokenizer.encode(text, add_special_tokens=False).ids + if prepend_bos: + ids.insert(0, self.bos_id) + return ids diff --git a/src/dropout_decay/tokenizer.py b/src/dropout_decay/tokenizer.py index 874803dee204888846dbb806124085efd4c6a136..97f35117a9adc9befdccd47892379f8d77f57320 100644 --- a/src/dropout_decay/tokenizer.py +++ b/src/dropout_decay/tokenizer.py @@ -16,90 +16,10 @@ The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. """ -from __future__ import annotations +from dropout_decay.tokenization.bpe import BpeTokenizer, SPECIAL_TOKENS, SPLIT_PATTERN -from pathlib import Path -from typing import Iterable - -from tokenizers import Regex -from tokenizers import Tokenizer as HFTokenizer -from tokenizers import decoders, pre_tokenizers -from tokenizers.models import BPE -from tokenizers.trainers import BpeTrainer - -SPECIAL_TOKENS = [ - "<|bos|>", - "<|user_start|>", - "<|user_end|>", - "<|assistant_start|>", - "<|assistant_end|>", - "<|python_start|>", - "<|python_end|>", - "<|output_start|>", - "<|output_end|>", +__all__ = [ + "BpeTokenizer", + "SPECIAL_TOKENS", + "SPLIT_PATTERN", ] - -SPLIT_PATTERN = ( - r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,2}|""" - r""" ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+""" -) - - -class BpeTokenizer: - """Minimal HuggingFace BPE wrapper following nanochat's GPT-4-style splitter.""" - - def __init__(self, tokenizer: HFTokenizer): - self.tokenizer = tokenizer - - @classmethod - def train_from_iterator( - cls, text_iterator: Iterable[str], vocab_size: int - ) -> "BpeTokenizer": - tokenizer = HFTokenizer(BPE(byte_fallback=True, unk_token=None, fuse_unk=False)) - tokenizer.normalizer = None - tokenizer.pre_tokenizer = pre_tokenizers.Sequence( - [ - pre_tokenizers.Split( - pattern=Regex(SPLIT_PATTERN), - behavior="isolated", - invert=False, - ), - pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False), - ] - ) - tokenizer.decoder = decoders.ByteLevel() - trainer = BpeTrainer( - vocab_size=vocab_size, - show_progress=True, - min_frequency=0, - initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), - special_tokens=SPECIAL_TOKENS, - ) - tokenizer.train_from_iterator(text_iterator, trainer) - return cls(tokenizer) - - @classmethod - def from_file(cls, path: str | Path) -> "BpeTokenizer": - return cls(HFTokenizer.from_file(str(path))) - - def save(self, path: str | Path) -> None: - path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - self.tokenizer.save(str(path)) - - @property - def vocab_size(self) -> int: - return self.tokenizer.get_vocab_size() - - @property - def bos_id(self) -> int: - bos = self.tokenizer.token_to_id("<|bos|>") - if bos is None: - raise ValueError("tokenizer is missing <|bos|>") - return bos - - def encode(self, text: str, prepend_bos: bool = False) -> list[int]: - ids = self.tokenizer.encode(text, add_special_tokens=False).ids - if prepend_bos: - ids.insert(0, self.bos_id) - return ids