saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-27 20:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions
onnxruntime/cmake/onnxruntime_rocm_hipify.cmake
Wei-Sheng Chin faef9c32fa
ONNX-Native Tensor Parallel: Using Distributed MatMul as Example (#17695) 
This PR introduces
- New data structure to represent kernel-level (aka node-level or
op-level) tensor sharding informaiton. I consider it as the
fundamentaion of ONNX distribtued inference.
- Building blocks for distribtued kernels implementation especially
stateless implementation for communication ops.
- Implementation of DistributedMatMul and its tests.

Code structure:
- sharding.h/.cc: Function to shard and reshard tensors (calling into
NCCL).
- sharding_spec.h/.cc: Representation of how a tensor is sharded.
- distributed_matmul.h/.cc: Implementation of tensor parallel MatMul.
Inputs and outputs are sharded across devices.
- onnxruntime_test_distributed.py: distributed operator tests.

Example of specifying sharding information
```python
        @onnxscript.script()
        def matmul_rs_sr_rr(tensor_x: FLOAT, tensor_w: FLOAT) -> FLOAT:
            # Run MatMul by sharding x along column axis and w along row axis on
            # 2 GPUs.
            return MICROSOFT_OPSET.DistributedMatMul(
                tensor_x,
                tensor_w,
                device_mesh_shape=[2],
                device_mesh_elements=[0, 1],
                input_shard_specs=["RS[0]", "S[0]R"],
                output_shard_specs=["RR"],
            )
        onnx_model = matmul_rs_sr_rr.to_model_proto(
            input_types=[FLOAT[2, "s"], FLOAT["s", 2]],
            output_types=[FLOAT[2, 2]],
        )
```

In this example, the device mesh can be visualized as 1-D tensor, `[0,
1]`. The 2nd axis of `tensor_x` is sharded across `[0, 1]` (i.e., the
0-axis of the device mesh). Similarly, the 1st axis of `tensor_w` is
sharded across `[0, 1]` as well.

C++ classes to represent tensor sharding (copied from sharding_spec.h):
```cpp
class DeviceMesh {
 public:
  // [Device Mesh and Tensor Sharding for Tensor Parallel]
  // Device mesh is a tensor of device indices.
  // A tensor can then be partitioned along specific mesh axes.
  //
  // Assume we have 4 GPUs indexed by 0, 1, 2, and 3.
  // Let's consider some examples.
  //  1. 1D device mesh [0, 1, 2, 3]. In this case,
  //     device_mesh_shape is [4] and device_mesh_elements
  //     is [0, 1, 2, 3].
  //     If we want to shard a 2-D tensor along its axis 1, the
  //     corresponding sharding spec is a string "RS[0]".
  //  2. 2D device mesh [[0, 1], [2, 3]]. In this case,
  //     device_mesh_shape is [2, 2] and device_mesh_elements
  //     is [0, 1, 2, 3].
  //     If we want to shard a 2-D tensor's
  //     rows along mesh axis 1 and
  //     columns along mesh axis 0, the
  //     corresponding sharding spec is a string "S[1]S[0]".
  //     If that 2-D tensor's value is np.array([[5, 6], [7, 8]]),
  //     GPU 0/1/2/3 owns 5/7/6/8.  Below is a visualization the sharding
  //     proccess.
  //     - Start with a 2-D device mesh [[0, 1], [2, 3]] and
  //       a 2-D tensor [[5, 6], [7, 8]]
  //       - GPU: [[0, 1], [2, 3]], Tensor: [[5, 6], [7, 8]]
  //     - Split GPU mesh along axis 1 and tensor along
  //       axis 0 for "S[1]" in "S[1]S[0]"
  //       - GPU: [[0], [2]], Tensor: [[5, 6]]
  //         GPU: [[1], [3]], Tensor: [[7, 8]]
  //     - Split GPU mesh along axis 0 and tensor along
  //       axis 1 for "S[0]" in "S[1]S[0]"
  //       - GPU: [[0]], Tensor: [[5]]
  //       - GPU: [[2]], Tensor: [[6]]
  //       - GPU: [[1]], Tensor: [[7]]
  //       - GPU: [[3]], Tensor: [[8]]

  // Actual shape of device mesh represented by `device_mesh_elements`.
  std::vector<int64_t> device_mesh_shape;

  // Flattened device mesh.
  std::vector<int64_t> device_mesh_elements;
};

class AxisPartitionSpec {
  // [Device Mesh and Tensor Sharding for Tensor Parallel]
  // This class is the in-memory representation of
  //  1. if a tensor is sharded or not (aka replica), and
  //  2. which tensor axis is shard by which device mesh axis.
  // Let's consider sharding 2-D tensor along column axis on
  // device mesh [0, 1] as an example.
  // The required sharding spec RS[0] can be represented by
  // - AxisPartitionSpec(Condition::Replica, -1)
  // - AxisPartitionSpec(Condition::Shard, 0)
 public:
  // Status of a tensor axis.
  // A tensor axis can be either sharded or replicated
  // along a device mesh axis.
  enum class Condition { Replica,
                         Shard };

  // This field tells if a tensor axis is sharded or not.
  Condition cond;

  // If a tensor axis is sharded, this field tells which device
  // mesh axis to distribute the shards along.
  // If a tensor axis is not sharded, this field is ignored.
  int device_mesh_axis;

  // A helper to construct a replica spec for a tensor axis.
  static AxisPartitionSpec CreateReplica() {
    return AxisPartitionSpec(Condition::Replica, -1);
  }

  // A helper to construct a sharding spec for a tensor axis.
  // This tensor axis is sharded along `device_mesh_axis` in device mesh.
  static AxisPartitionSpec CreateShard(int device_mesh_axis) {
    return AxisPartitionSpec(Condition::Shard, device_mesh_axis);
  }
};

class TensorPartitionSpec {
  // [Device Mesh and Tensor Sharding for Tensor Parallel]
  // TensorPartitionSpec holds a collection of AxisPartitionSpec and an
  // associated DeviceMesh. It is responsible for determining how a tensor
  // should be partitioned across a device mesh.
  //
  // Example 1: RS[0]
  // In this scenario, `axis_specs` would contain two `AxisPartitionSpec` objects.
  // - The first object is a Replica, denoting that the first axis of the tensor is
  //   not sharded but is instead replicated.
  // - The second object is a Shard along the 0-th axis of the device mesh. It denotes
  //   that the second axis of the tensor is sharded along the first axis of the
  //   device mesh.
  //
  // Example 2: S[0]RR
  // In this scenario, `axis_specs` would contain three `AxisPartitionSpec` objects.
  // - The first object is a Shard along the 0-th axis of the device mesh, indicating
  //   that the first axis of the tensor is sharded along the first axis of the
  //   device mesh.
  // - The second and third objects are Replicas, indicating that the second and third
  //   axes of the tensor are not sharded but are instead replicated.
 public:
  // axis_specs[i]: AxisPartitionSpec for tensor axis i. For a 2-D tensor,
  //                axis_specs[0] is for row axis and axis_specs[1] is for
  //                column axis. axis_specs[i].device_mesh_axis = j means that
  //                tensor axis i is sharded along device mesh axis j.
  std::vector<AxisPartitionSpec> axis_specs;

  // device_mesh: DeviceMesh for sharding the associated tensor.
  // Read [Device Mesh and Tensor Sharding for Tensor Parallel] in DeviceMesh's comment.
  DeviceMesh device_mesh;
};
```
2023-10-05 14:22:25 -07:00

268 lines
8.7 KiB
CMake
Raw Blame History

# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
find_package(Python3 COMPONENTS Interpreter REQUIRED)
# GLOB pattern of file to be excluded
set(contrib_ops_excluded_files
"bert/attention.cc"
"bert/attention.h"
"bert/attention_impl.cu"
"bert/attention_softmax.h"
"bert/attention_softmax.cu"
"bert/attention_prepare_qkv.cu"
"bert/decoder_attention_impl.h"
"bert/decoder_attention_impl.cu"
"bert/decoder_masked_multihead_attention.h"
"bert/decoder_masked_multihead_attention.cc"
"bert/decoder_masked_self_attention.h"
"bert/decoder_masked_self_attention.cc"
"bert/fastertransformer_decoder_attention/*"
"bert/multihead_attention.cc"
"bert/multihead_attention.h"
"bert/fast_gelu_impl.cu"
"bert/fast_gelu_impl.h"
"bert/fast_gelu.cc"
"bert/fast_gelu.h"
"bert/relative_attn_bias.cc"
"bert/relative_attn_bias.h"
"bert/relative_attn_bias_impl.cu"
"bert/relative_attn_bias_impl.h"
"bert/skip_layer_norm.cc"
"bert/skip_layer_norm.h"
"bert/skip_layer_norm_impl.cu"
"bert/skip_layer_norm_impl.h"
"bert/cutlass_fmha/*"
"bert/tensorrt_fused_multihead_attention/*"
"bert/transformer_common.h"
"bert/transformer_common.cc"
"bert/packed_attention.h"
"bert/packed_attention.cc"
"bert/packed_attention_impl.h"
"bert/packed_attention_impl.cu"
"bert/packed_multihead_attention.h"
"bert/packed_multihead_attention.cc"
"bert/packed_multihead_attention_impl.h"
"bert/packed_multihead_attention_impl.cu"
"diffusion/group_norm.cc"
"diffusion/group_norm_impl.cu"
"diffusion/group_norm_impl.h"
"diffusion/nhwc_conv.cc"
"math/complex_mul.cc"
"math/complex_mul.h"
"math/complex_mul_impl.cu"
"math/complex_mul_impl.h"
"math/cufft_plan_cache.h"
"math/fft_ops.cc"
"math/fft_ops.h"
"math/fft_ops_impl.cu"
"math/fft_ops_impl.h"
"quantization/attention_quantization.cc"
"quantization/attention_quantization.h"
"quantization/attention_quantization_impl.cu"
"quantization/attention_quantization_impl.cuh"
"quantization/quantize_dequantize_linear.cc"
"quantization/qordered_ops/qordered_attention_impl.cu"
"quantization/qordered_ops/qordered_attention_impl.h"
"quantization/qordered_ops/qordered_attention_input_enum.h"
"quantization/qordered_ops/qordered_attention.cc"
"quantization/qordered_ops/qordered_attention.h"
"quantization/qordered_ops/qordered_common.cuh"
"quantization/qordered_ops/qordered_layer_norm.h"
"quantization/qordered_ops/qordered_layer_norm.cc"
"quantization/qordered_ops/qordered_layer_norm_impl.h"
"quantization/qordered_ops/qordered_layer_norm_impl.cu"
"quantization/qordered_ops/qordered_longformer_attention.cc"
"quantization/qordered_ops/qordered_longformer_attention.h"
"quantization/qordered_ops/qordered_matmul.h"
"quantization/qordered_ops/qordered_matmul.cc"
"quantization/qordered_ops/qordered_matmul_utils.h"
"quantization/qordered_ops/qordered_matmul_utils.cc"
"quantization/qordered_ops/qordered_qdq_impl.cu"
"quantization/qordered_ops/qordered_qdq_impl.h"
"quantization/qordered_ops/qordered_qdq.cc"
"quantization/qordered_ops/qordered_qdq.h"
"quantization/qordered_ops/qordered_unary_ops.h"
"quantization/qordered_ops/qordered_unary_ops.cc"
"quantization/qordered_ops/qordered_unary_ops_impl.h"
"quantization/qordered_ops/qordered_unary_ops_impl.cu"
"tensor/crop.cc"
"tensor/crop.h"
"tensor/crop_impl.cu"
"tensor/crop_impl.h"
"tensor/dynamicslice.cc"
"tensor/image_scaler.cc"
"tensor/image_scaler.h"
"tensor/image_scaler_impl.cu"
"tensor/image_scaler_impl.h"
"transformers/greedy_search.cc"
"transformers/greedy_search.h"
"conv_transpose_with_dynamic_pads.cc"
"conv_transpose_with_dynamic_pads.h"
"cuda_contrib_kernels.cc"
"cuda_contrib_kernels.h"
"inverse.cc"
"fused_conv.cc"
)
if (NOT onnxruntime_ENABLE_ATEN)
list(APPEND contrib_ops_excluded_files "aten_ops/aten_op.cc")
endif()
if (NOT onnxruntime_USE_NCCL)
# Those are string patterns to exclude. Do NOT use stars such as
# collective/*.cc or *.h.
list(APPEND contrib_ops_excluded_files "collective/nccl_kernels.cc")
list(APPEND contrib_ops_excluded_files "collective/sharding.cc")
list(APPEND contrib_ops_excluded_files "collective/sharding_spec.cc")
list(APPEND contrib_ops_excluded_files "collective/distributed_matmul.cc")
endif()
set(provider_excluded_files
"atomic/common.cuh"
"controlflow/loop.cc"
"controlflow/loop.h"
"controlflow/scan.cc"
"controlflow/scan.h"
"cu_inc/common.cuh"
"math/einsum_utils/einsum_auxiliary_ops.cc"
"math/einsum_utils/einsum_auxiliary_ops.h"
"math/einsum_utils/einsum_auxiliary_ops_diagonal.cu"
"math/einsum_utils/einsum_auxiliary_ops_diagonal.h"
"math/einsum.cc"
"math/einsum.h"
"math/gemm.cc"
"math/matmul.cc"
"math/softmax_impl.cu"
"math/softmax_warpwise_impl.cuh"
"math/softmax_common.cc"
"math/softmax_common.h"
"math/softmax.cc"
"math/softmax.h"
"nn/conv.cc"
"nn/conv.h"
"nn/conv_transpose.cc"
"nn/conv_transpose.h"
"nn/pool.cc"
"nn/pool.h"
"reduction/reduction_ops.cc"
"rnn/cudnn_rnn_base.cc"
"rnn/cudnn_rnn_base.h"
"rnn/gru.cc"
"rnn/gru.h"
"rnn/lstm.cc"
"rnn/lstm.h"
"rnn/rnn.cc"
"rnn/rnn.h"
"rnn/rnn_impl.cu"
"rnn/rnn_impl.h"
"shared_inc/cuda_call.h"
"shared_inc/fpgeneric.h"
"cuda_allocator.cc"
"cuda_allocator.h"
"cuda_call.cc"
"cuda_common.cc"
"cuda_common.h"
"cuda_execution_provider_info.cc"
"cuda_execution_provider_info.h"
"cuda_execution_provider.cc"
"cuda_execution_provider.h"
"cuda_memory_check.cc"
"cuda_memory_check.h"
"cuda_fence.cc"
"cuda_fence.h"
"cuda_fwd.h"
"cuda_kernel.h"
"cuda_pch.cc"
"cuda_pch.h"
"cuda_profiler.cc"
"cuda_profiler.h"
"cuda_provider_factory.cc"
"cuda_provider_factory.h"
"cuda_stream_handle.cc",
"cuda_stream_handle.h",
"cuda_utils.cu"
"cudnn_common.cc"
"cudnn_common.h"
"cupti_manager.cc"
"cupti_manager.h"
"fpgeneric.cu"
"gpu_data_transfer.cc"
"gpu_data_transfer.h"
"integer_gemm.cc"
"tunable/*"
)
set(training_ops_excluded_files
"activation/gelu_grad_impl_common.cuh" # uses custom tanh
"collective/adasum_kernels.cc"
"collective/adasum_kernels.h"
"math/div_grad.cc" # miopen API differs from cudnn, no double type support
"nn/batch_norm_grad.cc" # no double type support
"nn/batch_norm_grad.h" # miopen API differs from cudnn
"nn/batch_norm_internal.cc" # miopen API differs from cudnn, no double type support
"nn/batch_norm_internal.h" # miopen API differs from cudnn, no double type support
"nn/conv_grad.cc"
"nn/conv_grad.h"
"reduction/reduction_all.cc" # deterministic = true, ignore ctx setting
"reduction/reduction_ops.cc" # no double type support
"cuda_training_kernels.cc"
"cuda_training_kernels.h"
"nn/conv_shared.cc"
"nn/conv_shared.h"
"nn/conv_transpose_grad.cc"
"nn/conv_transpose_grad.h"
)
function(auto_set_source_files_hip_language)
foreach(f ${ARGN})
if(f MATCHES ".*\\.cu$")
set_source_files_properties(${f} PROPERTIES LANGUAGE HIP)
endif()
endforeach()
endfunction()
# cuda_dir must be relative to REPO_ROOT
function(hipify cuda_dir in_excluded_file_patterns out_generated_cc_files out_generated_cu_files)
set(hipify_tool ${REPO_ROOT}/tools/ci_build/amd_hipify.py)
file(GLOB_RECURSE srcs CONFIGURE_DEPENDS
"${REPO_ROOT}/${cuda_dir}/cuda/*.h"
"${REPO_ROOT}/${cuda_dir}/cuda/*.cc"
"${REPO_ROOT}/${cuda_dir}/cuda/*.cuh"
"${REPO_ROOT}/${cuda_dir}/cuda/*.cu"
)
# do exclusion
set(excluded_file_patterns ${${in_excluded_file_patterns}})
list(TRANSFORM excluded_file_patterns PREPEND "${REPO_ROOT}/${cuda_dir}/cuda/")
file(GLOB_RECURSE excluded_srcs CONFIGURE_DEPENDS ${excluded_file_patterns})
foreach(f ${excluded_srcs})
message(STATUS "Excluded from hipify: ${f}")
endforeach()
list(REMOVE_ITEM srcs ${excluded_srcs})
foreach(f ${srcs})
file(RELATIVE_PATH cuda_f_rel "${REPO_ROOT}" ${f})
string(REPLACE "cuda" "rocm" rocm_f_rel ${cuda_f_rel})
set(f_out "${CMAKE_CURRENT_BINARY_DIR}/amdgpu/${rocm_f_rel}")
add_custom_command(
OUTPUT ${f_out}
COMMAND Python3::Interpreter ${hipify_tool}
--hipify_perl ${onnxruntime_HIPIFY_PERL}
${f} -o ${f_out}
DEPENDS ${hipify_tool} ${f}
COMMENT "Hipify: ${cuda_f_rel} -> amdgpu/${rocm_f_rel}"
)
if(f MATCHES ".*\\.cuh?$")
list(APPEND generated_cu_files ${f_out})
else()
list(APPEND generated_cc_files ${f_out})
endif()
endforeach()
set_source_files_properties(${generated_cc_files} PROPERTIES GENERATED TRUE)
set_source_files_properties(${generated_cu_files} PROPERTIES GENERATED TRUE)
auto_set_source_files_hip_language(${generated_cu_files})
set(${out_generated_cc_files} ${generated_cc_files} PARENT_SCOPE)
set(${out_generated_cu_files} ${generated_cu_files} PARENT_SCOPE)
endfunction()
Reference in a new issue View git blame Copy permalink