{
  "@context": "https://schema.org",
  "@id": "https://aggr.aymericbeaumet.com/items/kimi-com/2025-02-23-muon-is-scalable-for-llm-training/#webpage",
  "@type": [
    "WebPage",
    "ArchiveComponent"
  ],
  "articleBody": "[![](https://github.com/MoonshotAI/Moonlight/raw/master/figures/banner.png)](https://github.com/MoonshotAI/Moonlight/blob/master/Moonlight.pdf)\n\n## Abstract\n\nRecently, the [Muon optimizer](https://github.com/KellerJordan/Muon) based on matrix orthogonalization has demonstrated strong results in training small-scale language models, but the scalability to larger models has not been proven. We identify two crucial techniques for scaling up Muon: (1) adding weight decay and (2) carefully adjusting the per-parameter update scale. These techniques allow Muon to work out-of-the-box on large-scale training without the need of hyper-parameter tuning. Scaling law experiments indicate that Muon achieves ∼ 2× computational efficiency compared to AdamW with compute optimal training.\n\nBased on these improvements, we introduce **Moonlight**, a 3B/16B-parameter Mixture-of-Expert (MoE) model trained with 5.7T tokens using Muon. Our model improves the current Pareto frontier, achieving better performance with much fewer training FLOPs compared to prior models.\n\nWe open-source our distributed Muon implementation that is memory optimal and communication efficient. We also release the pretrained, instruction-tuned, and intermediate checkpoints to support future research.\n\nOur code is available at [MoonshotAI/Moonlight](https://github.com/MoonshotAI/Moonlight).\n\n## Key Ingredients\n\nOur work builds upon Muon while systematically identifying and resolving its limitations in large-scale training scenarios. Our technical contributions include:\n\n- **Analysis for Effective Scaling of Muon**: Through extensive analysis, we identify that weight decay plays a crucial roles in Muon's scalability. Besides, we proposed to keep a consistent update root mean square (RMS) across different matrix and non-matrix parameters through parameter-wise update scale adjustments. Such adjustments significantly enhanced training stability.\n\n- **Efficient Distributed Implementation**: We develop a distributed version of Muon with ZeRO-1 style optimization, achieving optimal memory efficiency and reduced communication overhead while preserving the mathematical properties of the algorithm.\n\n- **Scaling Law Validation**: We performed scaling law research that compares Muon with strong AdamW baselines, and showed the superior performance of Muon (see Figure 1). Based on the scaling law results, Muon achieves comparable performance to AdamW trained counterparts while requiring only approximately 52% of the training FLOPs.\n\n[![](https://github.com/MoonshotAI/Moonlight/raw/master/figures/scaling.png)](https://github.com/MoonshotAI/Moonlight/blob/master/figures/scaling.png)\n\n*Scaling up with Muon. **(a)** Scaling law experiments comparing Muon and Adam. Muon is 2 times more sample efficient than Adam. **(b)** The MMLU performance of our Moonlight model optimized with Muon and other comparable models. Moonlight advances the Pareto frontier of performance vs training FLOPs.*\n\n## Performance\n\nWe named our lightweight model trained with Muon \"Moonlight\". We compared Moonlight with SOTA public models at similar scale:\n\n- **LLAMA3-3B** is a 3B-parameter dense model trained with 9T tokens\n- **Qwen2.5-3B** is a 3B-parameter dense model trained with 18T tokens\n- **Deepseek-v2-Lite** is a 2.4B/16B-parameter MOE model trained with 5.7T tokens\n\n|             | **Benchmark (Metric)** | **Llama3.2-3B** | **Qwen2.5-3B** | **DSV2-Lite** | **Moonlight** |\n| ----------- | ---------------------- | --------------- | -------------- | ------------- | ------------- |\n|             | Activated Param†       | 2.81B           | 2.77B          | 2.24B         | 2.24B         |\n|             | Total Params†          | 2.81B           | 2.77B          | 15.29B        | 15.29B        |\n|             | Training Tokens        | 9T              | 18T            | 5.7T          | 5.7T          |\n|             | Optimizer              | AdamW           | \\*             | AdamW         | Muon          |\n| **English** | MMLU                   | 54.75           | 65.6           | 58.3          | **70.0**      |\n|             | MMLU-pro               | 25.0            | 34.6           | 25.5          | **42.4**      |\n|             | BBH                    | 46.8            | 56.3           | 44.1          | **65.2**      |\n|             | TriviaQA‡              | 59.6            | 51.1           | 65.1          | **66.3**      |\n| **Code**    | HumanEval              | 28.0            | 42.1           | 29.9          | **48.1**      |\n|             | MBPP                   | 48.7            | 57.1           | 43.2          | **63.8**      |\n| **Math**    | GSM8K                  | 34.0            | **79.1**       | 41.1          | 77.4          |\n|             | MATH                   | 8.5             | 42.6           | 17.1          | **45.3**      |\n|             | CMath                  | -               | 80.0           | 58.4          | **81.1**      |\n| **Chinese** | C-Eval                 | -               | 75.0           | 60.3          | **77.2**      |\n|             | CMMLU                  | -               | 75.0           | 64.3          | **78.2**      |\n\n*Qwen 2 & 2.5 reports didn't disclose their optimizer information. †The reported parameter counts exclude the embedding parameters. ‡We test all listed models with the full set of TriviaQA.*\n\n## Example usage\n\n### Model Download\n\n| **Model**          | **#Total Params** | **#Activated Params** | **Context Length** | **Download Link**                                                               |\n| ------------------ | ----------------- | --------------------- | ------------------ | ------------------------------------------------------------------------------- |\n| Moonlight          | 16B               | 3B                    | 8K                 | [🤗 Hugging Face](https://huggingface.co/moonshotai/Moonlight-16B-A3B)          |\n| Moonlight-Instruct | 16B               | 3B                    | 8K                 | [🤗 Hugging Face](https://huggingface.co/moonshotai/Moonlight-16B-A3B-Instruct) |\n\n### Inference with Hugging Face Transformers\n\nWe introduce how to use our model at inference stage using transformers library. It is recommended to use python=3.10, torch>=2.1.0, and transformers=4.48.2 as the development environment.\n\nFor our pretrained model (Moonlight):\n\n```python\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nmodel_path = \"moonshotai/Moonlight-16B-A3B\"\nmodel = AutoModelForCausalLM.from_pretrained(\n    model_path,\n    torch_dtype=\"auto\",\n    device_map=\"auto\",\n    trust_remote_code=True,\n)\ntokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)\n\nprompt = \"1+1=2, 1+2=\"\ninputs = tokenizer(prompt, return_tensors=\"pt\", padding=True, truncation=True).to(model.device)\ngenerated_ids = model.generate(**inputs, max_new_tokens=100)\nresponse = tokenizer.batch_decode(generated_ids)[0]\nprint(response)\n```\n\nFor our instruct model (Moonlight-Instruct):\n\n```python\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nmodel_path = \"moonshotai/Moonlight-16B-A3B-Instruct\"\nmodel = AutoModelForCausalLM.from_pretrained(\n    model_path,\n    torch_dtype=\"auto\",\n    device_map=\"auto\",\n    trust_remote_code=True\n)\ntokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)\n\nmessages = [\n    {\"role\": \"system\", \"content\": \"You are a helpful assistant provided by Moonshot-AI.\"},\n    {\"role\": \"user\", \"content\": \"Is 123 a prime?\"}\n]\ninput_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors=\"pt\").to(model.device)\ngenerated_ids = model.generate(inputs=input_ids, max_new_tokens=500)\nresponse = tokenizer.batch_decode(generated_ids)[0]\nprint(response)\n```\n\nMoonlight has the same architecture as DeepSeek-V3, which is supported by many popular inference engines, such as VLLM and SGLang. As a result, our model can also be easily deployed using these tools.\n\n### Training\n\n```\n# train qwen-like dense model with muon\npython3 examples/toy_train.py --model qwen --optimizer muon --dataset openwebtext-100k --hidden_size 896 --lr 1e-3\n\n# train qwen-like dense model with adamw\npython3 examples/toy_train.py --model qwen --optimizer adamw --dataset openwebtext-100k --hidden_size 896 --lr 1e-3\n```\n\n## Intermediate Checkpoints\n\nTo support ongoing research efforts, we will soon release our intermediate checkpoints. Coming soon...\n\n## Citation\n\nIf you find Moonlight is useful or want to use in your projects, please kindly cite our paper:\n\n```\n@misc{liu2025muonscalablellmtraining,\n      title={Muon is Scalable for LLM Training}, \n      author={Jingyuan Liu and Jianlin Su and Xingcheng Yao and Zhejun Jiang and Guokun Lai and Yulun Du and Yidao Qin and Weixin Xu and Enzhe Lu and Junjie Yan and Yanru Chen and Huabin Zheng and Yibo Liu and Shaowei Liu and Bohong Yin and Weiran He and Han Zhu and Yuzhi Wang and Jianzhou Wang and Mengnan Dong and Zheng Zhang and Yongsheng Kang and Hao Zhang and Xinran Xu and Yutao Zhang and Yuxin Wu and Xinyu Zhou and Zhilin Yang},\n      year={2025},\n      eprint={2502.16982},\n      archivePrefix={arXiv},\n      primaryClass={cs.LG},\n      url={https://arxiv.org/abs/2502.16982}, \n}\n```\n",
  "authors": [],
  "captured_at": "2026-10-19T01:26:53.998049939+00:00",
  "category": "ai",
  "content_html": "<p><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/MoonshotAI/Moonlight/blob/master/Moonlight.pdf\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://github.com/MoonshotAI/Moonlight/raw/master/figures/banner.png\" alt=\"\" loading=\"eager\" decoding=\"async\" fetchpriority=\"high\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></p>\n<h2>Abstract</h2>\n<p>Recently, the <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/KellerJordan/Muon\">Muon optimizer</a> based on matrix orthogonalization has demonstrated strong results in training small-scale language models, but the scalability to larger models has not been proven. We identify two crucial techniques for scaling up Muon: (1) adding weight decay and (2) carefully adjusting the per-parameter update scale. These techniques allow Muon to work out-of-the-box on large-scale training without the need of hyper-parameter tuning. Scaling law experiments indicate that Muon achieves ∼ 2× computational efficiency compared to AdamW with compute optimal training.</p>\n<p>Based on these improvements, we introduce <strong>Moonlight</strong>, a 3B/16B-parameter Mixture-of-Expert (MoE) model trained with 5.7T tokens using Muon. Our model improves the current Pareto frontier, achieving better performance with much fewer training FLOPs compared to prior models.</p>\n<p>We open-source our distributed Muon implementation that is memory optimal and communication efficient. We also release the pretrained, instruction-tuned, and intermediate checkpoints to support future research.</p>\n<p>Our code is available at <a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/MoonshotAI/Moonlight\">MoonshotAI/Moonlight</a>.</p>\n<h2>Key Ingredients</h2>\n<p>Our work builds upon Muon while systematically identifying and resolving its limitations in large-scale training scenarios. Our technical contributions include:</p>\n<ul>\n<li>\n<p><strong>Analysis for Effective Scaling of Muon</strong>: Through extensive analysis, we identify that weight decay plays a crucial roles in Muon's scalability. Besides, we proposed to keep a consistent update root mean square (RMS) across different matrix and non-matrix parameters through parameter-wise update scale adjustments. Such adjustments significantly enhanced training stability.</p>\n</li>\n<li>\n<p><strong>Efficient Distributed Implementation</strong>: We develop a distributed version of Muon with ZeRO-1 style optimization, achieving optimal memory efficiency and reduced communication overhead while preserving the mathematical properties of the algorithm.</p>\n</li>\n<li>\n<p><strong>Scaling Law Validation</strong>: We performed scaling law research that compares Muon with strong AdamW baselines, and showed the superior performance of Muon (see Figure 1). Based on the scaling law results, Muon achieves comparable performance to AdamW trained counterparts while requiring only approximately 52% of the training FLOPs.</p>\n</li>\n</ul>\n<p><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://github.com/MoonshotAI/Moonlight/blob/master/figures/scaling.png\"><picture class=\"article-picture article-picture-fallback\" style=\"--image-ratio:16 / 9\"><img src=\"https://github.com/MoonshotAI/Moonlight/raw/master/figures/scaling.png\" alt=\"\" loading=\"lazy\" decoding=\"async\" fetchpriority=\"low\" referrerpolicy=\"no-referrer\" class=\"progressive-image\"></picture></a></p>\n<p><em>Scaling up with Muon. <strong>(a)</strong> Scaling law experiments comparing Muon and Adam. Muon is 2 times more sample efficient than Adam. <strong>(b)</strong> The MMLU performance of our Moonlight model optimized with Muon and other comparable models. Moonlight advances the Pareto frontier of performance vs training FLOPs.</em></p>\n<h2>Performance</h2>\n<p>We named our lightweight model trained with Muon &quot;Moonlight&quot;. We compared Moonlight with SOTA public models at similar scale:</p>\n<ul>\n<li><strong>LLAMA3-3B</strong> is a 3B-parameter dense model trained with 9T tokens</li>\n<li><strong>Qwen2.5-3B</strong> is a 3B-parameter dense model trained with 18T tokens</li>\n<li><strong>Deepseek-v2-Lite</strong> is a 2.4B/16B-parameter MOE model trained with 5.7T tokens</li>\n</ul>\n<table>\n<thead>\n<tr>\n<th></th>\n<th><strong>Benchmark (Metric)</strong></th>\n<th><strong>Llama3.2-3B</strong></th>\n<th><strong>Qwen2.5-3B</strong></th>\n<th><strong>DSV2-Lite</strong></th>\n<th><strong>Moonlight</strong></th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td></td>\n<td>Activated Param†</td>\n<td>2.81B</td>\n<td>2.77B</td>\n<td>2.24B</td>\n<td>2.24B</td>\n</tr>\n<tr>\n<td></td>\n<td>Total Params†</td>\n<td>2.81B</td>\n<td>2.77B</td>\n<td>15.29B</td>\n<td>15.29B</td>\n</tr>\n<tr>\n<td></td>\n<td>Training Tokens</td>\n<td>9T</td>\n<td>18T</td>\n<td>5.7T</td>\n<td>5.7T</td>\n</tr>\n<tr>\n<td></td>\n<td>Optimizer</td>\n<td>AdamW</td>\n<td>*</td>\n<td>AdamW</td>\n<td>Muon</td>\n</tr>\n<tr>\n<td><strong>English</strong></td>\n<td>MMLU</td>\n<td>54.75</td>\n<td>65.6</td>\n<td>58.3</td>\n<td><strong>70.0</strong></td>\n</tr>\n<tr>\n<td></td>\n<td>MMLU-pro</td>\n<td>25.0</td>\n<td>34.6</td>\n<td>25.5</td>\n<td><strong>42.4</strong></td>\n</tr>\n<tr>\n<td></td>\n<td>BBH</td>\n<td>46.8</td>\n<td>56.3</td>\n<td>44.1</td>\n<td><strong>65.2</strong></td>\n</tr>\n<tr>\n<td></td>\n<td>TriviaQA‡</td>\n<td>59.6</td>\n<td>51.1</td>\n<td>65.1</td>\n<td><strong>66.3</strong></td>\n</tr>\n<tr>\n<td><strong>Code</strong></td>\n<td>HumanEval</td>\n<td>28.0</td>\n<td>42.1</td>\n<td>29.9</td>\n<td><strong>48.1</strong></td>\n</tr>\n<tr>\n<td></td>\n<td>MBPP</td>\n<td>48.7</td>\n<td>57.1</td>\n<td>43.2</td>\n<td><strong>63.8</strong></td>\n</tr>\n<tr>\n<td><strong>Math</strong></td>\n<td>GSM8K</td>\n<td>34.0</td>\n<td><strong>79.1</strong></td>\n<td>41.1</td>\n<td>77.4</td>\n</tr>\n<tr>\n<td></td>\n<td>MATH</td>\n<td>8.5</td>\n<td>42.6</td>\n<td>17.1</td>\n<td><strong>45.3</strong></td>\n</tr>\n<tr>\n<td></td>\n<td>CMath</td>\n<td>-</td>\n<td>80.0</td>\n<td>58.4</td>\n<td><strong>81.1</strong></td>\n</tr>\n<tr>\n<td><strong>Chinese</strong></td>\n<td>C-Eval</td>\n<td>-</td>\n<td>75.0</td>\n<td>60.3</td>\n<td><strong>77.2</strong></td>\n</tr>\n<tr>\n<td></td>\n<td>CMMLU</td>\n<td>-</td>\n<td>75.0</td>\n<td>64.3</td>\n<td><strong>78.2</strong></td>\n</tr>\n</tbody>\n</table>\n<p><em>Qwen 2 &amp; 2.5 reports didn't disclose their optimizer information. †The reported parameter counts exclude the embedding parameters. ‡We test all listed models with the full set of TriviaQA.</em></p>\n<h2>Example usage</h2>\n<h3>Model Download</h3>\n<table>\n<thead>\n<tr>\n<th><strong>Model</strong></th>\n<th><strong>#Total Params</strong></th>\n<th><strong>#Activated Params</strong></th>\n<th><strong>Context Length</strong></th>\n<th><strong>Download Link</strong></th>\n</tr>\n</thead>\n<tbody>\n<tr>\n<td>Moonlight</td>\n<td>16B</td>\n<td>3B</td>\n<td>8K</td>\n<td><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/moonshotai/Moonlight-16B-A3B\">🤗 Hugging Face</a></td>\n</tr>\n<tr>\n<td>Moonlight-Instruct</td>\n<td>16B</td>\n<td>3B</td>\n<td>8K</td>\n<td><a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://huggingface.co/moonshotai/Moonlight-16B-A3B-Instruct\">🤗 Hugging Face</a></td>\n</tr>\n</tbody>\n</table>\n<h3>Inference with Hugging Face Transformers</h3>\n<p>We introduce how to use our model at inference stage using transformers library. It is recommended to use python=3.10, torch&gt;=2.1.0, and transformers=4.48.2 as the development environment.</p>\n<p>For our pretrained model (Moonlight):</p>\n<pre><code class=\"language-python\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-from syntax-python\">from</span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-meta syntax-import-source syntax-python\"> <span class=\"syntax-meta syntax-import-path syntax-python\"><span class=\"syntax-meta syntax-import-name syntax-python\">transformers</span></span> <span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-python\">import</span></span></span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"> <span class=\"syntax-meta syntax-generic-name syntax-python\">AutoModelForCausalLM</span><span class=\"syntax-punctuation syntax-separator syntax-import-list syntax-python\">,</span> <span class=\"syntax-meta syntax-generic-name syntax-python\">AutoTokenizer</span></span>\n\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">model_path</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">moonshotai/Moonlight-16B-A3B<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span>\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">model</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">AutoModelForCausalLM</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">from_pretrained</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\">\n    <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">model_path</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n    <span class=\"syntax-variable syntax-parameter syntax-python\">torch_dtype</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">auto<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n    <span class=\"syntax-variable syntax-parameter syntax-python\">device_map</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">auto<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n    <span class=\"syntax-variable syntax-parameter syntax-python\">trust_remote_code</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-language syntax-python\">True</span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">tokenizer</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">AutoTokenizer</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">from_pretrained</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">model_path</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">trust_remote_code</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-language syntax-python\">True</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">prompt</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">1+1=2, 1+2=<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span>\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">inputs</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">tokenizer</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">prompt</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">return_tensors</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">pt<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">padding</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-language syntax-python\">True</span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">truncation</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-language syntax-python\">True</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span><span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">to</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">model</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">device</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">generated_ids</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">model</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">generate</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-keyword syntax-operator syntax-arithmetic syntax-python\">*</span><span class=\"syntax-keyword syntax-operator syntax-arithmetic syntax-python\">*</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">inputs</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">max_new_tokens</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">100</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">response</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">tokenizer</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">batch_decode</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">generated_ids</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-begin syntax-python\">[</span></span><span class=\"syntax-meta syntax-item-access syntax-arguments syntax-python\"><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">0</span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-end syntax-python\">]</span></span>\n<span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-support syntax-function syntax-builtin syntax-python\">print</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">response</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n</span></span></code></pre>\n<p>For our instruct model (Moonlight-Instruct):</p>\n<pre><code class=\"language-python\"><span class=\"code-snippet\" data-language=\"Python\"><span class=\"syntax-source syntax-python\"><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-from syntax-python\">from</span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-meta syntax-import-source syntax-python\"> <span class=\"syntax-meta syntax-import-path syntax-python\"><span class=\"syntax-meta syntax-import-name syntax-python\">transformers</span></span> <span class=\"syntax-meta syntax-statement syntax-import syntax-python\"><span class=\"syntax-keyword syntax-control syntax-import syntax-python\">import</span></span></span></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"></span><span class=\"syntax-meta syntax-statement syntax-import syntax-python\"> <span class=\"syntax-meta syntax-generic-name syntax-python\">AutoModelForCausalLM</span><span class=\"syntax-punctuation syntax-separator syntax-import-list syntax-python\">,</span> <span class=\"syntax-meta syntax-generic-name syntax-python\">AutoTokenizer</span></span>\n\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">model_path</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">moonshotai/Moonlight-16B-A3B-Instruct<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span>\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">model</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">AutoModelForCausalLM</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">from_pretrained</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\">\n    <span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">model_path</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n    <span class=\"syntax-variable syntax-parameter syntax-python\">torch_dtype</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">auto<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n    <span class=\"syntax-variable syntax-parameter syntax-python\">device_map</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">auto<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span>\n    <span class=\"syntax-variable syntax-parameter syntax-python\">trust_remote_code</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-language syntax-python\">True</span>\n</span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">tokenizer</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">AutoTokenizer</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">from_pretrained</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">model_path</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">trust_remote_code</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-language syntax-python\">True</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">messages</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-structure syntax-list syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-list syntax-begin syntax-python\">[</span>\n    <span class=\"syntax-meta syntax-structure syntax-dictionary-or-set syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-dictionary-or-set syntax-begin syntax-python\">{</span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">role<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-python\">:</span> <span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">system<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-separator syntax-dictionary-or-set syntax-python\">,</span> <span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">content<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-python\">:</span> <span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">You are a helpful assistant provided by Moonshot-AI.<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-section syntax-dictionary-or-set syntax-end syntax-python\">}</span></span><span class=\"syntax-punctuation syntax-separator syntax-list syntax-python\">,</span>\n    <span class=\"syntax-meta syntax-structure syntax-dictionary-or-set syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-dictionary-or-set syntax-begin syntax-python\">{</span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">role<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-python\">:</span> <span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">user<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-separator syntax-dictionary-or-set syntax-python\">,</span> <span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">content<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-separator syntax-key-value syntax-python\">:</span> <span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">Is 123 a prime?<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span><span class=\"syntax-punctuation syntax-section syntax-dictionary-or-set syntax-end syntax-python\">}</span></span>\n<span class=\"syntax-punctuation syntax-section syntax-list syntax-end syntax-python\">]</span></span>\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">input_ids</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">tokenizer</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">apply_chat_template</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">messages</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">add_generation_prompt</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-language syntax-python\">True</span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">return_tensors</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\"><span class=\"syntax-punctuation syntax-definition syntax-string syntax-begin syntax-python\">&quot;</span></span></span><span class=\"syntax-meta syntax-string syntax-python\"><span class=\"syntax-string syntax-quoted syntax-double syntax-python\">pt<span class=\"syntax-punctuation syntax-definition syntax-string syntax-end syntax-python\">&quot;</span></span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span><span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">to</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">model</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span><span class=\"syntax-meta syntax-generic-name syntax-python\">device</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">generated_ids</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">model</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">generate</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-variable syntax-parameter syntax-python\">inputs</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">input_ids</span></span><span class=\"syntax-punctuation syntax-separator syntax-arguments syntax-python\">,</span> <span class=\"syntax-variable syntax-parameter syntax-python\">max_new_tokens</span><span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">500</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n<span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">response</span></span> <span class=\"syntax-keyword syntax-operator syntax-assignment syntax-python\">=</span> <span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">tokenizer</span><span class=\"syntax-punctuation syntax-accessor syntax-dot syntax-python\">.</span></span><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-variable syntax-function syntax-python\">batch_decode</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">generated_ids</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-begin syntax-python\">[</span></span><span class=\"syntax-meta syntax-item-access syntax-arguments syntax-python\"><span class=\"syntax-constant syntax-numeric syntax-integer syntax-decimal syntax-python\">0</span></span><span class=\"syntax-meta syntax-item-access syntax-python\"><span class=\"syntax-punctuation syntax-section syntax-brackets syntax-end syntax-python\">]</span></span>\n<span class=\"syntax-meta syntax-function-call syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-support syntax-function syntax-builtin syntax-python\">print</span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-begin syntax-python\">(</span><span class=\"syntax-meta syntax-function-call syntax-arguments syntax-python\"><span class=\"syntax-meta syntax-qualified-name syntax-python\"><span class=\"syntax-meta syntax-generic-name syntax-python\">response</span></span></span><span class=\"syntax-punctuation syntax-section syntax-arguments syntax-end syntax-python\">)</span></span>\n</span></span></code></pre>\n<p>Moonlight has the same architecture as DeepSeek-V3, which is supported by many popular inference engines, such as VLLM and SGLang. As a result, our model can also be easily deployed using these tools.</p>\n<h3>Training</h3>\n<pre><code><span class=\"code-snippet\"># train qwen-like dense model with muon\npython3 examples/toy_train.py --model qwen --optimizer muon --dataset openwebtext-100k --hidden_size 896 --lr 1e-3\n\n# train qwen-like dense model with adamw\npython3 examples/toy_train.py --model qwen --optimizer adamw --dataset openwebtext-100k --hidden_size 896 --lr 1e-3\n</span></code></pre>\n<h2>Intermediate Checkpoints</h2>\n<p>To support ongoing research efforts, we will soon release our intermediate checkpoints. Coming soon...</p>\n<h2>Citation</h2>\n<p>If you find Moonlight is useful or want to use in your projects, please kindly cite our paper:</p>\n<pre><code><span class=\"code-snippet\">@misc{liu2025muonscalablellmtraining,\n      title={Muon is Scalable for LLM Training}, \n      author={Jingyuan Liu and Jianlin Su and Xingcheng Yao and Zhejun Jiang and Guokun Lai and Yulun Du and Yidao Qin and Weixin Xu and Enzhe Lu and Junjie Yan and Yanru Chen and Huabin Zheng and Yibo Liu and Shaowei Liu and Bohong Yin and Weiran He and Han Zhu and Yuzhi Wang and Jianzhou Wang and Mengnan Dong and Zheng Zhang and Yongsheng Kang and Hao Zhang and Xinran Xu and Yutao Zhang and Yuxin Wu and Xinyu Zhou and Zhilin Yang},\n      year={2025},\n      eprint={2502.16982},\n      archivePrefix={arXiv},\n      primaryClass={cs.LG},\n      url={https://arxiv.org/abs/2502.16982}, \n}\n</span></code></pre>\n",
  "content_markdown": "[![](https://github.com/MoonshotAI/Moonlight/raw/master/figures/banner.png)](https://github.com/MoonshotAI/Moonlight/blob/master/Moonlight.pdf)\n\n## Abstract\n\nRecently, the [Muon optimizer](https://github.com/KellerJordan/Muon) based on matrix orthogonalization has demonstrated strong results in training small-scale language models, but the scalability to larger models has not been proven. We identify two crucial techniques for scaling up Muon: (1) adding weight decay and (2) carefully adjusting the per-parameter update scale. These techniques allow Muon to work out-of-the-box on large-scale training without the need of hyper-parameter tuning. Scaling law experiments indicate that Muon achieves ∼ 2× computational efficiency compared to AdamW with compute optimal training.\n\nBased on these improvements, we introduce **Moonlight**, a 3B/16B-parameter Mixture-of-Expert (MoE) model trained with 5.7T tokens using Muon. Our model improves the current Pareto frontier, achieving better performance with much fewer training FLOPs compared to prior models.\n\nWe open-source our distributed Muon implementation that is memory optimal and communication efficient. We also release the pretrained, instruction-tuned, and intermediate checkpoints to support future research.\n\nOur code is available at [MoonshotAI/Moonlight](https://github.com/MoonshotAI/Moonlight).\n\n## Key Ingredients\n\nOur work builds upon Muon while systematically identifying and resolving its limitations in large-scale training scenarios. Our technical contributions include:\n\n- **Analysis for Effective Scaling of Muon**: Through extensive analysis, we identify that weight decay plays a crucial roles in Muon's scalability. Besides, we proposed to keep a consistent update root mean square (RMS) across different matrix and non-matrix parameters through parameter-wise update scale adjustments. Such adjustments significantly enhanced training stability.\n\n- **Efficient Distributed Implementation**: We develop a distributed version of Muon with ZeRO-1 style optimization, achieving optimal memory efficiency and reduced communication overhead while preserving the mathematical properties of the algorithm.\n\n- **Scaling Law Validation**: We performed scaling law research that compares Muon with strong AdamW baselines, and showed the superior performance of Muon (see Figure 1). Based on the scaling law results, Muon achieves comparable performance to AdamW trained counterparts while requiring only approximately 52% of the training FLOPs.\n\n[![](https://github.com/MoonshotAI/Moonlight/raw/master/figures/scaling.png)](https://github.com/MoonshotAI/Moonlight/blob/master/figures/scaling.png)\n\n*Scaling up with Muon. **(a)** Scaling law experiments comparing Muon and Adam. Muon is 2 times more sample efficient than Adam. **(b)** The MMLU performance of our Moonlight model optimized with Muon and other comparable models. Moonlight advances the Pareto frontier of performance vs training FLOPs.*\n\n## Performance\n\nWe named our lightweight model trained with Muon \"Moonlight\". We compared Moonlight with SOTA public models at similar scale:\n\n- **LLAMA3-3B** is a 3B-parameter dense model trained with 9T tokens\n- **Qwen2.5-3B** is a 3B-parameter dense model trained with 18T tokens\n- **Deepseek-v2-Lite** is a 2.4B/16B-parameter MOE model trained with 5.7T tokens\n\n|             | **Benchmark (Metric)** | **Llama3.2-3B** | **Qwen2.5-3B** | **DSV2-Lite** | **Moonlight** |\n| ----------- | ---------------------- | --------------- | -------------- | ------------- | ------------- |\n|             | Activated Param†       | 2.81B           | 2.77B          | 2.24B         | 2.24B         |\n|             | Total Params†          | 2.81B           | 2.77B          | 15.29B        | 15.29B        |\n|             | Training Tokens        | 9T              | 18T            | 5.7T          | 5.7T          |\n|             | Optimizer              | AdamW           | \\*             | AdamW         | Muon          |\n| **English** | MMLU                   | 54.75           | 65.6           | 58.3          | **70.0**      |\n|             | MMLU-pro               | 25.0            | 34.6           | 25.5          | **42.4**      |\n|             | BBH                    | 46.8            | 56.3           | 44.1          | **65.2**      |\n|             | TriviaQA‡              | 59.6            | 51.1           | 65.1          | **66.3**      |\n| **Code**    | HumanEval              | 28.0            | 42.1           | 29.9          | **48.1**      |\n|             | MBPP                   | 48.7            | 57.1           | 43.2          | **63.8**      |\n| **Math**    | GSM8K                  | 34.0            | **79.1**       | 41.1          | 77.4          |\n|             | MATH                   | 8.5             | 42.6           | 17.1          | **45.3**      |\n|             | CMath                  | -               | 80.0           | 58.4          | **81.1**      |\n| **Chinese** | C-Eval                 | -               | 75.0           | 60.3          | **77.2**      |\n|             | CMMLU                  | -               | 75.0           | 64.3          | **78.2**      |\n\n*Qwen 2 & 2.5 reports didn't disclose their optimizer information. †The reported parameter counts exclude the embedding parameters. ‡We test all listed models with the full set of TriviaQA.*\n\n## Example usage\n\n### Model Download\n\n| **Model**          | **#Total Params** | **#Activated Params** | **Context Length** | **Download Link**                                                               |\n| ------------------ | ----------------- | --------------------- | ------------------ | ------------------------------------------------------------------------------- |\n| Moonlight          | 16B               | 3B                    | 8K                 | [🤗 Hugging Face](https://huggingface.co/moonshotai/Moonlight-16B-A3B)          |\n| Moonlight-Instruct | 16B               | 3B                    | 8K                 | [🤗 Hugging Face](https://huggingface.co/moonshotai/Moonlight-16B-A3B-Instruct) |\n\n### Inference with Hugging Face Transformers\n\nWe introduce how to use our model at inference stage using transformers library. It is recommended to use python=3.10, torch>=2.1.0, and transformers=4.48.2 as the development environment.\n\nFor our pretrained model (Moonlight):\n\n```python\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nmodel_path = \"moonshotai/Moonlight-16B-A3B\"\nmodel = AutoModelForCausalLM.from_pretrained(\n    model_path,\n    torch_dtype=\"auto\",\n    device_map=\"auto\",\n    trust_remote_code=True,\n)\ntokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)\n\nprompt = \"1+1=2, 1+2=\"\ninputs = tokenizer(prompt, return_tensors=\"pt\", padding=True, truncation=True).to(model.device)\ngenerated_ids = model.generate(**inputs, max_new_tokens=100)\nresponse = tokenizer.batch_decode(generated_ids)[0]\nprint(response)\n```\n\nFor our instruct model (Moonlight-Instruct):\n\n```python\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nmodel_path = \"moonshotai/Moonlight-16B-A3B-Instruct\"\nmodel = AutoModelForCausalLM.from_pretrained(\n    model_path,\n    torch_dtype=\"auto\",\n    device_map=\"auto\",\n    trust_remote_code=True\n)\ntokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)\n\nmessages = [\n    {\"role\": \"system\", \"content\": \"You are a helpful assistant provided by Moonshot-AI.\"},\n    {\"role\": \"user\", \"content\": \"Is 123 a prime?\"}\n]\ninput_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors=\"pt\").to(model.device)\ngenerated_ids = model.generate(inputs=input_ids, max_new_tokens=500)\nresponse = tokenizer.batch_decode(generated_ids)[0]\nprint(response)\n```\n\nMoonlight has the same architecture as DeepSeek-V3, which is supported by many popular inference engines, such as VLLM and SGLang. As a result, our model can also be easily deployed using these tools.\n\n### Training\n\n```\n# train qwen-like dense model with muon\npython3 examples/toy_train.py --model qwen --optimizer muon --dataset openwebtext-100k --hidden_size 896 --lr 1e-3\n\n# train qwen-like dense model with adamw\npython3 examples/toy_train.py --model qwen --optimizer adamw --dataset openwebtext-100k --hidden_size 896 --lr 1e-3\n```\n\n## Intermediate Checkpoints\n\nTo support ongoing research efforts, we will soon release our intermediate checkpoints. Coming soon...\n\n## Citation\n\nIf you find Moonlight is useful or want to use in your projects, please kindly cite our paper:\n\n```\n@misc{liu2025muonscalablellmtraining,\n      title={Muon is Scalable for LLM Training}, \n      author={Jingyuan Liu and Jianlin Su and Xingcheng Yao and Zhejun Jiang and Guokun Lai and Yulun Du and Yidao Qin and Weixin Xu and Enzhe Lu and Junjie Yan and Yanru Chen and Huabin Zheng and Yibo Liu and Shaowei Liu and Bohong Yin and Weiran He and Han Zhu and Yuzhi Wang and Jianzhou Wang and Mengnan Dong and Zheng Zhang and Yongsheng Kang and Hao Zhang and Xinran Xu and Yutao Zhang and Yuxin Wu and Xinyu Zhou and Zhilin Yang},\n      year={2025},\n      eprint={2502.16982},\n      archivePrefix={arXiv},\n      primaryClass={cs.LG},\n      url={https://arxiv.org/abs/2502.16982}, \n}\n```\n",
  "dateCreated": "2026-10-19T01:26:53.998049939+00:00",
  "date_created": "2025-02-23T00:00:00+00:00",
  "date_modified": null,
  "date_published": "2025-02-23T00:00:00+00:00",
  "description": "2025-02-23",
  "external_url": "https://github.com/MoonshotAI/Moonlight",
  "git": {
    "commit": "ebb0941c5c06a2ab2377342d607535262a3e969e",
    "path": "items/kimi-com/2025/02/2025-02-23-muon-is-scalable-for-llm-training.md",
    "permalink": "https://github.com/aymericbeaumet/aggr-instance/blob/ebb0941c5c06a2ab2377342d607535262a3e969e/items/kimi-com/2025/02/2025-02-23-muon-is-scalable-for-llm-training.md"
  },
  "id": "urn:aggr:item:79b48f27241a8dfbf4acb69eb367cfc8f8b1eb26",
  "isBasedOn": {
    "@id": "https://github.com/MoonshotAI/Moonlight"
  },
  "mainEntity": {
    "@id": "https://github.com/MoonshotAI/Moonlight",
    "@type": "CreativeWork",
    "archivedAt": {
      "@id": "https://aggr.aymericbeaumet.com/items/kimi-com/2025-02-23-muon-is-scalable-for-llm-training/#webpage"
    },
    "datePublished": "2025-02-23T00:00:00+00:00",
    "headline": "Muon is Scalable for LLM Training",
    "image": {
      "@type": "ImageObject",
      "caption": "Muon is Scalable for LLM Training",
      "height": 90,
      "url": "https://kimi-file.kimi.ai/prod-chat-kimi/kfs/4/2/2026-03-25/1d71teo1l51jas5dtnj8g?x-tos-process=image%2Fauto-orient%2C1%2Fstrip%2Fignore-error%2C1",
      "width": 160
    },
    "timeRequired": "PT5M",
    "url": "https://github.com/MoonshotAI/Moonlight",
    "wordCount": 949
  },
  "name": "Archived snapshot: Muon is Scalable for LLM Training",
  "preview": {
    "alt": "Muon is Scalable for LLM Training",
    "color": null,
    "height": 90,
    "placeholder": {
      "data_url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=",
      "hash": ""
    },
    "url": "https://kimi-file.kimi.ai/prod-chat-kimi/kfs/4/2/2026-03-25/1d71teo1l51jas5dtnj8g?x-tos-process=image%2Fauto-orient%2C1%2Fstrip%2Fignore-error%2C1",
    "width": 160
  },
  "reading_minutes": 5,
  "replicated_at": null,
  "source": "kimi-com",
  "source_name": "Kimi Research Blog | Moonshot AI",
  "summary": "2025-02-23",
  "tags": [],
  "temporalCoverage": "2026-10-19T01:26:53.998049939+00:00",
  "title": "Muon is Scalable for LLM Training",
  "url": "https://aggr.aymericbeaumet.com/items/kimi-com/2025-02-23-muon-is-scalable-for-llm-training/",
  "word_count": 949
}