`nemo_export.vllm.model_loader`#

Module Contents#

Classes#

NemoModelLoader

Implements a custom ModelLoader for vLLM that reads the weights from a Nemo checkpoint and converts them to a vLLM compatible format at load time.

Data#

API#

nemo_export.vllm.model_loader.LOGGER = 'getLogger(...)'#

class nemo_export.vllm.model_loader.NemoModelLoader(/, *args, **kw)#

Bases: vllm.model_executor.model_loader.BaseModelLoader

Implements a custom ModelLoader for vLLM that reads the weights from a Nemo checkpoint and converts them to a vLLM compatible format at load time.

Also supports an ahead-of-time conversion that stores new weights in a Safetensors file, see convert_and_store_nemo_weights(…)

Initialization

static _load_nemo_checkpoint_state( nemo_file: str, ) → Dict[str, Any]#

abstractmethod download_model(model_config: vllm.config.ModelConfig) → None#

load_model( *, vllm_config: nemo_export.vllm.model_config.NemoModelConfig, ) → torch.nn.Module#: Overrides the load_model function from BaseModelLoader to convert Nemo weights at load time.

static convert_and_store_nemo_weights( model_config: nemo_export.vllm.model_config.NemoModelConfig, safetensors_file: str, )#: Converts Nemo weights and stores the converted weights in a Safetensors file.

static _standardize_nemo2_naming( state_dict: Dict[str, Any], ) → Dict[str, Any]#