batched newtonschulz implementation

.gitignore

@@ -0,0 +1,4 @@
+fineweb10B/
+pylog124M/
+__pycache__/
+logs/

Dockerfile

@@ -0,0 +1,33 @@
+FROM nvidia/cuda:12.6.2-cudnn-devel-ubuntu24.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHON_VERSION=3.12.7
+ENV PATH=/usr/local/bin:$PATH
+
+RUN apt update && apt install -y --no-install-recommends build-essential libssl-dev zlib1g-dev \
+    libbz2-dev libreadline-dev libsqlite3-dev curl git libncursesw5-dev xz-utils tk-dev libxml2-dev \
+    libxmlsec1-dev libffi-dev liblzma-dev \
+    && apt clean && rm -rf /var/lib/apt/lists/*
+
+RUN curl -O https://www.python.org/ftp/python/${PYTHON_VERSION}/Python-${PYTHON_VERSION}.tgz && \
+    tar -xzf Python-${PYTHON_VERSION}.tgz && \
+    cd Python-${PYTHON_VERSION} && \
+    ./configure --enable-optimizations && \
+    make -j$(nproc) && \
+    make altinstall && \
+    cd .. && \
+    rm -rf Python-${PYTHON_VERSION} Python-${PYTHON_VERSION}.tgz
+
+RUN ln -s /usr/local/bin/python3.12 /usr/local/bin/python && \
+    ln -s /usr/local/bin/pip3.12 /usr/local/bin/pip
+
+COPY requirements.txt /modded-nanogpt/requirements.txt
+WORKDIR /modded-nanogpt
+
+RUN python -m pip install --upgrade pip && \
+    pip install -r requirements.txt
+
+RUN pip install --pre torch==2.6.0.dev20241203+cu124 --index-url https://download.pytorch.org/whl/nightly/cu124 --upgrade
+
+CMD ["bash"]
+ENTRYPOINT []

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Keller Jordan
+
+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.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

data/cached_fineweb100B.py

@@ -0,0 +1,16 @@
+import os
+import sys
+from huggingface_hub import hf_hub_download
+# Download the GPT-2 tokens of Fineweb100B from huggingface. This
+# saves about an hour of startup time compared to regenerating them.
+def get(fname):
+    local_dir = os.path.join(os.path.dirname(__file__), 'fineweb100B')
+    if not os.path.exists(os.path.join(local_dir, fname)):
+        hf_hub_download(repo_id="kjj0/fineweb100B-gpt2", filename=fname,
+                        repo_type="dataset", local_dir=local_dir)
+get("fineweb_val_%06d.bin" % 0)
+num_chunks = 1030 # full fineweb100B. Each chunk is 100M tokens
+if len(sys.argv) >= 2: # we can pass an argument to download less
+    num_chunks = int(sys.argv[1])
+for i in range(1, num_chunks+1):
+    get("fineweb_train_%06d.bin" % i)

data/cached_fineweb10B.py

@@ -0,0 +1,16 @@
+import os
+import sys
+from huggingface_hub import hf_hub_download
+# Download the GPT-2 tokens of Fineweb10B from huggingface. This
+# saves about an hour of startup time compared to regenerating them.
+def get(fname):
+    local_dir = os.path.join(os.path.dirname(__file__), 'fineweb10B')
+    if not os.path.exists(os.path.join(local_dir, fname)):
+        hf_hub_download(repo_id="kjj0/fineweb10B-gpt2", filename=fname,
+                        repo_type="dataset", local_dir=local_dir)
+get("fineweb_val_%06d.bin" % 0)
+num_chunks = 103 # full fineweb10B. Each chunk is 100M tokens
+if len(sys.argv) >= 2: # we can pass an argument to download less
+    num_chunks = int(sys.argv[1])
+for i in range(1, num_chunks+1):
+    get("fineweb_train_%06d.bin" % i)

data/requirements.txt

@@ -0,0 +1,2 @@
+datasets
+tiktoken

records/060624_AdamW/README.md

@@ -0,0 +1,8 @@
+This is the log for my baseline AdamW training to which I compared the new Muon and SOAP optimizers.
+
+just the log, which is in the old llm.c format ("tel" lines are val loss)
+
+this was batch size 2^19, so ~5B tokens
+
+was learning rate 0.0018, warmup=250, warmdown=2000, betas=(0.9, 0.95) IIRC
+