NeMo Megatron API#

Pretraining Model Classes#

Customization Model Classes#

Modules#

class nemo.collections.nlp.modules.common.megatron.module.MegatronModule(*args: Any, **kwargs: Any)[source]#

Bases: Module

Megatron specific extensions of torch Module with support for pipelining.

decoder_cross_attention_relative_position_embeddings_weight()[source]#
decoder_relative_position_embeddings_weight()[source]#
encoder_relative_position_embeddings_weight()[source]#
initialize_word_embeddings(init_method, vocab_size, hidden_size, param_dtype=torch.float32)[source]#
position_embeddings_weight()[source]#
state_dict_for_save_checkpoint(destination=None, prefix='', keep_vars=False)[source]#

Use this function to override the state dict for saving checkpoints.

sync_initial_decoder_cross_attention_relative_position_embeddings()[source]#
sync_initial_decoder_relative_position_embeddings()[source]#
sync_initial_encoder_relative_position_embeddings()[source]#
sync_initial_position_embeddings()[source]#
sync_initial_word_embeddings()[source]#
word_embeddings_weight()[source]#
class nemo.collections.nlp.modules.common.megatron.module.Float16Module(*args: Any, **kwargs: Any)[source]#

Bases: MegatronModule

decoder_cross_attention_relative_position_embeddings_weight()[source]#
decoder_relative_position_embeddings_weight()[source]#
encoder_relative_position_embeddings_weight()[source]#
forward(*inputs, **kwargs)[source]#
position_embeddings_weight()[source]#
set_input_tensor(input_tensor)[source]#
state_dict(destination=None, prefix='', keep_vars=False)[source]#
state_dict_for_save_checkpoint(destination=None, prefix='', keep_vars=False)[source]#

Use this function to override the state dict for saving checkpoints.

word_embeddings_weight()[source]#

Datasets#

class nemo.collections.nlp.data.language_modeling.megatron.gpt_dataset.GPTDataset(*args: Any, **kwargs: Any)[source]#

Bases: Dataset

create_data_mmap()[source]#
class nemo.collections.nlp.data.language_modeling.megatron.gpt_dataset.MockGPTDataset(*args: Any, **kwargs: Any)[source]#

Bases: Dataset

class nemo.collections.nlp.data.language_modeling.megatron.bert_dataset.BertDataset(*args: Any, **kwargs: Any)[source]#

Bases: Dataset

Exportable Model Classes#