Skip to content

Bug in inference mode? #54

@ZheyuAqaZhang

Description

@ZheyuAqaZhang

I found that, in the code, if hidden_states.shape[1] != 1: is used to detect prefill stage.

I thought this criteria sometimes goes wrong. In my experiments, the hidden_states.shape was:

  • hidden_states.shape[1] = n: prefill
  • hidden_states.shape[1] =1: when generating the 2nd token
  • hidden_states.shape[1] =2: when generating the 3rd token
  • hidden_states.shape[1] =3: when generating the 4th token
    ...

So I have two questions:

  1. Can hidden_states.shape[1] consistently equals to 1 (instead of 1,2,3,...) after prefilling if I use some generating setting?
  2. Would position_ids[0,0] == 0 be a robuster criteria than hidden_states.shape[1] != 1 to detect prefill?

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions