models package

This package contains most of the popular LLMs . Currently we have support for GPT, Bert, T5 and Retro . This is an ever growing list so keep an eye out.

Subpackages

• models.gpt package
  • Submodules
  • models.gpt.gpt_model module
    • GPTModel
      • build_schedule_plan()
      • forward()
      • set_input_tensor()
      • sharded_state_dict()
      • shared_embedding_or_output_weight()
  • Module contents
• models.t5 package
  • Submodules
  • models.t5.t5_model module
    • T5LMHead
      • forward()
    • T5Model
      • forward()
      • set_input_tensor()
      • sharded_state_dict()
      • shared_embedding_or_output_weight()
    • t5_extended_attention_mask()
    • t5_position_ids()
  • Module contents
• models.bert package
  • Submodules
  • models.bert.bert_model module
    • BertModel
      • bert_extended_attention_mask()
      • bert_position_ids()
      • forward()
      • set_input_tensor()
    • get_te_version()
  • Module contents

Module contents