`nemo_deploy.llm.inference.nemo_io`#

IO utilities for loading NeMo 2.0 checkpoints without a direct nemo import.

Copied from the NeMo project (https://github.com/NVIDIA/NeMo). Static from nemo import … statements are removed; the logic is otherwise identical to the upstream sources. When a NeMo checkpoint is actually loaded at runtime, NeMo classes are imported transitively through pydoc.locate — NeMo must therefore still be installed to read NeMo checkpoints.

Sources

IOProtocol : nemo/lightning/io/capture.py
IO helpers, load : nemo/lightning/io/mixin.py
load_context : nemo/lightning/io/api.py
Torch-dtype fiddle registration : nemo/lightning/io/fdl_torch.py

Module Contents#

Classes#

IOProtocol

Functions#

`_set_thread_local_output_dir`	Set output_dir in our thread-local and in NeMo’s (if already imported).
`_register_torch_dtypes`
`_io_transform_args`	Capture init arguments as a plain dict for fdl.Config creation.
`_io_init`	Create an fdl.Config for self from captured init kwargs.
`_io_wrap_init`	Wrap cls.init to populate io on every instance.
`_io_flatten_object`	Flatten an IOMixin object to a form fiddle can serialize.
`_io_unflatten_object`	Unflatten an IOMixin object; load from pickle if it was saved that way.
`_io_path_elements_fn`	Return the path elements for fiddle graph traversal.
`_io_register_serialization`	Register fiddle traversal functions for cls using our _thread_local.
`track_io`	Add fiddle IO functionality to a class or all eligible classes in a module.
`drop_unexpected_params`	Remove deprecated / unexpected parameters from a fiddle Config tree.
`_artifact_transform_load`	Rewrite artifact paths stored in a config to absolute paths.
`load`	Load a fiddle-serialised NeMo checkpoint context from an `io.json` file.
`load_context`	Load a NeMo TrainerContext (or a subpath of it) from a checkpoint directory.

Data#

`_logger`
`_thread_local`
`_TORCH_DTYPE_NAMES`
`SelfT`

API#

nemo_deploy.llm.inference.nemo_io._logger = 'getLogger(...)'#

nemo_deploy.llm.inference.nemo_io._thread_local = 'local(...)'#

nemo_deploy.llm.inference.nemo_io._set_thread_local_output_dir(path: pathlib.Path) → None#

Set output_dir in our thread-local and in NeMo’s (if already imported).

NeMo classes registered before our first load call will use NeMo’s _io_unflatten_object, which reads from NeMo’s own _thread_local. We mirror the value there so that pickle-based artifacts still resolve correctly even in that edge case.

nemo_deploy.llm.inference.nemo_io._TORCH_DTYPE_NAMES = ['bool', 'uint8', 'int8', 'int16', 'int32', 'int64', 'float16', 'bfloat16', 'float32', 'float64', 'c...#

nemo_deploy.llm.inference.nemo_io._register_torch_dtypes() → None#

nemo_deploy.llm.inference.nemo_io.SelfT = 'TypeVar(...)'#

class nemo_deploy.llm.inference.nemo_io.IOProtocol#

Bases: typing.Protocol, typing.Generic[nemo_deploy.llm.inference.nemo_io.SelfT]

property __io__: fiddle.Config[nemo_deploy.llm.inference.nemo_io.SelfT]#

nemo_deploy.llm.inference.nemo_io._io_transform_args(

self,

init_fn,

*args,

**kwargs,

) → Dict[str, Any]#: Capture init arguments as a plain dict for fdl.Config creation.

nemo_deploy.llm.inference.nemo_io._io_init(self, **kwargs) → fiddle.Config#: Create an fdl.Config for self from captured init kwargs.

nemo_deploy.llm.inference.nemo_io._io_wrap_init(cls)#: Wrap cls.init to populate io on every instance.

nemo_deploy.llm.inference.nemo_io._io_flatten_object(instance)#: Flatten an IOMixin object to a form fiddle can serialize.

nemo_deploy.llm.inference.nemo_io._io_unflatten_object(values, metadata)#: Unflatten an IOMixin object; load from pickle if it was saved that way.

nemo_deploy.llm.inference.nemo_io._io_path_elements_fn(x)#: Return the path elements for fiddle graph traversal.

nemo_deploy.llm.inference.nemo_io._io_register_serialization(cls) → None#: Register fiddle traversal functions for cls using our _thread_local.

nemo_deploy.llm.inference.nemo_io.track_io(target, artifacts=None)#

Add fiddle IO functionality to a class or all eligible classes in a module.

Copied from nemo.lightning.io.mixin.track_io.

nemo_deploy.llm.inference.nemo_io.drop_unexpected_params(config: fiddle.Config) → bool#

Remove deprecated / unexpected parameters from a fiddle Config tree.

Copied from nemo.lightning.io.mixin.drop_unexpected_params.

nemo_deploy.llm.inference.nemo_io._artifact_transform_load( cfg: fiddle.Config, path: pathlib.Path, ) → None#

Rewrite artifact paths stored in a config to absolute paths.

Copied from nemo.lightning.io.mixin._artifact_transform_load.

nemo_deploy.llm.inference.nemo_io.load( path: pathlib.Path, output_type: Any = None, subpath: Optional[str] = None, build: bool = True, ) → Any#

Load a fiddle-serialised NeMo checkpoint context from an io.json file.

Copied from nemo.lightning.io.mixin.load.

nemo_deploy.llm.inference.nemo_io.load_context( path: pathlib.Path, subpath: Optional[str] = None, build: bool = True, ) → Any#

Load a NeMo TrainerContext (or a subpath of it) from a checkpoint directory.

Copied from nemo.lightning.io.api.load_context.

nemo_deploy.llm.inference.nemo_io#