rulixiang commited on
Commit
c0da5fd
·
1 Parent(s): 3757e95

update modeling and readme

Browse files
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>&nbsp&nbsp | &nbsp&nbsp🤖 <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
- num_mtp_layers=0,
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.num_mtp_layers = num_mtp_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
 
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.num_mtp_layers = config.num_mtp_layers
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.num_mtp_layers):
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.num_mtp_layers] if self.num_mtp_layers > 0 else self.layers
1271
- mtp_layers = self.layers[-self.num_mtp_layers :] if self.num_mtp_layers > 0 else None
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.num_mtp_layers = config.num_mtp_layers
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.num_mtp_layers > 0:
1495
  mtp_hidden_states = outputs.mtp_hidden_states
1496
- shift_labels_mtp = labels.clone()
1497
- for i in range(self.num_mtp_layers):
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
- mtp_loss = self.loss_function(mtp_logits, shift_labels_mtp, self.config.vocab_size, **kwargs)
 
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
+