nn

Common torch.nn.Module implementations.

Submodules

• nn.attention
  • SlidingWindowAttentionConfig
  • GateGranularity
  • GateConfig
  • AttentionType
  • AttentionBackendName
  • AttentionBackend
  • TorchAttentionBackend
  • FlashAttention2Backend
  • FlashAttention3Backend
  • FlashAttention4Backend
  • TEAttentionBackend
  • AttentionConfig
  • Attention
  • FusedAttention
  • NormalizedAttention
  • RingAttentionLoadBalancerType
  • RingAttentionLoadBalancer
  • RingAttentionZigZagLoadBalancer
  • RingAttentionLlama3LoadBalancer
  • UlyssesLoadBalancer
  • RingContextParallelStyle
  • UlyssesContextParallelStyle
  • GatedDeltaNetConfig
  • GatedDeltaNet
• nn.conversion
  • StateConverter
  • StateMapping
  • StateMappingTemplate
  • TemplatePlaceholder
• nn.feed_forward
  • ActivationFunction
  • FeedForwardType
  • FeedForwardConfig
  • FeedForward
  • NormalizedFeedForward
• nn.functional
  • cross_entropy_loss()
  • fused_linear_cross_entropy_loss()
• nn.hf
  • convert_checkpoint_to_hf()
  • convert_hybrid_state_to_hf()
  • convert_state_from_hf()
  • convert_state_to_hf()
  • get_converter_from_hf()
  • get_converter_to_hf()
  • get_hf_config()
  • get_hybrid_hf_config()
  • get_hybrid_layer_types()
  • is_olmo_hybrid_model()
  • load_config()
  • load_hf_model()
  • save_hf_hybrid_model()
  • save_hf_model()
  • HF_TO_OLMO_CORE_WEIGHT_MAPPINGS
  • HF_TO_OLMO_CORE_MODULE_MAPPINGS
  • MODEL_TYPE_SPECIFIC_HF_TO_OLMO_CORE_WEIGHT_MAPPINGS
  • MODEL_TYPE_SPECIFIC_HF_TO_OLMO_CORE_MODULE_MAPPINGS
  • HF_TO_OLMO_CORE_TEMPLATE_MAPPINGS
  • OLMO_CORE_TO_HF_WEIGHT_MAPPINGS
  • OLMO_CORE_TO_HF_MODULE_MAPPINGS
  • OLMO_CORE_TO_HF_TEMPLATE_MAPPINGS
  • MODEL_TYPE_SPECIFIC_OLMO_CORE_TO_HF_TEMPLATE_MAPPINGS
• nn.layer_norm
  • LayerNormType
  • LayerNormConfig
  • LayerNorm
  • RMSNorm
  • QwenRMSNorm
  • CuTeRMSNorm
  • FusedRMSNorm
  • L2Norm
• nn.lm_head
  • LMHeadType
  • LMLossImplementation
  • LMHeadConfig
  • LMHead
  • NormalizedLMHead
  • LMOutputWithLoss
• nn.moe
  • MoEBase
  • DroplessMoE
  • MoEConfig
  • MoEType
  • MoEMLP
  • DroplessMoEMLP
  • MoERouter
  • MoELinearRouter
  • MoERouterConfig
  • MoERouterType
  • MoERouterGatingFunction
  • MoELoadBalancingLossGranularity
• nn.rope
  • RoPEType
  • RoPEConfig
  • RoPEScalingConfig
  • ABFRoPEScalingConfig
  • PIRoPEScalingConfig
  • StepwiseRoPEScalingConfig
  • YaRNRoPEScalingConfig
  • RotaryEmbeddingBase
  • RotaryEmbedding
  • FusedRotaryEmbedding
  • ComplexRotaryEmbedding
• nn.transformer
  • TransformerType
  • TransformerConfig
  • Transformer
  • NormalizedTransformer
  • MoETransformer
  • MoEHybridTransformerBlockBase
  • MoEHybridTransformerBlock
  • MoEHybridReorderedNormTransformerBlock
  • TransformerBlockType
  • TransformerBlockConfig
  • TransformerBlockBase
  • TransformerBlock
  • ReorderedNormTransformerBlock
  • LayerNormScaledTransformerBlock
  • PeriNormTransformerBlock
  • NormalizedTransformerBlock
  • MoETransformerBlock
  • MoEReorderedNormTransformerBlock
  • InitMethod