Framework-specific API

• pyTorch
  • Linear
    • forward
    • set_tensor_parallel_group
  • LayerNorm
  • RMSNorm
  • LayerNormLinear
    • forward
    • set_tensor_parallel_group
  • LayerNormMLP
    • forward
    • set_tensor_parallel_group
  • DotProductAttention
    • forward
    • set_context_parallel_group
  • MultiheadAttention
    • forward
    • set_context_parallel_group
    • set_tensor_parallel_group
  • TransformerLayer
    • forward
    • set_context_parallel_group
    • set_tensor_parallel_group
  • InferenceParams
  • CudaRNGStatesTracker
    • add
    • fork
    • get_states
    • reset
    • set_states
  • fp8_autocast
  • fp8_model_init
  • checkpoint
  • onnx_export
• Jax
  • MajorShardingType
  • ShardingType
  • TransformerLayerType
  • ShardingResource
  • fp8_autocast
  • update_collections
  • update_fp8_metas
  • LayerNorm
    • __call__
  • DenseGeneral
    • __call__
  • LayerNormDenseGeneral
    • __call__
  • LayerNormMLP
    • __call__
  • RelativePositionBiases
    • __call__
  • MultiHeadAttention
    • __call__
  • TransformerLayer
    • __call__
  • extend_logical_axis_rules
• paddle
  • Linear
    • forward
  • LayerNorm
  • LayerNormLinear
    • forward
  • LayerNormMLP
    • forward
  • FusedScaleMaskSoftmax
    • forward
  • DotProductAttention
    • forward
  • MultiHeadAttention
    • forward
  • TransformerLayer
    • forward
  • fp8_autocast
  • recompute