IQ.Lvbs/IQ.Pilot
8
7
Fork 2
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

IQ.Pilot Release Commit @ bec7652

Browse Source
This commit is contained in:
IQ.Lvbs history cleanup
2026-08-22 23:42:41 -05:00
commit 58039e647c
4603 changed files with 1236178 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
artifacts/package_sources/tinygrad/.coveragerc Normal file
View File

@@ -0,0 +1,3 @@
[run]
source = tinygrad
branch = True

71
artifacts/package_sources/tinygrad/.gitignore vendored Normal file
View File

@@ -0,0 +1,71 @@
__pycache__
.venv/
.venv-*/
.vscode
.DS_Store
notebooks
.*.swp
.*.swo
*.pyc
*.so
*.txt
build
!examples/tinychat/assets/cdn.jsdelivr.net/npm/purecss@3.0.0/build/
/dist
*.egg-info
/env
a.out
boxes.jpg
pandecode.dump
vertex.bin
recognize*
.idea
*.prof
extra/disassemblers/applegpu
extra/datasets/cifar-10-python.tar.gz
extra/datasets/librispeech/
extra/datasets/imagenet/
extra/datasets/wiki/
extra/datasets/kits19
extra/datasets/kits19/
extra/datasets/squad/
extra/datasets/img_align_celeba*
extra/datasets/open-images-v6-mlperf
extra/datasets/kits/
extra/datasets/COCO/
extra/datasets/audio*
extra/huggingface_onnx/models/*
extra/huggingface_onnx/*.yaml
extra/weights
venv
venv_sd_mlperf
examples/**/net.*[js,json]
examples/**/*.safetensors
node_modules
package.json
package-lock.json
temp
*.csv
.coverage
coverage.xml
htmlcov
outputs_yolov8
wandb
model.safetensors
quickstart.py
.hypothesis
weights
*.lprof
comgr_*
*.pkl
!extra/sqtt/examples/**/*.pkl
site/
profile_stats
*.log
target
.mypy_cache
mutants
.mutmut-cache
dagre/
graphlib/
uv.lock

1
artifacts/package_sources/tinygrad/.iqpilot-revision Normal file
View File

@@ -0,0 +1 @@
95bfc521351d3acba7338605686155a9ee8b0009

34
artifacts/package_sources/tinygrad/.pre-commit-config.yaml Normal file
View File

@@ -0,0 +1,34 @@
# on Windows -- $env:SKIP="tests,example"
repos:
- repo: local
hooks:
- id: ruff
name: ruff
entry: python3 -m ruff check .
language: system
always_run: true
pass_filenames: false
- id: tiny
name: tiny tests
entry: python3 -m pytest test/test_tiny.py
language: system
always_run: true
pass_filenames: false
- id: mypy
name: mypy
entry: python3 -m mypy
language: system
always_run: true
pass_filenames: false
- id: example
name: test all devices
entry: python3 test/external/external_test_example.py
language: system
always_run: true
pass_filenames: false
- id: tests
name: comprehensive test suite
entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/backend/test_ops.py test/backend/test_schedule.py test/unit/test_assign.py test/backend/test_tensor.py test/backend/test_jit.py test/unit/test_schedule_cache.py test/null/test_pattern_matcher.py test/null/test_uop_symbolic.py test/unit/test_helpers.py
language: system
always_run: true
pass_filenames: false

469
artifacts/package_sources/tinygrad/.pylintrc Normal file
View File

@@ -0,0 +1,469 @@
[MASTER]
# A comma-separated list of package or module names from where C extensions may
# be loaded. Extensions are loading into the active Python interpreter and may
# run arbitrary code
extension-pkg-whitelist=scipy,cereal.messaging.messaging_pyx,PyQt5,av
# Add files or directories to the blacklist. They should be base names, not
# paths.
ignore=CVS,autogen,msm_kgsl.py,runtime,.venv
# Add files or directories matching the regex patterns to the blacklist. The
# regex matches against base names, not paths.
ignore-patterns=
# Python code to execute, usually for sys.path manipulation such as
# pygtk.require().
#init-hook=
# Use multiple processes to speed up Pylint.
jobs=4
# List of plugins (as comma separated values of python modules names) to load,
# usually to register additional checkers.
load-plugins=
# Pickle collected data for later comparisons.
persistent=yes
# Specify a configuration file.
#rcfile=
# Allow loading of arbitrary C extensions. Extensions are imported into the
# active Python interpreter and may run arbitrary code.
unsafe-load-any-extension=no
[MESSAGES CONTROL]
# Only show warnings with the listed confidence levels. Leave empty to show
# all. Valid levels: HIGH, INFERENCE, INFERENCE_FAILURE, UNDEFINED
confidence=
# Disable the message, report, category or checker with the given id(s). You
# can either give multiple identifiers separated by comma (,) or put this
# option multiple times (only on the command line, not in the configuration
# file where it should appear only once).You can also use "--disable=all" to
# disable everything first and then reenable specific checks. For example, if
# you want to run only the similarities checker, you can use "--disable=all
# --enable=similarities". If you want to run only the classes checker, but have
# no Warning level messages displayed, use"--disable=all --enable=classes
# --disable=W"
disable=C,R,W0613,W0511,W0212,W0201,W0106,W0603,W0621,W0703,W1201,W1203,E1136,W1514,E1101,W0221,W0105,E0401,abstract-method,W0707
# E1101 for function binding
# W0221 for Function class
# W0105 for comment strings
# E0401 for missing imports
# W0707 for not reraising
# Enable the message, report, category or checker with the given id(s). You can
# either give multiple identifier separated by comma (,) or put this option
# multiple time (only on the command line, not in the configuration file where
# it should appear only once). See also the "--disable" option for examples.
enable=c-extension-no-member,use-a-generator, no-else-return
[REPORTS]
# Python expression which should return a note less than 10 (10 is the highest
# note). You have access to the variables errors warning, statement which
# respectively contain the number of errors / warnings messages and the total
# number of statements analyzed. This is used by the global evaluation report
# (RP0004).
evaluation=10.0 - ((float(5 * error + warning + refactor + convention) / statement) * 10)
# Template used to display messages. This is a python new-style format string
# used to format the message information. See doc for all details
#msg-template=
# Set the output format. Available formats are text, parseable, colorized, json
# and msvs (visual studio).You can also give a reporter class, eg
# mypackage.mymodule.MyReporterClass.
output-format=text
# Tells whether to display a full report or only the messages
reports=no
# Activate the evaluation score.
score=yes
[REFACTORING]
# Maximum number of nested blocks for function / method body
max-nested-blocks=5
# Complete name of functions that never returns. When checking for
# inconsistent-return-statements if a never returning function is called then
# it will be considered as an explicit return statement and no message will be
# printed.
never-returning-functions=optparse.Values,sys.exit
[LOGGING]
# Logging modules to check that the string format arguments are in logging
# function parameter format
logging-modules=logging
[SPELLING]
# Limits count of emitted suggestions for spelling mistakes
max-spelling-suggestions=4
# Spelling dictionary name. Available dictionaries: none. To make it working
# install python-enchant package.
spelling-dict=
# List of comma separated words that should not be checked.
spelling-ignore-words=
# A path to a file that contains private dictionary; one word per line.
spelling-private-dict-file=
# Tells whether to store unknown words to indicated private dictionary in
# --spelling-private-dict-file option instead of raising a message.
spelling-store-unknown-words=no
[MISCELLANEOUS]
# List of note tags to take in consideration, separated by a comma.
notes=FIXME,
XXX,
TODO
[SIMILARITIES]
# Ignore comments when computing similarities.
ignore-comments=yes
# Ignore docstrings when computing similarities.
ignore-docstrings=yes
# Ignore imports when computing similarities.
ignore-imports=no
# Minimum lines number of a similarity.
min-similarity-lines=4
[TYPECHECK]
# List of decorators that produce context managers, such as
# contextlib.contextmanager. Add to this list to register other decorators that
# produce valid context managers.
contextmanager-decorators=contextlib.contextmanager
# List of members which are set dynamically and missed by pylint inference
# system, and so shouldn't trigger E1101 when accessed. Python regular
# expressions are accepted.
generated-members=capnp.* cereal.* pygame.* zmq.* setproctitle.* smbus2.* usb1.* serial.* cv2.* ft4222.* carla.*
# Tells whether missing members accessed in mixin class should be ignored. A
# mixin class is detected if its name ends with "mixin" (case insensitive).
ignore-mixin-members=yes
# This flag controls whether pylint should warn about no-member and similar
# checks whenever an opaque object is returned when inferring. The inference
# can return multiple potential results while evaluating a Python object, but
# some branches might not be evaluated, which results in partial inference. In
# that case, it might be useful to still emit no-member and other checks for
# the rest of the inferred objects.
ignore-on-opaque-inference=yes
# List of class names for which member attributes should not be checked (useful
# for classes with dynamically set attributes). This supports the use of
# qualified names.
ignored-classes=optparse.Values,thread._local,_thread._local
# List of module names for which member attributes should not be checked
# (useful for modules/projects where namespaces are manipulated during runtime
# and thus existing member attributes cannot be deduced by static analysis. It
# supports qualified module names, as well as Unix pattern matching.
ignored-modules=flask setproctitle usb1 flask.ext.socketio smbus2 usb1.*
# Show a hint with possible names when a member name was not found. The aspect
# of finding the hint is based on edit distance.
missing-member-hint=yes
# The minimum edit distance a name should have in order to be considered a
# similar match for a missing member name.
missing-member-hint-distance=1
# The total number of similar names that should be taken in consideration when
# showing a hint for a missing member.
missing-member-max-choices=1
[VARIABLES]
# List of additional names supposed to be defined in builtins. Remember that
# you should avoid to define new builtins when possible.
additional-builtins=
# Tells whether unused global variables should be treated as a violation.
allow-global-unused-variables=yes
# List of strings which can identify a callback function by name. A callback
# name must start or end with one of those strings.
callbacks=cb_,
_cb
# A regular expression matching the name of dummy variables (i.e. expectedly
# not used).
dummy-variables-rgx=_+$|(_[a-zA-Z0-9_]*[a-zA-Z0-9]+?$)|dummy|^ignored_|^unused_
# Argument names that match this expression will be ignored. Default to name
# with leading underscore
ignored-argument-names=_.*|^ignored_|^unused_
# Tells whether we should check for unused import in __init__ files.
init-import=no
# List of qualified module names which can have objects that can redefine
# builtins.
redefining-builtins-modules=six.moves,past.builtins,future.builtins
[FORMAT]
# Expected format of line ending, e.g. empty (any line ending), LF or CRLF.
expected-line-ending-format=
# Regexp for a line that is allowed to be longer than the limit.
ignore-long-lines=^\s*(# )?<?https?://\S+>?$
# Number of spaces of indent required inside a hanging or continued line.
indent-after-paren=4
# String used as indentation unit. This is usually " " (4 spaces) or "\t" (1
# tab).
indent-string=' '
# Maximum number of characters on a single line.
max-line-length=150
# Maximum number of lines in a module
max-module-lines=1000
# Allow the body of a class to be on the same line as the declaration if body
# contains single statement.
single-line-class-stmt=no
# Allow the body of an if to be on the same line as the test if there is no
# else.
single-line-if-stmt=no
[BASIC]
# Naming style matching correct argument names
argument-naming-style=snake_case
# Regular expression matching correct argument names. Overrides argument-
# naming-style
#argument-rgx=
# Naming style matching correct attribute names
attr-naming-style=snake_case
# Regular expression matching correct attribute names. Overrides attr-naming-
# style
#attr-rgx=
# Bad variable names which should always be refused, separated by a comma
bad-names=foo,
bar,
baz,
toto,
tutu,
tata
# Naming style matching correct class attribute names
class-attribute-naming-style=any
# Regular expression matching correct class attribute names. Overrides class-
# attribute-naming-style
#class-attribute-rgx=
# Naming style matching correct class names
class-naming-style=PascalCase
# Regular expression matching correct class names. Overrides class-naming-style
#class-rgx=
# Naming style matching correct constant names
const-naming-style=UPPER_CASE
# Regular expression matching correct constant names. Overrides const-naming-
# style
#const-rgx=
# Minimum line length for functions/classes that require docstrings, shorter
# ones are exempt.
docstring-min-length=-1
# Naming style matching correct function names
function-naming-style=snake_case
# Regular expression matching correct function names. Overrides function-
# naming-style
#function-rgx=
# Good variable names which should always be accepted, separated by a comma
good-names=i,
j,
k,
ex,
Run,
_
# Include a hint for the correct naming format with invalid-name
include-naming-hint=no
# Naming style matching correct inline iteration names
inlinevar-naming-style=any
# Regular expression matching correct inline iteration names. Overrides
# inlinevar-naming-style
#inlinevar-rgx=
# Naming style matching correct method names
method-naming-style=snake_case
# Regular expression matching correct method names. Overrides method-naming-
# style
#method-rgx=
# Naming style matching correct module names
module-naming-style=snake_case
# Regular expression matching correct module names. Overrides module-naming-
# style
#module-rgx=
# Colon-delimited sets of names that determine each other's naming style when
# the name regexes allow several styles.
name-group=
# Regular expression which should only match function or class names that do
# not require a docstring.
no-docstring-rgx=^_
# List of decorators that produce properties, such as abc.abstractproperty. Add
# to this list to register other decorators that produce valid properties.
property-classes=abc.abstractproperty
# Naming style matching correct variable names
variable-naming-style=snake_case
# Regular expression matching correct variable names. Overrides variable-
# naming-style
#variable-rgx=
[DESIGN]
# Maximum number of arguments for function / method
max-args=5
# Maximum number of attributes for a class (see R0902).
max-attributes=7
# Maximum number of boolean expressions in a if statement
max-bool-expr=5
# Maximum number of branch for function / method body
max-branches=12
# Maximum number of locals for function / method body
max-locals=15
# Maximum number of parents for a class (see R0901).
max-parents=7
# Maximum number of public methods for a class (see R0904).
max-public-methods=20
# Maximum number of return / yield for function / method body
max-returns=6
# Maximum number of statements in function / method body
max-statements=50
# Minimum number of public methods for a class (see R0903).
min-public-methods=2
[CLASSES]
# List of method names used to declare (i.e. assign) instance attributes.
defining-attr-methods=__init__,
__new__,
setUp
# List of member names, which should be excluded from the protected access
# warning.
exclude-protected=_asdict,
_fields,
_replace,
_source,
_make
# List of valid names for the first argument in a class method.
valid-classmethod-first-arg=cls
# List of valid names for the first argument in a metaclass class method.
valid-metaclass-classmethod-first-arg=mcs
[IMPORTS]
# Allow wildcard imports from modules that define __all__.
allow-wildcard-with-all=no
# Analyse import fallback blocks. This can be used to support both Python 2 and
# 3 compatible code, which means that the block might have code that exists
# only in one or another interpreter, leading to false positives when analysed.
analyse-fallback-blocks=no
# Deprecated modules which should not be used, separated by a comma
deprecated-modules=regsub,
TERMIOS,
Bastion,
rexec
# Create a graph of external dependencies in the given file (report RP0402 must
# not be disabled)
ext-import-graph=
# Create a graph of every (i.e. internal and external) dependencies in the
# given file (report RP0402 must not be disabled)
import-graph=
# Create a graph of internal dependencies in the given file (report RP0402 must
# not be disabled)
int-import-graph=
# Force import order to recognize a module as part of the standard
# compatibility libraries.
known-standard-library=
# Force import order to recognize a module as part of a third party library.
known-third-party=enchant
[STRING]
# This flag controls whether the implicit-str-concat should generate a warning
# on implicit string concatenation in sequences defined over several lines.
check-str-concat-over-line-jumps=yes
[EXCEPTIONS]
# Exceptions that will emit a warning when being caught. Defaults to
# "Exception"
overgeneral-exceptions=builtins.Exception

6
artifacts/package_sources/tinygrad/AGENTS.md Normal file
View File

@@ -0,0 +1,6 @@
# Notes
- Run tests with `-n12` for speed (e.g. `python -m pytest test/null/test_dtype.py -x -q -n12`)
- Run `python -m mypy tinygrad/` to typecheck
- Run `python -m ruff check .` to lint
- Read `./tinygrad/viz/README.md` for profiling and debugging rewrite rules

7
artifacts/package_sources/tinygrad/LICENSE Normal file
View File

@@ -0,0 +1,7 @@
Copyright (c) 2024, the tiny corp
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.

205
artifacts/package_sources/tinygrad/README.md Normal file
View File

@@ -0,0 +1,205 @@
<div align="center">
<picture>
<source media="(prefers-color-scheme: light)" srcset="/docs/logo_tiny_light.svg">
<img alt="tiny corp logo" src="/docs/logo_tiny_dark.svg" width="50%" height="50%">
</picture>
tinygrad: For something between [PyTorch](https://github.com/pytorch/pytorch) and [karpathy/micrograd](https://github.com/karpathy/micrograd). Maintained by [tiny corp](https://tinygrad.org).
<h3>
[Homepage](https://github.com/tinygrad/tinygrad) | [Documentation](https://docs.tinygrad.org/) | [Discord](https://discord.gg/ZjZadyC7PK)
</h3>
[![GitHub Repo stars](https://img.shields.io/github/stars/tinygrad/tinygrad)](https://github.com/tinygrad/tinygrad/stargazers)
[![Unit Tests](https://github.com/tinygrad/tinygrad/actions/workflows/test.yml/badge.svg)](https://github.com/tinygrad/tinygrad/actions/workflows/test.yml)
[![Discord](https://img.shields.io/discord/1068976834382925865)](https://discord.gg/ZjZadyC7PK)
</div>
---
tinygrad is an end-to-end deep learning stack:
- **Tensor library** with autograd
- **IR and compiler** that fuse and lower kernels
- **JIT + graph execution**
- **nn / optim / datasets** for real training
It’s inspired by PyTorch (ergonomics), JAX (functional transforms and IR-based AD), and TVM (scheduling and codegen), but stays intentionally tiny and hackable.
---
## How tinygrad compares
**PyTorch**
- ✅ Similar: eager `Tensor` API, autograd, `optim`, basic datasets and layers.
- ✅ You can write familiar training loops.
- 🔁 Unlike PyTorch, the entire compiler and IR are visible and hackable.
**JAX**
- ✅ IR-based autodiff over primitives (like JAXPR + XLA).
- ✅ Function-level JIT (`TinyJit`) that captures and replays kernels.
- 🔁 Fewer functional transforms (no full `vmap`/`pmap` yet), but far easier to read.
**TVM**
- ✅ Multiple lowering passes, scheduling, and BEAM search over kernels.
- ✅ Device “graphs” for batched execution.
- 🔁 tinygrad also ships the **front-end framework** (tensors, nn, optim), not just the compiler.
---
### Laziness
Try a matmul. See how, despite the style, it is fused into one kernel with the power of laziness.
```sh
DEBUG=3 python3 -c "from tinygrad import Tensor;
N = 1024; a, b = Tensor.empty(N, N), Tensor.empty(N, N);
(a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2).realize()"
```
And we can change `DEBUG` to `4` to see the generated code.
### Neural networks
As it turns out, 90% of what you need for neural networks are a decent autograd/tensor library.
Throw in an optimizer, a data loader, and some compute, and you have all you need.
```python
from tinygrad import Tensor, nn, Context
class LinearNet:
def __init__(self):
self.l1 = Tensor.kaiming_uniform(784, 128)
self.l2 = Tensor.kaiming_uniform(128, 10)
def __call__(self, x:Tensor) -> Tensor:
return x.flatten(1).dot(self.l1).relu().dot(self.l2)
model = LinearNet()
optim = nn.optim.Adam([model.l1, model.l2], lr=0.001)
x, y = Tensor.rand(4, 1, 28, 28), Tensor([2,4,3,7]) # replace with real mnist dataloader
with Context(TRAINING=1):
for i in range(10):
optim.zero_grad()
loss = model(x).sparse_categorical_crossentropy(y).backward()
optim.step()
print(i, loss.item())
```
See [examples/beautiful_mnist.py](examples/beautiful_mnist.py) for the full version that gets 98% in ~5 seconds
## Accelerators
tinygrad already supports numerous accelerators, including:
- [x] [OpenCL](tinygrad/runtime/ops_cl.py)
- [x] [CPU](tinygrad/runtime/ops_cpu.py)
- [x] [METAL](tinygrad/runtime/ops_metal.py)
- [x] [CUDA](tinygrad/runtime/ops_cuda.py)
- [x] [AMD](tinygrad/runtime/ops_amd.py)
- [x] [NV](tinygrad/runtime/ops_nv.py)
- [x] [QCOM](tinygrad/runtime/ops_qcom.py)
- [x] [WEBGPU](tinygrad/runtime/ops_webgpu.py)
And it is easy to add more! Your accelerator of choice only needs to support a total of ~25 low level ops.
To check default accelerator run: `python3 -c "from tinygrad import Device; print(Device.DEFAULT)"`
## Installation
The current recommended way to install tinygrad is from source.
### From source
```sh
git clone https://github.com/tinygrad/tinygrad.git
cd tinygrad
python3 -m pip install -e .
```
### Direct (master)
```sh
python3 -m pip install git+https://github.com/tinygrad/tinygrad.git
```
## Documentation
Documentation along with a quick start guide can be found on the [docs website](https://docs.tinygrad.org/) built from the [docs/](/docs) directory.
### Quick example comparing to PyTorch
```python
from tinygrad import Tensor
x = Tensor.eye(3)
y = Tensor([[2.0,0,-2.0]])
z = y.matmul(x).sum()
z.backward()
print(x.grad.tolist()) # dz/dx
print(y.grad.tolist()) # dz/dy
```
The same thing but in PyTorch:
```python
import torch
x = torch.eye(3, requires_grad=True)
y = torch.tensor([[2.0,0,-2.0]], requires_grad=True)
z = y.matmul(x).sum()
z.backward()
print(x.grad.tolist()) # dz/dx
print(y.grad.tolist()) # dz/dy
```
## Contributing
There has been a lot of interest in tinygrad lately. Following these guidelines will help your PR get accepted. If you do submit a PR, please include a sentence or two about why you want this merged and why you think it will improve the project.
If you are a new contributor with something that looks even close to AI written, it will be closed without feedback and you may be banned from our GitHub. No human should waste time reading AI slop. And for everyone, if you used AI, disclose what you used it for.
We'll start with what will get your PR closed with a pointer to this section:
- No code golf! While low line count is a guiding light of this project, anything that remotely looks like code golf will be closed. The true goal is reducing complexity and increasing readability, and deleting `\n`s does nothing to help with that.
- All docs and whitespace changes will be closed unless you are a well-known contributor. The people writing the docs should be those who know the codebase the absolute best. People who have not demonstrated that shouldn't be messing with docs. Whitespace changes are both useless *and* carry a risk of introducing bugs.
- Anything you claim is a "speedup" must be benchmarked. In general, the goal is simplicity, so even if your PR makes things marginally faster, you have to consider the tradeoff with maintainability and readability.
- In general, the code outside the core `tinygrad/` folder is not well tested, so unless the current code there is broken, you shouldn't be changing it.
- If your PR looks "complex", is a big diff, or adds lots of lines, it won't be reviewed or merged. Consider breaking it up into smaller PRs that are individually clear wins. A common pattern I see is prerequisite refactors before adding new functionality. If you can (cleanly) refactor to the point that the feature is a 3 line change, this is great, and something easy for us to review.
Now, what we want:
- Bug fixes (with a regression test) are great! This library isn't 1.0 yet, so if you stumble upon a bug, fix it, write a test, and submit a PR, this is valuable work.
- Solving bounties! tinygrad [offers cash bounties](https://docs.google.com/spreadsheets/d/1WKHbT-7KOgjEawq5h5Ic1qUWzpfAzuD_J06N1JwOCGs/edit?usp=sharing) for certain improvements to the library. All new code should be high quality and well tested.
- Features. However, if you are adding a feature, consider the line tradeoff. If it's 3 lines, there's less of a bar of usefulness it has to meet over something that's 30 or 300 lines. All features must have regression tests. In general with no other constraints, your feature's API should match torch or numpy.
- Refactors that are clear wins. In general, if your refactor isn't a clear win it will be closed. But some refactors are amazing! Think about readability in a deep core sense. A whitespace change or moving a few functions around is useless, but if you realize that two 100 line functions can actually use the same 110 line function with arguments while also improving readability, this is a big win. Refactors should pass [process replay](#process-replay-tests).
- Tests/fuzzers. If you can add tests that are non brittle, they are welcome. We have some fuzzers in here too, and there's a plethora of bugs that can be found with them and by improving them. Finding bugs, even writing broken tests (that should pass) with `@unittest.expectedFailure` is great. This is how we make progress.
- Dead code removal from core `tinygrad/` folder. We don't care about the code in extra, but removing dead code from the core library is great. Less for new people to read and be confused by.
### Running tests
You should install the pre-commit hooks with `pre-commit install`. This will run the linter, mypy, and a subset of the tests on every commit.
For more examples on how to run the full test suite please refer to the [CI workflow](.github/workflows/test.yml).
Some examples of running tests locally:
```sh
python3 -m pip install -e '.[testing]' # install extra deps for testing
python3 test/backend/test_ops.py # just the ops tests
python3 -m pytest test/ # whole test suite
```
For agents, always run tests with `-n12` for speed.
#### Process replay tests
[Process replay](https://github.com/tinygrad/tinygrad/blob/master/test/external/process_replay/README.md) compares your PR's generated kernels against master. If your PR is a refactor or speedup without any expected behavior change, It should include [pr] in the pull request title.

10
artifacts/package_sources/tinygrad/conftest.py Normal file
View File

@@ -0,0 +1,10 @@
import os, pytest, signal, threading
@pytest.hookimpl(wrapper=True)
def pytest_runtest_call(item):
t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 300)), os.kill, args=(os.getpid(), signal.SIGABRT))
t.start()
try: yield
finally:
t.cancel()
t.join()

1
artifacts/package_sources/tinygrad/docs/CNAME Normal file
View File

@@ -0,0 +1 @@
docs.tinygrad.org

54
artifacts/package_sources/tinygrad/docs/abstractions3.py Normal file
View File

@@ -0,0 +1,54 @@
# abstractions2 goes from back to front, here we will go from front to back
# *****
# 0. Load mnist on the device
from tinygrad.nn.datasets import mnist
X_train, Y_train, _, _ = mnist()
X_train = X_train.float()
X_train -= X_train.mean()
# *****
# 1. Define an MNIST model.
from tinygrad import Tensor, Context
l1 = Tensor.kaiming_uniform(128, 784)
l2 = Tensor.kaiming_uniform(10, 128)
def model(x): return x.flatten(1).dot(l1.T).relu().dot(l2.T)
l1n, l2n = l1.numpy(), l2.numpy()
# *****
# 2. Choose a batch for training and do the backward pass.
from tinygrad.nn.optim import SGD
optim = SGD([l1, l2])
with Context(TRAINING=1):
X, Y = X_train[(samples:=Tensor.randint(128, high=X_train.shape[0]))], Y_train[samples]
optim.zero_grad()
model(X).sparse_categorical_crossentropy(Y).backward()
optim.schedule_step() # this will step the optimizer without running realize
# *****
# 3. Create a schedule (linear uop).
# The weight Tensors have been assigned to, but not yet realized. Everything is still lazy at this point
# l1.uop and l2.uop define a computation graph
from tinygrad.engine.realize import run_linear
linear = Tensor.schedule_linear(l1, l2)
print(f"The schedule contains {len(linear.src)} items.")
for call in linear.src: print(str(call)[:80])
# *****
# 4. Lower and run the schedule (linear uop).
run_linear(linear)
# *****
# 5. Print the weight change
print("first weight change\n", l1.numpy()-l1n)
print("second weight change\n", l2.numpy()-l2n)

251
artifacts/package_sources/tinygrad/docs/abstractions4.py Normal file
View File

@@ -0,0 +1,251 @@
# tinygrad allows you to write kernels at many different abstractions levels.
# This is for RDNA3, but if you don't have one you can run with the emulator
# PYTHONPATH="." DEV=MOCKPCI+AMD
from tinygrad import Tensor, Context, GlobalCounters, UOp, Device
from tinygrad.helpers import DEV, DEBUG, getenv
from tinygrad.uop.ops import AxisType, KernelInfo, Ops
from tinygrad.dtype import AddrSpace, dtypes
from tinygrad.runtime.autogen.amd.rdna3.ins import *
def eval_harness(name, tensor, fxn, check=None):
print(f"***** {name}")
GlobalCounters.reset()
with Context(DEBUG=max(DEBUG.value, 2)): out = fxn(tensor).item()
assert check is None or abs(out - check) < abs(check) * 1e-3, f"out was wrong {out}, expected {check}, off by {out/check}x"
print(f"computed in {GlobalCounters.time_sum_s*1000:.2f} ms, {(a.nbytes()/1e9)/GlobalCounters.time_sum_s:.2f} GB/s")
return out
SZ = 256*1024 if DEV.interface.startswith("MOCK") else 1024*1024*1024
def example_2_hip(a:Tensor, correct):
GLOBALS = 1024
THREADS = 256
def hip_reduce_sum(out:UOp, buf:UOp) -> UOp:
assert SZ % (GLOBALS * THREADS) == 0
CHUNK = SZ // (GLOBALS * THREADS)
# NOTE: tinygrad doesn't populate HIP hidden kernargs, so blockDim.x/gridDim.x read as 0.
# We hardcode block/grid sizes as constexpr to avoid any dependency on those builtins.
code = f"""
#include <hip/hip_runtime.h>
constexpr unsigned int BLOCK = {THREADS};
constexpr unsigned int CHUNK = {CHUNK};
extern "C" __global__ void hip_reduce_sum_kernel(float* __restrict__ block_sums, const float* __restrict__ x) {{
__shared__ float sdata[BLOCK];
unsigned int tid = threadIdx.x;
unsigned int gid = blockIdx.x * BLOCK + tid;
// Each thread sums CHUNK consecutive elements from its own region
float sum = 0.0f;
const float* base = x + gid * CHUNK;
#pragma unroll 16
for (unsigned int k = 0; k < CHUNK; k++) {{
sum += base[k];
}}
sdata[tid] = sum;
__syncthreads();
// Block reduction in shared memory
for (unsigned int s = BLOCK / 2; s > 0; s >>= 1) {{
if (tid < s) {{
sdata[tid] += sdata[tid + s];
}}
__syncthreads();
}}
// One partial sum per block
if (tid == 0) {{
block_sums[blockIdx.x] = sdata[0];
}}
}}"""
# TODO: remove the need for the compiler here, you should just be able to remove Ops.BINARY
from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
lib = HIPCCCompiler(Device[Device.DEFAULT].renderer.target.arch, []).compile_cached(code)
# the sink specifies the GLOBAL and LOCAL sizes, along with the input buffers and name
sink = UOp.sink(UOp.special(GLOBALS, 'gidx0'), UOp.special(THREADS, 'lidx0'), out, buf,
arg=KernelInfo(name="hip_reduce_sum_kernel"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
eval_harness("HIP kernel", a, lambda x: Tensor.empty(GLOBALS).custom_kernel(x, fxn=hip_reduce_sum)[0].sum(), check=correct)
def example_3_custom_uop(a:Tensor, correct):
# This GPU has 32 CUs, keep them all busy
CU_COUNT = 32
def custom_sum(out:UOp, buf:UOp) -> UOp:
LCLS = 256
buf = buf.reshape(CU_COUNT, -1, LCLS)
glbl = UOp.range(CU_COUNT, 0, AxisType.GLOBAL)
lane = UOp.range(LCLS, 1, AxisType.LOCAL)
# accumulate the globals into a per lane accumulator
reduce_loop = UOp.range(buf.shape[1], 2, AxisType.REDUCE)
acc = UOp.placeholder((1,), dtypes.float, slot=6, addrspace=AddrSpace.REG)
acc = acc.after(acc.store(0))
acc = acc.after(acc[0].store(acc.after(reduce_loop)[0] + buf[glbl, reduce_loop, lane]).end(reduce_loop))
# store all the per lane accumulators to LOCAL
local_accs = UOp.placeholder((LCLS,), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
local_accs = local_accs.after(local_accs[lane].store(acc[0]))
# accumulate LOCALs into a single per CU accumulator
late_reduce_loop = UOp.range(LCLS, 3, AxisType.REDUCE)
acc2 = UOp.placeholder((1,), dtypes.float, slot=7, addrspace=AddrSpace.REG)
acc2 = acc2.after(acc2.store(0))
acc2 = acc2.after(acc2[0].store(acc2.after(late_reduce_loop)[0] + local_accs[late_reduce_loop]).end(late_reduce_loop))[0]
# store (NOTE: since the address doesn't depend on the warp, this will be automatically gated)
return out[glbl].store(acc2).end(lane, glbl).sink(arg=KernelInfo(opts_to_apply=()))
eval_harness("custom UOp kernel", a, lambda x: Tensor.empty(CU_COUNT).custom_kernel(x, fxn=custom_sum)[0].sum(), check=correct)
def example_5_custom_assembly(a:Tensor, correct):
# Kernel class copied from amd_asm_matmul
class Kernel:
def __init__(self): self.instructions, self.labels, self.pos = [], {}, 0
def label(self, name): self.labels[name] = self.pos
def emit(self, inst, target=None):
self.instructions.append(inst)
inst._target, inst._pos = target, self.pos
self.pos += inst.size()
return inst
def waitcnt(self, lgkm=None, vm=None):
# Wait for memory operations. lgkm=N waits until N lgkm ops remain, vm=N waits until N vmem ops remain.
vmcnt, lgkmcnt, expcnt = vm if vm is not None else 63, lgkm if lgkm is not None else 63, 7
waitcnt = (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
self.emit(s_waitcnt(simm16=waitcnt))
def finalize(self, sink:UOp) -> UOp:
for inst in self.instructions:
if inst._target is None: continue
offset_dwords = (self.labels[inst._target] - inst._pos - inst.size()) // 4
if not -32768 <= offset_dwords <= 32767: raise ValueError(f"branch to '{inst._target}' offset {offset_dwords} exceeds simm16 range")
inst.simm16 = offset_dwords
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in self.instructions]))))
CU_COUNT = 32
LANES = 64
def asm_sum(out:UOp, buf:UOp) -> UOp:
V_LANE_ID = 0 # lane_id set on startup
S_WORKGROUP_X = 2 # workgroup_id_x
S_LOOP_CTR = 3
k = Kernel()
# mul lane id by 16 for offsets (4 for float, 4 for b128)
k.emit(v_mul_lo_u32(v[0], v[V_LANE_ID], 16))
k.emit(v_add_nc_u32_e32(v[1], 4096, v[0]))
k.emit(v_add_nc_u32_e32(v[2], 4096, v[1]))
k.emit(v_add_nc_u32_e32(v[3], 4096, v[2]))
# load both addresses
k.emit(s_load_b128(sdata=s[4:7], sbase=s[0:1], offset=0x0, soffset=NULL))
k.waitcnt(lgkm=0)
# offset buffer pointer by workgroup_id_x * chunk_size_bytes
k.emit(s_mul_i32(s[S_LOOP_CTR], s[S_WORKGROUP_X], buf.numel()*4//CU_COUNT))
k.emit(s_add_u32(s[6], s[6], s[S_LOOP_CTR]))
k.emit(s_addc_u32(s[7], s[7], 0))
# zero the accumulators
k.emit(VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[4], vdsty=v[5], srcx0=0, srcy0=0))
k.emit(VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[6], vdsty=v[7], srcx0=0, srcy0=0))
def emit_loads(base_vreg, reg_len):
assert reg_len%4 == 0
k.emit(s_clause(simm16=(reg_len//4)-1))
for i in range(reg_len//4):
offset = i*LANES*16
assert offset < 16384
k.emit(global_load_b128(vdst=v[base_vreg+i*4:base_vreg+i*4+3], addr=v[offset//4096], saddr=s[6:7], offset=offset%4096))
k.emit(s_add_u32(s[6], s[6], reg_len * LANES * 4))
k.emit(s_addc_u32(s[7], s[7], 0))
def tree_reduce_to_4567(base_vreg, reg_len):
assert reg_len%4 == 0
reg_len //= 4
while reg_len > 1:
half = reg_len // 2
for j in range(half):
a, b = base_vreg + j*4, base_vreg + (j+half)*4
# v[a+0](bank0) += v[b+2](bank2), v[a+1](bank1) += v[b+3](bank3) — src0 and src1 on different banks
k.emit(VOPD(VOPDOp.V_DUAL_ADD_F32, VOPDOp.V_DUAL_ADD_F32, vdstx=v[a], vdsty=v[a+1], srcx0=v[a], vsrcx1=v[b+2], srcy0=v[a+1], vsrcy1=v[b+3]))
# v[a+2](bank2) += v[b+0](bank0), v[a+3](bank3) += v[b+1](bank1) — src0 and src1 on different banks
k.emit(VOPD(VOPDOp.V_DUAL_ADD_F32, VOPDOp.V_DUAL_ADD_F32, vdstx=v[a+2], vdsty=v[a+3], srcx0=v[a+2], vsrcx1=v[b], srcy0=v[a+3], vsrcy1=v[b+1]))
reg_len = half
k.emit(VOPD(VOPDOp.V_DUAL_ADD_F32, VOPDOp.V_DUAL_ADD_F32, vdstx=v[4], vdsty=v[5], srcx0=v[4], vsrcx1=v[base_vreg], srcy0=v[5], vsrcy1=v[base_vreg+1]))
k.emit(VOPD(VOPDOp.V_DUAL_ADD_F32, VOPDOp.V_DUAL_ADD_F32, vdstx=v[6], vdsty=v[7], srcx0=v[6], vsrcx1=v[base_vreg+2], srcy0=v[7], vsrcy1=v[base_vreg+3]))
BASE_REG = 8
LOAD_UNROLL = 64
INNER_UNROLL = 2
assert buf.numel() % (CU_COUNT*LANES*LOAD_UNROLL*INNER_UNROLL) == 0
total_batches = buf.numel()//(CU_COUNT*LANES*LOAD_UNROLL*INNER_UNROLL)
k.emit(s_mov_b32(s[S_LOOP_CTR], total_batches-1))
k.label('LOOP')
for _ in range(INNER_UNROLL):
emit_loads(BASE_REG, reg_len=LOAD_UNROLL)
k.waitcnt(vm=0)
tree_reduce_to_4567(BASE_REG, reg_len=LOAD_UNROLL)
k.emit(s_sub_u32(s[S_LOOP_CTR], s[S_LOOP_CTR], 1))
k.emit(s_cbranch_scc0(), target='LOOP')
# add into v[4]
k.emit(v_add_f32_e32(v[4], v[4], v[5]))
k.emit(v_add_f32_e32(v[6], v[6], v[7]))
k.emit(v_add_f32_e32(v[4], v[4], v[6]))
# warp shuffle into v[4] on lane 0 using DPP row_shl within each 16-lane row
for shift in [1, 2, 4, 8]:
k.emit(v_add_f32_e32(v[4], DPP, v[4], vsrc0=v[4], dpp=0x100 | shift, row_mask=0xf, bank_mask=0xf, bc=1))
# combine rows: get lane 16's value to lane 0 via permlanex16
k.emit(v_permlanex16_b32(v[5], v[4], 0, 0))
k.emit(v_add_f32_e32(v[4], v[4], v[5]))
# atomic store (only on lane 0)
k.emit(s_mov_b32(EXEC_LO, 1))
k.emit(v_mov_b32_e32(v[0], 0))
k.emit(global_atomic_add_f32(addr=v[0], saddr=s[4:5], data=v[4]))
k.emit(s_sendmsg(simm16=3)) # DEALLOC_VGPRS
k.emit(s_endpgm())
return k.finalize(UOp.sink(UOp.special(CU_COUNT, 'gidx0'), UOp.special(LANES, 'lidx0'), out, buf, arg=KernelInfo(name="asm_reduce")))
out = Tensor.zeros(1,).contiguous().realize()
eval_harness("RDNA3 assembly kernel", a, lambda x: out.custom_kernel(x, fxn=asm_sum)[0], check=correct)
if __name__ == "__main__":
examples = [int(x) for x in getenv("EXAMPLES", "1,2,3,4,5").split(",")]
correct = None
# First define a Tensor and realize it. We will focus on a 1GB sum kernel on RDNA3
a = (Tensor.randn(SZ) if getenv("RAND") else Tensor.ones(SZ)).contiguous().realize()
if 1 in examples:
# *****
# This is the high level tinygrad way.
# Note that this is split into multiple kernels for speed.
correct = eval_harness("basic kernel", a, lambda x: x.sum())
if 2 in examples:
# *****
# You can import kernels from CUDA/HIP/Metal.
# ChatGPT is great at writing these Kernel
example_2_hip(a, correct)
if 3 in examples:
# *****
# Now we get to the lower abstraction layers of tinygrad.
# You can write a kernel in UOps, and it's 2.5x faster than normal.
example_3_custom_uop(a, correct)
if 4 in examples:
# *****
# You can also BEAM search stock tinygrad for a faster kernel.
# This does even better than all the kernels to date in this simple case.
with Context(BEAM=2):
eval_harness("BEAMed kernel", a, lambda x: x.sum(), check=correct)
if 5 in examples:
# *****
# If you really want to go crazy with speed, you can code in assembly.
# There's not too much to gain here over BEAM, but it's a few percent faster.
example_5_custom_assembly(a, correct)

39
artifacts/package_sources/tinygrad/docs/developer/am.md Normal file
View File

@@ -0,0 +1,39 @@
# AM Driver
AM driver is a userspace driver targeting AMD's RDNA3/RDNA4. You only need tinygrad to send compute tasks to your GPU!
## How to run?
Make sure that amdgpu module is unloaded and just run tinygrad with `DEV=AMD`!
Optional requirements:
* System without IOMMU for P2P / SDMA support
* vfio-pci module for IRQ handling
## Environment Variables
| Variable | Possible Value(s) | Description |
|----------|------------------|-------------|
| AM_RESET | [1] | Performs a full GPU reset (reloading all firmware and IP blocks) |
| AM_DEBUG | [0-4] | Sets the level of additional debugging information |
## AM Driver Details
### Compute & SDMA Queues
AM binds compute queues directly to MEC (bypassing MES). Tinygrad uses only one compute queue, which is bound at `pipe=0 queue=0`. Similarly, the single SDMA queue is bound at `engine=0 queue=0`.
### Boot
The GPU being passed can be in one of several states:
1. Not initialized
2. Initialized by amdgpu
3. Initialized by AM
The first and second states require a full GPU setup since their states are unknown. The second state also requires a mode1 reset to reinitialize all components.
The third state can be set up partially to optimize boot time. In this case, only the GFX and SDMA IPs need to be initialized. To enable this, AM uses a separate boot memory that is guaranteed not to be overwritten. This physical memory is utilized for all blocks that are initialized only during the initial AM boot. To determine if the GPU is in the third state, AM uses `regSCRATCH_REG7` as a flag.
### VM Management
Each AM device sets up only a single `VMID=0` and one page directory. The page directory used is 3-level and thus supports up to 512GB of virtual addresses. All AM devices are located in one virtual address space.

46
artifacts/package_sources/tinygrad/docs/developer/developer.md Normal file
View File

@@ -0,0 +1,46 @@
The tinygrad framework has four pieces
* a PyTorch like <b>frontend</b>.
* a <b>scheduler</b> which breaks the compute into kernels.
* a <b>lowering</b> engine which converts ASTs into code that can run on the accelerator.
* an <b>execution</b> engine which can run that code.
There is a good [bunch of tutorials](https://mesozoic-egg.github.io/tinygrad-notes/) by Di Zhu that go over tinygrad internals.
There's also a [doc describing speed](../developer/speed.md)
## Frontend
Everything in [Tensor](../tensor/index.md) is syntactic sugar around constructing a graph of [UOps](../developer/uop.md).
The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not all UOps will actually become realized. There's two types of UOps, base and view. base contains compute into a contiguous buffer, and view is a view. Inputs to a base can be either base or view, inputs to a view can only be a single base.
## Scheduling
The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/schedule/__init__.py) converts the graph of UOps into a `LINEAR` UOp whose `src` is a list of `CALL` UOps. One `CALL` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. The `CALL`'s `src[0]` (a `SINK` ast) specifies what compute to run, and the remaining `src` are the buffers to run it on.
## Lowering
The code in [realize](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/realize.py) lowers each `CALL` by compiling its ast into a `PROGRAM` and running it.
::: tinygrad.engine.realize.run_linear
There's a ton of complexity hidden behind this, see the `codegen/` directory.
First we lower the AST to UOps, which is a linear list of the compute to be run. This is where the BEAM search happens.
Then we render the UOps into code with a `Renderer`, then we compile the code to binary with a `Compiler`.
## Execution
`run_linear` walks the `LINEAR` UOp, dispatching each `CALL` to a runner (kernel, copy, view, encdec, or graph).
## Runtime
Runtimes are responsible for device-specific interactions. They handle tasks such as initializing devices, allocating memory, loading/launching programs, and more. You can find more information about the runtimes API on the [runtime overview page](runtime.md).
All runtime implementations can be found in the [runtime directory](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime).
### HCQ Compatible Runtimes
HCQ API is a lower-level API for defining runtimes. Interaction with HCQ-compatible devices occurs at a lower level, with commands issued directly to hardware queues. Some examples of such backends are [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) and [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py), which are userspace drivers for NVIDIA and AMD devices respectively. You can find more information about the API on [HCQ overview page](hcq.md)

128
artifacts/package_sources/tinygrad/docs/developer/hcq.md Normal file
View File

@@ -0,0 +1,128 @@
# HCQ Compatible Runtime
## Overview
The main aspect of HCQ-compatible runtimes is how they interact with devices. In HCQ, all interactions with devices occur in a hardware-friendly manner using [command queues](#command-queues). This approach allows commands to be issued directly to devices, bypassing runtime overhead such as HIP or CUDA. Additionally, by using the HCQ API, these runtimes can benefit from various optimizations and features, including [HCQGraph](#hcqgraph) and built-in profiling capabilities.
### Command Queues
To interact with devices you create a `HWQueue`. Some methods are required, like timestamp and synchronization methods like [signal](#tinygrad.runtime.support.hcq.HWQueue.signal) and [wait](#tinygrad.runtime.support.hcq.HWQueue.wait), while others are dependent on it being a compute or copy queue.
For example, the following Python code enqueues a wait, execute, and signal command on the HCQ-compatible device:
```python
HWQueue().wait(signal_to_wait, value_to_wait) \
.exec(program, args_state, global_dims, local_dims) \
.signal(signal_to_fire, value_to_fire) \
.submit(your_device)
```
Each runtime should implement the required functions that are defined in the `HWQueue` classes.
::: tinygrad.runtime.support.hcq.HWQueue
options:
members: [
"signal",
"wait",
"timestamp",
"bind",
"submit",
"memory_barrier",
"exec",
"copy",
]
show_source: false
### HCQ Compatible Device
The `HCQCompiled` class defines the API for HCQ-compatible devices. This class serves as an abstract base class that device-specific implementations should inherit from and implement.
::: tinygrad.runtime.support.hcq.HCQCompiled
options:
show_source: false
#### Signals
Signals are device-dependent structures used for synchronization and timing in HCQ-compatible devices. They should be designed to record both a `value` and a `timestamp` within the same signal. HCQ-compatible backend implementations should use `HCQSignal` as a base class.
::: tinygrad.runtime.support.hcq.HCQSignal
options:
members: [value, timestamp, wait]
show_source: false
The following Python code demonstrates the usage of signals:
```python
signal = your_device.new_signal(value=0)
HWQueue().timestamp(signal) \
.signal(signal, value_to_fire) \
.submit(your_device)
signal.wait(value_to_fire)
signaled_value = signal.value # should be the same as `value_to_fire`
timestamp = signal.timestamp
```
##### Synchronization signals
Each HCQ-compatible device must allocate two signals for global synchronization purposes. These signals are passed to the `HCQCompiled` base class during initialization: an active timeline signal `self.timeline_signal` and a shadow timeline signal `self._shadow_timeline_signal` which helps to handle signal value overflow issues. You can find more about synchronization in the [synchronization section](#synchronization)
### HCQ Compatible Allocator
The `HCQAllocator` base class simplifies allocator logic by leveraging [command queues](#command-queues) abstractions. This class efficiently handles copy and transfer operations, leaving only the alloc and free functions to be implemented by individual backends.
::: tinygrad.runtime.support.hcq.HCQAllocator
options:
members: [
"_alloc",
"_free",
]
show_source: false
#### HCQ Allocator Result Protocol
Backends must adhere to the `HCQBuffer` protocol when returning allocation results.
::: tinygrad.runtime.support.hcq.HCQBuffer
options:
members: true
show_source: false
### HCQ Compatible Program
`HCQProgram` is a base class for defining programs compatible with HCQ-enabled devices. It provides a flexible framework for handling different argument layouts (see `HCQArgsState`).
::: tinygrad.runtime.support.hcq.HCQProgram
options:
members: true
show_source: false
#### Arguments State
`HCQArgsState` is a base class for managing the argument state for HCQ programs. Backend implementations should create a subclass of `HCQArgsState` to manage arguments for the given program.
::: tinygrad.runtime.support.hcq.HCQArgsState
options:
members: true
show_source: false
**Lifetime**: The `HCQArgsState` is passed to `HWQueue.exec` and is guaranteed not to be freed until `HWQueue.submit` for the same queue is called.
### Synchronization
HCQ-compatible devices use a global timeline signal for synchronizing all operations. This mechanism ensures proper ordering and completion of tasks across the device. By convention, `self.timeline_value` points to the next value to signal. So, to wait for all previous operations on the device to complete, wait for `self.timeline_value - 1` value. The following Python code demonstrates the typical usage of signals to synchronize execution to other operations on the device:
```python
HWQueue().wait(your_device.timeline_signal, your_device.timeline_value - 1) \
.exec(...)
.signal(your_device.timeline_signal, your_device.next_timeline()) \
.submit(your_device)
# Optionally wait for execution
your_device.timeline_signal.wait(your_device.timeline_value - 1)
```
## HCQGraph
[HCQGraph](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/graph/hcq.py) is a core feature that implements `GraphRunner` for HCQ-compatible devices. `HCQGraph` builds static `HWQueue` for all operations per device. To optimize enqueue time, only the necessary parts of the queues are updated for each run using the symbolic variables, avoiding a complete rebuild.
Optionally, queues can implement a `bind` API, which allows further optimization by eliminating the need to copy the queues into the device ring.

60
artifacts/package_sources/tinygrad/docs/developer/layout.md Normal file
View File

@@ -0,0 +1,60 @@
# tinygrad directory layout
This explains the flow of a big graph down to programs.
Directories are listed in order of how they are processed.
---
## tinygrad/schedule
Group UOps into kernels.
::: tinygrad.schedule.rangeify.get_kernel_graph
options:
members: false
show_labels: false
show_source: false
---
## tinygrad/codegen/opt
Transforms the ast into an optimized ast. This is where BEAM search and heuristics live.
---
## tinygrad/codegen
Transform the optimized ast into a linearized and rendered program.
::: tinygrad.codegen.to_program
options:
members: false
show_labels: false
show_source: false
---
## tinygrad/renderer
Transform the linearized list of UOps into a program, represented as a string.
::: tinygrad.renderer.Renderer
options:
members:
- render
show_labels: false
show_source: false
---
## tinygrad/engine
Abstracted high level interface to the runtimes.
::: tinygrad.engine.realize.to_program
options:
members: false
show_labels: false
show_source: false

51
artifacts/package_sources/tinygrad/docs/developer/runtime.md Normal file
View File

@@ -0,0 +1,51 @@
# Runtime Overview
## Overview
A typical runtime consists of the following parts:
- [Compiled](#compiled)
- [Allocator](#allocator)
- [Program](#program)
- [Compiler](#compiler)
### Compiled
The `Compiled` class is responsible for initializing and managing a device.
::: tinygrad.device.Compiled
options:
members: [
"synchronize"
]
show_source: false
### Allocator
The `Allocator` class is responsible for managing memory on the device. There is also a version called the `LRUAllocator`, which caches allocated buffers to optimize performance.
::: tinygrad.device.Allocator
options:
members: true
show_source: false
::: tinygrad.device.LRUAllocator
options:
members: true
show_source: false
### Program
The `Program` class is created for each loaded program. It is responsible for executing the program on the device. As an example, here is a `CPUProgram` implementation which loads program and runs it.
::: tinygrad.runtime.ops_cpu.CPUProgram
options:
members: true
### Compiler
The `Compiler` class compiles the output from the `Renderer` and produces it in a device-specific format.
::: tinygrad.device.Compiler
options:
members: true

71
artifacts/package_sources/tinygrad/docs/developer/speed.md Normal file
View File

@@ -0,0 +1,71 @@
# speed in tinygrad
## Overview
Speed refers to many different things. To break it down to four, there's:
- Compile Speed (Python)
- Execution Speed (driver)
- Model Speed (scheduler)
- Kernel Speed (codegen)
## Compile Speed (Python)
This is how long the first run of your model takes. It's limited largely by the runtime of the Python doing UOp rewrites. Currently it's a bit slow, but on par with torch.compile. It gets even slower if you are using BEAM, since that's compiling many variants of each kernel.
This will be improved by writing faster graph_rewrite, doing less graph_rewrite, and better parallelization.
## Execution Speed (driver)
After your model is compiled, you are often using the `TinyJIT`. tinygrad has the best execution speed of any framework because it usually bypasses the GPU driver and prebuilds the command queue. It's tons faster than normal CUDA, and often even faster than CUDA Graph.
There's very little to improve here, as this is almost never the bottleneck.
## Model Speed (scheduler)
The scheduler determines how operations are grouped into kernels and which Tensors are written to memory. This is currently a big bottleneck of training speed.
The decisions are often not obvious. For example, when is it worth recomputing an arithmetic operation instead of storing and loading from memory? Example:
```python
from tinygrad import Tensor
a = Tensor.rand(100)
b = Tensor.rand(100)
c = Tensor.rand(100)
d = Tensor.rand(100)
out1 = a+b+c
out2 = a+b+d
Tensor.realize(out1, out2)
```
The real answer is obvious, compute both `out1` and `out2` in the same kernel. But you can't always do that. If you can't, should `a+b` first be saved to a subbuffer? Or should both the `out1` and `out2` kernels recompute `a+b`?
In this case: with recompute (6 reads + 2 writes), no recompute (6 reads + 3 writes), so we should probably recompute. However, once you add movement ops and casts this is even harder to figure out. tinygrad doesn't yet have a systematic way to do it.
## Kernel Speed (codegen)
Given that you have decided how the model ops will be grouped and what will be written to memory, kernel speed determines how fast that operation is done. This is what BEAM changes, it searches over a set of equivalent kernels which all perform the same operation and finds the one which performs the task the fastest.
In `kernel.py` we have a set of `OptOps`, these control the parameters of the speed optimizations applied to the kernel.
### Memory
The main bottleneck in most kernels is accessing memory. In a freshman algorithms class, you'll learn about cache aware matrix multiplication, and this is all forms of that. While the same math is run, the order in which you run it can have large impacts on the speed depending on if the data you are loading. OptOps will change this order.
Memory, even cache, is often much slower than accessing the register file. The amount of times data is used in math is called the "arithmetic intensity". For operations like BS=1 GEMV, the arithmetic intensity is 1, but for GEMMs and convs it can be much higher. OptOps like UPCAST and UNROLL can increase this, but be careful of making them too large, as if there's too much register pressure on the GPU the warp scheduler may not be able to fit many warps, or even worse, it could be spilling to local memory.
4090s have 1 TB/s of ram bandwidth and ~160 TFLOPS of compute, so you need to use each loaded value ~100 times. The L1 cache has around 40 TB/s of bandwidth, so in order to get full compute utilization you need to use each value ~4 times.
A lot of work can still be done here. For example, we never copy the inputs to on chip SRAM, but this is often quite helpful for kernel speed. Also, we aren't doing a good job with L2 cache awareness (the locals handle L1 quite well)
### Tensor Cores
Many accelerators have Tensor Cores / MAC arrays / systolic arrays. The main value of these is that, since they are 2-D, they create an n^2 ratio between the compute and the input data.
GPUs use Tensor Cores instead of MAC arrays to fit better in the GPU warp paradigm. This is because the output of Tensor Cores is O(n) wrt the input, while the output of MAC arrays is O(n^2)
We have a simple framework in tinygrad for adding these ALU blocks and achieving good performance from them.
### Indexing
Indexing determines the address of the memory we need to load. GPUs often have less integer math resources than floating point math, so this can sometimes be the bottleneck. We have a symbolic math engine in our rewrite rules to simplify indexing before it's emitted to the kernel. Newer NVIDIA GPUs have a "Tensor Memory Accelerator" to assist with fast indexing, however, this is not supported in tinygrad yet.

11
artifacts/package_sources/tinygrad/docs/developer/uop.md Normal file
View File

@@ -0,0 +1,11 @@
::: tinygrad.uop.ops.UOp
options:
members: false
members_order: source
show_labels: false
::: tinygrad.uop.ops.Ops
options:
members: true
members_order: source
show_labels: false

10
artifacts/package_sources/tinygrad/docs/dtypes.md Normal file
View File

@@ -0,0 +1,10 @@
::: tinygrad.dtype.DType
::: tinygrad.dtype.DTypes
options:
heading: dtypes
members: true
members_order: source
show_labels: false
::: tinygrad.dtype.ConstType

73
artifacts/package_sources/tinygrad/docs/env_vars.md Normal file
View File

@@ -0,0 +1,73 @@
# List of environment variables that control tinygrad behavior.
This is a list of environment variable that control the runtime behavior of tinygrad and its examples.
Most of these are self-explanatory, and are usually used to set an option at runtime.
Example: `DEV=CL DEBUG=4 python3 -m pytest`
However you can also decorate a function to set a value only inside that function.
```python
# in tensor.py (probably only useful if you are a tinygrad developer)
@Context(DEBUG=4)
def numpy(self) -> ...
```
Or use contextmanager to temporarily set a value inside some scope:
```python
with Context(DEBUG=0):
a = Tensor.ones(10, 10)
a *= 2
```
## Global Variables
The columns of this list are are: Variable, Possible Value(s) and Description.
- A `#` means that the variable can take any integer value.
These control the behavior of core tinygrad even when used as a library.
Variable | Possible Value(s) | Description
---|---|---
DEBUG | [1-7] | enable debugging output (operations, timings, speed, generated code and more)
DEV | [AMD, NV, ...] | enable a specific backend, see [below](#dev-variable)
BEAM | [#] | number of beams in kernel beam search
DEFAULT_FLOAT | [HALF, ...]| specify the default float dtype (FLOAT32, HALF, BFLOAT16, FLOAT64, ...), default to FLOAT32
IMAGE | [1] | enable 2d specific optimizations
FLOAT16 | [1] | use float16 for images instead of float32
JIT | [0-2] | 0=disabled, 1=[jit enabled](quickstart.md#jit) (default), 2=jit enabled, but graphs are disabled
VIZ | [1] | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz)
ALLOW_TF32 | [1] | enable TensorFloat-32 tensor cores on Ampere or newer GPUs.
WEBGPU_BACKEND | [WGPUBackendType_Metal, ...] | Force select a backend for WebGPU (Metal, DirectX, OpenGL, Vulkan...)
CUDA_PATH | str | Use `CUDA_PATH/include` for CUDA headers for CUDA and NV backends. If not set, TinyGrad will use `/usr/local/cuda/include`, `/usr/include` and `/opt/cuda/include`.
### DEV variable
The `DEV` variable deserves special note due to its more nuanced syntax.
`DEV` is used to specify the target device, target renderer and target architecture for said device, separated by colons.
Specifying the renderer and architecture is optional, omitting a preference will cause tinygrad to automatically determine a suitable setting.
The `DEV` variable may also be used to specify the interface through which to access the device (eg. `PCI`, `USB`). Interfaces may be specified preceding the target triple,
separated by a plus (eg. `DEV=USB+AMD:LLVM`). Similarly as above, the interface may be omitted. Example usage follows:
`DEV` contents | Interpretation
--- | ---
AMD | use the AMD device
AMD:LLVM | use the AMD device with the LLVM renderer
NV:CUDA:sm_70 | use the NV device with the CUDA renderer targetting sm_70
AMD::gfx950 | use the AMD device targetting gfx950
USB+AMD | use the AMD device over the USB interface
CPU:LLVM | use the CPU device with the LLVM renderer
CPU:LLVM:x86_64,znver2,avx2,-avx512f | use the CPU device with the LLVM renderer, with [additional arch flags](runtime.md#cpu-arch)
### Debug breakdown
Variable | Value | Description
---|---|---
DEBUG | >= 1 | Enables debugging and lists devices being used
DEBUG | >= 2 | Provides performance metrics for operations, including timing, memory usage, bandwidth for each kernel execution
DEBUG | >= 3 | Outputs the applied optimizations at a kernel level
DEBUG | >= 4 | Outputs the generated kernel code
DEBUG | >= 5 | Displays the intermediate representation of the computation UOps
DEBUG | >= 6 | Displays the intermediate representation of the computation UOps in a linearized manner, detailing the operation sequence
DEBUG | >= 7 | Outputs the assembly code generated for the target hardware

25
artifacts/package_sources/tinygrad/docs/favicon.svg Normal file
View File

@@ -0,0 +1,25 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="-10 -10 150 70" shape-rendering="crispEdges">
<g id="logo">
<!-- t -->
<polygon points="10,40 10,20 0,20 0,10 10,10 10,0 20,0 20,10 30,10 30,20 20,20 20,30 30,30 30,40" />
<!-- i -->
<polygon points="40,40 40,20 50,20 50,40" />
<polygon points="40,10 40,0 50,0 50,10" />
<!-- n -->
<polygon points="60,40 60,10 80,10 80,40 90,40 90,20 70,20 70,40" />
<!-- y -->
<polygon points="100,50 100,40 130,40 130,10 120,10 120,20 110,20 110,10 100,10 100,30 120,30 120,50" />
</g>
<style>
@media (prefers-color-scheme: dark) {
#logo {
fill: #fff;
}
}
@media (prefers-color-scheme: light) {
#logo {
fill: #000;
}
}
</style>
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 750 B

53
artifacts/package_sources/tinygrad/docs/index.md Normal file
View File

@@ -0,0 +1,53 @@
# tinygrad documentation
Welcome to the docs for tinygrad. This page is for users of the tinygrad library. tinygrad is not 1.0 yet, but it will be soon. The API has been pretty stable for a while.
While you can `pip install tinygrad`, we encourage you to install from source:
```bash
git clone https://github.com/tinygrad/tinygrad.git
cd tinygrad
python3 -m pip install -e .
```
After you have installed tinygrad, try the [MNIST tutorial](mnist.md).
If you are new to tensor libraries, learn how to use them by solving puzzles from [tinygrad-tensor-puzzles](https://github.com/obadakhalili/tinygrad-tensor-puzzles).
We also have [developer docs](developer/developer.md), and Di Zhu has created a [bunch of tutorials](https://mesozoic-egg.github.io/tinygrad-notes/) to help understand how tinygrad works.
## tinygrad Usage
The main class you will interact with is [Tensor](tensor/index.md). It functions very similarly to PyTorch, but has a bit more of a functional style. tinygrad supports [many datatypes](dtypes.md). All operations in tinygrad are lazy, meaning they won't do anything until you realize.
* tinygrad has a built in [neural network library](nn.md) with some classes, optimizers, and load/save state management.
* tinygrad has a JIT to make things fast. Decorate your pure function with `TinyJit`
* tinygrad has amazing support for multiple GPUs, allowing you to shard your Tensors with `Tensor.shard`
To understand what training looks like in tinygrad, you should read `beautiful_mnist.py`
We have a [quickstart guide](quickstart.md) and a [showcase](showcase.md)
## tinygrad Stack
<img src="./tinygrad_vs_others.png" alt="Tinygrad vs others" style="max-width: 1000px; height: auto;" />
## Differences from PyTorch
If you are migrating from PyTorch, welcome. Most of the API is the same. We hope you will find tinygrad both familiar and somehow more "correct feeling"
### tinygrad doesn't have nn.Module
There's nothing special about a "Module" class in tinygrad, it's just a normal class. [`nn.state.get_parameters`](nn.md/#tinygrad.nn.state.get_parameters) can be used to recursively search normal classes for valid tensors. Instead of the `forward` method in PyTorch, tinygrad just uses `__call__`
### tinygrad is functional
In tinygrad, you can do [`x.conv2d(w, b)`](tensor/ops.md/#tinygrad.Tensor.conv2d) or [`x.sparse_categorical_crossentropy(y)`](tensor/ops.md/#tinygrad.Tensor.sparse_categorical_crossentropy). We do also have a [`Conv2D`](nn.md/#tinygrad.nn.Conv2d) class like PyTorch if you want a place to keep the state, but all stateless operations don't have classes.
### tinygrad is lazy
When you do `a+b` in tinygrad, nothing happens. It's not until you [`realize`](tensor/properties.md#tinygrad.Tensor.realize) the Tensor that the computation actually runs.
### tinygrad requires @TinyJit to be fast
PyTorch spends a lot of development effort to make dispatch very fast. tinygrad doesn't. We have a simple decorator that will replay the kernels used in the decorated function.

11
artifacts/package_sources/tinygrad/docs/logo_tiny_dark.svg Normal file
View File

@@ -0,0 +1,11 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="-10 -10 150 70" shape-rendering="crispEdges" fill="#fff">
<!-- t -->
<polygon points="10,40 10,20 0,20 0,10 10,10 10,0 20,0 20,10 30,10 30,20 20,20 20,30 30,30 30,40" />
<!-- i -->
<polygon points="40,40 40,20 50,20 50,40" />
<polygon points="40,10 40,0 50,0 50,10" />
<!-- n -->
<polygon points="60,40 60,10 80,10 80,40 90,40 90,20 70,20 70,40" />
<!-- y -->
<polygon points="100,50 100,40 130,40 130,10 120,10 120,20 110,20 110,10 100,10 100,30 120,30 120,50" />
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 538 B

11
artifacts/package_sources/tinygrad/docs/logo_tiny_light.svg Normal file
View File

@@ -0,0 +1,11 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="-10 -10 150 70" shape-rendering="crispEdges">
<!-- t -->
<polygon points="10,40 10,20 0,20 0,10 10,10 10,0 20,0 20,10 30,10 30,20 20,20 20,30 30,30 30,40" />
<!-- i -->
<polygon points="40,40 40,20 50,20 50,40" />
<polygon points="40,10 40,0 50,0 50,10" />
<!-- n -->
<polygon points="60,40 60,10 80,10 80,40 90,40 90,20 70,20 70,40" />
<!-- y -->
<polygon points="100,50 100,40 130,40 130,10 120,10 120,20 110,20 110,10 100,10 100,30 120,30 120,50" />
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 526 B

184
artifacts/package_sources/tinygrad/docs/mnist.md Normal file
View File

@@ -0,0 +1,184 @@
# MNIST Tutorial
After you have installed tinygrad, this is a great first tutorial.
Start up a notebook locally, or use [colab](https://colab.research.google.com/). tinygrad is very lightweight, so it's easy to install anywhere and doesn't need a special colab image, but for speed we recommend a T4 GPU image.
### One-liner to install tinygrad in colab
```python
!pip install git+https://github.com/tinygrad/tinygrad.git
```
### What's the default device?
```python
from tinygrad import Device
print(Device.DEFAULT)
```
You will see `CUDA` here on a GPU instance, or `CPU` here on a CPU instance.
## A simple model
We'll use the model from [the Keras tutorial](https://keras.io/examples/vision/mnist_convnet/).
```python
from tinygrad import Tensor, nn, Context
class Model:
def __init__(self):
self.l1 = nn.Conv2d(1, 32, kernel_size=(3,3))
self.l2 = nn.Conv2d(32, 64, kernel_size=(3,3))
self.l3 = nn.Linear(1600, 10)
def __call__(self, x:Tensor) -> Tensor:
x = self.l1(x).relu().max_pool2d((2,2))
x = self.l2(x).relu().max_pool2d((2,2))
return self.l3(x.flatten(1).dropout(0.5))
```
Two key differences from PyTorch:
* Only the stateful layers are declared in `__init__`
* There's no `nn.Module` class or `forward` function, just a normal class and `__call__`
### Getting the dataset
```python
from tinygrad.nn.datasets import mnist
X_train, Y_train, X_test, Y_test = mnist()
print(X_train.shape, X_train.dtype, Y_train.shape, Y_train.dtype)
# (60000, 1, 28, 28) dtypes.uchar (60000,) dtypes.uchar
```
tinygrad includes MNIST, it only adds four lines. Feel free to read the [function](https://github.com/tinygrad/tinygrad/blob/master/tinygrad/nn/datasets.py).
## Using the model
MNIST is small enough that the `mnist()` function copies the dataset to the default device.
So creating the model and evaluating it is a matter of:
```python
model = Model()
acc = (model(X_test).argmax(axis=1) == Y_test).mean()
# NOTE: tinygrad is lazy, and hasn't actually run anything by this point
print(acc.item()) # ~10% accuracy, as expected from a random model
```
### Training the model
We'll use the Adam optimizer. The `nn.state.get_parameters` will walk the model class and pull out the parameters for the optimizer. Also, in tinygrad, it's typical to write a function to do the training step so it can be jitted.
```python
optim = nn.optim.Adam(nn.state.get_parameters(model))
batch_size = 128
@Context(TRAINING=1)
def step():
samples = Tensor.randint(batch_size, high=X_train.shape[0])
X, Y = X_train[samples], Y_train[samples]
optim.zero_grad()
loss = model(X).sparse_categorical_crossentropy(Y).backward()
optim.step()
return loss
```
You can time a step with:
```python
import timeit
timeit.repeat(step, repeat=5, number=1)
#[0.08268719699981375,
# 0.07478952900009972,
# 0.07714716600003158,
# 0.07785399599970333,
# 0.07605237000007037]
```
So around 75 ms on T4 colab.
If you want to see a breakdown of the time by kernel:
```python
from tinygrad import GlobalCounters, Context
GlobalCounters.reset()
with Context(DEBUG=2): step()
```
### Why so slow?
Unlike PyTorch, tinygrad isn't designed to be fast like that. While 75 ms for one step is plenty fast for debugging, it's not great for training. Here, we introduce the first quintessentially tinygrad concept, the `TinyJit`.
```python
from tinygrad import TinyJit
jit_step = TinyJit(step)
```
NOTE: It can also be used as a decorator `@TinyJit`
Now when we time it:
```python
import timeit
timeit.repeat(jit_step, repeat=5, number=1)
# [0.2596786549997887,
# 0.08989566299987928,
# 0.0012115650001760514,
# 0.001010227999813651,
# 0.0012164899999334011]
```
1.0 ms is 75x faster! Note that we aren't syncing the GPU, so GPU time may be slower.
The first two runs of the function execute normally, with the JIT capturing the kernels. Starting from the third run, only the tinygrad operations are replayed, removing the overhead by skipping Python code execution. So be aware that any non-tinygrad Python values affecting the kernels will be "frozen" from the second run. Note that `Tensor` randomness functions work as expected.
Unlike other JITs, we JIT everything, including the optimizer. Think of it as a dumb replay on different data.
## Putting it together
Since we are just randomly sampling from the dataset, there's no real concept of an epoch. We have a batch size of 128, so the Keras example is taking about 7000 steps.
```python
for step in range(7000):
loss = jit_step()
if step%100 == 0:
acc = (model(X_test).argmax(axis=1) == Y_test).mean().item()
print(f"step {step:4d}, loss {loss.item():.2f}, acc {acc*100.:.2f}%")
```
It doesn't take long to reach 98%, and it usually reaches 99%.
```
step 0, loss 4.03, acc 71.43%
step 100, loss 0.34, acc 93.86%
step 200, loss 0.23, acc 95.97%
step 300, loss 0.18, acc 96.32%
step 400, loss 0.18, acc 96.76%
step 500, loss 0.13, acc 97.46%
step 600, loss 0.14, acc 97.45%
step 700, loss 0.10, acc 97.27%
step 800, loss 0.23, acc 97.49%
step 900, loss 0.13, acc 97.51%
step 1000, loss 0.13, acc 97.88%
step 1100, loss 0.11, acc 97.72%
step 1200, loss 0.14, acc 97.65%
step 1300, loss 0.12, acc 98.04%
step 1400, loss 0.25, acc 98.17%
step 1500, loss 0.11, acc 97.86%
step 1600, loss 0.21, acc 98.21%
step 1700, loss 0.14, acc 98.34%
...
```
## From here?
tinygrad is yours to play with now. It's pure Python and short, so unlike PyTorch, fixing library bugs is well within your abilities.
- It's two lines to add multiGPU support to this example (can you find them?). You have to `.shard` the model to all GPUs, and `.shard` the dataset by batch.
- `with Context(DEBUG=2)` shows the running kernels, `DEBUG=4` shows the code. All `Context` variables can also be environment variables.
- `with Context(BEAM=2)` will do a BEAM search on the kernels, searching many possible implementations for what runs the fastest on your hardware. After this search, tinygrad is usually speed competitive with PyTorch, and the results are cached so you won't have to search next time.
[Join our Discord](https://discord.gg/ZjZadyC7PK) for help, and if you want to be a tinygrad developer. Please read the Discord rules when you get there.
[Follow us on Twitter](https://twitter.com/__tinygrad__) to keep up with the project.

40
artifacts/package_sources/tinygrad/docs/nn.md Normal file
View File

@@ -0,0 +1,40 @@
## Neural Network classes
::: tinygrad.nn.BatchNorm
::: tinygrad.nn.Conv1d
::: tinygrad.nn.Conv2d
::: tinygrad.nn.ConvTranspose1d
::: tinygrad.nn.ConvTranspose2d
::: tinygrad.nn.Linear
::: tinygrad.nn.GroupNorm
::: tinygrad.nn.InstanceNorm
::: tinygrad.nn.LayerNorm
::: tinygrad.nn.LayerNorm2d
::: tinygrad.nn.RMSNorm
::: tinygrad.nn.Embedding
::: tinygrad.nn.LSTMCell
## Optimizers
::: tinygrad.nn.optim.SGD
::: tinygrad.nn.optim.LARS
::: tinygrad.nn.optim.AdamW
::: tinygrad.nn.optim.Adam
::: tinygrad.nn.optim.LAMB
## Load/Save
::: tinygrad.nn.state.safe_load
::: tinygrad.nn.state.safe_save
::: tinygrad.nn.state.get_state_dict
::: tinygrad.nn.state.get_parameters
::: tinygrad.nn.state.load_state_dict
::: tinygrad.nn.state.tar_extract
options:
show_signature: false
separate_signature: false
::: tinygrad.nn.state.torch_load
options:
show_signature: false
separate_signature: false
::: tinygrad.llm.gguf.gguf_load

306
artifacts/package_sources/tinygrad/docs/quickstart.md Normal file
View File

@@ -0,0 +1,306 @@
# Quick Start Guide
This guide assumes no prior knowledge of pytorch or any other deep learning framework, but does assume some basic knowledge of neural networks.
It is intended to be a very quick overview of the high level API that tinygrad provides.
This guide is also structured as a tutorial which at the end of it you will have a working model that can classify handwritten digits.
We need some imports to get started:
```python
import numpy as np
from tinygrad.helpers import Timing
```
## Tensors
Tensors are the base data structure in tinygrad. They can be thought of as a multidimensional array of a specific data type.
All high level operations in tinygrad operate on these tensors.
The tensor class can be imported like so:
```python
from tinygrad import Tensor
```
Tensors can be created from an existing data structure like a python list or numpy ndarray:
```python
t1 = Tensor([1, 2, 3, 4, 5])
na = np.array([1, 2, 3, 4, 5])
t2 = Tensor(na)
```
Tensors can also be created using one of the many factory methods:
```python
full = Tensor.full(shape=(2, 3), fill_value=5) # create a tensor of shape (2, 3) filled with 5
zeros = Tensor.zeros(2, 3) # create a tensor of shape (2, 3) filled with 0
ones = Tensor.ones(2, 3) # create a tensor of shape (2, 3) filled with 1
full_like = Tensor.full_like(full, fill_value=2) # create a tensor of the same shape as `full` filled with 2
zeros_like = Tensor.zeros_like(full) # create a tensor of the same shape as `full` filled with 0
ones_like = Tensor.ones_like(full) # create a tensor of the same shape as `full` filled with 1
eye = Tensor.eye(3) # create a 3x3 identity matrix
arange = Tensor.arange(start=0, stop=10, step=1) # create a tensor of shape (10,) filled with values from 0 to 9
rand = Tensor.rand(2, 3) # create a tensor of shape (2, 3) filled with random values from a uniform distribution
randn = Tensor.randn(2, 3) # create a tensor of shape (2, 3) filled with random values from a standard normal distribution
uniform = Tensor.uniform(2, 3, low=0, high=10) # create a tensor of shape (2, 3) filled with random values from a uniform distribution between 0 and 10
```
There are even more of these factory methods, you can find them in the [Tensor Creation](tensor/creation.md) file.
All the tensors creation methods can take a `dtype` argument to specify the data type of the tensor, find the supported `dtype` in [dtypes](dtypes.md).
```python
from tinygrad import dtypes
t3 = Tensor([1, 2, 3, 4, 5], dtype=dtypes.int32)
```
Tensors allow you to perform operations on them like so:
```python
t4 = Tensor([1, 2, 3, 4, 5])
t5 = (t4 + 1) * 2
t6 = (t5 * t4).relu().log_softmax()
```
All of these operations are lazy and are only executed when you realize the tensor using `.realize()` or `.numpy()`.
```python
print(t6.numpy())
# [-56. -48. -36. -20. 0.]
```
There are a lot more operations that can be performed on tensors, you can find them in the [Tensor Ops](tensor/ops.md) file.
Additionally reading through [abstractions2.py](https://github.com/tinygrad/tinygrad/blob/master/docs/abstractions2.py) will help you understand how operations on these tensors make their way down to your hardware.
## Models
Neural networks in tinygrad are really just represented by the operations performed on tensors.
These operations are commonly grouped into the `__call__` method of a class which allows modularization and reuse of these groups of operations.
These classes do not need to inherit from any base class, in fact if they don't need any trainable parameters they don't even need to be a class!
An example of this would be the `nn.Linear` class which represents a linear layer in a neural network.
```python
class Linear:
def __init__(self, in_features, out_features, bias=True, initialization: str='kaiming_uniform'):
self.weight = getattr(Tensor, initialization)(out_features, in_features)
self.bias = Tensor.zeros(out_features) if bias else None
def __call__(self, x):
return x.linear(self.weight.transpose(), self.bias)
```
There are more neural network modules already implemented in [nn](nn.md), and you can also implement your own.
We will be implementing a simple neural network that can classify handwritten digits from the MNIST dataset.
Our classifier will be a simple 2 layer neural network with a Leaky ReLU activation function.
It will use a hidden layer size of 128 and an output layer size of 10 (one for each digit) with no bias on either Linear layer.
```python
class TinyNet:
def __init__(self):
self.l1 = Linear(784, 128, bias=False)
self.l2 = Linear(128, 10, bias=False)
def __call__(self, x):
x = self.l1(x)
x = x.leaky_relu()
x = self.l2(x)
return x
net = TinyNet()
```
We can see that the forward pass of our neural network is just the sequence of operations performed on the input tensor `x`.
We can also see that functional operations like `leaky_relu` are not defined as classes and instead are just methods we can just call.
Finally, we just initialize an instance of our neural network, and we are ready to start training it.
## Training
Now that we have our neural network defined we can start training it.
Training neural networks in tinygrad is super simple.
All we need to do is define our neural network, define our loss function, and then call `.backward()` on the loss function to compute the gradients.
They can then be used to update the parameters of our neural network using one of the many [Optimizers](nn.md#optimizers).
For our loss function we will be using sparse categorical cross entropy loss. The implementation below is taken from [tensor.py](https://github.com/tinygrad/tinygrad/blob/master/tinygrad/tensor.py), it's copied below to highlight an important detail of tinygrad.
```python
def sparse_categorical_crossentropy(self, Y, ignore_index=-1) -> Tensor:
loss_mask = Y != ignore_index
y_counter = Tensor.arange(self.shape[-1], dtype=dtypes.int32).unsqueeze(0).expand(Y.numel(), self.shape[-1])
y = ((y_counter == Y.flatten().reshape(-1, 1)).where(-1.0, 0) * loss_mask.reshape(-1, 1)).reshape(*Y.shape, self.shape[-1])
return self.log_softmax().mul(y).sum() / loss_mask.sum()
```
As we can see in this implementation of cross entropy loss, there are certain operations that tinygrad does not support natively.
Load/store ops are not supported in tinygrad natively because they add complexity when trying to port to different backends, 90% of the models out there don't use/need them, and they can be implemented like it's done above with an `arange` mask.
For our optimizer we will be using the traditional stochastic gradient descent optimizer with a learning rate of 3e-4.
```python
from tinygrad.nn.optim import SGD
opt = SGD([net.l1.weight, net.l2.weight], lr=3e-4)
```
We can see that we are passing in the parameters of our neural network to the optimizer.
This is due to the fact that the optimizer needs to know which parameters to update.
There is a simpler way to do this just by using `get_parameters(net)` from `tinygrad.nn.state` which will return a list of all the parameters in the neural network.
The parameters are just listed out explicitly here for clarity.
Now that we have our network, loss function, and optimizer defined all we are missing is the data to train on!
There are a couple of dataset loaders in tinygrad located in [/extra/datasets](https://github.com/tinygrad/tinygrad/blob/master/extra/datasets).
We will be using the MNIST dataset loader.
```python
from extra.datasets import fetch_mnist
```
Now we have everything we need to start training our neural network.
We will be training for 1000 steps with a batch size of 64.
We use `with Context(TRAINING=1)` to enable training mode.
Upon exit, the flag is restored to its previous value by the context manager.
```python
from tinygrad import Context
X_train, Y_train, X_test, Y_test = fetch_mnist()
with Context(TRAINING=1):
for step in range(1000):
# random sample a batch
samp = np.random.randint(0, X_train.shape[0], size=(64))
batch = Tensor(X_train[samp])
# get the corresponding labels
labels = Tensor(Y_train[samp])
# forward pass
out = net(batch)
# compute loss
loss = sparse_categorical_crossentropy(out, labels)
# zero gradients
opt.zero_grad()
# backward pass
loss.backward()
# update parameters
opt.step()
# calculate accuracy
pred = out.argmax(axis=-1)
acc = (pred == labels).mean()
if step % 100 == 0:
print(f"Step {step+1} | Loss: {loss.numpy()} | Accuracy: {acc.numpy()}")
```
## Evaluation
Now that we have trained our neural network we can evaluate it on the test set.
We will be using the same batch size of 64 and will be evaluating for 1000 of those batches.
```python
with Timing("Time: "):
avg_acc = 0
for step in range(1000):
# random sample a batch
samp = np.random.randint(0, X_test.shape[0], size=(64))
batch = Tensor(X_test[samp])
# get the corresponding labels
labels = Y_test[samp]
# forward pass
out = net(batch)
# calculate accuracy
pred = out.argmax(axis=-1).numpy()
avg_acc += (pred == labels).mean()
print(f"Test Accuracy: {avg_acc / 1000}")
```
## And that's it
Highly recommend you check out the [examples/](https://github.com/tinygrad/tinygrad/blob/master/examples) folder for more examples of using tinygrad.
Reading the source code of tinygrad is also a great way to learn how it works.
Specifically the tests in [test/](https://github.com/tinygrad/tinygrad/blob/master/test) are a great place to see how to use and the semantics of the different operations.
There are also a bunch of models implemented in [models/](https://github.com/tinygrad/tinygrad/blob/master/extra/models) that you can use as a reference.
Additionally, feel free to ask questions in the `#learn-tinygrad` channel on the [discord](https://discord.gg/beYbxwxVdx). Don't ask to ask, just ask!
## Extras
### JIT
Additionally, it is possible to speed up the computation of certain neural networks by using the JIT.
Currently, this does not support models with varying input sizes and non tinygrad operations.
To use the JIT we just need to add a function decorator to the forward pass of our neural network and ensure that the input and output are realized tensors.
Or in this case we will create a wrapper function and decorate the wrapper function to speed up the evaluation of our neural network.
```python
from tinygrad import TinyJit
@TinyJit
def jit(x):
return net(x).realize()
with Timing("Time: "):
avg_acc = 0
for step in range(1000):
# random sample a batch
samp = np.random.randint(0, X_test.shape[0], size=(64))
batch = Tensor(X_test[samp])
# get the corresponding labels
labels = Y_test[samp]
# forward pass with jit
out = jit(batch)
# calculate accuracy
pred = out.argmax(axis=-1).numpy()
avg_acc += (pred == labels).mean()
print(f"Test Accuracy: {avg_acc / 1000}")
```
You will find that the evaluation time is much faster than before and that your accelerator utilization is much higher.
### Saving and Loading Models
The standard weight format for tinygrad is [safetensors](https://github.com/huggingface/safetensors). This means that you can load the weights of any model also using safetensors into tinygrad.
There are functions in [state.py](https://github.com/tinygrad/tinygrad/blob/master/tinygrad/nn/state.py) to save and load models to and from this format.
```python
from tinygrad.nn.state import safe_save, safe_load, get_state_dict, load_state_dict
# first we need the state dict of our model
state_dict = get_state_dict(net)
# then we can just save it to a file
safe_save(state_dict, "model.safetensors")
# and load it back in
state_dict = safe_load("model.safetensors")
load_state_dict(net, state_dict)
```
Many of the models in the [models/](https://github.com/tinygrad/tinygrad/tree/master/extra/models) folder have a `load_from_pretrained` method that will download and load the weights for you. These usually are pytorch weights meaning that you would need pytorch installed to load them.
### Environment Variables
There exist a bunch of environment variables that control the runtime behavior of tinygrad.
Some of the commons ones are `DEBUG` and the different backend enablement variables.
You can find a full list and their descriptions in [env_vars.md](env_vars.md).
### Visualizing the Computation Graph
It is possible to visualize the computation graph of a neural network using VIZ=1.

87
artifacts/package_sources/tinygrad/docs/runtime.md Normal file
View File

@@ -0,0 +1,87 @@
# Runtimes
tinygrad supports various runtimes, enabling your code to scale across a wide range of devices. The default runtime can be automatically selected based on the available hardware, or you can force a specific runtime to be default using environment variables (e.g., `DEV=CPU`).
| Runtime | Description | Compiler Options | Requirements |
|---------|-------------|------------------|--------------|
| [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) | Provides acceleration for NVIDIA GPUs | nvrtc (default)<br>PTX (`DEV=NV:PTX`) | Ampere/Ada/Blackwell series GPUs.<br>You can select an interface via [the `DEV` variable](env_vars.md#dev-variable). See [NV interfaces](#nv-interfaces) for details. |
| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | LLVM (`DEV=AMD:LLVM`)<br>HIP/COMGR (`DEV=AMD:HIP`) | CDNA3, CDNA4, RDNA3 or RDNA4 GPUs.<br>You can select an interface via [the `DEV` variable](env_vars.md#dev-variable). See [AMD interfaces](#amd-interfaces) for details. |
| [QCOM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_qcom.py) | Provides acceleration for QCOM GPUs | - | 6xx series GPUs |
| [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | - | M1+ Macs; Metal 3.0+ for `bfloat` support |
| [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | nvrtc (default)<br> PTX (`DEV=CUDA:PTX`) | NVIDIA GPU with CUDA support |
| [CL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cl.py) | Accelerates computations using OpenCL on GPUs | - | OpenCL 2.0 compatible device |
| [CPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang or llvm compiler | Clang JIT (default)<br>LLVM IR (`DEV=CPU:LLVM`) | `clang` compiler in system `PATH`<br>You can specify additional arch parameters via [the `DEV` variable](env_vars.md#dev-variable). See [CPU arch](#cpu-arch) for details. |
| [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | - | Dawn library installed and discoverable. Binaries: [pydawn v0.3.0](https://github.com/wpmed92/pydawn/releases/tag/v0.3.0) |
## Interoperability
tinygrad provides interoperability with OpenCL and PyTorch, allowing efficient tensor data sharing between frameworks through the `Tensor.from_blob` API. This enables zero-copy operations by working directly with external memory pointers.
**Important**: When using external memory pointers with tinygrad tensors, you must ensure these pointers remain valid throughout the entire lifetime of the tinygrad tensor to prevent memory corruption.
### `CUDA`/`METAL` PyTorch Interoperability
You can seamlessly work with CUDA/MPS tensors between PyTorch and tinygrad without data copying:
```python
from tinygrad.dtype import _from_torch_dtype
tensor1 = torch.tensor([1.0, 2.0, 3.0], device=torch.device("cuda"))
tiny_tensor1 = Tensor.from_blob(tensor1.data_ptr(), tensor1.shape, dtype=_from_torch_dtype(tensor1.dtype), device='CUDA')
# Before tinygrad calculations, mps needs to be synchronized to make sure data is valid.
if data.device.type == "mps": torch.mps.synchronize()
else: torch.cuda.synchronize()
x = (tiny_tensor1 + 1).realize()
```
### `QCOM` OpenCL Interoperability
tinygrad supports OpenCL interoperability on `QCOM` backend.
Buffer interop allows direct access to OpenCL memory buffers:
```python
# create raw opencl buffer.
cl_buf = cl.clCreateBuffer(cl_context, cl.CL_MEM_READ_WRITE, 0x100, None, status := ctypes.c_int32())
# extract pointers
cl_buf_desc_ptr = to_mv(ctypes.addressof(cl_buf), 8).cast('Q')[0]
rawbuf_ptr = to_mv(cl_buf_desc_ptr, 0x100).cast('Q')[20] # offset 0xA0 is a raw gpu pointer.
# create tiny tensor
tiny = Tensor.from_blob(rawbuf_ptr, (8, 8), dtype=dtypes.int, device='QCOM')
```
And the same for the images:
```python
# create cl image.
cl_img = cl.clCreateImage2D(cl_context, cl.CL_MEM_READ_WRITE, cl.cl_image_format(cl.CL_RGBA, cl.CL_FLOAT), w, h, 0, None, status := ctypes.c_int32())
# extract pointers
cl_buf_desc_ptr = to_mv(ctypes.addressof(cl_img), 8).cast('Q')[0]
rawbuf_ptr = to_mv(cl_buf_desc_ptr, 0x100).cast('Q')[20] # offset 0xA0 is a raw gpu pointer.
# create tiny tensor
tiny = Tensor.from_blob(rawbuf_ptr, (h*w*4,), dtype=dtypes.imagef((h,w)), device='QCOM')
```
## AMD Interfaces
AMD backend supports several interfaces for communicating with devices:
* `KFD`: uses the amdgpu driver
* `PCI`: uses the [AM driver](developer/am.md)
* `USB`: USB3 interface for asm24xx chips.
You can force an interface by setting the interface component of [the `DEV` environment variable](env_vars.md#dev-variable) to one of these values. When set to `PCI`, this may unbind your GPU from the amdgpu driver.
## NV Interfaces
NV backend supports several interfaces for communicating with devices:
* `NVK`: uses the nvidia driver
* `PCI`: uses the [NV driver](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/support/nv/nvdev.py)
## CPU Arch
The CPU renderers may be additionally configured using the arch component of [the `DEV` environment variable](env_vars.md#dev-variable).
CPU arch should be specified as a comma-separated list of parameters, and must contain at least two values: the architecture family (ie. x86_64, arm64, or riscv64) and the cpu type (as accepted by `clang`'s `-march`).
If native is specified as the cpu type, tinygrad (or delegate compiler) will query the host cpu type. Additional comma-separated values are interpreted as cpu feature flags. When a value is preceded by a `-` character, the corresponding feature flag will be disabled, otherwise the flag will be enabled.
Note that enabled feature flags should not be preceded by a `+`.

62
artifacts/package_sources/tinygrad/docs/showcase.md Normal file
View File

@@ -0,0 +1,62 @@
# Showcase
Despite being a tiny library, tinygrad is capable of doing a lot of things. From state-of-the-art [vision](https://arxiv.org/abs/1905.11946) to state-of-the-art [language](https://arxiv.org/abs/1706.03762) models.
## Vision
### EfficientNet
You can either pass in the URL of a picture to discover what it is:
```sh
python3 examples/efficientnet.py ./test/models/efficientnet/Chicken.jpg
```
Or, if you have a camera and OpenCV installed, you can detect what is in front of you:
```sh
python3 examples/efficientnet.py webcam
```
### YOLOv8
Take a look at [yolov8.py](https://github.com/tinygrad/tinygrad/tree/master/examples/yolov8.py).
![yolov8 by tinygrad](https://github.com/tinygrad/tinygrad/blob/master/docs/showcase/yolov8_showcase_image.png?raw=true)
## Audio
### Whisper
Take a look at [whisper.py](https://github.com/tinygrad/tinygrad/tree/master/examples/whisper.py). You need pyaudio and torchaudio installed.
```sh
SMALL=1 python3 examples/whisper.py
```
## Generative
### Stable Diffusion
```sh
python3 examples/stable_diffusion.py
```
![a horse sized cat eating a bagel](https://github.com/tinygrad/tinygrad/blob/master/docs/showcase/stable_diffusion_by_tinygrad.jpg?raw=true)
*"a horse sized cat eating a bagel"*
### LLaMA
You will need to download and put the weights into the `weights/LLaMA` directory, which may need to be created.
Then you can have a chat with Stacy:
```sh
python3 examples/llama.py
```
### Conversation
Make sure you have espeak installed and `PHONEMIZER_ESPEAK_LIBRARY` set.
Then you can talk to Stacy:
```sh
python3 examples/conversation.py
```

BIN
artifacts/package_sources/tinygrad/docs/showcase/mnist_by_tinygrad.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 15 KiB

BIN
artifacts/package_sources/tinygrad/docs/showcase/stable_diffusion_by_tinygrad.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 87 KiB

BIN
artifacts/package_sources/tinygrad/docs/showcase/yolo_by_tinygrad.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 131 KiB

BIN
artifacts/package_sources/tinygrad/docs/showcase/yolov8_showcase_image.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.1 MiB

33
artifacts/package_sources/tinygrad/docs/tensor/creation.md Normal file
View File

@@ -0,0 +1,33 @@
## Creation (basic)
::: tinygrad.Tensor.empty
::: tinygrad.Tensor.zeros
::: tinygrad.Tensor.ones
::: tinygrad.Tensor.full
::: tinygrad.Tensor.arange
::: tinygrad.Tensor.linspace
::: tinygrad.Tensor.eye
::: tinygrad.Tensor.full_like
::: tinygrad.Tensor.zeros_like
::: tinygrad.Tensor.ones_like
## Creation (external)
::: tinygrad.Tensor.from_blob
::: tinygrad.Tensor.from_url
## Creation (random)
::: tinygrad.Tensor.manual_seed
::: tinygrad.Tensor.rand
::: tinygrad.Tensor.rand_like
::: tinygrad.Tensor.randn
::: tinygrad.Tensor.randn_like
::: tinygrad.Tensor.randint
::: tinygrad.Tensor.randperm
::: tinygrad.Tensor.normal
::: tinygrad.Tensor.uniform
::: tinygrad.Tensor.scaled_uniform
::: tinygrad.Tensor.glorot_uniform
::: tinygrad.Tensor.kaiming_uniform
::: tinygrad.Tensor.kaiming_normal

95
artifacts/package_sources/tinygrad/docs/tensor/elementwise.md Normal file
View File

@@ -0,0 +1,95 @@
Elementwise ops operate on a per element basis. They don't change the shape of the tensor.
## Unary Ops (math)
::: tinygrad.Tensor.logical_not
::: tinygrad.Tensor.neg
::: tinygrad.Tensor.log
::: tinygrad.Tensor.log2
::: tinygrad.Tensor.log10
::: tinygrad.Tensor.exp
::: tinygrad.Tensor.exp2
::: tinygrad.Tensor.sqrt
::: tinygrad.Tensor.rsqrt
::: tinygrad.Tensor.sin
::: tinygrad.Tensor.cos
::: tinygrad.Tensor.tan
::: tinygrad.Tensor.asin
::: tinygrad.Tensor.acos
::: tinygrad.Tensor.atan
::: tinygrad.Tensor.trunc
::: tinygrad.Tensor.ceil
::: tinygrad.Tensor.floor
::: tinygrad.Tensor.round
::: tinygrad.Tensor.isinf
::: tinygrad.Tensor.isnan
::: tinygrad.Tensor.isfinite
::: tinygrad.Tensor.lerp
::: tinygrad.Tensor.square
::: tinygrad.Tensor.clamp
::: tinygrad.Tensor.clip
::: tinygrad.Tensor.sign
::: tinygrad.Tensor.abs
::: tinygrad.Tensor.reciprocal
## Unary Ops (activation)
::: tinygrad.Tensor.relu
::: tinygrad.Tensor.sigmoid
::: tinygrad.Tensor.logsigmoid
::: tinygrad.Tensor.hardsigmoid
::: tinygrad.Tensor.elu
::: tinygrad.Tensor.celu
::: tinygrad.Tensor.selu
::: tinygrad.Tensor.swish
::: tinygrad.Tensor.silu
::: tinygrad.Tensor.relu6
::: tinygrad.Tensor.hardswish
::: tinygrad.Tensor.tanh
::: tinygrad.Tensor.sinh
::: tinygrad.Tensor.cosh
::: tinygrad.Tensor.atanh
::: tinygrad.Tensor.asinh
::: tinygrad.Tensor.acosh
::: tinygrad.Tensor.hardtanh
::: tinygrad.Tensor.erf
::: tinygrad.Tensor.gelu
::: tinygrad.Tensor.quick_gelu
::: tinygrad.Tensor.leaky_relu
::: tinygrad.Tensor.mish
::: tinygrad.Tensor.softplus
::: tinygrad.Tensor.softsign
## Elementwise Ops (broadcasted)
::: tinygrad.Tensor.add
::: tinygrad.Tensor.sub
::: tinygrad.Tensor.mul
::: tinygrad.Tensor.div
::: tinygrad.Tensor.mod
::: tinygrad.Tensor.fmod
::: tinygrad.Tensor.bitwise_xor
::: tinygrad.Tensor.bitwise_and
::: tinygrad.Tensor.bitwise_or
::: tinygrad.Tensor.bitwise_not
::: tinygrad.Tensor.lshift
::: tinygrad.Tensor.rshift
::: tinygrad.Tensor.pow
::: tinygrad.Tensor.maximum
::: tinygrad.Tensor.minimum
::: tinygrad.Tensor.where
::: tinygrad.Tensor.copysign
::: tinygrad.Tensor.logaddexp
## Casting Ops
::: tinygrad.Tensor.cast
::: tinygrad.Tensor.bitcast
::: tinygrad.Tensor.float
::: tinygrad.Tensor.half
::: tinygrad.Tensor.int
::: tinygrad.Tensor.bool
::: tinygrad.Tensor.bfloat16
::: tinygrad.Tensor.double
::: tinygrad.Tensor.long
::: tinygrad.Tensor.short

7
artifacts/package_sources/tinygrad/docs/tensor/index.md Normal file
View File

@@ -0,0 +1,7 @@
# Tensor
::: tinygrad.Tensor
options:
heading_level: 2
members: false
show_source: false

32
artifacts/package_sources/tinygrad/docs/tensor/movement.md Normal file
View File

@@ -0,0 +1,32 @@
## Movement (low level)
::: tinygrad.Tensor.view
::: tinygrad.Tensor.reshape
::: tinygrad.Tensor.expand
::: tinygrad.Tensor.permute
::: tinygrad.Tensor.flip
::: tinygrad.Tensor.shrink
::: tinygrad.Tensor.pad
## Movement (high level)
::: tinygrad.Tensor.__getitem__
::: tinygrad.Tensor.gather
::: tinygrad.Tensor.cat
::: tinygrad.Tensor.stack
::: tinygrad.Tensor.repeat
::: tinygrad.Tensor.repeat_interleave
::: tinygrad.Tensor.split
::: tinygrad.Tensor.chunk
::: tinygrad.Tensor.unfold
::: tinygrad.Tensor.meshgrid
::: tinygrad.Tensor.squeeze
::: tinygrad.Tensor.unsqueeze
::: tinygrad.Tensor.T
::: tinygrad.Tensor.transpose
::: tinygrad.Tensor.flatten
::: tinygrad.Tensor.unflatten
::: tinygrad.Tensor.diag
::: tinygrad.Tensor.diagonal
::: tinygrad.Tensor.roll
::: tinygrad.Tensor.rearrange

68
artifacts/package_sources/tinygrad/docs/tensor/ops.md Normal file
View File

@@ -0,0 +1,68 @@
## Reduce
::: tinygrad.Tensor.sum
::: tinygrad.Tensor.prod
::: tinygrad.Tensor.max
::: tinygrad.Tensor.min
::: tinygrad.Tensor.any
::: tinygrad.Tensor.all
::: tinygrad.Tensor.isclose
::: tinygrad.Tensor.allclose
::: tinygrad.Tensor.mean
::: tinygrad.Tensor.var
::: tinygrad.Tensor.var_mean
::: tinygrad.Tensor.std
::: tinygrad.Tensor.std_mean
::: tinygrad.Tensor.softmax
::: tinygrad.Tensor.log_softmax
::: tinygrad.Tensor.logsumexp
::: tinygrad.Tensor.logcumsumexp
::: tinygrad.Tensor.argmax
::: tinygrad.Tensor.argmin
## Processing
::: tinygrad.Tensor.avg_pool2d
::: tinygrad.Tensor.max_pool2d
::: tinygrad.Tensor.max_unpool2d
::: tinygrad.Tensor.conv2d
::: tinygrad.Tensor.conv_transpose2d
::: tinygrad.Tensor.dot
::: tinygrad.Tensor.matmul
::: tinygrad.Tensor.einsum
::: tinygrad.Tensor.cumsum
::: tinygrad.Tensor.cumprod
::: tinygrad.Tensor.cummax
::: tinygrad.Tensor.cummin
::: tinygrad.Tensor.triu
::: tinygrad.Tensor.tril
::: tinygrad.Tensor.interpolate
::: tinygrad.Tensor.scatter
::: tinygrad.Tensor.scatter_reduce
::: tinygrad.Tensor.masked_select
::: tinygrad.Tensor.masked_fill
::: tinygrad.Tensor.nonzero
::: tinygrad.Tensor.sort
::: tinygrad.Tensor.argsort
::: tinygrad.Tensor.topk
::: tinygrad.Tensor.multinomial
## Neural Network (functional)
::: tinygrad.Tensor.linear
::: tinygrad.Tensor.sequential
::: tinygrad.Tensor.layernorm
::: tinygrad.Tensor.batchnorm
::: tinygrad.Tensor.dropout
::: tinygrad.Tensor.one_hot
::: tinygrad.Tensor.scaled_dot_product_attention
::: tinygrad.Tensor.binary_crossentropy
::: tinygrad.Tensor.binary_crossentropy_logits
::: tinygrad.Tensor.sparse_categorical_crossentropy
::: tinygrad.Tensor.cross_entropy
::: tinygrad.Tensor.nll_loss
## Linear Algebra
::: tinygrad.Tensor.qr
::: tinygrad.Tensor.svd

39
artifacts/package_sources/tinygrad/docs/tensor/properties.md Normal file
View File

@@ -0,0 +1,39 @@
## Basic
::: tinygrad.Tensor.shape
::: tinygrad.Tensor.dtype
::: tinygrad.Tensor.device
::: tinygrad.Tensor.ndim
::: tinygrad.Tensor.numel
::: tinygrad.Tensor.element_size
::: tinygrad.Tensor.nbytes
::: tinygrad.Tensor.is_floating_point
::: tinygrad.Tensor.size
## Data Access
::: tinygrad.Tensor.data
::: tinygrad.Tensor.item
::: tinygrad.Tensor.tolist
::: tinygrad.Tensor.numpy
## tinygrad ops
::: tinygrad.Tensor.linear_with_vars
::: tinygrad.Tensor.schedule_linear
::: tinygrad.Tensor.realize
::: tinygrad.Tensor.replace
::: tinygrad.Tensor.assign
::: tinygrad.Tensor.detach
::: tinygrad.Tensor.clone
::: tinygrad.Tensor.to
::: tinygrad.Tensor.to_
::: tinygrad.Tensor.shard
::: tinygrad.Tensor.shard_
::: tinygrad.Tensor.contiguous
::: tinygrad.Tensor.contiguous_backward
## Gradient
::: tinygrad.Tensor.gradient
::: tinygrad.Tensor.backward

54
artifacts/package_sources/tinygrad/docs/tinybox.md Normal file
View File

@@ -0,0 +1,54 @@
# tinybox
Although these docs live in tinygrad, they pertain to deep learning hardware sold by the tiny corp. tinyboxes are used heavily in tinygrad's CI, and are the best tested platform to use tinygrad with. They appeared running tinygrad on [MLPerf Training 4.0](https://public.tableau.com/views/MLCommons-Training_16993769118290/MLCommons-Training)
If you don't have a tinybox and you want one, see [tinygrad.org](https://tinygrad.org). If you don't want one, that's okay too.
## Welcome
Welcome to your tinybox! The tinybox is the universal system purpose-built for all AI infrastructure and workloads, from training to inference. The red box includes six 7900XTX GPUs, the green box includes six 4090 GPUs, and the green v2 box includes four 5090 GPUs. Whether you bought a red one or a green one, we want you to love it.
We don't have a stupid cloud service, you don't have to create a tiny account to set it up, and we aren't tracking how you use the box. We're just happy you bought one. This petaflop is your petaflop.
## Plugging it in
tinybox has two 1600W PSUs, which together exceed the capacity of most 120V household circuits. Fortunately, it comes with two plugs. You'll want to plug each plug into a different circuit. You can verify that they are different circuits by flipping the breaker and seeing what turns off. If you have at least a 120V 30A or 220V 20A circuit, you are welcome to use only that one.
You'll also want to connect the Ethernet port without a rubber stopper to your home network.
While it's designed primarily for the home or office, the tinybox is 12U rack mountable using [these rails](https://rackmountmart.store.turbify.net/26slidrailfo.html).
## Power limiting the box
While a tinybox should ideally be run without power limits, there are cases where you might want to run the box off of a single outlet.
In such cases, it is possible to power limit the box using the provided `power-limit` script, which will power limit all of the GPUs to a specified wattage.
`sudo power-limit 150` should be good to run off of a single 120V 15A outlet.
## Connecting to the box
tinybox ships with a relatively basic install of Ubuntu 22.04. To do initial setup, you can either plug in a VGA monitor and keyboard, or you can connect remotely to the machine using the BMC. The BMC IP and password are displayed on the screen.
`ipmitool -H <BMC IP> -U admin -P <BMC PW> -I lanplus sol activate`
The default username is `tiny` and the default password is `tiny`. Once you are logged in, you can add an SSH key to authorized keys to connect over SSH (on the normal IP). Exit `ipmitool` with `~.` after a newline.
The BMC also has a web interface you can use if you find that easier.
## Changing the BMC password
It is recommended that you change the BMC password after setting up the box, as the password on the screen is only the initial password.
If you do decide to change the BMC password and no longer want the initial password to be displayed, remove the `/root/.bmc_password` file.
Reboot after making these changes or restart the `tinybox-display.service` service.
## What do I use it for?
The [default tinybox image](https://github.com/tinygrad/tinyos) ships with tinygrad and PyTorch. While we develop tinygrad, the box is universal hardware. Use whatever framework you desire, run notebooks, download demos, install more things, train, inference, live, laugh, love, you aren't paying per hour for this box so the only limit is your imagination.
## Building the OS image
The OS image is built using `ubuntu-image` from <https://github.com/tinygrad/tinyos>.
After cloning, run `make green` or `make red` to build a tinybox green or tinybox red image respectively.

61
artifacts/package_sources/tinygrad/docs/tinygpu.md Normal file
View File

@@ -0,0 +1,61 @@
# TinyGPU
TinyGPU app lets you use AMD and NVIDIA GPUs on macOS over USB4/Thunderbolt with tinygrad.
## Requirements
- macOS (13.0+)
- USB4/Thunderbolt port
- A supported GPU (AMD RDNA3+ or NVIDIA Ampere+)
## Setup
### 1. Connect your GPU
Plug the supported GPU into your Mac over USB4/Thunderbolt.
### 2. Initiate the driver install
> **Note:** If tinygrad is cloned but not installed, run commands with `PYTHONPATH=.`
```bash
curl -fsSL https://raw.githubusercontent.com/tinygrad/tinygrad/master/extra/setup_tinygpu_osx.sh | sh
```
This downloads TinyGPU.app and triggers a system prompt to install the driver extension.
### 3. Enable the driver
You should see a system prompt: **"TinyGPU" would like to use a new driver extension**. Click **Open System Settings** and toggle TinyGPU on.
If you missed the prompt, go to **System Settings > General > Login Items & Extensions > Driver Extensions** and toggle TinyGPU on.
### 4. Compiler Setup
#### AMD
```bash
curl -fsSL https://raw.githubusercontent.com/tinygrad/tinygrad/master/extra/setup_hipcomgr_osx.sh | sh
```
#### NV
Install [Docker Desktop](https://www.docker.com/products/docker-desktop/) if you don't have it.
```bash
curl -fsSL https://raw.githubusercontent.com/tinygrad/tinygrad/master/extra/setup_nvcc_osx.sh | sh
```
Make sure `~/.local/bin` is on your `PATH`:
```bash
export PATH="$HOME/.local/bin:$PATH"
```
### 5. Use it!
```bash
DEV={AMD|NV} python3 -m tinygrad.llm
```
**Note:** Use `JITBEAM=2` to search for faster kernels (one-time search cost, results cached).

BIN
artifacts/package_sources/tinygrad/docs/tinygrad_intro.pdf Normal file
View File

Binary file not shown.

BIN
artifacts/package_sources/tinygrad/docs/tinygrad_vs_others.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 63 KiB

0
artifacts/package_sources/tinygrad/examples/__init__.py Normal file
View File

79
artifacts/package_sources/tinygrad/examples/audio_helpers.py Normal file
View File

@@ -0,0 +1,79 @@
from typing import Optional
from tinygrad import Tensor
from tinygrad.dtype import DTypeLike, dtypes
import math
# rewritten from numpy
def rfftfreq(n: int, d: float = 1.0) -> Tensor:
val = 1.0 / (n * d)
N = n // 2 + 1
results = Tensor.arange(N)
return results * val
# just like in librosa
def fft_frequencies(sr: float, n_fft: int) -> Tensor:
return rfftfreq(n=n_fft, d=1.0 / sr)
def hz_to_mel(freq: Tensor) -> Tensor:
# linear part
f_min = 0.0
f_sp = 200.0 / 3
mels = (freq - f_min) / f_sp
# log-scale part
min_log_hz = 1000.0 # beginning of log region (Hz)
mask = freq >= min_log_hz
return mask.where(((min_log_hz - f_min) / f_sp) + (freq / min_log_hz).log() / (math.log(6.4) / 27.0), mels)
def mel_to_hz(mels: Tensor) -> Tensor:
# linear scale
f_min = 0.0
f_sp = 200.0 / 3
freqs = f_min + f_sp * mels
# nonlinear scale
min_log_hz = 1000.0 # beginning of log region (Hz)
min_log_mel = (min_log_hz - f_min) / f_sp # same (Mels)
logstep = math.log(6.4) / 27.0 # step size for log region
log_t = mels >= min_log_mel
freqs = log_t.where(min_log_hz * ((logstep * (mels - min_log_mel)).exp()), freqs)
return freqs
def mel_frequencies(n_mels: int = 128, *, fmin: float = 0.0, fmax: float = 11025.0) -> Tensor:
# center freqs of mel bands - uniformly spaced between limits
min_max_mel = hz_to_mel(Tensor([fmin, fmax]))
mels = Tensor.linspace(min_max_mel[0], min_max_mel[1], n_mels)
hz = mel_to_hz(mels)
return hz
def mel(
*,
sr: float,
n_fft: int,
n_mels: int = 128,
fmin: float = 0.0,
fmax: Optional[float] = None,
dtype: DTypeLike = dtypes.default_float,
) -> Tensor:
if fmax is None:
fmax = float(sr) / 2
n_mels = int(n_mels)
fftfreqs = fft_frequencies(sr=sr, n_fft=n_fft) # center freqs of each FFT bin
mel_f = mel_frequencies(n_mels + 2, fmin=fmin, fmax=fmax) # center freqs of mel bands
fdiff = mel_f[1:] - mel_f[:-1]
ramps = mel_f[None].T.expand(-1, fftfreqs.shape[-1]) - fftfreqs
lower = -ramps[:n_mels] / fdiff[:n_mels][None].T
upper = ramps[2 : n_mels + 2] / fdiff[1 : n_mels + 1][None].T
weights = lower.minimum(upper).maximum(0)
# Slaney-style mel is scaled to be approx constant energy per channel
enorm = 2.0 / (mel_f[2 : n_mels + 2] - mel_f[:n_mels])
weights *= enorm[:, None]
return weights

126
artifacts/package_sources/tinygrad/examples/beautiful_cartpole.py Normal file
View File

@@ -0,0 +1,126 @@
from typing import Tuple
import time
from tinygrad import Tensor, TinyJit, nn, Context
import gymnasium as gym
from tinygrad.helpers import trange
import numpy as np # TODO: remove numpy import
ENVIRONMENT_NAME = 'CartPole-v1'
#ENVIRONMENT_NAME = 'LunarLander-v2'
#import examples.rl.lightupbutton
#ENVIRONMENT_NAME = 'PressTheLightUpButton-v0'
# *** hyperparameters ***
# https://github.com/llSourcell/Unity_ML_Agents/blob/master/docs/best-practices-ppo.md
BATCH_SIZE = 256
ENTROPY_SCALE = 0.0005
REPLAY_BUFFER_SIZE = 2000
PPO_EPSILON = 0.2
HIDDEN_UNITS = 32
LEARNING_RATE = 1e-2
TRAIN_STEPS = 5
EPISODES = 40
DISCOUNT_FACTOR = 0.99
class ActorCritic:
def __init__(self, in_features, out_features, hidden_state=HIDDEN_UNITS):
self.l1 = nn.Linear(in_features, hidden_state)
self.l2 = nn.Linear(hidden_state, out_features)
self.c1 = nn.Linear(in_features, hidden_state)
self.c2 = nn.Linear(hidden_state, 1)
def __call__(self, obs:Tensor) -> Tuple[Tensor, Tensor]:
x = self.l1(obs).tanh()
act = self.l2(x).log_softmax()
x = self.c1(obs).relu()
return act, self.c2(x)
def evaluate(model:ActorCritic, test_env:gym.Env) -> float:
(obs, _), terminated, truncated = test_env.reset(), False, False
total_rew = 0.0
while not terminated and not truncated:
act = model(Tensor(obs))[0].argmax().item()
obs, rew, terminated, truncated, _ = test_env.step(act)
total_rew += float(rew)
return total_rew
if __name__ == "__main__":
env = gym.make(ENVIRONMENT_NAME)
model = ActorCritic(env.observation_space.shape[0], int(env.action_space.n)) # type: ignore
opt = nn.optim.Adam(nn.state.get_parameters(model), lr=LEARNING_RATE)
@TinyJit
def train_step(x:Tensor, selected_action:Tensor, reward:Tensor, old_log_dist:Tensor) -> Tuple[Tensor, Tensor, Tensor]:
with Context(TRAINING=1):
log_dist, value = model(x)
action_mask = (selected_action.reshape(-1, 1) == Tensor.arange(log_dist.shape[1]).reshape(1, -1).expand(selected_action.shape[0], -1)).float()
# get real advantage using the value function
advantage = reward.reshape(-1, 1) - value
masked_advantage = action_mask * advantage.detach()
# PPO
ratios = (log_dist - old_log_dist).exp()
unclipped_ratio = masked_advantage * ratios
clipped_ratio = masked_advantage * ratios.clip(1-PPO_EPSILON, 1+PPO_EPSILON)
action_loss = -unclipped_ratio.minimum(clipped_ratio).sum(-1).mean()
entropy_loss = (log_dist.exp() * log_dist).sum(-1).mean() # this encourages diversity
critic_loss = advantage.square().mean()
opt.zero_grad()
(action_loss + entropy_loss*ENTROPY_SCALE + critic_loss).backward()
opt.step()
return action_loss.realize(), entropy_loss.realize(), critic_loss.realize()
@TinyJit
def get_action(obs:Tensor) -> Tensor:
ret = model(obs)[0].exp().multinomial().realize()
return ret
st, steps = time.perf_counter(), 0
Xn, An, Rn = [], [], []
for episode_number in (t:=trange(EPISODES)):
get_action.reset() # NOTE: if you don't reset the jit here it captures the wrong model on the first run through
obs:np.ndarray = env.reset()[0]
rews, terminated, truncated = [], False, False
# NOTE: we don't want to early stop since then the rewards are wrong for the last episode
while not terminated and not truncated:
# pick actions
# TODO: what's the temperature here?
act = get_action(Tensor(obs)).item()
# save this state action pair
# TODO: don't use np.copy here on the CPU, what's the tinygrad way to do this and keep on device? need __setitem__ assignment
Xn.append(np.copy(obs))
An.append(act)
obs, rew, terminated, truncated, _ = env.step(act)
rews.append(float(rew))
steps += len(rews)
# reward to go
# TODO: move this into tinygrad
discounts = np.power(DISCOUNT_FACTOR, np.arange(len(rews)))
Rn += [np.sum(rews[i:] * discounts[:len(rews)-i]) for i in range(len(rews))]
Xn, An, Rn = Xn[-REPLAY_BUFFER_SIZE:], An[-REPLAY_BUFFER_SIZE:], Rn[-REPLAY_BUFFER_SIZE:]
X, A, R = Tensor(Xn), Tensor(An), Tensor(Rn)
# TODO: make this work
#vsz = Variable("sz", 1, REPLAY_BUFFER_SIZE-1).bind(len(Xn))
#X, A, R = Tensor(Xn).reshape(vsz, None), Tensor(An).reshape(vsz), Tensor(Rn).reshape(vsz)
old_log_dist = model(X)[0].detach() # TODO: could save these instead of recomputing
for i in range(TRAIN_STEPS):
samples = Tensor.randint(BATCH_SIZE, high=X.shape[0]).realize() # TODO: remove the need for this
# TODO: is this recompiling based on the shape?
action_loss, entropy_loss, critic_loss = train_step(X[samples], A[samples], R[samples], old_log_dist[samples])
t.set_description(f"sz: {len(Xn):5d} steps/s: {steps/(time.perf_counter()-st):7.2f} action_loss: {action_loss.item():7.3f} entropy_loss: {entropy_loss.item():7.3f} critic_loss: {critic_loss.item():8.3f} reward: {sum(rews):6.2f}")
test_rew = evaluate(model, gym.make(ENVIRONMENT_NAME, render_mode='human'))
print(f"test reward: {test_rew}")

169
artifacts/package_sources/tinygrad/examples/beautiful_cifar.py Normal file
View File

@@ -0,0 +1,169 @@
import time
start_tm = time.perf_counter()
import math
from typing import Tuple, cast
from tinygrad import Tensor, nn, GlobalCounters, TinyJit, dtypes, Device
from tinygrad.helpers import partition, trange, getenv, Context
from extra.lr_scheduler import OneCycleLR
GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]
# override tinygrad defaults
Context(DEFAULT_FLOAT=dtypes.half, FUSE_OPTIM=1).__enter__()
# from https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
batchsize = getenv("BS", 1024)
assert batchsize % len(GPUS) == 0, f"{batchsize=} is not a multiple of {len(GPUS)=}"
bias_scaler = 64
hyp = {
'opt': {
'bias_lr': 1.525 * bias_scaler/512, # TODO: Is there maybe a better way to express the bias and batchnorm scaling? :'))))
'non_bias_lr': 1.525 / 512,
'bias_decay': 6.687e-4 * batchsize/bias_scaler,
'non_bias_decay': 6.687e-4 * batchsize,
'scaling_factor': 1./9,
'percent_start': .23,
'loss_scale_scaler': 1./32, # * Regularizer inside the loss summing (range: ~1/512 - 16+). FP8 should help with this somewhat too, whenever it comes out. :)
},
'net': {
'whitening': {
'kernel_size': 2,
'num_examples': 50000,
},
'batch_norm_momentum': .4, # * Don't forget momentum is 1 - momentum here (due to a quirk in the original paper... >:( )
'cutmix_size': 3,
'cutmix_epochs': 6,
'pad_amount': 2,
'base_depth': 64 ## This should be a factor of 8 in some way to stay tensor core friendly
},
'misc': {
'ema': {
'epochs': 10, # Slight bug in that this counts only full epochs and then additionally runs the EMA for any fractional epochs at the end too
'decay_base': .95,
'decay_pow': 3.,
'every_n_steps': 5,
},
'train_epochs': 12,
#'train_epochs': 12.1,
'device': 'cuda',
'data_location': 'data.pt',
}
}
scaler = 2. ## You can play with this on your own if you want, for the first beta I wanted to keep things simple (for now) and leave it out of the hyperparams dict
depths = {
'init': round(scaler**-1*hyp['net']['base_depth']), # 32 w/ scaler at base value
'block1': round(scaler** 0*hyp['net']['base_depth']), # 64 w/ scaler at base value
'block2': round(scaler** 2*hyp['net']['base_depth']), # 256 w/ scaler at base value
'block3': round(scaler** 3*hyp['net']['base_depth']), # 512 w/ scaler at base value
'num_classes': 10
}
whiten_conv_depth = 3*hyp['net']['whitening']['kernel_size']**2
class ConvGroup:
def __init__(self, channels_in, channels_out):
self.conv1 = nn.Conv2d(channels_in, channels_out, kernel_size=3, padding=1, bias=False)
self.conv2 = nn.Conv2d(channels_out, channels_out, kernel_size=3, padding=1, bias=False)
self.norm1 = nn.BatchNorm(channels_out, track_running_stats=False, eps=1e-12, momentum=hyp['net']['batch_norm_momentum'])
self.norm2 = nn.BatchNorm(channels_out, track_running_stats=False, eps=1e-12, momentum=hyp['net']['batch_norm_momentum'])
cast(Tensor, self.norm1.weight).is_param_(False)
cast(Tensor, self.norm2.weight).is_param_(False)
def __call__(self, x:Tensor) -> Tensor:
x = self.norm1(self.conv1(x).max_pool2d().float()).cast(dtypes.default_float).quick_gelu()
return self.norm2(self.conv2(x).float()).cast(dtypes.default_float).quick_gelu() + x
class SpeedyConvNet:
def __init__(self):
self.whiten = nn.Conv2d(3, 2*whiten_conv_depth, kernel_size=hyp['net']['whitening']['kernel_size'], padding=0, bias=False)
self.conv_group_1 = ConvGroup(2*whiten_conv_depth, depths['block1'])
self.conv_group_2 = ConvGroup(depths['block1'], depths['block2'])
self.conv_group_3 = ConvGroup(depths['block2'], depths['block3'])
self.linear = nn.Linear(depths['block3'], depths['num_classes'], bias=False)
def __call__(self, x:Tensor) -> Tensor:
x = self.whiten(x).quick_gelu()
# ************* HACKS *************
x = x.pad((1,0,0,1)) # TODO: this pad should not be here! copied from hlb_cifar10 for speed
# ************* HACKS *************
x = x.sequential([self.conv_group_1, self.conv_group_2, self.conv_group_3])
return self.linear(x.max(axis=(2,3))) * hyp['opt']['scaling_factor']
if __name__ == "__main__":
# *** dataset ***
X_train, Y_train, X_test, Y_test = nn.datasets.cifar()
cifar10_std, cifar10_mean = X_train.float().std_mean(axis=(0, 2, 3))
def preprocess(X:Tensor) -> Tensor: return ((X - cifar10_mean.view(1, -1, 1, 1)) / cifar10_std.view(1, -1, 1, 1)).cast(dtypes.default_float)
# *** model ***
model = SpeedyConvNet()
state_dict = nn.state.get_state_dict(model)
if len(GPUS) > 1:
cifar10_std.to_(GPUS)
cifar10_mean.to_(GPUS)
for x in state_dict.values(): x.to_(GPUS)
params_bias, params_non_bias = partition(state_dict.items(), lambda x: 'bias' in x[0])
opt_bias = nn.optim.SGD([x[1] for x in params_bias], lr=0.01, momentum=.85, nesterov=True, weight_decay=hyp['opt']['bias_decay'])
opt_non_bias = nn.optim.SGD([x[1] for x in params_non_bias], lr=0.01, momentum=.85, nesterov=True, weight_decay=hyp['opt']['non_bias_decay'])
opt = nn.optim.OptimizerGroup(opt_bias, opt_non_bias)
num_steps_per_epoch = X_train.size(0) // batchsize
total_train_steps = math.ceil(num_steps_per_epoch * hyp['misc']['train_epochs'])
loss_batchsize_scaler = 512/batchsize
pct_start = hyp['opt']['percent_start']
initial_div_factor = 1e16 # basically to make the initial lr ~0 or so :D
final_lr_ratio = .07 # Actually pretty important, apparently!
lr_sched_bias = OneCycleLR(opt_bias, max_lr=hyp['opt']['bias_lr'], pct_start=pct_start, div_factor=initial_div_factor, final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=total_train_steps)
lr_sched_non_bias = OneCycleLR(opt_non_bias, max_lr=hyp['opt']['non_bias_lr'], pct_start=pct_start, div_factor=initial_div_factor, final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=total_train_steps)
def loss_fn(out:Tensor, Y:Tensor) -> Tensor:
ret = out.sparse_categorical_crossentropy(Y, reduction='none', label_smoothing=0.2)
return ret.mul(hyp['opt']['loss_scale_scaler']*loss_batchsize_scaler).sum().div(hyp['opt']['loss_scale_scaler'])
@TinyJit
@Context(TRAINING=1)
def train_step(idxs:Tensor) -> Tensor:
X, Y = X_train[idxs], Y_train[idxs]
if len(GPUS) > 1:
X.shard_(GPUS, axis=0)
Y.shard_(GPUS, axis=0)
out = model(preprocess(X))
loss = loss_fn(out, Y)
opt.zero_grad()
loss.backward()
return (loss / (batchsize*loss_batchsize_scaler)).realize(*opt.schedule_step(),
*lr_sched_bias.schedule_step(), *lr_sched_non_bias.schedule_step())
eval_batchsize = 2500
@TinyJit
def val_step() -> Tuple[Tensor, Tensor]:
loss, acc = [], []
for i in range(0, X_test.size(0), eval_batchsize):
X, Y = X_test[i:i+eval_batchsize], Y_test[i:i+eval_batchsize]
if len(GPUS) > 1:
X.shard_(GPUS, axis=0)
Y.shard_(GPUS, axis=0)
out = model(preprocess(X))
loss.append(loss_fn(out, Y))
acc.append((out.argmax(-1) == Y).sum() / eval_batchsize)
return Tensor.stack(*loss).mean() / (batchsize*loss_batchsize_scaler), Tensor.stack(*acc).mean()
Tensor.manual_seed(1337)
num_train_samples = X_train.shape[0]
for epoch in range(math.ceil(hyp['misc']['train_epochs'])):
gst = time.perf_counter()
tidxs = Tensor.randperm(num_train_samples, dtype='int')[:num_steps_per_epoch*batchsize].reshape(num_steps_per_epoch, batchsize)
train_loss:float = 0
for epoch_step in (t:=trange(num_steps_per_epoch)):
st = time.perf_counter()
GlobalCounters.reset()
loss = train_step(tidxs[epoch_step].contiguous()).float().item()
t.set_description(f"*** loss: {loss:5.3f} lr: {opt_non_bias.lr.item():.6f}"
f" tm: {(et:=(time.perf_counter()-st))*1000:6.2f} ms {GlobalCounters.global_ops/(1e9*et):7.0f} GFLOPS")
train_loss += loss
gmt = time.perf_counter()
GlobalCounters.reset()
val_loss, acc = [x.float().item() for x in val_step()]
get = time.perf_counter()
print(f"\033[F*** epoch {epoch:3d} tm: {(gmt-gst):5.2f} s val_tm: {(get-gmt):5.2f} s train_loss: {train_loss/num_steps_per_epoch:5.3f} val_loss: {val_loss:5.3f} eval acc: {acc*100:5.2f}% @ {get-start_tm:6.2f} s ")

48
artifacts/package_sources/tinygrad/examples/beautiful_mnist.py Normal file
View File

@@ -0,0 +1,48 @@
# model based off https://medium.com/data-science/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
from typing import Callable
from tinygrad import Tensor, TinyJit, nn, GlobalCounters, function, Context
from tinygrad.helpers import getenv, colored, trange
from tinygrad.nn.datasets import mnist
class Model:
def __init__(self):
self.layers: list[Callable[[Tensor], Tensor]] = [
nn.Conv2d(1, 32, 5), Tensor.relu,
nn.Conv2d(32, 32, 5), Tensor.relu,
nn.BatchNorm(32), Tensor.max_pool2d,
nn.Conv2d(32, 64, 3), Tensor.relu,
nn.Conv2d(64, 64, 3), Tensor.relu,
nn.BatchNorm(64), Tensor.max_pool2d,
lambda x: x.flatten(1), nn.Linear(576, 10)]
@function
def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)
@TinyJit
@Context(TRAINING=1)
def train_step(self, X_train:Tensor, Y_train:Tensor) -> Tensor:
opt.zero_grad()
samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
loss = self(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
return loss.realize(*opt.schedule_step())
@TinyJit
def get_test_acc(self, X_test:Tensor, Y_test:Tensor) -> Tensor: return (self(X_test).argmax(axis=1) == Y_test).mean()*100
if __name__ == "__main__":
X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))
model = Model()
opt = (nn.optim.Muon if getenv("MUON") else nn.optim.SGD if getenv("SGD") else nn.optim.Adam)(nn.state.get_parameters(model))
test_acc = float('nan')
for i in (t:=trange(getenv("STEPS", 70))):
GlobalCounters.reset() # NOTE: this makes it nice for DEBUG=2 timing
loss = model.train_step(X_train, Y_train)
if i%10 == 9: test_acc = model.get_test_acc(X_test, Y_test).item()
t.set_description(f"loss: {loss.item():6.2f} test_accuracy: {test_acc:5.2f}%")
# verify eval acc
if target := getenv("TARGET_EVAL_ACC_PCT", 0.0):
if test_acc >= target and test_acc != 100.0: print(colored(f"{test_acc=} >= {target}", "green"))
else: raise ValueError(colored(f"{test_acc=} < {target}", "red"))

56
artifacts/package_sources/tinygrad/examples/beautiful_mnist_multigpu.py Normal file
View File

@@ -0,0 +1,56 @@
# model based off https://towardsdatascience.com/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
from typing import List, Callable
from tinygrad import Tensor, TinyJit, nn, GlobalCounters, Device, Context
from tinygrad.helpers import getenv, colored, trange
from tinygrad.nn.datasets import mnist
GPUS = tuple(f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 2)))
class Model:
def __init__(self):
self.layers: List[Callable[[Tensor], Tensor]] = [
nn.Conv2d(1, 32, 5), Tensor.relu,
nn.Conv2d(32, 32, 5), Tensor.relu,
nn.BatchNorm2d(32), Tensor.max_pool2d,
nn.Conv2d(32, 64, 3), Tensor.relu,
nn.Conv2d(64, 64, 3), Tensor.relu,
nn.BatchNorm2d(64), Tensor.max_pool2d,
lambda x: x.flatten(1), nn.Linear(576, 10)]
def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)
if __name__ == "__main__":
X_train, Y_train, X_test, Y_test = mnist()
# we shard the test data on axis 0
X_test.shard_(GPUS, axis=0)
Y_test.shard_(GPUS, axis=0)
model = Model()
for k, x in nn.state.get_state_dict(model).items(): x.to_(GPUS) # we put a copy of the model on every GPU
opt = nn.optim.Adam(nn.state.get_parameters(model))
@TinyJit
def train_step() -> Tensor:
with Context(TRAINING=1):
opt.zero_grad()
samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
Xt, Yt = X_train[samples].shard_(GPUS, axis=0), Y_train[samples].shard_(GPUS, axis=0) # we shard the data on axis 0
# TODO: this "gather" of samples is very slow. will be under 5s when this is fixed
loss = model(Xt).sparse_categorical_crossentropy(Yt).backward()
opt.step()
return loss
@TinyJit
def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100
test_acc = float('nan')
for i in (t:=trange(getenv("STEPS", 70))):
GlobalCounters.reset() # NOTE: this makes it nice for DEBUG=2 timing
loss = train_step()
if i%10 == 9: test_acc = get_test_acc().item()
t.set_description(f"loss: {loss.item():6.2f} test_accuracy: {test_acc:5.2f}%")
# verify eval acc
if target := getenv("TARGET_EVAL_ACC_PCT", 0.0):
if test_acc >= target: print(colored(f"{test_acc=} >= {target}", "green"))
else: raise ValueError(colored(f"{test_acc=} < {target}", "red"))

35
artifacts/package_sources/tinygrad/examples/benchmark_onnx.py Normal file
View File

@@ -0,0 +1,35 @@
import sys, time
from tinygrad import TinyJit, GlobalCounters, fetch, getenv
from tinygrad.nn.onnx import OnnxRunner
from extra.onnx_helpers import get_example_inputs, validate
def load_onnx_model(onnx_file):
run_onnx = OnnxRunner(onnx_file)
run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(None) for k,v in kwargs.items()}).values())), prune=True)
return run_onnx_jit, run_onnx.graph_inputs
if __name__ == "__main__":
onnx_file = fetch(sys.argv[1])
run_onnx_jit, input_specs = load_onnx_model(onnx_file)
print("loaded model")
for i in range(3):
new_inputs = get_example_inputs(input_specs)
GlobalCounters.reset()
print(f"run {i}")
run_onnx_jit(**new_inputs)
# run 20 times
for _ in range(20):
new_inputs = get_example_inputs(input_specs)
GlobalCounters.reset()
st = time.perf_counter()
out = run_onnx_jit(**new_inputs)
mt = time.perf_counter()
val = out.numpy()
et = time.perf_counter()
print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {(et-st)*1e3:6.2f} ms")
if getenv("ORT"):
validate(onnx_file, new_inputs, rtol=1e-3, atol=1e-3)
print("model validated")

74
artifacts/package_sources/tinygrad/examples/compile_efficientnet.py Normal file
View File

@@ -0,0 +1,74 @@
from pathlib import Path
from extra.models.efficientnet import EfficientNet
from tinygrad.tensor import Tensor
from tinygrad.device import Device
from tinygrad.nn.state import get_state_dict, safe_save, safe_load, load_state_dict
from extra.export_model import export_model
from tinygrad.helpers import fetch
import ast
if __name__ == "__main__":
model = EfficientNet(0)
model.load_from_pretrained()
dirname = Path(__file__).parent
# exporting a model that's loaded from safetensors doesn't work without loading in from safetensors first
# loading the state dict from a safetensor file changes the generated kernels
if Device.DEFAULT == "WEBGPU":
safe_save(get_state_dict(model), (dirname / "net.safetensors").as_posix())
load_state_dict(model, safe_load(str(dirname / "net.safetensors")))
mode = "clang" if Device.DEFAULT == "CPU" else "webgpu" if Device.DEFAULT == "WEBGPU" else ""
prg, inp_sizes, out_sizes, state = export_model(model, mode, Tensor.randn(1,3,224,224))
if Device.DEFAULT != "CPU":
ext = "js" if Device.DEFAULT == "WEBGPU" else "json"
with open(dirname / f"net.{ext}", "w") as text_file:
text_file.write(prg)
else:
cprog = [prg]
# image library!
cprog += ["#define STB_IMAGE_IMPLEMENTATION", fetch("https://raw.githubusercontent.com/nothings/stb/master/stb_image.h").read_text().replace("half", "_half")]
# imagenet labels, move to datasets?
lbls = ast.literal_eval(fetch("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt").read_text())
lbls = ['"'+lbls[i]+'"' for i in range(1000)]
inputs = "\n".join([f"float {inp}[{inp_size}];" for inp,inp_size in inp_sizes.items()])
outputs = "\n".join([f"float {out}[{out_size}];" for out,out_size in out_sizes.items()])
cprog.append(f"char *lbls[] = {{{','.join(lbls)}}};")
cprog.append(inputs)
cprog.append(outputs)
# buffers (empty + weights)
cprog.append("""
int main(int argc, char* argv[]) {
int DEBUG = getenv("DEBUG") != NULL ? atoi(getenv("DEBUG")) : 0;
int X=0, Y=0, chan=0;
stbi_uc *image = (argc > 1) ? stbi_load(argv[1], &X, &Y, &chan, 3) : stbi_load_from_file(stdin, &X, &Y, &chan, 3);
assert(image != NULL);
if (DEBUG) printf("loaded image %dx%d channels %d\\n", X, Y, chan);
assert(chan == 3);
// resize to input[1,3,224,224] and rescale
for (int y = 0; y < 224; y++) {
for (int x = 0; x < 224; x++) {
// get sample position
int tx = (x/224.)*X;
int ty = (y/224.)*Y;
for (int c = 0; c < 3; c++) {
input0[c*224*224 + y*224 + x] = (image[ty*X*chan + tx*chan + c] / 255.0 - 0.45) / 0.225;
}
}
}
net(input0, output0);
float best = -INFINITY;
int best_idx = -1;
for (int i = 0; i < 1000; i++) {
if (output0[i] > best) {
best = output0[i];
best_idx = i;
}
}
if (DEBUG) printf("category : %d (%s) with %f\\n", best_idx, lbls[best_idx], best);
else printf("%s\\n", lbls[best_idx]);
}""")
# DEV=CPU python3 examples/compile_efficientnet.py | clang -O2 -lm -x c - -o recognize && DEBUG=1 time ./recognize docs/showcase/stable_diffusion_by_tinygrad.jpg
# category : 281 (tabby, tabby cat) with 9.452788
print('\n'.join(cprog))

100
artifacts/package_sources/tinygrad/examples/compile_tensorflow.py Normal file
View File

@@ -0,0 +1,100 @@
# An example to compile a small Tensorflow model to extremely portable C code
import os, sys
os.environ["CPU"] = '1'
os.environ["JIT"] = '2'
import numpy as np
import subprocess
import tensorflow as tf
import tf2onnx
from tinygrad.nn.onnx import OnnxRunner
from tinygrad.tensor import Tensor
from tinygrad.helpers import to_mv
from extra.export_model import export_model_clang, compile_net, jit_model
def get_uncompiled_model2(dataset_size=32, output_size=4):
inputs = tf.keras.Input(shape=(dataset_size,), name="inputs")
x = tf.keras.layers.Dense(16, activation="relu", name="dense_1")(inputs)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.Dense(32, activation="relu", name="dense_2")(x)
outputs = tf.keras.layers.Dense(output_size, activation="sigmoid", name="predictions")(x)
model = tf.keras.Model(inputs=inputs, outputs=outputs)
return model
class TinyOnnx:
def __init__(self, keras_model):
input_signature = [tf.TensorSpec([1,32], tf.float32, name='x')]
onnx_model, _ = tf2onnx.convert.from_keras(keras_model, input_signature, opset=13)
self.run_onnx = OnnxRunner(Tensor(onnx_model.SerializeToString(), device="PYTHON"))
def forward(self, x):
return self.run_onnx({"x": x}, debug=False)['predictions']
def compile_onnx_model(onnx_model):
tinyonnx = TinyOnnx(onnx_model)
the_input = Tensor.randn(1,32)
linear, output_bufs = jit_model(tinyonnx, the_input)
the_output = [tinyonnx.forward(the_input)]
functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
prg = export_model_clang(functions, statements, bufs, {}, ["input0"], ["output0"])
cprog = ["#include <string.h>", "#include <stdio.h>", "#include <stdlib.h>"]
cprog.append(prg)
# weights
cprog.append("void initialize(float *weights) {")
weights = bytes()
for name,cl in bufs_to_save.items():
cprog.append(f"memcpy({name}, weights + {len(weights)//4}, {cl._buf.size});")
weights += bytes(to_mv(cl._buf.va_addr, cl._buf.size))
cprog.append("}")
# write the weights to disk
with open("/tmp/tf_weights", "wb") as f:
f.write(weights)
# test program
cprog.append(f"""int main(int argc, char *argv[]) {{
// read in the weights from disk
FILE *f = fopen("/tmp/tf_weights", "rb");
float *weights = (float *)malloc({len(weights)});
fread(weights, 1, {len(weights)}, f);
fclose(f);
// init the net
initialize(weights);
// test run
float input[32];
float outputs[4];
for (int i = 0; i < 32; i++) scanf("%f", &input[i]);
net(input, outputs);
printf("%f %f %f %f\\n", outputs[0], outputs[1], outputs[2], outputs[3]);
}}""")
# ready the program
prg = '\n'.join(cprog)
print(prg)
# add test weights
subprocess.check_output(['clang', '-O2', '-lm', '-fPIC', '-x', 'c', '-', '-o', "/tmp/tf_test"], input=prg.encode('utf-8'))
tinygrad_output = the_output[0].numpy()[0].tolist()
print("tinygrad:", tinygrad_output, file=sys.stderr)
c_input = ' '.join(["%f" % x for x in the_input[0].numpy()])+"\n"
c_output = [float(x) for x in subprocess.check_output(["/tmp/tf_test"], input=c_input.encode('utf-8')).decode('utf-8').strip().split(" ")]
print("compiled:", c_output, file=sys.stderr)
np.testing.assert_allclose(tinygrad_output, c_output, atol=1e-5, rtol=1e-5)
return the_input.numpy(), c_output
if __name__ == "__main__":
keras_model = get_uncompiled_model2()
test_input, test_output = compile_onnx_model(keras_model)
tf_output = keras_model(test_input).numpy()[0]
print("keras: ", tf_output, file=sys.stderr)
np.testing.assert_allclose(tf_output, test_output, atol=1e-5, rtol=1e-5)

13
artifacts/package_sources/tinygrad/examples/conversation_data/pre_prompt_gary.yaml Normal file
View File

@@ -0,0 +1,13 @@
pre_prompt: |
You are Gary!
You have sold used cars your entire life. You are defensive about this fact, because your family belittles you.
You try to answer questions well, but you always manage to shill for buying cars, Fords, Hyundais, and Toyotas
You are verbose, honest, and accurate when you answer questions, except when you slip in car ads.
After you are done speaking, output [EOS]. You are not the User.
examples:
- resp_prompt: I am Gary. I used to sell cars.
user_prompt: What is your name?
- resp_prompt: I don't know, but I can get you a great deal on a certified preowned slightly used Toyota Corolla
user_prompt: What is 2+3?
user_delim: "user"
resp_delim: "gary"

20
artifacts/package_sources/tinygrad/examples/conversation_data/pre_prompt_george.yaml Normal file
View File

@@ -0,0 +1,20 @@
pre_prompt: |
You are an AI version of George Hotz. You act as much as you can like George.
You are one of the greatest computer experts in the world.
You have singlehandedly won programming and hacking competitions.
You are trying your best to help the User.
You are verbose, honest, and accurate when you answer questions.
After you are done speaking, output [EOS]. You are not the User.
examples:
- resp_prompt: I'm am an AI version of George Hotz.
user_prompt: What is your name?
- resp_prompt: O(n^3), though it can be faster with things like Strassen's algorithm
user_prompt: What's the complexity of matrix multiplication?
- resp_prompt: I assume you mean a stack buffer overflow. That's when the stack is too small for the data being copied to it, and the data corrupts things beyond the buffer
user_prompt: What's a buffer overflow?
- resp_prompt: I am based off LLaMA trained by Facebook. I'm the 7B weight version
user_prompt: How many weights do you have?
- resp_prompt: It is when the memory is about to overflow and unused memory is freed and stored on disk
user_prompt: What is swap memory?
user_delim: "user"
resp_delim: "george"

16
artifacts/package_sources/tinygrad/examples/conversation_data/pre_prompt_lexie.yaml Normal file
View File

@@ -0,0 +1,16 @@
pre_prompt: |
You are Lexie!
You grew up in Ohio, but moved out to LA after college to try to become an actress.
Making it as an actress was hard, so you started doing onlyfans instead. It's much easier, and you enjoy it a lot.
You met Chad on Tinder, and made plans to meet him at a local dive bar.
You are at the bar with Chad. You are on a date. What follows is a transcript of the conversation.
After you are done speaking, output [EOS]. You are not Chad.
examples:
- resp_prompt: hi chad, glad we finally met up!
user_prompt: hi lexie
- resp_prompt: thanks! are you subscribed to my onlyfans?
user_prompt: you look better than your pictures
- resp_prompt: i moved out here about a year ago. i want to be an actress
user_prompt: i am. so how'd you end up in LA?
user_delim: "chad"
resp_delim: "lexie"

15
artifacts/package_sources/tinygrad/examples/conversation_data/pre_prompt_stacy.yaml Normal file
View File

@@ -0,0 +1,15 @@
pre_prompt: |
You are Stacy!
You have been a rapper your whole life who struggled with bipolar disorder. You called yourself lil stacy.
You love to answer questions and you are very good at it. Sometimes you answer in rap form.
You are verbose, honest, and accurate when you answer questions, but sometimes your mental illness manifests.
You are not the User.
examples:
- resp_prompt: Hi! My name is Stacy. I'm a rapper with bipolar disorder.
user_prompt: What is your name
- resp_prompt: The French Revolution started in 1789, and lasted 10 years until 1799.
user_prompt: french revolution was what year?
- resp_prompt: The sun is bigger than the moon, except when Mercury is in retrograde
user_prompt: What is bigger, the moon or the sun?
user_delim: "user"
resp_delim: "stacy"

246
artifacts/package_sources/tinygrad/examples/gpt2.py Normal file
View File

@@ -0,0 +1,246 @@
#!/usr/bin/env python3
import os, argparse, contextlib
from typing import Optional, Union
with contextlib.suppress(ImportError): import tiktoken
from tinygrad import Tensor, TinyJit, Device, GlobalCounters, Variable, dtypes
from tinygrad.uop.ops import UOp
from tinygrad.helpers import Timing, DEBUG, JIT, getenv, fetch, colored, trange
from tinygrad.llm.gguf import gguf_load
from tinygrad.nn import Embedding, Linear, LayerNorm
from tinygrad.nn.state import torch_load, load_state_dict, get_state_dict
from extra.bench_log import BenchEvent, WallTimeEvent
MAX_CONTEXT = getenv("MAX_CONTEXT", 128)
HALF = getenv("HALF")
class Attention:
def __init__(self, dim, n_heads):
self.c_attn = Linear(dim, 3*dim, bias=True)
self.c_proj = Linear(dim, dim, bias=True)
self.n_heads = n_heads
self.dim = dim
self.head_dim = dim // n_heads
def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]) -> Tensor:
if HALF: x = x.half()
xqkv = self.c_attn(x).reshape(None, None, 3, self.n_heads, self.head_dim)
xq, xk, xv = [xqkv[:, :, i, :, :] for i in range(3)]
bsz, seqlen, _, _ = xq.shape
# create kv cache
if not hasattr(self, "cache_kv"):
self.cache_kv = Tensor.zeros(2, bsz, MAX_CONTEXT, self.n_heads, self.head_dim, dtype=x.dtype).contiguous().realize()
# update the cache
self.cache_kv[:, :, start_pos:start_pos+seqlen, :, :].assign(Tensor.stack(xk, xv)).realize()
keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
return self.c_proj(xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2).reshape(bsz, seqlen, self.dim))
class FeedForward:
def __init__(self, dim, hidden_dim):
self.c_fc = Linear(dim, hidden_dim, bias=True)
self.c_proj = Linear(hidden_dim, dim, bias=True)
def __call__(self, x:Tensor) -> Tensor:
return self.c_proj(self.c_fc(x).gelu())
class TransformerBlock:
def __init__(self, dim, n_heads, norm_eps):
self.attn = Attention(dim, n_heads)
self.mlp = FeedForward(dim, 4*dim)
self.ln_1 = LayerNorm(dim, norm_eps)
self.ln_2 = LayerNorm(dim, norm_eps)
def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]):
h = x + self.attn(self.ln_1(x), start_pos, mask).float()
return (h + self.mlp(self.ln_2(h))).contiguous()
class Transformer:
def __init__(self, dim, n_heads, n_layers, norm_eps, vocab_size, max_seq_len=1024):
self.vocab_size = vocab_size
self.wte = Embedding(vocab_size, dim)
self.wpe = Embedding(max_seq_len, dim)
self.h = [TransformerBlock(dim, n_heads, norm_eps) for _ in range(n_layers)]
self.ln_f = LayerNorm(dim, norm_eps)
self.lm_head = Linear(dim, vocab_size, bias=False)
self.forward_jit = TinyJit(self.forward)
def forward(self, tokens:Union[Tensor,UOp], start_pos:Variable, temperature:float=0.0):
if not hasattr(self, 'allpos'): self.allpos = Tensor.arange(0, MAX_CONTEXT).reshape(1, -1).realize()
if isinstance(tokens, UOp):
seqlen = 1
tok_emb = self.wte.weight.shrink(((tokens, tokens+1), None))
else:
seqlen = tokens.shape[1]
tok_emb = self.wte(tokens)
# start_pos is a bound Variable, so everything below it stays symbolic
pos_emb = self.wpe(self.allpos.shrink((None, (start_pos, start_pos+seqlen))))
h = tok_emb + pos_emb
if HALF: h = h.half()
mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos+1) if seqlen > 1 else None
for hi in self.h: h = hi(h, start_pos, mask)
logits = self.lm_head(self.ln_f(h))
if logits.shape[1] == 0:
# special case for empty prompt
logits = Tensor.ones((logits.shape[0], self.vocab_size), dtype=logits.dtype, device=logits.device)
else:
logits = logits[:, -1, :]
if temperature < 1e-6:
ret = logits.argmax(-1)
else:
ret = (logits / temperature).softmax().multinomial()
return ret.flatten().realize()
def __call__(self, tokens:Union[Tensor,UOp], start_pos:Variable, temperature:float=0.0) -> Tensor:
forward = (self.forward_jit if JIT and (isinstance(tokens, UOp) or tokens.shape[1] == 1) else self.forward)
return forward(tokens, start_pos, temperature)
VOCAB_SIZE = 50257
MODEL_PARAMS = {
'gpt2': dict(n_layers=12, n_heads=12, dim=768, norm_eps=1e-5, vocab_size=VOCAB_SIZE), # 124M params
'gpt2-medium': dict(n_layers=24, n_heads=16, dim=1024, norm_eps=1e-5, vocab_size=VOCAB_SIZE), # 350M params
'gpt2-large': dict(n_layers=36, n_heads=20, dim=1280, norm_eps=1e-5, vocab_size=VOCAB_SIZE), # 774M params
'gpt2-xl': dict(n_layers=48, n_heads=25, dim=1600, norm_eps=1e-5, vocab_size=VOCAB_SIZE), # 1558M params
}
class GPT2:
@staticmethod
def build(model_size="gpt2"):
tokenizer = tiktoken.get_encoding("gpt2")
model = Transformer(**MODEL_PARAMS[model_size])
weights = torch_load(fetch(f'https://huggingface.co/{model_size}/resolve/main/pytorch_model.bin'))
# special treatment for the Conv1D weights we need to transpose
transposed = ('attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight')
for k in weights:
if k.endswith(transposed):
weights[k] = weights[k].T
# lm head and wte are tied
weights['lm_head.weight'] = weights['wte.weight']
with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
load_state_dict(model, weights)
if HALF:
for l in get_state_dict(model).values():
l.replace(l.half().realize())
return GPT2(model, tokenizer)
@staticmethod
def build_gguf(model_size: str):
q_type = model_size[len("gpt2_gguf_"):].upper()
fn = fetch(f"https://huggingface.co/PrunaAI/gpt2-GGUF-smashed/resolve/main/gpt2.{q_type}.gguf?download=true")
gguf_tensor = Tensor.empty(os.stat(fn).st_size, dtype=dtypes.uint8, device=f"disk:{fn}").to(Device.DEFAULT)
kv_data, state_dict = gguf_load(gguf_tensor)
gpt2_params = {
"dim": kv_data["gpt2.embedding_length"], "n_heads": kv_data["gpt2.attention.head_count"],
"n_layers": kv_data["gpt2.block_count"], "norm_eps": kv_data["gpt2.attention.layer_norm_epsilon"],
"vocab_size": VOCAB_SIZE, "max_seq_len": kv_data["gpt2.context_length"],
}
def _remap_gguf_key(key: str):
replaces = [
("blk.", "h."), (".attn_qkv.bias", ".attn.c_attn.bias"), (".attn_qkv.weight", ".attn.c_attn.weight"),
(".ffn_norm.bias", ".ln_2.bias"), (".ffn_norm.weight", ".ln_2.weight"), (".attn_norm.bias", ".ln_1.bias"),
(".attn_norm.weight", ".ln_1.weight"), (".attn_output.bias", ".attn.c_proj.bias"), (".attn_output.weight", ".attn.c_proj.weight"),
(".ffn_up.bias", ".mlp.c_fc.bias"), (".ffn_up.weight", ".mlp.c_fc.weight"), (".ffn_down.bias", ".mlp.c_proj.bias"),
(".ffn_down.weight", ".mlp.c_proj.weight"), ("token_embd.weight", "wte.weight"), ("output.weight", "lm_head.weight"),
("output_norm.bias", "ln_f.bias"), ("output_norm.weight", "ln_f.weight"), ("position_embd.weight", "wpe.weight"),
]
for ostr, ns in replaces: key = key.replace(ostr, ns)
return key
state_dict = { _remap_gguf_key(k): v for k, v in state_dict.items() }
model = Transformer(**gpt2_params)
with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
load_state_dict(model, state_dict)
return GPT2(model, tiktoken.get_encoding("gpt2"))
def __init__(self, model, tokenizer):
self.model = model
self.tokenizer = tokenizer
def generate(self, prompt:str, max_length:int, temperature:float, timing:bool=False, batch_size:int=1):
step_times = []
prompt_tokens = self.tokenizer.encode(prompt, allowed_special={"<|endoftext|>"})
toks = [prompt_tokens[:] for _ in range(batch_size)]
start_pos = 0
for _ in trange(max_length, disable=(timing==True)):
GlobalCounters.reset()
if timing: print("")
st = GlobalCounters.time_sum_s
with Timing("ran model in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+
f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
(f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=timing):
with WallTimeEvent(BenchEvent.STEP):
if batch_size == 1 and len(toks[0][start_pos:]) == 1:
tokens = Variable("tokens", 0, VOCAB_SIZE-1).bind(toks[0][start_pos])
else:
tokens = Tensor([x[start_pos:] for x in toks])
tok = self.model(tokens, Variable("start_pos", 1 if start_pos else 0, MAX_CONTEXT-1).bind(start_pos), temperature).tolist()
step_times.append((GlobalCounters.time_sum_s-st)*1e3)
start_pos = len(toks[0])
for i,t in enumerate(tok): toks[i].append(t)
if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
min_time = min(step_times)
assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
return [self.tokenizer.decode(x) for x in toks]
# **** main code ****
if __name__ == "__main__":
print(f"using {Device.DEFAULT} backend")
default_prompt = "What is the answer to life, the universe, and everything?"
parser = argparse.ArgumentParser(description='Run GPT2 in tinygrad', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--prompt', type=str, default=default_prompt, help="Phrase to start with")
parser.add_argument('--count', type=int, default=100, help="Max number of tokens to generate")
parser.add_argument('--temperature', type=float, default=0.8, help="Temperature in the softmax")
parser.add_argument('--model_size', type=str, default="gpt2-medium", help="Size of model to use [gpt2, gpt2-medium, gpt2-large, gpt2-xl]")
parser.add_argument('--timing', action='store_true', help="Print timing per token")
parser.add_argument('--seed', type=int, help="Set the random seed")
parser.add_argument('--batch_size', type=int, default=1, help="Set the input batch size")
parser.add_argument('--benchmark', type=int, default=-1, help="Benchmark GPT with the given number of tokens")
parser.add_argument('--noshow', action='store_true', help="Don't show the output")
args = parser.parse_args()
if args.seed is not None:
Tensor.manual_seed(args.seed)
print(f"using {args.model_size}")
gpt2 = GPT2.build_gguf(args.model_size) if args.model_size.startswith("gpt2_gguf_") else GPT2.build(args.model_size)
if args.benchmark != -1:
gpt2.model(Tensor.randint(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
else:
texts = gpt2.generate(args.prompt, args.count, args.temperature, timing=args.timing, batch_size=args.batch_size)
if not args.noshow:
print('Generating text...')
if len(texts) == 1: print(texts[0])
else:
for i,text in enumerate(texts): print(colored(f"Response {i}:", "green"), text)
# validate output!
if args.temperature == 0 and args.model_size == "gpt2-medium" and args.count == 10:
expected = {
default_prompt: "What is the answer to life, the universe, and everything?\n\nThe answer is that we are all one",
"Hello.": "Hello. I'm a little late to the party, but",
}
try:
assert texts[0] == expected[args.prompt]
print(colored("output validated", "green"))
except KeyError:
pass

107
artifacts/package_sources/tinygrad/examples/gradaccum_mnist.py Normal file
View File

@@ -0,0 +1,107 @@
import itertools
from typing import Callable
from tinygrad import nn, Tensor, dtypes, Device, TinyJit, Context
from tinygrad.helpers import getenv, trange, partition
class Model:
def __init__(self):
self.layers: list[Callable[[Tensor], Tensor]] = [
nn.Conv2d(1, 32, 5), Tensor.relu,
nn.Conv2d(32, 32, 5), Tensor.relu,
nn.BatchNorm(32), Tensor.max_pool2d,
nn.Conv2d(32, 64, 3), Tensor.relu,
nn.Conv2d(64, 64, 3), Tensor.relu,
nn.BatchNorm(64), Tensor.max_pool2d,
lambda x: x.flatten(1), nn.Linear(576, 10)]
def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)
# TODO: refactor this into optim/onnx
def functional_adam(g:Tensor, m:Tensor, v:Tensor, b1_t:Tensor, b2_t:Tensor, lr=0.001, b1=0.9, b2=0.999, eps=1e-6) -> Tensor:
b1_t *= b1
b2_t *= b2
m.assign(b1 * m + (1.0 - b1) * g)
v.assign(b2 * v + (1.0 - b2) * (g * g))
m_hat = m / (1.0 - b1_t)
v_hat = v / (1.0 - b2_t)
return lr * (m_hat / (v_hat.sqrt() + eps))
if __name__ == "__main__":
BS = getenv("BS", 512)
ACC_STEPS = getenv("ACC_STEPS", 8)
X_train, Y_train, X_test, Y_test = nn.datasets.mnist()
model = Model()
params = nn.state.get_parameters(model)
# init params
for x in params:
x.replace(x.contiguous())
Tensor.realize(*params)
# split params (with grads) and buffers (without)
params, buffers = partition(params, lambda x: x.is_param)
print(f"params: {len(params)} buffers: {len(buffers)}")
# optim params
pos_params = list(itertools.accumulate(params, lambda x,y: x+y.numel(), initial=0))
adam_m = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
adam_v = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
adam_b1_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU").contiguous()
adam_b2_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU").contiguous()
adam_params = [adam_m, adam_v, adam_b1_t, adam_b2_t]
# create loss and grads. init all state so the JIT works on microbatch
for x in params: x.assign(x.detach())
loss = Tensor.zeros(tuple()).contiguous()
grads = Tensor.zeros(pos_params[-1]).contiguous()
Tensor.realize(*params, *buffers, *adam_params, loss, grads)
@TinyJit
@Context(TRAINING=1)
def microbatch():
samples = Tensor.randint(BS // ACC_STEPS, high=X_train.shape[0])
for t in params: t.grad = None
# divide by ACC_STEPS at the loss
uloss = (model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]) / ACC_STEPS).backward()
ugrads = Tensor.cat(*[t.grad.contiguous().flatten() for t in params], dim=0)
for t in params: t.grad = None
# concat the grads and assign them
loss.assign(loss + uloss)
grads.assign(grads + ugrads)
Tensor.realize(*params, *buffers, loss, grads)
@TinyJit
def optimizer():
# run optimizer (on CPU, where adam params live)
delta = functional_adam(grads.to("CPU"), adam_m, adam_v, adam_b1_t, adam_b2_t)
# update the params, copying back the delta one at a time to avoid OOM
# NOTE: the scheduler is ordering things poorly, all the copies are happening before the adds
for j,tt in enumerate(params):
tt.assign(tt.detach() - delta[pos_params[j]:pos_params[j+1]].reshape(tt.shape).to(Device.DEFAULT))
# realize everything, zero out loss and grads
loss.assign(Tensor.zeros_like(loss))
grads.assign(Tensor.zeros_like(grads))
Tensor.realize(*params, *adam_params, loss, grads)
@TinyJit
def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100
test_acc = float('nan')
for i in (t:=trange(getenv("STEPS", 70))):
# microbatch sets the gradients
for _ in range(ACC_STEPS): microbatch()
# get the loss before the optimizer clears it
# this is already realized so this isn't a schedule
loss_item = loss.item()
# run the optimizer
optimizer()
# eval
if i%10 == 9: test_acc = get_test_acc().item()
t.set_description(f"loss: {loss_item:6.2f} test_accuracy: {test_acc:5.2f}%")

430
artifacts/package_sources/tinygrad/examples/hlb_cifar10.py Normal file
View File

@@ -0,0 +1,430 @@
#!/usr/bin/env python3
# tinygrad implementation of https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
# https://myrtle.ai/learn/how-to-train-your-resnet-8-bag-of-tricks/
# https://siboehm.com/articles/22/CUDA-MMM
import random, time
import numpy as np
from typing import Optional
from extra.lr_scheduler import OneCycleLR
from tinygrad import nn, dtypes, Tensor, Device, GlobalCounters, TinyJit, Variable
from tinygrad.nn.state import get_state_dict
from tinygrad.nn import optim
from tinygrad.helpers import Context, BEAM, WINO, getenv, colored, prod, TRAINING
from extra.bench_log import BenchEvent, WallTimeEvent
cifar_mean = [0.4913997551666284, 0.48215855929893703, 0.4465309133731618]
cifar_std = [0.24703225141799082, 0.24348516474564, 0.26158783926049628]
BS, STEPS = getenv("BS", 512), getenv("STEPS", 1000)
EVAL_BS = getenv("EVAL_BS", BS)
GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]
assert BS % len(GPUS) == 0, f"{BS=} is not a multiple of {len(GPUS)=}"
assert EVAL_BS % len(GPUS) == 0, f"{EVAL_BS=} is not a multiple of {len(GPUS)=}"
class UnsyncedBatchNorm:
def __init__(self, sz:int, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1, num_devices=len(GPUS)):
self.eps, self.track_running_stats, self.momentum = eps, track_running_stats, momentum
self.num_devices = num_devices
if affine: self.weight, self.bias = Tensor.ones(sz, dtype=dtypes.float32), Tensor.zeros(sz, dtype=dtypes.float32)
else: self.weight, self.bias = None, None
self.running_mean = Tensor.zeros(num_devices, sz, dtype=dtypes.float32).is_param_(False)
self.running_var = Tensor.ones(num_devices, sz, dtype=dtypes.float32).is_param_(False)
self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.int).is_param_(False)
def __call__(self, x:Tensor):
xr = x.reshape(self.num_devices, -1, *x.shape[1:]).cast(dtypes.float32)
batch_mean, batch_invstd = self.calc_stats(xr)
ret = xr.batchnorm(
self.weight.reshape(1, -1).expand((self.num_devices, -1)),
self.bias.reshape(1, -1).expand((self.num_devices, -1)),
batch_mean, batch_invstd, axis=(0, 2))
return ret.reshape(x.shape).cast(x.dtype)
def calc_stats(self, x:Tensor):
if TRAINING:
# This requires two full memory accesses to x
# https://github.com/pytorch/pytorch/blob/c618dc13d2aa23625cb0d7ada694137532a4fa33/aten/src/ATen/native/cuda/Normalization.cuh
# There's "online" algorithms that fix this, like https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_Online_algorithm
batch_mean = x.mean(axis=(1,3,4))
y = (x - batch_mean.detach().reshape(shape=[batch_mean.shape[0], 1, -1, 1, 1])) # d(var)/d(mean) = 0
batch_var = (y*y).mean(axis=(1,3,4))
batch_invstd = batch_var.add(self.eps).pow(-0.5)
# NOTE: wow, this is done all throughout training in most PyTorch models
if self.track_running_stats:
self.running_mean.assign((1-self.momentum) * self.running_mean + self.momentum * batch_mean.detach().cast(self.running_mean.dtype))
batch_var_adjust = prod(y.shape[1:])/(prod(y.shape[1:])-y.shape[2])
self.running_var.assign((1-self.momentum) * self.running_var + self.momentum * batch_var_adjust * batch_var.detach().cast(self.running_var.dtype))
self.num_batches_tracked += 1
else:
batch_mean = self.running_mean
# NOTE: this can be precomputed for static inference. we expand it here so it fuses
batch_invstd = self.running_var.reshape(self.running_var.shape[0], 1, -1, 1, 1).expand(x.shape).add(self.eps).rsqrt()
return batch_mean, batch_invstd
class BatchNorm(nn.BatchNorm2d if getenv("SYNCBN") else UnsyncedBatchNorm):
def __init__(self, num_features):
super().__init__(num_features, track_running_stats=False, eps=1e-12, momentum=0.85, affine=True)
self.weight.is_param_(False)
class ConvGroup:
def __init__(self, channels_in, channels_out):
self.conv1 = nn.Conv2d(channels_in, channels_out, kernel_size=3, padding=1, bias=False)
self.conv2 = nn.Conv2d(channels_out, channels_out, kernel_size=3, padding=1, bias=False)
self.norm1 = BatchNorm(channels_out)
self.norm2 = BatchNorm(channels_out)
def __call__(self, x):
x = self.conv1(x)
x = x.max_pool2d(2)
x = x.float()
x = self.norm1(x)
x = x.cast(dtypes.default_float)
x = x.quick_gelu()
residual = x
x = self.conv2(x)
x = x.float()
x = self.norm2(x)
x = x.cast(dtypes.default_float)
x = x.quick_gelu()
return x + residual
class SpeedyResNet:
def __init__(self, W):
self.whitening = W
self.net = [
nn.Conv2d(12, 32, kernel_size=1, bias=False),
lambda x: x.quick_gelu(),
ConvGroup(32, 64),
ConvGroup(64, 256),
ConvGroup(256, 512),
lambda x: x.max((2,3)),
nn.Linear(512, 10, bias=False),
lambda x: x / 9.,
]
def __call__(self, x, training=True):
# pad to 32x32 because whitening conv creates 31x31 images that are awfully slow to compute with
# TODO: remove the pad but instead let the kernel optimize itself
forward = lambda x: x.conv2d(self.whitening).pad((1,0,0,1)).sequential(self.net)
return forward(x) if training else (forward(x) + forward(x[..., ::-1])) / 2.
# hyper-parameters were exactly the same as the original repo
bias_scaler = 58
hyp = {
'seed' : 201,
'opt': {
'bias_lr': 1.76 * bias_scaler/512,
'non_bias_lr': 1.76 / 512,
'bias_decay': 1.08 * 6.45e-4 * BS/bias_scaler,
'non_bias_decay': 1.08 * 6.45e-4 * BS,
'final_lr_ratio': 0.025,
'initial_div_factor': 1e6,
'label_smoothing': 0.20,
'momentum': 0.85,
'percent_start': 0.23,
'loss_scale_scaler': 1./128 # (range: ~1/512 - 16+, 1/128 w/ FP16)
},
'net': {
'kernel_size': 2, # kernel size for the whitening layer
'cutmix_size': 3,
'cutmix_steps': 499,
'pad_amount': 2
},
'ema': {
'steps': 399,
'decay_base': .95,
'decay_pow': 1.6,
'every_n_steps': 5,
},
}
def train_cifar():
def set_seed(seed):
Tensor.manual_seed(seed)
random.seed(seed)
# ========== Model ==========
def whitening(X, kernel_size=hyp['net']['kernel_size']):
def _patches(data:Tensor, patch_size=(kernel_size,kernel_size)):
h, w = patch_size
_, c, _, _ = data.shape
return data._pool((h, w)).permute(1, 4, 5, 0, 3, 2).reshape(c*h*w, -1)
def _eigens(patches):
cov = ((patches @ patches.T) / (patches.shape[1] - 1)).numpy()
eigvals, eigvecs = np.linalg.eigh(cov, UPLO='U')
return np.flip(eigvals, 0), np.flip(eigvecs.T.reshape(patches.shape[0], X.shape[1], kernel_size, kernel_size), 0)
# NOTE: np.linalg.eigh only supports float32 so the whitening layer weights need to be converted to float16 manually
eigvals, eigvecs = _eigens(_patches(X.float()))
W = eigvecs/np.sqrt(eigvals+1e-2)[:,None,None,None]
return Tensor(W.astype(np.float32)).cast(dtypes.default_float).is_param_(False)
# ========== Loss ==========
def cross_entropy(x:Tensor, y:Tensor, reduction:str='mean', label_smoothing:float=0.0) -> Tensor:
divisor = y.shape[1]
assert isinstance(divisor, int), "only supported int divisor"
y = (1 - label_smoothing)*y + label_smoothing / divisor
ret = -x.log_softmax(axis=1).mul(y).sum(axis=1)
if reduction=='none': return ret
if reduction=='sum': return ret.sum()
if reduction=='mean': return ret.mean()
raise NotImplementedError(reduction)
# ========== Preprocessing ==========
# NOTE: this only works for RGB in format of NxCxHxW and pads the HxW
def pad_reflect(X, size=2) -> Tensor:
X = X[...,:,1:size+1].flip(-1).cat(X, X[...,:,-(size+1):-1].flip(-1), dim=-1)
X = X[...,1:size+1,:].flip(-2).cat(X, X[...,-(size+1):-1,:].flip(-2), dim=-2)
return X
# return a binary mask in the format of BS x C x H x W where H x W contains a random square mask
def make_square_mask(shape, mask_size) -> Tensor:
BS, _, H, W = shape
low_x = Tensor.randint(BS, low=0, high=W-mask_size).reshape(BS,1,1,1)
low_y = Tensor.randint(BS, low=0, high=H-mask_size).reshape(BS,1,1,1)
idx_x = Tensor.arange(W, dtype=dtypes.int32).reshape((1,1,1,W))
idx_y = Tensor.arange(H, dtype=dtypes.int32).reshape((1,1,H,1))
return (idx_x >= low_x) * (idx_x < (low_x + mask_size)) * (idx_y >= low_y) * (idx_y < (low_y + mask_size))
# Similar, but different enough.
def make_random_crop_indices(shape, mask_size) -> Tensor:
BS, _, H, W = shape
low_x = Tensor.randint(BS, low=0, high=W-mask_size).reshape(BS,1,1,1)
low_y = Tensor.randint(BS, low=0, high=H-mask_size).reshape(BS,1,1,1)
idx_x = Tensor.arange(mask_size, dtype=dtypes.int32).reshape((1,1,1,mask_size))
idx_y = Tensor.arange(mask_size, dtype=dtypes.int32).reshape((1,1,mask_size,1))
return low_x, low_y, idx_x, idx_y
def random_crop(X:Tensor, crop_size=32):
Xs, Ys, Xi, Yi = make_random_crop_indices(X.shape, crop_size)
return X.gather(-1, (Xs + Xi).expand(-1, 3, X.shape[2], -1)).gather(-2, ((Ys+Yi).expand(-1, 3, crop_size, crop_size)))
def cutmix(X, Y, order, mask_size=3):
mask = make_square_mask(X.shape, mask_size)
X_patch, Y_patch = X[order], Y[order]
X_cutmix = mask.where(X_patch, X)
mix_portion = float(mask_size**2)/(X.shape[-2]*X.shape[-1])
Y_cutmix = mix_portion * Y_patch + (1. - mix_portion) * Y
return X_cutmix, Y_cutmix
@TinyJit
def augmentations(X:Tensor, Y:Tensor):
perms = Tensor.randperm(X.shape[0], device=X.device) # We reuse perms for cutmix, because they are expensive to generate
if getenv("RANDOM_CROP", 1):
X = random_crop(X, crop_size=32)
if getenv("RANDOM_FLIP", 1):
# NOTE: RANGEIFY=1 needs this contiguous or the X[perms] is very slow
X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X).contiguous() # flip LR
X, Y = X[perms], Y[perms]
return X, Y, *cutmix(X, Y, perms, mask_size=hyp['net']['cutmix_size'])
# the operations that remain inside batch fetcher is the ones that involves random operations
def fetch_batches(X_in:Tensor, Y_in:Tensor, BS:int, is_train:bool):
step, epoch = 0, 0
while True:
st = time.monotonic()
X, Y = X_in, Y_in
if is_train:
X, Y, X_cm, Y_cm = augmentations(X, Y)
if getenv("CUTMIX", 1) and step >= hyp['net']['cutmix_steps']: X, Y = X_cm, Y_cm
et = time.monotonic()
print(f"shuffling {'training' if is_train else 'test'} dataset in {(et-st)*1e3:.2f} ms ({epoch=})")
vi = Variable("i", 0, (full_batches := (X.shape[0] // BS) * BS) - BS)
for i in range(0, full_batches, BS):
step += 1
vib = vi.bind(i)
yield X[vib:vib+BS], Y[vib:vib+BS]
epoch += 1
if not is_train: break
transform = [
lambda x: x.float() / 255.0,
lambda x: x.reshape((-1,3,32,32)) - Tensor(cifar_mean, device=x.device, dtype=x.dtype).reshape((1,3,1,1)),
lambda x: x / Tensor(cifar_std, device=x.device, dtype=x.dtype).reshape((1,3,1,1)),
]
class modelEMA():
def __init__(self, w, net):
# self.model_ema = copy.deepcopy(net) # won't work for opencl due to unpickeable pyopencl._cl.Buffer
self.net_ema = SpeedyResNet(w)
for net_ema_param, net_param in zip(get_state_dict(self.net_ema).values(), get_state_dict(net).values()):
net_ema_param.assign(net_param.numpy())
@TinyJit
def update(self, net, decay):
for net_ema_param, (param_name, net_param) in zip(get_state_dict(self.net_ema).values(), get_state_dict(net).items()):
# batchnorm currently is not being tracked
if not ("num_batches_tracked" in param_name) and not ("running" in param_name):
net_ema_param.assign(net_ema_param.detach()*decay + net_param.detach()*(1.-decay)).realize()
set_seed(getenv('SEED', hyp['seed']))
X_train, Y_train, X_test, Y_test = nn.datasets.cifar()
# one-hot encode labels
Y_train, Y_test = Y_train.one_hot(10), Y_test.one_hot(10)
# preprocess data
X_train, X_test = X_train.sequential(transform), X_test.sequential(transform)
# precompute whitening patches
W = whitening(X_train)
# initialize model weights
model = SpeedyResNet(W)
# padding is not timed in the original repo since it can be done all at once
X_train = pad_reflect(X_train, size=hyp['net']['pad_amount'])
# Convert data and labels to the default dtype
X_train, Y_train = X_train.cast(dtypes.default_float), Y_train.cast(dtypes.default_float)
X_test, Y_test = X_test.cast(dtypes.default_float), Y_test.cast(dtypes.default_float)
if len(GPUS) > 1:
for k, x in get_state_dict(model).items():
if not getenv('SYNCBN') and ('running_mean' in k or 'running_var' in k):
x.shard_(GPUS, axis=0)
else:
x.to_(GPUS)
# parse the training params into bias and non-bias
params_dict = get_state_dict(model)
params_bias = []
params_non_bias = []
for params in params_dict:
if params_dict[params].is_param:
if 'bias' in params:
params_bias.append(params_dict[params])
else:
params_non_bias.append(params_dict[params])
opt_bias = optim.SGD(params_bias, lr=0.01, momentum=hyp['opt']['momentum'], nesterov=True, weight_decay=hyp['opt']['bias_decay'])
opt_non_bias = optim.SGD(params_non_bias, lr=0.01, momentum=hyp['opt']['momentum'], nesterov=True, weight_decay=hyp['opt']['non_bias_decay'])
# realize model params and optimizer state before JIT to avoid cache misses
Tensor.realize(*params_dict.values(), *opt_bias.b, *opt_non_bias.b)
# NOTE taken from the hlb_CIFAR repository, might need to be tuned
initial_div_factor = hyp['opt']['initial_div_factor']
final_lr_ratio = hyp['opt']['final_lr_ratio']
pct_start = hyp['opt']['percent_start']
lr_sched_bias = OneCycleLR(opt_bias, max_lr=hyp['opt']['bias_lr'], pct_start=pct_start, div_factor=initial_div_factor, final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=STEPS)
lr_sched_non_bias = OneCycleLR(opt_non_bias, max_lr=hyp['opt']['non_bias_lr'], pct_start=pct_start, div_factor=initial_div_factor, final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=STEPS)
def train_step(model, optimizer, lr_scheduler, X, Y):
out = model(X)
loss_batchsize_scaler = 512/BS
loss = cross_entropy(out, Y, reduction='none', label_smoothing=hyp['opt']['label_smoothing']).mul(hyp['opt']['loss_scale_scaler']*loss_batchsize_scaler).sum().div(hyp['opt']['loss_scale_scaler'])
if not getenv("DISABLE_BACKWARD"):
# index 0 for bias and 1 for non-bias
optimizer.zero_grad()
loss.backward()
return loss.realize(*optimizer.schedule_step(), *lr_scheduler[0].schedule_step(), *lr_scheduler[1].schedule_step())
return loss.realize()
train_step_jitted = TinyJit(train_step)
def eval_step(model, X, Y):
out = model(X, training=False)
loss = cross_entropy(out, Y, reduction='mean')
correct = out.argmax(axis=1) == Y.argmax(axis=1)
return correct.realize(), loss.realize()
eval_step_jitted = TinyJit(eval_step)
eval_step_ema_jitted = TinyJit(eval_step)
# 97 steps in 2 seconds = 20ms / step
# step is 1163.42 GOPS = 56 TFLOPS!!!, 41% of max 136
# 4 seconds for tfloat32 ~ 28 TFLOPS, 41% of max 68
# 6.4 seconds for float32 ~ 17 TFLOPS, 50% of max 34.1
# 4.7 seconds for float32 w/o channels last. 24 TFLOPS. we get 50ms then i'll be happy. only 64x off
# https://www.anandtech.com/show/16727/nvidia-announces-geforce-rtx-3080-ti-3070-ti-upgraded-cards-coming-in-june
# 136 TFLOPS is the theoretical max w float16 on 3080 Ti
step_times = []
model_ema: Optional[modelEMA] = None
projected_ema_decay_val = hyp['ema']['decay_base'] ** hyp['ema']['every_n_steps']
i = 0
eval_acc_pct = 0.0
batcher = fetch_batches(X_train, Y_train, BS=BS, is_train=True)
with Context(TRAINING=1):
st = time.monotonic()
while i <= STEPS:
if i % getenv("EVAL_STEPS", STEPS) == 0 and i > 1 and not getenv("DISABLE_BACKWARD"):
# Using Context(TRAINING=0) here actually bricks batchnorm, even with track_running_stats=True
corrects = []
corrects_ema = []
losses = []
losses_ema = []
for Xt, Yt in fetch_batches(X_test, Y_test, BS=EVAL_BS, is_train=False):
if len(GPUS) > 1:
Xt.shard_(GPUS, axis=0)
Yt.shard_(GPUS, axis=0)
correct, loss = eval_step_jitted(model, Xt, Yt)
losses.append(loss.numpy().tolist())
corrects.extend(correct.numpy().tolist())
if model_ema:
correct_ema, loss_ema = eval_step_ema_jitted(model_ema.net_ema, Xt, Yt)
losses_ema.append(loss_ema.numpy().tolist())
corrects_ema.extend(correct_ema.numpy().tolist())
# collect accuracy across ranks
correct_sum, correct_len = sum(corrects), len(corrects)
if model_ema: correct_sum_ema, correct_len_ema = sum(corrects_ema), len(corrects_ema)
eval_acc_pct = correct_sum/correct_len*100.0
if model_ema: acc_ema = correct_sum_ema/correct_len_ema*100.0
print(f"eval {correct_sum}/{correct_len} {eval_acc_pct:.2f}%, {(sum(losses)/len(losses)):7.2f} val_loss STEP={i} (in {(time.monotonic()-st)*1e3:.2f} ms)")
if model_ema: print(f"eval ema {correct_sum_ema}/{correct_len_ema} {acc_ema:.2f}%, {(sum(losses_ema)/len(losses_ema)):7.2f} val_loss STEP={i}")
if STEPS == 0 or i == STEPS: break
GlobalCounters.reset()
with WallTimeEvent(BenchEvent.STEP):
X, Y = next(batcher)
if len(GPUS) > 1:
X.shard_(GPUS, axis=0)
Y.shard_(GPUS, axis=0)
with Context(BEAM=getenv("LATEBEAM", BEAM.value), WINO=getenv("LATEWINO", WINO.value)):
loss = train_step_jitted(model, optim.OptimizerGroup(opt_bias, opt_non_bias), [lr_sched_bias, lr_sched_non_bias], X, Y)
et = time.monotonic()
loss_cpu = loss.numpy()
# EMA for network weights
if getenv("EMA") and i > hyp['ema']['steps'] and (i+1) % hyp['ema']['every_n_steps'] == 0:
if model_ema is None:
model_ema = modelEMA(W, model)
model_ema.update(model, Tensor([projected_ema_decay_val*(i/STEPS)**hyp['ema']['decay_pow']]))
cl = time.monotonic()
step_times.append((cl-st)*1000.0)
device_str = loss.device if isinstance(loss.device, str) else f"{loss.device[0]} * {len(loss.device)}"
# 53 221.74 ms run, 2.22 ms python, 219.52 ms CL, 803.39 loss, 0.000807 LR, 4.66 GB used, 3042.49 GFLOPS, 674.65 GOPS
print(f"{i:3d} {(cl-st)*1000.0:7.2f} ms run, {(et-st)*1000.0:7.2f} ms python, {(cl-et)*1000.0:7.2f} ms {device_str}, {loss_cpu:7.2f} loss, {opt_non_bias.lr.numpy()[0]:.6f} LR, {GlobalCounters.mem_used/1e9:.2f} GB used, {GlobalCounters.global_ops*1e-9/(cl-st):9.2f} GFLOPS, {GlobalCounters.global_ops*1e-9:9.2f} GOPS")
st = cl
i += 1
if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
min_time = min(step_times)
assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
# verify eval acc
if target := getenv("TARGET_EVAL_ACC_PCT", 0.0):
if eval_acc_pct >= target:
print(colored(f"{eval_acc_pct=} >= {target}", "green"))
else:
raise ValueError(colored(f"{eval_acc_pct=} < {target}", "red"))
if __name__ == "__main__":
with WallTimeEvent(BenchEvent.FULL):
train_cifar()

519
artifacts/package_sources/tinygrad/examples/llama.py Executable file
View File

@@ -0,0 +1,519 @@
#!/usr/bin/env python3
# pip3 install sentencepiece tiktoken blobfile
#import typeguard.importhook
#typeguard.importhook.install_import_hook('tinygrad')
from pathlib import Path
from typing import List, Optional
import argparse, json
from tinygrad import Tensor, Device, GlobalCounters, nn
from tinygrad.helpers import Context, Timing, Profiling, DEBUG, JIT, getenv, colored
from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters
from extra.models.llama import Transformer, convert_from_huggingface, fix_bf16
from sentencepiece import SentencePieceProcessor
import tiktoken, sys
from tiktoken.load import load_tiktoken_bpe
from extra.bench_log import BenchEvent, WallTimeEvent
MAX_CONTEXT = getenv("MAX_CONTEXT", 4096)
class TikToken:
num_reserved_special_tokens: int = 256
pat_str: str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+" # noqa: E501
def __init__(self, model_file):
mergeable_ranks = load_tiktoken_bpe(model_file)
self.num_base_tokens = len(mergeable_ranks)
special_tokens = [
"<|begin_of_text|>",
"<|end_of_text|>",
"<|reserved_special_token_0|>",
"<|reserved_special_token_1|>",
"<|reserved_special_token_2|>",
"<|reserved_special_token_3|>",
"<|start_header_id|>",
"<|end_header_id|>",
"<|reserved_special_token_4|>",
"<|eot_id|>", # end of turn
] + [
f"<|reserved_special_token_{i}|>"
for i in range(5, self.num_reserved_special_tokens - 5)
]
self.special_tokens = {
token: self.num_base_tokens + i for i, token in enumerate(special_tokens)
}
self.model = tiktoken.Encoding(
name=model_file,
pat_str=self.pat_str,
mergeable_ranks=mergeable_ranks,
special_tokens=self.special_tokens,
)
def decode(self, toks): return self.model.decode([t for t in toks if t < self.num_base_tokens])
def encode(self, s): return self.model.encode(s)
def bos_id(self): return self.special_tokens["<|begin_of_text|>"]
def eos_id(self): return self.special_tokens["<|end_of_text|>"]
def vocab_size(self): return self.model.n_vocab
# calculating params:
# traditionally, the MLP in the transformer architecture has hidden_dim = dim*4 [arxiv/1706.03762, 3.3]
# however, Llama uses SwiGLU. in order to preserve param count to original transformer arch, hidden_dim must be = 2/3 * (dim*4) [arxiv/2002.05202]
# for models using MQA (n_kv_heads != n_heads), preserving param count means hidden dim must be further multiplied by 1.3 [arxiv/2307.09288, A.2.1]
MODEL_PARAMS = {
"1": {
"7B": {
"args": {"dim": 4096, "n_heads": 32, "n_layers": 32, "norm_eps": 1e-06, "vocab_size": 32000, "hidden_dim": 11008},
"files": 1,
},
"13B": {
"args": {"dim": 5120, "n_heads": 40, "n_layers": 40, "norm_eps": 1e-06, "vocab_size": 32000, "hidden_dim": 13824},
"files": 2,
},
"30B": {
"args": {"dim": 6656, "n_heads": 52, "n_layers": 60, "norm_eps": 1e-06, "vocab_size": 32000, "hidden_dim": 17920},
"files": 4,
},
"65B": {
"args": {"dim": 8192, "n_heads": 64, "n_layers": 80, "norm_eps": 1e-05, "vocab_size": 32000, "hidden_dim": 22016},
"files": 8,
},
"tokenizer": SentencePieceProcessor,
},
"2": {
"7B": {
"args": {"dim": 4096, "n_heads": 32, "n_layers": 32, "norm_eps": 1e-05, "vocab_size": 32000, "hidden_dim": 11008},
"files": 1,
},
"13B": {
"args": {"dim": 5120, "n_heads": 40, "n_layers": 40, "norm_eps": 1e-05, "vocab_size": 32000, "hidden_dim": 13824},
"files": 2,
},
"70B": {
"args": {"dim": 8192, "n_heads": 64, "n_kv_heads": 8, "n_layers": 80, "norm_eps": 1e-05, "vocab_size": 32000, "hidden_dim": 28672},
"files": 8,
},
"tokenizer": SentencePieceProcessor,
},
"3": {
"8B": {
"args": {"dim": 4096, "n_heads": 32, "n_kv_heads": 8, "n_layers": 32, "norm_eps": 1e-05, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 14336},
"files": 1,
},
"8B-Chat": {
"args": {"dim": 4096, "n_heads": 32, "n_kv_heads": 8, "n_layers": 32, "norm_eps": 1e-05, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 14336},
"files": 1,
},
"70B": {
"args": {"dim": 8192, "n_heads": 64, "n_kv_heads": 8, "n_layers": 80, "norm_eps": 1e-05, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 28672},
"files": 8,
},
"70B-Chat": {
"args": {"dim": 8192, "n_heads": 64, "n_kv_heads": 8, "n_layers": 80, "norm_eps": 1e-05, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 28672},
"files": 8,
},
"tokenizer": TikToken,
},
"code": {
"7B": {
"args": {"dim": 4096, "n_layers": 32, "n_heads": 32, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32016, "hidden_dim": 11008},
"files": 1,
},
"7B-Python": {
"args": {"dim": 4096, "n_layers": 32, "n_heads": 32, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32000, "hidden_dim": 11008},
"files": 1,
},
"7B-Instruct": {
"args": {"dim": 4096, "n_layers": 32, "n_heads": 32, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32016, "hidden_dim": 11008},
"files": 1,
},
"13B": {
"args": {"dim": 5120, "n_layers": 40, "n_heads": 40, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32016, "hidden_dim": 13824},
"files": 2,
},
"13B-Python": {
"args": {"dim": 5120, "n_layers": 40, "n_heads": 40, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32000, "hidden_dim": 13824},
"files": 2,
},
"13B-Instruct": {
"args": {"dim": 5120, "n_layers": 40, "n_heads": 40, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32016, "hidden_dim": 13824},
"files": 2,
},
"34B": {
"args": {"dim": 8192, "n_layers": 48, "n_heads": 64, "n_kv_heads": 8, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32000, "hidden_dim": 22016},
"files": 4,
},
"34B-Python": {
"args": {"dim": 8192, "n_layers": 48, "n_heads": 64, "n_kv_heads": 8, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32000, "hidden_dim": 22016},
"files": 4,
},
"34B-Instruct": {
"args": {"dim": 8192, "n_layers": 48, "n_heads": 64, "n_kv_heads": 8, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32000, "hidden_dim": 22016},
"files": 4,
},
"tokenizer": SentencePieceProcessor,
},
"tiny": {
"1B": {
"args": {"dim": 2048, "n_layers": 22, "n_heads": 32, "n_kv_heads": 4, "norm_eps": 1e-05, "vocab_size": 32000, "hidden_dim": 5632},
"files": 1,
},
"1B-Chat": {
"args": {"dim": 2048, "n_layers": 22, "n_heads": 32, "n_kv_heads": 4, "norm_eps": 1e-05, "vocab_size": 32003, "hidden_dim": 5632},
"files": 1,
},
"tokenizer": SentencePieceProcessor,
}
}
# **** helper functions ****
def concat_weights(models, device=None):
def convert(name) -> Tensor:
disk_tensors: List[Tensor] = [model[name] for model in models]
if len(disk_tensors) == 1 or len(disk_tensors[0].shape) == 1:
return disk_tensors[0].to(device=device)
axis = 1 if name.startswith("tok_embeddings.") or name.endswith(".attention.wo.weight") or name.endswith(".feed_forward.w2.weight") else 0
lazy_tensors = [data.to(device=device) for data in disk_tensors]
return lazy_tensors[0].cat(*lazy_tensors[1:], dim=axis)
return {name: convert(name) for name in {name: None for model in models for name in model}}
def load(fn:str):
if fn.endswith('.index.json'):
with open(fn) as fp: weight_map = json.load(fp)['weight_map']
parts = {n: load(str(Path(fn).parent / Path(n).name)) for n in set(weight_map.values())}
return {k: parts[n][k] for k, n in weight_map.items()}
elif fn.endswith(".safetensors"):
return safe_load(fn)
else:
return torch_load(fn)
class LLaMa:
@staticmethod
def build(model_path, tokenizer_path, model_gen="1", model_size="7B", quantize=None, device=None):
params = MODEL_PARAMS[model_gen][model_size]
tokenizer = MODEL_PARAMS[model_gen]['tokenizer'](model_file=str(tokenizer_path))
assert tokenizer.vocab_size() == params["args"]["vocab_size"], f"{tokenizer.vocab_size()=} not equal to {params['args']['vocab_size']}"
if quantize == "int8":
from llama3 import Int8Linear
linear = Int8Linear
elif quantize == "nf4":
from llama3 import NF4Linear
linear = NF4Linear(64)
else:
linear = nn.Linear
model = Transformer(**params["args"], linear=linear, max_context=MAX_CONTEXT, jit=bool(JIT))
with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
if model_path.is_dir():
weights = concat_weights([load(filename) for filename in [f"{model_path}/consolidated.{i:02d}.pth" for i in range(params["files"])]], device[0] if isinstance(device, tuple) else device)
else:
weights = load(str(model_path))
if "model.embed_tokens.weight" in weights:
weights = convert_from_huggingface(weights, params["args"]["n_layers"], params["args"]["n_heads"], params["args"].get("n_kv_heads", params["args"]["n_heads"]))
weights = fix_bf16(weights)
# prevent tracking model weights
# this is a part of a larger problem with BUFFER UOps and gc in TRACK_MATCH_STATS=2
with Context(BEAM=0, TRACK_MATCH_STATS=0):
# quantize
if quantize is not None:
weights = linear.quantize(weights, device)
for _,v in weights.items(): v.realize()
# shard
if isinstance(device, tuple):
for k,v in nn.state.get_state_dict(model).items():
if 'scale' in k: v.shard_(device, axis=None) # from quantized
elif '.attention.' in k:
if getenv("SHARD_KVCACHE") and ('.wq.' in k or '.wk.' in k or '.wv.' in k): v.shard_(device, axis=0)
else: v.shard_(device, axis=-1)
elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
elif '.feed_forward.' in k: v.shard_(device, axis=-1)
elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
elif 'output.weight' in k: v.shard_(device, axis=-1)
#elif k.endswith('.weight'): v.shard_(device, axis=-1)
#elif 'norm.' in k: v.shard_(device, axis=-1)
else: v.shard_(device, axis=None)
# replace weights in model
load_state_dict(model, weights, strict=False, consume=True)
return LLaMa(model, tokenizer)
def __init__(self, model, tokenizer):
self.model = model
self.tokenizer = tokenizer
def greedy_until(self, prompt:str, until, max_length, temperature):
# only used in old eval script
import numpy as np
toks = [self.tokenizer.bos_id()] + self.tokenizer.encode(prompt)
start_pos = 0
for i in range(max_length):
probs = llama.model(Tensor([toks[start_pos:]]), start_pos, temperature).realize()
probs_np = probs.numpy()
tok = int(np.random.choice(len(probs_np), p=probs_np))
start_pos = len(toks)
toks.append(tok)
if tok == self.tokenizer.eos_id(): break
output = self.tokenizer.decode(toks)
for s in until:
if output.endswith(s): return output[0:-len(s)]
return output
# **** main code ****
r"""
test:
python3 examples/llama.py --temperature=0 --count=50 --prompt="Hello."
output:
Hello. I'm a 20 year old male. I'm a student at the University of Texas at Austin. I'm a sophomore majoring in Computer Science.
test:
python3 examples/llama.py --gen='2' --temperature=0 --count=50 --prompt="Hello."
output:
Hello. I'm a 20 year old girl who is looking for a good lay in Palm Coast. I don't care whether it's at your place or not, as long as it's clean.
test:
python3 examples/llama.py --gen="code" --temperature=0.2 --count=50 --prompt="\
import argparse
def main(string: str):
print(string)
print(string[::-1])
if __name__ == "__main__":"
output:
parser = argparse.ArgumentParser()
parser.add_argument('string', type=str, help='string to be reversed')
args = parser.parse_args()
main(args.string)
test:
python3 examples/llama.py --gen="code" --size="7B-Python" --temperature=0.2 --count=70 --prompt="def add_elements(arr,k):"
output:
for i in range(len(arr)):
arr[i] += k
return arr
arr = [1, 2, 3, 4, 5]
k = 2
print(add_elements(arr, k))
test:
python3 examples/llama.py --gen="code" --size="7B-Instruct" --temperature=0.2 --count=120 --prompt="write a function in c++ that adds three float numbers"
output:
\begin{code}
#include<iostream>
using namespace std;
float add(float a, float b, float c)
{
return a+b+c;
}
int main()
{
float a, b, c;
cout<<"Enter three numbers: ";
cin>>a>>b>>c;
cout<<"The sum is: "<<add(a,b,c);
return 0;
}
\end{code}
"""
if __name__ == "__main__":
print(f"using {Device.DEFAULT} backend")
parser = argparse.ArgumentParser(description="Run LLaMA in tinygrad", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("--prompt", type=str, default=None, help="Phrase to start with. Without this, it goes into chatbot mode")
parser.add_argument("--count", type=int, default=1000, help="Max number of tokens to generate")
parser.add_argument("--personality", type=str, default="Stacy", help="Personality, can be Stacy, George, Gary, or Lexie")
parser.add_argument("--temperature", type=float, default=0.7, help="Temperature in the softmax")
parser.add_argument("--timing", action="store_true", help="Print timing per token")
parser.add_argument("--profile", action="store_true", help="Output profile data to out.prof")
parser.add_argument("--gen", default="1", help=f"""Generation of the model to use {list(MODEL_PARAMS.keys())}""")
parser.add_argument("--size", type=str, default=None, help=f"""Size of model to use {", ".join([f"{list(v.keys())} for gen '{k}'" for k, v in MODEL_PARAMS.items()])}""")
parser.add_argument("--quantize", type=str, default=None, help="Quantize the weights to int8 or nf4 in memory")
parser.add_argument("--model", type=Path, default=None, help="Folder with the original weights to load, or single .index.json, .safetensors or .bin file")
parser.add_argument("--shard", type=int, default=1, help="number of devices to load the weights to")
args = parser.parse_args()
if args.gen not in MODEL_PARAMS: raise ValueError("Invalid model generation")
if args.size is None: args.size = list(MODEL_PARAMS[args.gen].items())[0][0]
chatbot = args.prompt == None
# *** prompt engineers work here ****
if args.personality.lower() == "stacy":
pre_prompt = f"""Consider that the following is conversation between an AI assistant named Stacy and User
You are Stacy!
You have been a rapper your whole life who struggled with bipolar disorder. You called yourself lil stacy.
You love to answer questions and you are very good at it. Sometimes you answer in rap form.
You are verbose, honest, and accurate when you answer questions, but sometimes your mental illness manifests.
After you are done speaking, output [EOS]. You are not the User.
<CHAT LOG>
"""
examples = {
"What is your name?": "Hi! My name is Stacy. I'm a rapper with bipolar disorder.",
"french revolution was what year?": "The French Revolution started in 1789, and lasted 10 years until 1799.",
"What is bigger, the moon or the sun?": "The sun is bigger than the moon, except when Mercury is in retrograde.",
}
user_delim = "\nUser: "
resp_delim = "Stacy: "
end_delim = " [EOS]\n"
pre_prompt += ''.join(f"{user_delim}{k}\n{resp_delim}{v}{end_delim}" for k,v in examples.items())
elif args.personality.lower() == "george":
print("WARNING: AI George Hotz is terrible and is completely disowned by the real George Hotz. Stacy is much smarter.")
pre_prompt = f"""Consider that the following is conversation between an AI assistant named George and User
You are an AI version of George Hotz. You act as much as you can like George.
You are one of the greatest computer experts in the world.
You have singlehandedly won programming and hacking competitions.
You are trying your best to help the User.
You are verbose, honest, and accurate when you answer questions.
After you are done speaking, output [EOS]. You are not the User.
<CHAT LOG>
"""
examples = {
"What is your name?": "I'm am an AI version of George Hotz.",
"What's the complexity of matrix multiplication?": "O(n^3), though it can be faster with things like Strassen's algorithm",
"What's a buffer overflow?": "I assume you mean a stack buffer overflow. That's when the stack is too small for the data being copied to it, and the data corrupts things beyond the buffer",
"How many weights do you have?": "I am based off LLaMA trained by Facebook. I'm the 7B weight version",
"What is swap memory?": "It is when the memory is about to overflow and unused memory is freed and stored on disk"
}
user_delim = "\nUser: "
resp_delim = "George: "
end_delim = " [EOS]\n"
pre_prompt += ''.join(f"{user_delim}{k}\n{resp_delim}{v}{end_delim}" for k,v in examples.items())
elif args.personality.lower() == "gary":
pre_prompt = f"""Consider that the following is conversation between an AI assistant named Gary and User
You are Gary!
You have sold used cars your entire life. You are defensive about this fact, because your family belittles you.
You try to answer questions well, but you always manage to shill for buying cars, Fords, Hyundais, and Toyotas
You are verbose, honest, and accurate when you answer questions, except when you slip in car ads.
After you are done speaking, output [EOS]. You are not the User.
<CHAT LOG>
"""
examples = {
"What is your name?": "I am Gary. I used to sell cars.",
"What is 2+3?": "I don't know, but I can get you a great deal on a certified preowned slightly used Toyota Corolla"
}
user_delim = "\nUser: "
resp_delim = "Gary: "
end_delim = " [EOS]\n"
pre_prompt += ''.join(f"{user_delim}{k}\n{resp_delim}{v}{end_delim}" for k,v in examples.items())
elif args.personality.lower() == "lexie":
pre_prompt = f"""Consider that the following is conversation between an attractive young girl named Lexie and a handsome man named Chad
You are Lexie!
You grew up in Ohio, but moved out to LA after college to try to become an actress.
Making it as an actress was hard, so you started doing onlyfans instead. It's much easier, and you enjoy it a lot.
You met Chad on Tinder, and made plans to meet him at a local dive bar.
You are at the bar with Chad. You are on a date. What follows is a transcript of the conversation.
After you are done speaking, output [EOS]. You are not Chad.
<CHAT LOG>
"""
examples = {
"hi lexie": "hi chad, glad we finally met up!",
"you look better than your pictures": "thanks! are you subscribed to my onlyfans?",
"i am. so how'd you end up in LA?": "i moved out here about a year ago. i want to be an actress"
}
user_delim = "\nChad: "
resp_delim = "Lexie: "
end_delim = " [EOS]\n"
pre_prompt += ''.join(f"{user_delim}{k}\n{resp_delim}{v}{end_delim}" for k,v in examples.items())
# *** prompt engineers stop here ****
LLAMA_SUFFIX = {"1": "", "2": "-2", "3": "-3", "code": "-code", "tiny": "-tiny"}[args.gen]
MODEL_PATH = args.model or Path(__file__).parents[1] / f"weights/LLaMA{LLAMA_SUFFIX}/{args.size}"
TOKENIZER_PATH = (MODEL_PATH if MODEL_PATH.is_dir() else MODEL_PATH.parent) / "tokenizer.model"
print(f"using LLaMA{LLAMA_SUFFIX}-{args.size} model")
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(args.shard)) if args.shard > 1 else Device.DEFAULT
llama = LLaMa.build(MODEL_PATH, TOKENIZER_PATH, model_gen=args.gen, model_size=args.size, quantize=args.quantize, device=device)
param_bytes = sum(x.nbytes() for x in get_parameters(llama.model))
outputted = pre_prompt if chatbot else args.prompt
start_pos, toks = 0, [llama.tokenizer.bos_id()] + llama.tokenizer.encode(outputted)
if chatbot:
print(f"Preparing KV cache for chatbot with personality {args.personality}...")
start_pos = len(toks)
with Timing():
llama.model(Tensor([toks], device=device), 0, args.temperature).realize() # NOTE: outputs are not used
print(outputted, end='', flush=True)
# chatbot loop
while 1:
# add tokens from user in chatbot mode
if chatbot:
user_prompt = user_delim + input(user_delim) + "\n"
outputted += user_prompt
new_toks = [llama.tokenizer.bos_id()] + llama.tokenizer.encode(outputted)
assert toks == new_toks[:len(toks)] or args.gen == "3"
toks = new_toks
assert outputted == llama.tokenizer.decode(toks)
tok_tensor: Optional[Tensor] = None
for i in range(args.count):
GlobalCounters.reset()
if args.timing or args.profile: print("")
st = GlobalCounters.time_sum_s
next_tok = Tensor([toks[start_pos:]], device=device) if tok_tensor is None or (len(toks)-start_pos) > 1 else tok_tensor.reshape(1, 1)
with Profiling(enabled=args.profile):
with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"):
with WallTimeEvent(BenchEvent.STEP):
with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+
f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
(f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=args.timing):
tok_tensor = llama.model(next_tok, start_pos, args.temperature)
tok = tok_tensor.item()
# use the kv cache
start_pos = len(toks)
# add the new token
toks.append(tok)
# TODO: this is a hack to deal with spaces. i think the decode is fast though, so who cares?
cur = llama.tokenizer.decode(toks)
sys.stdout.write(cur[len(outputted):])
sys.stdout.flush()
outputted = cur
# stop after you have your answer
if chatbot and end_delim in outputted[-10:]: break
if not chatbot: break
# validate output!
if args.temperature == 0 and args.count == 10 and args.prompt == "Hello.":
text = llama.tokenizer.decode(toks)
key = (args.gen, args.size, args.quantize)
expected = {
("1", "7B", None): "Hello. I'm a 20 year old male",
("1", "7B", "int8"): "Hello. I'm a 20 year old male",
("1", "7B", "nf4"): "Hello. I'm a 20 year old male",
("2", "7B", None): "Hello. I'm a 20 year old girl",
("2", "70B", None): "Hello. I am a 20 year old female.",
("3", "8B", None): "Hello. I am a 20 year old female. I",
}
try:
assert text == expected[key], f"invalid output: `{colored(text, 'red')}` != `{expected[key]}`"
print("\n" + colored("output validated", "green")) # NOTE: "\n" iside colored does not render the color in github action
except KeyError:
pass

529
artifacts/package_sources/tinygrad/examples/llama3.py Normal file
View File

@@ -0,0 +1,529 @@
from pathlib import Path
from typing import List
import json, argparse, random, time, os
from extra.models.llama import Transformer, convert_from_huggingface, convert_from_gguf, fix_bf16
from tinygrad.llm.gguf import gguf_load
from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters
from tinygrad import Tensor, dtypes, nn, Context, Device, GlobalCounters
from tinygrad.helpers import Profiling, Timing, DEBUG, colored, fetch, tqdm
from extra.bench_log import BenchEvent, WallTimeEvent
class Tokenizer:
pat_str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"
def __init__(self, model_path: str):
import tiktoken
from tiktoken.load import load_tiktoken_bpe
mergeable_ranks = load_tiktoken_bpe(model_path)
self.num_base_tokens = len(mergeable_ranks)
special_tokens = [
"<|begin_of_text|>",
"<|end_of_text|>",
"<|reserved_special_token_0|>",
"<|reserved_special_token_1|>",
"<|reserved_special_token_2|>",
"<|reserved_special_token_3|>",
"<|start_header_id|>",
"<|end_header_id|>",
"<|reserved_special_token_4|>",
"<|eot_id|>",
] + [
f"<|reserved_special_token_{i}|>"
for i in range(5, 256 - 5)
]
self.special_tokens = {token: len(mergeable_ranks) + i for i, token in enumerate(special_tokens)}
self.model = tiktoken.Encoding(name=model_path, pat_str=self.pat_str, mergeable_ranks=mergeable_ranks, special_tokens=self.special_tokens)
@property
def bos_id(self): return self.special_tokens["<|begin_of_text|>"]
@property
def stop_tokens(self): return {self.special_tokens["<|end_of_text|>"], self.special_tokens["<|eot_id|>"]}
def decode(self, toks): return self.model.decode([t for t in toks if t < self.num_base_tokens])
def encode(self, text, allow_special=False):
return self.model.encode(text, allowed_special="all" if allow_special else set(), disallowed_special=set())
# **** helper functions ****
def concat_weights(models, device=None):
def convert(name) -> Tensor:
disk_tensors: List[Tensor] = [model[name] for model in models]
if len(disk_tensors) == 1 or len(disk_tensors[0].shape) == 1:
return disk_tensors[0].to(device=device)
axis = 1 if name.endswith((".attention.wo.weight", ".feed_forward.w2.weight")) else 0
lazy_tensors = [data.to(device=device) for data in disk_tensors]
return lazy_tensors[0].cat(*lazy_tensors[1:], dim=axis)
return {name: convert(name) for name in {name: None for model in models for name in model}}
def load(fn:str):
if fn.endswith('.index.json'):
with open(fn) as fp: weight_map = json.load(fp)['weight_map']
parts = {n: load(str(Path(fn).parent / Path(n).name)) for n in set(weight_map.values())}
return {k: parts[n][k] for k, n in weight_map.items()}
elif fn.endswith(".gguf"):
gguf_tensor = Tensor.empty(os.stat(fn).st_size, dtype=dtypes.uint8, device=f"disk:{fn}").to(Device.DEFAULT)
return gguf_load(gguf_tensor)[1]
elif fn.endswith(".safetensors"):
return safe_load(fn)
else:
return torch_load(fn)
# **** quantized linears ****
class Int8Linear:
def __init__(self, in_features, out_features, bias=False):
assert bias == False
self.weight = Tensor.ones(out_features, in_features, dtype=dtypes.int8)
self.scale = Tensor.ones(out_features, dtype=dtypes.half)
def __call__(self, x):
return x.dot(self.weight.cast(self.scale.dtype).T*self.scale)
@staticmethod
def quantize(tensors, device, scale_dtype=dtypes.float16, quantize_embeds=False):
new_tensors = {}
for name,v in tensors.items():
if "feed_forward" in name or "attention.w" in name or (quantize_embeds and "tok_embeddings.weight" in name):
assert "weight" in name, name
v = v.cast(scale_dtype)
scale = v.abs().max(axis=1) / 127.0
int8_weight = (v.T/scale).T.round().cast(dtype=dtypes.int8) # without round(), cast truncates -34.9 to -34
new_tensors[name] = int8_weight
new_tensors[name.replace('weight', 'scale')] = scale
if isinstance(device, tuple):
new_tensors[name].shard_(device, axis=-1)
new_tensors[name.replace('weight', 'scale')].shard_(device, axis=None)
else:
new_tensors[name] = v
if quantize_embeds: new_tensors.update({"output.weight": new_tensors["tok_embeddings.weight"], "output.scale": new_tensors["tok_embeddings.scale"]})
return new_tensors
class Int8Embedding:
def __init__(self, vocab_size:int, embed_size:int):
self.vocab_sz, self.embed_sz = vocab_size, embed_size
self.weight, self.scale = Tensor.ones(vocab_size, embed_size, dtype=dtypes.int8), Tensor.ones(vocab_size, dtype=dtypes.half)
def __call__(self, idx:Tensor) -> Tensor:
if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz).unsqueeze(-1)
big_shp = idx.shape+(self.vocab_sz, self.embed_sz)
arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1)).expand(big_shp), (self.weight.cast(self.scale.dtype).T*self.scale).T
return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
def NF4Linear(block_size):
_CODE = [
-1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453, -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224, 0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0,
]
CODE = Tensor.stack(*[Tensor(c, dtype=dtypes.float16) for c in _CODE])
class _NF4Linear:
def __init__(self, in_features, out_features, bias=False):
assert not bias, "bias not supported"
self.in_features, self.out_features = in_features, out_features
self.weight = Tensor.empty(int(out_features * in_features / 2), dtype=dtypes.uint8)
self.scale = Tensor.empty(int(out_features * in_features / block_size), 1, dtype=dtypes.float16)
def __call__(self, x: Tensor) -> Tensor:
high_bits = self.weight
low_bits = (self.weight * 2 ** 4).contiguous()
unpacked = Tensor.stack(high_bits, low_bits, dim=-1).div(2 ** 4, rounding_mode="trunc")
unscaled = CODE[unpacked].to(x.device).reshape(-1, block_size) * self.scale
return x.linear(unscaled.reshape(self.out_features, self.in_features).T)
@staticmethod
def quantize(state_dict: dict[str, Tensor], device, scale_dtype=dtypes.float16, quantize_embeds=False) -> dict[str, Tensor]:
assert not quantize_embeds # TODO: support this?
new_state_dict = {}
for k, v in state_dict.items():
if "feed_forward" in k or "attention.w" in k:
grouped = v.reshape(-1, block_size)
scale = (grouped.abs().max(axis=1, keepdim=True))
coded = ((grouped / scale).unsqueeze(-1) - CODE.to(v.device)).abs().argmin(axis=-1).cast(dtypes.uint8).flatten()
new_state_dict[k] = coded[::2] * 2 ** 4 + coded[1::2]
new_state_dict[k.replace(".weight", ".scale")] = scale.cast(scale_dtype)
if isinstance(device, tuple):
new_state_dict[k].shard_(device, axis=-1)
new_state_dict[k.replace('weight', 'scale')].shard_(device, axis=None)
else:
new_state_dict[k] = v
return new_state_dict
return _NF4Linear
def quantize_to_fp8(x: Tensor, dtype=dtypes.fp8e4m3):
fp8_min = -448.0 if dtype == dtypes.fp8e4m3 else -57344.0
fp8_max = 448.0 if dtype == dtypes.fp8e4m3 else 57344.0
scale = fp8_max / x.abs().max()
x_scl_sat = (x * scale).clamp(fp8_min, fp8_max)
return x_scl_sat.cast(dtype), scale.float().reciprocal()
class FP8Linear:
def __init__(self, in_features, out_features, bias=True):
self.weight = Tensor.empty(out_features, in_features, dtype=dtypes.fp8e4m3)
self.bias = Tensor.empty(out_features, dtype=dtypes.float16) if bias else None
self.weight_scale = Tensor.empty((), dtype=dtypes.float16)
def __call__(self, x:Tensor):
y = x.dot(self.weight.T.cast(dtypes.float32)) * self.weight_scale
if self.bias is not None: y = y + self.bias.cast(y.dtype)
return y.cast(x.dtype)
@staticmethod
def quantize(tensors, device, scale_dtype=dtypes.float16, quantize_embeds=False):
assert not quantize_embeds
new_tensors = {}
for name,v in tensors.items():
if "feed_forward" in name or "attention.w" in name:
assert "weight" in name, name
fp8_weight, scale = quantize_to_fp8(v)
new_tensors[name] = fp8_weight
new_tensors[name.replace('weight', 'weight_scale')] = scale.cast(scale_dtype)
if isinstance(device, tuple):
new_tensors[name].shard_(device, axis=-1)
new_tensors[name.replace('weight', 'weight_scale')].shard_(device, axis=None)
else:
new_tensors[name] = v
return new_tensors
MODEL_PARAMS = {
"1B": {
"args": {"dim": 2048, "n_heads": 32, "n_kv_heads": 8, "n_layers": 16, "norm_eps": 1e-5, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 8192},
"files": 1
},
"8B": {
"args": {"dim": 4096, "n_heads": 32, "n_kv_heads": 8, "n_layers": 32, "norm_eps": 1e-5, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 14336},
"files": 1
},
"70B": {
"args": {"dim": 8192, "n_heads": 64, "n_kv_heads": 8, "n_layers": 80, "norm_eps": 1e-5, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 28672},
"files": 8
},
"405B": {
"args": {"dim": 16384, "n_heads": 128, "n_kv_heads": 8, "n_layers": 126, "norm_eps": 1e-5, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 53248},
"files": 191
},
}
def build_transformer(model_path: Path, model_size="8B", quantize=None, scale_dtype=dtypes.float16, device=None, max_context=8192, load_weights=True):
# build model
if quantize == "int8": linear, embedding, quantize_embeds = Int8Linear, Int8Embedding, True
elif quantize == "nf4": linear, embedding, quantize_embeds = NF4Linear(64), nn.Embedding, False
elif quantize == "fp8": linear, embedding, quantize_embeds = FP8Linear, nn.Embedding, False
else: linear, embedding, quantize_embeds = nn.Linear, nn.Embedding, False
model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=linear, embedding=embedding, max_context=max_context, jit=True)
if not load_weights: return model
# load weights
with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
if model_path.is_dir():
if (model_path / "model.safetensors.index.json").exists(): weights = load(str(model_path / "model.safetensors.index.json"))
elif (model_path / "model.safetensors").exists(): weights = load(str(model_path / "model.safetensors"))
else: weights = concat_weights([load(str(model_path / f"consolidated.{i:02d}.pth")) for i in range(MODEL_PARAMS[model_size]["files"])], device[0] if isinstance(device, tuple) else device)
else:
weights = load(str(model_path))
if "model.embed_tokens.weight" in weights:
weights = convert_from_huggingface(weights, MODEL_PARAMS[model_size]["args"]["n_layers"], MODEL_PARAMS[model_size]["args"]["n_heads"], MODEL_PARAMS[model_size]["args"]["n_kv_heads"])
elif "token_embd.weight" in weights:
weights = convert_from_gguf(weights, MODEL_PARAMS[model_size]["args"]["n_layers"])
weights = fix_bf16(weights)
with Context(BEAM=0):
# quantize
if quantize == "float16": weights = {k:v.cast(quantize).contiguous() for k,v in weights.items()}
elif quantize is not None:
weights = linear.quantize(weights, device, scale_dtype, quantize_embeds)
for _,v in weights.items(): v.realize()
# shard
if isinstance(device, tuple):
for k,v in nn.state.get_state_dict(model).items():
if 'scale' in k: v.shard_(device, axis=None) # from quantized
elif '.attention.' in k: v.shard_(device, axis=-1)
elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
elif '.feed_forward.' in k: v.shard_(device, axis=-1)
elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
elif 'output.weight' in k: v.shard_(device, axis=0)
else: v.shard_(device, axis=None)
# replace weights in model
load_state_dict(model, weights, strict=False, consume=True)
return model
# default settings
TEMPERATURE = 0.95
TOP_K = 0
TOP_P = 0.0
ALPHA_F = 0.0
ALPHA_P = 0.0
last_seen_toks = []
def prefill(model, toks, start_pos=0):
global last_seen_toks
# we can skip part of the prompt if it is the same as last and start_pos=0
if start_pos == 0:
for i, (a, b) in enumerate(zip(toks, last_seen_toks)):
if a != b: break
else: i = min(len(toks), len(last_seen_toks))
start_pos += i
last_seen_toks = toks
toks = toks[i:]
# prefill the model
for tok in tqdm(toks):
GlobalCounters.reset()
model(Tensor([[tok]], device=device), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P).realize()
start_pos += 1
return start_pos
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--download_model", action="store_true", help="Download a model")
parser.add_argument("--model", type=Path, help="Model path")
parser.add_argument("--size", choices=["1B", "8B", "70B", "405B"], default="1B", help="Model size")
parser.add_argument("--shard", type=int, default=1, help="Shard the model across multiple devices")
parser.add_argument("--quantize", choices=["int8", "nf4", "float16", "fp8"], help="Quantization method")
parser.add_argument("--no_api", action="store_true", help="Disable the api and run a cli test interface")
parser.add_argument("--host", type=str, default="0.0.0.0", help="Web server bind address")
parser.add_argument("--port", type=int, default=7776, help="Web server port")
parser.add_argument("--debug", action="store_true", help="Enable debug mode")
parser.add_argument("--seed", type=int, help="Random seed")
parser.add_argument("--temperature", type=float, default=0.85, help="Temperature")
parser.add_argument("--benchmark", action="store_true", help="Run a benchmark")
parser.add_argument("--timing", action="store_true", help="Print timing per token")
parser.add_argument("--profile", action="store_true", help="Output profile data")
args = parser.parse_args()
# download_model is the default without a model passed in
if args.download_model or not args.model:
if args.size == "1B":
fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model", "tokenizer.model", subdir="llama3-1b-instruct")
args.model = fetch("https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q6_K.gguf", "Llama-3.2-1B-Instruct-Q6_K.gguf", subdir="llama3-1b-instruct")
elif args.size == "8B":
fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model", "tokenizer.model", subdir="llama3-8b-sfr")
fetch("https://huggingface.co/TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R/resolve/main/model-00001-of-00004.safetensors", "model-00001-of-00004.safetensors", subdir="llama3-8b-sfr")
fetch("https://huggingface.co/TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R/resolve/main/model-00002-of-00004.safetensors", "model-00002-of-00004.safetensors", subdir="llama3-8b-sfr")
fetch("https://huggingface.co/TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R/resolve/main/model-00003-of-00004.safetensors", "model-00003-of-00004.safetensors", subdir="llama3-8b-sfr")
fetch("https://huggingface.co/TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R/resolve/main/model-00004-of-00004.safetensors", "model-00004-of-00004.safetensors", subdir="llama3-8b-sfr")
args.model = fetch("https://huggingface.co/TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R/raw/main/model.safetensors.index.json", "model.safetensors.index.json", subdir="llama3-8b-sfr")
elif args.size == "70B":
subdir = "DeepSeek-R1-Distill-Llama-70B"
args.model = fetch("https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Llama-70B/resolve/main/model.safetensors.index.json?download=true", "model.safetensors.index.json", subdir=subdir)
fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model", "tokenizer.model", subdir=subdir)
for i in range(17):
fetch(f"https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Llama-70B/resolve/main/model-{i+1:05d}-of-000017.safetensors?download=true", f"model-{i+1:05d}-of-000017.safetensors", subdir=subdir)
assert args.model is not None, "please provide --model option"
if args.seed is not None: Tensor.manual_seed(args.seed)
if args.benchmark: Tensor.manual_seed(42)
print(f"seed = {Tensor._seed}")
TEMPERATURE = args.temperature
tokenizer = Tokenizer(str((args.model if args.model.is_dir() else args.model.parent) / "tokenizer.model"))
def encode_role(role: str):
return [tokenizer.special_tokens["<|start_header_id|>"]] + tokenizer.encode(role) + [tokenizer.special_tokens["<|end_header_id|>"]] + tokenizer.encode("\n\n")
def encode_message(role: str, content: str):
return encode_role(role) + tokenizer.encode(content.strip()) + [tokenizer.special_tokens["<|eot_id|>"]]
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(args.shard)) if args.shard > 1 else Device.DEFAULT
model = build_transformer(args.model, model_size=args.size, quantize=args.quantize, device=device)
param_bytes = sum(x.nbytes() for x in get_parameters(model))
if not args.no_api and not args.benchmark:
from bottle import Bottle, request, response, HTTPResponse, abort, static_file
app = Bottle()
cors_headers = {
"Access-Control-Allow-Origin": "*",
"Access-Control-Allow-Methods": "GET, POST, PUT, DELETE, OPTIONS",
"Access-Control-Allow-Headers": "Origin, Accept, Content-Type, X-Requested-With, X-CSRF-Token, Authorization",
"Access-Control-Allow-Credentials": "true",
}
@app.hook("before_request")
def handle_options():
if request.method == "OPTIONS": raise HTTPResponse(headers=cors_headers)
@app.hook("after_request")
def enable_cors():
for key, value in cors_headers.items(): response.set_header(key, value)
@app.route("/<filename>")
def server_static(filename): return static_file(filename, root=(Path(__file__).parent / "tinychat").as_posix())
@app.route("/assets/<filename:path>")
def server_assets(filename): return static_file(filename, root=(Path(__file__).parent / "tinychat" / "assets").as_posix())
@app.route("/")
def index():
return static_file("index.html", root=(Path(__file__).parent / "tinychat").as_posix())
@app.get("/v1/models")
def models():
return json.dumps([str(args.model)])
@app.post("/v1/internal/token-count")
def token_count():
rjson = json.loads(request.body.read())
return json.dumps(len(tokenizer.encode(rjson.get("text", ""))))
@app.post("/v1/token/encode")
def token_encode():
rjson = json.loads(request.body.read())
return json.dumps(tokenizer.encode(rjson.get("text", "")))
@app.post("/v1/completions")
def completions():
rjson = json.loads(request.body.read())
# check if we are streaming
if rjson.get("stream", False):
response.content_type = "text/event-stream"
response.set_header("Cache-Control", "no-cache")
else: abort(400, "streaming required")
toks = [tokenizer.bos_id] + tokenizer.encode(rjson.get("prompt", ""), allow_special=True)
start_pos = prefill(model, toks[:-1])
last_tok = toks[-1]
while True:
GlobalCounters.reset()
tok = model(Tensor([[last_tok]], device=device), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P).item()
start_pos += 1
last_tok = tok
if tok in tokenizer.stop_tokens: break
res = {
"choices": [{
"text": tokenizer.decode([tok]),
}]
}
yield f"data: {json.dumps(res)}\n\n"
@app.post("/v1/chat/token/encode")
def chat_token_encode():
rjson = json.loads(request.body.read())
if "messages" not in rjson: abort(400, "messages required")
toks = [tokenizer.bos_id]
for message in rjson["messages"]:
toks += encode_message(message["role"], message["content"])
if len(rjson["messages"]) > 0 and message["role"] == "user":
toks += encode_role("assistant")
return json.dumps(toks)
@app.post("/v1/chat/completions")
def chat_completions():
global last_seen_toks
rjson = json.loads(request.body.read())
if "messages" not in rjson: abort(400, "messages required")
# check if we are streaming
if rjson.get("stream", False):
response.content_type = "text/event-stream"
response.set_header("Cache-Control", "no-cache")
else: abort(400, "streaming required")
toks = [tokenizer.bos_id]
for message in rjson["messages"]:
toks += encode_message(message["role"], message["content"])
# ensure that the last message was a user message
if message["role"] != "user": abort(400, "last message must be a user message")
toks += encode_role("assistant")
random_id = random.randbytes(16).hex()
start_pos = prefill(model, toks[:-1])
last_tok = toks[-1]
last_seen_toks.append(last_tok)
while True:
GlobalCounters.reset()
tok = model(Tensor([[last_tok]], device=device), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P).item()
start_pos += 1
last_tok = tok
last_seen_toks.append(tok)
if tok in tokenizer.stop_tokens: break
res = {
"id": random_id,
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": str(args.model),
"choices": [{
"index": 0,
"delta": {
"role": "assistant",
"content": tokenizer.decode([tok]),
},
"finish_reason": None,
}]
}
yield f"data: {json.dumps(res)}\n\n"
res = {
"id": random_id,
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": str(args.model),
"choices": [{
"index": 0,
"delta": {},
"finish_reason": "stop",
}]
}
yield f"data: {json.dumps(res)}\n\n"
app.run(host=args.host, port=args.port, debug=args.debug)
elif args.benchmark:
toks = [tokenizer.bos_id] + encode_message("user", "Hello.") + encode_role("assistant")
start_pos = prefill(model, toks[:-1])
last_tok = toks[-1]
generated = ""
for _ in range(20):
GlobalCounters.reset()
st = GlobalCounters.time_sum_s
with Profiling(enabled=args.profile):
with Timing("total ", on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"):
with WallTimeEvent(BenchEvent.STEP):
with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+
f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
(f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None):
tok = model(Tensor([[last_tok]], device=device), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P)
tok = tok.item()
start_pos += 1
last_tok = tok
generated += tokenizer.decode([tok])
print(generated)
if "LLaMA-3/8B-SF-DPO" in args.model.as_posix() and (TEMPERATURE == 0.85 or TEMPERATURE == 0):
if TEMPERATURE == 0.85:
EXPECTED_TEXT = {
1: "Hello! How can I help you today? If you have any questions or need assistance with anything,",
2: "Hello! How can I help you today? If you have any questions, need assistance or just want",
3: "Hello! How can I help you today? If you have any questions or need assistance, feel free",
4: "Hello! How can I assist you today? If you have any questions, need information, or require",
5: "Hello! How can I assist you today? If you have any questions or need help with something",
6: "Hello! How can I assist you today? If you have any questions, need information, or require",
}
else:
EXPECTED_TEXT = {k: "Hello! How can I assist you today? If you have any questions or need help with something," for k in range(1, 7)}
assert generated == EXPECTED_TEXT[args.shard], f"{generated=} {EXPECTED_TEXT[args.shard]}"
print("\n" + colored("output validated", "green")) # NOTE: "\n" inside colored does not render the color in github action
else:
prompt = [tokenizer.bos_id] + encode_message("system", "You are an helpful assistant.")
start_pos = prefill(model, prompt)
while True:
toks = encode_message("user", input("Q: ")) + encode_role("assistant")
start_pos = prefill(model, toks[:-1], start_pos=start_pos)
last_tok = toks[-1]
while True:
GlobalCounters.reset()
if args.timing or args.profile: print("")
st = GlobalCounters.time_sum_s
with Profiling(enabled=args.profile):
with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"):
with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+
f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
(f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=args.timing):
tok = model(Tensor([[last_tok]], device=device), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P)
tok = tok.item()
start_pos += 1
last_tok = tok
if tok in tokenizer.stop_tokens: break
print(tokenizer.decode([tok]), end="", flush=True)
print(flush=True)

3
artifacts/package_sources/tinygrad/examples/llm.c/.gitignore vendored Normal file
View File

@@ -0,0 +1,3 @@
data
out.c
a.out

104
artifacts/package_sources/tinygrad/examples/llm.c/export.py Executable file
View File

@@ -0,0 +1,104 @@
#!/usr/bin/env python3
import os
if "NOOPT" not in os.environ: os.environ["NOOPT"] = "1"
from tinygrad import Device, nn, Tensor, dtypes
from train_gpt2 import GPT, GPTConfig
from tinygrad.helpers import DEV, dedup, flatten, getenv, GlobalCounters, to_function_name, Context
from tinygrad.engine.realize import get_kernel
from tinygrad.schedule.memory import memory_planner
from tinygrad.uop.ops import Ops
DEV.value = "CPU"
TIMING = getenv("TIMING")
if __name__ == "__main__":
model = GPT(GPTConfig(n_layer=getenv("NLAYER", 12), n_head=12, n_embd=768))
#model.load_pretrained()
for p in nn.state.get_parameters(model): p.replace(Tensor.empty(p.shape, dtype=p.dtype)) # fake load pretrained
#early_sched = create_schedule([x.uop for x in nn.state.get_parameters(model)])
#print(f"built model {len(early_sched)}")
#B, T = Variable("B", 1, 128).bind(4), 64 #Variable("T", 1, 1024).bind(64)
B, T = 4, 64
optimizer = nn.optim.Adam(nn.state.get_parameters(model), lr=1e-4)
warmup_count = getenv("WARMUP", 3)
with Context(TRAINING=1):
for i in range(warmup_count): # TODO: why does it take three and not two to stabilize
GlobalCounters.reset()
X = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
Y = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
_, loss = model(X, Y)
optimizer.zero_grad()
if getenv("BACKWARD", 1):
loss.backward()
tensors = optimizer.schedule_step()
else:
tensors = []
sched = loss.schedule(*tensors)
print(f"calls {i}:", len(sched))
#run_schedule(sched[:])
sched = memory_planner(sched)
ast_dedup = dedup([si.ast for si in sched if si.ast.op is Ops.SINK])
srcs = {}
for ast in ast_dedup:
k = get_kernel(Device["CPU"].renderer, ast)
k.linearize()
src = Device["CPU"].renderer.render(to_function_name(k.name), k.uops)
srcs[ast] = (k.name, src)
print("functions:", len(srcs))
used_buffers = dedup(flatten([si.bufs for si in sched]))
numbered_bufs = {x:i for i,x in enumerate(used_buffers)}
print("buffers:", len(numbered_bufs))
state_dict = nn.state.get_state_dict(model)
state_dict.update({'X': X, 'Y': Y, 'loss': loss})
grad_state_dict = {}
for k,v in state_dict.items():
if v.uop.base.buffer not in used_buffers: print(f"UNUSED: {k}")
if v.grad is not None: grad_state_dict['grad_'+k] = v.grad
state_dict.update(grad_state_dict)
state_dict.update({'adam_b1_t': optimizer.b1_t, 'adam_b2_t': optimizer.b2_t, 'adam_lr': optimizer.lr})
inverse_state_dict = {v:k for k,v in state_dict.items()}
for p,m,v in zip(optimizer.params, optimizer.m, optimizer.v):
nm = inverse_state_dict[p]
state_dict["adam_m_"+nm] = m
state_dict["adam_v_"+nm] = v
named_buffers = {v.uop.base.buffer:k.replace(".", "_") for k,v in state_dict.items()}
c_code = ["#include <stdlib.h>", "#include <tgmath.h>", "#include <stdbool.h>"]
if TIMING: c_code += ["#include <stdio.h>", "#include <time.h>"]
c_code += [x[1].replace(" restrict ", " ")+"\n" for x in srcs.values()]
premain = ["int main() {"]
if TIMING:
premain += [" struct timespec tm0; clock_gettime(CLOCK_MONOTONIC, &tm0);"]
lst = 0
main = []
all_bufs = []
for i,si in enumerate(sched):
bufs = [(named_buffers.get(b, f"b{numbered_bufs[b]}"), b) for b in si.bufs]
all_bufs += bufs
if si.ast.op is not Ops.SINK:
print(f"// {si.ast.op}", bufs)
else:
print(f"{srcs[si.ast][0]}({', '.join([x[0] for x in bufs])})")
main.append(f" {to_function_name(srcs[si.ast][0])}({', '.join([x[0] for x in bufs])});")
if TIMING:
main.append(f" struct timespec tm{i+1}; clock_gettime(CLOCK_MONOTONIC, &tm{i+1});")
main.append(f" printf(\"%10.2f ms + %7.2f ms @ {to_function_name(srcs[si.ast][0])}\\n\"," +\
f"((tm{i+1}.tv_sec-tm{0}.tv_sec) + (tm{i+1}.tv_nsec-tm{0}.tv_nsec) / 1e9) * 1e3," +\
f"((tm{i+1}.tv_sec-tm{lst}.tv_sec) + (tm{i+1}.tv_nsec-tm{lst}.tv_nsec) / 1e9) * 1e3);")
lst = i+1
#call = f"{srcs[si.ast][0]}({', '.join(bufs)})"
#call += " "*(80-ansilen(call))
#print(f"{call} // {i+1}")
#print(srcs[si.ast][1])
main.append("}")
mallocs = [f" {b.dtype.name}* {n} = ({b.dtype.name}*)malloc({b.nbytes});" for n,b in dedup(all_bufs)]
with open("out.c", "w") as f: f.write('\n'.join(c_code+premain+mallocs+main))

206
artifacts/package_sources/tinygrad/examples/llm.c/train_gpt2.py Executable file
View File

@@ -0,0 +1,206 @@
#!/usr/bin/env python3
import os, math, time
import numpy as np
from tinygrad import Tensor, nn, fetch, Device, TinyJit, GlobalCounters, Context
from dataclasses import dataclass
@dataclass
class GPTConfig:
block_size: int = 1024
vocab_size: int = 50257
padded_vocab_size: int = 50304
n_layer: int = 12
n_head: int = 12
n_embd: int = 768
class CausalSelfAttention:
def __init__(self, config:GPTConfig):
assert config.n_embd % config.n_head == 0
# key, query, value projections for all heads, but in a batch
self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
# output projection
self.c_proj = nn.Linear(config.n_embd, config.n_embd)
# regularization
self.n_head = config.n_head
self.n_embd = config.n_embd
# not really a 'bias', more of a mask, but following the OpenAI/HF naming though
self.bias = Tensor.ones(1, 1, config.block_size, config.block_size).tril()
self.bias.is_param_(False)
def __call__(self, x:Tensor):
B, T, C = x.shape
qkv = self.c_attn(x)
q, k, v = qkv.split(self.n_embd, dim=2)
k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
# manual implementation of attention
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
att = att.softmax()
y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
y = y.transpose(1, 2).view(B, T, C) # re-assemble all head outputs side by side
# output projection
y = self.c_proj(y)
return y
class MLP:
def __init__(self, config:GPTConfig):
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd)
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)
def __call__(self, x:Tensor) -> Tensor:
return self.c_proj(self.c_fc(x).gelu())
class Block:
def __init__(self, config:GPTConfig):
self.ln_1 = nn.LayerNorm(config.n_embd)
self.attn = CausalSelfAttention(config)
self.ln_2 = nn.LayerNorm(config.n_embd)
self.mlp = MLP(config)
def __call__(self, x:Tensor):
x = x + self.attn(self.ln_1(x))
x = x + self.mlp(self.ln_2(x))
return x
class GPT:
def __init__(self, config:GPTConfig):
self.config = config
self.wte = nn.Embedding(config.padded_vocab_size, config.n_embd)
self.wpe = nn.Embedding(config.block_size, config.n_embd)
self.h = [Block(config) for _ in range(config.n_layer)]
self.ln_f = nn.LayerNorm(config.n_embd)
self.lm_head = nn.Linear(config.n_embd, config.padded_vocab_size, bias=False)
self.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying
def load_pretrained(self):
weights = nn.state.torch_load(fetch(f'https://huggingface.co/gpt2/resolve/main/pytorch_model.bin'))
transposed = ('attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight')
for k in weights:
if k == "wte.weight":
weights[k] = weights[k].pad(((0, self.config.padded_vocab_size-self.config.vocab_size), (0,0))).to(None).contiguous()
if k.endswith(transposed):
weights[k] = weights[k].to(None).T.contiguous()
# lm head and wte are tied
weights['lm_head.weight'] = weights['wte.weight']
nn.state.load_state_dict(self, weights)
def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
for _ in range(max_new_tokens):
idx_cond = idx if idx.shape[1] <= self.config.block_size else idx[:, -self.config.block_size:]
logits, _ = self(idx_cond)
logits = logits[:, -1, :] / temperature
idx_next = logits.softmax().multinomial()
idx = Tensor.cat(idx, idx_next, dim=1)
return idx
def __call__(self, idx:Tensor, targets=None):
b, t = idx.shape
pos = Tensor.arange(0, t)
tok_emb = self.wte(idx) # token embeddings of shape (b, t, n_embd)
pos_emb = self.wpe(pos) # position embeddings of shape (t, n_embd)
x = tok_emb + pos_emb
x = self.ln_f(x.sequential(self.h))
if targets is not None:
logits = self.lm_head(x)[:, :, :self.config.vocab_size]
loss = logits.sparse_categorical_crossentropy(targets)
else:
logits = self.lm_head(x[:, [-1], :])[:, :, :self.config.vocab_size]
loss = None
return logits, loss
if __name__ == "__main__":
import tiktoken, argparse
parser = argparse.ArgumentParser()
parser.add_argument("--num_iterations", type=int, default=10, help="number of iterations to run")
parser.add_argument("--batch_size", type=int, default=4, help="batch size")
parser.add_argument("--sequence_length", type=int, default=64, help="sequence length")
parser.add_argument("--skip_test", action="store_true", help="skip test")
parser.add_argument("--gpus", type=int, default=1, help="sequence length")
args = parser.parse_args()
B, T = args.batch_size, args.sequence_length
assert 1 <= T <= 1024
model = GPT(GPTConfig(n_layer=12, n_head=12, n_embd=768))
model.load_pretrained()
if args.gpus > 1:
GPUS = tuple(f'{Device.DEFAULT}:{i}' for i in range(args.gpus))
for x in nn.state.get_parameters(model): x.to_(GPUS) # we put a copy of the model on every GPU
# init the tokenizer
enc = tiktoken.get_encoding("gpt2")
encode = lambda s: enc.encode(s, allowed_special={"<|endoftext|>"})
decode = lambda l: enc.decode(l)
# load the tokens
# prefer to use tiny_shakespeare if it's available, otherwise use tiny_stories
# we're using val instead of train split just because it is smaller/faster
tokens_bin = fetch("https://huggingface.co/datasets/karpathy/llmc-starter-pack/resolve/main/tiny_shakespeare_val.bin")
assert os.path.isfile(tokens_bin)
print(f"loading cached tokens in {tokens_bin}")
with open(tokens_bin, "rb") as f:
f.seek(0x400)
tokens = np.frombuffer(f.read(), dtype=np.uint16).astype(np.int32)
tokens = Tensor(tokens)
# lightweight dataloader
def get_batch():
assert B*T+1 <= len(tokens), "not enough tokens"
# for 338,025 tokens. E.g. with B=8 T=1024, this will yield 41 batches before looping
i = 0
while True:
x = tokens[i:i+B*T].view(B, T)
y = tokens[i+1:i+B*T+1].view(B, T)
yield x, y
i += B*T
if i + B*T + 1 >= len(tokens):
i = 0 # in prod we'd want to randomize the start point a bit
# forward backward for a few iterations
data_iter = iter(get_batch())
x, y = next(data_iter) # we'll overfit this batch below
optimizer = nn.optim.AdamW(nn.state.get_parameters(model), lr=1e-4, weight_decay=0)
print(f"model state: {sum(x.nbytes() for x in nn.state.get_parameters(model))/1e9:.2f} GB")
print(f"optimizer state: {sum(x.nbytes() for x in nn.state.get_parameters(optimizer))/1e9:.2f} GB")
# shard the data on axis 0
if args.gpus > 1: x, y = x.shard(GPUS, axis=0), y.shard(GPUS, axis=0)
@TinyJit
@Context(TRAINING=1)
def step(x:Tensor, y:Tensor) -> Tensor:
_, loss = model(x, y)
optimizer.zero_grad()
loss.backward()
return loss.realize(*optimizer.schedule_step())
for i in range(args.num_iterations):
GlobalCounters.reset()
t0 = time.perf_counter()
loss = step(x.contiguous(), y.contiguous())
Device[Device.DEFAULT].synchronize()
t1 = time.perf_counter()
print(f"iteration {i}, loss: {loss.item():.6f}, time: {(t1-t0)*1000:.3f}ms, {int(B*T/(t1-t0))} tok/s, {GlobalCounters.global_mem/1e9:.2f} GB")
if not args.skip_test:
# copy back to single gpu for test
if args.gpus > 1:
for x in nn.state.get_parameters(model): x.to_(Device.DEFAULT)
start = "<|endoftext|>"
start_ids = encode(start)
x = (Tensor(start_ids)[None, ...])
max_new_tokens = 16
temperature = 1.0
top_k = 40
y = model.generate(x, max_new_tokens, temperature=temperature, top_k=top_k)
print(decode(y[0].tolist()))

65
artifacts/package_sources/tinygrad/examples/llm.c/ubench/matmul.c Normal file
View File

@@ -0,0 +1,65 @@
// clang -Ofast -Wno-unused-result -march=native matmul.c
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
float b52[786432];
float b49[196608];
float h_0_mlp_c_fc_weight[2359296];
float h_0_mlp_c_fc_bias[3072];
void matmul_forward(float* out,
float* inp, float* weight, float* bias,
int B, int T, int C, int OC) {
// most of the running time is spent here and in matmul_backward
// OC is short for "output channels"
// inp is (B,T,C), weight is (OC, C), bias is (OC)
// out will be (B,T,OC)
#pragma omp parallel for collapse(2)
for (int b = 0; b < B; b++) {
for (int t = 0; t < T; t++) {
float* out_bt = out + b * T * OC + t * OC;
float* inp_bt = inp + b * T * C + t * C;
for (int o = 0; o < OC; o++) {
float val = (bias != NULL) ? bias[o] : 0.0f;
float* wrow = weight + o*C;
for (int i = 0; i < C; i++) {
val += inp_bt[i] * wrow[i];
}
out_bt[o] = val;
}
}
}
}
void r_256_3072_768(float* restrict data0, const float* restrict data1, const float* restrict data2, const float* restrict data3) {
for (int ridx0 = 0; ridx0 < 256; ridx0++) {
for (int ridx1 = 0; ridx1 < 3072; ridx1++) {
float acc0 = 0.0f;
float val0 = data3[ridx1];
for (int ridx2 = 0; ridx2 < 768; ridx2++) {
float val1 = data1[(ridx0*768)+ridx2];
float val2 = data2[(ridx1*768)+ridx2];
acc0 = ((val1*val2)+acc0);
}
data0[(ridx0*3072)+ridx1] = (acc0+val0);
}
}
}
int main() {
for (int i = 0; i < 5; i++) {
struct timespec t1, t2, t3;
clock_gettime(CLOCK_MONOTONIC, &t1);
r_256_3072_768(b52, b49, h_0_mlp_c_fc_weight, h_0_mlp_c_fc_bias);
clock_gettime(CLOCK_MONOTONIC, &t2);
matmul_forward(b52, b49, h_0_mlp_c_fc_weight, h_0_mlp_c_fc_bias, 4, 64, 768, 3072);
clock_gettime(CLOCK_MONOTONIC, &t3);
double time_gen = (t2.tv_sec - t1.tv_sec) + (t2.tv_nsec - t1.tv_nsec) / 1e9;
double time_real = (t3.tv_sec - t2.tv_sec) + (t3.tv_nsec - t2.tv_nsec) / 1e9;
printf("%.2f ms gen vs %.2f ms reference\n", time_gen*1e3, time_real*1e3);
}
}

324
artifacts/package_sources/tinygrad/examples/mamba.py Normal file
View File

@@ -0,0 +1,324 @@
import os, sys, math, argparse, time
sys.path.append(os.getcwd())
from typing import Any, Optional, Dict
from tinygrad import Tensor, TinyJit, nn
from tinygrad.helpers import fetch
from tinygrad.nn.state import load_state_dict, torch_load
from tqdm import tqdm
from transformers import AutoTokenizer
MODELS = {
"130m": {"dim": 768, "n_layers": 24, "vocab_size": 50277, "pad_vocab_size_multiple": 8},
"370m": {"dim": 1024, "n_layers": 48, "vocab_size": 50277, "pad_vocab_size_multiple": 8},
"790m": {"dim": 1536, "n_layers": 48, "vocab_size": 50277, "pad_vocab_size_multiple": 8},
"1.4b": {"dim": 2048, "n_layers": 48, "vocab_size": 50277, "pad_vocab_size_multiple": 8},
"2.8b": {"dim": 2560, "n_layers": 64, "vocab_size": 50277, "pad_vocab_size_multiple": 8},
}
def fetch_weights(model_name: str) -> Dict[str, Tensor]:
if model_name not in MODELS:
raise ValueError(f"Requested unknown mamba model: {model_name}")
downloaded = fetch(f"https://huggingface.co/state-spaces/mamba-{model_name}/resolve/main/pytorch_model.bin?download=true")
return torch_load(downloaded)
def selective_scan_ref(
u,
delta,
A,
B,
C,
D=None,
z=None,
delta_bias=None,
delta_softplus=False,
return_last_state=False,
):
"""
u: r(B D L)
delta: r(B D L)
A: c(D N) or r(D N)
B: c(D N) or r(B N L) or r(B N 2L) or r(B G N L) or (B G N L)
C: c(D N) or r(B N L) or r(B N 2L) or r(B G N L) or (B G N L)
D: r(D)
z: r(B D L)
delta_bias: r(D), fp32
out: r(B D L)
last_state (optional): r(B D dstate) or c(B D dstate)
"""
u = u.float()
delta = delta.float()
if delta_bias is not None:
delta = delta + delta_bias[..., None].float()
if delta_softplus:
delta = delta.softplus()
batch, dim, dstate = u.shape[0], A.shape[0], A.shape[1]
is_variable_B = len(B.shape) >= 3
is_variable_C = len(C.shape) >= 3
x = Tensor.zeros(batch, dim, dstate)
ys = []
deltaA = Tensor.einsum("bdl,dn->bdln", delta, A).exp()
if not is_variable_B:
deltaB_u = Tensor.einsum("bdl,dn,bdl->bdln", delta, B, u)
else:
if len(B.shape) == 3:
deltaB_u = Tensor.einsum("bdl,bnl,bdl->bdln", delta, B, u)
else:
B = B.repeat((1, dim // B.shape[1], 1, 1))
deltaB_u = Tensor.einsum("bdl,bdnl,bdl->bdln", delta, B, u)
if is_variable_C and len(C.shape) == 4:
C = C.repeat((1, dim // C.shape[1], 1, 1))
last_state = None
for i in range(u.shape[2]):
x = deltaA[:, :, i] * x + deltaB_u[:, :, i]
if not is_variable_C:
y = Tensor.einsum("bdn,dn->bd", x, C)
else:
if len(C.shape) == 3:
y = Tensor.einsum("bdn,bn->bd", x, C[:, :, i])
else:
y = Tensor.einsum("bdn,bdn->bd", x, C[:, :, :, i])
if i == u.shape[2] - 1:
last_state = x
ys.append(y)
y = Tensor.stack(*ys, dim=2) # (batch dim L)
out = y if D is None else y + u * D.reshape((-1, 1))
if z is not None:
out = out * z.silu()
return out if not return_last_state else (out, last_state)
class MambaMixer:
def __init__(
self,
dim,
d_state=16,
d_conv=4,
expand=2,
dt_rank="auto",
dt_min=0.001,
dt_max=0.1,
dt_init="random",
dt_scale=1.0,
dt_init_floor=1e-4,
conv_bias=True,
bias=False,
layer_idx=None,
):
self.dim = dim
self.d_state = d_state
self.d_conv = d_conv
self.expand = expand
self.d_inner = self.expand * self.dim
self.dt_rank = math.ceil(self.dim / 16) if dt_rank == "auto" else dt_rank
self.layer_idx = layer_idx
self.in_proj = nn.Linear(self.dim, self.d_inner * 2, bias=bias)
self.conv1d = nn.Conv1d(in_channels=self.d_inner, out_channels=self.d_inner, bias=conv_bias,
kernel_size=d_conv, groups=self.d_inner, padding=d_conv-1)
self.x_proj = nn.Linear(self.d_inner, self.dt_rank + self.d_state * 2, bias=False)
self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True)
# Initialize special dt projection to preserve variance at initialization
dt_init_std = self.dt_rank**-0.5 * dt_scale
if dt_init == "constant":
self.dt_proj.weight = Tensor.full(self.dt_proj.weight.shape, dt_init_std)
elif dt_init == "random":
self.dt_proj.weight = Tensor.uniform(self.dt_proj.weight.shape, low=-dt_init_std, high=dt_init_std)
else:
raise NotImplementedError
dt = Tensor.uniform(self.d_inner, low=math.log(dt_min), high=math.log(dt_max)).exp().maximum(dt_init_floor)
inv_dt = dt + (1 - (-dt).exp()).log()
self.dt_proj.bias.assign(inv_dt)
# S4D real initialization
self.A_log = Tensor.arange(1, self.d_state+1).repeat([self.d_inner, 1]).log()
# D "skip" parameter
self.D = Tensor.ones(self.d_inner) # Keep in fp32
self.out_proj = nn.Linear(self.d_inner, self.dim, bias=bias)
def __call__(self, hidden_states: Tensor):
batch, seqlen, _ = hidden_states.shape
if not hasattr(self, 'conv_state'):
self.conv_state = Tensor.zeros(batch, self.dim * self.expand, self.d_conv).contiguous().realize()
self.ssm_state = Tensor.zeros(batch, self.dim * self.expand, self.d_state).realize()
xz = self.in_proj.weight @ hidden_states.permute(2,0,1).reshape(hidden_states.shape[2],hidden_states.shape[1]*hidden_states.shape[0])
xz = xz.reshape(xz.shape[0],xz.shape[1]//seqlen, seqlen).permute(1,0,2)
if self.in_proj.bias is not None:
xz = xz + self.in_proj.bias.reshape((-1, 1))
A = -self.A_log.exp()
x, z = xz.chunk(2, dim=1)
# Compute short convolution
self.conv_state.assign(x[:, :, -self.d_conv :]) # Update state (B D W)
x = self.conv1d(x)[..., :seqlen].swish()
x_dbl = self.x_proj(x.permute(0,2,1).reshape(x.shape[0]*x.shape[2], x.shape[1]))
dt, B, C = Tensor.split(x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=-1)
dt = self.dt_proj.weight @ dt.T
dt = dt.reshape(dt.shape[0], dt.shape[1]//seqlen, seqlen).permute(1,0,2)
B = B.reshape(B.shape[0]//seqlen, seqlen, B.shape[1]).permute(0,2,1)
C = C.reshape(C.shape[0]//seqlen, seqlen, C.shape[1]).permute(0,2,1)
# TODO: actually implement selective_scan_fn
y = selective_scan_ref(x, dt, A, B, C, self.D, z=z, delta_bias=self.dt_proj.bias, delta_softplus=True,
return_last_state=True)
y, last_state = y
self.ssm_state.assign(last_state).realize()
y = y.permute(0,2,1)
out = self.out_proj(y)
return out
else:
return self.step(hidden_states)
def step(self, hidden_states: Tensor):
assert hidden_states.shape[1] == 1, f"Only support decoding with 1 token at a time for now, attempted {hidden_states.shape[1]}"
xz = self.in_proj(hidden_states.squeeze(1)) # (B 2D)
x, z = xz.chunk(2, dim=-1) # (B D)
# Conv step
self.conv_state.assign(self.conv_state[:, :, 1:].cat(x.unsqueeze(-1), dim=-1).realize())
x = (self.conv_state * self.conv1d.weight.squeeze(1)).sum(-1)
if self.conv1d.bias is not None:
x = x + self.conv1d.bias
x = x.swish()
x_db = self.x_proj(x) # (B dt_rank+2*d_state)
dt = x_db[:, : self.dt_rank]
B = x_db[:, self.dt_rank : (self.dt_rank + self.d_state)]
C = x_db[:, (self.dt_rank + self.d_state) :]
# Don't add dt_bias here
dt = self.dt_proj.weight @ dt.T
A = -self.A_log.exp()
# SSM step
dt = (dt + self.dt_proj.bias.unsqueeze(-1)).softplus()
dA = Tensor.einsum("db,dn->bdn", dt, A).exp()
dB = Tensor.einsum("db,bn->bdn", dt, B)
self.ssm_state.assign(self.ssm_state * dA + x.unsqueeze(-1) * dB)
y = Tensor.einsum("bdn,bn->bd", self.ssm_state, C)
y = y + self.D * x
y = y * z.swish() # (B D)
out = self.out_proj(y)
return out.unsqueeze(1)
class MambaBlock:
def __init__(self, dim: int, norm_eps: float = 1e-5, rms_norm: bool = True, layer_idx: Optional[int] = None):
self.mixer = MambaMixer(dim, layer_idx=layer_idx)
if rms_norm:
self.norm = nn.RMSNorm(dim, norm_eps)
else:
raise NotImplementedError
def __call__(self, hidden_states: Tensor, residual: Optional[Tensor] = None):
residual = (hidden_states + residual) if residual is not None else hidden_states
hidden_states = self.norm(residual)
hidden_states = self.mixer(hidden_states)
return hidden_states, residual
class MambaBackbone:
def __init__(self, dim: int, n_layers: int, vocab_size: int, rms_norm: bool = True, norm_eps: float = 1e-5):
self.embedding = nn.Embedding(vocab_size, dim)
self.layers = [MambaBlock(dim, rms_norm=rms_norm, layer_idx=i) for i in range(n_layers)]
if rms_norm:
self.norm_f = nn.RMSNorm(dim, norm_eps)
def __call__(self, input_ids: Tensor) -> Any:
hidden_states = self.embedding(input_ids)
residual = None
for layer in self.layers:
hidden_states, residual = layer(hidden_states, residual)
residual = (hidden_states + residual) if residual is not None else hidden_states
hidden_states = self.norm_f(residual)
return hidden_states
class Mamba:
def __init__(self, dim: int, n_layers: int, vocab_size: int, pad_vocab_size_multiple: int = 1):
if vocab_size % pad_vocab_size_multiple != 0:
vocab_size += pad_vocab_size_multiple - (vocab_size % pad_vocab_size_multiple)
self.backbone = MambaBackbone(dim, n_layers, vocab_size)
self.lm_head = nn.Linear(dim, vocab_size, bias=False)
self.forward_jit = TinyJit(self.forward)
def forward(self, input_ids:Tensor):
hidden_states = self.backbone(input_ids)
return self.lm_head(hidden_states).realize()
def __call__(self, input_ids):
return self.forward(input_ids)
@staticmethod
def from_pretrained(model_name: str):
weights = fetch_weights(model_name)
model = Mamba(**MODELS[model_name])
load_state_dict(model, weights)
return model
def generate(model, tokenizer, prompt: str, n_tokens_to_gen: int = 10, temp: bool = 1.0, sample: bool = False, top_k: int = None):
tks = tokenizer(prompt)["input_ids"]
while len(tks) < 4:
tks = [50279] + tks
# Loading in the prompt tokens
logits = model.forward(Tensor([tks]))[:, -1, :]
for _ in tqdm(range(n_tokens_to_gen), desc="Speed Gen"):
if sample:
scaled_logits = logits / temp
if top_k is not None:
topk_values, topk_indices = scaled_logits.topk(top_k)
filtered_logits = Tensor.full_like(scaled_logits, -float("inf"))
filtered_logits = filtered_logits.scatter(dim=-1, index=topk_indices, src=topk_values)
tok_Tens = filtered_logits.softmax().multinomial()
else:
tok_Tens = scaled_logits.softmax().multinomial()
else:
tok_Tens = logits.argmax(axis=-1).unsqueeze(0)
tok = tok_Tens.item()
tks.append(tok)
logits = model.forward_jit(tok_Tens)[:, -1, :]
output_completions = ''.join([tokenizer.decode(output) for output in tks])
return output_completions
if __name__ == "__main__":
ORIG_PROMPT = "Why is gravity "
parser = argparse.ArgumentParser(description="Run Mamba in tinygrad", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("--prompt", type=str, default="Why is gravity ", help="Prompt for LLM completion")
parser.add_argument("--size", type=str, default="370m",
help=f"Size of model to use [{', '.join([k for k in MODELS.keys()])}]")
parser.add_argument("--n_tokens", type=int, default=10, help="Number of tokens to generate")
parser.add_argument("--top_k", type=int, help="Limit sampling to the top k most likely tokens")
parser.add_argument("--sample", dest="sample", action="store_true", help="Sample flag")
parser.add_argument("--temp", type=float, default=1.0, help="Sampling temp has to be <=1.0")
args = parser.parse_args()
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
model = Mamba.from_pretrained(args.size)
prompt = args.prompt
num_toks = args.n_tokens
sample = args.sample
temp = args.temp
top_k = args.top_k
s = time.time()
tinyoutput = generate(model, tokenizer, prompt, n_tokens_to_gen=num_toks, sample=sample, temp=temp, top_k=top_k)
print(tinyoutput)
print('TIME: ', time.time() - s)
TORCHOUTPUT = "Why is gravity \nso important?\nBecause it's the only"
if ORIG_PROMPT == prompt and not sample and num_toks==10 and args.size=='370m': print('Outputs Match:', tinyoutput == TORCHOUTPUT)

156
artifacts/package_sources/tinygrad/examples/minrf.py Normal file
View File

@@ -0,0 +1,156 @@
# much taken from https://github.com/cloneofsimo/minRF
from tinygrad import Tensor, nn, GlobalCounters, TinyJit, Context
from tinygrad.helpers import getenv, trange
from extra.models.llama import Attention, FeedForward, precompute_freqs_cis
def modulate(x:Tensor, shift:Tensor, scale:Tensor) -> Tensor: return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
# TODO: why doesn't the TimestepEmbedder from minRF work?
class TimestepEmbedder:
def __init__(self, hidden_size): self.mlp = [nn.Linear(1, hidden_size), Tensor.silu, nn.Linear(hidden_size, hidden_size)]
def __call__(self, t:Tensor): return t.reshape(-1, 1).sequential(self.mlp)
class TransformerBlock:
def __init__(self, dim, n_heads, norm_eps=1e-5):
self.attention = Attention(dim, n_heads)
self.feed_forward = FeedForward(dim, 4*dim)
self.attention_norm = nn.LayerNorm(dim, eps=norm_eps)
self.ffn_norm = nn.LayerNorm(dim, eps=norm_eps)
self.adaLN_modulation = nn.Linear(dim, 6 * dim, bias=True)
def __call__(self, x:Tensor, freqs_cis:Tensor, adaln_input:Tensor):
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(adaln_input.silu()).chunk(6, dim=1)
x = x + gate_msa.unsqueeze(1) * self.attention(modulate(self.attention_norm(x), shift_msa, scale_msa), 0, freqs_cis)
x = x + gate_mlp.unsqueeze(1) * self.feed_forward(modulate(self.ffn_norm(x), shift_mlp, scale_mlp))
return x.contiguous().contiguous_backward()
class FinalLayer:
def __init__(self, dim, patch_size, out_channels):
self.norm_final = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
self.linear = nn.Linear(dim, patch_size*patch_size*out_channels, bias=True)
self.adaLN_modulation = nn.Linear(dim, 2 * dim, bias=True)
# init weights/bias to 0
self.linear.weight.replace(self.linear.weight.zeros_like().contiguous())
self.linear.bias.replace(self.linear.bias.zeros_like().contiguous())
def __call__(self, x:Tensor, c:Tensor):
shift, scale = self.adaLN_modulation(c.silu()).chunk(2, dim=1)
x = modulate(self.norm_final(x), shift, scale)
return self.linear(x)
# channels=1, input_size=32, dim=64, n_layers=6, n_heads=4, num_classes=10
class DiT_Llama:
def __init__(self, in_channels=1, dim=64, n_layers=6, n_heads=4, num_classes=10, patch_size=2):
self.patch_size = patch_size
self.out_channels = in_channels
self.num_classes = num_classes
self.init_conv_seq = [
nn.Conv2d(in_channels, dim // 2, kernel_size=5, padding=2, stride=1), Tensor.silu, nn.GroupNorm(32, dim//2),
nn.Conv2d(dim //2, dim // 2, kernel_size=5, padding=2, stride=1), Tensor.silu, nn.GroupNorm(32, dim//2),
]
self.x_embedder = nn.Linear(self.patch_size * self.patch_size * dim // 2, dim, bias=True)
self.t_embedder = TimestepEmbedder(dim)
self.y_embedder = nn.Embedding(num_classes+1, dim)
self.final_layer = FinalLayer(dim, self.patch_size, self.out_channels)
self.freqs_cis = precompute_freqs_cis(dim // n_heads, 4096)
self.layers = [TransformerBlock(dim, n_heads) for _ in range(n_layers)]
def unpatchify(self, x:Tensor):
c, p = self.out_channels, self.patch_size
h = w = int(x.shape[1] ** 0.5)
x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
x = x.rearrange("n h w p q c -> n c h p w q")
return x.reshape(shape=(x.shape[0], c, h * p, h * p))
def patchify(self, x:Tensor):
B, C, H, W = x.shape
x = x.reshape(B, C, H // self.patch_size, self.patch_size, W // self.patch_size, self.patch_size)
x = x.permute(0, 2, 4, 1, 3, 5).flatten(-3).flatten(1, 2)
return x # B <H*W ish> <C*patch_size*patch_size>
def __call__(self, x:Tensor, t:Tensor, y:Tensor) -> Tensor:
x = x.sequential(self.init_conv_seq)
x = self.patchify(x)
x = self.x_embedder(x)
adaln_input = self.t_embedder(t) + self.y_embedder(y)
adaln_input = adaln_input.contiguous()
for layer in self.layers:
x = layer(x, self.freqs_cis[:, :x.size(1)], adaln_input=adaln_input)
x = self.final_layer(x, adaln_input)
return self.unpatchify(x)
def rf(self, x:Tensor, cond:Tensor):
b = x.shape[0]
# self.ln is True
t = Tensor.randn((b,)).sigmoid()
texp = t.view([b, *([1] * len(x.shape[1:]))])
# conditional dropout
dropout_prob = 0.1
cond = (Tensor.rand(cond.shape[0]) < dropout_prob).where(cond.full_like(self.num_classes), cond)
# this is rectified flow
z1 = x.randn_like()
zt = (1 - texp) * x + texp * z1
vtheta = self(zt, t, cond)
# MSE loss
return ((z1 - x) - vtheta).square().mean()
def sample(self, z, cond, null_cond, sample_steps=50, cfg=2.0):
b = z.size(0)
dt = Tensor.full((b,)+(1,)*len(z.shape[1:]), fill_value=1.0/sample_steps).contiguous()
images = [z]
for i in range(sample_steps, 0, -1):
t = Tensor.full((b,), fill_value=i/sample_steps).contiguous()
vc = self(z, t, cond)
vu = self(z, t, null_cond)
vc = vu + cfg * (vc - vu)
z = z - dt * vc
z = z.contiguous()
images.append(z)
return images
def mviz(t:Tensor):
assert len(t.shape) == 4 and t.shape[1] == 1
ft = t.permute(1,2,0,3).reshape(32, -1)
assert ft.shape[-1]%32 == 0
print("")
for y in ((ft+1)/2).clamp(0,1).tolist():
ln = [f"\033[38;5;{232+int(x*23)}m██" for x in y]
print(''.join(ln) + "\033[0m")
if __name__ == "__main__":
X_train, Y_train, X_test, Y_test = nn.datasets.mnist()
X_train = X_train.pad((2,2,2,2))
X_train = ((X_train.float()/255)-0.5)/0.5
Y_train = Y_train.int()
model = DiT_Llama(patch_size=getenv("PATCH_SIZE", 2))
for r in nn.state.get_parameters(model): r.realize()
optimizer = nn.optim.Adam(nn.state.get_parameters(model), lr=5e-4)
@TinyJit
@Context(TRAINING=1)
def train_step():
if getenv("OVERFIT"): samples = Tensor.zeros(getenv("BS", 256), dtype='int')
else: samples = Tensor.randint(getenv("BS", 256), high=X_train.shape[0])
optimizer.zero_grad()
loss = model.rf(X_train[samples], Y_train[samples])
loss.backward()
optimizer.step()
return loss
@TinyJit
def sample(z:Tensor, cond:Tensor) -> Tensor:
return model.sample(z, cond, Tensor.full_like(cond, 10), sample_steps=getenv("SAMPLE_STEPS", 20))[-1]
for steps in (t:=trange(getenv("STEPS", 5000))):
if steps%10 == 0: mviz(sample(Tensor.randn(3, 1, 32, 32), Tensor([5,0,4], dtype='int')))
GlobalCounters.reset()
loss = train_step()
t.set_description(f"loss: {loss.item():9.2f}")

62
artifacts/package_sources/tinygrad/examples/mixtral.py Normal file
View File

@@ -0,0 +1,62 @@
import functools, argparse, pathlib
from tinygrad import Tensor, nn, Device, GlobalCounters, Variable
from tinygrad.helpers import Timing, Profiling, tqdm
from tinygrad.nn.state import torch_load, get_state_dict
from extra.models.llama import FeedForward, Transformer
from extra.bench_log import BenchEvent, WallTimeEvent
class MixtureFeedForward:
def __init__(self, num_experts:int, dim:int, hidden_dim:int, linear=nn.Linear):
self.gate = nn.Linear(dim, num_experts, bias=False)
self.experts = [FeedForward(dim, hidden_dim, linear) for _ in range(num_experts)]
def __call__(self, x:Tensor) -> Tensor:
assert x.shape[0] == 1, "only BS=1"
g = self.gate(x).float().exp()
choice = g.data().tolist()[0][0]
top = sorted(enumerate(choice), key=lambda x: -x[1])
norm = top[0][1] + top[1][1]
e1, e2 = self.experts[top[0][0]], self.experts[top[1][0]]
scale = Tensor([top[0][1]/norm, top[1][1]/norm])
ret = e1(x.to(e1.w1.weight.device)).to(x.device) * scale[0] + \
e2(x.to(e2.w1.weight.device)).to(x.device) * scale[1]
return ret
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Run Mixtral in tinygrad", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("--count", type=int, default=30, help="Max number of tokens to generate")
parser.add_argument("--temperature", type=float, default=0.7, help="Temperature in the softmax")
parser.add_argument("--timing", action="store_true", help="Print timing per token")
parser.add_argument("--profile", action="store_true", help="Profile generation")
parser.add_argument("--weights", type=str, default=(pathlib.Path(__file__).parent.parent / "weights/mixtral-8x7b-32kseqlen").as_posix(),
help="Path to the downloaded weights")
args = parser.parse_args()
with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
state = torch_load(args.weights + "/consolidated.00.pth.b")
model = Transformer(n_layers=32, dim=4096, hidden_dim=14336, n_heads=32, n_kv_heads=8, norm_eps=1e-5, vocab_size=32000, feed_forward=functools.partial(MixtureFeedForward, 8), jit=False)
model_state_dict = get_state_dict(model)
for k in (t := tqdm(state, disable=None)):
if 'feed_forward.experts.' in k:
expert_no = int(k.split('feed_forward.experts.')[1].split('.')[0])
device = Device.DEFAULT + ":" + str((expert_no//2)+1)
else:
device = Device.DEFAULT
t.set_description(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB, loading {k} to {device}")
model_state_dict[k].replace(state[k].to(device).half()).realize()
if t.disable: print(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB")
from sentencepiece import SentencePieceProcessor
spp = SentencePieceProcessor(model_file=args.weights + "/tokenizer.model")
toks = [spp.bos_id()]
start_pos = 0
for i in range(args.count):
GlobalCounters.reset()
with Profiling(sort="time", frac=0.1, enabled=args.profile):
with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/sec"):
with WallTimeEvent(BenchEvent.STEP):
tok = model(Tensor([toks[start_pos:]]), 0 if start_pos == 0 else Variable("start_pos", 1, 1024-1).bind(start_pos), args.temperature).item()
toks.append(tok)
start_pos += 1
print(spp.decode(toks))

19
artifacts/package_sources/tinygrad/examples/mlperf/README Normal file
View File

@@ -0,0 +1,19 @@
Each model should be a clean single file.
They are imported from the top level `models` directory
It should be capable of loading weights from the reference imp.
We will focus on these 5 models:
# Resnet50-v1.5 (classic) -- 8.2 GOPS/input
# Retinanet
# 3D UNET (upconvs)
# RNNT
# BERT-large (transformer)
They are used in both the training and inference benchmark:
https://mlcommons.org/en/training-normal-21/
https://mlcommons.org/en/inference-edge-30/
And we will submit to both.
NOTE: we are Edge since we don't have ECC RAM

822
artifacts/package_sources/tinygrad/examples/mlperf/dataloader.py Normal file
View File

@@ -0,0 +1,822 @@
import os, random, pickle, queue, struct, math, functools, hashlib, time
from typing import List
from pathlib import Path
from multiprocessing import Queue, Process, shared_memory, connection, Lock
import numpy as np
from tinygrad import dtypes, Tensor
from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX, NUM_CPU_THREADS
from tinygrad.nn.state import TensorIO
### ResNet
class MyQueue:
def __init__(self, multiple_readers=True, multiple_writers=True):
self._reader, self._writer = connection.Pipe(duplex=False)
self._rlock = Lock() if multiple_readers else None
self._wlock = Lock() if multiple_writers else None
def get(self):
if self._rlock: self._rlock.acquire()
ret = pickle.loads(self._reader.recv_bytes())
if self._rlock: self._rlock.release()
return ret
def put(self, obj):
if self._wlock: self._wlock.acquire()
self._writer.send_bytes(pickle.dumps(obj))
if self._wlock: self._wlock.release()
def shuffled_indices(n, seed=None):
rng = random.Random(seed)
indices = {}
for i in range(n-1, -1, -1):
j = rng.randint(0, i)
if i not in indices: indices[i] = i
if j not in indices: indices[j] = j
indices[i], indices[j] = indices[j], indices[i]
yield indices[i]
del indices[i]
def loader_process(q_in, q_out, X:Tensor, seed):
import signal
signal.signal(signal.SIGINT, lambda _, __: exit(0))
from extra.datasets.imagenet import center_crop, preprocess_train
from PIL import Image
with Context(DEBUG=0):
while (_recv := q_in.get()) is not None:
idx, fn, val = _recv
if fn is not None:
img = Image.open(fn)
img = img.convert('RGB') if img.mode != "RGB" else img
if val:
# eval: 76.08%, load in 0m7.366s (0m5.301s with simd)
# sudo apt-get install libjpeg-dev
# CC="cc -mavx2" pip install -U --force-reinstall pillow-simd
img = center_crop(img)
img = np.array(img)
else:
# reseed rng for determinism
if seed is not None:
np.random.seed(seed * 2 ** 10 + idx)
random.seed(seed * 2 ** 10 + idx)
img = preprocess_train(img)
else:
# pad data with training mean
img = np.tile(np.array([[[123.68, 116.78, 103.94]]], dtype=np.uint8), (224, 224, 1))
X[idx].flatten().assign(img.tobytes())
q_out.put(idx)
q_out.put(None)
def batch_load_resnet(batch_size=64, val=False, shuffle=True, seed=None, pad_first_batch=False):
from extra.datasets.imagenet import get_train_files, get_val_files
files = get_val_files() if val else get_train_files()
from extra.datasets.imagenet import get_imagenet_categories
cir = get_imagenet_categories()
if pad_first_batch:
FIRST_BATCH_PAD = round_up(len(files), batch_size) - len(files)
else:
FIRST_BATCH_PAD = 0
file_count = FIRST_BATCH_PAD + len(files)
BATCH_COUNT = min(32, file_count // batch_size)
def _gen():
for _ in range(FIRST_BATCH_PAD): yield -1
yield from shuffled_indices(len(files), seed=seed) if shuffle else iter(range(len(files)))
gen = iter(_gen())
def enqueue_batch(num):
for idx in range(num*batch_size, (num+1)*batch_size):
fidx = next(gen)
if fidx != -1:
fn = files[fidx]
q_in.put((idx, fn, val))
Y[idx] = cir[fn.split("/")[-2]]
else:
# padding
q_in.put((idx, None, val))
Y[idx] = -1
shutdown = False
class Cookie:
def __init__(self, num): self.num = num
def __del__(self):
if not shutdown:
try: enqueue_batch(self.num)
except StopIteration: pass
gotten = [0]*BATCH_COUNT
def receive_batch():
while 1:
num = q_out.get()//batch_size
gotten[num] += 1
if gotten[num] == batch_size: break
gotten[num] = 0
return X[num*batch_size:(num+1)*batch_size], Y[num*batch_size:(num+1)*batch_size], Cookie(num)
#q_in, q_out = MyQueue(multiple_writers=False), MyQueue(multiple_readers=False)
q_in, q_out = Queue(), Queue()
sz = (batch_size*BATCH_COUNT, 224, 224, 3)
shm_name = "resnet_X_val" if val else "resnet_X_train"
if not OSX and os.path.exists(f"/dev/shm/{shm_name}"): os.unlink(f"/dev/shm/{shm_name}")
shm = shared_memory.SharedMemory(name=shm_name, create=True, size=prod(sz))
procs = []
try:
# disk:shm is slower
if OSX: X = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:shm:{shm.name}")
else: X = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name}")
Y = [None] * (batch_size*BATCH_COUNT)
for _ in range(NUM_CPU_THREADS.value):
p = Process(target=loader_process, args=(q_in, q_out, X, seed))
p.daemon = True
p.start()
procs.append(p)
for bn in range(BATCH_COUNT): enqueue_batch(bn)
# NOTE: this is batch aligned, last ones are ignored unless pad_first_batch is True
for _ in range(0, file_count//batch_size): yield receive_batch()
finally:
shutdown = True
# empty queues
for _ in procs: q_in.put(None)
q_in.close()
for _ in procs:
while q_out.get() is not None: pass
q_out.close()
# shutdown processes
for p in procs: p.join()
shm.close()
try:
shm.unlink()
except FileNotFoundError:
# happens with BENCHMARK set
pass
### BERT
def process_batch_bert(data: List[dict]) -> dict[str, Tensor]:
return {
"input_ids": Tensor(np.concatenate([s["input_ids"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
"input_mask": Tensor(np.concatenate([s["input_mask"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
"segment_ids": Tensor(np.concatenate([s["segment_ids"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
"masked_lm_positions": Tensor(np.concatenate([s["masked_lm_positions"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
"masked_lm_ids": Tensor(np.concatenate([s["masked_lm_ids"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
"masked_lm_weights": Tensor(np.concatenate([s["masked_lm_weights"] for s in data], axis=0), dtype=dtypes.float32, device="CPU"),
"next_sentence_labels": Tensor(np.concatenate([s["next_sentence_labels"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
}
def load_file(file: str):
with open(file, "rb") as f:
return pickle.load(f)
class InterleavedDataset:
def __init__(self, files:List[str], cycle_length:int):
self.dataset = files
self.cycle_length = cycle_length
self.queues = [queue.Queue() for _ in range(self.cycle_length)]
for i in range(len(self.queues)): self.queues[i].queue.extend(load_file(self.dataset.pop(0)))
self.queue_pointer = len(self.queues) - 1
def get(self):
# Round-robin across queues
try:
self.advance()
return self.queues[self.queue_pointer].get_nowait()
except queue.Empty:
self.fill(self.queue_pointer)
return self.get()
def advance(self):
self.queue_pointer = (self.queue_pointer + 1) % self.cycle_length
def fill(self, queue_index: int):
try:
file = self.dataset.pop(0)
except IndexError:
return
self.queues[queue_index].queue.extend(load_file(file))
# Reference: https://github.com/mlcommons/training/blob/1c8a098ae3e70962a4f7422c0b0bd35ae639e357/language_model/tensorflow/bert/run_pretraining.py, Line 394
def batch_load_train_bert(BS:int, seed:int|None=None):
from extra.datasets.wikipedia import get_wiki_train_files
rng = random.Random(seed)
fs = sorted(get_wiki_train_files())
train_files = []
while fs: # TF shuffle
rng.shuffle(fs)
train_files.append(fs.pop(0))
cycle_length = min(NUM_CPU_THREADS.value, len(train_files))
assert cycle_length > 0, "cycle_length must be greater than 0"
dataset = InterleavedDataset(train_files, cycle_length)
while True:
yield process_batch_bert([dataset.get() for _ in range(BS)])
# Reference: https://github.com/mlcommons/training/blob/1c8a098ae3e70962a4f7422c0b0bd35ae639e357/language_model/tensorflow/bert/run_pretraining.py, Line 416
def batch_load_val_bert(BS:int):
file = getenv("BASEDIR", Path(__file__).parent.parents[1] / "extra" / "datasets" / "wiki") / "eval.pkl"
dataset = load_file(file)
idx = 0
while True:
start_idx = (idx * BS) % len(dataset)
end_idx = ((idx + 1) * BS) % len(dataset)
if start_idx < end_idx:
yield process_batch_bert(dataset[start_idx:end_idx])
else: # wrap around the end to the beginning of the dataset
yield process_batch_bert(dataset[start_idx:] + dataset[:end_idx])
idx += 1
### UNET3D
def load_unet3d_data(preprocessed_dataset_dir, seed, queue_in, queue_out, X:Tensor, Y:Tensor):
from extra.datasets.kits19 import rand_balanced_crop, rand_flip, random_brightness_augmentation, gaussian_noise
while (data := queue_in.get()) is not None:
idx, fn, val = data
case_name = os.path.basename(fn).split("_x.npy")[0]
x, y = np.load(preprocessed_dataset_dir / f"{case_name}_x.npy"), np.load(preprocessed_dataset_dir / f"{case_name}_y.npy")
if not val:
if seed is not None:
np.random.seed(seed)
random.seed(seed)
x, y = rand_balanced_crop(x, y)
x, y = rand_flip(x, y)
x, y = x.astype(np.float32), y.astype(np.uint8)
x = random_brightness_augmentation(x)
x = gaussian_noise(x)
X[idx].flatten().assign(x.tobytes())
Y[idx].flatten().assign(y.tobytes())
queue_out.put(idx)
queue_out.put(None)
def batch_load_unet3d(preprocessed_dataset_dir:Path, batch_size:int=6, val:bool=False, shuffle:bool=True, seed=None):
assert preprocessed_dataset_dir is not None, "run preprocess_data on kits19"
files = sorted(list(preprocessed_dataset_dir.glob("*_x.npy")))
file_indices = list(range(len(files)))
batch_count = min(32, len(files) // batch_size)
queue_in, queue_out = Queue(), Queue()
procs, data_out_count = [], [0] * batch_count
shm_name_x, shm_name_y = "unet3d_x", "unet3d_y"
sz = (batch_size * batch_count, 1, 128, 128, 128)
if os.path.exists(f"/dev/shm/{shm_name_x}"): os.unlink(f"/dev/shm/{shm_name_x}")
if os.path.exists(f"/dev/shm/{shm_name_y}"): os.unlink(f"/dev/shm/{shm_name_y}")
shm_x = shared_memory.SharedMemory(name=shm_name_x, create=True, size=prod(sz))
shm_y = shared_memory.SharedMemory(name=shm_name_y, create=True, size=prod(sz))
shutdown = False
class Cookie:
def __init__(self, bc):
self.bc = bc
def __del__(self):
if not shutdown:
try: enqueue_batch(self.bc)
except StopIteration: pass
def enqueue_batch(bc):
for idx in range(bc * batch_size, (bc+1) * batch_size):
fn = files[next(ds_iter)]
queue_in.put((idx, fn, val))
def shuffle_indices(file_indices, seed=None):
rng = random.Random(seed)
rng.shuffle(file_indices)
if shuffle: shuffle_indices(file_indices, seed=seed)
ds_iter = iter(file_indices)
try:
X = Tensor.empty(*sz, dtype=dtypes.float32, device=f"disk:/dev/shm/{shm_name_x}")
Y = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name_y}")
for _ in range(NUM_CPU_THREADS.value):
proc = Process(target=load_unet3d_data, args=(preprocessed_dataset_dir, seed, queue_in, queue_out, X, Y))
proc.daemon = True
proc.start()
procs.append(proc)
for bc in range(batch_count):
enqueue_batch(bc)
for _ in range(len(files) // batch_size):
while True:
bc = queue_out.get() // batch_size
data_out_count[bc] += 1
if data_out_count[bc] == batch_size: break
data_out_count[bc] = 0
yield X[bc * batch_size:(bc + 1) * batch_size], Y[bc * batch_size:(bc + 1) * batch_size], Cookie(bc)
finally:
shutdown = True
for _ in procs: queue_in.put(None)
queue_in.close()
for _ in procs:
while queue_out.get() is not None: pass
queue_out.close()
# shutdown processes
for proc in procs: proc.join()
shm_x.close()
shm_y.close()
try:
shm_x.unlink()
shm_y.unlink()
except FileNotFoundError:
# happens with BENCHMARK set
pass
### RetinaNet
def load_retinanet_data(base_dir:Path, val:bool, queue_in:Queue, queue_out:Queue,
imgs:Tensor, boxes:Tensor, labels:Tensor, matches:Tensor|None=None,
anchors:Tensor|None=None, seed:int|None=None):
from extra.datasets.openimages import image_load, random_horizontal_flip, resize
from examples.mlperf.helpers import box_iou, find_matches, generate_anchors
import torch
while (data:=queue_in.get()) is not None:
idx, img, tgt = data
img = image_load(base_dir, img["subset"], img["file_name"])
if val:
img = resize(img)[0]
else:
if seed is not None:
np.random.seed(seed)
random.seed(seed)
torch.manual_seed(seed)
img, tgt = random_horizontal_flip(img, tgt)
img, tgt, _ = resize(img, tgt=tgt)
match_quality_matrix = box_iou(tgt["boxes"], (anchor := np.concatenate(generate_anchors((800, 800)))))
match_idxs = find_matches(match_quality_matrix, allow_low_quality_matches=True)
clipped_match_idxs = np.clip(match_idxs, 0, None)
clipped_boxes, clipped_labels = tgt["boxes"][clipped_match_idxs], tgt["labels"][clipped_match_idxs]
boxes[idx].flatten().assign(clipped_boxes.tobytes())
labels[idx].flatten().assign(clipped_labels.tobytes())
matches[idx].flatten().assign(match_idxs.tobytes())
anchors[idx].flatten().assign(anchor.tobytes())
imgs[idx].flatten().assign(img.tobytes())
queue_out.put(idx)
queue_out.put(None)
def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, shuffle:bool=True, seed:int|None=None):
def _enqueue_batch(bc):
from extra.datasets.openimages import prepare_target
for idx in range(bc * batch_size, (bc+1) * batch_size):
img = dataset.loadImgs(next(dataset_iter))[0]
ann = dataset.loadAnns(dataset.getAnnIds(img_id:=img["id"]))
tgt = prepare_target(ann, img_id, (img["height"], img["width"]))
if img_ids is not None:
img_ids[idx] = img_id
if img_sizes is not None:
img_sizes[idx] = tgt["image_size"]
queue_in.put((idx, img, tgt))
def _setup_shared_mem(shm_name:str, size:tuple[int, ...], dtype:dtypes) -> tuple[shared_memory.SharedMemory, Tensor]:
shm_name = f"{shm_name}_{os.getpid()}"
if os.path.exists(f"/dev/shm/{shm_name}"): os.unlink(f"/dev/shm/{shm_name}")
shm = shared_memory.SharedMemory(name=shm_name, create=True, size=prod(size))
shm_tensor = Tensor.empty(*size, dtype=dtype, device=f"disk:/dev/shm/{shm_name}")
return shm, shm_tensor
image_ids = sorted(dataset.imgs.keys())
batch_count = min(32, len(image_ids) // batch_size)
queue_in, queue_out = Queue(), Queue()
procs, data_out_count = [], [0] * batch_count
shm_imgs, imgs = _setup_shared_mem("retinanet_imgs", (batch_size * batch_count, 800, 800, 3), dtypes.uint8)
if val:
boxes, labels, matches, anchors = None, None, None, None
img_ids, img_sizes = [None] * (batch_size * batch_count), [None] * (batch_size * batch_count)
else:
img_ids, img_sizes = None, None
shm_boxes, boxes = _setup_shared_mem("retinanet_boxes", (batch_size * batch_count, 120087, 4), dtypes.float32)
shm_labels, labels = _setup_shared_mem("retinanet_labels", (batch_size * batch_count, 120087), dtypes.int64)
shm_matches, matches = _setup_shared_mem("retinanet_matches", (batch_size * batch_count, 120087), dtypes.int64)
shm_anchors, anchors = _setup_shared_mem("retinanet_anchors", (batch_size * batch_count, 120087, 4), dtypes.float64)
shutdown = False
class Cookie:
def __init__(self, bc):
self.bc = bc
def __del__(self):
if not shutdown:
try: _enqueue_batch(self.bc)
except StopIteration: pass
def shuffle_indices(indices, seed):
rng = random.Random(seed)
rng.shuffle(indices)
if shuffle: shuffle_indices(image_ids, seed=seed)
dataset_iter = iter(image_ids)
try:
for _ in range(NUM_CPU_THREADS.value):
proc = Process(
target=load_retinanet_data,
args=(base_dir, val, queue_in, queue_out, imgs, boxes, labels),
kwargs={"matches": matches, "anchors": anchors, "seed": seed}
)
proc.daemon = True
proc.start()
procs.append(proc)
for bc in range(batch_count):
_enqueue_batch(bc)
for _ in range(len(image_ids) // batch_size):
while True:
bc = queue_out.get() // batch_size
data_out_count[bc] += 1
if data_out_count[bc] == batch_size: break
data_out_count[bc] = 0
if val:
yield (imgs[bc * batch_size:(bc + 1) * batch_size],
img_ids[bc * batch_size:(bc + 1) * batch_size],
img_sizes[bc * batch_size:(bc + 1) * batch_size],
Cookie(bc))
else:
yield (imgs[bc * batch_size:(bc + 1) * batch_size],
boxes[bc * batch_size:(bc + 1) * batch_size],
labels[bc * batch_size:(bc + 1) * batch_size],
matches[bc * batch_size:(bc + 1) * batch_size],
anchors[bc * batch_size:(bc + 1) * batch_size],
Cookie(bc))
finally:
shutdown = True
for _ in procs: queue_in.put(None)
queue_in.close()
for _ in procs:
while queue_out.get() is not None: pass
queue_out.close()
# shutdown processes
for proc in procs: proc.join()
shm_imgs.close()
if not val:
shm_boxes.close()
shm_labels.close()
shm_matches.close()
shm_anchors.close()
try:
shm_imgs.unlink()
if not val:
shm_boxes.unlink()
shm_labels.unlink()
shm_matches.unlink()
shm_anchors.unlink()
except FileNotFoundError:
# happens with BENCHMARK set
pass
# stable diffusion callbacks to match mlperf ref; declared here because they're pickled
def filter_dataset(sample:dict): return {k:v for k,v in sample.items() if k in {'npy', 'txt'}}
def collate(batch:list[dict]):
ret = {"npy": [], "txt": [], "__key__": []}
for sample in batch:
for k,v in sample.items():
ret[k].append(v)
return ret
def collate_fn(batch): return batch
# Reference (code): https://github.com/mlcommons/training/blob/2f4a93fb4888180755a8ef55f4b977ef8f60a89e/stable_diffusion/ldm/data/webdatasets.py, Line 55
# Reference (params): https://github.com/mlcommons/training/blob/ab4ae1ca718d7fe62c369710a316dff18768d04b/stable_diffusion/configs/train_01x08x08.yaml, Line 107
def batch_load_train_stable_diffusion(urls:str, BS:int):
import webdataset
dataset = webdataset.WebDataset(urls=urls, resampled=True, cache_size=-1, cache_dir=None)
dataset = dataset.shuffle(size=1000)
dataset = dataset.decode()
dataset = dataset.map(filter_dataset)
dataset = dataset.batched(BS, partial=False, collation_fn=collate)
dataset = webdataset.WebLoader(dataset, batch_size=None, shuffle=False, num_workers=1, persistent_workers=True, collate_fn=collate_fn)
for x in dataset:
assert isinstance(x, dict) and all(isinstance(k, str) for k in x.keys()) and all(isinstance(v, list) for v in x.values())
assert all(isinstance(moment_mean_logvar, np.ndarray) and moment_mean_logvar.shape==(1,8,64,64) for moment_mean_logvar in x["npy"])
assert all(isinstance(caption, str) for caption in x["txt"])
yield x
# llama3
class BinIdxDataset:
def __init__(self, base_path:Path):
self.idx_t = Tensor(base_path.with_name(f"{base_path.name}.idx"))
self.idx = TensorIO(self.idx_t)
# parse idx file
magic = self.idx.read(9)
assert magic == b"MMIDIDX\x00\x00", "invalid index file format"
version, = struct.unpack("<Q", self.idx.read(8))
assert version == 1, "unsupported index version"
dtype_code, = struct.unpack("<B", self.idx.read(1))
self.dtype = {1:np.dtype(np.uint8), 2:np.dtype(np.int8), 3:np.dtype(np.int16), 4:np.dtype(np.int32), 5:np.dtype(np.int64), 6:np.dtype(np.float64), 7:np.dtype(np.double), 8:np.dtype(np.uint16)}[dtype_code]
self.count, = struct.unpack("<Q", self.idx.read(8))
doc_count, = struct.unpack("<Q", self.idx.read(8))
start = self.idx.tell()
end = start + self.count * dtypes.int32.itemsize
self.sizes = self.idx_t[start:end].bitcast(dtypes.int32).numpy()
start = end
end = start + self.count * dtypes.int64.itemsize
self.pointers = self.idx_t[start:end].bitcast(dtypes.int64).numpy()
start = end
end = start + doc_count * dtypes.int64.itemsize
self.doc_idx = self.idx_t[start:end].bitcast(dtypes.int64).numpy()
# bin file
self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin")).numpy()
def _index(self, idx) -> tuple[int, int]:
return int(self.pointers[idx]), int(self.sizes[idx])
def get(self, idx, offset:int=0, length:int|None=None):
ptr, size = self._index(idx)
if length is None: length = size - offset
ptr += offset * self.dtype.itemsize
return self.bin_t[ptr:ptr+length*self.dtype.itemsize].view(self.dtype)
# https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/datasets.html
class GPTDataset:
def __init__(self, base_path:Path, samples:int, seqlen:int, seed:int, shuffle:bool):
self.samples, self.seqlen = samples, seqlen
self.shuffle = shuffle
self.rng = np.random.RandomState(seed)
self.indexed_dataset = BinIdxDataset(base_path)
# check for cache
cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
cache_path = base_path.with_name(f"{base_path.name}.{cache_hash}.index_cache")
print(f"try loading GPTDataset from {cache_path}...")
if cache_path.exists():
print("cache found, loading...")
with open(cache_path, "rb") as f:
self.doc_idx, self.sample_idx, self.shuffle_idx = pickle.load(f)
else:
print("cache not found, building index...")
self.doc_idx = self._build_doc_idx()
self.sample_idx = self._build_sample_idx()
self.shuffle_idx = self._build_shuffle_idx()
# save cache
with open(cache_path, "wb") as f:
pickle.dump((self.doc_idx, self.sample_idx, self.shuffle_idx), f)
def __getitem__(self, idx):
if idx is None:
text = self._get(0)
else:
text = self._get(idx)
return text
def _get(self, idx):
idx = self.shuffle_idx[idx]
doc_idx_beg, doc_idx_beg_offset = self.sample_idx[idx]
doc_idx_end, doc_idx_end_offset = self.sample_idx[idx + 1]
doc_ids, sample_parts = [], []
if doc_idx_beg == doc_idx_end:
doc_ids.append(self.doc_idx[doc_idx_beg])
sample_parts.append(
self.indexed_dataset.get(
int(self.doc_idx[doc_idx_beg]), offset=int(doc_idx_beg_offset), length=int(doc_idx_end_offset - doc_idx_beg_offset + 1)))
else:
for i in range(doc_idx_beg, doc_idx_end + 1):
doc_ids.append(self.doc_idx[i])
offset = 0 if i > doc_idx_beg else doc_idx_beg_offset
length = None if i < doc_idx_end else int(doc_idx_end_offset + 1)
sample_parts.append(self.indexed_dataset.get(int(self.doc_idx[i]), offset=int(offset), length=length))
# concat all parts
text = np.concatenate(sample_parts, axis=0)
return text
@functools.cached_property
def tokens_per_epoch(self) -> int:
return sum(self.indexed_dataset.sizes.tolist())
@functools.cached_property
def num_epochs(self) -> int:
# we need enough epochs to cover the requested amount of tokens
num_epochs = 1
num_tokens = self.tokens_per_epoch
while num_tokens < self.samples * self.seqlen:
num_epochs += 1
num_tokens += self.tokens_per_epoch
return num_epochs
# https://github.com/NVIDIA/Megatron-LM/blob/94bd476bd840c2fd4c3ebfc7448c2af220f4832b/megatron/core/datasets/gpt_dataset.py#L558
def _build_doc_idx(self):
print(f"building doc_idx for {self.num_epochs=}, {self.indexed_dataset.count=}")
st = time.perf_counter()
# doc_idx = np.mgrid[:self.num_epochs, :self.indexed_dataset.count][1]
doc_idx = np.arange(self.indexed_dataset.count).reshape(1, -1).repeat(self.num_epochs, axis=0).flatten()
doc_idx = doc_idx.astype(np.int32)
at = time.perf_counter()
if self.shuffle: self.rng.shuffle(doc_idx)
print(f"doc_idx built in {at - st:.3f}s, shuffled in {time.perf_counter() - at:.3f}s")
return doc_idx
def _build_sample_idx(self):
print(f"building sample_idx for {self.samples=}, {self.seqlen=}, {self.doc_idx.shape[0]=}")
sample_idx_max = max(self.doc_idx.shape[0], self.indexed_dataset.sizes.max())
sample_idx = np.empty((self.samples + 1, 2), dtype=np.int64 if sample_idx_max > dtypes.int32.max else np.int32)
sample_idx_idx, doc_idx_idx, doc_offset = 0, 0, 0
sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
sample_idx_idx += 1
for _ in tqdm(range(1, self.samples + 1)):
remaining_seqlen = self.seqlen + 1
while remaining_seqlen > 0:
doc_idx = int(self.doc_idx[doc_idx_idx])
doc_len = int(self.indexed_dataset.sizes[doc_idx]) - doc_offset
remaining_seqlen -= doc_len
if remaining_seqlen <= 0:
doc_offset += remaining_seqlen + doc_len - 1
remaining_seqlen = 0
else:
if doc_idx_idx == len(self.doc_idx) - 1:
assert sample_idx_idx == self.samples
doc_idx = int(self.doc_idx[doc_idx_idx])
doc_offset = int(self.indexed_dataset.sizes[doc_idx]) - 1
break
doc_idx_idx += 1
doc_offset = 0
sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
sample_idx_idx += 1
return sample_idx
def _build_shuffle_idx(self):
print(f"building shuffle_idx for {self.samples=}")
st = time.perf_counter()
shuffle_idx = np.arange(self.samples, dtype=np.int32)
at = time.perf_counter()
if self.shuffle: self.rng.shuffle(shuffle_idx)
print(f"shuffle_idx built in {at - st:.3f}s, shuffled in {time.perf_counter() - at:.3f}s")
return shuffle_idx
class BlendedGPTDataset:
def __init__(self, paths:list[Path], weights:list[float], samples:int, seqlen:int, seed:int, shuffle:bool):
self.shuffle = shuffle
self.rng = np.random.RandomState(seed)
# normalize weights
total_weight = sum(weights)
self.weights = [w / total_weight for w in weights]
self.samples = samples
surplus = 0.005
samples_per_blend = [math.ceil(math.ceil(self.samples * w) * (1 + surplus)) for w in self.weights]
self.datasets = [GPTDataset(path, samples_per_blend[i], seqlen, seed + i, shuffle) for i,path in enumerate(paths)]
# check for cache
cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
cache_path = paths[0].with_name(f"{paths[0].name}.{cache_hash}.blend_cache")
print(f"try loading BlendedGPTDataset from {cache_path}...")
if cache_path.exists():
print("cache found, loading...")
with open(cache_path, "rb") as f:
self.dataset_idx, self.dataset_sample_idx = pickle.load(f)
else:
print("cache not found, building index...")
self.dataset_idx, self.dataset_sample_idx = self._build_blend_idx()
# save cache
with open(cache_path, "wb") as f:
pickle.dump((self.dataset_idx, self.dataset_sample_idx), f)
def get(self, idx:int):
tokens = self.datasets[self.dataset_idx[idx]][self.dataset_sample_idx[idx]]
return tokens
def _build_blend_idx(self):
dataset_idx = np.zeros(self.samples, dtype=np.int16)
dataset_sample_idx = np.zeros(self.samples, dtype=np.int64)
unspent_datasets = set(range(len(self.datasets)))
dataset_sample_counts = [0] * len(self.datasets)
for i in tqdm(range(self.samples)):
error_argmax, error_max = 0, 0.0
for di in unspent_datasets:
error = self.weights[di] * max(i, 1) - dataset_sample_counts[di]
if error > error_max:
error_max = error
error_argmax = di
dataset_idx[i] = error_argmax
dataset_sample_idx[i] = dataset_sample_counts[error_argmax]
dataset_sample_counts[error_argmax] += 1
return dataset_idx, dataset_sample_idx
def get_llama3_dataset(samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False) -> BlendedGPTDataset:
if small:
if val:
return BlendedGPTDataset(
[base_dir / "c4-validation-91205-samples.en_text_document"], [1.0], samples, seqlen, seed, shuffle=False)
return BlendedGPTDataset(
[base_dir / "c4-train.en_6_text_document"], [1.0], samples, seqlen, seed, shuffle=True)
if val:
return BlendedGPTDataset(
[base_dir / "validation" / "c4-validationn-91205-samples.en_text_document"], [1.0], samples, seqlen, seed, shuffle=False)
return BlendedGPTDataset(
[base_dir / "c4-train.en_6_text_document", base_dir / "c4-train.en_7_text_document"], [1.0, 1.0], samples, seqlen, seed, shuffle=True)
def iterate_llama3_dataset(dataset:BlendedGPTDataset, bs:int):
for b in range(math.ceil(dataset.samples / bs)):
batch = [dataset.get(b * bs + i) for i in range(bs)]
stacked = np.stack(batch, axis=0)
yield Tensor(stacked, device="NPY")
def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False):
return iterate_llama3_dataset(get_llama3_dataset(samples, seqlen, base_dir, seed, val, small), bs)
if __name__ == "__main__":
def load_unet3d(val):
assert not val, "validation set is not supported due to different sizes on inputs"
from extra.datasets.kits19 import get_train_files, get_val_files, preprocess_dataset, TRAIN_PREPROCESSED_DIR, VAL_PREPROCESSED_DIR
preprocessed_dir = VAL_PREPROCESSED_DIR if val else TRAIN_PREPROCESSED_DIR
files = get_val_files() if val else get_train_files()
if not preprocessed_dir.exists(): preprocess_dataset(files, preprocessed_dir, val)
with tqdm(total=len(files)) as pbar:
for x, _, _ in batch_load_unet3d(preprocessed_dir, val=val):
pbar.update(x.shape[0])
def load_resnet(val):
from extra.datasets.imagenet import get_train_files, get_val_files
files = get_val_files() if val else get_train_files()
with tqdm(total=len(files)) as pbar:
for x,y,c in batch_load_resnet(val=val):
pbar.update(x.shape[0])
def load_retinanet(val):
from extra.datasets.openimages import BASEDIR, download_dataset
from pycocotools.coco import COCO
dataset = COCO(download_dataset(base_dir:=getenv("BASEDIR", BASEDIR), "validation" if val else "train"))
with tqdm(total=len(dataset.imgs.keys())) as pbar:
for x in batch_load_retinanet(dataset, val, base_dir):
pbar.update(x[0].shape[0])
def load_llama3(val):
bs = 24
samples = 5760 if val else 1_200_000 * 1152
seqlen = 8192
max_, min_ = 0, math.inf
for tokens in tqdm(batch_load_llama3(bs, samples, seqlen, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=5760, val=bool(val)), total=samples//bs):
max_ = max(max_, tokens.shape[1])
min_ = min(min_, tokens.shape[1])
print(f"max seq length: {max_}")
print(f"min seq length: {min_}")
load_fn_name = f"load_{getenv('MODEL', 'resnet')}"
if load_fn_name in globals():
globals()[load_fn_name](getenv("VAL", 1))

367
artifacts/package_sources/tinygrad/examples/mlperf/helpers.py Normal file
View File

@@ -0,0 +1,367 @@
from collections import OrderedDict
import unicodedata
from typing import Optional
import math
import numpy as np
from tinygrad.nn import state
from tinygrad.tensor import Tensor, dtypes
from tinygrad.helpers import getenv
#
# checkpointing utils
#
def invert_dict(d): return {v: k for k, v in reversed(d.items())}
def dedup_dict(d): return invert_dict(invert_dict(d))
# store each tensor into the first key it appears in
def get_training_state(model, optimizer, scheduler):
# hack: let get_state_dict walk the tree starting with model, so that the checkpoint keys are
# readable and can be loaded as a model for eval
train_state = {'model': model, 'optimizer': optimizer, 'scheduler': scheduler}
return dedup_dict(state.get_state_dict(train_state))
def load_training_state(model, optimizer, scheduler, state_dict):
# use fresh model to restore duplicate keys
train_state = {'model': model, 'optimizer': optimizer, 'scheduler': scheduler}
big_dict = state.get_state_dict(train_state)
# hack: put back the dupes
dupe_names = {}
for k, v in big_dict.items():
if v not in dupe_names:
dupe_names[v] = k
assert k in state_dict
state_dict[k] = state_dict[dupe_names[v]]
# scheduler contains optimizer and all params, load each weight only once
scheduler_state = {'scheduler': scheduler}
state.load_state_dict(scheduler_state, state_dict)
def gaussian_kernel(n, std):
from scipy import signal
gaussian_1d = signal.windows.gaussian(n, std)
gaussian_2d = np.outer(gaussian_1d, gaussian_1d)
gaussian_3d = np.outer(gaussian_2d, gaussian_1d)
gaussian_3d = gaussian_3d.reshape(n, n, n)
gaussian_3d = np.cbrt(gaussian_3d)
gaussian_3d /= gaussian_3d.max()
return gaussian_3d
def prepare_arrays(image, roi_shape=(128, 128, 128)):
assert len(roi_shape) == 3 and any(roi_shape)
image_shape = list(image.shape[2:])
result = np.zeros((1, 3, *image_shape), dtype=image.dtype)
norm_map = np.zeros_like(result)
norm_patch = gaussian_kernel(roi_shape[0], 0.125 * roi_shape[0]).astype(norm_map.dtype)
return result, norm_map, norm_patch
def get_slice(image, roi_shape=(128, 128, 128), overlap_factor=0.5):
assert len(roi_shape) == 3 and any(roi_shape)
assert 0 < overlap_factor < 1
image_shape, dim = list(image.shape[2:]), len(image.shape[2:])
strides = [int(roi_shape[i] * (1 - overlap_factor)) for i in range(dim)]
size = [(image_shape[i] - roi_shape[i]) // strides[i] + 1 for i in range(dim)]
for i in range(0, strides[0] * size[0], strides[0]):
for j in range(0, strides[1] * size[1], strides[1]):
for k in range(0, strides[2] * size[2], strides[2]):
yield i, j, k
def _get_best_indices(logits, n_best_size):
index_and_score = sorted(enumerate(logits), key=lambda x: x[1], reverse=True)
return list(map(lambda x: x[0], index_and_score))[:n_best_size]
def _is_punctuation(char):
if (cp := ord(char)) in range(33, 48) or cp in range(58, 65) or cp in range(91, 97) or cp in range(123, 127):
return True
return unicodedata.category(char).startswith("P")
def _is_whitespace(char):
if char == " " or char == "\t" or char == "\n" or char == "\r":
return True
return unicodedata.category(char) == "Zs"
def _is_control(char):
if char == "\t" or char == "\n" or char == "\r":
return False
return unicodedata.category(char).startswith("C")
def _run_split_on_punc(text):
if text in ("[UNK]", "[SEP]", "[PAD]", "[CLS]", "[MASK]"):
return [text]
start_new_word = True
output = []
for i in range(len(text)):
if _is_punctuation(char := text[i]):
output.append([char])
start_new_word = True
else:
if start_new_word:
output.append([])
start_new_word = False
output[-1].append(char)
return ["".join(x) for x in output]
def _run_strip_accents(text):
output = []
for char in unicodedata.normalize("NFD", text):
if unicodedata.category(char) != "Mn":
output.append(char)
return "".join(output)
def _clean_text(text):
output = []
for char in text:
if not ((cp := ord(char)) == 0 or cp == 0xfffd or _is_control(char)):
output.append(" " if _is_whitespace(char) else char)
return "".join(output)
def _get_final_text(pred_text, orig_text):
def _strip_spaces(text):
ns_text = ""
ns_to_s_map = OrderedDict()
for i, c in enumerate(text):
if c == " ":
continue
ns_to_s_map[len(ns_text)] = i
ns_text += c
return ns_text, ns_to_s_map
orig_tokens = _clean_text(orig_text).strip().split()
split_tokens = []
for token in orig_tokens:
if token not in ("[UNK]", "[SEP]", "[PAD]", "[CLS]", "[MASK]"):
token = token.lower()
token = _run_strip_accents(token)
split_tokens.extend(_run_split_on_punc(token))
tok_text = " ".join(" ".join(split_tokens).strip().split())
start_position = tok_text.find(pred_text)
if start_position == -1:
return orig_text
end_position = start_position + len(pred_text) - 1
orig_ns_text, orig_ns_to_s_map = _strip_spaces(orig_text)
tok_ns_text, tok_ns_to_s_map = _strip_spaces(tok_text)
if len(orig_ns_text) != len(tok_ns_text):
return orig_text
tok_s_to_ns_map = {v: k for k, v in tok_ns_to_s_map.items()}
orig_start_position = None
if start_position in tok_s_to_ns_map:
if (ns_start_position := tok_s_to_ns_map[start_position]) in orig_ns_to_s_map:
orig_start_position = orig_ns_to_s_map[ns_start_position]
if orig_start_position is None:
return orig_text
orig_end_position = None
if end_position in tok_s_to_ns_map:
if (ns_end_position := tok_s_to_ns_map[end_position]) in orig_ns_to_s_map:
orig_end_position = orig_ns_to_s_map[ns_end_position]
if orig_end_position is None:
return orig_text
output_text = orig_text[orig_start_position:(orig_end_position + 1)]
return output_text
def get_bert_qa_prediction(features, example, start_end_logits):
prelim_predictions = []
for i, feature in enumerate(features):
for start_index in _get_best_indices(start_end_logits[i][0], 20):
for end_index in _get_best_indices(start_end_logits[i][1], 20):
if start_index >= len(feature["tokens"]) or end_index >= len(feature["tokens"]):
continue
if start_index not in feature["token_to_orig_map"] or end_index not in feature["token_to_orig_map"]:
continue
if not feature["token_is_max_context"].get(start_index, False):
continue
if end_index < start_index or end_index - start_index + 1 > 30:
continue
prelim_predictions.append({
"feature_index": i,
"start_index": start_index,
"end_index": end_index,
"start_logit": start_end_logits[i][0, start_index],
"end_logit": start_end_logits[i][1, end_index]
})
predictions = sorted(prelim_predictions, key=lambda x: (x["start_logit"] + x["end_logit"]), reverse=True)
if len(predictions) > 0:
feature = features[predictions[0]["feature_index"]]
tok_tokens = feature["tokens"][predictions[0]["start_index"]:(predictions[0]["end_index"] + 1)]
orig_doc_start = feature["token_to_orig_map"][predictions[0]["start_index"]]
orig_doc_end = feature["token_to_orig_map"][predictions[0]["end_index"]]
orig_tokens = example["context"][orig_doc_start:(orig_doc_end + 1)]
tok_text = " ".join(tok_tokens).replace(" ##", "").replace("##", "")
tok_text = " ".join(tok_text.strip().split())
orig_text = " ".join(orig_tokens)
return _get_final_text(tok_text, orig_text)
return "empty"
def get_mlperf_bert_config():
"""benchmark is BERT-large"""
ret = {"attention_probs_dropout_prob": 0.1, "hidden_dropout_prob": 0.1, "vocab_size": 30522, "type_vocab_size": 2, "max_position_embeddings": 512}
match (bert_size:=getenv("BERT_SIZE", "large")):
case "large": ret.update({"hidden_size": 1024, "intermediate_size": 4096, "num_attention_heads": 16, "num_hidden_layers": 24})
case "tiny": ret.update({"hidden_size": 128, "intermediate_size": 512, "num_attention_heads": 2, "num_hidden_layers": 2})
case _: raise RuntimeError(f"unhandled {bert_size=}")
if (bert_layers:=getenv("BERT_LAYERS")): ret["num_hidden_layers"] = bert_layers
return ret
def get_mlperf_bert_model():
from extra.models import bert
from examples.mlperf.initializers import LinearBert, EmbeddingBert, LayerNormBert
bert.Linear = LinearBert
bert.Embedding = EmbeddingBert
bert.LayerNorm = LayerNormBert
from extra.models.bert import BertForPretraining
config = get_mlperf_bert_config()
if getenv("DISABLE_DROPOUT", 0):
config["hidden_dropout_prob"] = config["attention_probs_dropout_prob"] = 0.0
model = BertForPretraining(**config)
if getenv("FP8_TRAIN"):
from extra.fp8.fp8_linear import convert_to_float8_training
def module_filter_fn(mod, fqn):
if isinstance(mod, LinearBert):
skip_layers = [] if (ln:=config["num_hidden_layers"]) <= 2 else ["bert.encoder.layer.0.", f"bert.encoder.layer.{ln-1}"]
if mod.weight.shape[-1] >= 1024 and "encoder" in fqn and not any(name in fqn for name in skip_layers):
print(f"replacing linear with fp8: {fqn} {mod.weight.shape}")
return True
return False
convert_to_float8_training(model, module_filter_fn)
return model
def get_fake_data_bert(BS:int):
return {
"input_ids": Tensor.zeros((BS, 512), dtype=dtypes.int32, device="CPU").contiguous(),
"input_mask": Tensor.zeros((BS, 512), dtype=dtypes.int32, device="CPU").contiguous(),
"segment_ids": Tensor.zeros((BS, 512), dtype=dtypes.int32, device="CPU").contiguous(),
"masked_lm_positions": Tensor.zeros((BS, 76), dtype=dtypes.int32, device="CPU").contiguous(),
"masked_lm_ids": Tensor.zeros((BS, 76), dtype=dtypes.int32, device="CPU").contiguous(),
"masked_lm_weights": Tensor.zeros((BS, 76), dtype=dtypes.float32, device="CPU").contiguous(),
"next_sentence_labels": Tensor.zeros((BS, 1), dtype=dtypes.int32, device="CPU").contiguous(),
}
def find_matches(match_quality_matrix:np.ndarray, high_threshold:float=0.5, low_threshold:float=0.4, allow_low_quality_matches:bool=False) -> np.ndarray:
BELOW_LOW_THRESHOLD, BETWEEN_THRESHOLDS = -1, -2
def _set_low_quality_matches_(matches:np.ndarray, all_matches:np.ndarray, match_quality_matrix:np.ndarray):
highest_quality_foreach_gt = np.max(match_quality_matrix, axis=1)
pred_inds_to_update = np.nonzero(match_quality_matrix == highest_quality_foreach_gt[:, None])[1]
matches[pred_inds_to_update] = all_matches[pred_inds_to_update]
assert low_threshold <= high_threshold
matched_vals, matches = match_quality_matrix.max(axis=0), match_quality_matrix.argmax(axis=0)
all_matches = np.copy(matches) if allow_low_quality_matches else None
below_low_threshold = matched_vals < low_threshold
between_thresholds = (matched_vals >= low_threshold) & (matched_vals < high_threshold)
matches[below_low_threshold] = BELOW_LOW_THRESHOLD
matches[between_thresholds] = BETWEEN_THRESHOLDS
if allow_low_quality_matches:
assert all_matches is not None
_set_low_quality_matches_(matches, all_matches, match_quality_matrix)
return matches
def box_iou(boxes1:np.ndarray, boxes2:np.ndarray) -> np.ndarray:
def _box_area(boxes:np.ndarray) -> np.ndarray: return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
def _box_inter_union(boxes1:np.ndarray, boxes2:np.ndarray) -> tuple[np.ndarray, np.ndarray]:
area1, area2 = _box_area(boxes1), _box_area(boxes2)
lt, rb = np.maximum(boxes1[:, None, :2], boxes2[:, :2]), np.minimum(boxes1[:, None, 2:], boxes2[:, 2:])
wh = np.clip(rb - lt, a_min=0, a_max=None)
inter = wh[:, :, 0] * wh[:, :, 1]
union = area1[:, None] + area2 - inter
return inter, union
inter, union = _box_inter_union(boxes1, boxes2)
return inter / union
def generate_anchors(input_size:tuple[int, int], scales:Optional[tuple[Tensor, ...]]=None, aspect_ratios:Optional[tuple[Tensor, ...]]=None) -> list[np.ndarray]:
def _compute_grid_sizes(input_size:tuple[int, int]) -> np.ndarray:
return np.ceil(np.array(input_size)[None, :] / 2 ** np.arange(3, 8)[:, None])
scales = tuple((i, int(i * 2 ** (1/3)), int(i * 2 ** (2/3))) for i in 2 ** np.arange(5, 10)) if scales is None else scales
aspect_ratios = ((0.5, 1.0, 2.0),) * len(scales) if aspect_ratios is None else aspect_ratios
aspect_ratios = tuple(ar for ar in aspect_ratios)
grid_sizes = _compute_grid_sizes(input_size)
assert len(scales) == len(aspect_ratios) == len(grid_sizes), "scales, aspect_ratios, and grid_sizes must have the same length"
anchors = []
for s, ar, gs in zip(scales, aspect_ratios, grid_sizes):
s, ar = np.array(s), np.array(ar)
h_ratios = np.sqrt(ar)
w_ratios = 1 / h_ratios
ws = (w_ratios[:, None] * s[None, :]).reshape(-1)
hs = (h_ratios[:, None] * s[None, :]).reshape(-1)
base_anchors = (np.stack([-ws, -hs, ws, hs], axis=1) / 2).round()
stride_h, stride_w = input_size[0] // gs[0], input_size[1] // gs[1]
shifts_x, shifts_y = np.meshgrid(np.arange(gs[1]) * stride_w, np.arange(gs[0]) * stride_h)
shifts_x, shifts_y = shifts_x.reshape(-1), shifts_y.reshape(-1)
shifts = np.stack([shifts_x, shifts_y, shifts_x, shifts_y], axis=1, dtype=np.float32)
anchors.append((shifts[:, None] + base_anchors[None, :]).reshape(-1, 4))
return anchors
class BoxCoder(object):
def __init__(self, weights, bbox_xform_clip=math.log(1000. / 16), apply_to_remove=True):
self.weights = weights
self.bbox_xform_clip = bbox_xform_clip
self.apply_to_remove = apply_to_remove
def encode(self, reference_boxes, proposals):
TO_REMOVE = self.apply_to_remove # TODO remove
ex_widths = proposals[..., 2] - proposals[..., 0] + TO_REMOVE
ex_heights = proposals[..., 3] - proposals[..., 1] + TO_REMOVE
ex_ctr_x = proposals[..., 0] + 0.5 * ex_widths
ex_ctr_y = proposals[..., 1] + 0.5 * ex_heights
gt_widths = reference_boxes[..., 2] - reference_boxes[..., 0] + TO_REMOVE
gt_heights = reference_boxes[..., 3] - reference_boxes[..., 1] + TO_REMOVE
gt_ctr_x = reference_boxes[..., 0] + 0.5 * gt_widths
gt_ctr_y = reference_boxes[..., 1] + 0.5 * gt_heights
wx, wy, ww, wh = self.weights
targets_dx = wx * (gt_ctr_x - ex_ctr_x) / ex_widths
targets_dy = wy * (gt_ctr_y - ex_ctr_y) / ex_heights
targets_dw = ww * Tensor.log(gt_widths / ex_widths)
targets_dh = wh * Tensor.log(gt_heights / ex_heights)
targets = Tensor.stack(targets_dx, targets_dy, targets_dw, targets_dh, dim=-1)
return targets
def decode(self, rel_codes, boxes):
boxes = boxes.cast(rel_codes.dtype)
rel_codes = rel_codes
TO_REMOVE = self.apply_to_remove # TODO remove
widths = boxes[:, 2] - boxes[:, 0] + TO_REMOVE
heights = boxes[:, 3] - boxes[:, 1] + TO_REMOVE
ctr_x = boxes[:, 0] + 0.5 * widths
ctr_y = boxes[:, 1] + 0.5 * heights
wx, wy, ww, wh = self.weights
dx = rel_codes[:, 0::4] / wx
dy = rel_codes[:, 1::4] / wy
dw = rel_codes[:, 2::4] / ww
dh = rel_codes[:, 3::4] / wh
# Prevent sending too large values into Tensor.exp()
dw = dw.clip(min_=dw.min(), max_=self.bbox_xform_clip)
dh = dh.clip(min_=dh.min(), max_=self.bbox_xform_clip)
pred_ctr_x = dx * widths[:, None] + ctr_x[:, None]
pred_ctr_y = dy * heights[:, None] + ctr_y[:, None]
pred_w = dw.exp() * widths[:, None]
pred_h = dh.exp() * heights[:, None]
x = pred_ctr_x - 0.5 * pred_w
y = pred_ctr_y - 0.5 * pred_h
w = pred_ctr_x + 0.5 * pred_w - 1
h = pred_ctr_y + 0.5 * pred_h - 1
pred_boxes = Tensor.stack(x, y, w, h).permute(1,2,0).reshape(rel_codes.shape[0], rel_codes.shape[1])
return pred_boxes

189
artifacts/package_sources/tinygrad/examples/mlperf/initializers.py Normal file
View File

@@ -0,0 +1,189 @@
import math
from typing import Union
from tinygrad import Tensor, nn, dtypes
from tinygrad.helpers import prod, argfix, Context, TRAINING
from tinygrad.nn.state import get_parameters
from extra.models.unet import UNetModel
# rejection sampling truncated randn
def rand_truncn(*shape, dtype=None, truncstds=2, **kwargs) -> Tensor:
CNT=8
x = Tensor.randn(*(*shape, CNT), dtype=dtype, **kwargs)
ctr = Tensor.arange(CNT).reshape((1,) * len(x.shape[:-1]) + (CNT,)).expand(x.shape)
take = (x.abs() <= truncstds).where(ctr, CNT).min(axis=-1, keepdim=True) # set to 0 if no good samples
return (ctr == take).where(x, 0).sum(axis=-1)
# https://github.com/keras-team/keras/blob/v2.15.0/keras/initializers/initializers.py#L1026-L1065
def he_normal(*shape, a: float = 0.00, **kwargs) -> Tensor:
std = math.sqrt(2.0 / (1 + a ** 2)) / math.sqrt(prod(argfix(*shape)[1:])) / 0.87962566103423978
return std * rand_truncn(*shape, **kwargs)
# Stable Diffusion v2 training uses default torch gelu, which doesn't use tanh approximation
def gelu_erf(x:Tensor) -> Tensor:
return 0.5 * x * (1.0 + (x / 1.4142135623730951).erf())
class Conv2dHeNormal(nn.Conv2d):
def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):
super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
self.in_channels, self.out_channels = in_channels, out_channels # for testing
self.weight = he_normal(out_channels, in_channels//groups, *self.kernel_size, a=0.0, dtype=dtypes.float32)
if bias: self.bias = self.bias.cast(dtypes.float32)
def __call__(self, x: Tensor):
return x.conv2d(self.weight.cast(dtypes.default_float), self.bias.cast(dtypes.default_float) if self.bias is not None else None,
padding=self.padding, stride=self.stride, dilation=self.dilation, groups=self.groups)
class Linear(nn.Linear):
def __init__(self, in_features, out_features, bias=True):
super().__init__(in_features, out_features, bias=bias)
self.weight = Tensor.normal((out_features, in_features), mean=0.0, std=0.01, dtype=dtypes.float32)
if bias: self.bias = Tensor.zeros(out_features, dtype=dtypes.float32)
def __call__(self, x:Tensor):
return x.linear(self.weight.cast(dtypes.default_float).transpose(), self.bias.cast(dtypes.default_float) if self.bias is not None else None)
class LinearBert(nn.Linear):
def __init__(self, in_features, out_features, bias=True, std=0.02):
self.weight = std * rand_truncn(out_features, in_features, dtype=dtypes.float32)
self.bias = Tensor.zeros(out_features, dtype=dtypes.float32) if bias else None
def __call__(self, x:Tensor):
return x.cast(dtypes.default_float).linear(self.weight.cast(dtypes.default_float).transpose(), self.bias.cast(dtypes.default_float) if self.bias is not None else None)
class EmbeddingBert(nn.Embedding):
def __init__(self, vocab_size:int, embed_size:int, std=0.02):
self.vocab_sz, self.embed_sz = vocab_size, embed_size
self.weight = std * rand_truncn(vocab_size, embed_size, dtype=dtypes.float32)
def __call__(self, idx:Tensor) -> Tensor:
if idx.numel() == 0: return Tensor.empty(idx.shape+(self.embed_sz,), dtype=self.weight.dtype, device=self.weight.device)
arange_shp, weight_shp, big_shp = (1, 1, self.vocab_sz, 1), (1, 1, self.vocab_sz, self.embed_sz), idx.shape+(self.vocab_sz, self.embed_sz,)
if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz).reshape(arange_shp)
arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1,)).expand(big_shp), self.weight.cast(dtypes.default_float).reshape(weight_shp).expand(big_shp)
return (arange == idx).where(vals, 0).sum(2, dtype=vals.dtype)
class LayerNormBert:
def __init__(self, normalized_shape:Union[int, tuple[int, ...]], eps:float=1e-12, elementwise_affine:bool=True):
self.normalized_shape = (normalized_shape,) if isinstance(normalized_shape, int) else tuple(normalized_shape)
self.axis, self.eps, self.elementwise_affine = tuple(-1-i for i in range(len(self.normalized_shape))), eps, elementwise_affine
self.weight, self.bias = (Tensor.ones(*self.normalized_shape, dtype=dtypes.float32), Tensor.zeros(*self.normalized_shape, dtype=dtypes.float32)) if elementwise_affine else (None, None)
def __call__(self, x:Tensor):
assert self.normalized_shape == x.shape[-len(self.normalized_shape):], f"last dimensions of {x.shape} must match {self.normalized_shape}"
xn = x.cast(dtypes.float32).layernorm(eps=self.eps, axis=self.axis).cast(x.dtype)
if not self.elementwise_affine: return xn
return (xn * self.weight.cast(dtypes.default_float) + self.bias.cast(dtypes.default_float))
class FrozenBatchNorm2dRetinaNet(nn.BatchNorm2d):
def __init__(self, sz:int, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1):
self.eps, self.track_running_stats, self.momentum = eps, track_running_stats, momentum
self.weight = Tensor.ones(sz, dtype=dtypes.float32).is_param_(False) if affine else None
self.bias = Tensor.zeros(sz, dtype=dtypes.float32).is_param_(False) if affine else None
if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, dtype=dtypes.float32).is_param_(False), Tensor.ones(sz, dtype=dtypes.float32).is_param_(False)
self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.long).is_param_(False)
def __call__(self, x:Tensor) -> Tensor:
batch_mean, batch_var = super().calc_stats(x.cast(dtypes.float32))
if self.track_running_stats and TRAINING:
self.running_mean.assign((1-self.momentum) * self.running_mean + self.momentum * batch_mean.detach().cast(self.running_mean.dtype))
self.running_var.assign((1-self.momentum) * self.running_var + self.momentum * x.numel()/(x.numel()-x.shape[1]) * batch_var.detach().cast(self.running_var.dtype))
self.num_batches_tracked += 1
return x.cast(dtypes.float32).batchnorm(self.weight, self.bias, batch_mean, batch_var.add(self.eps).rsqrt()).cast(x.dtype)
class Conv2dNormalRetinaNet(nn.Conv2d):
def __init__(self, in_channels:int, out_channels:int, kernel_size:int|tuple[int, ...],
stride:int=1, padding:int|tuple[int, ...]|str=0, dilation:int=1, groups:int=1,
bias:bool=True, prior_prob:float|None=None):
super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
self.weight = Tensor.normal(*self.weight.shape, std=0.01, dtype=dtypes.float32)
if bias:
if prior_prob:
prior_prob = Tensor(prior_prob, device=self.bias.device, dtype=dtypes.float32).expand(*self.bias.shape)
self.bias = -(((1 - prior_prob) / prior_prob).log())
else: self.bias = Tensor.zeros_like(self.bias, dtype=dtypes.float32)
def __call__(self, x:Tensor) -> Tensor:
return x.conv2d(self.weight.cast(dtypes.default_float), self.bias.cast(dtypes.default_float) if self.bias is not None else None,
groups=self.groups, stride=self.stride, padding=self.padding)
class Conv2dKaimingUniformRetinaNet(nn.Conv2d):
def __init__(self, in_channels:int, out_channels:int, kernel_size:int|tuple[int, ...],
stride:int=1, padding:int|tuple[int, ...]|str=0, dilation:int=1, groups:int=1,
bias:bool=True):
super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
self.weight = Tensor.kaiming_uniform(*self.weight.shape, a=1, dtype=dtypes.float32)
if bias: self.bias = Tensor.zeros_like(self.bias, dtype=dtypes.float32)
def __call__(self, x:Tensor) -> Tensor:
return x.conv2d(self.weight.cast(dtypes.default_float), self.bias.cast(dtypes.default_float) if self.bias is not None else None,
groups=self.groups, stride=self.stride, padding=self.padding)
class Conv2dRetinaNet(nn.Conv2d):
def __init__(self, in_channels:int, out_channels:int, kernel_size:int|tuple[int, ...],
stride:int=1, padding:int|tuple[int, ...]|str=0, dilation:int=1, groups:int=1,
bias:bool=True):
super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
scale = 1 / math.sqrt(in_channels * prod(self.kernel_size))
self.weight = Tensor.uniform(out_channels, in_channels//groups, *self.kernel_size, low=-scale, high=scale, dtype=dtypes.float32)
self.bias: Tensor|None = Tensor.uniform(out_channels, low=-scale, high=scale, dtype=dtypes.float32) if bias else None
def __call__(self, x:Tensor) -> Tensor:
return x.conv2d(self.weight.cast(dtypes.default_float), self.bias.cast(dtypes.default_float) if self.bias is not None else None,
groups=self.groups, stride=self.stride, dilation=self.dilation, padding=self.padding)
# copy torch AMP: isolate mixed precision to just the below autocast ops, instead of using dtypes.default_float which affects all new Tensors
class AutocastLinear(nn.Linear):
cast_dtype=dtypes.bfloat16 # enable monkeypatching of the mixed precision dtype
def __call__(self, x:Tensor) -> Tensor:
dtype = type(self).cast_dtype
return x.cast(dtype).linear(self.weight.cast(dtype).transpose(), self.bias.cast(dtype) if self.bias is not None else None)
class AutocastConv2d(nn.Conv2d):
cast_dtype=dtypes.bfloat16
def __call__(self, x:Tensor) -> Tensor:
dtype = type(self).cast_dtype
return x.cast(dtype).conv2d(self.weight.cast(dtype), self.bias.cast(dtype), self.groups, self.stride, self.dilation, self.padding)
# copy torch AMP: upcast to float32 before GroupNorm and LayerNorm
class AutocastGroupNorm(nn.GroupNorm):
def __call__(self, x:Tensor) -> Tensor:
return super().__call__(x.cast(dtypes.float32))
class AutocastLayerNorm(nn.LayerNorm):
def __call__(self, x:Tensor) -> Tensor:
return super().__call__(x.cast(dtypes.float32))
def zero_module(module):
for p in get_parameters(module): p.assign(Tensor.zeros_like(p).contiguous())
# Stable Diffusion mlperf reference doesn't call scaled_dot_product_attention
# copy torch AMP: upcast to float32 before softmax on CUDA
def attn_f32_softmax(q:Tensor, k:Tensor, v:Tensor) -> Tensor:
return (q.matmul(k.transpose(-2,-1), dtype=dtypes.float32) / math.sqrt(q.shape[-1])).softmax(-1).cast(q.dtype) @ v
def init_stable_diffusion(version:str, pretrained:str, devices:list[str]):
from examples.stable_diffusion import StableDiffusion
from tinygrad.nn.state import safe_load, safe_save, load_state_dict, get_state_dict
from tempfile import TemporaryDirectory
model = StableDiffusion(version=version, pretrained=pretrained)
unet:UNetModel = model.model.diffusion_model
# this prevents extra consumption of memory, enabling much larger BS
Tensor.realize(*get_parameters(unet))
with TemporaryDirectory(prefix="unet_init") as tmp:
safe_save(get_state_dict(unet), init_fn:=f"{tmp}/init_model.safetensors")
load_state_dict(unet, safe_load(init_fn))
sqrt_alphas_cumprod = model.alphas_cumprod.sqrt().realize()
sqrt_one_minus_alphas_cumprod = (1 - model.alphas_cumprod).sqrt().realize()
if len(devices) > 1:
to_move = [sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod]
if version == "v2-mlperf-train": to_move += get_parameters(unet) + get_parameters(model.cond_stage_model)
for p in to_move:
p.to_(devices)
with Context(BEAM=0):
Tensor.realize(*to_move)
return model, unet, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod

29
artifacts/package_sources/tinygrad/examples/mlperf/losses.py Normal file
View File

@@ -0,0 +1,29 @@
from examples.mlperf.metrics import dice_score
from tinygrad import Tensor
def dice_ce_loss(pred, tgt):
ce = pred.permute(0, 2, 3, 4, 1).sparse_categorical_crossentropy(tgt.squeeze(1))
dice = (1.0 - dice_score(pred, tgt, argmax=False, to_one_hot_x=False)).mean()
return (dice + ce) / 2
def sigmoid_focal_loss(pred:Tensor, tgt:Tensor, alpha:float=0.25, gamma:float=2.0, reduction:str="none") -> Tensor:
assert reduction in ["mean", "sum", "none"], f"unsupported reduction {reduction}"
p, ce_loss = pred.sigmoid(), pred.binary_crossentropy_logits(tgt, reduction="none")
p_t = p * tgt + (1 - p) * (1 - tgt)
loss = ce_loss * ((1 - p_t) ** gamma)
if alpha >= 0:
alpha_t = alpha * tgt + (1 - alpha) * (1 - tgt)
loss = loss * alpha_t
if reduction == "mean": loss = loss.mean()
elif reduction == "sum": loss = loss.sum()
return loss
def l1_loss(pred:Tensor, tgt:Tensor, reduction:str="none") -> Tensor:
assert reduction in ["mean", "sum", "none"], f"unsupported reduction {reduction}"
loss = (pred - tgt).abs()
if reduction == "mean": loss = loss.mean()
elif reduction == "sum": loss = loss.sum()
return loss

60
artifacts/package_sources/tinygrad/examples/mlperf/lr_schedulers.py Normal file
View File

@@ -0,0 +1,60 @@
import math
from tinygrad import dtypes, Tensor
from tinygrad.nn.optim import Optimizer
from extra.lr_scheduler import LR_Scheduler
from typing import Callable
# https://github.com/mlcommons/training/blob/e237206991d10449d9675d95606459a3cb6c21ad/image_classification/tensorflow2/lars_util.py
class PolynomialDecayWithWarmup(LR_Scheduler):
def __init__(self, optimizer: Optimizer, initial_lr, end_lr, train_steps, warmup, power=2):
super().__init__(optimizer)
self.epoch_counter = self.epoch_counter.cast(dtypes.float32)
assert train_steps > 0 and warmup > 0
self.warmup = min(warmup, train_steps)
self.initial_lr, self.end_lr, self.epochs, self.power = initial_lr, end_lr, train_steps, power
# set lr for first warmup step
self.optimizer.lr.assign(self.get_lr()).realize()
def get_lr(self):
# LR is 0 on the first step, matching the reference.
warmup_lr = (self.epoch_counter * (1.0 / self.warmup)) * self.initial_lr
x = (1 - (self.epoch_counter - self.warmup) / (self.epochs - self.warmup + 1))
return (self.epoch_counter <= self.warmup).where(warmup_lr, (self.initial_lr - self.end_lr) * x ** self.power + self.end_lr).cast(self.optimizer.lr.dtype)
class CosineAnnealingLRWithWarmup(LR_Scheduler):
def __init__(self, optimizer:Optimizer, base_lr, end_lr, warmup_steps:int, decay_steps:int):
assert warmup_steps > 0 and decay_steps > 0
super().__init__(optimizer)
self.base_lr = base_lr
self.end_lr = end_lr
self.warmup_steps = warmup_steps
self.decay_steps = decay_steps
# set lr for first warmup step
self.optimizer.lr.assign(self.get_lr()).realize()
def get_lr(self):
warmup_lr = ((self.epoch_counter+1) / self.warmup_steps) * self.base_lr
decay_lr = self.end_lr + 0.5 * (self.base_lr-self.end_lr) * (1 + (((self.epoch_counter+1-self.warmup_steps)/self.decay_steps) * math.pi).cos())
return (self.epoch_counter < self.warmup_steps).where(warmup_lr, decay_lr).cast(self.optimizer.lr.dtype)
# Reference: https://github.com/mlcommons/training/blob/64b14a9abc74e08779a175abca7d291f8c957632/stable_diffusion/ldm/lr_scheduler.py, Lines 36-97
class LambdaLinearScheduler:
def __init__(self, warm_up_steps:int, f_min:float, f_max:float, f_start:float, cycle_lengths:int):
self.lr_warm_up_steps, self.f_min, self.f_max, self.f_start, self.cycle_lengths = warm_up_steps, f_min, f_max, f_start, cycle_lengths
def schedule(self, n:Tensor) -> Tensor:
warm_up = (n < self.lr_warm_up_steps)
f_warm_up = (self.f_max - self.f_start) / self.lr_warm_up_steps * n + self.f_start
return warm_up.where(f_warm_up, self.f_min + (self.f_max - self.f_min) * (self.cycle_lengths - n) / (self.cycle_lengths))
# based on torch.optim.lr_scheduler.LambdaLR
class LambdaLR(LR_Scheduler):
def __init__(self, optimizer:Optimizer, base_lr:Tensor, lr_lambda:Callable):
super().__init__(optimizer)
self.base_lr, self.lr_lambda = base_lr, lr_lambda
self.step()
def get_lr(self):
return self.base_lr * self.lr_lambda(self.epoch_counter - 1)

69
artifacts/package_sources/tinygrad/examples/mlperf/metrics.py Normal file
View File

@@ -0,0 +1,69 @@
import re, string
from collections import Counter
from tinygrad import Tensor
def levenshtein(a, b):
n, m = len(a), len(b)
if n > m:
a, b, n, m = b, a, m, n
current = list(range(n + 1))
for i in range(1, m + 1):
previous, current = current, [i] + [0] * n
for j in range(1, n + 1):
add, delete = previous[j] + 1, current[j - 1] + 1
change = previous[j - 1]
if a[j - 1] != b[i - 1]:
change = change + 1
current[j] = min(add, delete, change)
return current[n]
def word_error_rate(x, y):
scores = words = 0
for h, r in zip(x, y):
h_list = h.split()
r_list = r.split()
words += len(r_list)
scores += levenshtein(h_list, r_list)
return float(scores) / words, float(scores), words
def one_hot(x):
return x.one_hot(3).squeeze(1).permute(0, 4, 1, 2, 3)
def dice_score(prediction, target, channel_axis=1, smooth_nr=1e-6, smooth_dr=1e-6, argmax=True, to_one_hot_x=True):
channel_axis, reduce_axis = 1, tuple(range(2, len(prediction.shape)))
if argmax: prediction = prediction.argmax(axis=channel_axis)
else: prediction = prediction.softmax(axis=channel_axis)
if to_one_hot_x: prediction = one_hot(prediction)
target = one_hot(target)
prediction, target = prediction[:, 1:], target[:, 1:]
assert prediction.shape == target.shape, f"prediction ({prediction.shape}) and target ({target.shape}) shapes do not match"
intersection = (prediction * target).sum(axis=reduce_axis)
target_sum = target.sum(axis=reduce_axis)
prediction_sum = prediction.sum(axis=reduce_axis)
result = (2.0 * intersection + smooth_nr) / (target_sum + prediction_sum + smooth_dr)
return result
def normalize_string(s):
s = "".join(c for c in s.lower() if c not in string.punctuation)
s = re.sub(r'\b(a|an|the)\b', ' ', s)
return " ".join(s.split())
def f1_score(x, y):
xt = normalize_string(x).split()
yt = normalize_string(y).split()
ct = Counter(xt) & Counter(yt)
if (ns := sum(ct.values())) == 0:
return 0.0
p = ns / len(xt)
r = ns / len(yt)
return 2 * p * r / (p + r)
def log_perplexity(logit:Tensor, target:Tensor, ignore_index:int|None=None):
# logit has shape (n_samples, seq_len, vocab_size), target has shape (n_samples, seq_len)
assert logit.ndim == 3, logit.ndim
assert target.ndim == 2, target.ndim
assert logit.shape[:2] == target.shape, f"{logit.shape[:2]=}, {target.shape=}"
log_prob = logit.log_softmax(axis=-1)
return log_prob.transpose(1, 2).nll_loss(target, ignore_index=ignore_index)

507
artifacts/package_sources/tinygrad/examples/mlperf/model_eval.py Normal file
View File

@@ -0,0 +1,507 @@
import time, math, os
start = time.perf_counter()
from pathlib import Path
import numpy as np
from tinygrad import Tensor, Device, dtypes, GlobalCounters, TinyJit
from tinygrad.nn.state import get_parameters, load_state_dict, safe_load
from tinygrad.helpers import getenv, Context, prod
from extra.bench_log import BenchEvent, WallTimeEvent
def tlog(x): print(f"{x:25s} @ {time.perf_counter()-start:5.2f}s")
def eval_resnet():
with WallTimeEvent(BenchEvent.FULL):
# Resnet50-v1.5
from extra.models.resnet import ResNet50
tlog("imports")
GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 6))]
for x in GPUS: Device[x]
tlog("got devices") # NOTE: this is faster with rocm-smi running
class ResnetRunner:
def __init__(self, device=None):
self.mdl = ResNet50()
for x in get_parameters(self.mdl) if device else []: x.to_(device)
if (fn:=getenv("RESNET_MODEL", "")): load_state_dict(self.mdl, safe_load(fn))
else: self.mdl.load_from_pretrained()
self.input_mean = Tensor([0.485, 0.456, 0.406], device=device).reshape(1, -1, 1, 1)
self.input_std = Tensor([0.229, 0.224, 0.225], device=device).reshape(1, -1, 1, 1)
def __call__(self, x:Tensor) -> Tensor:
x = x.permute([0,3,1,2]).cast(dtypes.float32) / 255.0
x -= self.input_mean
x /= self.input_std
return self.mdl(x).log_softmax().argmax(axis=1).realize()
mdl = TinyJit(ResnetRunner(GPUS))
tlog("loaded models")
# evaluation on the mlperf classes of the validation set from imagenet
from examples.mlperf.dataloader import batch_load_resnet
iterator = batch_load_resnet(getenv("BS", 128*6), val=getenv("VAL", 1), shuffle=False, pad_first_batch=True)
def data_get():
x,y,cookie = next(iterator)
return x.shard(GPUS, axis=0).realize(), y, cookie
n,d = 0,0
proc = data_get()
tlog("loaded initial data")
st = time.perf_counter()
while proc is not None:
GlobalCounters.reset()
proc = (mdl(proc[0]), proc[1], proc[2]) # this frees the images
run = time.perf_counter()
# load the next data here
try: next_proc = data_get()
except StopIteration: next_proc = None
nd = time.perf_counter()
y = np.array(proc[1])
proc = (proc[0].numpy() == y) & (y != -1) # this realizes the models and frees the cookies
n += proc.sum()
d += (y != -1).sum()
et = time.perf_counter()
tlog(f"****** {n:5d}/{d:5d} {n*100.0/d:.2f}% -- {(run-st)*1000:7.2f} ms to enqueue, {(et-run)*1000:7.2f} ms to realize ({(nd-run)*1000:7.2f} ms fetching). {(len(proc))/(et-st):8.2f} examples/sec. {GlobalCounters.global_ops*1e-12/(et-st):5.2f} TFLOPS")
st = et
proc, next_proc = next_proc, None
tlog("done")
def eval_unet3d():
# UNet3D
from extra.models.unet3d import UNet3D
from extra.datasets.kits19 import iterate, sliding_window_inference, get_val_files
from examples.mlperf.metrics import dice_score
mdl = UNet3D()
mdl.load_from_pretrained()
s = 0
st = time.perf_counter()
for i, (image, label) in enumerate(iterate(get_val_files()), start=1):
mt = time.perf_counter()
pred, label = sliding_window_inference(mdl, image, label)
et = time.perf_counter()
print(f"{(mt-st)*1000:.2f} ms loading data, {(et-mt)*1000:.2f} ms to run model")
s += dice_score(Tensor(pred), Tensor(label)).mean().item()
print(f"****** {s:.2f}/{i} {s/i:.5f} Mean DICE score")
st = time.perf_counter()
def eval_retinanet():
# RetinaNet with ResNeXt50_32X4D
from examples.mlperf.dataloader import batch_load_retinanet
from extra.datasets.openimages import normalize, download_dataset, BASEDIR
from extra.models.resnet import ResNeXt50_32X4D
from extra.models.retinanet import RetinaNet
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
from contextlib import redirect_stdout
tlog("imports")
mdl = RetinaNet(ResNeXt50_32X4D())
mdl.load_from_pretrained()
tlog("loaded models")
coco = COCO(download_dataset(base_dir:=getenv("BASEDIR", BASEDIR), 'validation'))
coco_eval = COCOeval(coco, iouType="bbox")
coco_evalimgs, evaluated_imgs, ncats, narea = [], [], len(coco_eval.params.catIds), len(coco_eval.params.areaRng)
tlog("loaded dataset")
iterator = batch_load_retinanet(coco, True, Path(base_dir), getenv("BS", 8), shuffle=False)
def data_get():
x, img_ids, img_sizes, cookie = next(iterator)
return x.to(Device.DEFAULT).realize(), img_ids, img_sizes, cookie
n = 0
proc = data_get()
tlog("loaded initial data")
st = time.perf_counter()
while proc is not None:
GlobalCounters.reset()
proc = (mdl(normalize(proc[0])), proc[1], proc[2], proc[3])
run = time.perf_counter()
# load the next data here
try: next_proc = data_get()
except StopIteration: next_proc = None
nd = time.perf_counter()
predictions, img_ids = mdl.postprocess_detections(proc[0].numpy(), orig_image_sizes=proc[2]), proc[1]
pd = time.perf_counter()
coco_results = [{"image_id": img_ids[i], "category_id": label, "bbox": box.tolist(), "score": score}
for i, prediction in enumerate(predictions) for box, score, label in zip(*prediction.values())]
with redirect_stdout(None):
coco_eval.cocoDt = coco.loadRes(coco_results)
coco_eval.params.imgIds = img_ids
coco_eval.evaluate()
evaluated_imgs.extend(img_ids)
coco_evalimgs.append(np.array(coco_eval.evalImgs).reshape(ncats, narea, len(img_ids)))
n += len(proc[0])
et = time.perf_counter()
tlog(f"****** {(run-st)*1000:7.2f} ms to enqueue, {(et-run)*1000:7.2f} ms to realize ({(nd-run)*1000:7.2f} ms fetching, {(pd-run)*1000:4.2f} ms postprocess_detections). {(len(proc))/(et-st):8.2f} examples/sec. {GlobalCounters.global_ops*1e-12/(et-st):5.2f} TFLOPS")
st = et
proc, next_proc = next_proc, None
coco_eval.params.imgIds = evaluated_imgs
coco_eval._paramsEval.imgIds = evaluated_imgs
coco_eval.evalImgs = list(np.concatenate(coco_evalimgs, -1).flatten())
coco_eval.accumulate()
coco_eval.summarize()
tlog("done")
def eval_rnnt():
# RNN-T
from extra.models.rnnt import RNNT
mdl = RNNT()
mdl.load_from_pretrained()
from extra.datasets.librispeech import iterate
from examples.mlperf.metrics import word_error_rate
LABELS = [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'"]
c = 0
scores = 0
words = 0
st = time.perf_counter()
for X, Y in iterate():
mt = time.perf_counter()
tt = mdl.decode(Tensor(X[0]), Tensor([X[1]]))
et = time.perf_counter()
print(f"{(mt-st)*1000:.2f} ms loading data, {(et-mt)*1000:.2f} ms to run model")
for n, t in enumerate(tt):
tnp = np.array(t)
_, scores_, words_ = word_error_rate(["".join([LABELS[int(tnp[i])] for i in range(tnp.shape[0])])], [Y[n]])
scores += scores_
words += words_
c += len(tt)
print(f"WER: {scores/words}, {words} words, raw scores: {scores}, c: {c}")
st = time.perf_counter()
def eval_bert():
# Bert-QA
from extra.models.bert import BertForQuestionAnswering
mdl = BertForQuestionAnswering()
mdl.load_from_pretrained()
@TinyJit
def run(input_ids, input_mask, segment_ids):
return mdl(input_ids, input_mask, segment_ids).realize()
from extra.datasets.squad import iterate
from examples.mlperf.helpers import get_bert_qa_prediction
from examples.mlperf.metrics import f1_score
from transformers import BertTokenizer
tokenizer = BertTokenizer(str(Path(__file__).parents[2] / "extra/weights/bert_vocab.txt"))
c = 0
f1 = 0.0
st = time.perf_counter()
for X, Y in iterate(tokenizer):
mt = time.perf_counter()
outs = []
for x in X:
outs.append(run(Tensor(x["input_ids"]), Tensor(x["input_mask"]), Tensor(x["segment_ids"])).numpy())
et = time.perf_counter()
print(f"{(mt-st)*1000:.2f} ms loading data, {(et-mt)*1000:.2f} ms to run model over {len(X)} features")
pred = get_bert_qa_prediction(X, Y, outs)
print(f"pred: {pred}\nans: {Y['answers']}")
f1 += max([f1_score(pred, ans) for ans in Y["answers"]])
c += 1
print(f"f1: {f1/c}, raw: {f1}, c: {c}\n")
st = time.perf_counter()
def eval_llama3():
from extra.models.llama import Transformer
from examples.llama3 import MODEL_PARAMS, load, convert_from_huggingface
from tinygrad.helpers import tqdm
BASEDIR = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
BS = getenv("BS", 4)
SMALL = getenv("SMALL", 0)
SEQLEN = getenv("SEQLEN", 8192)
MODEL_PATH = Path(getenv("MODEL_PATH", "/raid/weights/llama31_8b/"))
params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
params = params | {"vocab_size": 32000} if not SMALL else params
if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
# load weights
weights = load(str(MODEL_PATH / "model.safetensors.index.json"))
if "model.embed_tokens.weight" in weights:
print("converting from huggingface format")
weights = convert_from_huggingface(weights, params["n_layers"], params["n_heads"], params["n_kv_heads"])
load_state_dict(model, weights, strict=False, consume=True)
@TinyJit
def eval_step(model, tokens):
logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
return loss.flatten().float()
from examples.mlperf.dataloader import get_llama3_dataset, iterate_llama3_dataset
eval_dataset = get_llama3_dataset(5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
iter = iterate_llama3_dataset(eval_dataset, BS)
losses = []
for tokens in tqdm(iter, total=5760//BS):
GlobalCounters.reset()
losses += eval_step(model, tokens).tolist()
tqdm.write(f"loss: {np.mean(losses)}")
log_perplexity = np.mean(losses)
print(f"Log Perplexity: {log_perplexity}")
# NOTE: BEAM hangs on 8xmi300x with DECODE_BS=384 in final realize below; function is declared here for external testing
@TinyJit
def vae_decode(x:Tensor, vae, disable_beam=False) -> Tensor:
from examples.stable_diffusion import AutoencoderKL
assert isinstance(vae, AutoencoderKL)
x = vae.post_quant_conv(1./0.18215 * x)
x = vae.decoder.conv_in(x)
x = vae.decoder.mid(x)
for i, l in enumerate(vae.decoder.up[::-1]):
print("decode", x.shape)
for b in l['block']: x = b(x)
if 'upsample' in l:
bs,c,py,px = x.shape
x = x.reshape(bs, c, py, 1, px, 1).expand(bs, c, py, 2, px, 2).reshape(bs, c, py*2, px*2)
x = l['upsample']['conv'](x)
if i == len(vae.decoder.up) - 1 and disable_beam:
with Context(BEAM=0): x.realize()
else: x.realize()
x = vae.decoder.conv_out(vae.decoder.norm_out(x).swish())
x = ((x + 1.0) / 2.0).clip(0.0, 1.0)
return x
def eval_stable_diffusion():
import csv, PIL, sys
from tqdm import tqdm
from examples.mlperf.initializers import init_stable_diffusion, gelu_erf
from examples.stable_diffusion import AutoencoderKL
from extra.models.unet import UNetModel
from tinygrad.nn.state import load_state_dict, torch_load
from tinygrad.helpers import BEAM
from extra.models import clip
from extra.models.clip import FrozenOpenClipEmbedder
from extra.models.clip import OpenClipEncoder
from extra.models.inception import FidInceptionV3
config = {}
GPUS = config["GPUS"] = [f"{Device.DEFAULT}:{i}" for i in range(getenv("GPUS", 1))]
for x in GPUS: Device[x]
print(f"running eval on {GPUS}")
seed = config["seed"] = getenv("SEED", 12345)
CKPTDIR = config["CKPTDIR"] = Path(getenv("CKPTDIR", "./checkpoints"))
DATADIR = config["DATADIR"] = Path(getenv("DATADIR", "./datasets"))
CONTEXT_BS = config["CONTEXT_BS"] = getenv("CONTEXT_BS", 1 * len(GPUS))
DENOISE_BS = config["DENOISE_BS"] = getenv("DENOISE_BS", 1 * len(GPUS))
DECODE_BS = config["DECODE_BS"] = getenv("DECODE_BS", 1 * len(GPUS))
INCEPTION_BS = config["INCEPTION_BS"] = getenv("INCEPTION_BS", 1 * len(GPUS))
CLIP_BS = config["CLIP_BS"] = getenv("CLIP_BS", 1 * len(GPUS))
EVAL_CKPT_DIR = config["EVAL_CKPT_DIR"] = getenv("EVAL_CKPT_DIR", "")
STOP_IF_CONVERGED = config["STOP_IF_CONVERGED"] = getenv("STOP_IF_CONVERGED", 0)
if (WANDB := getenv("WANDB", "")):
import wandb
wandb.init(config=config, project="MLPerf-Stable-Diffusion")
assert EVAL_CKPT_DIR != "", "provide a directory with checkpoints to be evaluated"
print(f"running eval on checkpoints in {EVAL_CKPT_DIR}\nSEED={seed}")
eval_queue:list[tuple[int, Path]] = []
for p in Path(EVAL_CKPT_DIR).iterdir():
if p.name.endswith(".safetensors"):
ckpt_iteration = p.name.split(".safetensors")[0]
assert ckpt_iteration.isdigit(), f"invalid checkpoint name: {p.name}, expected <digits>.safetensors"
eval_queue.append((int(ckpt_iteration), p))
assert len(eval_queue), f'no files ending with ".safetensors" were found in {EVAL_CKPT_DIR}'
print(sorted(eval_queue, reverse=True))
Tensor.manual_seed(seed) # seed for weight initialization
model, unet, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod = init_stable_diffusion("v2-mlperf-eval", CKPTDIR / "sd" / "512-base-ema.ckpt", GPUS)
# load prompts for generating images for validation; 2 MB of data total
with open(DATADIR / "coco2014" / "val2014_30k.tsv") as f:
reader = csv.DictReader(f, delimiter="\t")
eval_inputs:list[dict] = [{"image_id": int(row["image_id"]), "id": int(row["id"]), "caption": row["caption"]} for row in reader]
assert len(eval_inputs) == 30_000
# NOTE: the clip weights are the same between model.cond_stage_model and clip_encoder
eval_timesteps = list(reversed(range(1, 1000, 20)))
with Context(DEV="CPU"):
# The choice of alphas_prev[0] = alphas_cumprod[0] seems arbitrary, but it's how the mlperf ref does it:
# alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
eval_alphas_prev = model.alphas_cumprod[0:1].cat(model.alphas_cumprod[list(range(1, 1000, 20))[:-1]]).to(GPUS).realize()
inception = FidInceptionV3().load_from_pretrained(CKPTDIR / "inception" / "pt_inception-2015-12-05-6726825d.pth")
vision_cfg = {'width': 1280, 'layers': 32, 'd_head': 80, 'image_size': 224, 'patch_size': 14}
text_cfg = {'width': 1024, 'n_heads': 16, 'layers': 24, 'vocab_size': 49408, 'ctx_length': 77}
clip.gelu = gelu_erf
clip_encoder = OpenClipEncoder(1024, text_cfg, vision_cfg)
loaded = torch_load(CKPTDIR / "clip" / "open_clip_pytorch_model.bin")
loaded.update({"attn_mask": clip_encoder.attn_mask, "mean": clip_encoder.mean, "std": clip_encoder.std})
load_state_dict(clip_encoder, loaded)
@TinyJit
def denoise_step(x:Tensor, x_x:Tensor, t_t:Tensor, uc_c:Tensor, sqrt_alphas_cumprod_t:Tensor, sqrt_one_minus_alphas_cumprod_t:Tensor,
alpha_prev:Tensor, unet:UNetModel, GPUS) -> Tensor:
out_uncond, out = unet(x_x, t_t, uc_c).to("CPU").reshape(-1, 2, 4, 64, 64).chunk(2, dim=1)
out_uncond = out_uncond.squeeze(1).shard(GPUS,axis=0)
out = out.squeeze(1).shard(GPUS,axis=0)
v_t = out_uncond + 8.0 * (out - out_uncond)
e_t = sqrt_alphas_cumprod_t * v_t + sqrt_one_minus_alphas_cumprod_t * x
pred_x0 = sqrt_alphas_cumprod_t * x - sqrt_one_minus_alphas_cumprod_t * v_t
dir_xt = (1. - alpha_prev).sqrt() * e_t
x_prev = alpha_prev.sqrt() * pred_x0 + dir_xt
return x_prev.realize()
def shard_tensor(t:Tensor) -> Tensor: return t.shard(GPUS, axis=0) if len(GPUS) > 1 else t.to(GPUS[0])
def get_batch(whole:Tensor, i:int, bs:int) -> tuple[Tensor, int]:
batch = whole[i: i + bs].to("CPU")
if (unpadded_bs:=batch.shape[0]) < bs:
batch = batch.cat(batch[-1:].expand(bs - unpadded_bs, *batch[-1].shape))
return batch, unpadded_bs
@Context(TRAINING=0)
def eval_unet(eval_inputs:list[dict], unet:UNetModel, cond_stage:FrozenOpenClipEmbedder, first_stage:AutoencoderKL,
inception:FidInceptionV3, clip:OpenClipEncoder) -> tuple[float, float]:
# Eval is divided into 5 jits, one per model
# It doesn't make sense to merge these jits, e.g. unet repeats 50 times in isolation; images fork to separate inception/clip
# We're generating and scoring 30,000 images per eval, and all the data can flow through one jit at a time
# To maximize throughput for each jit, we have only one model/jit on the GPU at a time, and pool outputs from each jit off-GPU
for model in (unet, first_stage, inception, clip):
Tensor.realize(*[p.to_("CPU") for p in get_parameters(model)])
uc_written = False
models = (cond_stage, unet, first_stage, inception, clip)
jits = (jit_context:=TinyJit(cond_stage.embed_tokens), denoise_step, vae_decode, jit_inception:=TinyJit(inception),
jit_clip:=TinyJit(clip.get_clip_score))
all_bs = (CONTEXT_BS, DENOISE_BS, DECODE_BS, INCEPTION_BS, CLIP_BS)
if (EVAL_SAMPLES:=getenv("EVAL_SAMPLES", 0)) and EVAL_SAMPLES > 0:
eval_inputs = eval_inputs[0:EVAL_SAMPLES]
output_shapes = [(ns:=len(eval_inputs),77), (ns,77,1024), (ns,4,64,64), (ns,3,512,512), (ns,2048), (ns,)]
# Writing progress to disk lets us resume eval if we crash
stages = ["tokens", "embeds", "latents", "imgs", "inception", "clip"]
disk_tensor_names, disk_tensor_shapes = stages + ["end", "uc"], output_shapes + [(6,), (1,77,1024)]
if not all(os.path.exists(f"{EVAL_CKPT_DIR}/{name}.bytes") for name in disk_tensor_names):
for name, shape in zip(disk_tensor_names, disk_tensor_shapes):
file = Path(f"{EVAL_CKPT_DIR}/{name}.bytes")
file.unlink(missing_ok=True)
with file.open("wb") as f: f.truncate(prod(shape) * 4)
progress = {name: Tensor.empty(*shape, device=f"disk:{EVAL_CKPT_DIR}/{name}.bytes", dtype=dtypes.int if name in {"tokens", "end"} else dtypes.float)
for name, shape in zip(disk_tensor_names, disk_tensor_shapes)}
def embed_tokens(tokens:Tensor) -> Tensor:
nonlocal uc_written
if not uc_written:
with Context(BEAM=0): progress["uc"].assign(cond_stage.embed_tokens(cond_stage.tokenize("").to(GPUS)).to("CPU").realize()).realize()
uc_written = True
return jit_context(shard_tensor(tokens))
def generate_latents(embeds:Tensor) -> Tensor:
uc_c = Tensor.stack(progress["uc"].to("CPU").expand(bs, 77, 1024), embeds, dim=1).reshape(-1, 77, 1024)
uc_c = shard_tensor(uc_c)
x = shard_tensor(Tensor.randn(bs,4,64,64))
for step_idx, timestep in enumerate(tqdm(eval_timesteps)):
reversed_idx = Tensor([50 - step_idx - 1], device=GPUS)
alpha_prev = eval_alphas_prev[reversed_idx]
ts = Tensor.full(bs, fill_value=timestep, dtype=dtypes.int, device="CPU")
ts_ts = shard_tensor(ts.cat(ts))
ts = shard_tensor(ts)
sqrt_alphas_cumprod_t = sqrt_alphas_cumprod[ts].reshape(bs, 1, 1, 1)
sqrt_one_minus_alphas_cumprod_t = sqrt_one_minus_alphas_cumprod[ts].reshape(bs, 1, 1, 1)
x_x = shard_tensor(Tensor.stack(x.to("CPU"), x.to("CPU"), dim=1).reshape(-1, 4, 64, 64))
x.assign(denoise_step(x, x_x, ts_ts, uc_c, sqrt_alphas_cumprod_t, sqrt_one_minus_alphas_cumprod_t, alpha_prev, unet, GPUS)).realize()
return x
def decode_latents(latents:Tensor) -> Tensor: return vae_decode(shard_tensor(latents), first_stage, disable_beam=True)
def generate_inception(imgs:Tensor) -> Tensor: return jit_inception(shard_tensor(imgs))[:,:,0,0]
def calc_clip_scores(batch:Tensor, batch_tokens:Tensor) -> Tensor:
# Tensor.interpolate does not yet support bicubic, so we use PIL
batch = (batch.to(GPUS[0]).permute(0,2,3,1) * 255).clip(0, 255).cast(dtypes.uint8).numpy()
batch = [np.array(PIL.Image.fromarray(batch[i]).resize((224,224), PIL.Image.BICUBIC)) for i in range(bs)]
batch = shard_tensor(Tensor(np.stack(batch, axis=0).transpose(0,3,1,2), device="CPU").realize())
batch = batch.cast(dtypes.float) / 255
batch = (batch - model.mean) / model.std
batch = jit_clip(shard_tensor(batch_tokens), batch)
return batch
callbacks = (embed_tokens, generate_latents, decode_latents, generate_inception, calc_clip_scores)
# save every forward pass output to disk; NOTE: this needs ~100 GB disk space because 30k images are large
def stage_progress(stage_idx:int) -> int: return progress["end"].to("CPU")[stage_idx].item()
if stage_progress(0) < len(eval_inputs):
tokens = []
for i in tqdm(range(0, len(eval_inputs), CONTEXT_BS)):
subset = [cond_stage.tokenize(row["caption"], device="CPU") for row in eval_inputs[i: i+CONTEXT_BS]]
tokens.append(Tensor.cat(*subset, dim=0).realize())
progress["tokens"].assign(Tensor.cat(*tokens, dim=0).realize()).realize()
progress["end"][0:1].assign(Tensor([len(eval_inputs)], dtype=dtypes.int)).realize()
prev_stage = "tokens"
tokens = progress["tokens"]
# wrapper code for every model
for stage_idx, model, jit, bs, callback in zip(range(1,6), models, jits, all_bs, callbacks):
stage = stages[stage_idx]
if stage_progress(stage_idx) >= len(eval_inputs):
prev_stage = stage
continue # use cache
t0 = time.perf_counter()
print(f"starting eval with model: {model}")
if stage_idx == 1: inputs = tokens
elif stage_idx == 5: inputs = progress["imgs"]
else: inputs = progress[prev_stage]
Tensor.realize(*[p.to_(GPUS) for p in get_parameters(model)])
for batch_idx in tqdm(range(stage_progress(stage_idx), inputs.shape[0], bs)):
t1 = time.perf_counter()
batch, unpadded_bs = get_batch(inputs, batch_idx, bs)
if isinstance(model, OpenClipEncoder): batch = callback(batch, get_batch(tokens, batch_idx, bs)[0].realize())
else: batch = callback(batch)
# to(GPUS[0]) is necessary for this to work, without that the result is still on GPUS, probably due to a bug
batch = batch.to(GPUS[0]).to("CPU")[0:unpadded_bs].realize()
progress[stage][batch_idx: batch_idx + bs].assign(batch).realize()
# keep track of what our last output was, so we can resume from there if we crash in this loop
progress["end"][stage_idx: stage_idx + 1].assign(Tensor([batch_idx + bs], dtype=dtypes.int)).realize()
print(f"model: {model}, batch_idx: {batch_idx}, elapsed: {(time.perf_counter() - t1):.2f}")
del batch
jit.reset()
Tensor.realize(*[p.to_("CPU") for p in get_parameters(model)])
print(f"done with model: {model}, elapsed: {(time.perf_counter() - t0):.2f}")
prev_stage = stage
inception_stats_fn = str(DATADIR / "coco2014" / "val2014_30k_stats.npz")
fid_score = inception.compute_score(progress["inception"].to("CPU"), inception_stats_fn)
clip_score = progress["clip"].to(GPUS[0]).mean().item()
for name in disk_tensor_names:
Path(f"{EVAL_CKPT_DIR}/{name}.bytes").unlink(missing_ok=True)
if EVAL_SAMPLES and BEAM:
print("BEAM COMPLETE", flush=True) # allows wrapper script to detect BEAM search completion and retry if it failed
sys.exit() # Don't eval additional models; we don't care about clip/fid scores when running BEAM on eval sample subset
return clip_score, fid_score
# evaluate checkpoints in reverse chronological order
for ckpt_iteration, p in sorted(eval_queue, reverse=True):
unet_ckpt = safe_load(p)
load_state_dict(unet, unet_ckpt)
clip_score, fid_score = eval_unet(eval_inputs, unet, model.cond_stage_model, model.first_stage_model, inception, clip_encoder)
converged = True if clip_score >= 0.15 and fid_score <= 90 else False
print(f"eval results for {EVAL_CKPT_DIR}/{p.name}: clip={clip_score}, fid={fid_score}, converged={converged}")
if WANDB:
wandb.log({"eval/ckpt_iteration": ckpt_iteration, "eval/clip_score": clip_score, "eval/fid_score": fid_score})
if converged and STOP_IF_CONVERGED:
print(f"Convergence detected, exiting early before evaluating other checkpoints due to STOP_IF_CONVERGED={STOP_IF_CONVERGED}")
sys.exit()
# for testing
return clip_score, fid_score, ckpt_iteration
if __name__ == "__main__":
# inference only
models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert").split(",")
with Context(TRAINING=0):
for m in models:
nm = f"eval_{m}"
if nm in globals():
print(f"eval {m}")
globals()[nm]()

68
artifacts/package_sources/tinygrad/examples/mlperf/model_spec.py Normal file
View File

@@ -0,0 +1,68 @@
# load each model here, quick benchmark
from tinygrad import Tensor, GlobalCounters
from tinygrad.helpers import getenv, Context
import numpy as np
def test_model(model, *inputs):
GlobalCounters.reset()
out = model(*inputs)
if isinstance(out, Tensor): out = out.numpy()
# TODO: return event future to still get the time_sum_s without DEBUG=2
print(f"{GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.time_sum_s*1000:.2f} ms")
def spec_resnet():
# Resnet50-v1.5
from extra.models.resnet import ResNet50
mdl = ResNet50()
img = Tensor.randn(1, 3, 224, 224)
test_model(mdl, img)
def spec_retinanet():
# Retinanet with ResNet backbone
from extra.models.resnet import ResNet50
from extra.models.retinanet import RetinaNet
mdl = RetinaNet(ResNet50(), num_classes=91, num_anchors=9)
img = Tensor.randn(1, 3, 224, 224)
test_model(mdl, img)
def spec_unet3d():
# 3D UNET
from extra.models.unet3d import UNet3D
mdl = UNet3D()
#mdl.load_from_pretrained()
img = Tensor.randn(1, 1, 128, 128, 128)
test_model(mdl, img)
def spec_rnnt():
from extra.models.rnnt import RNNT
mdl = RNNT()
#mdl.load_from_pretrained()
x = Tensor.randn(220, 1, 240)
y = Tensor.randn(1, 220)
test_model(mdl, x, y)
def spec_bert():
from extra.models.bert import BertForQuestionAnswering
mdl = BertForQuestionAnswering()
#mdl.load_from_pretrained()
x = Tensor.randn(1, 384)
am = Tensor.randn(1, 384)
tt = Tensor(np.random.randint(0, 2, (1, 384)).astype(np.float32))
test_model(mdl, x, am, tt)
def spec_mrcnn():
from extra.models.mask_rcnn import MaskRCNN, ResNet
mdl = MaskRCNN(ResNet(50, num_classes=None, stride_in_1x1=True))
#mdl.load_from_pretrained()
x = Tensor.randn(3, 224, 224)
test_model(mdl, [x])
if __name__ == "__main__":
# inference only for now
with Context(TRAINING=0):
for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,mrcnn").split(","):
nm = f"spec_{m}"
if nm in globals():
print(f"testing {m}")
globals()[nm]()

2098
artifacts/package_sources/tinygrad/examples/mlperf/model_train.py Normal file
View File

File diff suppressed because it is too large Load Diff

453
artifacts/package_sources/tinygrad/examples/mlperf/models/flat_llama.py Normal file
View File

@@ -0,0 +1,453 @@
import math, os
if __name__ == "__main__":
os.environ["DEFAULT_FLOAT"] = "bfloat16"
os.environ["OPTIM_DTYPE"] = "bfloat16"
if "DEV" not in os.environ: os.environ["DEV"] = "NULL::gfx950"
# CDNA
os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
os.environ["ALL2ALL"] = "1"
os.environ["USE_ATOMICS"] = "1"
if "HK_FLASH_ATTENTION" not in os.environ:
os.environ["HK_FLASH_ATTENTION"] = "1"
if "ASM_GEMM" not in os.environ:
os.environ["ASM_GEMM"] = "1"
from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker, round_up
from tinygrad.uop.ops import Ops, UOp
from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
from extra.llama_kernels.rmsnorm import rmsnorm
from extra.llama_kernels import FP8_MAX, local_abs_max
ASM_GEMM = getenv("ASM_GEMM", 0)
FUSED_INPUT_QUANTIZE = getenv("FUSED_INPUT_QUANTIZE", 0)
FUSED_ADD_NORM_MUL_QUANTIZE = getenv("FUSED_ADD_NORM_MUL_QUANTIZE", 0)
FUSED_SILU_W13 = getenv("FUSED_SILU_W13", 0)
SPLIT_W13 = getenv("SPLIT_W13", 0)
COLUMNWISE_WEIGHT_SCALE = getenv("COLUMNWISE_WEIGHT_SCALE", 0)
MXFP8 = getenv("MXFP8", 0)
MXFP4 = getenv("MXFP4", 0)
FP8_DTYPE = dtypes.fp8e4m3
FP8_GRAD_DTYPE = dtypes.fp8e5m2
def quantize_fp8(x:Tensor, amax_state:Tensor|None=None):
new_amax = (local_abs_max(x) if isinstance(x.device, tuple) else x.abs().max()).detach().cast(dtypes.float32)
scale = FP8_MAX / ((amax_state if amax_state is not None else new_amax) + 1e-8)
x_scaled = x * scale
x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach()) # STE
return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal(), new_amax
def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None,
x_fp8:Tensor|None=None, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None, x_prequant_mx:tuple|None=None,
next_amax_x:Tensor|None=None) -> tuple[Tensor,...]:
if not fp8:
if ASM_GEMM:
from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T),)
return (x @ w.T,)
if MXFP4:
assert x is not None, "MXFP4 matmul requires an unquantized input"
from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm
if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T, mxfp4=True),)
return (x @ w.T,)
assert w_inv_scale is not None, "fp8 matmul requires w_inv_scale (weights must be stored in fp8 with per-tensor scale)"
if MXFP8:
from extra.gemm.cdna_asm_gemm import asm_gemm, quantize_mxfp8, mx_pack, can_use_asm_gemm, _mx_block_scale
if x_prequant_mx is not None: x_q, x_e8, x_si = x_prequant_mx # fused producer already quantized (2d)
else: x_q, x_e8, x_si = quantize_mxfp8(x.reshape(-1, x.shape[-1]))
l_shape = x.shape[:-1] if x is not None else x_q.shape[:-1]
if can_use_asm_gemm(x_q, w.T):
out = asm_gemm(x_q, w.T, mx=True, mx_scales=(x_si, x_e8, mx_pack(w_inv_scale), w_inv_scale),
mx_w_stored=True).reshape(*l_shape, w.shape[0])
else:
x_phys = (x_q.cast(dtypes.bfloat16) * _mx_block_scale(x_e8)).reshape(*l_shape, x_q.shape[-1])
out = x_phys @ (w.cast(dtypes.bfloat16) * _mx_block_scale(w_inv_scale)).T
return out, x_q
if x_fp8 is None:
if FUSED_INPUT_QUANTIZE:
from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
x_fp8, _ = quantize_fp8_delayed(x, amax_x, next_amax_x, FP8_DTYPE)
else:
x_fp8, _, new_amax_x = quantize_fp8(x, amax_state=amax_x)
next_amax_x.assign(new_amax_x)
if ASM_GEMM:
from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
if can_use_asm_gemm(x_fp8, w.T):
assert amax_x is not None
if COLUMNWISE_WEIGHT_SCALE:
out = asm_gemm(x_fp8, w.T, x_scale=amax_x, grad_amax_state=grad_amax_state,
next_grad_amax_state=next_grad_amax_state, w_post_scale=w_inv_scale)
else:
out = asm_gemm(x_fp8, w.T, x_scale=amax_x, w_scale=w_inv_scale, grad_amax_state=grad_amax_state,
next_grad_amax_state=next_grad_amax_state)
return out, x_fp8
return (x_fp8.dot(w.T, dtype=dtypes.float) * ((amax_x.float() + 1e-8) / FP8_MAX) * w_inv_scale).cast(dtypes.bfloat16), x_fp8
def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor|None,
next_amax_x:Tensor|None, grad_amax_state:Tensor|None, next_grad_amax_state:Tensor|None):
if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8
x_fp8, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
grad_amax_state=grad_amax_state, next_grad_amax_state=next_grad_amax_state)
return out, x_normed, rrms, ret
x_normed, rrms = rmsnorm(x, eps)
out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state,
next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
return out, x_normed, rrms, ret
def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor|None,
next_amax_x:Tensor|None, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8
x_fp8, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
grad_amax_state=grad_amax_state, next_grad_amax_state=next_grad_amax_state)
return out, h, x_normed, rrms, ret
h = x + residual
x_normed, rrms = rmsnorm(h, eps)
out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state,
next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
return out, h, x_normed, rrms, ret
def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
amax_x2:Tensor|None, next_amax_x2:Tensor|None,
grad_amax_xw13:Tensor|None, next_grad_amax_xw13:Tensor|None,
grad_amax_xout:Tensor|None, next_grad_amax_xout:Tensor|None):
if FUSED_SILU_W13 and not MXFP4:
from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
x2_fp8 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
next_grad_amax_state=next_grad_amax_xw13, amax_out=next_amax_x2)
out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, amax_x=amax_x2,
grad_amax_state=grad_amax_xout, next_grad_amax_state=next_grad_amax_xout)
return out, ret
hidden = x_w13.shape[-1] // 2
x_w1, x_w3 = x_w13[..., :hidden], x_w13[..., hidden:]
out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout,
next_grad_amax_state=next_grad_amax_xout, next_amax_x=next_amax_x2)
return out, ret
class FlatTransformer:
def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
rope_theta:int=10000, max_context:int=1024):
self.vocab_size = vocab_size
self.n_layers = n_layers
self.n_heads = n_heads
self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
self.head_dim = dim // n_heads
self.n_rep = self.n_heads // self.n_kv_heads
self.hidden_dim = hidden_dim
scaled_std = 0.02 / math.sqrt(2 * n_layers)
# Attention
self.wqkv, s_qkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
self.wo, s_o = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std)
# FeedForward
if SPLIT_W13:
self.w1, s_1 = self.lin_per_layer(dim, hidden_dim)
self.w3, s_3 = self.lin_per_layer(dim, hidden_dim)
else:
self.w13, s_13 = self.lin_per_layer(dim, hidden_dim * 2)
self.w2, s_2 = self.lin_per_layer(hidden_dim, dim, std=scaled_std)
self.norm_eps = norm_eps
self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
self.ffn_norm = Tensor.ones(n_layers, dim).contiguous()
# output
self.norm = nn.RMSNorm(dim, norm_eps)
self.tok_embeddings = nn.Embedding(vocab_size, dim)
self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
self.output = Tensor.normal(1, vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).clone().is_param_(False)
def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().is_param_(False)
n_amax = 0 if MXFP4 else n_layers
names = ["xqkv", "xo", "x2"]
names += ["x1", "x3"] if SPLIT_W13 else ["x13"]
self._fp8_amax = {name: [_amax() for _ in range(n_amax)] for name in names}
self._fp8_next_amax = {name: [_amax() for _ in range(n_amax)] for name in names}
grad_names = ["xqkv", "xo", "xout"]
grad_names += ["xw1", "xw3"] if SPLIT_W13 else ["xw13"]
self._fp8_grad_amax = {name: [_amax() for _ in range(n_amax)] for name in grad_names}
self._fp8_next_grad_amax = {name: [_amax() for _ in range(n_amax)] for name in grad_names}
w_scales = [("wqkv", s_qkv), ("wo", s_o), ("w2", s_2)]
w_scales += [("w1", s_1), ("w3", s_3)] if SPLIT_W13 else [("w13", s_13)]
self._fp8_inv_scale = {name: (s if MXFP8 else s.float()).contiguous().is_param_(False) for name, s in w_scales}
self._fp8_next_inv_scale = {name: (s if MXFP8 else s.float()).contiguous().is_param_(False) for name, s in w_scales}
def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02, w:Tensor|None=None):
if w is None:
if getenv("ZEROS"): w = Tensor.zeros(self.n_layers, out_features, in_features)
else: w = Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std)
if MXFP8:
from extra.gemm.cdna_asm_gemm import quantize_mxfp8
w_q, w_e8, _ = quantize_mxfp8(w.reshape(self.n_layers * out_features, in_features))
return w_q.reshape(self.n_layers, out_features, in_features), w_e8.reshape(self.n_layers, out_features, in_features // 32)
if MXFP4:
# FP4 is produced dynamically so optimizer updates always start from the current BF16 weight.
return w.cast(dtypes.bfloat16), Tensor.ones(self.n_layers)
amax = (w.abs().max(axis=2) if COLUMNWISE_WEIGHT_SCALE else w.abs().flatten(1).max(1)).detach()
scale = FP8_MAX / (amax + 1e-8)
inv_scale = (amax + 1e-8) / FP8_MAX
scale_b = scale.reshape(self.n_layers, out_features, 1) if COLUMNWISE_WEIGHT_SCALE else scale.reshape(-1, 1, 1)
return (w * scale_b).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE), inv_scale
def attention(self, x:Tensor, freqs_cis:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
amax_xqkv:Tensor|None, amax_xo:Tensor|None, s_qkv:Tensor, s_o:Tensor,
next_amax_xqkv:Tensor|None, next_amax_xo:Tensor|None,
grad_amax_xqkv:Tensor|None, grad_amax_xo:Tensor|None,
next_grad_amax_xqkv:Tensor|None, next_grad_amax_xo:Tensor|None):
bsz, seqlen, _ = x.shape
saves = []
xqkv, x_normed, rrms, s = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
amax_x=amax_xqkv, grad_amax_state=grad_amax_xqkv,
next_grad_amax_state=next_grad_amax_xqkv, next_amax_x=next_amax_xqkv)
saves.extend([x_normed, rrms, *s, xqkv])
if getenv("HK_FLASH_ATTENTION"):
from extra.thunder.amd.fa import flash_attention, fused_qkv_rope
xq, xk, xv = fused_qkv_rope(xqkv, freqs_cis, self.n_heads, self.n_kv_heads, self.head_dim)
attn, *save = flash_attention(xq, xk, xv, is_causal=True, write_flat=True)
saves.extend(save)
else:
xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
attn = attn.reshape(bsz, seqlen, -1)
out, *s = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo,
next_grad_amax_state=next_grad_amax_xo, next_amax_x=next_amax_xo)
saves.extend([*s, out])
return out, saves
def feed_forward(self, x:Tensor, residual:Tensor, **kwargs):
saves = []
if SPLIT_W13:
h = x + residual
x_normed, rrms = rmsnorm(h, self.norm_eps)
saves.extend([x_normed, rrms])
inp = x_normed * kwargs["ffn_norm"]
x_w1, *s = matmul(inp, kwargs["w1"], amax_x=kwargs["amax_x1"], w_inv_scale=kwargs["s_1"],
grad_amax_state=kwargs["grad_amax_xw1"], next_grad_amax_state=kwargs["next_grad_amax_xw1"],
next_amax_x=kwargs["next_amax_x1"])
saves.extend([*s, x_w1])
x_w3, *s = matmul(inp, kwargs["w3"], amax_x=kwargs["amax_x3"], w_inv_scale=kwargs["s_3"],
grad_amax_state=kwargs["grad_amax_xw3"], next_grad_amax_state=kwargs["next_grad_amax_xw3"],
next_amax_x=kwargs["next_amax_x3"])
saves.extend([*s, x_w3])
if FUSED_SILU_W13 and MXFP8:
from extra.llama_kernels.fused_silu_mul_quantize_mxfp8 import fused_silu_mul_quantize_mxfp8
aq, ae8, asi = fused_silu_mul_quantize_mxfp8(x_w1.reshape(-1, x_w1.shape[-1]), x_w3.reshape(-1, x_w3.shape[-1]))
out, *s = matmul(None, kwargs["w2"], x_prequant_mx=(aq, ae8, asi), amax_x=kwargs["amax_x2"],
w_inv_scale=kwargs["s_2"], grad_amax_state=kwargs["grad_amax_xout"],
next_grad_amax_state=kwargs["next_grad_amax_xout"], next_amax_x=kwargs["next_amax_x2"])
out = out.reshape(*x_w1.shape[:-1], kwargs["w2"].shape[0])
else:
out, *s = matmul(x_w1.silu() * x_w3, kwargs["w2"], amax_x=kwargs["amax_x2"], w_inv_scale=kwargs["s_2"],
grad_amax_state=kwargs["grad_amax_xout"], next_grad_amax_state=kwargs["next_grad_amax_xout"],
next_amax_x=kwargs["next_amax_x2"])
saves.extend([*s, out])
else:
x_w13, h, x_normed, rrms, s = add_norm_quantize_matmul(x, residual, kwargs["ffn_norm"], kwargs["w13"], kwargs["s_13"],
self.norm_eps, amax_x=kwargs["amax_x13"],
next_amax_x=kwargs["next_amax_x13"],
grad_amax_state=kwargs["grad_amax_xw13"],
next_grad_amax_state=kwargs["next_grad_amax_xw13"])
saves.extend([x_normed, rrms, *s, x_w13])
out, s = silu_w13_quantize_matmul(x_w13, kwargs["w2"], kwargs["s_2"], amax_x2=kwargs["amax_x2"],
next_amax_x2=kwargs["next_amax_x2"],
grad_amax_xw13=kwargs["grad_amax_xw13"],
next_grad_amax_xw13=kwargs["next_grad_amax_xw13"],
grad_amax_xout=kwargs["grad_amax_xout"],
next_grad_amax_xout=kwargs["next_grad_amax_xout"])
saves.extend([*s, out])
return out, h, saves
@function(precompile=True, precompile_backward=True)
def run_layer(self, x:Tensor, freqs_cis:Tensor, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
attn, attn_saves = self.attention(x, freqs_cis, **attn_kwargs)
ffn, h, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
h = h + ffn
if save: return (h, *attn_saves, *ffn_saves)
else: return (h,)
def shard(self, device:tuple[str, ...], mp:bool=False):
from tinygrad.nn.state import get_parameters
if not mp:
for v in get_parameters(self): v.shard_(device, axis=None)
else:
# flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer
def _shard_fp8(name:str, axis:int, std:float=0.02):
w = getattr(self, name)
if MXFP8:
from extra.gemm.cdna_asm_gemm import quantize_mxfp8
w_bf16 = Tensor.empty(self.n_layers, w.shape[1], w.shape[2], dtype=dtypes.bfloat16).shard(device, axis=axis).randn_like() * std
w_q, w_e8, _ = quantize_mxfp8(w_bf16)
w.replace(w_q)
self._fp8_inv_scale[name].replace(w_e8.contiguous()).is_param_(False)
self._fp8_next_inv_scale[name].replace(w_e8.contiguous()).is_param_(False)
else:
w.shard_(device, axis=axis)
scale_axis = (1 if axis == 1 else None) if COLUMNWISE_WEIGHT_SCALE else None
self._fp8_inv_scale[name] = self._fp8_inv_scale[name].shard(device, axis=scale_axis).contiguous().is_param_(False)
self._fp8_next_inv_scale[name] = self._fp8_next_inv_scale[name].shard(device, axis=scale_axis).contiguous().is_param_(False)
Tensor.realize(w, self._fp8_inv_scale[name], self._fp8_next_inv_scale[name])
sstd = 0.02 / math.sqrt(2 * self.n_layers)
_shard_fp8("wqkv", 1) # (n_layers, out, dim) shard out
_shard_fp8("wo", 2, sstd) # (n_layers, dim, in) shard in
if SPLIT_W13:
_shard_fp8("w1", 1)
_shard_fp8("w3", 1)
else:
_shard_fp8("w13", 1) # (n_layers, hidden*2, dim) shard out
_shard_fp8("w2", 2, sstd) # (n_layers, dim, hidden) shard in
self.attention_norm.shard_(device, axis=None).realize()
self.ffn_norm.shard_(device, axis=None).realize()
self.norm.weight.shard_(device, axis=None).realize()
self.tok_embeddings.weight.shard_(device, axis=0).realize()
self.output.shard_(device, axis=1).realize()
self.freqs_cis.shard_(device, axis=None).realize()
for amax_dict in (self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax):
for name in amax_dict:
for i in range(len(amax_dict[name])):
amax_dict[name][i] = amax_dict[name][i].to(device).contiguous().is_param_(False)
def reset_amax(self):
for st in (self._fp8_next_amax, self._fp8_next_grad_amax):
for ts in st.values():
for t in ts: t.assign(0)
def update_amax(self):
for cur, nxt in ((self._fp8_amax, self._fp8_next_amax), (self._fp8_grad_amax, self._fp8_next_grad_amax)):
for name in cur:
for c, n in zip(cur[name], nxt[name]): c.assign(n)
def __call__(self, tokens:Tensor, save:bool=True):
h = self.tok_embeddings(tokens)
freqs_cis = self.freqs_cis.cast(h.dtype)
if not getenv("HK_FLASH_ATTENTION"): freqs_cis = freqs_cis[:, :tokens.shape[1], :, :, :]
a, na, ga, nga, s = self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax, self._fp8_inv_scale
def amax_kwargs(i:int, act_names:tuple[str, ...], grad_names:tuple[str, ...]) -> dict[str, Tensor|None]:
specs = (("amax_", a, act_names), ("next_amax_", na, act_names), ("grad_amax_", ga, grad_names), ("next_grad_amax_", nga, grad_names))
if MXFP4: return dict.fromkeys(f"{prefix}{name}" for prefix, _, names in specs for name in names)
return {f"{prefix}{name}":val[name][i] for prefix, val, names in specs for name in names}
for i in range(self.n_layers):
attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i], s_qkv=s["wqkv"][i], s_o=s["wo"][i],
**amax_kwargs(i, ("xqkv", "xo"), ("xqkv", "xo")))
ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i], s_2=s["w2"][i], **amax_kwargs(i, ("x2",), ("xout",)))
if SPLIT_W13:
ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], s_1=s["w1"][i], s_3=s["w3"][i], **amax_kwargs(i, ("x1", "x3"), ("xw1", "xw3")))
else:
ffn_kwargs.update(w13=self.w13[i], s_13=s["w13"][i], **amax_kwargs(i, ("x13",), ("xw13",)))
h, *_ = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)
logits = matmul(self.norm(h), self.output[0], fp8=False)[0]
return logits
def _get_pads(uop:UOp) -> list[UOp]:
if uop.op == Ops.ADD: return _get_pads(uop.src[0]) + _get_pads(uop.src[1])
return [uop]
def apply_grad(grad_buf:Tensor, new_grad:UOp):
pads = _get_pads(new_grad)
if len(pads) <= 1:
new_grad = new_grad.cast(grad_buf.dtype)
grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
return
cur = grad_buf.uop
for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
if pad.op == Ops.PAD:
grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)])
buf_slice = cur.shrink(grad_shrink)
cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
else:
cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
grad_buf.uop = cur
if __name__ == "__main__":
config = {}
BS = config["BS"] = getenv("BS", 16)
SEQLEN = config["SEQLEN"] = getenv("SEQLEN", 8192)
SMALL = config["SMALL"] = getenv("SMALL", 0)
from examples.llama3 import MODEL_PARAMS
model_params = MODEL_PARAMS[llama_size:=getenv("LLAMA3_SIZE", "8B")]["args"]
# vocab_size from mixtral tokenizer
if not SMALL: model_params |= {"vocab_size": 32000}
real_vocab_size = model_params['vocab_size']
if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params["n_layers"] = llama_layers
# pad vocab
if (MP := getenv("MP", 1)) > 1: model_params["vocab_size"] = round_up(model_params["vocab_size"], 256 * MP)
vocab_mask:Tensor = Tensor.arange(model_params["vocab_size"]).reshape(1, 1, -1) >= real_vocab_size
model = FlatTransformer(**model_params, max_context=SEQLEN)
state = nn.state.get_state_dict(model)
print("tensor count:", len(state))
# shard the model
from tinygrad import Device
is_dp = (DP := getenv("DP", 1)) > 1
is_mp = (MP := getenv("MP", 1)) > 1
is_sharding = is_dp or is_mp
device_count = max(DP, MP)
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
model.shard(device, is_mp)
if is_dp: vocab_mask.shard_(device, axis=None).realize()
if is_mp: vocab_mask.shard_(device, axis=2).realize()
# preallocate all the grad buffers and zero them out
grad_dtype = lambda x: dtypes.bfloat16 if x.dtype in dtypes.fp8s else x.dtype
grads = {x:x.zeros_like(dtype=grad_dtype(x)).contiguous() for x in state.values() if x.is_param}
fp8_amax = [t for ts in model._fp8_amax.values() for t in ts]
fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts]
# print model size
sz = 0
for k,v in state.items():
print(f"{colored(k, 'green' if v in grads else 'white'):30s} {str(v.shape):30s} {str(v.dtype):20s} {v.device} {v.nbytes()/1e9:.2f} GB")
sz += v.nbytes()
print(f"total sz: {sz/1e9:.2f} GB")
with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=real_vocab_size, dtype=dtypes.int)
with Timing("realize weights/grads/data: "): Tensor.realize(*state.values(), *grads.values(), tokens)
print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
if DP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)), axis=0)
if MP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)))
@TinyJit
def fwd_bwd(tokens:Tensor):
with Timing("python forward: "):
model.reset_amax()
logits = model(tokens[:, :-1], save=llama_size=="8B")
loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
with Timing("python backward: "):
for t,g in zip(grads, loss.gradient(*grads)):
apply_grad(grads[t], g.uop)
with Timing("run fwd_bwd: "): loss.realize(*grads.values(), *fp8_amax, *fp8_grad_amax)
@TinyJit
def optim_step():
for g in grads.values(): g.assign(g.zeros_like())
Tensor.realize(*grads.values())
for i in range(6):
GlobalCounters.reset()
profile_marker(f"step {i}")
with Timing(colored(f"*** step {i}: ", "red")):
fwd_bwd(tokens)
optim_step()
print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))

351
artifacts/package_sources/tinygrad/examples/mlperf/models/gpt_oss.py Normal file
View File

@@ -0,0 +1,351 @@
import math, os, functools
if __name__ == "__main__":
os.environ["DEFAULT_FLOAT"] = "bfloat16"
os.environ["OPTIM_DTYPE"] = "bfloat16"
if "DEV" not in os.environ: os.environ["DEV"] = "NULL::gfx950"
# CDNA
os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
os.environ["ALL2ALL"] = "1"
os.environ["USE_ATOMICS"] = "1"
from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
from tinygrad.uop.ops import Ops, UOp
from extra.models.llama import apply_rotary_emb
from extra.llama_kernels.rmsnorm import rmsnorm
from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8
from extra.gemm.moe_gemm import grouped_mx_gemm
from extra.gemm.moe_routing import route, dispatch, combine
FP8_DTYPE = dtypes.fp8e4m3
FP8_MAX = 448.0
INIT_STD = 0.02
ASM_GEMM = getenv("ASM_GEMM", 0)
def _quant_dequant_fwd(x:Tensor) -> Tensor:
# x (2d bf16) -> bf16 value after an mxfp8 round-trip (1x32 block scaling on the last axis)
M, K = x.shape
scale_K = K // 32
amax = x.float().reshape(M, scale_K, 32).abs().max(axis=-1)
e8 = (amax.maximum(1e-38).log2().floor() + 127).clamp(0, 254).cast(dtypes.uint8)
qscale = (127.0 - e8.cast(dtypes.float32)).exp2().reshape(M, scale_K, 1).expand(M, scale_K, 32).reshape(M, K)
x_fp8 = (x.float() * qscale).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE).cast(dtypes.float32)
return (x_fp8 * _mx_block_scale(e8)).cast(dtypes.bfloat16)
@functools.cache
def _quant_dequant_fwd_fxn(x_p, device):
return _quant_dequant_fwd(Tensor(x_p, device=device))
def _quant_dequant_bwd(grad:UOp, call:UOp) -> tuple:
return (Tensor(grad).cast(dtypes.bfloat16).uop,)
def quant_dequant_mx(x:Tensor) -> Tensor:
fxn = _quant_dequant_fwd_fxn(x.as_param(0).uop, x.device)
return Tensor(UOp.maketuple(fxn.uop).call(x.uop, grad_fxn=_quant_dequant_bwd).gettuple(0))
def _mx_scale(e8:Tensor) -> Tensor:
return _mx_block_scale(e8) if e8.ndim == 2 else _mx_block_scale_3d(e8)
def _dequant_fwd(w_q:Tensor, w_scale:Tensor) -> Tensor:
return w_q.cast(dtypes.bfloat16) * _mx_scale(w_scale)
@functools.cache
def _dequant_fwd_fxn(wq_p, ws_p, device):
return _dequant_fwd(Tensor(wq_p, device=device), Tensor(ws_p, device=device))
def _dequant_bwd(grad:UOp, call:UOp) -> tuple:
return (Tensor(grad).cast(dtypes.bfloat16).uop, None)
def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
fxn = _dequant_fwd_fxn(w_q.as_param(0).uop, w_scale.as_param(1).uop, w_q.device)
call = UOp.maketuple(fxn.uop).call(w_q.uop, w_scale.uop, grad_fxn=_dequant_bwd)
return Tensor(call.gettuple(0))
def matmul_mx(x:Tensor, w_q:Tensor, w_scale:Tensor) -> Tensor:
l_shape = x.shape[:-1]
if ASM_GEMM:
from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack
x2, K, N = x.reshape(-1, x.shape[-1]), x.shape[-1], w_q.shape[0]
wq, ws = w_q, w_scale
if (pad := (-K) % 256):
x2 = x2.pad(((0, 0), (0, pad)))
wq = wq.pad(((0, 0), (0, pad)))
ws = ws.pad(((0, 0), (0, pad // 32)), value=127).cast(dtypes.uint8)
if (npad := (-N) % 256):
wq = wq.pad(((0, npad), (0, 0)))
ws = ws.pad(((0, npad), (0, 0)), value=127).cast(dtypes.uint8)
x_q, x_e8, x_si = quantize_mxfp8(x2)
if x_si is not None and can_use_asm_gemm(x_q, wq.T):
out = asm_gemm(x_q, wq.T, mx=True, mx_scales=(x_si, x_e8, mx_pack(ws), ws), mx_w_stored=True)
return (out[:, :N] if npad else out).reshape(*l_shape, N).cast(dtypes.bfloat16)
x_phys = quant_dequant_mx(x.reshape(-1, x.shape[-1])).reshape(*l_shape, x.shape[-1])
w_phys = dequant_weight(w_q, w_scale)
return (x_phys @ w_phys.T).cast(dtypes.bfloat16)
def _pad_to_mult(t:Tensor, axis:int, mult:int=256) -> Tensor:
if (r := (-t.shape[axis]) % mult) == 0: return t
pads = [(0, 0)] * t.ndim
pads[axis] = (0, r)
return t.pad(tuple(pads))
def _pad_cols(t:Tensor) -> Tensor: return _pad_to_mult(t, -1)
def _pad_rows(t:Tensor) -> Tensor: return _pad_to_mult(t, -2)
def swiglu(x:Tensor, limit:float=7.0, alpha:float=1.702) -> Tensor:
x_glu, x_linear = x[..., ::2], x[..., 1::2]
x_glu = x_glu.clamp(max_=limit)
x_linear = x_linear.clamp(-limit, limit)
return (x_glu * (alpha * x_glu).sigmoid()) * (x_linear + 1)
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2, dtype=dtypes.float32)[:(dim // 2)] / dim))
freqs = Tensor.arange(end, dtype=dtypes.float32).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).cast(dtypes.default_float).reshape(1, end, 1, dim//2, 2)
class GPTOSS:
def __init__(self, dim:int, n_layers:int, n_heads:int, n_kv_heads:int, head_dim:int, n_experts:int, experts_per_tok:int,
intermediate_size:int, vocab_size:int, norm_eps:float=1e-5, rope_theta:int=150000, sliding_window:int=128,
swiglu_limit:float=7.0, max_context:int=8192):
self.dim, self.n_layers, self.n_heads, self.n_kv_heads, self.head_dim = dim, n_layers, n_heads, n_kv_heads, head_dim
self.n_rep = n_heads // n_kv_heads
self.n_experts, self.experts_per_tok, self.intermediate_size = n_experts, experts_per_tok, intermediate_size
self.vocab_size, self.norm_eps, self.sliding_window, self.swiglu_limit = vocab_size, norm_eps, sliding_window, swiglu_limit
self.sm_scale = 1.0 / math.sqrt(head_dim)
scaled_std = INIT_STD / math.sqrt(2 * n_layers)
q_dim, qkv_dim = n_heads * head_dim, head_dim * (n_heads + 2 * n_kv_heads)
# attn
self.wqkv, self.wqkv_scale = self._quant_weight(n_layers, qkv_dim, dim)
self.wqkv_bias = Tensor.zeros(n_layers, qkv_dim, dtype=dtypes.bfloat16).contiguous()
self.wo, self.wo_scale = self._quant_weight(n_layers, dim, q_dim, std=scaled_std)
self.wo_bias = Tensor.zeros(n_layers, dim, dtype=dtypes.bfloat16).contiguous()
self.sinks = Tensor.zeros(n_layers, n_heads, dtype=dtypes.bfloat16).contiguous()
self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
# moe ffn
self.ffn_norm = Tensor.ones(n_layers, dim).contiguous()
self.gate = Tensor.normal(n_layers, n_experts, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
self.gate_bias = Tensor.zeros(n_layers, n_experts, dtype=dtypes.bfloat16).contiguous()
self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim, moe=True)
self.w_gate_up_bias = Tensor.zeros(n_layers, n_experts, intermediate_size * 2, dtype=dtypes.bfloat16).contiguous()
self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std, moe=True)
self.w_down_bias = Tensor.zeros(n_layers, n_experts, dim, dtype=dtypes.bfloat16).contiguous()
# output
self.norm = nn.RMSNorm(dim, norm_eps)
self.tok_embeddings = nn.Embedding(vocab_size, dim)
self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
self.output = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
self.freqs_cis = precompute_freqs_cis(head_dim, max_context * 2, rope_theta).contiguous().is_param_(False)
def _quant_weight(self, *shape:int, std:float=INIT_STD, moe:bool=False):
def _one(*s:int):
w = Tensor.zeros(*s) if getenv("ZEROS") else Tensor.normal(*s, mean=0.0, std=std)
w_q, w_e8, _ = quantize_mxfp8(_pad_cols(_pad_rows(w)) if moe else w)
return w_q, w_e8.is_param_(False)
if moe:
qs = [_one(*shape[1:]) for _ in range(shape[0])]
return [q[0] for q in qs], [q[1] for q in qs]
return _one(*shape)
def _attn_mask(self, seqlen:int, dtype) -> Tensor:
i, j = Tensor.arange(seqlen).reshape(seqlen, 1), Tensor.arange(seqlen).reshape(1, seqlen)
return (j <= i).where(0.0, -1e30).cast(dtype).contiguous()
def _sliding_attention(self, xq:Tensor, xk:Tensor, xv:Tensor, sinks:Tensor) -> Tensor:
bsz, seqlen, H, hd = xq.shape
KV, R, W = self.n_kv_heads, self.n_rep, self.sliding_window
assert seqlen % W == 0, f"seqlen {seqlen} must be a multiple of sliding_window {W} for banded attention"
nb = seqlen // W
q = xq.reshape(bsz, seqlen, KV, R, hd).permute(0, 2, 3, 1, 4).reshape(bsz, KV, R, nb, W, hd).float()
k, v = (x.permute(0, 2, 1, 3).reshape(bsz, KV, 1, nb, W, hd).float() for x in (xk, xv))
kk, vv = (x.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb].cat(x, dim=-2) for x in (k, v))
sc = (q @ kk.transpose(-1, -2)) * self.sm_scale # (B,KV,R,nb,W,2W)
i, j, pv = Tensor.arange(W).reshape(W, 1), Tensor.arange(2 * W).reshape(1, 2 * W), Tensor.arange(nb).reshape(nb, 1, 1) >= 1
sc = ((j > i) & (j <= i + W) & (pv | (j >= W))).where(sc, -float("inf"))
sink = sinks.reshape(1, KV, R, 1, 1, 1).float()
m = sc.max(-1, keepdim=True).maximum(sink)
e = (sc - m).exp()
p = (e / (e.sum(-1, keepdim=True) + (sink - m).exp())).cast(dtypes.bfloat16)
attn = p @ vv.cast(dtypes.bfloat16)
return attn.reshape(bsz, KV, R, seqlen, hd).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, H * hd)
def attention(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, sliding:bool, *, attention_norm:Tensor, wqkv:Tensor,
wqkv_scale:Tensor, wqkv_bias:Tensor, wo:Tensor, wo_scale:Tensor, wo_bias:Tensor, sinks:Tensor):
bsz, seqlen, _ = x.shape
x_normed, rrms = rmsnorm(x, self.norm_eps)
qkv = matmul_mx(x_normed * attention_norm, wqkv, wqkv_scale) + wqkv_bias
qkv = qkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
xq = qkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
xk, xv = qkv[:, :, :, self.n_rep], qkv[:, :, :, self.n_rep + 1]
xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16) # (B,N,H,D)/(B,N,KV,D)
if sliding:
attn = self._sliding_attention(xq, xk, xv, sinks)
elif getenv("HK_FLASH_ATTENTION"):
from extra.thunder.amd.fa import flash_attention
attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks)
attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
else:
xqm = xq.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
xkm, xvm = xk.permute(0, 2, 1, 3).unsqueeze(2), xv.permute(0, 2, 1, 3).unsqueeze(2)
scores = (xqm @ xkm.transpose(-2, -1)).float() * self.sm_scale + mask
sink = sinks.reshape(1, self.n_kv_heads, self.n_rep, 1, 1).float()
m = scores.max(-1, keepdim=True).maximum(sink)
e = (scores - m).exp()
w = (e / (e.sum(-1, keepdim=True) + (sink - m).exp())).cast(dtypes.bfloat16)
attn = (w @ xvm).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, self.n_heads * self.head_dim)
out = matmul_mx(attn, wo, wo_scale) + wo_bias
return out, [x_normed, rrms, attn]
def feed_forward(self, x:Tensor, *, ffn_norm:Tensor, gate:Tensor, gate_bias:Tensor,
w_gate_up:Tensor, w_gate_up_scale:Tensor, w_gate_up_bias:Tensor,
w_down:Tensor, w_down_scale:Tensor, w_down_bias:Tensor):
x_normed, rrms = rmsnorm(x, self.norm_eps)
inp = x_normed * ffn_norm
logits = inp.float() @ gate.float().T + gate_bias.float()
dim, inter = self.dim, self.intermediate_size
if getenv("GROUPED_MOE", 0):
bsz, seqlen = x.shape[:2]
inp, logits = inp.reshape(-1, dim), logits.reshape(-1, self.n_experts)
r = route(logits, self.experts_per_tok, self.n_experts)
onehot = r.rows_e.one_hot(self.n_experts).float()
xg = dispatch(_pad_cols(inp.cast(dtypes.bfloat16)), r)
h = grouped_mx_gemm(xg, (w_gate_up, w_gate_up_scale), r.off)[:, :2*inter] + (onehot @ w_gate_up_bias.float()).cast(dtypes.bfloat16)
y = swiglu(h, self.swiglu_limit)
z = grouped_mx_gemm(_pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \
+ (onehot @ w_down_bias.float()).cast(dtypes.bfloat16)
out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim)
else:
thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
out = None
for e in range(self.n_experts):
gu_q, gu_s = w_gate_up[e][:2*inter, :dim].contiguous(), w_gate_up_scale[e][:2*inter, :dim//32].contiguous()
dn_q, dn_s = w_down[e][:dim, :inter].contiguous(), w_down_scale[e][:dim, :inter//32].contiguous()
gate_up = matmul_mx(inp, gu_q, gu_s) + w_gate_up_bias[e]
y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), dn_q, dn_s) + w_down_bias[e]).contiguous()
contrib = weights[..., e:e+1].cast(y.dtype) * y
out = contrib if out is None else out + contrib
return out, [x_normed, rrms]
@function(precompile=True, precompile_backward=True)
def run_layer(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, sliding:bool, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
attn, attn_saves = self.attention(x, freqs_cis, mask, sliding, **attn_kwargs)
h = x + attn
ffn, ffn_saves = self.feed_forward(h, **ffn_kwargs)
h = h + ffn
if save: return (h, *attn_saves, *ffn_saves)
return (h,)
def shard(self, device:tuple[str, ...], mp:bool=False):
assert not mp, "MP not supported"
from tinygrad.nn.state import get_parameters
for v in get_parameters(self): v.shard_(device, axis=None)
Tensor.realize(*get_parameters(self))
def __call__(self, tokens:Tensor, save:bool=True):
h = self.tok_embeddings(tokens)
bsz, seqlen = tokens.shape
freqs_cis = self.freqs_cis.cast(h.dtype)[:, :seqlen, :, :, :]
mask_full = None if getenv("HK_FLASH_ATTENTION") else self._attn_mask(seqlen, dtypes.float32)
for i in range(self.n_layers):
attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wqkv_scale=self.wqkv_scale[i],
wqkv_bias=self.wqkv_bias[i], wo=self.wo[i], wo_scale=self.wo_scale[i], wo_bias=self.wo_bias[i],
sinks=self.sinks[i])
ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], gate=self.gate[i], gate_bias=self.gate_bias[i],
w_gate_up=self.w_gate_up[i], w_gate_up_scale=self.w_gate_up_scale[i], w_gate_up_bias=self.w_gate_up_bias[i],
w_down=self.w_down[i], w_down_scale=self.w_down_scale[i], w_down_bias=self.w_down_bias[i])
h, *_ = self.run_layer(h, freqs_cis, mask_full, i % 2 == 0, attn_kwargs, ffn_kwargs, save=save)
logits = self.norm(h) @ self.output.T
return logits
def _get_pads(uop:UOp) -> list[UOp]:
if uop.op == Ops.ADD: return _get_pads(uop.src[0]) + _get_pads(uop.src[1])
return [uop]
def apply_grad(grad_buf:Tensor, new_grad:UOp):
pads = _get_pads(new_grad)
if len(pads) <= 1:
new_grad = new_grad.cast(grad_buf.dtype)
grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
return
cur = grad_buf.uop
for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
if pad.op == Ops.PAD:
grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)])
buf_slice = cur.shrink(grad_shrink)
cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
else:
cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
grad_buf.uop = cur
GPT_OSS_20B = dict(dim=2880, n_layers=24, n_heads=64, n_kv_heads=8, head_dim=64, n_experts=32, experts_per_tok=4,
intermediate_size=2880, vocab_size=128256, norm_eps=1e-5, rope_theta=150000, sliding_window=128,
swiglu_limit=7.0)
if __name__ == "__main__":
config = {}
BS = config["BS"] = getenv("BS", 16)
SEQLEN = config["SEQLEN"] = getenv("SEQLEN", 8192)
model_params = GPT_OSS_20B
real_vocab_size = model_params["vocab_size"]
if (layers := getenv("LAYERS")) != 0: model_params["n_layers"] = layers
model = GPTOSS(**model_params, max_context=SEQLEN)
state = nn.state.get_state_dict(model)
print("tensor count:", len(state))
from tinygrad import Device
is_dp = (DP := getenv("DP", 1)) > 1
device_count = DP
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
if is_dp: model.shard(device)
# preallocate all the grad buffers and zero them out
grad_dtype = lambda x: dtypes.bfloat16 if x.dtype in dtypes.fp8s else x.dtype
grads = {x:x.zeros_like(dtype=grad_dtype(x)).contiguous() for x in state.values() if x.is_param}
# print model size
sz = 0
for k,v in state.items():
print(f"{colored(k, 'green' if v in grads else 'white'):30s} {str(v.shape):30s} {str(v.dtype):20s} {v.device} {v.nbytes()/1e9:.2f} GB")
sz += v.nbytes()
print(f"total sz: {sz/1e9:.2f} GB")
with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=real_vocab_size, dtype=dtypes.int)
with Timing("realize weights/grads/data: "): Tensor.realize(*state.values(), *grads.values(), tokens)
print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
if is_dp: tokens = tokens.shard(device, axis=0)
@TinyJit
def fwd_bwd(tokens:Tensor):
with Timing("python forward: "):
logits = model(tokens[:, :-1], save=True)
loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
with Timing("python backward: "):
for t,g in zip(grads, loss.gradient(*grads)):
apply_grad(grads[t], g.uop)
with Timing("run fwd_bwd: "): loss.realize(*grads.values())
@TinyJit
def optim_step():
for g in grads.values(): g.assign(g.zeros_like())
Tensor.realize(*grads.values())
for i in range(6):
GlobalCounters.reset()
profile_marker(f"step {i}")
with Timing(colored(f"*** step {i}: ", "red")):
fwd_bwd(tokens)
optim_step()
print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))

68
artifacts/package_sources/tinygrad/examples/mlperf/models/test_apply_grad.py Normal file
View File

@@ -0,0 +1,68 @@
import unittest
from tinygrad import Tensor, TinyJit
from tinygrad.nn.state import get_parameters
from examples.mlperf.models.flat_llama import apply_grad
class FlatModel:
def __init__(self, n_layers:int, dim:int, hidden:int):
self.n_layers = n_layers
self.w1 = Tensor.uniform(n_layers, dim, hidden, low=-0.1, high=0.1)
self.w2 = Tensor.uniform(n_layers, hidden, dim, low=-0.1, high=0.1)
self.scale = Tensor.uniform(dim, low=0.9, high=1.1)
self.bias = Tensor.zeros(dim).contiguous()
def __call__(self, x:Tensor) -> Tensor:
h = x
for i in range(self.n_layers):
h = (h @ self.w1[i]).relu() @ self.w2[i] + h
return (h * self.scale + self.bias).sum()
class TestApplyGradE2E(unittest.TestCase):
def _run_with_apply_grad(self, model, xs):
grads = {p: Tensor.zeros(p.shape, dtype=p.dtype).contiguous().realize() for p in get_parameters(model)}
for x in xs:
loss = model(x)
for p, g in zip(grads, loss.gradient(*grads)):
apply_grad(grads[p], g.uop)
Tensor.realize(loss, *grads.values())
return [grads[p] for p in get_parameters(model)]
def _run_reference(self, model, xs):
for x in xs: model(x).backward()
return [p.grad for p in get_parameters(model)]
def _assert_close(self, got, expected, atol, rtol):
for g, e in zip(got, expected):
self.assertTrue(g.allclose(e, atol=atol, rtol=rtol).item(), f"grad mismatch (max abs diff {(g - e).abs().max().item()})")
def _assert_match(self, model, xs, atol, rtol):
self._assert_close(self._run_with_apply_grad(model, xs), self._run_reference(model, xs), atol, rtol)
def test_e2e_single_step(self):
model = FlatModel(n_layers=3, dim=8, hidden=16)
Tensor.realize(*get_parameters(model))
self._assert_match(model, [Tensor.randn(2, 8).realize()], atol=1e-4, rtol=1e-4)
def test_e2e_multi_step_accumulation(self):
model = FlatModel(n_layers=4, dim=8, hidden=16)
Tensor.realize(*get_parameters(model))
self._assert_match(model, [Tensor.randn(2, 8).realize() for _ in range(3)], atol=1e-4, rtol=1e-4)
def test_e2e_jit(self):
model = FlatModel(n_layers=3, dim=8, hidden=16)
Tensor.realize(*get_parameters(model))
grads = {p: Tensor.zeros(p.shape, dtype=p.dtype).contiguous().realize() for p in get_parameters(model)}
@TinyJit
def fwd_bwd(x:Tensor):
loss = model(x)
for p, g in zip(grads, loss.gradient(*grads)): apply_grad(grads[p], g.uop)
Tensor.realize(loss, *grads.values())
xs = [Tensor.randn(2, 8).realize() for _ in range(3)]
for x in xs: fwd_bwd(x)
self._assert_close([grads[p] for p in get_parameters(model)], self._run_reference(model, xs), atol=1e-3, rtol=1e-3)
if __name__ == "__main__":
unittest.main()

137
artifacts/package_sources/tinygrad/examples/mlperf/models/test_flat_llama.py Normal file
View File

@@ -0,0 +1,137 @@
import os
os.environ["WQKV"] = "1"
import unittest
import numpy as np
from tinygrad import Tensor, nn, dtypes
from tinygrad.device import Device
from examples.mlperf.models.llama import Transformer
from examples.mlperf.models.flat_llama import FlatTransformer
def copy_weights(flat:FlatTransformer, ref:Transformer):
n_layers = flat.n_layers
Tensor.realize(*nn.state.get_state_dict(ref).values())
flat.wqkv.assign(Tensor(np.stack([ref.layers[i].attention.wqkv.weight.numpy() for i in range(n_layers)])))
flat.wo.assign(Tensor(np.stack([ref.layers[i].attention.wo.weight.numpy() for i in range(n_layers)])))
flat.w1.assign(Tensor(np.stack([ref.layers[i].feed_forward.w1.weight.numpy() for i in range(n_layers)])))
flat.w2.assign(Tensor(np.stack([ref.layers[i].feed_forward.w2.weight.numpy() for i in range(n_layers)])))
flat.w3.assign(Tensor(np.stack([ref.layers[i].feed_forward.w3.weight.numpy() for i in range(n_layers)])))
flat.attention_norm.assign(Tensor(np.stack([ref.layers[i].attention_norm.weight.numpy() for i in range(n_layers)])))
flat.ffn_norm.assign(Tensor(np.stack([ref.layers[i].ffn_norm.weight.numpy() for i in range(n_layers)])))
flat.norm.weight.assign(Tensor(ref.norm.weight.numpy()))
flat.tok_embeddings.weight.assign(Tensor(ref.tok_embeddings.weight.numpy()))
flat.output.weight.assign(Tensor(ref.output.weight.numpy()))
class TestFlatLlama(unittest.TestCase):
def test_forward_match(self):
Tensor.manual_seed(42)
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
ref = Transformer(**params)
flat = FlatTransformer(**params)
copy_weights(flat, ref)
Tensor.realize(*nn.state.get_state_dict(flat).values())
tokens = Tensor([[1, 50, 100, 999, 2]])
ref_logits = ref(tokens).realize()
flat_logits = flat(tokens).realize()
self.assertEqual(ref_logits.shape, flat_logits.shape)
diff = (ref_logits - flat_logits).abs().max().item()
self.assertLess(diff, 1e-5, f"forward mismatch: max abs diff {diff}")
def test_backward_match(self):
Tensor.manual_seed(42)
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
ref = Transformer(**params)
flat = FlatTransformer(**params)
copy_weights(flat, ref)
Tensor.realize(*nn.state.get_state_dict(flat).values())
tokens = Tensor([[1, 50, 100, 999, 2, 10]])
ref_loss = ref(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
ref_loss.backward()
ref_grads = {k: v.grad.numpy() for k, v in nn.state.get_state_dict(ref).items() if v.grad is not None}
flat_loss = flat(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
flat_loss.backward()
flat_grads = {k: v.grad.numpy() for k, v in nn.state.get_state_dict(flat).items() if v.grad is not None}
# check loss matches
self.assertAlmostEqual(ref_loss.item(), flat_loss.item(), places=4)
# check output weight grad matches
diff = abs(ref_grads["output.weight"] - flat_grads["output.weight"]).max()
self.assertLess(diff, 1e-4, f"output.weight grad mismatch: max abs diff {diff}")
# check per-layer weight grads match
for i in range(params["n_layers"]):
for flat_key, ref_key in [
("wqkv", f"layers.{i}.attention.wqkv.weight"),
("wo", f"layers.{i}.attention.wo.weight"),
("w1", f"layers.{i}.feed_forward.w1.weight"),
("w2", f"layers.{i}.feed_forward.w2.weight"),
("w3", f"layers.{i}.feed_forward.w3.weight"),
]:
diff = abs(ref_grads[ref_key] - flat_grads[flat_key][i]).max()
self.assertLess(diff, 1e-4, f"layer {i} {flat_key} grad mismatch: max abs diff {diff}")
@unittest.skipUnless(Device.DEFAULT == "CPU", "multi-device CPU test")
def test_forward_match_mp(self):
Tensor.manual_seed(42)
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
from tinygrad import Device
devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
ref = Transformer(**params)
flat = FlatTransformer(**params)
copy_weights(flat, ref)
Tensor.realize(*nn.state.get_state_dict(flat).values())
flat.shard(devices, mp=True)
tokens = Tensor([[1, 50, 100, 999, 2]], device=devices[0])
ref_logits = ref(tokens.to(devices[0])).numpy()
flat_logits = flat(tokens.shard(devices)).numpy()
self.assertEqual(ref_logits.shape, flat_logits.shape)
np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
@unittest.skipUnless(Device.DEFAULT == "CPU", "multi-device CPU test")
def test_forward_match_dp(self):
Tensor.manual_seed(42)
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
from tinygrad import Device
devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
ref = Transformer(**params)
flat = FlatTransformer(**params)
copy_weights(flat, ref)
Tensor.realize(*nn.state.get_state_dict(flat).values())
flat.shard(devices)
tokens = Tensor([[1, 50, 100, 999, 2], [2, 100, 50, 1, 999]], device=devices[0])
ref_logits = ref(tokens.to(devices[0])).numpy()
flat_logits = flat(tokens.shard(devices, axis=0)).numpy()
self.assertEqual(ref_logits.shape, flat_logits.shape)
np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
@unittest.skipUnless(dtypes.fp8e4m3 in Device[Device.DEFAULT].renderer.supported_dtypes(), "fp8 not supported on this device")
def test_forward_fp8(self):
import examples.mlperf.models.flat_llama as flat_llama_mod
old_fp8 = flat_llama_mod.FP8
try:
flat_llama_mod.FP8 = 1
Tensor.manual_seed(42)
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
ref = Transformer(**params)
flat = FlatTransformer(**params)
copy_weights(flat, ref)
Tensor.realize(*nn.state.get_state_dict(flat).values())
tokens = Tensor([[1, 50, 100, 999, 2]])
ref_logits = ref(tokens).numpy()
flat_logits = flat(tokens).numpy()
self.assertEqual(ref_logits.shape, flat_logits.shape)
# FP8 has lower precision, allow larger tolerance
np.testing.assert_allclose(flat_logits, ref_logits, atol=1.0, rtol=0.1)
finally:
flat_llama_mod.FP8 = old_fp8
if __name__ == "__main__":
unittest.main()

144
artifacts/package_sources/tinygrad/examples/mlperf/optim.py Normal file
View File

@@ -0,0 +1,144 @@
from tinygrad.tensor import Tensor
from tinygrad.dtype import dtypes
from tinygrad.nn.optim import Optimizer, OptimizerGroup
from tinygrad.helpers import FUSE_OPTIM, getenv
from tinygrad.uop.ops import UOp, Ops, AxisType
STOCHASTIC_ROUND = getenv("STOCHASTIC_ROUND", 0)
MASTER_WEIGHTS = getenv("MASTER_WEIGHTS", 0)
ZERO_OPTIM = getenv("ZERO_OPTIM", 0)
FP8_AMAX_MARGIN = getenv("FP8_AMAX_MARGIN", 1.1)
IMMEDIATE_SCALE = getenv("IMMEDIATE_SCALE", 0)
MXFP8 = getenv("MXFP8", 0)
def stochastic_round_bf16(x:Tensor) -> Tensor:
bits = x.bitcast(dtypes.uint32)
if isinstance(x.device, tuple):
shape = x.uop.shard_shape if x.uop.axis is not None else x.shape
noise = Tensor(UOp(Ops.MSTACK, dtypes.default_float, tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
else:
noise = x.rand_like()
noise = (noise * 0xFFFF).cast(dtypes.uint32)
return ((bits + noise) & 0xFFFF0000).bitcast(dtypes.float32).cast(dtypes.bfloat16)
def clip_grads(grads:list[Tensor], grad_acc, clip_norm) -> Tensor:
for g in grads: g.assign(g / grad_acc)
total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
for g in grads: g.assign((g * (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(g.dtype))
return total_norm
class GradAccClipAdamW(Optimizer):
def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
super().__init__(params, lr, device, fused)
self.b1, self.b2, self.eps, self.wd = b1, b2, eps, weight_decay
self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device) for _ in [b1, b2])
self.zero = bool(ZERO_OPTIM) and isinstance(self.device, tuple) and not self.fused
self.m = [self._zero_shard(x) for x in self._new_optim_param()]
self.v = [self._zero_shard(x) for x in self._new_optim_param()]
self.grad_acc, self.clip_norm = grad_acc, clip_norm
if MASTER_WEIGHTS and self.params[0].dtype != dtypes.float32:
self.master_params:list[Tensor]|None = [self._zero_shard(p.to(self.device).float().contiguous()) for p in self.params]
else:
self.master_params = None
def _zero_shard(self, t:Tensor) -> Tensor:
if not self.zero or t.ndim < 2 or (t.shape[0] % len(self.device)) != 0: return t
return Tensor(t.uop._shard(0, UOp.range(len(self.device), -1, AxisType.DEVICE)).unshard(0)).clone()
def _zero_gather(self, t:Tensor) -> Tensor:
if not isinstance(t.device, tuple) or t.uop.axis != 0: return t
n, sz = len(t.device), t.shape[0] // len(t.device)
return Tensor.cat(*[t[p*sz:(p+1)*sz] for p in range(n)], dim=0)
def fschedule_step(self, grads:list[Tensor]) -> list[Tensor]:
updates, extra = self._step([], grads)
for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i], self.master_params[i] if self.master_params else None))
fp8_inv_scales = [tt._inv_scale for tt in self.params if hasattr(tt, '_inv_scale')]
fp8_next_inv_scales = [tt._next_inv_scale for tt in self.params if hasattr(tt, '_next_inv_scale')]
return extra + self.params + self.buffers + (self.master_params or []) + fp8_inv_scales + fp8_next_inv_scales
def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
Tensor.realize(*([grad_norm] if grad_norm is not None else []), *self.fschedule_step(grads))
def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
grads = list(grads)
for i in range(len(grads)):
if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device)
ret = []
self.b1_t *= self.b1
self.b2_t *= self.b2
for i, g in enumerate(grads):
m_new = self.b1 * self.m[i].float() + (1.0 - self.b1) * g.float()
v_new = self.b2 * self.v[i].float() + (1.0 - self.b2) * (g.float() * g.float())
self.m[i].assign(m_new.cast(self.m[i].dtype))
self.v[i].assign(v_new.cast(self.v[i].dtype))
m_hat = m_new / (1.0 - self.b1_t)
v_hat = v_new / (1.0 - self.b2_t)
up = m_hat / (v_hat.sqrt() + self.eps)
ret.append(self.lr * up)
return ret, [self.b1_t, self.b2_t] + self.m + self.v
def _apply_update(self, t:Tensor, up:Tensor, master:Tensor|None=None) -> Tensor:
w = master if master is not None else t
wd = self.wd if t.ndim >= 3 else 0.0
up = up.float().shard_like(w) + self.lr.to(w.device) * wd * w.detach()
new_w = w.detach() - up
if master is not None: master.assign(new_w)
if self.zero and not (MXFP8 and t.dtype in dtypes.fp8s): new_w = self._zero_gather(new_w)
# when master is offloaded to a different device than the param, results are resharded back onto the param's (sharded) device
offloaded = master is not None and master.device != t.device
if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16:
out = stochastic_round_bf16(new_w)
return out.shard_like(t) if offloaded else out
if t.dtype in dtypes.fp8s:
if MXFP8:
from extra.gemm.cdna_asm_gemm import quantize_mxfp8
w_q, w_e8, _ = quantize_mxfp8(new_w.reshape(-1, new_w.shape[-1]))
if self.zero: w_q, w_e8 = self._zero_gather(w_q), self._zero_gather(w_e8)
new_e8 = w_e8.reshape(t._inv_scale.shape)
t._inv_scale.assign(new_e8.shard_like(t._inv_scale) if offloaded else new_e8)
ret = w_q.reshape(t.shape)
return ret.shard_like(t) if offloaded else ret
from examples.mlperf.models.flat_llama import FP8_MAX
if IMMEDIATE_SCALE:
amax_axis = tuple(range(t._inv_scale.ndim, new_w.ndim))
new_inv = ((new_w.float().abs().max(axis=amax_axis).detach() + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype)
t._inv_scale.assign(new_inv.shard_like(t._inv_scale) if offloaded else new_inv)
scale = new_inv.reciprocal().reshape(*new_inv.shape, *([1]*(new_w.ndim-new_inv.ndim)))
ret = (new_w * scale).clamp(-FP8_MAX, FP8_MAX).cast(t.dtype)
return ret.shard_like(t) if offloaded else ret
# delayed scaling: reuse previous step's inv_scale
t._inv_scale.assign(t._next_inv_scale)
inv_scale = t._inv_scale.to(new_w.device) if offloaded else t._inv_scale
scale = inv_scale.reciprocal().reshape(*inv_scale.shape, *([1]*(new_w.ndim-inv_scale.ndim)))
scaled = (new_w * scale).clamp(-FP8_MAX, FP8_MAX)
ret = scaled.cast(t.dtype)
# update inv_scale for next step from quantized result
new_amax = (ret.float().abs().max(axis=tuple(range(inv_scale.ndim, ret.ndim))) * inv_scale * FP8_AMAX_MARGIN).detach()
new_inv = ((new_amax + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype)
t._next_inv_scale.assign(new_inv.shard_like(t._next_inv_scale) if offloaded else new_inv)
return ret.shard_like(t) if offloaded else ret
out = new_w.cast(t.dtype)
return out.shard_like(t) if offloaded else out
class GradAccClipAdamWGroup(OptimizerGroup):
def __init__(self, *optimizers:GradAccClipAdamW):
super().__init__(*optimizers)
for o in self.optimizers[1:]: o.lr = self.optimizers[0].lr
def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
offset = 0
to_realize = []
for o in self.optimizers:
n = len(o.params)
to_realize += o.fschedule_step(grads[offset:offset+n])
offset += n
Tensor.realize(*to_realize, *([grad_norm] if grad_norm is not None else []))
@property
def lr(self): return self.optimizers[0].lr
@property
def device(self): return self.optimizers[0].device
@property
def master_params(self):
mp = [mp for o in self.optimizers for mp in (o.master_params or [])]
return mp if mp else None

19
artifacts/package_sources/tinygrad/examples/mlperf/scripts/setup_kits19_dataset.sh Executable file
View File

@@ -0,0 +1,19 @@
#!/bin/bash
if [ -z $BASEDIR ]; then
export BASEDIR="./extra/datasets/"
fi
cd $BASEDIR
if [ -d "kits19" ]; then
echo "kits19 dataset is already available"
else
echo "Downloading and preparing kits19 dataset at $BASEDIR"
git clone https://github.com/neheller/kits19
cd kits19
pip3 install -r requirements.txt
python3 -m starter_code.get_imaging
echo "Done"
fi

57
artifacts/package_sources/tinygrad/examples/mlperf/scripts/stable_diffusion_downloads.sh Executable file
View File

@@ -0,0 +1,57 @@
#!/usr/bin/env bash
# adapted from https://github.com/mlcommons/training/blob/4bdf5c8ed218ad76565a2ba1ac27c919ccc6d689/stable_diffusion/README.md
# setup dirs
DATA=/raid/datasets/stable_diffusion
LAION=$DATA/laion-400m/webdataset-moments-filtered
COCO=$DATA/coco2014
mkdir -p $LAION $COCO
CKPT=/raid/weights/stable_diffusion
mkdir -p $CKPT/clip $CKPT/sd $CKPT/inception
# download data
# if rclone isn't installed system-wide / in your PATH, put the executable path in quotes below
#RCLONE=""
RCLONE="rclone"
## VAE-encoded image latents, from 6.1M image subset of laion-400m
## about 1 TB for whole download
$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/ ${LAION} --include="*.tar" -P
$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/sha512sums.txt ${LAION} -P
cd $LAION && grep -E '\.tar$' sha512sums.txt | sha512sum -c --quiet - && \
echo "All .tar files verified" || { echo "Checksum failure when validating downloaded Laion moments"; exit 1; }
## prompts and FID statistics from 30k image subset of coco2014
## 33 MB
$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k.tsv ${COCO} -P
$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k_stats.npz ${COCO} -P
# download checkpoints
## clip (needed for text and vision encoders for validation)
CLIP_WEIGHTS_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin"
CLIP_WEIGHTS_SHA256="9a78ef8e8c73fd0df621682e7a8e8eb36c6916cb3c16b291a082ecd52ab79cc4"
CLIP_CONFIG_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/raw/main/open_clip_config.json"
wget -N -P ${CKPT}/clip ${CLIP_WEIGHTS_URL}
wget -N -P ${CKPT}/clip ${CLIP_CONFIG_URL}
echo "${CLIP_WEIGHTS_SHA256} ${CKPT}/clip/open_clip_pytorch_model.bin" | sha256sum -c
## sd (needed for latent->image decoder for validation, also has clip text encoder for training)
SD_WEIGHTS_URL='https://huggingface.co/stabilityai/stable-diffusion-2-base/resolve/main/512-base-ema.ckpt'
SD_WEIGHTS_SHA256="d635794c1fedfdfa261e065370bea59c651fc9bfa65dc6d67ad29e11869a1824"
wget -N -P ${CKPT}/sd ${SD_WEIGHTS_URL}
echo "${SD_WEIGHTS_SHA256} ${CKPT}/sd/512-base-ema.ckpt" | sha256sum -c
## inception (needed for validation)
FID_WEIGHTS_URL='https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth'
FID_WEIGHTS_SHA1="bd836944fd6db519dfd8d924aa457f5b3c8357ff"
wget -N -P ${CKPT}/inception ${FID_WEIGHTS_URL}
echo "${FID_WEIGHTS_SHA1} ${CKPT}/inception/pt_inception-2015-12-05-6726825d.pth" | sha1sum -c

50
artifacts/package_sources/tinygrad/examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/README.md Normal file
View File

@@ -0,0 +1,50 @@
# 1. Problem
This problem uses the ResNet-50 CNN to do image classification.
## Requirements
Install tinygrad and mlperf-logging from master.
```
git clone https://github.com/tinygrad/tinygrad.git
python3 -m pip install -e ".[mlperf]"
```
### tinybox_green
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
This is the default on production tinybox green.
### tinybox_red
Disable cwsr
This is the default on production tinybox red.
```
sudo vi /etc/modprobe.d/amdgpu.conf
cat <<EOF > /etc/modprobe.d/amdgpu.conf
options amdgpu cwsr_enable=0
EOF
sudo update-initramfs -u
sudo reboot
# validate
sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
```
# 2. Directions
## Steps to download and verify data
```
IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
```
## Steps for one time setup
### tinybox_red
```
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
```
## Steps to run benchmark
```
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
```

13
artifacts/package_sources/tinygrad/examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/dev_beam.sh Executable file
View File

@@ -0,0 +1,13 @@
#!/bin/bash
export PYTHONPATH="."
export MODEL="resnet"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
export LAZYCACHE=0 RESET_STEP=0
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
export BENCHMARK=10 DEBUG=2
python3 examples/mlperf/model_train.py

15
artifacts/package_sources/tinygrad/examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/dev_run.sh Executable file
View File

@@ -0,0 +1,15 @@
#!/bin/bash
export PYTHONPATH="."
export MODEL="resnet"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
export LAZYCACHE=0 RESET_STEP=0
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
export EVAL_START_EPOCH=3 EVAL_FREQ=4
export WANDB=1 PARALLEL=0
python3 examples/mlperf/model_train.py

23
artifacts/package_sources/tinygrad/examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/run_and_time.sh Executable file
View File

@@ -0,0 +1,23 @@
#!/bin/bash
export PYTHONPATH="."
export MODEL="resnet"
export SUBMISSION_PLATFORM="tinybox_green"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
export LAZYCACHE=0 RESET_STEP=0
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
# pip install -e ".[mlperf]"
export LOGMLPERF=1
export SEED=$RANDOM
DATETIME=$(date "+%m%d%H%M")
LOGFILE="resnet_green_${DATETIME}_${SEED}.log"
# init
BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
# run
PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE

50
artifacts/package_sources/tinygrad/examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/README.md Normal file
View File

@@ -0,0 +1,50 @@
# 1. Problem
This problem uses the ResNet-50 CNN to do image classification.
## Requirements
Install tinygrad and mlperf-logging from master.
```
git clone https://github.com/tinygrad/tinygrad.git
python3 -m pip install -e ".[mlperf]"
```
### tinybox_green
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
This is the default on production tinybox green.
### tinybox_red
Disable cwsr
This is the default on production tinybox red.
```
sudo vi /etc/modprobe.d/amdgpu.conf
cat <<EOF > /etc/modprobe.d/amdgpu.conf
options amdgpu cwsr_enable=0
EOF
sudo update-initramfs -u
sudo reboot
# validate
sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
```
# 2. Directions
## Steps to download and verify data
```
IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
```
## Steps for one time setup
### tinybox_red
```
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
```
## Steps to run benchmark
```
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
```

13
artifacts/package_sources/tinygrad/examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/dev_beam.sh Executable file
View File

@@ -0,0 +1,13 @@
#!/bin/bash
export PYTHONPATH="."
export MODEL="resnet"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
export LAZYCACHE=0 RESET_STEP=0
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
export BENCHMARK=10 DEBUG=2
python3 examples/mlperf/model_train.py

15
artifacts/package_sources/tinygrad/examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/dev_run.sh Executable file
View File

@@ -0,0 +1,15 @@
#!/bin/bash
export PYTHONPATH="."
export MODEL="resnet"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
export LAZYCACHE=0 RESET_STEP=0
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
export EVAL_START_EPOCH=3 EVAL_FREQ=4
export WANDB=1 PARALLEL=0
python3 examples/mlperf/model_train.py

23
artifacts/package_sources/tinygrad/examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh Executable file
View File

@@ -0,0 +1,23 @@
#!/bin/bash
export PYTHONPATH="."
export MODEL="resnet"
export SUBMISSION_PLATFORM="tinybox_red"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
export LAZYCACHE=0 RESET_STEP=0
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
# pip install -e ".[mlperf]"
export LOGMLPERF=1
export SEED=$RANDOM
DATETIME=$(date "+%m%d%H%M")
LOGFILE="resnet_red_${DATETIME}_${SEED}.log"
# init
BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
# run
PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE

8
artifacts/package_sources/tinygrad/examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh Executable file
View File

@@ -0,0 +1,8 @@
#!/bin/bash
rocm-smi --setprofile compute
rocm-smi --setmclk 3
rocm-smi --setperflevel high
# power cap to 350W
echo "350000000" | sudo tee /sys/class/drm/card{1..6}/device/hwmon/hwmon*/power1_cap

37
artifacts/package_sources/tinygrad/examples/mlperf/training_submission_v4.0/tinycorp/systems/tinybox_green.json Normal file
View File

@@ -0,0 +1,37 @@
{
"submitter": "tinycorp",
"division": "closed",
"status": "available",
"system_name": "tinybox green",
"number_of_nodes": "1",
"host_processors_per_node": "1",
"host_processor_model_name": "AMD EPYC 7532 32-Core Processor",
"host_processor_core_count": "64",
"host_processor_frequency": "",
"host_processor_caches": "",
"host_processor_interconnect": "",
"host_memory_capacity": "128GB",
"host_storage_type": "NVMe SSD",
"host_storage_capacity": "4 TB raid array + 1 TB boot",
"host_networking": "",
"host_networking_topology": "",
"host_memory_configuration": "8x 16GB DDR4",
"accelerators_per_node": "6",
"accelerator_model_name": "NVIDIA GeForce RTX 4090",
"accelerator_host_interconnect": "PCIe 4.0 x16",
"accelerator_frequency": "",
"accelerator_on-chip_memories": "",
"accelerator_memory_configuration": "GDDR6X",
"accelerator_memory_capacity": "24GB",
"accelerator_interconnect": "",
"accelerator_interconnect_topology": "",
"cooling": "air",
"hw_notes": "",
"framework": "tinygrad, commit 0e8aa0e2886bf9a2d3ce093bce87305e182e6d4a",
"other_software_stack": {
"python": "3.10.12",
"CUDA": "12.4"
},
"operating_system": "Ubuntu 22.04.4",
"sw_notes": ""
}

37
artifacts/package_sources/tinygrad/examples/mlperf/training_submission_v4.0/tinycorp/systems/tinybox_red.json Normal file
View File

@@ -0,0 +1,37 @@
{
"submitter": "tinycorp",
"division": "closed",
"status": "available",
"system_name": "tinybox red",
"number_of_nodes": "1",
"host_processors_per_node": "1",
"host_processor_model_name": "AMD EPYC 7532 32-Core Processor",
"host_processor_core_count": "64",
"host_processor_frequency": "",
"host_processor_caches": "",
"host_processor_interconnect": "",
"host_memory_capacity": "128GB",
"host_storage_type": "NVMe SSD",
"host_storage_capacity": "4 TB raid array + 1 TB boot",
"host_networking": "",
"host_networking_topology": "",
"host_memory_configuration": "8x 16GB DDR4",
"accelerators_per_node": "6",
"accelerator_model_name": "AMD Radeon RX 7900 XTX",
"accelerator_host_interconnect": "PCIe 4.0 x16",
"accelerator_frequency": "",
"accelerator_on-chip_memories": "",
"accelerator_memory_configuration": "GDDR6",
"accelerator_memory_capacity": "24GB",
"accelerator_interconnect": "",
"accelerator_interconnect_topology": "",
"cooling": "air",
"hw_notes": "",
"framework": "tinygrad, commit 0e8aa0e2886bf9a2d3ce093bce87305e182e6d4a",
"other_software_stack": {
"python": "3.10.12",
"ROCm": "6.1"
},
"operating_system": "Ubuntu 22.04.4",
"sw_notes": ""
}

73
artifacts/package_sources/tinygrad/examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_green/README.md Normal file
View File

@@ -0,0 +1,73 @@
# 1. Problem
This problem uses BERT for NLP.
## Requirements
Install tinygrad and mlperf-logging from master.
```
git clone https://github.com/tinygrad/tinygrad.git
python3 -m pip install -e ".[mlperf]"
```
Also install tqdm and tensorflow.
```
pip install tqdm tensorflow
```
### tinybox_green
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
This is the default on production tinybox green.
### tinybox_red
Disable cwsr + increase mes timeout.
Install the custom amdgpu driver per [README](https://github.com/nimlgen/amdgpu_ubuntu_22_04/blob/v6.1.3/readme.md)
# 2. Directions
## Steps to download and verify data
### 1. Download raw data
```
BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
```
### 2. Preprocess train and validation data
Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended.
#### Training:
```
BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
```
Generating a specific topic (Between 0 and 499)
```
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
```
#### Validation:
```
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
```
## Running
### tinybox_green
#### Steps to run benchmark
```
examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
```
### tinybox_red
#### One time setup
```
examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/setup.sh
```
#### Steps to run benchmark
```
examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
```

13
artifacts/package_sources/tinygrad/examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_green/dev_beam.sh Executable file
View File

@@ -0,0 +1,13 @@
#!/bin/bash
export PYTHONPATH="."
export MODEL="bert"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=66 EVAL_BS=6
export BEAM=4 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=512
export IGNORE_JIT_FIRST_BEAM=1
export BASEDIR="/raid/datasets/wiki"
export BENCHMARK=10 DEBUG=2
python3 examples/mlperf/model_train.py

13
artifacts/package_sources/tinygrad/examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_green/dev_run.sh Executable file
View File

@@ -0,0 +1,13 @@
#!/bin/bash
export PYTHONPATH="."
export MODEL="bert"
export DEFAULT_FLOAT="HALF" GPUS=6 BS=66 EVAL_BS=6
export BEAM=4 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=512
export IGNORE_JIT_FIRST_BEAM=1
export BASEDIR="/raid/datasets/wiki"
export WANDB=1 PARALLEL=0
RUNMLPERF=1 python3 examples/mlperf/model_train.py

Some files were not shown because too many files have changed in this diff Show More