search

SafeAILab / EAGLE

Official Implementation of EAGLE-1 and EAGLE-2
https://arxiv.org/pdf/2406.16858
Apache License 2.0
679 stars 67 forks source link

can eagle support #90

Open zhangxy1234 opened 3 weeks ago

zhangxy1234 commented 3 weeks ago

when I inference multi batch, it raise an Error

RuntimeError: output with shape [1, 32, 84, 128] doesn't match the broadcast shape [2, 32, 84, 128]

hongyanz commented 2 weeks ago

Only EAGLE-1 supports bs > 1 now.

haiduo commented 2 weeks ago

Is it because of the dynamic tree attention mechanism that Eagle2 does not support batch size>1?

Liyuhui-12 commented 2 weeks ago

EAGLE-2 requires some code modifications to support bs>1, so it temporarily does not support bs>1.

zhangxy1234 commented 1 week ago

Only EAGLE-1 supports bs > 1 now.

can you show me an example to use bs = 2 ?

I create model and prompts this way but it raise this error model = EaModel.from_pretrained( base_model_path=base_model_path, ea_model_path=EAGLE_model_path, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, device_map="auto", pretraining_tp = 1, # use if tp>1 )

input_ids=model.tokenizer([prompt1,prompt2]).input_ids

RuntimeError: output with shape [1, 32, 44, 128] doesn't match the broadcast shape [2, 32, 44, 128]

@hongyanz

hongyanz commented 1 week ago

You can go to the v1 branch and check here: https://github.com/SafeAILab/EAGLE/tree/v1?tab=readme-ov-file#batch-size--1