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.
- initialize_word_embeddings(init_method, vocab_size, hidden_size, param_dtype=torch.float32)[source]#
- class nemo.collections.nlp.modules.common.megatron.module.Float16Module(*args: Any, **kwargs: Any)[source]#
Bases:
MegatronModule
Datasets#
- class nemo.collections.nlp.data.language_modeling.megatron.gpt_dataset.GPTDataset(*args: Any, **kwargs: Any)[source]#
Bases:
Dataset