transformers - 💡(How to fix) Fix Qwen3.5 cannot generate normally with flash-attention [2 comments, 2 participants]

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import pandas as pd
import torch
from transformers import AutoTokenizer,Qwen3ForCausalLM,AutoConfig,AutoModelForCausalLM,GenerationConfig

if __name__ == '__main__':
    model_path="/share/models/Qwen3.5-4B"#"/share/models/Qwen3-4B-Instruct-2507"#

    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, add_bos_token=False,
                                              add_eos_token=False)

    model = AutoModelForCausalLM.from_pretrained(model_path,
                                             dtype=torch.bfloat16,
                                             trust_remote_code=True,
                                             attn_implementation="flash_attention_2",
                                             device_map="cuda"
                                             ).eval()

    prompt="How are you today?"
    chat_prompt=tokenizer.apply_chat_template([{"role":"user","content":prompt}],tokenize=False,add_generation_prompt=True,enable_thinking=False)

    generation_config=GenerationConfig(
        do_sample=False,
        temperature=1.0,
        top_p=1.0,
        num_return_sequences=1,
        max_new_tokens=100,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id,
        use_cache=True,
        return_dict_in_generate=True,
        output_logits=True,
    )


      chat_prompt_ids = tokenizer(chat_prompt, return_tensors="pt")["input_ids"].to(model.device)
      output = model.generate(input_ids=chat_prompt_ids,
                                      generation_config=generation_config,)
      output_text = tokenizer.decode(output['sequences'][0][chat_prompt_ids.size(1):],skip_special_tokens=False)
      print(output_text)