megatron.model.module.MegatronModule

class megatron.model.module.MegatronModule(share_word_embeddings=True)

Bases: Module

Megatron specific extensions of torch Module with support for pipelining.

state_dict_for_save_checkpoint(prefix='', keep_vars=False)

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