{"id":28100732,"url":"https://github.com/mit-han-lab/neurips-micronet","last_synced_at":"2025-07-07T05:03:48.530Z","repository":{"id":89746954,"uuid":"220020937","full_name":"mit-han-lab/neurips-micronet","owner":"mit-han-lab","description":"[JMLR'20] NeurIPS 2019 MicroNet Challenge Efficient Language Modeling, Champion","archived":false,"fork":false,"pushed_at":"2021-02-26T02:55:44.000Z","size":68779,"stargazers_count":40,"open_issues_count":0,"forks_count":8,"subscribers_count":9,"default_branch":"master","last_synced_at":"2025-05-13T18:49:39.774Z","etag":null,"topics":["efficient-model","knowledge-distillation","language-modeling","natural-language-processing","pruning","quantization"],"latest_commit_sha":null,"homepage":"https://spatten.mit.edu","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/mit-han-lab.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","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,"zenodo":null}},"created_at":"2019-11-06T14:52:05.000Z","updated_at":"2024-07-20T02:56:53.000Z","dependencies_parsed_at":null,"dependency_job_id":"7c7385b1-722c-40a9-9334-1bd68a091c4f","html_url":"https://github.com/mit-han-lab/neurips-micronet","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/mit-han-lab/neurips-micronet","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mit-han-lab%2Fneurips-micronet","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mit-han-lab%2Fneurips-micronet/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mit-han-lab%2Fneurips-micronet/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mit-han-lab%2Fneurips-micronet/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/mit-han-lab","download_url":"https://codeload.github.com/mit-han-lab/neurips-micronet/tar.gz/refs/heads/master","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mit-han-lab%2Fneurips-micronet/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":264016712,"owners_count":23544624,"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":["efficient-model","knowledge-distillation","language-modeling","natural-language-processing","pruning","quantization"],"created_at":"2025-05-13T18:38:31.614Z","updated_at":"2025-07-07T05:03:48.494Z","avatar_url":"https://github.com/mit-han-lab.png","language":"Jupyter Notebook","funding_links":[],"categories":["Jupyter Notebook"],"sub_categories":[],"readme":"# MicroNet: Team MIT-HAN-Lab\n\n## News\nOur work has been accepted by PMLR!  [PMLR link](http://proceedings.mlr.press/v123/yan20a.html); [arXiv link](https://arxiv.org/abs/2005.07877).\n\nHanrui and Zhongxia gave a [talk](https://slideslive.com/38922007/competition-track-day-13) (Starts from 26:17) on the challenge in NeurIPS 2019, Vancouver.\n\n\u003cimg src=\"neurips_micronet_challenge/talk_photo.png\" alt=\"drawing\" width=\"240\"/\u003e\n\n## Introduction\nThis codebase provides the code, configurations, and commands for our submission to PMLR for representing our work in the NeurIPS 2019 MicroNet Challenge on the [WikiText-103 Language Modeling task](https://micronet-challenge.github.io/index.html). The information for our submission to NeurIPS 2019 MicroNet Challenge can be found [here](https://github.com/mit-han-lab/neurips-micronet/tree/master/neurips_micronet_challenge).\n\nTeam members: Zhongxia Yan, Hanrui Wang, Demi Guo, Song Han.\n\nOur work implements or make improvements to the following methods, integrating them to create an efficient language model for the Wikitext-103 task\n* [Transformer-XL](https://arxiv.org/abs/1901.02860)\n* [Adaptive embedding and softmax](http://arxiv.org/abs/1809.10853)\n* [Non-parametric cache](http://arxiv.org/abs/1612.04426)\n* [Hebbian softmax](http://arxiv.org/abs/1803.10049)\n* [Knowledge distillation with teacher annealing](https://arxiv.org/abs/1907.04829)\n* Pruning: we use the [Distiller implementation](https://nervanasystems.github.io/distiller/algo_pruning.html#automated-gradual-pruner-agp) of Automated Gradual Pruning\n* Quantization: we use the [Distiller implementation](https://nervanasystems.github.io/distiller/algo_quantization.html) of the quantization aware training of the symmetric range-based linear quantizer.\n\nOur best model achieves a validation perplexity of 34.1 and test perplexity of 35.0 on Wikitext-103, while using 1.8M 32-bit parameters and 8.8M 32-bit math operations.\n\n## MicroNet\nThe MicroNet challenge website can be found [here](https://micronet-challenge.github.io/). Our best model discussed below gets a MicroNet score of `0.0387`. Note that this is better than our score `0.0475` on the MicroNet website, since we previously had an evaluation error which miscalculated the number of math operations.\n\n## Pipeline\nWe show our pipeline with incremental performance results. Each row in each column is an ablation on our best configuration [`quantize_prune0.358_distill_8.3M_cache2000_hebbian_step175000_cache3000_bits9`](results/quantize_prune0.358_distill_8.3M_cache2000_hebbian_step175000_cache3000_bits9/config.yaml). The left column does not use compression techniques, while the right column does. From top to bottom, each stack displays the progression of techniques. Each row displays associated metrics: parameters (top left), operations (top right), validation perplexity (bottom left), and estimated processing time (bottom right). Metrics are displayed when changed from the previous row, with green for desirable change and red for undesirable. Red rows represent Core LM techniques, blue rows represent compression techniques, and gray rows represent cache search; joined rows represent joint training.\n\u003cp align=\"center\"\u003e\u003cimg align=\"center\" src=\"figures/pipeline.svg\" width=\"500\"/\u003e\u003c/p\u003e\n\n## Installation\nWe run our code on Python 3.6.8 and PyTorch 1.1.0+. We set up our environment using a mixture of Conda and Pip, though in theory Conda shouldn't be necessary. Our code has submodules, so make sure to use `--recursive` while cloning.\n```bash\ngit clone --recursive https://github.com/mit-han-lab/neurips-micronet.git\n\n# If you need to install conda first, follow the instructions from https://docs.conda.io/en/latest/miniconda.html\nconda create -n micronet python=3.6\nconda activate micronet\n\n# Install distiller (pruning and quantization) requirements\npip install -r distiller/requirements.txt\n\n# For using mixed precision training with https://github.com/NVIDIA/apex\n# This is not necessary but some experiments may benefit larger batch sizes with mixed precision training\n# Depending on what CUDA version your PyTorch uses, you may have to change the CUDA_HOME environment\n# variable in the command below\ncd apex \u0026\u0026 CUDA_HOME=/usr/local/cuda-10.0 pip install -v --no-cache-dir --global-option=\"--cpp_ext\" --global-option=\"--cuda_ext\" ./\n\npip install enlighten future gitpython==3.1.2\n```\n\n## Dataset\nOur code in [`setup_data.ipynb`](setup_data.ipynb) directly downloads and preprocesses the [Wikitext-103 dataset](https://blog.einstein.ai/the-wikitext-long-term-dependency-language-modeling-dataset/). Just run all the cells!\n\n## Models and Configurations\nAll of our configurations are already in the [`results/`](results/) directory. For example, our best configuration evaluated by the MicroNet criteria is [`results/quantize_prune0.358_distill_8.3M_cache2000_hebbian_step175000_cache3000_bits9`](results/quantize_prune0.358_distill_8.3M_cache2000_hebbian_step175000_cache3000_bits9/config.yaml). You can download our trained models from [here](https://www.dropbox.com/sh/8b37zkfvuyog4tu/AAB1wH9GgQgVgO1b7Lh0Pap4a?dl=0) into your directory. The necessary files to evaluate a particular configuration are\n```bash\nresults/\u003cconfiguration_name\u003e/config.yaml # this contains readable hyperparameters that we use\nresults/\u003cconfiguration_name\u003e/cache_step\u003csearched_checkpoint_step\u003e_n\u003csearch_cache_size\u003e.yaml # this contains local searched cache parameters for the particular checkpoint step and local search cache size\nresults/\u003cconfiguration_name\u003e/models/model-\u003cstep\u003e.pth # this contains the step number, saved weights of the network, and saved weights of any optimizer\n```\n\nConfiguration names are mostly intuitive. We use `attn129` and `attn257` to denote `C = 129` and `C = 257`, respectively, otherwise the default is `C = 97`. For quantization configurations, for example `quantize_prune0.358_distill_8.3M_cache2000_hebbian_step175000_cache3000_bits9`, the first `cache\u003csize\u003e` refers to the training cache size, the second `cache\u003csize\u003e` refers to the local search cache size, and `step\u003cnumber\u003e` refers to the checkpoint step of the model before quantization.\n\n### Evaluation\nTo evaluation our trained model, make sure to download it as mentioned above, then go to [`micronet_evaluate.ipynb`](micronet_evaluate.ipynb) and substitute in the configuration name. This gives the validation and test perplexities as well as the number of parameters and math operations after pruning and quantization.\n\n### Running the Pipeline\nIn general, we run the pipeline from the desired configuration directory (e.g. [`results/quantize_prune0.358_distill_8.3M_cache2000_hebbian_step175000_cache3000_bits9/`](results/quantize_prune0.358_distill_8.3M_cache2000_hebbian_step175000_cache3000_bits9/)).\n\n#### Training the Core LM with or without Distillation\nYou can refer to [`results/`](results/) for examples of configurations for training with / without adaptive softmax, training cache, Hebbian softmax, or distillation. You may also use this to train a teacher model. To start training, you can use an existing configuration directory or create a new one, then run\n```bash\ncd \u003cconfiguration_directory\u003e\n# Make sure your directory has the correct config.yaml file\n\nCUDA_VISIBLE_DEVICES=\u003cdevice\u003e python ../../main.py .\n```\nWe recommend that you modify the `train_batch` hyperparameter in the `config.yaml` to be the maximum that fits in memory. If you'd like to use mixed precision training with [Apex](https://github.com/NVIDIA/apex) Amp, make sure to install Apex as in the instruction, then `opt_level=O1` as an argument to the command. Note that sometimes the cache or Hebbian softmax have instabilities with using mixed precision training. In that case you could try to debug the code or just use full precision.\n\n#### Generating Soft Labels for Distillation\nIf you'd like to use distillation and you have already trained a teacher model, you may generate the top 30 soft labels for the training set tokens by running this from the teacher's `\u003cconfiguration_directory\u003e`\n```bash\ncd \u003cconfiguration_directory\u003e\nCUDA_VISIBLE_DEVICES=\u003cdevice\u003e python ../../gen_soft_labels.py .\n```\nNote that this takes around 40G of disk storage and may take several hours.\n\n#### Pruning\nAfter you train a model for a configuration `\u003cconfiguration_directory\u003e` with the above instructions, you can run\n```bash\ncd \u003cconfiguration_directory\u003e\nCUDA_VISIBLE_DEVICES=\u003cdevice\u003e python ../../setup_prune.py .\n```\nThis will automatically create a new configuration directory `\u003cprune_configuration_directory\u003e` for you. Just follow the printed instruction to run pruning. Note that by default this uses the pruning configuration for sparsity `0.358`. You can also use pruninf configuration for sparsity `0.239` and `0.477` by replacing the `distiller_prune.yaml` in your `\u003cprune_configuration_directory\u003e` with [`distiller_prune0.239.yaml`](distiller_prune0.239.yaml) or [`distiller_prune0.477.yaml`](distiller_prune0.477.yaml).\n\n#### Cache Local Search\nYou may run local search on either a pruned model or unpruned (but trained) model. This will generate a new `\u003cconfiguration_directory\u003e/cache_step\u003csearched_checkpoint_step\u003e_n\u003csearch_cache_size\u003e.yaml` with your searched cache configuration.\n```bash\ncd \u003cconfiguration_directory\u003e\nCUDA_VISIBLE_DEVICES=\u003cdevice\u003e python ../../cache_search.py . n_cache=\u003csearch_cache_size\u003e\n```\nThis will by default run search on the checkpoint trained for the largest number of steps. If you want to run local search on an arbitrary saved checkpoint step, then add the argument `step=\u003ctrained_step_that_you_want_to_search\u003e`.\n\n#### Quantization\nYou may quantize a trained model, pruned model, or a local-searched model.\n```bash\ncd \u003cconfiguration_directory\u003e\nCUDA_VISIBLE_DEVICES=\u003cdevice\u003e python ../../setup_quantize.py . bits=\u003cbits_to_quantize_to\u003e\n```\nThis will by default quantize the largest checkpoint with the `n_cache` value in `\u003cconfiguration_directory\u003e/config.yaml`. If you'd like to quantize with a different cache size (possibly with local searched parameters), you can add the argument `n_cache=\u003ccache_size\u003e`. If you'd like to quantize a different step, add the argument `step=\u003ctrained_step_that_you_want_to_quantize\u003e`, this will by default use your cache parameters if you performed local search with that step previously.\n\n## Common Questions / Errors\n\u003eAssertionError: Training already exists\n\nUsually this happens if you try to run training again after your training crashes. This is because we have a guard against accidentally running multiple trainings as the same time. Just remove the `\u003cconfiguration_directory\u003e/is_training` guard file and you should be fine.\n\n\u003eWhy is the perplexity `nan`?\n\nTo prevent obscenely large perplexities at the beginning of training, we set perplexities greater than `e ** 5` to be `nan`.\n\n## FAQ\nIf you have any further questions about our submission, please don't hesitate reaching out to us through Github Issues :)\nThanks!\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmit-han-lab%2Fneurips-micronet","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fmit-han-lab%2Fneurips-micronet","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmit-han-lab%2Fneurips-micronet/lists"}