Moonlight-16B-A3B-Instruct-transformers-5.13.0

Fork: transformers==5.13.0 compatibility

This is a fork of moonshotai/Moonlight-16B-A3B-Instruct whose only differences from upstream are edits to the custom modeling_deepseek.py file so that the model loads and generates correctly on transformers==5.13.0. The safetensors weights are identical to upstream. Verified working with torch==2.6.0+cu124 and transformers==5.13.0.

The upstream file targets transformers==4.48.2. Between 4.48 and 5.13 several internal APIs changed or were removed, and the transformers-5.x meta-device-first load path caused an uninitialized RoPE buffer that silently broke the model. Concretely, the following was changed in modeling_deepseek.py:

• RoPE inv_freq buffer (the critical fix). In transformers ≥ 5, from_pretrained initializes modules on the meta device before loading the checkpoint. Buffers registered with persistent=False (like DeepseekV3RotaryEmbedding.inv_freq) are not populated from state_dict, so inv_freq remained as meta-device zeros — every position produced cos=1, sin=0 and RoPE was silently disabled. DeepseekV3RotaryEmbedding now lazily recomputes inv_freq on the real device inside _set_cos_sin_cache if it is on meta or all-zero. Without this fix the model produces fluent but position-blind gibberish.
• rope_scaling default schema. transformers ≥ 4.45 auto-populates config.rope_scaling with {'rope_type': 'default', ...} even when the config file has no rope_scaling key. _init_rope (and the softmax-scale branch) now call a new _rope_scaling_active() helper that treats rope_type == 'default' (or a missing factor) as "no scaling," matching the pre-5.x behavior for Moonlight.
• Cache.seen_tokens removed. prepare_inputs_for_generation now uses getattr(past_key_values, "seen_tokens", cache_length) where cache_length comes from get_seq_length().
• Cache.get_usable_length(...) removed. All call sites were replaced with the equivalent get_seq_length(...) (per-layer where applicable). DynamicCache has no cache size limit, so the replacement is exact.
• DynamicCache.get_max_cache_shape() changed return value. It now returns -1 (previously None) for unbounded caches. prepare_inputs_for_generation treats any non-positive value as "unbounded," which prevents an erroneous attention_mask[:, -(-1):] truncation that dropped the first token.
• Full-sequence position_ids in the decode loop. The new _sample loop passes position IDs covering the whole sequence even when only one new token is being processed. prepare_inputs_for_generation now trims position_ids down to input_ids' length so RoPE indexing stays consistent.

The README's two "Inference with Hugging Face Transformers" snippets below have been updated for the new API (the second snippet uses return_dict=True with apply_chat_template, which is required in transformers ≥ 5) and both were re-run end-to-end against this fork — the actual outputs are shown after each snippet.

---

<div align="center"> <a href="https://github.com/MoonshotAI/Moonlight"><img width="80%" src="figures/banner.png"></a> </div> <!-- # Muon is Scalable For LLM Training --> <div align="center"> <a href="https://github.com/MoonshotAI/Moonlight/blob/master/Moonlight.pdf" ><img src="figures/logo.png" height="16" width="16" style="display: inline-block; vertical-align: middle; margin: 2px;"><b style="display: inline-block;"> Tech Report</b></a> | <a href="https://huggingface.co/moonshotai/Moonlight-16B-A3B"><img src="https://huggingface.co/front/assets/huggingface_logo-noborder.svg" height="16" width="16" style="display: inline-block; vertical-align: middle; margin: 2px;"><b style="display: inline-block;"> HuggingFace</b></a> | <a href="#"><img src="figures/megatron.png" height="16" width="16" style="display: inline-block; vertical-align: middle; margin: 2px;"><b style="display: inline-block;">Megatron(coming soon)</b></a> </div>

Abstract

Recently, the Muon optimizer 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:

• Weight Decay: Critical for scaling to larger models
• Consistent RMS Updates: Enforcing a consistent root mean square on model updates

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 is $\sim2\times$ more sample efficient than Adam with compute optimal training.

Based 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.

We open-source our Muon implementation that is memory optimal and communication efficient. We also release the pretrained, instruction-tuned, and intermediate checkpoints to support future research.

Our code is available at MoonshotAI/Moonlight.

Key Ingredients

Our work builds upon Muon while systematically identifying and resolving its limitations in large-scale training scenarios. Our technical contributions include:

• 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.

• 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.

• 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.

<div align="center"> <img width="90%" src="figures/scaling.png"> <p><em>Scaling up with Muon. <b>(a)</b> Scaling law experiments comparing Muon and Adam. Muon is 2 times more sample efficient than Adam. <b>(b)</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.</em></p> </div>

Performance

We compared Moonlight with SOTA public models at similar scale:

• LLAMA3-3B is a 3B-parameter dense model trained with 9T tokens
• Qwen2.5-3B is a 3B-parameter dense model trained with 18T tokens
• Deepseek-v2-Lite is a 2.4B/16B-parameter MOE model trained with 5.7T tokens

<div align="center"> </div>

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.

Example usage

Model Download

<div align="center"> </div>

Inference with Hugging Face Transformers

This fork targets transformers==5.13.0. Recommended environment: python=3.10+, torch>=2.1 (verified on torch==2.6.0+cu124), transformers==5.13.0. Both snippets below were re-run end-to-end against this fork and the verified outputs are shown after each.

For the base pretrained model (Moonlight-16B-A3B) — apply the same modeling_deepseek.py patches described above if you want to run the base model on transformers==5.13.0:

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "melhoushi/Moonlight-16B-A3B-Instruct-transformers-5.13.0"  # this fork; the base variant follows the same pattern
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto",
    trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

prompt = "1+1=2, 1+2="
inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to(model.device)
generated_ids = model.generate(**inputs, max_new_tokens=100)
response = tokenizer.batch_decode(generated_ids)[0]
print(response)

Verified output (transformers==5.13.0, torch==2.6.0+cu124, single H100):

1+1=2, 1+2=3, 2+1=3, 2+2=4, 3+1=4, 3+2=5, 3+3=6, 4+1=5, 4+2=6, 4+3=7, 4+4=8, 5+1=6, 5+2=7, 5+3=8, 5+4=9, 5+5=10,

For the instruct model (this fork). Note the change from the upstream snippet: apply_chat_template(..., return_tensors="pt") returns a BatchEncoding in transformers>=5, not a bare tensor, so we pass return_dict=True and unpack it into generate(**inputs, ...):

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "melhoushi/Moonlight-16B-A3B-Instruct-transformers-5.13.0"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto",
    trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

messages = [
    {"role": "system", "content": "You are a helpful assistant provided by Moonshot-AI."},
    {"role": "user", "content": "Is 123 a prime?"},
]
inputs = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
).to(model.device)
generated_ids = model.generate(**inputs, max_new_tokens=500)
response = tokenizer.batch_decode(generated_ids)[0]
print(response)

Verified output (transformers==5.13.0, torch==2.6.0+cu124, single H100):

<|im_system|>system<|im_middle|>You are a helpful assistant provided by Moonshot-AI.<|im_end|><|im_user|>user<|im_middle|>Is 123 a prime?<|im_end|><|im_assistant|>assistant<|im_middle|>To determine if 123 is a prime number, we need to check if it has any divisors other than 1 and itself. A prime number is a number greater than 1 that has no positive divisors other than 1 and itself.

Let's check for divisibility by smaller prime numbers:

1. **Divisibility by 2**: 123 is an odd number, so it is not divisible by 2.
2. **Divisibility by 3**: Sum the digits of 123: \(1 + 2 + 3 = 6\). Since 6 is divisible by 3, 123 is also divisible by 3.

Since 123 is divisible by 3, it is not a prime number.<|im_end|>

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.

Citation

If you find Moonlight is useful or want to use in your projects, please kindly cite our paper:

@misc{liu2025muonscalablellmtraining,
      title={Muon is Scalable for LLM Training}, 
      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},
      year={2025},
      eprint={2502.16982},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2502.16982}, 
}