Framework-specific API

• pyTorch
  • Linear
    • Linear.forward()
    • Linear.set_tensor_parallel_group()
  • GroupedLinear
    • GroupedLinear.forward()
    • GroupedLinear.set_tensor_parallel_group()
  • LayerNorm
  • RMSNorm
  • LayerNormLinear
    • LayerNormLinear.forward()
    • LayerNormLinear.set_tensor_parallel_group()
  • LayerNormMLP
    • LayerNormMLP.forward()
    • LayerNormMLP.set_tensor_parallel_group()
  • DotProductAttention
    • DotProductAttention.forward()
    • DotProductAttention.set_context_parallel_group()
  • MultiheadAttention
    • MultiheadAttention.forward()
    • MultiheadAttention.set_context_parallel_group()
    • MultiheadAttention.set_tensor_parallel_group()
  • TransformerLayer
    • TransformerLayer.forward()
    • TransformerLayer.set_context_parallel_group()
    • TransformerLayer.set_tensor_parallel_group()
  • CudaRNGStatesTracker
    • CudaRNGStatesTracker.add()
    • CudaRNGStatesTracker.fork()
    • CudaRNGStatesTracker.get_states()
    • CudaRNGStatesTracker.reset()
    • CudaRNGStatesTracker.set_states()
  • fp8_autocast()
  • fp8_model_init()
  • checkpoint()
  • make_graphed_callables()
  • get_cpu_offload_context()
  • moe_permute()
  • moe_permute_with_probs()
  • moe_unpermute()
  • moe_sort_chunks_by_index()
  • moe_sort_chunks_by_index_with_probs()
  • initialize_ub()
  • destroy_ub()
• Jax
  • Pre-defined Variable of Logical Axes
  • Modules
    • TransformerLayerType
    • MeshResource
      • MeshResource.dp_resource
      • MeshResource.tp_resource
      • MeshResource.fsdp_resource
      • MeshResource.pp_resource
      • MeshResource.cp_resource
    • fp8_autocast()
    • update_collections()
    • LayerNorm
      • LayerNorm.__call__()
    • DenseGeneral
      • DenseGeneral.__call__()
    • LayerNormDenseGeneral
      • LayerNormDenseGeneral.__call__()
    • LayerNormMLP
      • LayerNormMLP.__call__()
    • RelativePositionBiases
      • RelativePositionBiases.__call__()
    • DotProductAttention
      • DotProductAttention.__call__()
    • MultiHeadAttention
      • MultiHeadAttention.__call__()
    • TransformerLayer
      • TransformerLayer.__call__()
    • extend_logical_axis_rules()