update modeling and readme
Browse files- README.md +274 -3
- configuration_bailing_moe_v2.py +2 -2
- modeling_bailing_moe_v2.py +13 -11
README.md
CHANGED
|
@@ -1,3 +1,274 @@
|
|
| 1 |
-
---
|
| 2 |
-
license: mit
|
| 3 |
-
---
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: mit
|
| 3 |
+
---
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
|
| 7 |
+
<p align="center">
|
| 8 |
+
<img src="https://mdn.alipayobjects.com/huamei_qa8qxu/afts/img/A*4QxcQrBlTiAAAAAAQXAAAAgAemJ7AQ/original" width="100"/>
|
| 9 |
+
<p>
|
| 10 |
+
|
| 11 |
+
<p align="center">🤗 <a href="https://huggingface.co/inclusionAI">Hugging Face</a>   |   🤖 <a href="https://modelscope.cn/organization/inclusionAI">ModelScope</a></p>
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
## Introduction
|
| 15 |
+
|
| 16 |
+
Today, we are excited to announce the open-sourcing of __Ling 2.0__ — a family of MoE-based large language models that combine __SOTA performance__ with __high efficiency__.
|
| 17 |
+
The first released version, Ling-mini-2.0, is compact yet powerful. It has __16B total parameters__, but only __1.4B__ are activated per input token (non-embedding 789M). Trained on more than __20T tokens__ of high-quality data and enhanced through multi-stage supervised fine-tuning and reinforcement learning, Ling-mini-2.0 achieves remarkable improvements in complex reasoning and instruction following. With just 1.4B activated parameters, it still reaches the top-tier level of sub-10B dense LLMs and even matches or surpasses much larger MoE models.
|
| 18 |
+
|
| 19 |
+
<p align="center"><img src="https://mdn.alipayobjects.com/huamei_fi95qp/afts/img/2NKZS5LVXzcAAAAASBAAAAgADkZ7AQFr/fmt.webp" /></p>
|
| 20 |
+
|
| 21 |
+
### Strong General and Professional Reasoning
|
| 22 |
+
|
| 23 |
+
We evaluated Ling-mini-2.0 on challenging general reasoning tasks in coding (LiveCodeBench, CodeForces) and mathematics (AIME 2025, HMMT 2025), as well as knowledge-intensive reasoning tasks across multiple domains (MMLU-Pro, Humanity's Last Exam). Compared with sub-10B dense models (e.g., Qwen3-4B-instruct-2507, Qwen3-8B-nothinking) and larger-scale MoE models (Ernie-4.5-21B-A3B-PT, GPT-OSS-20B/low), Ling-mini-2.0 demonstrated outstanding overall reasoning capabilities.
|
| 24 |
+
|
| 25 |
+
### 7× Equivalent Dense Performance Leverage
|
| 26 |
+
|
| 27 |
+
Guided by [Ling Scaling Laws](https://arxiv.org/abs/2507.17702), Ling 2.0 adopts a __1/32 activation ratio__ MoE architecture, with empirically optimized design choices in expert granularity, shared expert ratio, attention ratio, aux-loss free + sigmoid routing strategy, MTP loss, QK-Norm, half RoPE, and more. This enables small-activation MoE models to achieve over __7× equivalent dense performance__. In other words, __Ling-mini-2.0 with only 1.4B activated parameters (non-embedding 789M) can deliver performance equivalent to a 7–8B dense model__.
|
| 28 |
+
|
| 29 |
+
### High-speed Generation at 300+ token/s
|
| 30 |
+
|
| 31 |
+
<p align="center"><img src="https://mdn.alipayobjects.com/huamei_fi95qp/afts/img/bnxIRaK9tzcAAAAAgSAAAAgADkZ7AQFr/original" /></p>
|
| 32 |
+
|
| 33 |
+
The highly sparse small-activation MoE architecture also delivers significant training and inference efficiency. In simple QA scenarios (within 2000 tokens), __Ling-mini-2.0 generates at 300+ token/s (on H20 deployment)__ — more than __2× faster__ than an 8B dense model. Ling-mini-2.0 is able to handle __128K context length__ with YaRN, as sequence length increases, the relative speedup can reach __over 7×__.
|
| 34 |
+
|
| 35 |
+
<p align="center"><img src="https://raw.githubusercontent.com/inclusionAI/Ling-V2/refs/heads/main/figures/needle_in_a_haystack.webp" /></p>
|
| 36 |
+
|
| 37 |
+
### Open-sourced FP8 Efficient Training Solution
|
| 38 |
+
|
| 39 |
+
Ling 2.0 employs __FP8 mixed-precision training__ throughout. Compared with BF16, experiments with over 1T training tokens show nearly identical loss curves and downstream benchmark performance. To support the community in efficient continued pretraining and fine-tuning under limited compute, we are also open-sourcing our __FP8 training solution__. Based on tile/blockwise FP8 scaling, it further introduces FP8 optimizer, FP8 on-demand transpose weight, and FP8 padding routing map for extreme memory optimization. On 8/16/32 80G GPUs, compared with LLaMA 3.1 8B and Qwen3 8B, __Ling-mini-2.0 achieved 30–60% throughput gains with MTP enabled, and 90–120% throughput gains with MTP disabled__.
|
| 40 |
+
|
| 41 |
+
### A More Open Opensource Strategy
|
| 42 |
+
|
| 43 |
+
We believe Ling-mini-2.0 is an ideal starting point for MoE research. For the first time at this scale, it integrates 1/32 sparsity, MTP layers, and FP8 training — achieving both strong effectiveness and efficient training/inference performance, making it a prime candidate for the small-size LLM segment.
|
| 44 |
+
To further foster community research, in addition to releasing the post-trained version, we are also open-sourcing __five pretraining checkpoints__: the pre-finetuning Ling-mini-2.0-base, along with four base models trained on 5T, 10T, 15T, and 20T tokens, enabling deeper research and broader applications.
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
## Model Downloads
|
| 48 |
+
|
| 49 |
+
You can download the following table to see the various stage of Ling-mini-2.0 models(1.43B activated of 16.26B total params). If you are located in mainland China, we also provide the model on ModelScope.cn to speed up the download process.
|
| 50 |
+
|
| 51 |
+
<center>
|
| 52 |
+
|
| 53 |
+
| **Model** | **Context Length** | **Download** |
|
| 54 |
+
|:----------------------:| :----------------: |:--------------------------------------------------------------------------------------------------------------------------------------------------------------------:|
|
| 55 |
+
| Ling-mini-base-2.0 | 32K -> 128K (YaRN) | [🤗 HuggingFace](https://huggingface.co/inclusionAI/Ling-mini-base-2.0) <br>[🤖 ModelScope](https://www.modelscope.cn/models/inclusionAI/Ling-mini-base-2.0) |
|
| 56 |
+
| Ling-mini-base-2.0-5T | 4K | [🤗 HuggingFace](https://huggingface.co/inclusionAI/Ling-mini-base-2.0-5T) <br>[🤖 ModelScope](https://www.modelscope.cn/models/inclusionAI/Ling-mini-base-2.0-5T) |
|
| 57 |
+
| Ling-mini-base-2.0-10T | 4K | [🤗 HuggingFace](https://huggingface.co/inclusionAI/Ling-mini-base-2.0-10T) <br>[🤖 ModelScope](https://www.modelscope.cn/models/inclusionAI/Ling-mini-base-2.0-10T) |
|
| 58 |
+
| Ling-mini-base-2.0-15T | 4K | [🤗 HuggingFace](https://huggingface.co/inclusionAI/Ling-mini-base-2.0-15T) <br>[🤖 ModelScope](https://www.modelscope.cn/models/inclusionAI/Ling-mini-base-2.0-15T) |
|
| 59 |
+
| Ling-mini-base-2.0-20T | 4K | [🤗 HuggingFace](https://huggingface.co/inclusionAI/Ling-mini-base-2.0-20T) <br>[🤖 ModelScope](https://www.modelscope.cn/models/inclusionAI/Ling-mini-base-2.0-20T) |
|
| 60 |
+
| Ling-mini-2.0 | 32K -> 128K (YaRN) | [🤗 HuggingFace](https://huggingface.co/inclusionAI/Ling-mini-2.0) <br>[🤖 ModelScope](https://www.modelscope.cn/models/inclusionAI/Ling-mini-2.0) |
|
| 61 |
+
|
| 62 |
+
</center>
|
| 63 |
+
|
| 64 |
+
Note: If you are interested in previous version, please visit the past model collections in [Huggingface](https://huggingface.co/inclusionAI) or [ModelScope](https://modelscope.cn/organization/inclusionAI).
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
## Quickstart
|
| 68 |
+
|
| 69 |
+
### Convert to safetensors
|
| 70 |
+
|
| 71 |
+
Models with safetensors format can be downloaded from [HuggingFace](https://huggingface.co/inclusionAI) or [ModelScope](https://modelscope.cn/organization/inclusionAI).
|
| 72 |
+
If you want to train your model and eval it, you can convert from dcp produced by training.
|
| 73 |
+
```shell
|
| 74 |
+
python tools/convert_dcp_to_safe_tensors.py --checkpoint-path ${DCP_PATH} --target-path ${SAFETENSORS_PATH}
|
| 75 |
+
```
|
| 76 |
+
|
| 77 |
+
Currently, BF16 and FP8 formats are supported, you can use convert parameter to handle it:
|
| 78 |
+
- `--force-bf16` for BF16 format.
|
| 79 |
+
- `--force-fp8` for FP8 format.
|
| 80 |
+
|
| 81 |
+
### 🤗 Hugging Face Transformers
|
| 82 |
+
|
| 83 |
+
Here is a code snippet to show you how to use the chat model with `transformers`:
|
| 84 |
+
|
| 85 |
+
```python
|
| 86 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
| 87 |
+
|
| 88 |
+
model_name = "inclusionAI/Ling-mini-2.0"
|
| 89 |
+
|
| 90 |
+
model = AutoModelForCausalLM.from_pretrained(
|
| 91 |
+
model_name,
|
| 92 |
+
dtype="auto",
|
| 93 |
+
device_map="auto",
|
| 94 |
+
trust_remote_code=True,
|
| 95 |
+
)
|
| 96 |
+
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
| 97 |
+
|
| 98 |
+
prompt = "Give me a short introduction to large language models."
|
| 99 |
+
messages = [
|
| 100 |
+
{"role": "system", "content": "You are Ling, an assistant created by inclusionAI"},
|
| 101 |
+
{"role": "user", "content": prompt}
|
| 102 |
+
]
|
| 103 |
+
text = tokenizer.apply_chat_template(
|
| 104 |
+
messages,
|
| 105 |
+
tokenize=False,
|
| 106 |
+
add_generation_prompt=True
|
| 107 |
+
)
|
| 108 |
+
model_inputs = tokenizer([text], return_tensors="pt", return_token_type_ids=False).to(model.device)
|
| 109 |
+
|
| 110 |
+
generated_ids = model.generate(
|
| 111 |
+
**model_inputs,
|
| 112 |
+
max_new_tokens=512
|
| 113 |
+
)
|
| 114 |
+
generated_ids = [
|
| 115 |
+
output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
|
| 116 |
+
]
|
| 117 |
+
|
| 118 |
+
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
|
| 119 |
+
```
|
| 120 |
+
|
| 121 |
+
### 🤖 ModelScope
|
| 122 |
+
|
| 123 |
+
If you're in mainland China, we strongly recommend you to use our model from 🤖 <a href="https://modelscope.cn/organization/inclusionAI">ModelScope</a>.
|
| 124 |
+
|
| 125 |
+
## Deployment
|
| 126 |
+
|
| 127 |
+
### vLLM
|
| 128 |
+
|
| 129 |
+
vLLM supports offline batched inference or launching an OpenAI-Compatible API Service for online inference.
|
| 130 |
+
|
| 131 |
+
#### Environment Preparation
|
| 132 |
+
|
| 133 |
+
Since the Pull Request (PR) has not been submitted to the vLLM community at this stage, please prepare the environment by following the steps below:
|
| 134 |
+
|
| 135 |
+
```bash
|
| 136 |
+
git clone -b v0.10.0 https://github.com/vllm-project/vllm.git
|
| 137 |
+
cd vllm
|
| 138 |
+
git apply Ling-V2/inference/vllm/bailing_moe_v2.patch
|
| 139 |
+
pip install -e .
|
| 140 |
+
```
|
| 141 |
+
|
| 142 |
+
#### Offline Inference:
|
| 143 |
+
|
| 144 |
+
```bash
|
| 145 |
+
from transformers import AutoTokenizer
|
| 146 |
+
from vllm import LLM, SamplingParams
|
| 147 |
+
|
| 148 |
+
tokenizer = AutoTokenizer.from_pretrained("inclusionAI/Ling-mini-2.0")
|
| 149 |
+
|
| 150 |
+
sampling_params = SamplingParams(temperature=0.7, top_p=0.8, repetition_penalty=1.05, max_tokens=16384)
|
| 151 |
+
|
| 152 |
+
llm = LLM(model="inclusionAI/Ling-mini-2.0", dtype='bfloat16')
|
| 153 |
+
prompt = "Give me a short introduction to large language models."
|
| 154 |
+
messages = [
|
| 155 |
+
{"role": "system", "content": "You are Ling, an assistant created by inclusionAI"},
|
| 156 |
+
{"role": "user", "content": prompt}
|
| 157 |
+
]
|
| 158 |
+
|
| 159 |
+
text = tokenizer.apply_chat_template(
|
| 160 |
+
messages,
|
| 161 |
+
tokenize=False,
|
| 162 |
+
add_generation_prompt=True
|
| 163 |
+
)
|
| 164 |
+
outputs = llm.generate([text], sampling_params)
|
| 165 |
+
|
| 166 |
+
```
|
| 167 |
+
|
| 168 |
+
#### Online Inference:
|
| 169 |
+
|
| 170 |
+
```bash
|
| 171 |
+
vllm serve inclusionAI/Ling-mini-2.0 \
|
| 172 |
+
--tensor-parallel-size 2 \
|
| 173 |
+
--pipeline-parallel-size 1 \
|
| 174 |
+
--use-v2-block-manager \
|
| 175 |
+
--gpu-memory-utilization 0.90
|
| 176 |
+
```
|
| 177 |
+
|
| 178 |
+
To handle long context in vLLM using YaRN, we need to follow these two steps:
|
| 179 |
+
1. Add a `rope_scaling` field to the model's `config.json` file, for example:
|
| 180 |
+
```json
|
| 181 |
+
{
|
| 182 |
+
...,
|
| 183 |
+
"rope_scaling": {
|
| 184 |
+
"factor": 4.0,
|
| 185 |
+
"original_max_position_embeddings": 32768,
|
| 186 |
+
"type": "yarn"
|
| 187 |
+
}
|
| 188 |
+
}
|
| 189 |
+
```
|
| 190 |
+
2. Use an additional parameter `--max-model-len` to specify the desired maximum context length when starting the vLLM service.
|
| 191 |
+
|
| 192 |
+
For detailed guidance, please refer to the vLLM [`instructions`](https://docs.vllm.ai/en/latest/).
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
### SGLang
|
| 196 |
+
|
| 197 |
+
#### Environment Preparation
|
| 198 |
+
|
| 199 |
+
We will later submit our model to SGLang official release, now we can prepare the environment following steps:
|
| 200 |
+
```shell
|
| 201 |
+
pip3 install sglang==0.5.2rc0 sgl-kernel==0.3.7.post1
|
| 202 |
+
```
|
| 203 |
+
You can use docker image as well:
|
| 204 |
+
```shell
|
| 205 |
+
docker pull lmsysorg/sglang:v0.5.2rc0-cu126
|
| 206 |
+
```
|
| 207 |
+
Then you should apply patch to sglang installation:
|
| 208 |
+
```shell
|
| 209 |
+
# patch command is needed, run `yum install -y patch` if needed
|
| 210 |
+
patch -d `python -c 'import sglang;import os; print(os.path.dirname(sglang.__file__))'` -p3 < inference/sglang/bailing_moe_v2.patch
|
| 211 |
+
```
|
| 212 |
+
|
| 213 |
+
#### Run Inference
|
| 214 |
+
|
| 215 |
+
BF16 and FP8 models are supported by SGLang now, it depends on the dtype of the model in ${MODEL_PATH}. They both share the same command in the following:
|
| 216 |
+
|
| 217 |
+
- Start server:
|
| 218 |
+
```shell
|
| 219 |
+
python -m sglang.launch_server \
|
| 220 |
+
--model-path $MODLE_PATH \
|
| 221 |
+
--host 0.0.0.0 --port $PORT \
|
| 222 |
+
--trust-remote-code \
|
| 223 |
+
--attention-backend fa3
|
| 224 |
+
```
|
| 225 |
+
MTP is supported for base model, and not yet for chat model. You can add parameter `--speculative-algorithm NEXTN`
|
| 226 |
+
to start command.
|
| 227 |
+
|
| 228 |
+
- Client:
|
| 229 |
+
```shell
|
| 230 |
+
curl -s http://localhost:${PORT}/v1/chat/completions \
|
| 231 |
+
-H "Content-Type: application/json" \
|
| 232 |
+
-d '{"model": "auto", "messages": [{"role": "user", "content": "What is the capital of France?"}]}'
|
| 233 |
+
"""
|
| 234 |
+
```
|
| 235 |
+
More usage can be found [here](https://docs.sglang.ai/basic_usage/send_request.html)
|
| 236 |
+
|
| 237 |
+
## Training
|
| 238 |
+
|
| 239 |
+
We also provide a complete and efficient training framework that covers both pre-training and finetune. Based on this framework, continue training can be performed on the Ling-mini-2.0 checkpoint. With our training framework, the training throughput of the Ling-mini-2.0 model is significantly better than that of the existing Dense 8B model (Qwen3-8B, Llama3-8B).
|
| 240 |
+
|
| 241 |
+
### Pre-training
|
| 242 |
+
|
| 243 |
+
[Pretraining demo](https://github.com/inclusionAI/Ling-V2/blob/main/docs/gpu_based_training.md) to Continue pretraining Ling models.
|
| 244 |
+
|
| 245 |
+
#### Performance Benchmark
|
| 246 |
+
|
| 247 |
+
The table below shows the pre-training performance of several models, measured in **tokens per second** on 8, 16, and 32 80G GPUs. Ling-mini-2.0 achieves significantly higher training efficiency compared to the baseline, making it easier and more cost-effective to continue pre-training with our [demo scripts](https://github.com/inclusionAI/Ling-V2/blob/main/docs/gpu_based_training.md).
|
| 248 |
+
|
| 249 |
+
<center>
|
| 250 |
+
|
| 251 |
+
| **Model** | **8 x 80G GPUs (GBS=128)** | **16 x 80G GPUs (GBS=256)** | **32 x 80G GPUs (GBS=512)** |
|
| 252 |
+
|:-----------------------:| :--------------------: | :---------------------: | :---------------------: |
|
| 253 |
+
| LLaMA 3.1 8B (baseline) | 81222 | 161319 | 321403 |
|
| 254 |
+
| Qwen3 8B | 55775 (-31.33%) | 109799 (-31.94%) | 219943 (-31.57%) |
|
| 255 |
+
| Ling-mini-2.0 | 109532 (+34.86%) | 221585 (+37.36%) | 448726 (+39.61%) |
|
| 256 |
+
| Ling-mini-2.0 w/o MTP | 128298 (+57.96%) | 307264 (+90.47%) | 611466 (+90.25%) |
|
| 257 |
+
|
| 258 |
+
</center>
|
| 259 |
+
|
| 260 |
+
### Finetuning
|
| 261 |
+
|
| 262 |
+
We recommend you to use [Llama-Factory](https://github.com/hiyouga/LLaMA-Factory) to [finetune Ling](https://github.com/inclusionAI/Ling-V2/blob/main/docs/llamafactory_finetuning.md). In addition to that, you can also use [Megatron for finetuning](https://github.com/inclusionAI/Ling-V2/blob/main/docs/megatron_sft_training.md).
|
| 263 |
+
|
| 264 |
+
## License
|
| 265 |
+
|
| 266 |
+
This code repository is licensed under [the MIT License](https://github.com/inclusionAI/Ling-V2/blob/master/LICENCE).
|
| 267 |
+
|
| 268 |
+
## Citation
|
| 269 |
+
|
| 270 |
+
If you find our work helpful, feel free to give us a cite.
|
| 271 |
+
|
| 272 |
+
```
|
| 273 |
+
|
| 274 |
+
```
|
configuration_bailing_moe_v2.py
CHANGED
|
@@ -39,7 +39,7 @@ class BailingMoeV2Config(PretrainedConfig):
|
|
| 39 |
head_dim=128,
|
| 40 |
output_router_logits=False,
|
| 41 |
use_qk_norm=True,
|
| 42 |
-
|
| 43 |
mtp_loss_scaling_factor=0,
|
| 44 |
moe_router_enable_expert_bias=True,
|
| 45 |
routed_scaling_factor=1.0,
|
|
@@ -58,7 +58,7 @@ class BailingMoeV2Config(PretrainedConfig):
|
|
| 58 |
self.embedding_dropout = embedding_dropout
|
| 59 |
self.attention_dropout = attention_dropout
|
| 60 |
self.output_dropout = output_dropout
|
| 61 |
-
self.
|
| 62 |
self.mtp_loss_scaling_factor = mtp_loss_scaling_factor
|
| 63 |
self.initializer_range = initializer_range
|
| 64 |
self.max_position_embeddings = max_position_embeddings
|
|
|
|
| 39 |
head_dim=128,
|
| 40 |
output_router_logits=False,
|
| 41 |
use_qk_norm=True,
|
| 42 |
+
num_nextn_predict_layers=0,
|
| 43 |
mtp_loss_scaling_factor=0,
|
| 44 |
moe_router_enable_expert_bias=True,
|
| 45 |
routed_scaling_factor=1.0,
|
|
|
|
| 58 |
self.embedding_dropout = embedding_dropout
|
| 59 |
self.attention_dropout = attention_dropout
|
| 60 |
self.output_dropout = output_dropout
|
| 61 |
+
self.num_nextn_predict_layers = num_nextn_predict_layers
|
| 62 |
self.mtp_loss_scaling_factor = mtp_loss_scaling_factor
|
| 63 |
self.initializer_range = initializer_range
|
| 64 |
self.max_position_embeddings = max_position_embeddings
|
modeling_bailing_moe_v2.py
CHANGED
|
@@ -25,9 +25,7 @@ from typing import List, Optional, Tuple, Union
|
|
| 25 |
|
| 26 |
import torch
|
| 27 |
import torch.nn.functional as F
|
| 28 |
-
import torch.utils.checkpoint
|
| 29 |
from torch import nn
|
| 30 |
-
from torch.nn import CrossEntropyLoss
|
| 31 |
|
| 32 |
from transformers.activations import ACT2FN
|
| 33 |
from transformers.cache_utils import Cache, DynamicCache
|
|
@@ -1157,11 +1155,11 @@ class BailingMoeV2Model(BailingMoeV2PreTrainedModel):
|
|
| 1157 |
super().__init__(config)
|
| 1158 |
self.padding_idx = config.pad_token_id
|
| 1159 |
self.vocab_size = config.vocab_size
|
| 1160 |
-
self.
|
| 1161 |
|
| 1162 |
self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
|
| 1163 |
self.layers = []
|
| 1164 |
-
for layer_idx in range(config.num_hidden_layers + config.
|
| 1165 |
layer_cls = BailingMoeV2DecoderLayer if layer_idx < config.num_hidden_layers else BailingMoeV2MTPLayer
|
| 1166 |
self.layers.append(layer_cls(config, layer_idx))
|
| 1167 |
|
|
@@ -1267,8 +1265,8 @@ class BailingMoeV2Model(BailingMoeV2PreTrainedModel):
|
|
| 1267 |
all_self_attns = () if output_attentions else None
|
| 1268 |
all_router_logits = () if output_router_logits else None
|
| 1269 |
next_decoder_cache = None
|
| 1270 |
-
layers = self.layers[: -self.
|
| 1271 |
-
mtp_layers = self.layers[-self.
|
| 1272 |
|
| 1273 |
for decoder_layer in layers:
|
| 1274 |
if output_hidden_states:
|
|
@@ -1391,7 +1389,7 @@ class BailingMoeV2ForCausalLM(BailingMoeV2PreTrainedModel, GenerationMixin):
|
|
| 1391 |
self.model = BailingMoeV2Model(config)
|
| 1392 |
self.vocab_size = config.vocab_size
|
| 1393 |
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
| 1394 |
-
self.
|
| 1395 |
self.mtp_loss_scaling_factor = config.mtp_loss_scaling_factor
|
| 1396 |
|
| 1397 |
# Initialize weights and apply final processing
|
|
@@ -1491,18 +1489,21 @@ class BailingMoeV2ForCausalLM(BailingMoeV2PreTrainedModel, GenerationMixin):
|
|
| 1491 |
loss = self.loss_function(logits, labels, self.config.vocab_size, **kwargs)
|
| 1492 |
|
| 1493 |
all_mtp_logits = None
|
| 1494 |
-
if self.
|
| 1495 |
mtp_hidden_states = outputs.mtp_hidden_states
|
| 1496 |
-
shift_labels_mtp =
|
| 1497 |
-
for i in range(self.
|
| 1498 |
mtp_hidden_states = mtp_hidden_states[i]
|
| 1499 |
mtp_logits = self.lm_head(mtp_hidden_states).float()
|
| 1500 |
if all_mtp_logits is None:
|
| 1501 |
all_mtp_logits = []
|
| 1502 |
all_mtp_logits.append(mtp_logits)
|
| 1503 |
if labels is not None:
|
|
|
|
|
|
|
| 1504 |
shift_labels_mtp, _ = roll_tensor(shift_labels_mtp, shifts=-1, dims=-1, fill_value=-100)
|
| 1505 |
-
|
|
|
|
| 1506 |
if loss is not None:
|
| 1507 |
loss += self.mtp_loss_scaling_factor * mtp_loss
|
| 1508 |
else:
|
|
@@ -1529,3 +1530,4 @@ class BailingMoeV2ForCausalLM(BailingMoeV2PreTrainedModel, GenerationMixin):
|
|
| 1529 |
attentions=outputs.attentions,
|
| 1530 |
router_logits=outputs.router_logits,
|
| 1531 |
)
|
|
|
|
|
|
| 25 |
|
| 26 |
import torch
|
| 27 |
import torch.nn.functional as F
|
|
|
|
| 28 |
from torch import nn
|
|
|
|
| 29 |
|
| 30 |
from transformers.activations import ACT2FN
|
| 31 |
from transformers.cache_utils import Cache, DynamicCache
|
|
|
|
| 1155 |
super().__init__(config)
|
| 1156 |
self.padding_idx = config.pad_token_id
|
| 1157 |
self.vocab_size = config.vocab_size
|
| 1158 |
+
self.num_nextn_predict_layers = config.num_nextn_predict_layers
|
| 1159 |
|
| 1160 |
self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
|
| 1161 |
self.layers = []
|
| 1162 |
+
for layer_idx in range(config.num_hidden_layers + config.num_nextn_predict_layers):
|
| 1163 |
layer_cls = BailingMoeV2DecoderLayer if layer_idx < config.num_hidden_layers else BailingMoeV2MTPLayer
|
| 1164 |
self.layers.append(layer_cls(config, layer_idx))
|
| 1165 |
|
|
|
|
| 1265 |
all_self_attns = () if output_attentions else None
|
| 1266 |
all_router_logits = () if output_router_logits else None
|
| 1267 |
next_decoder_cache = None
|
| 1268 |
+
layers = self.layers[: -self.num_nextn_predict_layers] if self.num_nextn_predict_layers > 0 else self.layers
|
| 1269 |
+
mtp_layers = self.layers[-self.num_nextn_predict_layers :] if self.num_nextn_predict_layers > 0 else None
|
| 1270 |
|
| 1271 |
for decoder_layer in layers:
|
| 1272 |
if output_hidden_states:
|
|
|
|
| 1389 |
self.model = BailingMoeV2Model(config)
|
| 1390 |
self.vocab_size = config.vocab_size
|
| 1391 |
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
| 1392 |
+
self.num_nextn_predict_layers = config.num_nextn_predict_layers
|
| 1393 |
self.mtp_loss_scaling_factor = config.mtp_loss_scaling_factor
|
| 1394 |
|
| 1395 |
# Initialize weights and apply final processing
|
|
|
|
| 1489 |
loss = self.loss_function(logits, labels, self.config.vocab_size, **kwargs)
|
| 1490 |
|
| 1491 |
all_mtp_logits = None
|
| 1492 |
+
if self.num_nextn_predict_layers > 0:
|
| 1493 |
mtp_hidden_states = outputs.mtp_hidden_states
|
| 1494 |
+
shift_labels_mtp = None
|
| 1495 |
+
for i in range(self.num_nextn_predict_layers):
|
| 1496 |
mtp_hidden_states = mtp_hidden_states[i]
|
| 1497 |
mtp_logits = self.lm_head(mtp_hidden_states).float()
|
| 1498 |
if all_mtp_logits is None:
|
| 1499 |
all_mtp_logits = []
|
| 1500 |
all_mtp_logits.append(mtp_logits)
|
| 1501 |
if labels is not None:
|
| 1502 |
+
if shift_labels_mtp is None:
|
| 1503 |
+
shift_labels_mtp = labels.clone()
|
| 1504 |
shift_labels_mtp, _ = roll_tensor(shift_labels_mtp, shifts=-1, dims=-1, fill_value=-100)
|
| 1505 |
+
mtp_logits_ = mtp_logits.view(-1, self.config.vocab_size)
|
| 1506 |
+
mtp_loss = self.loss_function(mtp_logits_, shift_labels_mtp.to(mtp_logits_.device).view(-1), self.config.vocab_size, **kwargs)
|
| 1507 |
if loss is not None:
|
| 1508 |
loss += self.mtp_loss_scaling_factor * mtp_loss
|
| 1509 |
else:
|
|
|
|
| 1530 |
attentions=outputs.attentions,
|
| 1531 |
router_logits=outputs.router_logits,
|
| 1532 |
)
|
| 1533 |
+
|