ValueError: block_mask was created for block_mask.shape=(1, 1, 20096, 40192) but got q_len=20000 and kv_len=40000.

[Jim2016713]

When I was training, I encountered a problem like this.
[rank3]: ValueError: block_mask was created for block_mask.shape=(1, 1, 20096, 40192) but got q_len=20000 and kv_len=40000. As the block mask was created for a larger length than you're using it for, you can either 1. create a new block mask with the correct length, or 2. 'adjust' the existing block mask to the correct length by calling block_mask._adjust(q_len, kv_len). This essentially 'crops' the block mask to the upper left corner, which does not work for all mask_mods!
[rank2]: Traceback (most recent call last):

training scripts：
python ./scripts/data_generation_offline.py
--target-model-path $TARGET_MODEL_NAME_OR_PATH
--train-data-path $TRAIN_DATA_PATH
--seq-length 20000
--hf-cache-dir ./output/eagle3_data_gen/cache
--output-dir ./eagle3_data_gen/training_data
--max-model-len 20000
--batch-size 32
--num-preprocessing-workers 32

[fynnsu]

Hi @Jim2016713, thanks for submitting the issue. This looks like an issue with the FlexAttention integration, which we can help you debug if you provide a bit more information!
Could you provide:

• pytorch and vllm versions
• speculators commit hash
• The name of the target model you're using if using a public model on huggingface. Otherwise if it's a local model, some basic information on the model architecture would be helpful (like hidden dim, number of attention heads, moe or dense, if special attention layers like sliding window attention are used).

You also mentioned this is something you ran into during training but the script call you provided is for the datagen step, could you clarify if the script call is correct or otherwise share your ./scripts/train.py call?