NVIDIA-NeMo / NVIDIA-NeMo/RL

Sequence Packing Bug: cu_seqlens not propagated to AutoModel

Open
#2,105 0 comments 0 reactions 1 assignee Claimed by @zpqiu View on GitHub
bug
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

**Describe the bug**

When sequence packing is enabled with AutoModel custom/registered models (e.g., GLM-4.7-Flash, Qwen3.5-MoE), cu_seqlens is silently dropped due to an interface mismatch between NeMo-RL and AutoModel. This causes cross-sequence attention leakage in packed batches — the model attends across sequence boundaries as if the entire packed tensor were a single sequence.

**Steps/Code to reproduce bug**

Trigger two runs with an AutoModel registered model, one with sequence packing enabled, one not.
Compare the `gen_kl_error`. The symptom is the value of `gen_kl_error` gets much higher when packing is enabled. e.g. On glm4.7-flash, from 0.x down to 0.00x.

All models registered in AutoModel's `ModelRegistry` are affected. As sequence packing is enabled by default. It is of high impact.

**Additional context**

Root cause:

NeMo-RL constructs a `FlashAttentionKwargs` dataclass and passes it as a single named kwarg:
```
# nemo_rl/models/automodel/train.py (model_forward)
model_args["flash_attn_kwargs"] = processed_inputs.flash_attn_kwargs
# ...
model(**model_args)
```
This results in the model receiving:
```
**attn_kwargs = {"flash_attn_kwargs": FlashAttentionKwargs(cu_seqlens_q=..., ...)}
```
HF models understand this format — their internal attention layers know to unpack the flash_attn_kwargs dataclass.

AutoModel custom models (GLM4, Qwen3.5) expect flat kwargs:
```
**attn_kwargs = {"cu_seqlens": tensor, "max_seqlen": int, "qkv_format": "thd"}
```
The AutoModel attention layers look for `cu_seqlens` as a top-level key:
```
# nemo_automodel/components/models/deepseek_v3/layers.py:159
cu_seqlens = attn_kwargs.get("cu_seqlens", None) # Returns None!

# nemo_automodel/components/attention/utils.py:92
elif "cu_seqlens" in kwargs: # False! Key is "flash_attn_kwargs"
```
And the THD format trigger is also missed:
```
# glm4_moe_lite/model.py:267
if "qkv_format" in attn_kwargs and attn_kwargs["qkv_format"] == "thd": # False!
# squeeze_input_for_thd() never called
```
Result: TE's DotProductAttention runs with default causal attention on the entire packed tensor, treating all sequences as one continuous sequence.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.