class_docstr: Base Megatron module inhertied by all Models.
  
  Megatron specific extensions of torch Module with support
  for pipelining
  
  Args:
      config (TransformerConfig): Transformer config
docstr: Override sharded state dict with Dist Checkpointing.
  
  Override sharded_state_dict when using distributed checkpointing. keep_vars must always be set to True so that optimizer states can be sharded.
  
  Args:
      prefix (str, optional): _description_. Defaults to ''.
  
  Returns:
      _type_: _description_
docstr_tok: ["Override","sharded","state","dict","with","Dist","Checkpointing",".","Override","sharded_state_dict","when","using","distributed","checkpointing",".","keep_vars","must","always","be","set","to","True","so","that","optimizer","states","can","be","sharded",".","Args",":","prefix","(","str",",","optional",")",":","_description_",".","Defaults","to","``",".","Returns",":","_type_",":","_description_"]
code: def sharded_state_dict(self, prefix: str = ""):
      """Override sharded state dict with Dist Checkpointing.
  
      Override sharded_state_dict when using distributed checkpointing. keep_vars must always be set to True so that optimizer states can be sharded.
  
      Args:
          prefix (str, optional): _description_. Defaults to ''.
  
      Returns:
          _type_: _description_
      """
      return self.state_dict(prefix=prefix, keep_vars=True)
  
imports: import torch
  from torch.autograd import Variable
  from torch.nn.parameter import Parameter
  from megatron.core import parallel_state
  from megatron.core.transformer.transformer_config import TransformerConfig
label_desc: Use image node_id 3 for calling the MegatronModule obj's underlying member method code with example usage: obj.sharded_state_dict(prefix) and returns: self