Training loss crashes when sequence length exceeds 20k

[LuLuLuyi]

Hi, I encountered an issue when using llamafactor to train SDAR-30B-A3B-Sci.
Whenever the training sequence length exceeds 20k (e.g., 32k), the training crashes and the loss becomes 0.

Below is my training config:

### model
model_name_or_path: SDAR-30B-A3B-Sci
train_from_scratch: false
trust_remote_code: true

### method
stage: sft
do_train: true
finetuning_type: full
freeze_vision_tower: true
freeze_multi_modal_projector: false
freeze_language_model: true
deepspeed: examples/deepspeed/ds_z3_config.json
# disable_gradient_checkpointing: true
gradient_checkpointing: true

### dataset
dataset: open_r1_math
template: qwen3
block_length: 4
cutoff_len: 32768
truncate_mode: drop
overwrite_cache: false
preprocessing_num_workers: 96
dataloader_num_workers: 4
neat_packing: true

### output
output_dir: /sft_test/sdar_30ba3b_math_r1_cot/
logging_steps: 5
save_steps: 256
save_total_limit: 10
plot_loss: true
overwrite_output_dir: true
save_only_model: false
report_to: wandb

### train
run_name: sdar_30ba3b_math_r1_cot
include_effective_tokens_per_second: true
per_device_train_batch_size: 1
gradient_accumulation_steps: 1
learning_rate: 1.0e-5
num_train_epochs: 10.0
lr_scheduler_type: constant_with_warmup
warmup_ratio: 0.03
bf16: true
ddp_timeout: 180000000
resume_from_checkpoint: null

this is my training loss curve:

0%|          | 1/5190 [06:08<531:06:38, 368.47s/it]
  0%|          | 2/5190 [20:47<964:11:30, 669.06s/it]
  0%|          | 3/5190 [21:59<570:44:41, 396.12s/it]
  0%|          | 4/5190 [22:59<379:23:41, 263.37s/it]
  0%|          | 5/5190 [23:52<270:08:47, 187.57s/it]
                                                     
{'loss': 2.1261, 'grad_norm': 1.4142135623730951, 'learning_rate': 2.564102564102564e-07, 'epoch': 0.01}

  0%|          | 5/5190 [23:52<270:08:47, 187.57s/it]
  0%|          | 6/5190 [24:46<204:34:09, 142.06s/it]
  0%|          | 7/5190 [25:41<163:40:46, 113.69s/it]
  0%|          | 8/5190 [26:35<136:10:48, 94.61s/it] 
  0%|          | 9/5190 [27:28<117:36:57, 81.73s/it]
  0%|          | 10/5190 [28:22<105:26:29, 73.28s/it]
                                                     
{'loss': 0.0, 'grad_norm': 1.4142135623730951, 'learning_rate': 5.76923076923077e-07, 'epoch': 0.02}

  0%|          | 10/5190 [28:22<105:26:29, 73.28s/it]
  0%|          | 11/5190 [29:19<98:09:25, 68.23s/it] 
  0%|          | 12/5190 [30:14<92:17:26, 64.16s/it]
  0%|          | 13/5190 [31:07<87:32:16, 60.87s/it]
  0%|          | 14/5190 [32:19<92:15:20, 64.17s/it]
  0%|          | 15/5190 [33:22<91:52:30, 63.91s/it]
                                                    
{'loss': 0.0, 'grad_norm': 1.4142135623730951, 'learning_rate': 8.974358974358975e-07, 'epoch': 0.03}

  0%|          | 15/5190 [33:22<91:52:30, 63.91s/it]
  0%|          | 16/5190 [34:18<88:26:25, 61.54s/it]
  0%|          | 17/5190 [35:10<83:56:17, 58.41s/it]
  0%|          | 18/5190 [36:03<81:35:43, 56.79s/it]
  0%|          | 19/5190 [36:55<79:51:08, 55.59s/it]
  0%|          | 20/5190 [37:49<79:04:08, 55.06s/it]
                                                    
{'loss': 0.0, 'grad_norm': 1.4142135623730951, 'learning_rate': 1.217948717948718e-06, 'epoch': 0.04}

  0%|          | 20/5190 [37:49<79:04:08, 55.06s/it]
  0%|          | 21/5190 [38:40<77:23:28, 53.90s/it]
  0%|          | 22/5190 [39:34<77:18:16, 53.85s/it]
  0%|          | 23/5190 [40:29<77:44:53, 54.17s/it]
  0%|          | 24/5190 [41:21<76:54:32, 53.60s/it]
  0%|          | 25/5190 [42:15<77:06:53, 53.75s/it]

I suspect this issue may be related to the model’s maximum supported context length = 40k:
• When training with 20k sequences, the attention mask becomes 20k clean + 20k noise = 40k, which fits within the limit.
• When training with 32k sequences, the attention mask becomes 32k clean + 32k noise = 64k, which exceeds the 40k limit and causes training to crash.

However, I still have some questions: The model’s RoPE / position encoding should only use positions from 0 to cutoff_len, so in principle it shouldn’t trigger out-of-range errors.

Could you clarify what exactly causes the training to crash when the sequence length exceeds around 20k, and whether there is any way to support longer training sequences such as 32k or 64k? Any guidance would be greatly appreciated. Thank you!

[chengshuang18]

For MoE training, some special kernels are required. We currently haven’t implemented SDAR MoE training based on llamafactory — the released SDAR models were trained using our internal training framework.

From a theoretical perspective, the position_ids / RoPE encoding themselves shouldn’t be the root cause of the crash; we’ve already validated long-sequence (20k+) training on 4B/8B dense models without issues. In your case, the problem is most likely related to the MoE modeling implementation in llamafactory, rather than the context-length limit of the model itself.

I may try implementing MoE training for SDAR models in llamafactory in the near future. If there’s progress, I’ll update this issue with the results. That said, this might take some time — hope you can understand.

[LuLuLuyi]

Thanks for the clarification! I followed the advice in another issue (#18 (comment)) and implemented the SDAR-30B-A3B modeling file by referring to the 4B/8B dense model’s modeling code. With this implementation, training works correctly when the sequence length is ≤20k, so the basic modeling alignment seems fine.

I’d like to confirm one thing: when you mentioned that “For MoE training, some special kernels are required,” does this mean that I cannot fully rely on the dense 4B/8B modeling code as a reference when implementing the MoE version? In other words, are there MoE-specific components that must be implemented differently for long-sequence training?

Additionally, could you point out which parts of the modeling code are most likely related to the failure when training beyond 20k sequence length? Knowing which modules (attention，bd_atten_mask, fused kernels, etc.) may interact with long-sequence execution would help me investigate further.

Thanks again for your suggestions!

[unsosp1ro]

To clarify, the attn modules (including the attn mask) and the loss objective (specifically, the main loss in MoE) in the Dense model and the MoE are implemented in exactly the same way.

First, please check whether the dense models exhibit the same crash issue when length > 20k. If it only occurs in MoE, there are some points need to double-check. These may not necessarily be the cause of your issue, but it’s worth verifying:

1. MLP part: The MLP component can be directly reused from the standard AR model’s open-source training kernels. This is orthogonal to our block-diffusion logic. (Given that your training runs successfully with 20k length, this is not the main cause probably.)
2. Balancing_loss: During training, noisy and clean sequences are concatenated together. However, we only use the router_logits from the noisy part to compute the balancing loss. If the input_ids length is 2n, only n tokens are actually used for the aux loss. (Note: the noisy part includes both masked tokens and non-masked tokens.)

Hopefully checking these points will be helpful :)

[LuLuLuyi]

Thank you for your suggestions!

I have tried both 4B and 8B dense models, and they exhibit the same crash issue when the sequence length exceeds 20k. So it seems that this problem is not specific to MoE.

SingleProcess AUTOTUNE benchmarking takes 32.3676 seconds and 58.9653 seconds precompiling for 20 choices

  0%|          | 1/1890 [07:35<238:50:46, 455.19s/it]
  0%|          | 2/1890 [08:56<123:28:00, 235.42s/it]
  0%|          | 3/1890 [10:18<86:35:08, 165.19s/it] 
  0%|          | 4/1890 [11:41<69:29:20, 132.64s/it]
  0%|          | 5/1890 [13:02<59:43:29, 114.06s/it]
                                                    
{'loss': 2.4734, 'grad_norm': 1.4142135381698608, 'learning_rate': 7.017543859649123e-07, 'epoch': 0.03}

  0%|          | 5/1890 [13:02<59:43:29, 114.06s/it]
  0%|          | 6/1890 [14:24<54:00:50, 103.21s/it]
  0%|          | 7/1890 [15:47<50:31:26, 96.59s/it] 
  0%|          | 8/1890 [17:09<48:03:37, 91.93s/it]
  0%|          | 9/1890 [18:31<46:26:21, 88.88s/it]
  1%|          | 10/1890 [19:54<45:30:20, 87.14s/it]
                                                    
{'loss': 0.0, 'grad_norm': 1.4142135381698608, 'learning_rate': 1.5789473684210526e-06, 'epoch': 0.05}

  1%|          | 10/1890 [19:54<45:30:20, 87.14s/it]
  1%|          | 11/1890 [21:17<44:42:51, 85.67s/it]
  1%|          | 12/1890 [22:38<44:02:57, 84.44s/it]
  1%|          | 13/1890 [24:01<43:43:13, 83.85s/it]
  1%|          | 14/1890 [25:29<44:28:21, 85.34s/it]
  1%|          | 15/1890 [26:51<43:53:54, 84.28s/it]

This is my training config for SDAR-8B-Chat：

### model
model_name_or_path: dllm_training/model/SDAR-8B-Chat
train_from_scratch: false
trust_remote_code: true

### method
stage: sft
do_train: true
finetuning_type: full
freeze_vision_tower: true
freeze_multi_modal_projector: false
freeze_language_model: true
deepspeed: dllm_training/llama_factory_sdar/examples/deepspeed/ds_z3_config.json
# disable_gradient_checkpointing: true
gradient_checkpointing: true

### dataset
dataset: open_r1_math
template: qwen3 # qwen2_vl
block_length: 4
cutoff_len: 32768
truncate_mode: drop # choices: [drop, cut] 多模态推荐 drop
overwrite_cache: false
# tokenized_path: /cache_dir/for/tokenized_data
preprocessing_num_workers: 96
dataloader_num_workers: 4
neat_packing: true
# streaming: true
# max_steps: 1000

### output
output_dir: sft_test/sdar_8b_math_r1_cot/
logging_steps: 5
save_steps: 512
save_total_limit: 5
plot_loss: true
overwrite_output_dir: true
save_only_model: false
report_to: wandb  # choices: [none, wandb, tensorboard, swanlab, mlflow]

### train
run_name: sdar_8b_math_r1_cot
include_effective_tokens_per_second: true
per_device_train_batch_size: 1
gradient_accumulation_steps: 1
learning_rate: 1.0e-5
num_train_epochs: 10.0
lr_scheduler_type: constant_with_warmup
warmup_ratio: 0.03
bf16: true
ddp_timeout: 180000000
resume_from_checkpoint: null

### eval
# val_size: 0.1
# per_device_eval_batch_size: 1
# eval_strategy: steps
# eval_steps: 500

I did not make any changes to your 4B/8B modeling files. Previously, you mentioned that “we’ve already validated long-sequence (20k+) training on 4B/8B dense models without issues.” May I ask if this validation was performed directly on the provided llamafactory repository?

My training process for SDAR-8B fully reproduces the workflow in your repository, except for setting cut_off_len=32768; no other modifications were made. If there are any specific differences between your validation setup and mine that I should pay attention to, please let me know. Thank you!

[ayushnangia]

@chengshuang18 Do you have plans to open source the MoE training pipeline and dataset?
If one want to replicate the behaviour of MoE model in 8B or 4B what dataset and how much would you recommend?

Thanks for the great work!

[chengshuang18]

Regarding the training pipeline — the framework we’ve released, based on llamafactory, should work fine for general purposes. The >32k context length issue you mentioned hasn’t been investigated yet, and we haven’t implemented the MoE version publicly. The core of block diffusion lies in the attention mechanism, which is different from AR. If you’re aiming for higher training efficiency, you might consider implementing it with Megatron or XTuner. Our current work was trained using an internal framework, which we’re not able to open source at the moment (hope you can understand).

As for the dataset — our internal data quality is roughly on par with InternLM2.5. For replication, I’d recommend looking at the recently released Olmo 3. If you want to strictly reproduce our pipeline, you could start from the Olmo 3 base model, sample up to ~50B tokens from their long-context data, then use Olmo 3’s SFT dataset for fine-tuning. This setup would allow a direct comparison with the Olmo 3-instruct model to see if there’s any performance drop (with matched training data, the drop should be minimal).

For general academic research, I believe Olmo 3 provides sufficient open data and checkpoints for experiments. If strict reproduction is not necessary, you could simply take a strong open-source AR model (e.g., Qwen3-Instruct [Think]), acquire suitable SFT or “thinking” datasets, and perform SFT fine-tuning directly.

[LuLuLuyi]

Thank you for the detailed explanation. My issue is whenever the training sequence length exceeds 20k (e.g., 24k, 28k, or 32k), the training crashes and the loss becomes 0 for both dense and moe models.

My goal is to reproduce the long CoT SFT training on longer input sequences, because SDAR-30B-A3B-Sci was trained with long CoT at ~32k context length. Training on shorter sequences would likely lead to performance degradation.

Could you please try training with longer sequence lengths (>20k) on your side? I’m not sure whether the issue comes from LlamaFactory itself or from the SDAR–LlamaFactory integration.

I went through the LlamaFactory issues and didn’t find cases where increasing the training sequence length results in the loss dropping to 0. Could you share some implementation details or best practices about adapting SDAR to LlamaFactory? I’m not sure which part of the workflow might be causing the failure.

[chengshuang18]

I noticed that the scalar computation in the loss section might have a problem.

Current code:

labels = labels[logits_to_keep_half].contiguous()
loss_fct = FusedLinearDiffusionCrossEntropyLoss(reduction='sum')
loss = loss_fct(  # returns (sum_loss, unreduced_loss)
    # performs `view(-1, V)` inside the function
    x=hidden_states,
    target=labels,
    weight=self.lm_head.weight,
    bias=self.lm_head.bias,
    p_mask=p_mask,
)
loss = loss / labels.numel()

Expected code (see here):

answer_len = (labels != -100).sum()
loss_fct = FusedLinearDiffusionCrossEntropyLoss(reduction='sum')
loss = loss_fct(  # returns (sum_loss, unreduced_loss)
    # performs `view(-1, V)` inside the function
    x=hidden_states,
    target=labels[logits_to_keep_half].contiguous(),
    weight=self.lm_head.weight,
    bias=self.lm_head.bias,
    p_mask=p_mask,
)
loss = loss / answer_len

The difference is that in the current implementation, the scaler changes from 1/t to 1/t², which could make training unstable and potentially cause numerical overflow.
It’s unclear if this is the exact cause, but I wanted to share this observation in case it helps.

[LuLuLuyi]

I’ve basically identified why training crashes when sequence length >20k.
In my SFT experiments, I used the r1-math-94k dataset, whose maximum sample length is around 23k. When I set max_seq_length above 23k (e.g., 24k / 28k / 32k), the training loss collapses to 0. I believe this is caused by neat_packing.

For example, when max_seq_length=24k, two samples may get packed into one final training sample, which eventually triggers the crash.

I ran an ablation:
• I sampled 50 samples <20k tokens and trained with max_seq_length=32k → loss collapsed.
• If I add one additional sample >32k tokens to the above 50 samples → training becomes normal.

This strongly suggests the bug is related to packing.
I noticed you have made some packing-related modifications on top of the LlamaFactory implementation, but I haven’t pinpointed where the bug is. Since you’re more familiar with LlamaFactory and your packing logic, you will probably locate the issue much faster.

---

Regarding the loss scale issue, I have the following questions:

current implementation: loss = sum_loss / labels.numel()
expected implementation: loss = sum_loss / answer_len

What exactly is wrong in the original implementation?
You mentioned the scaler changes from 1/t to 1/t², but conceptually I think:
• labels[logits_to_keep_half]
• and (labels != -100)
should produce equivalent filtered targets.
Could you elaborate how the 1/t² issue arises?

---

I also have two MoE-related questions (I can open a separate issue if necessary):

In my MoE modeling implementation:

labels = labels[logits_to_keep_half].contiguous()
loss_fct = FusedLinearDiffusionCrossEntropyLoss(reduction='sum')
loss = loss_fct(
        x=hidden_states,
        target=labels,
        weight=self.lm_head.weight,
        bias=self.lm_head.bias,
        p_mask=p_mask,
    )
loss = loss / labels.numel()
logits = None

aux_loss = None
if output_router_logits:
    aux_loss = load_balancing_loss_func(
        outputs.router_logits,
        self.num_experts,
        self.num_experts_per_tok,
        attention_mask_for_router,
    )
    if labels is not None:
        loss += self.router_aux_loss_coef * aux_loss.to(loss.device)

(a) When training MoE models:
• during continued pretrain stage, and
• during SFT stage,

do you enable output_router_logits in both stages?
(i.e., router is trained jointly the whole time)

(b) What is the masking rule for attention_mask_for_router?
In this comment (#26 (comment)) you mentioned it only applies to the noisy part, but I’d like to understand the precise definition.

During SFT, the structure is generally:

[prompt tokens] [response tokens] [padding tokens]

So do you only mask the padding tokens? or only the prompt tokens? or prompt + padding, and allow only response tokens to participate in router training?

If possible, could you share the relevant part of your MoE modeling code?
Even a small snippet would be extremely helpful. Thank you!

---

If you need logs or further diagnostics, I’m happy to help.