Merge pull request #58 from graphcore-research/umup-updates

Updates to support u-muP, as the new default behaviour

README.md

@@ -1,10 +1,9 @@
-# Unit Scaling
+# Unit-Scaled Maximal Update Parameterization (u-μP)
 
-A library for unit scaling in PyTorch, based on the paper
-[Unit Scaling: Out-of-the-Box Low-Precision Training](https://arxiv.org/abs/2303.11257).
+A library for unit scaling in PyTorch, based on the paper [Unit-Scaled Maximal Update Parametrization](https://openreview.net/forum?id=44NKKzz1n5) and previous work [Unit Scaling: Out-of-the-Box Low-Precision Training](https://arxiv.org/abs/2303.11257).
 
 Documentation can be found at
-[https://graphcore-research.github.io/unit-scaling](https://graphcore-research.github.io/unit-scaling).
+[https://graphcore-research.github.io/unit-scaling](https://graphcore-research.github.io/unit-scaling) and an example notebook at [examples/demo.ipynb](examples/demo.ipynb).
 
 **Note:** The library is currently in its _beta_ release.
 Some features have yet to be implemented and occasional bugs may be present.

dev

@@ -85,7 +85,8 @@ def format(check: bool) -> None:
 def copyright() -> None:
     """check for Graphcore copyright headers on relevant files"""
     command = (
-        f"find {' '.join(PYTHON_ROOTS)} -type f -not -name *.pyc"
+        f"find {' '.join(PYTHON_ROOTS)} -type f"
+        " -not -name *.pyc -not -name *.json -not -name .gitignore"
         " | xargs grep -L 'Copyright (c) 202. Graphcore Ltd[.] All rights reserved[.]'"
     )
     print(f"$ {command}", file=sys.stderr)

docs/api_reference.rst

@@ -39,7 +39,9 @@ Click below for the full documentation:
    unit_scaling.constraints
    unit_scaling.formats
    unit_scaling.functional
+   unit_scaling.optim
    unit_scaling.scale
    unit_scaling.transforms
    unit_scaling.transforms.utils
    unit_scaling.utils
+   unit_scaling.core.functional

examples/.gitignore

@@ -0,0 +1 @@
+*.json