mirror/OpenLLM
1
0
Fork 0
mirror of https://github.com/bentoml/OpenLLM.git synced 2026-02-19 15:18:12 -05:00
Code Issues Packages Projects Releases Wiki Activity

fix(torch_dtype): load eagerly (#631)

Browse Source 
Signed-off-by: Aaron <29749331+aarnphm@users.noreply.github.com>
This commit is contained in:
Aaron Pham
2023-11-13 13:48:04 -05:00
committed by GitHub
parent 0924c0b34d
commit d358e68539
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
openllm-python/src/openllm/_llm.py
View File

@@ -249,11 +249,13 @@ class LLM(t.Generic[M, T], ReprMixin):
import transformers
if not isinstance(self.__llm_torch_dtype__, torch.dtype):
config_dtype = getattr(
transformers.AutoConfig.from_pretrained(self.bentomodel.path, trust_remote_code=self.trust_remote_code),
'torch_dtype',
None,
)
try:
hf_config = transformers.AutoConfig.from_pretrained(
self.bentomodel.path, trust_remote_code=self.trust_remote_code
)
except OpenLLMException:
hf_config = transformers.AutoConfig.from_pretrained(self.model_id, trust_remote_code=self.trust_remote_code)
config_dtype = getattr(hf_config, 'torch_dtype', None)
if config_dtype is None:
config_dtype = torch.float32
if self.__llm_torch_dtype__ == 'auto':