{"id":21135244,"url":"https://github.com/cloneofsimo/ezmup","last_synced_at":"2025-10-16T21:31:59.933Z","repository":{"id":214889516,"uuid":"737494102","full_name":"cloneofsimo/ezmup","owner":"cloneofsimo","description":"Simple implementation of muP, based on Spectral Condition for Feature Learning. The implementation is SGD only, dont use it for Adam ","archived":false,"fork":false,"pushed_at":"2024-07-28T07:01:45.000Z","size":1423,"stargazers_count":73,"open_issues_count":4,"forks_count":3,"subscribers_count":5,"default_branch":"master","last_synced_at":"2025-01-23T14:13:22.033Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":"https://arxiv.org/abs/2310.17813","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/cloneofsimo.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2023-12-31T08:48:53.000Z","updated_at":"2024-12-26T17:14:08.000Z","dependencies_parsed_at":"2025-01-16T13:35:15.096Z","dependency_job_id":"0d9a6257-2c34-4313-a117-af3a298753c7","html_url":"https://github.com/cloneofsimo/ezmup","commit_stats":{"total_commits":14,"total_committers":2,"mean_commits":7.0,"dds":0.2857142857142857,"last_synced_commit":"90f9fd2f446c9a5d2858035787ed36c573debdc8"},"previous_names":["cloneofsimo/ezmup"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cloneofsimo%2Fezmup","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cloneofsimo%2Fezmup/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cloneofsimo%2Fezmup/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/cloneofsimo%2Fezmup/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/cloneofsimo","download_url":"https://codeload.github.com/cloneofsimo/ezmup/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":236748988,"owners_count":19198622,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-11-20T06:46:54.199Z","updated_at":"2025-10-16T21:31:59.550Z","avatar_url":"https://github.com/cloneofsimo.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Minimal Implementation of muP (Maximal Update Parametrization) that happens to be also Easy\n\n\u003e This is radical implementation of the muP algorithm (Maximal Update Parametrization) for the paper [Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer](https://arxiv.org/abs/2203.03466) and [A Spectral Condition for Feature Learning](https://arxiv.org/abs/2310.17813), series of research driven by [Greg Yang](https://thegregyang.com/). *This is not an official implementation, which can be found [here](https://github.com/microsoft/mup).*\n\n## What is muP?\n\n1. You want to train a large NN with pytorch (you have a *width* as a hyperparameter)\n2. You want to find optimal hyperparameters, such as learning rate, adam beta, lr scheduler, init varaince, etc (see list of hparams that are relevent [here](#list-of-hyperparameters). Note that not all hparams are relevant to muP!)\n3. So of course you are going to search for the optimal hparams using the *small model*, and apply it to the *big model*... right?\n\n![Alt text](contents/meme.png)\n\nBut as you all know, with *default pytorch settings*, you can't do this! You often have to **search the hparams for the big model, because hparams you find on the small model is often not applicable to the big model.**\n\nHowever, this can be fixed! *muP* is a method of **scaling width** so that you can search the hyperparams on the smaller model (small width), and then apply it to the big model (large width): **you can transfer the hparams from small model to big model!, that saves you a lot of time and money!**\n\n![Alt text](contents/image.png)\n\n\u003e You see? in practice, your optimal Learing rate found with width 128 does not transfer to width 8192. But with muP, you can transfer the hyperparameters from small model to big model! Figure directly from [Tensor Programs V](https://arxiv.org/abs/2203.03466)\n\n## How to use this repo?\n\nThis package suggest the following approach:\n\nFirst, in your project, change the config so that some weird large enough prime number represents the *varying width*. By *weird*, I mean such number should not be used in your other hyperparameters of model shapes. 47 is such good number.\n\n```python\n\nfrom ezmup import Ezmup\n\nmodel = MyModel(ff_width = 47 * 32, hidden_dim = 47 * 8, head_dim = 47, num_layers = 4, ...)\n\n```\n\nThen, you can use the ezmup package so that it changes all of the width variable to value you want.\n\n\n```python\nmup_engine = Ezmup(width_basis = 47, model = model) # this will treat 47 as a width variable. i.e., replace 47 occuring in all of the model as a variable length.\n\n# do hyperparameter-search, such as scaling, with fixed width. here, we take 32.\nmup_engine.change_width_as(32)\n\noptimizer = mup_engine.get_optimizer(lr = test_lr)\n\n# now after you get the best parameters, you can apply them to your desired width.\nmup_engine.change_width_as(1024)\n\n# if you want to use the model for other purposes, say, other frameworks just save them as state-dict, safetensors, etc.\ntorch.save(model.state_dict(), 'model.pt')\n\n```\n\n\nIn that way, we can\n\n1. Use muP in our own projects without too much hassle (just set the varying-width you desire to be prime number like 47! Explained later)\n2. Check the correctness of muP via Coord-Checking\n\nOh well, this is exactly how you use this package! The code is very minimal, so do have a look as well.\n\n## Installation\n\n```bash\npython3 -m pip install ezmup+https://github.com/cloneofsimo/ezmup.git\n```\n\n## Other methods:\n\nThe code does two things:\n\n1. Without needing to replace the code-layers in your implementation, it changes the width variable to the value you want, by actually replacing the neural network parameters into desired ones. While doing that, it will use a standard scaling InitStd, which is also some value you can change (per layer!)\n\n2. It also tells you the learning-rate scaling factor you need per layer, which is a varying factor depending on the width.\n\nSay you have know that you want different learning rate for different layers such as Embedding and FFN.\n\nThen, you can do the following:\n\n```python\n# You can do it manually by getting the parameter_name, lr_scaling dictionary.\nmup_engine = Ezmup(width_basis = 47, model = model)\nmup_scaling = mup_engine.lr_scaling_dict\n\nlr_dict = {\n    'embedding' : 1e-3,\n    'ffn' : 1e-4\n}\n\noptimizer_groups = [\n    {'params' : [p] , 'lr' : lr_dict[name] * mup_scaling[p_name] } for p_name, p in model.named_parameters()\n]\noptimizer = Adam(optimizer_groups, **kwargs)\n```\n\n\n# List of hyperparameters\n\nNot all hyperparameters are relevant to muP. Here is the list of hyperparameters that are relevant to muP:\n\n![Alt text](contents/hyperparams.png)\n\n## So what is with the prime number?\n\nOk so the difficulty of muP implementation for the drop-in replacement is that you need to change the width variable *automagically*. muP requires you to identify which `fan_in` or `fan_out` refers to *infinite-width*.\n\nSo, the idea is to use a prime number as a width variable, which is then used as an indicator via dividability : if *some* shape is multiple of this prime number, then it is infinite-width. With design choice, users don't have to manually specify which shape is infinite-width, nor change the code.\n\n\n## Coord-Checking\n\nYou can see how to run the coord checking in the `example.py` file. The result should looks as:\n\n![Alt text](contents/coord-check.png)\n\n\n# For Developers\n## Installation\n\n```bash\ngit clone https://github.com/cloneofsimo/ezmup.git\ncd ezmup\npython3 -m pip install \".[dev]\"\n```\n\n## Build docs\n```bash\nsphinx-apidoc -f -o docs/source ezmup\ncd docs\nmake html\n```\n\n## Update `requirements.txt` and `requirements-dev.txt` from `pyproject.toml` (Use `pip-compile`)\n```bash\npip-compile -o requirements.txt pyproject.toml\npip-compile --extra dev -o requirements-dev.txt pyproject.toml\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fcloneofsimo%2Fezmup","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fcloneofsimo%2Fezmup","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fcloneofsimo%2Fezmup/lists"}