Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 17 additions & 3 deletions tensorrt_llm/bench/build/build.py
Original file line number Diff line number Diff line change
Expand Up @@ -86,9 +86,23 @@ def get_model_config(model_name: str, model_path: Path = None) -> ModelConfig:
Raises:
ValueError: When model is not supported.
"""
if is_nemotron_hybrid(
AutoConfig.from_pretrained(model_path or model_name,
trust_remote_code=True)):
# TODO: remove this once the transformers can support all of those models in _CONFIG_REGISTRY
import transformers
from tensorrt_llm._torch.model_config import _CONFIG_REGISTRY
model_name_or_path = model_path or model_name
config_dict, _ = transformers.PretrainedConfig.get_config_dict(
model_name_or_path,
trust_remote_code=True)
model_type = config_dict.get("model_type")
if model_type in _CONFIG_REGISTRY:
config_class = _CONFIG_REGISTRY[model_type]
model_config = config_class.from_pretrained(
model_name_or_path,
trust_remote_code=True)
else:
model_config = AutoConfig.from_pretrained(model_name_or_path,
trust_remote_code=True)
if is_nemotron_hybrid(model_config):
Comment thread
lfr-0531 marked this conversation as resolved.
Outdated
return NemotronHybridConfig.from_hf(model_name, model_path)
return ModelConfig.from_hf(model_name, model_path)

Expand Down
18 changes: 16 additions & 2 deletions tensorrt_llm/bench/build/dataclasses.py
Original file line number Diff line number Diff line change
Expand Up @@ -192,9 +192,23 @@ def get_param_count(cls, model_hf_name, hf_model_path):

@classmethod
def from_hf(cls, model_hf_name, hf_model_path):
# TODO: remove this once the transformers can support all of those models in _CONFIG_REGISTRY
import transformers
from tensorrt_llm._torch.model_config import _CONFIG_REGISTRY
model_name_or_path = hf_model_path or model_hf_name
hf_config = AutoConfig.from_pretrained(
model_name_or_path, trust_remote_code=True).to_dict()
config_dict, _ = transformers.PretrainedConfig.get_config_dict(
model_name_or_path,
trust_remote_code=True)
model_type = config_dict.get("model_type")
if model_type in _CONFIG_REGISTRY:
config_class = _CONFIG_REGISTRY[model_type]
model_config = config_class.from_pretrained(
model_name_or_path,
trust_remote_code=True)
else:
model_config = AutoConfig.from_pretrained(model_name_or_path,
trust_remote_code=True)
hf_config = model_config.to_dict()
Comment thread
lfr-0531 marked this conversation as resolved.
Outdated
param_count = cls.get_param_count(model_hf_name, hf_model_path)

return cls(name=model_hf_name, param_count=param_count, **hf_config)
Expand Down
Loading