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> \

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 any-en 1B

mnmt_deesfr_en_transformer12x2

Transformer 12x2

German (de), Spanish (es), French (fr)

English (en)

NMT Multilingual De/Es/Fr En Transformer12x2

mnmt_en_deesfr_transformer12x2

Transformer 12x2

English (en)

German (de), Spanish (es), French (fr)

NMT Multilingual En De/Es/Fr Transformer12x2

mnmt_deesfr_en_transformer24x6

Transformer 24x6

German (de), Spanish (es), French (fr)

English (en)

NMT Multilingual De/Es/Fr En Transformer24x6

mnmt_en_deesfr_transformer24x6

Transformer 24x6

English (en)

German (de), Spanish (es), French (fr)

NMT Multilingual En De/Es/Fr Transformer24x6

en_de_24x6

Transformer 24x6

English (en)

German (de)

NMT En - De Transformer24x6

en_es_24x6

Transformer 24x6

English (en)

Spanish (es)

NMT En Es Transformer24x6

en_fr_24x6

Transformer 24x6

English (en)

French (fr)

NMT En Fr Transformer24x6

en_ru_24x6

Transformer 24x6

English (en)

Russian (ru)

NMT En Ru Transformer24x6

en_zh_24x6

Transformer 24x6

English (en)

Simplified Chinese (zh)

NMT En Zh Transformer24x6

de_en_24x6

Transformer 24x6

German (de)

English (en)

NMT De En Transformer24x6

es_en_24x6

Transformer 24x6

Spanish (es)

English (en)

NMT Es En Transformer24x6

fr_en_24x6

Transformer 24x6

French (fr)

English (en)

NMT Fr En Transformer24x6

ru_en_24x6

Transformer 24x6

Russian (ru)

English (en)

NMT Ru En Transformer24x6

zh_en_24x6

Transformer 24x6

Simplified Chinese (zh)

English (en)

NMT Zh En Transformer24x6