Custom Models#

The Riva translation supports bilingual and multilingual models trained in NeMo. Each model must have 1G shared memory available. If not using the quick start path, specify --shm-size for Docker to ensure the models have enough memory to run.

NeMo models can be converted into Riva models using nemo2riva. For example:

nemo2riva 
<nemo_filename> \
--out=<riva_filename> \
--max-dim=<max-dim>

Translation models can be deployed in Riva using riva-build and riva-deploy.

The translation pipeline has a single optional parameter for <model name> called --name. For example:

riva-build translation \
--name <model name> \
<rmir_filename>:<encryption_key> \
<riva_filename>:<encryption_key>

For NeMo Megatron-LLM trained models, replace translation with megatron_translation. For example:

riva-build megatron_translation \
--name <model name> \
<rmir_filename>:<encryption_key> \
<riva_filename>:<encryption_key>

Both model types accept the following flags:

  • <rmir_filename> is the Riva rmir file that is generated

  • <riva_filename> is the name of the riva file to use as input

  • <encryption_key> is the encryption key used during the export of the .riva file

  • <model name> is how to differentiate the model at inference time. The default is riva-nmt.

Example#

riva-build translation \
--name mnmt_en_deesfr_transformer12x2 \
/data/mnmt_en_deesfr_transformer12x2nmt.rmir \
/data/mnmt_en_deesfr_transformer12x2.riva

riva-deploy  -f /data/mnmt_en_deesfr_transformer12x2nmt.rmir  /data/models

Supported Models#

Model

Architecture

Source Language(s)

Target Language(s)

NGC link

megatronnmt_en_any_500m_32

Transformer Encoder-Decoder

English (en)

Any language

Riva Megatron NMT en-any 500M

megatronnmt_any_en_500m_32

Transformer Encoder-Decoder

Any language

English (en)

Riva Megatron NMT any-en 500M

megatronnmt_any_en_1b

Transformer Encoder-Decoder

Any language

English (en)

Riva Megatron NMT any-en 1B

megatronnmt_en_any_1b

Transformer Encoder-Decoder

Any language

English (en)

Riva Megatron NMT en-any 1B

megatronnmt_any_any_1b

Transformer Encoder-Decoder

Any language

Any language

Riva Megatron NMT any-any 1B