gpt_bigcode could also go pageattn

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>
This commit is contained in:
Wang, Yi A 2025-03-18 23:59:31 -07:00
parent 073f793976
commit 2cde30de24

View File

@ -486,8 +486,6 @@ def get_model(
model_type = config_dict["model_type"] model_type = config_dict["model_type"]
if model_type == "gpt_bigcode":
return StarCoder(model_id=model_id, revision=revision, dtype=dtype)
kv_cache_dtype = dtype kv_cache_dtype = dtype
if FLASH_ATTENTION: if FLASH_ATTENTION:
@ -871,6 +869,8 @@ def get_model(
trust_remote_code=trust_remote_code, trust_remote_code=trust_remote_code,
) )
adapt_transformers_to_gaudi() adapt_transformers_to_gaudi()
if model_type == "gpt_bigcode":
return StarCoder(model_id=model_id, revision=revision, dtype=dtype)
if model_type == "bloom": if model_type == "bloom":
return BLOOM( return BLOOM(
model_id=model_id, model_id=model_id,