[BUG]Huggingface版本推理流式输出报错

[cauwulixuan]

我在用以下代码进行流式推理的时候，参考fastchat-inference.py 的这一部分stream_generate

for i in range(max_new_tokens):
    if i == 0:  # prefill
        out = model(input_ids=start_ids, use_cache=True)
        logits = out.logits
        past_key_values = out.past_key_values
        ...
    else:  # decoding
        out = model(
            input_ids=torch.as_tensor(
                [[token] if not sent_interrupt else output_ids],
                device=device,
            ),
            use_cache=True,
            past_key_values=past_key_values if not sent_interrupt else None,
        )
        sent_interrupt = False
        logits = out.logits
        past_key_values = out.past_key_values
    ...
    probs = torch.softmax(last_token_logits, dim=-1)
    indices = torch.multinomial(probs, num_samples=2)
    tokens = [int(token) for token in indices.tolist()]
    token = tokens[0]
    output_ids.append(token)
    ...
...

1. use_flash_attention=True的时候，是可以正常推理的；
2. use_flash_attention=False的时候，报错了，报错信息如下：

You are using the default legacy behaviour of the <class 'transformers.models.llama.tokenization_llama.LlamaTokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565
{'torch_dtype': torch.float16, 'revision': 'main'}
YuanForCausalLM(
  (model): YuanModel(
    (embed_tokens): Embedding(135040, 2048, padding_idx=77185)
    (layers): ModuleList(
      (0-23): 24 x YuanDecoderLayer(
        (self_attn): YuanAttention(
          (v_proj): Linear(in_features=2048, out_features=2048, bias=False)
          (o_proj): Linear(in_features=2048, out_features=2048, bias=False)
          (rotary_emb): LlamaRotaryEmbedding()
          (lf_gate): LocalizedFiltering(
            (conv1): Conv2d(2048, 1024, kernel_size=(2, 1), stride=(1, 1), padding=(1, 0))
            (conv2): Conv2d(1024, 2048, kernel_size=(2, 1), stride=(1, 1), padding=(1, 0))
            (output_layernorm): LlamaRMSNorm()
          )
          (q_proj): Linear(in_features=2048, out_features=2048, bias=False)
          (k_proj): Linear(in_features=2048, out_features=2048, bias=False)
        )
        (mlp): YuanMLP(
          (up_proj): Linear(in_features=2048, out_features=8192, bias=False)
          (gate_proj): Linear(in_features=2048, out_features=8192, bias=False)
          (down_proj): Linear(in_features=8192, out_features=2048, bias=False)
          (act_fn): SiLU()
        )
        (input_layernorm): LlamaRMSNorm()
        (post_attention_layernorm): LlamaRMSNorm()
      )
    )
    (norm): LlamaRMSNorm()
  )
  (lm_head): Linear(in_features=2048, out_features=135040, bias=False)
)
user: yuan2.0是谁开发的？
assistant: Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/github/FastChat/fastchat/serve/cli.py", line 304, in <module>
    main(args)
  File "/github/FastChat/fastchat/serve/cli.py", line 227, in main
    chat_loop(
  File "/github/FastChat/fastchat/serve/inference.py", line 532, in chat_loop
    outputs = chatio.stream_output(output_stream)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/github/FastChat/fastchat/serve/cli.py", line 63, in stream_output
    for outputs in output_stream:
  File "/opt/conda/envs/fc/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 56, in generator_context
    response = gen.send(request)
               ^^^^^^^^^^^^^^^^^
  File "/github/FastChat/fastchat/serve/inference.py", line 160, in generate_stream
    out = model(
          ^^^^^^
  File "/opt/conda/envs/fc/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/huggingface/modules/transformers_modules/yuan_hf_model.py", line 938, in forward
    outputs = self.model(
              ^^^^^^^^^^^
  File "/opt/conda/envs/fc/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/huggingface/modules/transformers_modules/yuan_hf_model.py", line 768, in forward
    layer_outputs = decoder_layer(
                    ^^^^^^^^^^^^^^
  File "/opt/conda/envs/fc/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/huggingface/modules/transformers_modules/yuan_hf_model.py", line 426, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
                                                          ^^^^^^^^^^^^^^^
  File "/opt/conda/envs/fc/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/huggingface/modules/transformers_modules/yuan_hf_model.py", line 358, in forward
    raise ValueError(
ValueError: Attention mask should be of size (1, 1, 1, 10), but is torch.Size([1, 1, 1, 1]

是否是和yuan_hf_model.py脚本里相关模块的处理有关？

我上述使用推理脚本还是比较常见的，所以如果可以的话，是否可以修复这个问题？

[ljg-ieisystem]

在开发huggingface没有考虑到这种推理情况，可以通过yuan_hf_model.py 中更改以下代码段解决，之后我们会考虑该情况更新对应的代码

if self.training or self.reset_position_ids and attention_mask is not None：
            attention_mask, _ = self._prepare_decoder_attention_mask_training(input_ids1, inputs_embeds, self.eod_token, reset_mask_flag, self.reset_attention_mask, self.reset_position_ids)

[cauwulixuan]

这种推理情况似乎还挺常见的，我本地修改了这个文件，确实可以了，谢谢您。

这种情况仅限于我已经下载了huggingface的模型。如果直接from_pretrained("IEITYuan/Yuan2-2B-hf")，就没法手工修改了吧？后续官方会统一更新吗？

[Shawn-IEITSystems]

@ljg-ieisystem