nemo_deploy.llm.inference.nemo_utils

NeMo utility code copied from the NeMo project.

All utilities here are copied directly from NeMo and have no static dependency on the nemo package. When a NeMo checkpoint is loaded at runtime, NeMo classes are imported transitively through pydoc.locate inside nemo_io.load_context — NeMo must therefore still be installed to read NeMo checkpoints.

Sources:

• MCoreTokenizerWrappper : nemo/collections/llm/inference/base.py

• ckpt_to_dir, idempotent_path_append, ckpt_to_context_subdir : nemo/lightning/ckpt_utils.py

• ckpt_to_weights_subdir : nemo/lightning/io/pl.py

• constants : nemo/lightning/ckpt_utils.py

• set_modelopt_spec_* : nemo/collections/llm/modelopt/model_utils.py

• load_context, io : nemo_io.py (copied from nemo/lightning/io/)

Module Contents

Classes

MCoreTokenizerWrappper

Thin wrapper that adapts a NeMo tokenizer to the MCore generate API.

Functions

ckpt_to_dir	Return the checkpoint directory path for a given filepath.
idempotent_path_append	Append suffix to base_dir only when it is not already the last component.
ckpt_to_context_subdir	Return the context sub-directory of a NeMo-2 checkpoint.
ckpt_to_weights_subdir	Return the weights sub-directory of a NeMo-2 checkpoint.
set_modelopt_spec_if_exists_in_ckpt	Set model.config.transformer_layer_spec to a modelopt spec if the checkpoint contains a modelopt_state directory.

Data

_logger
io
GPTConfig
T5Config
WEIGHTS_PATH
CONTEXT_PATH
ADAPTER_META_FILENAME

API

nemo_deploy.llm.inference.nemo_utils._logger = 'getLogger(...)'

nemo_deploy.llm.inference.nemo_utils.io = 'SimpleNamespace(...)'

nemo_deploy.llm.inference.nemo_utils.GPTConfig = None

nemo_deploy.llm.inference.nemo_utils.T5Config = None

nemo_deploy.llm.inference.nemo_utils.WEIGHTS_PATH: str = 'weights'

nemo_deploy.llm.inference.nemo_utils.CONTEXT_PATH: str = 'context'

nemo_deploy.llm.inference.nemo_utils.ADAPTER_META_FILENAME: str = 'adapter_metadata.json'

nemo_deploy.llm.inference.nemo_utils.ckpt_to_dir(filepath: Union[str, pathlib.Path]) → pathlib.Path

Return the checkpoint directory path for a given filepath.

PTL treats checkpoints as .ckpt files. This helper strips the extension (appending it first when absent) and returns a :class:Path suitable for use as a distributed-checkpoint directory.

Copied from nemo.lightning.ckpt_utils.ckpt_to_dir with the AdapterPath and MultiStorageClient branches removed.

nemo_deploy.llm.inference.nemo_utils.idempotent_path_append(

base_dir: Union[str, pathlib.Path],

suffix: str,

) → pathlib.Path

Append suffix to base_dir only when it is not already the last component.

Copied from nemo.lightning.ckpt_utils.idempotent_path_append with the AdapterPath and MultiStorageClient branches removed.

nemo_deploy.llm.inference.nemo_utils.ckpt_to_context_subdir(

filepath: Union[str, pathlib.Path],

) → pathlib.Path

Return the context sub-directory of a NeMo-2 checkpoint.

Copied from nemo.lightning.ckpt_utils.ckpt_to_context_subdir.

nemo_deploy.llm.inference.nemo_utils.ckpt_to_weights_subdir(

filepath: Union[str, pathlib.Path],

is_saving: bool,

) → pathlib.Path

Return the weights sub-directory of a NeMo-2 checkpoint.

Copied from nemo.lightning.io.pl.ckpt_to_weights_subdir with the AdapterPath branch removed.

class nemo_deploy.llm.inference.nemo_utils.MCoreTokenizerWrappper(tokenizer, vocab_size=None)

Thin wrapper that adapts a NeMo tokenizer to the MCore generate API.

MCore’s generate pipeline expects tokenizer.detokenize, tokenizer.tokenize, tokenizer.bos, and tokenizer.pad – this wrapper maps those calls to the corresponding NeMo tokenizer methods/properties.

Copied verbatim from nemo.collections.llm.inference.base.MCoreTokenizerWrappper.

Initialization

detokenize(tokens, remove_special_tokens=False)

Detokenize tokens into a string.

tokenize(prompt)

Tokenize prompt into a list of token IDs.

property additional_special_tokens_ids

IDs of additional special tokens.

property bos

Beginning-of-sequence token ID.

property pad

Padding token ID.

nemo_deploy.llm.inference.nemo_utils.set_modelopt_spec_if_exists_in_ckpt(model, path: str) → None

Set model.config.transformer_layer_spec to a modelopt spec if the checkpoint contains a modelopt_state directory.

Copied from nemo.collections.llm.modelopt.model_utils.set_modelopt_spec_if_exists_in_ckpt with NeMo isinstance checks replaced by class-name comparisons.