Skip to content

fix: Add missing triton warning #392

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: Add missing triton warning #392

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
8 changes: 7 additions & 1 deletion nemo_deploy/nlp/inference/inference_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@
from megatron.core.transformer.module import MegatronModule
from packaging import version

from nemo_export_deploy_common.import_utils import MISSING_NEMO_MSG, UnavailableError
from nemo_export_deploy_common.import_utils import MISSING_NEMO_MSG, MISSING_TRITON_MSG, UnavailableError

from .tron_utils import (
DistributedInitConfig,
Expand Down Expand Up @@ -201,6 +201,8 @@ def load_nemo_checkpoint_to_tron_model(model: List[MegatronModule], path: Path,
path (Path): Path to NeMo checkpoint directory
legacy_ckpt (bool): Whether to use legacy checkpoint format
"""
if not HAVE_TRITON:
raise UnavailableError(MISSING_TRITON_MSG)
if not HAVE_NEMO:
raise UnavailableError(MISSING_NEMO_MSG)
weights_dir = ckpt_to_weights_subdir(path, is_saving=False)
Expand Down Expand Up @@ -302,6 +304,8 @@ def setup_model_and_tokenizer_for_inference(
Raises:
ValueError: If checkpoint_path is not a valid NeMo-2.0 checkpoint
"""
if not HAVE_TRITON:
raise UnavailableError(MISSING_TRITON_MSG)
if not HAVE_NEMO:
raise UnavailableError(MISSING_NEMO_MSG)

Expand Down Expand Up @@ -469,6 +473,8 @@ def create_mcore_engine(
- GPTInferenceWrapper: Inference-wrapped model
- Union[MCoreTokenizerWrappper, MegatronTokenizer]: Tokenizer instance
"""
if not HAVE_TRITON:
raise UnavailableError(MISSING_TRITON_MSG)
if not HAVE_NEMO:
raise UnavailableError(MISSING_NEMO_MSG)

Expand Down
Loading