From 00212f57d2a4044d5dbfad36d5f5efd469013b3d Mon Sep 17 00:00:00 2001 From: luyang Date: Thu, 20 Feb 2025 08:46:59 +0000 Subject: [PATCH 01/16] adaptation of multi devices based on refactor_xccl_primitive --- ...erarchical_sub_task_graph_builder_impl.cpp | 22 +-- .../boxing/sub_task_graph_builder_util.cpp | 4 +- .../boxing/sub_task_graph_builder_util.h | 2 +- oneflow/core/job/eager_ccl_comm_manager.cpp | 2 + oneflow/core/job/eager_ccl_comm_manager.h | 43 ++++++ oneflow/core/job/resource_desc.cpp | 2 +- oneflow/core/job/runtime.cpp | 4 +- .../insert_nccl_logical_op_pass.cpp | 7 +- oneflow/core/job_rewriter/job_completer.cpp | 4 +- .../nccl_logical_chain_strict_order_pass.cpp | 4 +- .../nccl_logical_op_fusion_pass.cpp | 4 +- .../kernel/nccl_send_recv_boxing_kernel.cpp | 62 ++++---- oneflow/user/kernels/eager_nccl_s2s_kernel.cu | 34 ++-- .../kernels/nccl_logical_2d_sbp_kernels.cpp | 145 +++++++++++------- .../kernels/nccl_logical_fusion_kernel.cpp | 53 ++++--- oneflow/user/kernels/nccl_logical_kernels.cpp | 132 +++++++++++----- .../kernels/nccl_logical_send_recv_kernel.cpp | 36 +++-- 17 files changed, 350 insertions(+), 210 deletions(-) diff --git a/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp b/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp index 0562a60f34a..c5e7181de8e 100644 --- a/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp +++ b/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp @@ -99,10 +99,12 @@ class NDNcclSendRecvBoxingSubTskGphBuilder final : public HierarchicalSubTskGphB const LogicalBlobId& lbi, const BlobDesc& logical_blob_desc, const NdSbp& in_nd_sbp, const NdSbp& out_nd_sbp, const Shape& time_shape) const override { - if (in_parallel_desc.device_type() == DeviceType::kCUDA - && out_parallel_desc.device_type() == DeviceType::kCUDA + if (in_parallel_desc.device_type() != DeviceType::kCPU + && out_parallel_desc.device_type() != DeviceType::kCPU && !NdSbpHasPartialParallel(out_nd_sbp)) { -#if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 + // TODO: (zhaoluyang) use WITH_DEVICE or another way to support multi devices(e.g. + // cuda/npu/xpu....) + // #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 ParallelConf merged_parallel_conf; MergeParallelConf(in_parallel_desc.parallel_conf(), out_parallel_desc.parallel_conf(), &merged_parallel_conf); @@ -136,9 +138,9 @@ class NDNcclSendRecvBoxingSubTskGphBuilder final : public HierarchicalSubTskGphB if (has_output) { sorted_out_tasks->push_back(node); } } return BuildSubTskGphBuilderStatus("NDNcclSendRecvBoxingSubTskGphBuilder", ""); -#else - return Error::BoxingNotSupportedError() << "No CUDA or low NCCL version"; -#endif + // #else + // return Error::BoxingNotSupportedError() << "No CUDA or low NCCL version"; + // #endif } else { return Error::BoxingNotSupportedError() << "Partial SBP in the consumer or not running on CUDA"; @@ -267,8 +269,8 @@ Maybe DispatchHierarchicalSubTskGphBuilder::Build( const auto& in_hierarchy = reduced_in_parallel_desc.hierarchy(); const auto& out_hierarchy = reduced_out_parallel_desc.hierarchy(); if ((in_hierarchy->NumAxes() > 2 || out_hierarchy->NumAxes() > 2) - && reduced_in_parallel_desc.device_type() == DeviceType::kCUDA - && reduced_out_parallel_desc.device_type() == DeviceType::kCUDA) { + && reduced_in_parallel_desc.device_type() != DeviceType::kCPU + && reduced_out_parallel_desc.device_type() != DeviceType::kCPU) { return impl_->nd_nccl_send_recv_boxing_sub_tsk_gph_builder_->Build( ctx, sorted_in_tasks, sorted_out_tasks, sorted_ctrl_tasks, reduced_in_parallel_desc, reduced_out_parallel_desc, lbi, logical_blob_desc, reduced_in_nd_sbp, reduced_out_nd_sbp, @@ -285,8 +287,8 @@ Maybe DispatchHierarchicalSubTskGphBuilder::Build( ctx, sorted_in_tasks, sorted_out_tasks, sorted_ctrl_tasks, reduced_in_parallel_desc, reduced_out_parallel_desc, lbi, logical_blob_desc, reduced_in_nd_sbp, reduced_out_nd_sbp, time_shape); - } else if (reduced_in_parallel_desc.device_type() == DeviceType::kCUDA - && reduced_out_parallel_desc.device_type() == DeviceType::kCUDA) { + } else if (reduced_in_parallel_desc.device_type() != DeviceType::kCPU + && reduced_out_parallel_desc.device_type() != DeviceType::kCPU) { return impl_->nd_nccl_send_recv_boxing_sub_tsk_gph_builder_->Build( ctx, sorted_in_tasks, sorted_out_tasks, sorted_ctrl_tasks, reduced_in_parallel_desc, reduced_out_parallel_desc, lbi, logical_blob_desc, reduced_in_nd_sbp, reduced_out_nd_sbp, diff --git a/oneflow/core/graph/boxing/sub_task_graph_builder_util.cpp b/oneflow/core/graph/boxing/sub_task_graph_builder_util.cpp index be88a23693b..c98038878c0 100644 --- a/oneflow/core/graph/boxing/sub_task_graph_builder_util.cpp +++ b/oneflow/core/graph/boxing/sub_task_graph_builder_util.cpp @@ -35,9 +35,9 @@ bool SubTskGphBuilderUtil::HasEmptySliceIfSplit(int64_t parallel_num, } } -bool SubTskGphBuilderUtil::IsOnSameCUDADevice(const TaskNode* lhs, const TaskNode* rhs) { +bool SubTskGphBuilderUtil::IsOnSameDevice(const TaskNode* lhs, const TaskNode* rhs) { return lhs->stream_id().device_id() == rhs->stream_id().device_id() - && lhs->stream_id().device_id().device_type() == DeviceType::kCUDA; + && lhs->stream_id().device_id().device_type() != DeviceType::kCPU; } bool SubTskGphBuilderUtil::IsBoxingS2S(const SbpParallel& src, const SbpParallel& dst) { diff --git a/oneflow/core/graph/boxing/sub_task_graph_builder_util.h b/oneflow/core/graph/boxing/sub_task_graph_builder_util.h index 96b77df44e6..4ff39fff1f5 100644 --- a/oneflow/core/graph/boxing/sub_task_graph_builder_util.h +++ b/oneflow/core/graph/boxing/sub_task_graph_builder_util.h @@ -32,7 +32,7 @@ struct SubTskGphBuilderUtil { static bool IsDeviceTypeCPUOr(const ParallelDesc& parallel_desc, DeviceType device_type); static bool HasEmptySliceIfSplit(int64_t parallel_num, const SbpParallel& sbp_parallel, const BlobDesc& blob_desc); - static bool IsOnSameCUDADevice(const TaskNode* lhs, const TaskNode* rhs); + static bool IsOnSameDevice(const TaskNode* lhs, const TaskNode* rhs); static bool IsBoxingS2S(const SbpParallel& src, const SbpParallel& dst); static bool IsBoxingS2B(const SbpParallel& src, const SbpParallel& dst); static bool IsBoxingP2S(const SbpParallel& src, const SbpParallel& dst); diff --git a/oneflow/core/job/eager_ccl_comm_manager.cpp b/oneflow/core/job/eager_ccl_comm_manager.cpp index c7283caf150..8429e75136a 100644 --- a/oneflow/core/job/eager_ccl_comm_manager.cpp +++ b/oneflow/core/job/eager_ccl_comm_manager.cpp @@ -18,6 +18,8 @@ limitations under the License. namespace oneflow { +const std::string EagerCclCommMgr::kDefaultCclStreamName = "DEFAULT"; + EagerCclCommMgrBuilder& EagerCclCommMgrBuilder::Get() { static EagerCclCommMgrBuilder mgr; return mgr; diff --git a/oneflow/core/job/eager_ccl_comm_manager.h b/oneflow/core/job/eager_ccl_comm_manager.h index 647b4d0a132..3773799f8b9 100644 --- a/oneflow/core/job/eager_ccl_comm_manager.h +++ b/oneflow/core/job/eager_ccl_comm_manager.h @@ -24,6 +24,7 @@ namespace oneflow { class EagerCclCommMgr { public: + static const std::string kDefaultCclStreamName; OF_DISALLOW_COPY_AND_MOVE(EagerCclCommMgr); virtual ~EagerCclCommMgr() = default; @@ -82,6 +83,48 @@ class EagerCclCommMgrBuilder { #define REGISTER_CCL_COMM_MGR(device, Derived) \ COMMAND(EagerCclCommMgrBuilder::Get().RegisterEagerCclCommMgrType(device)) +class UserKernelUnifiedCclCommInitRegistry final { + public: + struct Trigger { + explicit Trigger(const std::string& key) { + UserKernelUnifiedCclCommInitRegistry::Instance().Register(key); + } + }; + + static UserKernelUnifiedCclCommInitRegistry& Instance() { + static UserKernelUnifiedCclCommInitRegistry reg; + return reg; + } + + OF_DISALLOW_COPY_AND_MOVE(UserKernelUnifiedCclCommInitRegistry); + ~UserKernelUnifiedCclCommInitRegistry() = default; + + void Register(const std::string& key) { + bool insert_success = reg_set_.insert(key).second; + if (!insert_success) { + std::cerr << key << " was already registered in CclCommRegistry" << std::endl; + abort(); + } + } + + bool IsRegistered(const std::string& key) const { return reg_set_.find(key) != reg_set_.end(); } + + private: + UserKernelUnifiedCclCommInitRegistry() = default; + std::set reg_set_; +}; + +static const std::string kSystemCclOpPrefix = "sys_op_"; + +#define REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT(op_type_name) \ + static auto OF_PP_CAT(g_nccl_comm_reg_, __COUNTER__) = \ + ::oneflow::UserKernelUnifiedCclCommInitRegistry::Trigger(op_type_name) + +#define REGISTER_SYSTEM_OP_KERNEL_UNIFIED_CCL_COMM_INIT(op_type_case) \ + static auto OF_PP_CAT(g_nccl_comm_reg_, __COUNTER__) = \ + ::oneflow::UserKernelUnifiedCclCommInitRegistry::Trigger(::oneflow::kSystemCclOpPrefix \ + + std::to_string(op_type_case)) + } // namespace oneflow #endif // ONEFLOW_CORE_JOB_EAGER_CCL_COMM_MANAGER_H_ diff --git a/oneflow/core/job/resource_desc.cpp b/oneflow/core/job/resource_desc.cpp index be9940fe571..8137530592b 100644 --- a/oneflow/core/job/resource_desc.cpp +++ b/oneflow/core/job/resource_desc.cpp @@ -74,7 +74,7 @@ bool ResourceDesc::nccl_use_compute_stream() const { #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 return resource_.nccl_use_compute_stream(); #else - return false; + return true; // TODO: find a batter way for multi devices #endif } diff --git a/oneflow/core/job/runtime.cpp b/oneflow/core/job/runtime.cpp index 068484c7a78..a124a87f3d5 100644 --- a/oneflow/core/job/runtime.cpp +++ b/oneflow/core/job/runtime.cpp @@ -70,9 +70,9 @@ Runtime::Runtime( Singleton::Get()->AddPlan(plan); collective_boxing_scheduler_plan_token_ = Singleton::Get()->AddPlan(plan); -#ifdef WITH_CUDA + // #ifdef WITH_CUDA Singleton::Get()->CreateCommFromPlan(plan); -#endif // WITH_CUDA + // #endif // WITH_CUDA } std::vector source_tasks; source_tasks.reserve(plan.task().size()); diff --git a/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp b/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp index abdb0b596fd..563170f15fd 100644 --- a/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp +++ b/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp @@ -16,7 +16,7 @@ limitations under the License. #include "oneflow/core/auto_parallel/auto_memory.h" #include "oneflow/core/common/util.h" #include "oneflow/core/job/nd_sbp_util.h" -#ifdef WITH_CUDA +// #ifdef WITH_CUDA #include "oneflow/core/framework/framework.h" #include "oneflow/core/framework/nd_sbp.h" #include "oneflow/core/framework/instructions_builder.h" @@ -146,7 +146,8 @@ void FindAllConnectedSubgraphForGpuExecOrder(std::vector> CHECK(visited.insert(seed_node).second); const ParallelDesc& seed_parallel_desc = seed_node->parallel_desc(); // NOTE(chengcheng): ONLY consider GPU op and parallel num > 1. - if (seed_parallel_desc.device_type() != DeviceType::kCUDA) { continue; } + // if (seed_parallel_desc.device_type() != DeviceType::kCUDA) { continue; } + if (seed_parallel_desc.device_type() == DeviceType::kCPU) { continue; } if (seed_parallel_desc.parallel_num() <= 1) { continue; } // NOTE(chengcheng): using fastest time shape for merge acc into bw subgraph. if (!SharedPtrShapeEqual(GetOpNodeFastestTimeShape(seed_node), seed_time_shape)) { continue; } @@ -883,4 +884,4 @@ REGISTER_JOB_PASS("InsertNcclLogicalOpPass", InsertNcclLogicalOpPass); } // namespace oneflow -#endif // WITH_CUDA +// #endif // WITH_CUDA diff --git a/oneflow/core/job_rewriter/job_completer.cpp b/oneflow/core/job_rewriter/job_completer.cpp index dc605d8fb96..c3b638d9f22 100644 --- a/oneflow/core/job_rewriter/job_completer.cpp +++ b/oneflow/core/job_rewriter/job_completer.cpp @@ -153,7 +153,7 @@ Maybe JobCompleter::Complete(Job* job) { compile_tc->Count("[GraphCompile]" + job_name + " SystemOpFillJobNamePass", 1, true); JUST(JobPass4Name("DumpBlobParallelConfPass")(job, &job_pass_ctx)); compile_tc->Count("[GraphCompile]" + job_name + " DumpBlobParallelConfPass", 1, true); -#ifdef WITH_CUDA + // #ifdef WITH_CUDA if (Singleton::Get()->nccl_use_compute_stream()) { // NOTE(chengcheng): this pass need as last pass for insert correct op with nccl boxing. JUST(JobPass4Name("InsertNcclLogicalOpPass")(job, &job_pass_ctx)); @@ -169,7 +169,7 @@ Maybe JobCompleter::Complete(Job* job) { JUST(JobPass4Name("DumpBlobParallelConfPass")(job, &job_pass_ctx)); compile_tc->Count("[GraphCompile]" + job_name + " DumpBlobParallelConfPass", 1, true); } -#endif // WITH_CUDA + // #endif // WITH_CUDA JUST(JobPass4Name("LogicalChainPass")(job, &job_pass_ctx)); JUST(JobPass4Name("DumpBlobParallelConfPass")(job, &job_pass_ctx)); diff --git a/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp b/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp index 5fd82c36a2b..90d5bd81908 100644 --- a/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp +++ b/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp @@ -13,7 +13,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. */ -#ifdef WITH_CUDA +// #ifdef WITH_CUDA #include "oneflow/core/auto_parallel/auto_memory.h" #include "oneflow/core/job/nd_sbp_util.h" #include "oneflow/core/framework/framework.h" @@ -210,4 +210,4 @@ REGISTER_JOB_PASS("NcclLogicalChainStrictOrderPass", NcclLogicalChainStrictOrder } // namespace oneflow -#endif // WITH_CUDA +// #endif // WITH_CUDA diff --git a/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp b/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp index 7556d91d035..96faefdebd4 100644 --- a/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp +++ b/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp @@ -13,7 +13,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. */ -#ifdef WITH_CUDA +// #ifdef WITH_CUDA #include "oneflow/core/auto_parallel/auto_memory.h" #include "oneflow/core/job/nd_sbp_util.h" #include "oneflow/core/framework/framework.h" @@ -293,4 +293,4 @@ REGISTER_JOB_PASS("NcclLogicalOpFusionPass", NcclLogicalOpFusionPass); } // namespace oneflow -#endif // WITH_CUDA +// #endif // WITH_CUDA diff --git a/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp b/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp index 825c43a36a6..a4e4bb886fe 100644 --- a/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp +++ b/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp @@ -22,15 +22,15 @@ limitations under the License. #include "oneflow/core/operator/nccl_send_recv_boxing_op_util.h" #include "oneflow/user/kernels/collective_communication/include/all_to_all.h" -#if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 +// #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 namespace oneflow { -class NcclSendRecvBoxingKernel final : public Kernel { +class CclSendRecvBoxingKernel final : public Kernel { public: - OF_DISALLOW_COPY_AND_MOVE(NcclSendRecvBoxingKernel); - NcclSendRecvBoxingKernel() = default; - ~NcclSendRecvBoxingKernel() override = default; + OF_DISALLOW_COPY_AND_MOVE(CclSendRecvBoxingKernel); + CclSendRecvBoxingKernel() = default; + ~CclSendRecvBoxingKernel() override = default; const std::vector>& in_tensor_slice_copier_vec() const { return in_tensor_slice_copier_vec_; @@ -38,15 +38,15 @@ class NcclSendRecvBoxingKernel final : public Kernel { const std::vector>& out_tensor_slice_copier_vec() const { return out_tensor_slice_copier_vec_; } - const std::vector& send_elem_cnts() const { return send_elem_cnts_; } - const std::vector& recv_elem_cnts() const { return recv_elem_cnts_; } + const std::vector& send_elem_cnts() const { return send_elem_cnts_; } + const std::vector& recv_elem_cnts() const { return recv_elem_cnts_; } const bool has_input() const { return has_input_; } const bool has_output() const { return has_output_; } ccl::CclComm ccl_comm() const { return GetOrCreate().ccl_comm; } private: struct Comm { - Comm(ccl::CclComm comm) : ccl_comm(comm) {} + explicit Comm(ccl::CclComm comm) : ccl_comm(comm) {} ccl::CclComm ccl_comm; }; @@ -72,23 +72,24 @@ class NcclSendRecvBoxingKernel final : public Kernel { bool src_nd_sbp_no_partial_parallel_; std::vector> in_tensor_slice_copier_vec_; std::vector> out_tensor_slice_copier_vec_; - std::vector send_elem_cnts_; - std::vector recv_elem_cnts_; + std::vector send_elem_cnts_; + std::vector recv_elem_cnts_; bool has_input_; bool has_output_; }; -void NcclSendRecvBoxingKernel::ForwardDataContent(KernelContext* ctx) const { +void CclSendRecvBoxingKernel::ForwardDataContent(KernelContext* ctx) const { Blob* buf = ctx->BnInOp2Blob("buf"); ccl::CclComm ccl_comm = this->ccl_comm(); - const std::vector& send_elem_cnts = this->send_elem_cnts(); - const std::vector& recv_elem_cnts = this->recv_elem_cnts(); + const std::vector& send_elem_cnts = this->send_elem_cnts(); + const std::vector& recv_elem_cnts = this->recv_elem_cnts(); const int64_t parallel_num = this->kernel_conf().parallel_ctx().parallel_num(); const DataType data_type = buf->data_type(); + const size_t dtype_size = GetSizeOfDataType(data_type); std::vector send_in_ptr; std::vector recv_out_ptr; - std::vector send_offsets; - std::vector recv_offsets; + std::vector send_offsets; + std::vector recv_offsets; char* buf_ptr = buf->mut_dptr(); uint64_t offset = 0; if (this->has_input()) { @@ -96,7 +97,7 @@ void NcclSendRecvBoxingKernel::ForwardDataContent(KernelContext* ctx) const { void* send_ptr = reinterpret_cast(buf_ptr + offset); send_in_ptr.push_back(send_ptr); send_offsets.push_back(offset); - offset += send_elem_cnts.at(i) * GetSizeOfDataType(data_type); + offset += send_elem_cnts.at(i) * dtype_size; } } const uint64_t recv_offset = offset; @@ -105,7 +106,7 @@ void NcclSendRecvBoxingKernel::ForwardDataContent(KernelContext* ctx) const { void* recv_ptr = reinterpret_cast(buf_ptr + offset); recv_out_ptr.push_back(recv_ptr); recv_offsets.push_back(offset - recv_offset); - offset += recv_elem_cnts.at(i) * GetSizeOfDataType(data_type); + offset += recv_elem_cnts.at(i) * dtype_size; } } if (this->has_input()) { @@ -141,10 +142,10 @@ void NcclSendRecvBoxingKernel::ForwardDataContent(KernelContext* ctx) const { } } } else { - std::unique_ptr primitive = + std::unique_ptr add_primitive = ep::primitive::NewPrimitive(ctx->stream()->device_type(), out->data_type()); - CHECK(primitive); + CHECK(add_primitive); std::unique_ptr memset_primitive = ep::primitive::NewPrimitive(ctx->stream()->device_type()); CHECK(memset_primitive); @@ -156,21 +157,21 @@ void NcclSendRecvBoxingKernel::ForwardDataContent(KernelContext* ctx) const { if (recv_elem_cnts.at(i) != out->shape().elem_cnt()) { // if not same shape, memset out memset_primitive->Launch(ctx->stream(), out->mut_dptr(), 0, - out->shape().elem_cnt() * GetSizeOfDataType(data_type)); + out->shape().elem_cnt() * dtype_size); } out_tensor_slice_copier_vec.at(i)->Copy(ctx->stream(), out->mut_dptr(), recv_out_ptr.at(i)); } else { if (recv_elem_cnts.at(i) == out->shape().elem_cnt()) { - primitive->Launch(ctx->stream(), out->dptr(), recv_out_ptr.at(i), out->mut_dptr(), - out->shape().elem_cnt()); + add_primitive->Launch(ctx->stream(), out->dptr(), recv_out_ptr.at(i), out->mut_dptr(), + out->shape().elem_cnt()); } else { void* out_buf = reinterpret_cast(buf_ptr + offset); memset_primitive->Launch(ctx->stream(), out_buf, 0, - out->shape().elem_cnt() * GetSizeOfDataType(data_type)); + out->shape().elem_cnt() * dtype_size); out_tensor_slice_copier_vec.at(i)->Copy(ctx->stream(), out_buf, recv_out_ptr.at(i)); - primitive->Launch(ctx->stream(), out->dptr(), out_buf, out->mut_dptr(), - out->shape().elem_cnt()); + add_primitive->Launch(ctx->stream(), out->dptr(), out_buf, out->mut_dptr(), + out->shape().elem_cnt()); } } } @@ -178,12 +179,12 @@ void NcclSendRecvBoxingKernel::ForwardDataContent(KernelContext* ctx) const { } } -void NcclSendRecvBoxingKernel::VirtualKernelInit(KernelContext* ctx) { +void CclSendRecvBoxingKernel::VirtualKernelInit(KernelContext* ctx) { const NcclSendRecvBoxingOpConf& conf = this->op_conf().nccl_send_recv_boxing_conf(); if (this->op_conf().has_stream_name_hint()) { stream_name_ = this->op_conf().stream_name_hint(); } else { - stream_name_ = EagerNcclCommMgr::kDefaultStreamName; + stream_name_ = EagerCclCommMgr::kDefaultCclStreamName; } parallel_conf_ = conf.parallel_conf(); const int64_t parallel_id = this->kernel_conf().parallel_ctx().parallel_id(); @@ -246,10 +247,11 @@ void NcclSendRecvBoxingKernel::VirtualKernelInit(KernelContext* ctx) { } } -REGISTER_KERNEL(OperatorConf::kNcclSendRecvBoxingConf, NcclSendRecvBoxingKernel); +// TODO: replace all kNcclxxxConf with kCclxxxConf(for multi devices) +REGISTER_KERNEL(OperatorConf::kNcclSendRecvBoxingConf, CclSendRecvBoxingKernel); -REGISTER_SYSTEM_OP_KERNEL_UNIFIED_NCCL_COMM_INIT(OperatorConf::kNcclSendRecvBoxingConf); +REGISTER_SYSTEM_OP_KERNEL_UNIFIED_CCL_COMM_INIT(OperatorConf::kNcclSendRecvBoxingConf); } // namespace oneflow -#endif // WITH_CUDA && NCCL_VERSION_CODE > 2700 +// #endif // WITH_CUDA && NCCL_VERSION_CODE > 2700 diff --git a/oneflow/user/kernels/eager_nccl_s2s_kernel.cu b/oneflow/user/kernels/eager_nccl_s2s_kernel.cu index 7518efb0c33..9485a83a6b5 100644 --- a/oneflow/user/kernels/eager_nccl_s2s_kernel.cu +++ b/oneflow/user/kernels/eager_nccl_s2s_kernel.cu @@ -23,16 +23,16 @@ limitations under the License. #include "oneflow/core/ep/cuda/cuda_stream.h" #include "oneflow/user/kernels/collective_communication/include/all_to_all.h" -#if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 +// #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 namespace oneflow { namespace { -class EagerNcclOpKernelCache final : public user_op::OpKernelCache { +class EagerCclOpKernelCache final : public user_op::OpKernelCache { public: - explicit EagerNcclOpKernelCache(user_op::KernelCacheContext* ctx) { Init(ctx); } - ~EagerNcclOpKernelCache() override = default; + explicit EagerCclOpKernelCache(user_op::KernelCacheContext* ctx) { Init(ctx); } + ~EagerCclOpKernelCache() override = default; Symbol parallel_desc() const { return parallel_desc_; } const ccl::CclComm& ccl_comm() const { return ccl_comm_; } @@ -51,7 +51,7 @@ class EagerNcclOpKernelCache final : public user_op::OpKernelCache { ccl::CclComm ccl_comm_{}; }; -size_t InferEagerNcclS2SKernelTmpBufferSize(user_op::InferContext* ctx) { +size_t InferEagerCclS2SKernelTmpBufferSize(user_op::InferContext* ctx) { const user_op::TensorDesc& in_tensor = ctx->InputTensorDesc("in", 0); size_t tensor_byte_size = GetCudaAlignedSize(in_tensor.shape().elem_cnt() * GetSizeOfDataType(in_tensor.data_type())); @@ -60,31 +60,31 @@ size_t InferEagerNcclS2SKernelTmpBufferSize(user_op::InferContext* ctx) { return tensor_byte_size * 2; } -void InitEagerNcclOpKernelCache(user_op::KernelCacheContext* ctx, - std::shared_ptr* cache_ptr) { +void InitEagerCclOpKernelCache(user_op::KernelCacheContext* ctx, + std::shared_ptr* cache_ptr) { // NOTE(jianhao): the cache only depends on parallel_conf, and the kernel is singleton // once parallel_conf is determined, so only init the cache at the first time. - if (*cache_ptr == nullptr) { *cache_ptr = std::make_shared(ctx); } + if (*cache_ptr == nullptr) { *cache_ptr = std::make_shared(ctx); } } } // namespace template -class EagerNcclS2SKernel final : public user_op::OpKernel { +class EagerCclS2SKernel final : public user_op::OpKernel { public: - EagerNcclS2SKernel() = default; - ~EagerNcclS2SKernel() override = default; + EagerCclS2SKernel() = default; + ~EagerCclS2SKernel() override = default; void InitOpKernelCacheWithFlags( user_op::KernelCacheContext* ctx, int8_t flag, std::shared_ptr* cache_ptr) const override { - InitEagerNcclOpKernelCache(ctx, cache_ptr); + InitEagerCclOpKernelCache(ctx, cache_ptr); } private: using user_op::OpKernel::Compute; void Compute(user_op::KernelComputeContext* ctx, user_op::OpKernelState*, const user_op::OpKernelCache* cache) const override { - auto* kernel_cache = dynamic_cast(cache); + auto* kernel_cache = dynamic_cast(cache); CHECK(kernel_cache != nullptr); // NOTE(hanbinbin): Compute logic copy from _nccl_logical_s2s const user_op::Tensor* in = ctx->Tensor4ArgNameAndIndex("in", 0); @@ -145,7 +145,7 @@ class EagerNcclS2SKernel final : public user_op::OpKernel { const int64_t elem_per_chunk = elem_cnt / num_ranks; std::unique_ptr all_to_all = ccl::NewCollectiveCommunication( ctx->stream()->device_type(), in->data_type(), in->data_type(), num_ranks); - auto& ccl_comm = kernel_cache->ccl_comm(); + const auto& ccl_comm = kernel_cache->ccl_comm(); all_to_all->Launch(ctx->stream(), const_cast(pack_to_ptr), elem_per_chunk, unpack_from_ptr, elem_per_chunk, ccl_comm); } @@ -174,11 +174,11 @@ class EagerNcclS2SKernel final : public user_op::OpKernel { #define REGISTER_CUDA_EAGER_CCL_S2S_KERNEL(dtype) \ REGISTER_USER_KERNEL("eager_ccl_s2s") \ - .SetCreateFn>() \ + .SetCreateFn>() \ .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCUDA) \ && (user_op::HobDataType("in", 0) == GetDataType::value) \ && (user_op::HobDataType("out", 0) == GetDataType::value)) \ - .SetInferTmpSizeFn(InferEagerNcclS2SKernelTmpBufferSize); + .SetInferTmpSizeFn(InferEagerCclS2SKernelTmpBufferSize); REGISTER_CUDA_EAGER_CCL_S2S_KERNEL(int8_t) REGISTER_CUDA_EAGER_CCL_S2S_KERNEL(int32_t) @@ -189,4 +189,4 @@ REGISTER_CUDA_EAGER_CCL_S2S_KERNEL(double) REGISTER_CUDA_EAGER_CCL_S2S_KERNEL(float16) } // namespace oneflow -#endif // WITH_CUDA && NCCL_VERSION_CODE > 2700 +// #endif // WITH_CUDA && NCCL_VERSION_CODE > 2700 diff --git a/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp b/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp index 03cecc9e427..64e1f1bd3be 100644 --- a/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp +++ b/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp @@ -25,22 +25,49 @@ limitations under the License. #include "oneflow/user/kernels/collective_communication/include/all_gather.h" #include "oneflow/user/kernels/collective_communication/include/all_to_all.h" -#if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 +// #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 namespace oneflow { namespace { -class NcclLogical2DSameDim0KernelCommState : public user_op::OpKernelState { +auto AllReduceCollectiveCommunicationExists() { + return hob::make_custom("AllReduceCollectiveCommunicationExists", + [=](const user_op::KernelRegContext& ctx) { + DeviceType device_type = ctx.device_type(); + return ccl::IsCommunicationContextRegistered(device_type) + && ccl::IsAllReduceRegistered(device_type); + }); +} + +auto AllGatherCollectiveCommunicationExists() { + return hob::make_custom("AllGatherCollectiveCommunicationExists", + [=](const user_op::KernelRegContext& ctx) { + DeviceType device_type = ctx.device_type(); + return ccl::IsCommunicationContextRegistered(device_type) + && ccl::IsAllGatherRegistered(device_type); + }); +} + +auto AllToAllCollectiveCommunicationExists() { + return hob::make_custom("AllToAllCollectiveCommunicationExists", + [=](const user_op::KernelRegContext& ctx) { + DeviceType device_type = ctx.device_type(); + return ccl::IsCommunicationContextRegistered(device_type) + && ccl::IsAllToAllRegistered(device_type); + }); +} + +class CclLogical2DSameDim0KernelCommState : public user_op::OpKernelState { public: - explicit NcclLogical2DSameDim0KernelCommState(user_op::KernelInitContext* ctx) + explicit CclLogical2DSameDim0KernelCommState(user_op::KernelInitContext* ctx) : is_init_(false), - stream_name_(EagerNcclCommMgr::kDefaultStreamName), + stream_name_(EagerCclCommMgr::kDefaultCclStreamName), parallel_desc_(ctx->parallel_desc()), this_parallel_id_(ctx->parallel_ctx().parallel_id()) { if (ctx->op_conf().has_stream_name_hint()) { stream_name_ = ctx->op_conf().stream_name_hint(); } } - ~NcclLogical2DSameDim0KernelCommState() override = default; + ~CclLogical2DSameDim0KernelCommState() override = default; const ccl::CclComm& ccl_comm() { if (!is_init_) { Init(); } @@ -75,12 +102,12 @@ class NcclLogical2DSameDim0KernelCommState : public user_op::OpKernelState { ccl::CclComm ccl_comm_{}; }; -class NcclLogical2DSameDim0AllGatherNoncontinuousKernelState - : public NcclLogical2DSameDim0KernelCommState { +class CclLogical2DSameDim0AllGatherNoncontinuousKernelState + : public CclLogical2DSameDim0KernelCommState { public: - explicit NcclLogical2DSameDim0AllGatherNoncontinuousKernelState(user_op::KernelInitContext* ctx) - : NcclLogical2DSameDim0KernelCommState(ctx), src_split_axis_(-1) {} - ~NcclLogical2DSameDim0AllGatherNoncontinuousKernelState() override = default; + explicit CclLogical2DSameDim0AllGatherNoncontinuousKernelState(user_op::KernelInitContext* ctx) + : CclLogical2DSameDim0KernelCommState(ctx), src_split_axis_(-1) {} + ~CclLogical2DSameDim0AllGatherNoncontinuousKernelState() override = default; int64_t src_split_axis() const { return src_split_axis_; } void set_src_split_axis(int64_t split_axis) { src_split_axis_ = split_axis; } @@ -89,11 +116,11 @@ class NcclLogical2DSameDim0AllGatherNoncontinuousKernelState int64_t src_split_axis_; }; -class NcclLogical2DSameDim0All2AllKernelState : public NcclLogical2DSameDim0KernelCommState { +class CclLogical2DSameDim0All2AllKernelState : public CclLogical2DSameDim0KernelCommState { public: - explicit NcclLogical2DSameDim0All2AllKernelState(user_op::KernelInitContext* ctx) - : NcclLogical2DSameDim0KernelCommState(ctx), src_split_axis_(-1), dst_split_axis_(-1) {} - ~NcclLogical2DSameDim0All2AllKernelState() override = default; + explicit CclLogical2DSameDim0All2AllKernelState(user_op::KernelInitContext* ctx) + : CclLogical2DSameDim0KernelCommState(ctx), src_split_axis_(-1), dst_split_axis_(-1) {} + ~CclLogical2DSameDim0All2AllKernelState() override = default; int64_t src_split_axis() const { return src_split_axis_; } void set_src_split_axis(int64_t split_axis) { src_split_axis_ = split_axis; } @@ -105,20 +132,20 @@ class NcclLogical2DSameDim0All2AllKernelState : public NcclLogical2DSameDim0Kern int64_t dst_split_axis_; }; -class NcclLogical2DSameDim0AllReduce final : public user_op::OpKernel { +class CclLogical2DSameDim0AllReduce final : public user_op::OpKernel { public: - NcclLogical2DSameDim0AllReduce() = default; - ~NcclLogical2DSameDim0AllReduce() override = default; + CclLogical2DSameDim0AllReduce() = default; + ~CclLogical2DSameDim0AllReduce() override = default; std::shared_ptr CreateOpKernelState( user_op::KernelInitContext* ctx) const override { - return std::make_shared(ctx); + return std::make_shared(ctx); } private: void Compute(user_op::KernelComputeContext* ctx, user_op::OpKernelState* state, const user_op::OpKernelCache*) const override { - auto* comm_state = dynamic_cast(state); + auto* comm_state = dynamic_cast(state); CHECK(comm_state != nullptr); const user_op::Tensor* in = ctx->Tensor4ArgNameAndIndex("in", 0); user_op::Tensor* out = ctx->Tensor4ArgNameAndIndex("out", 0); @@ -142,20 +169,20 @@ class NcclLogical2DSameDim0AllReduce final : public user_op::OpKernel { } }; -class NcclLogical2DSameDim0AllGather final : public user_op::OpKernel { +class CclLogical2DSameDim0AllGather final : public user_op::OpKernel { public: - NcclLogical2DSameDim0AllGather() = default; - ~NcclLogical2DSameDim0AllGather() override = default; + CclLogical2DSameDim0AllGather() = default; + ~CclLogical2DSameDim0AllGather() override = default; std::shared_ptr CreateOpKernelState( user_op::KernelInitContext* ctx) const override { - return std::make_shared(ctx); + return std::make_shared(ctx); } private: void Compute(user_op::KernelComputeContext* ctx, user_op::OpKernelState* state, const user_op::OpKernelCache*) const override { - auto* comm_state = dynamic_cast(state); + auto* comm_state = dynamic_cast(state); CHECK(comm_state != nullptr); const user_op::Tensor* in = ctx->Tensor4ArgNameAndIndex("in", 0); user_op::Tensor* out = ctx->Tensor4ArgNameAndIndex("out", 0); @@ -180,14 +207,14 @@ class NcclLogical2DSameDim0AllGather final : public user_op::OpKernel { }; template -class NcclLogical2DSameDim0AllGatherNoncontinuous final : public user_op::OpKernel { +class CclLogical2DSameDim0AllGatherNoncontinuous final : public user_op::OpKernel { public: - NcclLogical2DSameDim0AllGatherNoncontinuous() = default; - ~NcclLogical2DSameDim0AllGatherNoncontinuous() override = default; + CclLogical2DSameDim0AllGatherNoncontinuous() = default; + ~CclLogical2DSameDim0AllGatherNoncontinuous() override = default; std::shared_ptr CreateOpKernelState( user_op::KernelInitContext* ctx) const override { - auto state = std::make_shared(ctx); + auto state = std::make_shared(ctx); NdSbp src_nd_sbp; CHECK_JUST(GetNcclLogicalNdSbpFromAttr(ctx, "src_reduced_nd_sbp", &src_nd_sbp)); CHECK_EQ(src_nd_sbp.sbp_parallel_size(), 2); @@ -200,7 +227,7 @@ class NcclLogical2DSameDim0AllGatherNoncontinuous final : public user_op::OpKern void Compute(user_op::KernelComputeContext* ctx, user_op::OpKernelState* state, const user_op::OpKernelCache*) const override { auto* kernel_state = - dynamic_cast(state); + dynamic_cast(state); CHECK_NOTNULL(kernel_state); const user_op::Tensor* in = ctx->Tensor4ArgNameAndIndex("in", 0); user_op::Tensor* out = ctx->Tensor4ArgNameAndIndex("out", 0); @@ -261,14 +288,14 @@ size_t Infer2DSameDim0AllGatherNoncontinuousKernelTmpBufferSize(user_op::InferCo } template -class NcclLogical2DSameDim0All2All final : public user_op::OpKernel { +class CclLogical2DSameDim0All2All final : public user_op::OpKernel { public: - NcclLogical2DSameDim0All2All() = default; - ~NcclLogical2DSameDim0All2All() override = default; + CclLogical2DSameDim0All2All() = default; + ~CclLogical2DSameDim0All2All() override = default; std::shared_ptr CreateOpKernelState( user_op::KernelInitContext* ctx) const override { - auto state = std::make_shared(ctx); + auto state = std::make_shared(ctx); NdSbp src_nd_sbp; NdSbp dst_nd_sbp; CHECK_JUST(GetNcclLogicalNdSbpFromAttr(ctx, "src_reduced_nd_sbp", &src_nd_sbp)); @@ -285,7 +312,7 @@ class NcclLogical2DSameDim0All2All final : public user_op::OpKernel { private: void Compute(user_op::KernelComputeContext* ctx, user_op::OpKernelState* state, const user_op::OpKernelCache*) const override { - auto* kernel_state = dynamic_cast(state); + auto* kernel_state = dynamic_cast(state); CHECK_NOTNULL(kernel_state); const user_op::Tensor* in = ctx->Tensor4ArgNameAndIndex("in", 0); user_op::Tensor* out = ctx->Tensor4ArgNameAndIndex("out", 0); @@ -396,16 +423,16 @@ size_t Infer2DSameDim0All2AllKernelTmpBufferSize(user_op::InferContext* ctx) { return ret; } -class NcclLogical2DSameDim1KernelCommState final : public user_op::OpKernelState { +class CclLogical2DSameDim1KernelCommState final : public user_op::OpKernelState { public: - explicit NcclLogical2DSameDim1KernelCommState(user_op::KernelInitContext* ctx) + explicit CclLogical2DSameDim1KernelCommState(user_op::KernelInitContext* ctx) : is_init_(false), - stream_name_(EagerNcclCommMgr::kDefaultStreamName), + stream_name_(EagerCclCommMgr::kDefaultCclStreamName), parallel_desc_(ctx->parallel_desc()), this_parallel_id_(ctx->parallel_ctx().parallel_id()) { if (ctx->op_conf().has_stream_name_hint()) { stream_name_ = ctx->op_conf().stream_name_hint(); } } - ~NcclLogical2DSameDim1KernelCommState() = default; + ~CclLogical2DSameDim1KernelCommState() = default; const ccl::CclComm& ccl_comm() { if (!is_init_) { @@ -429,20 +456,20 @@ class NcclLogical2DSameDim1KernelCommState final : public user_op::OpKernelState ccl::CclComm ccl_comm_{}; }; -class NcclLogical2DSameDim1AllReduce final : public user_op::OpKernel { +class CclLogical2DSameDim1AllReduce final : public user_op::OpKernel { public: - NcclLogical2DSameDim1AllReduce() = default; - ~NcclLogical2DSameDim1AllReduce() override = default; + CclLogical2DSameDim1AllReduce() = default; + ~CclLogical2DSameDim1AllReduce() override = default; std::shared_ptr CreateOpKernelState( user_op::KernelInitContext* ctx) const override { - return std::make_shared(ctx); + return std::make_shared(ctx); } private: void Compute(user_op::KernelComputeContext* ctx, user_op::OpKernelState* state, const user_op::OpKernelCache*) const override { - auto* comm_state = dynamic_cast(state); + auto* comm_state = dynamic_cast(state); CHECK(comm_state != nullptr); const user_op::Tensor* in = ctx->Tensor4ArgNameAndIndex("in", 0); user_op::Tensor* out = ctx->Tensor4ArgNameAndIndex("out", 0); @@ -470,17 +497,17 @@ class NcclLogical2DSameDim1AllReduce final : public user_op::OpKernel { } // namespace REGISTER_USER_KERNEL("_nccl_logical_2D_same_dim0_all_reduce") - .SetCreateFn() - .SetIsMatchedHob(user_op::HobDeviceType() == DeviceType::kCUDA); + .SetCreateFn() + .SetIsMatchedHob(AllReduceCollectiveCommunicationExists()); REGISTER_USER_KERNEL("_nccl_logical_2D_same_dim0_all_gather") - .SetCreateFn() - .SetIsMatchedHob(user_op::HobDeviceType() == DeviceType::kCUDA); + .SetCreateFn() + .SetIsMatchedHob(AllGatherCollectiveCommunicationExists()); #define REGISTER_2D_SAME_DIM0_ALLGATHER_NONCONTINUOUS_KERNEL(dtype) \ REGISTER_USER_KERNEL("_nccl_logical_2D_same_dim0_all_gather_noncontinuous") \ - .SetCreateFn>() \ - .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCUDA) \ + .SetCreateFn>() \ + .SetIsMatchedHob(AllGatherCollectiveCommunicationExists() \ && (user_op::HobDataType("in", 0) == GetDataType::value) \ && (user_op::HobDataType("out", 0) == GetDataType::value)) \ .SetInferTmpSizeFn(Infer2DSameDim0AllGatherNoncontinuousKernelTmpBufferSize); @@ -498,8 +525,8 @@ REGISTER_2D_SAME_DIM0_ALLGATHER_NONCONTINUOUS_KERNEL(nv_bfloat16) #define REGISTER_2D_SAME_DIM0_ALL2ALL_KERNEL(dtype) \ REGISTER_USER_KERNEL("_nccl_logical_2D_same_dim0_all2all") \ - .SetCreateFn>() \ - .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCUDA) \ + .SetCreateFn>() \ + .SetIsMatchedHob(AllToAllCollectiveCommunicationExists() \ && (user_op::HobDataType("in", 0) == GetDataType::value) \ && (user_op::HobDataType("out", 0) == GetDataType::value)) \ .SetInferTmpSizeFn(Infer2DSameDim0All2AllKernelTmpBufferSize); @@ -516,15 +543,15 @@ REGISTER_2D_SAME_DIM0_ALL2ALL_KERNEL(nv_bfloat16) #endif REGISTER_USER_KERNEL("_nccl_logical_2D_same_dim1_all_reduce") - .SetCreateFn() - .SetIsMatchedHob(user_op::HobDeviceType() == DeviceType::kCUDA); + .SetCreateFn() + .SetIsMatchedHob(AllReduceCollectiveCommunicationExists()); -REGISTER_USER_KERNEL_UNIFIED_NCCL_COMM_INIT("_nccl_logical_2D_same_dim0_all_reduce"); -REGISTER_USER_KERNEL_UNIFIED_NCCL_COMM_INIT("_nccl_logical_2D_same_dim0_all_gather"); -REGISTER_USER_KERNEL_UNIFIED_NCCL_COMM_INIT("_nccl_logical_2D_same_dim0_all_gather_noncontinuous"); -REGISTER_USER_KERNEL_UNIFIED_NCCL_COMM_INIT("_nccl_logical_2D_same_dim0_all2all"); -REGISTER_USER_KERNEL_UNIFIED_NCCL_COMM_INIT("_nccl_logical_2D_same_dim1_all_reduce"); +REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT("_nccl_logical_2D_same_dim0_all_reduce"); +REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT("_nccl_logical_2D_same_dim0_all_gather"); +REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT("_nccl_logical_2D_same_dim0_all_gather_noncontinuous"); +REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT("_nccl_logical_2D_same_dim0_all2all"); +REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT("_nccl_logical_2D_same_dim1_all_reduce"); } // namespace oneflow -#endif // WITH_CUDA && NCCL_VERSION_CODE > 2700 +// #endif // WITH_CUDA && NCCL_VERSION_CODE > 2700 diff --git a/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp b/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp index 4efe792d1a8..93c72a14c67 100644 --- a/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp +++ b/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp @@ -28,7 +28,7 @@ limitations under the License. #include "collective_communication/include/all_to_all.h" #include "collective_communication/include/reduce_scatter.h" -#if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 +// #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 namespace oneflow { @@ -74,11 +74,11 @@ size_t GetTensorByteSize(const user_op::Tensor& tensor) { return GetCudaAlignedSize(tensor.shape_view().elem_cnt() * GetSizeOfDataType(tensor.data_type())); } -class NcclLogicalFusionKernelState : public user_op::OpKernelState { +class CclLogicalFusionKernelState : public user_op::OpKernelState { public: - explicit NcclLogicalFusionKernelState(user_op::KernelInitContext* ctx) + explicit CclLogicalFusionKernelState(user_op::KernelInitContext* ctx) : is_init_(false), - stream_name_(EagerNcclCommMgr::kDefaultStreamName), + stream_name_(EagerCclCommMgr::kDefaultCclStreamName), parallel_desc_(ctx->parallel_desc()), this_parallel_id_(ctx->parallel_ctx().parallel_id()), num_ranks_(-1), @@ -87,7 +87,7 @@ class NcclLogicalFusionKernelState : public user_op::OpKernelState { if (ctx->op_conf().has_stream_name_hint()) { stream_name_ = ctx->op_conf().stream_name_hint(); } InitSplitAxisAndTmpBufferOffset(ctx); } - ~NcclLogicalFusionKernelState() override = default; + ~CclLogicalFusionKernelState() override = default; ccl::CclComm ccl_comm() { if (!is_init_) { InitComm(); } @@ -258,15 +258,15 @@ class NcclLogicalFusionKernelState : public user_op::OpKernelState { ccl::CclComm ccl_comm_{}; }; -class NcclLogicalFusionKernel final : public user_op::OpKernel { +class CclLogicalFusionKernel final : public user_op::OpKernel { public: - OF_DISALLOW_COPY_AND_MOVE(NcclLogicalFusionKernel); - NcclLogicalFusionKernel() = default; - ~NcclLogicalFusionKernel() override = default; + OF_DISALLOW_COPY_AND_MOVE(CclLogicalFusionKernel); + CclLogicalFusionKernel() = default; + ~CclLogicalFusionKernel() override = default; std::shared_ptr CreateOpKernelState( user_op::KernelInitContext* ctx) const override { - return std::make_shared(ctx); + return std::make_shared(ctx); } private: @@ -276,12 +276,13 @@ class NcclLogicalFusionKernel final : public user_op::OpKernel { bool IsKernelLaunchSynchronized() const override { EagerCclCommMgr* comm_mgr = CHECK_NOTNULL(Singleton::Get()); return comm_mgr->IsAsyncLaunchCclLogicalKernel(); + return true; } }; const void* UpdatePackToPtrByNcclType(const void* pack_to_ptr, const std::string& nccl_type, user_op::Tensor* tmp_buffer, - NcclLogicalFusionKernelState* kernel_state, const int32_t i) { + CclLogicalFusionKernelState* kernel_state, const int32_t i) { CHECK_NOTNULL(tmp_buffer); const void* tmp_dptr = static_cast(tmp_buffer->dptr() + kernel_state->tmp_buffer_offset(i)); @@ -301,7 +302,7 @@ const void* UpdatePackToPtrByNcclType(const void* pack_to_ptr, const std::string void* UpdateUnpackFromPtrByNcclType(void* unpack_from_ptr, const std::string& nccl_type, user_op::Tensor* tmp_buffer, const user_op::Tensor* in, - NcclLogicalFusionKernelState* kernel_state, const int32_t i) { + CclLogicalFusionKernelState* kernel_state, const int32_t i) { CHECK_NOTNULL(tmp_buffer); void* tmp_dptr = static_cast(tmp_buffer->mut_dptr() + kernel_state->tmp_buffer_offset(i)); @@ -327,7 +328,7 @@ void* UpdateUnpackFromPtrByNcclType(void* unpack_from_ptr, const std::string& nc void DoPackBeforeNcclGroup(void* pack_to_ptr, const std::string& nccl_type, const user_op::Tensor* in, user_op::KernelComputeContext* ctx, - NcclLogicalFusionKernelState* kernel_state, const int32_t i) { + CclLogicalFusionKernelState* kernel_state, const int32_t i) { if (nccl_type == "_nccl_logical_reduce_scatter_noncontinuous") { // Do pack before reduce scatter const int64_t num_ranks = kernel_state->num_ranks(); @@ -402,7 +403,7 @@ void DoPackBeforeNcclGroup(void* pack_to_ptr, const std::string& nccl_type, void DoNcclComputeByNcclTypeInGroup(const void* pack_to_ptr, void* unpack_from_ptr, const std::string& nccl_type, const user_op::Tensor* in, user_op::Tensor* out, user_op::KernelComputeContext* ctx, - NcclLogicalFusionKernelState* kernel_state, const int32_t i, + CclLogicalFusionKernelState* kernel_state, const int32_t i, ccl::CclComm ccl_comm) { std::unique_ptr ccl_send = ccl::NewCollectiveCommunication(ctx->stream()->device_type(), in->data_type()); @@ -527,7 +528,7 @@ void DoNcclComputeByNcclTypeInGroup(const void* pack_to_ptr, void* unpack_from_p void DoUnpackAfterNcclGroup(void* unpack_from_ptr, const std::string& nccl_type, const user_op::Tensor* in, user_op::Tensor* out, user_op::KernelComputeContext* ctx, - NcclLogicalFusionKernelState* kernel_state, const int32_t i) { + CclLogicalFusionKernelState* kernel_state, const int32_t i) { const int64_t num_ranks = kernel_state->num_ranks(); const int64_t in_split_axis = kernel_state->src_split_axis(i); const int64_t out_split_axis = kernel_state->dst_split_axis(i); @@ -611,10 +612,10 @@ void DoUnpackAfterNcclGroup(void* unpack_from_ptr, const std::string& nccl_type, } } -void NcclLogicalFusionKernel::Compute(user_op::KernelComputeContext* ctx, - user_op::OpKernelState* state, - const user_op::OpKernelCache*) const { - auto* kernel_state = dynamic_cast(state); +void CclLogicalFusionKernel::Compute(user_op::KernelComputeContext* ctx, + user_op::OpKernelState* state, + const user_op::OpKernelCache*) const { + auto* kernel_state = dynamic_cast(state); CHECK_NOTNULL(kernel_state); const int32_t nccl_num = kernel_state->nccl_num(); const std::vector& nccl_type_list = @@ -655,14 +656,15 @@ void NcclLogicalFusionKernel::Compute(user_op::KernelComputeContext* ctx, ccl::CclComm ccl_comm = kernel_state->ccl_comm(); // do nccl compute in group - OF_NCCL_CHECK(ncclGroupStart()); + // TODO:(zhaoluyang) replacre ncclGroupStart/ncclGroupEnd with ccl CclGroupStart/CclGroupEnd + // OF_NCCL_CHECK(ncclGroupStart()); for (int32_t i = 0; i < nccl_num; ++i) { const user_op::Tensor* in = ctx->Tensor4ArgNameAndIndex("in", i); user_op::Tensor* out = ctx->Tensor4ArgNameAndIndex("out", i); DoNcclComputeByNcclTypeInGroup(pack_to_ptr_list.at(i), unpack_from_ptr_list.at(i), nccl_type_list.at(i), in, out, ctx, kernel_state, i, ccl_comm); } - OF_NCCL_CHECK(ncclGroupEnd()); + // OF_NCCL_CHECK(ncclGroupEnd()); // try to do unpack after all nccl for (int32_t i = 0; i < nccl_num; ++i) { @@ -699,11 +701,14 @@ size_t InferNcclLogicalFusionKernelTmpBufferSize(user_op::InferContext* ctx) { } REGISTER_USER_KERNEL("_nccl_logical_fusion") - .SetCreateFn() - .SetIsMatchedHob(user_op::HobDeviceType() == DeviceType::kCUDA) + .SetCreateFn() + .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCUDA) + || (user_op::HobDeviceType() == DeviceType::kNPU)) .SetInferTmpSizeFn(InferNcclLogicalFusionKernelTmpBufferSize); + +// TODO: SetIsMatchedHob support multi devices(not including cpu) } // namespace } // namespace oneflow -#endif // WITH_CUDA && NCCL_VERSION_CODE > 2700 +// #endif // WITH_CUDA && NCCL_VERSION_CODE > 2700 diff --git a/oneflow/user/kernels/nccl_logical_kernels.cpp b/oneflow/user/kernels/nccl_logical_kernels.cpp index f965ac6c418..b4e519744ec 100644 --- a/oneflow/user/kernels/nccl_logical_kernels.cpp +++ b/oneflow/user/kernels/nccl_logical_kernels.cpp @@ -25,18 +25,74 @@ limitations under the License. #include "oneflow/user/kernels/collective_communication/include/all_reduce.h" #include "oneflow/user/kernels/collective_communication/include/all_gather.h" #include "oneflow/user/kernels/collective_communication/include/reduce_scatter.h" +#include "oneflow/user/kernels/collective_communication/include/broadcast.h" +#include "oneflow/user/kernels/collective_communication/include/reduce.h" -#if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 +// #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 namespace oneflow { namespace { +auto AllReduceCollectiveCommunicationExists() { + return hob::make_custom("AllReduceCollectiveCommunicationExists", + [=](const user_op::KernelRegContext& ctx) { + DeviceType device_type = ctx.device_type(); + return ccl::IsCommunicationContextRegistered(device_type) + && ccl::IsAllReduceRegistered(device_type); + }); +} + +auto ReduceScatterCollectiveCommunicationExists() { + return hob::make_custom("ReduceScatterCollectiveCommunicationExists", + [=](const user_op::KernelRegContext& ctx) { + DeviceType device_type = ctx.device_type(); + return ccl::IsCommunicationContextRegistered(device_type) + && ccl::IsReduceScatterRegistered(device_type); + }); +} + +auto AllGatherCollectiveCommunicationExists() { + return hob::make_custom("AllGatherCollectiveCommunicationExists", + [=](const user_op::KernelRegContext& ctx) { + DeviceType device_type = ctx.device_type(); + return ccl::IsCommunicationContextRegistered(device_type) + && ccl::IsAllGatherRegistered(device_type); + }); +} + +auto ReduceCollectiveCommunicationExists() { + return hob::make_custom("ReduceCollectiveCommunicationExists", + [=](const user_op::KernelRegContext& ctx) { + DeviceType device_type = ctx.device_type(); + return ccl::IsCommunicationContextRegistered(device_type) + && ccl::IsReduceRegistered(device_type); + }); +} + +auto BroadcastCollectiveCommunicationExists() { + return hob::make_custom("BroadcastCollectiveCommunicationExists", + [=](const user_op::KernelRegContext& ctx) { + DeviceType device_type = ctx.device_type(); + return ccl::IsCommunicationContextRegistered(device_type) + && ccl::IsBroadcastRegistered(device_type); + }); +} + +auto AllToAllCollectiveCommunicationExists() { + return hob::make_custom("AllToAllCollectiveCommunicationExists", + [=](const user_op::KernelRegContext& ctx) { + DeviceType device_type = ctx.device_type(); + return ccl::IsCommunicationContextRegistered(device_type) + && ccl::IsAllToAllRegistered(device_type); + }); +} + class NcclLogicalKernelCommState : public user_op::OpKernelState { public: explicit NcclLogicalKernelCommState(user_op::KernelInitContext* ctx) : is_init_(false), - stream_name_(EagerNcclCommMgr::kDefaultStreamName), + stream_name_(EagerCclCommMgr::kDefaultCclStreamName), parallel_desc_(ctx->parallel_desc()) { if (ctx->op_conf().has_stream_name_hint()) { stream_name_ = ctx->op_conf().stream_name_hint(); } } @@ -102,10 +158,10 @@ class NcclLogicalS2SKernelState : public NcclLogicalKernelCommState { int64_t dst_split_axis_; }; -class NcclLogicalAllReduceKernel final : public user_op::OpKernel { +class CclLogicalAllReduceKernel final : public user_op::OpKernel { public: - NcclLogicalAllReduceKernel() = default; - ~NcclLogicalAllReduceKernel() override = default; + CclLogicalAllReduceKernel() = default; + ~CclLogicalAllReduceKernel() override = default; std::shared_ptr CreateOpKernelState( user_op::KernelInitContext* ctx) const override { @@ -140,10 +196,10 @@ class NcclLogicalAllReduceKernel final : public user_op::OpKernel { } }; -class NcclLogicalReduceScatterKernel final : public user_op::OpKernel { +class CclLogicalReduceScatterKernel final : public user_op::OpKernel { public: - NcclLogicalReduceScatterKernel() = default; - ~NcclLogicalReduceScatterKernel() override = default; + CclLogicalReduceScatterKernel() = default; + ~CclLogicalReduceScatterKernel() override = default; std::shared_ptr CreateOpKernelState( user_op::KernelInitContext* ctx) const override { @@ -179,10 +235,10 @@ class NcclLogicalReduceScatterKernel final : public user_op::OpKernel { } }; -class NcclLogicalAllGatherKernel final : public user_op::OpKernel { +class CclLogicalAllGatherKernel final : public user_op::OpKernel { public: - NcclLogicalAllGatherKernel() = default; - ~NcclLogicalAllGatherKernel() override = default; + CclLogicalAllGatherKernel() = default; + ~CclLogicalAllGatherKernel() override = default; std::shared_ptr CreateOpKernelState( user_op::KernelInitContext* ctx) const override { @@ -217,10 +273,10 @@ class NcclLogicalAllGatherKernel final : public user_op::OpKernel { }; template -class NcclLogicalAllGatherNoncontinuous final : public user_op::OpKernel { +class CclLogicalAllGatherNoncontinuous final : public user_op::OpKernel { public: - NcclLogicalAllGatherNoncontinuous() = default; - ~NcclLogicalAllGatherNoncontinuous() override = default; + CclLogicalAllGatherNoncontinuous() = default; + ~CclLogicalAllGatherNoncontinuous() override = default; std::shared_ptr CreateOpKernelState( user_op::KernelInitContext* ctx) const override { @@ -295,10 +351,10 @@ size_t InferAllGatherNoncontinuousKernelTmpBufferSize(user_op::InferContext* ctx } template -class NcclLogicalReduceScatterNoncontinuous final : public user_op::OpKernel { +class CclLogicalReduceScatterNoncontinuous final : public user_op::OpKernel { public: - NcclLogicalReduceScatterNoncontinuous() = default; - ~NcclLogicalReduceScatterNoncontinuous() override = default; + CclLogicalReduceScatterNoncontinuous() = default; + ~CclLogicalReduceScatterNoncontinuous() override = default; std::shared_ptr CreateOpKernelState( user_op::KernelInitContext* ctx) const override { @@ -370,10 +426,10 @@ size_t InferReduceScatterNoncontinuousKernelTmpBufferSize(user_op::InferContext* } template -class NcclLogicalS2SKernel final : public user_op::OpKernel { +class CclLogicalS2SKernel final : public user_op::OpKernel { public: - NcclLogicalS2SKernel() = default; - ~NcclLogicalS2SKernel() override = default; + CclLogicalS2SKernel() = default; + ~CclLogicalS2SKernel() override = default; std::shared_ptr CreateOpKernelState( user_op::KernelInitContext* ctx) const override { @@ -508,21 +564,21 @@ size_t InferS2SKernelTmpBufferSize(user_op::InferContext* ctx) { } // namespace REGISTER_USER_KERNEL("_nccl_logical_all_reduce") - .SetCreateFn() - .SetIsMatchedHob(user_op::HobDeviceType() == DeviceType::kCUDA); + .SetCreateFn() + .SetIsMatchedHob(AllReduceCollectiveCommunicationExists()); REGISTER_USER_KERNEL("_nccl_logical_reduce_scatter") - .SetCreateFn() - .SetIsMatchedHob(user_op::HobDeviceType() == DeviceType::kCUDA); + .SetCreateFn() + .SetIsMatchedHob(ReduceScatterCollectiveCommunicationExists()); REGISTER_USER_KERNEL("_nccl_logical_all_gather") - .SetCreateFn() - .SetIsMatchedHob(user_op::HobDeviceType() == DeviceType::kCUDA); + .SetCreateFn() + .SetIsMatchedHob(AllGatherCollectiveCommunicationExists()); #define REGISTER_ALLGATHER_NONCONTINUOUS_KERNEL(dtype) \ REGISTER_USER_KERNEL("_nccl_logical_all_gather_noncontinuous") \ - .SetCreateFn>() \ - .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCUDA) \ + .SetCreateFn>() \ + .SetIsMatchedHob(AllGatherCollectiveCommunicationExists() \ && (user_op::HobDataType("in", 0) == GetDataType::value) \ && (user_op::HobDataType("out", 0) == GetDataType::value)) \ .SetInferTmpSizeFn(InferAllGatherNoncontinuousKernelTmpBufferSize); @@ -540,8 +596,8 @@ REGISTER_ALLGATHER_NONCONTINUOUS_KERNEL(nv_bfloat16) #define REGISTER_REDUCE_SCATTER_NONCONTINUOUS_KERNEL(dtype) \ REGISTER_USER_KERNEL("_nccl_logical_reduce_scatter_noncontinuous") \ - .SetCreateFn>() \ - .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCUDA) \ + .SetCreateFn>() \ + .SetIsMatchedHob(ReduceScatterCollectiveCommunicationExists() \ && (user_op::HobDataType("in", 0) == GetDataType::value) \ && (user_op::HobDataType("out", 0) == GetDataType::value)) \ .SetInferTmpSizeFn(InferReduceScatterNoncontinuousKernelTmpBufferSize); @@ -559,8 +615,8 @@ REGISTER_REDUCE_SCATTER_NONCONTINUOUS_KERNEL(nv_bfloat16) #define REGISTER_S2S_KERNEL(dtype) \ REGISTER_USER_KERNEL("_nccl_logical_s2s") \ - .SetCreateFn>() \ - .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCUDA) \ + .SetCreateFn>() \ + .SetIsMatchedHob(AllToAllCollectiveCommunicationExists() \ && (user_op::HobDataType("in", 0) == GetDataType::value) \ && (user_op::HobDataType("out", 0) == GetDataType::value)) \ .SetInferTmpSizeFn(InferS2SKernelTmpBufferSize); @@ -576,12 +632,12 @@ REGISTER_S2S_KERNEL(float16) REGISTER_S2S_KERNEL(nv_bfloat16) #endif -REGISTER_USER_KERNEL_UNIFIED_NCCL_COMM_INIT("_nccl_logical_all_reduce"); -REGISTER_USER_KERNEL_UNIFIED_NCCL_COMM_INIT("_nccl_logical_reduce_scatter"); -REGISTER_USER_KERNEL_UNIFIED_NCCL_COMM_INIT("_nccl_logical_all_gather"); -REGISTER_USER_KERNEL_UNIFIED_NCCL_COMM_INIT("_nccl_logical_all_gather_noncontinuous"); -REGISTER_USER_KERNEL_UNIFIED_NCCL_COMM_INIT("_nccl_logical_s2s"); +REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT("_nccl_logical_all_reduce"); +REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT("_nccl_logical_reduce_scatter"); +REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT("_nccl_logical_all_gather"); +REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT("_nccl_logical_all_gather_noncontinuous"); +REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT("_nccl_logical_s2s"); } // namespace oneflow -#endif // WITH_CUDA && NCCL_VERSION_CODE > 2700 +// #endif // WITH_CUDA && NCCL_VERSION_CODE > 2700 diff --git a/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp b/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp index fa06e22fd44..88ff9894070 100644 --- a/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp +++ b/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp @@ -30,13 +30,13 @@ limitations under the License. #include "oneflow/core/operator/nccl_send_recv_boxing_op_util.h" #include "oneflow/user/kernels/collective_communication/include/all_to_all.h" -#if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 +// #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 namespace oneflow { -class NcclLogicalSendRecvState final : public user_op::OpKernelState { +class CclLogicalSendRecvState final : public user_op::OpKernelState { public: - explicit NcclLogicalSendRecvState(user_op::KernelInitContext* ctx); + explicit CclLogicalSendRecvState(user_op::KernelInitContext* ctx); const std::vector>& in_tensor_slice_copier_vec() const { return in_tensor_slice_copier_vec_; } @@ -70,8 +70,8 @@ class NcclLogicalSendRecvState final : public user_op::OpKernelState { std::vector recv_elem_cnts_; }; -NcclLogicalSendRecvState::NcclLogicalSendRecvState(user_op::KernelInitContext* ctx) - : stream_name_(EagerNcclCommMgr::kDefaultStreamName) { +CclLogicalSendRecvState::CclLogicalSendRecvState(user_op::KernelInitContext* ctx) + : stream_name_(EagerCclCommMgr::kDefaultCclStreamName) { if (ctx->op_conf().has_stream_name_hint()) { stream_name_ = ctx->op_conf().stream_name_hint(); } const int64_t parallel_id = ctx->parallel_ctx().parallel_id(); parallel_desc_ = std::make_unique(ctx->parallel_desc()); @@ -125,22 +125,22 @@ NcclLogicalSendRecvState::NcclLogicalSendRecvState(user_op::KernelInitContext* c } } -void NcclLogicalSendRecvState::InitComm() const { +void CclLogicalSendRecvState::InitComm() const { EagerCclCommMgr* comm_mgr = CHECK_NOTNULL(Singleton::Get()); ccl::CclComm ccl_comm = comm_mgr->GetCclCommForParallelDescAndStreamName(*parallel_desc_.get(), stream_name_); ccl_comm_.reset(new Comm(ccl_comm)); } -class NcclLogicalSendRecv final : public user_op::OpKernel { +class CclLogicalSendRecv final : public user_op::OpKernel { public: - OF_DISALLOW_COPY_AND_MOVE(NcclLogicalSendRecv); - NcclLogicalSendRecv() = default; - ~NcclLogicalSendRecv() override = default; + OF_DISALLOW_COPY_AND_MOVE(CclLogicalSendRecv); + CclLogicalSendRecv() = default; + ~CclLogicalSendRecv() override = default; std::shared_ptr CreateOpKernelState( user_op::KernelInitContext* ctx) const override { - return std::make_shared(ctx); + return std::make_shared(ctx); } private: @@ -153,9 +153,9 @@ class NcclLogicalSendRecv final : public user_op::OpKernel { } }; -void NcclLogicalSendRecv::Compute(user_op::KernelComputeContext* ctx, user_op::OpKernelState* state, - const user_op::OpKernelCache*) const { - auto* kernel_state = dynamic_cast(state); +void CclLogicalSendRecv::Compute(user_op::KernelComputeContext* ctx, user_op::OpKernelState* state, + const user_op::OpKernelCache*) const { + auto* kernel_state = dynamic_cast(state); CHECK_NOTNULL(kernel_state); const user_op::Tensor* in = ctx->Tensor4ArgNameAndIndex("in", 0); user_op::Tensor* out = ctx->Tensor4ArgNameAndIndex("out", 0); @@ -284,11 +284,13 @@ size_t InferTmpBufferSize(user_op::InferContext* ctx) { return buf_count * GetSizeOfDataType(data_type); } +// TODO:(zhaoluyang) SetIsMatchedHob support multi devices(not including cpu) REGISTER_USER_KERNEL("_nccl_logical_send_recv") - .SetCreateFn() - .SetIsMatchedHob(user_op::HobDeviceType() == DeviceType::kCUDA) + .SetCreateFn() + .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCUDA) + || (user_op::HobDeviceType() == DeviceType::kNPU)) .SetInferTmpSizeFn(InferTmpBufferSize); } // namespace oneflow -#endif // WITH_CUDA && NCCL_VERSION_CODE > 2700 +// #endif // WITH_CUDA && NCCL_VERSION_CODE > 2700 From 6ad645c13ae6ed7a17f44a29ede56ab4d9b53ac3 Mon Sep 17 00:00:00 2001 From: luyang Date: Fri, 14 Mar 2025 06:47:23 +0000 Subject: [PATCH 02/16] support BUILD_NPU --- cmake/caches/ci/cpu.cmake | 1 + cmake/caches/cn/cpu.cmake | 1 + cmake/cuda.cmake | 1 - cmake/oneflow.cmake | 5 +++++ 4 files changed, 7 insertions(+), 1 deletion(-) diff --git a/cmake/caches/ci/cpu.cmake b/cmake/caches/ci/cpu.cmake index d9d40e8c055..95303189483 100644 --- a/cmake/caches/ci/cpu.cmake +++ b/cmake/caches/ci/cpu.cmake @@ -1,4 +1,5 @@ set(BUILD_CUDA NO CACHE BOOL "") +set(BUILD_NPU NO CACHE BOOL "") set(BUILD_GIT_VERSION YES CACHE BOOL "") set(BUILD_TESTING YES CACHE BOOL "") set(WITH_ONEDNN YES CACHE BOOL "") diff --git a/cmake/caches/cn/cpu.cmake b/cmake/caches/cn/cpu.cmake index 6fff5d3493e..b7fae9fece0 100644 --- a/cmake/caches/cn/cpu.cmake +++ b/cmake/caches/cn/cpu.cmake @@ -1,4 +1,5 @@ set(BUILD_CUDA NO CACHE BOOL "") +set(BUILD_NPU NO CACHE BOOL "") set(BUILD_SHARED_LIBS YES CACHE BOOL "") set(THIRD_PARTY_MIRROR aliyun CACHE STRING "") set(PIP_INDEX_MIRROR "https://pypi.tuna.tsinghua.edu.cn/simple" CACHE STRING "") diff --git a/cmake/cuda.cmake b/cmake/cuda.cmake index 1c57a5a6036..ea48f9535e4 100644 --- a/cmake/cuda.cmake +++ b/cmake/cuda.cmake @@ -19,7 +19,6 @@ if(BUILD_CUDA) message(FATAL_ERROR "CUDA_NVCC_GENCODES is deprecated, use CMAKE_CUDA_ARCHITECTURES instead") endif() add_definitions(-DWITH_CUDA) - add_definitions(-DWITH_NPU) # NOTE: For some unknown reason, CUDAToolkit_VERSION may become empty when running cmake again set(CUDA_VERSION ${CUDAToolkit_VERSION} CACHE STRING "") diff --git a/cmake/oneflow.cmake b/cmake/oneflow.cmake index b37535367e1..bfbf87232e8 100644 --- a/cmake/oneflow.cmake +++ b/cmake/oneflow.cmake @@ -353,6 +353,11 @@ if(BUILD_CUDA) PROPERTIES COMPILE_FLAGS "-DCUDA_REAL_ARCHS=\"${CUDA_REAL_ARCHS}\"") endif() +if(BUILD_NPU) + add_definitions(-DWITH_NPU) +endif() +message(STATUS "BUILD_NPU: ${BUILD_NPU}") + if(BUILD_CUDA AND WITH_CUTLASS) if(CUDA_VERSION VERSION_GREATER_EQUAL "10.1") add_definitions(-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1) From eb4318cd496eb5daf6dd97474ceee212f2bb11f9 Mon Sep 17 00:00:00 2001 From: luyang Date: Fri, 14 Mar 2025 06:47:23 +0000 Subject: [PATCH 03/16] support BUILD_NPU --- cmake/caches/ci/cpu.cmake | 1 + cmake/caches/cn/cpu.cmake | 1 + cmake/cuda.cmake | 1 - cmake/oneflow.cmake | 5 +++++ 4 files changed, 7 insertions(+), 1 deletion(-) diff --git a/cmake/caches/ci/cpu.cmake b/cmake/caches/ci/cpu.cmake index d9d40e8c055..95303189483 100644 --- a/cmake/caches/ci/cpu.cmake +++ b/cmake/caches/ci/cpu.cmake @@ -1,4 +1,5 @@ set(BUILD_CUDA NO CACHE BOOL "") +set(BUILD_NPU NO CACHE BOOL "") set(BUILD_GIT_VERSION YES CACHE BOOL "") set(BUILD_TESTING YES CACHE BOOL "") set(WITH_ONEDNN YES CACHE BOOL "") diff --git a/cmake/caches/cn/cpu.cmake b/cmake/caches/cn/cpu.cmake index 6fff5d3493e..b7fae9fece0 100644 --- a/cmake/caches/cn/cpu.cmake +++ b/cmake/caches/cn/cpu.cmake @@ -1,4 +1,5 @@ set(BUILD_CUDA NO CACHE BOOL "") +set(BUILD_NPU NO CACHE BOOL "") set(BUILD_SHARED_LIBS YES CACHE BOOL "") set(THIRD_PARTY_MIRROR aliyun CACHE STRING "") set(PIP_INDEX_MIRROR "https://pypi.tuna.tsinghua.edu.cn/simple" CACHE STRING "") diff --git a/cmake/cuda.cmake b/cmake/cuda.cmake index 1c57a5a6036..ea48f9535e4 100644 --- a/cmake/cuda.cmake +++ b/cmake/cuda.cmake @@ -19,7 +19,6 @@ if(BUILD_CUDA) message(FATAL_ERROR "CUDA_NVCC_GENCODES is deprecated, use CMAKE_CUDA_ARCHITECTURES instead") endif() add_definitions(-DWITH_CUDA) - add_definitions(-DWITH_NPU) # NOTE: For some unknown reason, CUDAToolkit_VERSION may become empty when running cmake again set(CUDA_VERSION ${CUDAToolkit_VERSION} CACHE STRING "") diff --git a/cmake/oneflow.cmake b/cmake/oneflow.cmake index b37535367e1..bfbf87232e8 100644 --- a/cmake/oneflow.cmake +++ b/cmake/oneflow.cmake @@ -353,6 +353,11 @@ if(BUILD_CUDA) PROPERTIES COMPILE_FLAGS "-DCUDA_REAL_ARCHS=\"${CUDA_REAL_ARCHS}\"") endif() +if(BUILD_NPU) + add_definitions(-DWITH_NPU) +endif() +message(STATUS "BUILD_NPU: ${BUILD_NPU}") + if(BUILD_CUDA AND WITH_CUTLASS) if(CUDA_VERSION VERSION_GREATER_EQUAL "10.1") add_definitions(-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1) From b803a22500c6e814ef91b5bc20519a7137848771 Mon Sep 17 00:00:00 2001 From: luyang Date: Mon, 17 Mar 2025 03:47:36 +0000 Subject: [PATCH 04/16] refine micro WITH_NPU --- oneflow/core/job/resource_desc.cpp | 9 ++++++++- .../core/job_rewriter/insert_nccl_logical_op_pass.cpp | 5 ++--- oneflow/core/job_rewriter/job_completer.cpp | 4 ++-- .../nccl_logical_chain_strict_order_pass.cpp | 4 ++-- .../core/job_rewriter/nccl_logical_op_fusion_pass.cpp | 4 ++-- oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp | 4 ++-- oneflow/user/kernels/eager_nccl_s2s_kernel.cu | 4 ++-- oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp | 4 ++-- oneflow/user/kernels/nccl_logical_fusion_kernel.cpp | 4 ++-- oneflow/user/kernels/nccl_logical_kernels.cpp | 4 ++-- oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp | 4 ++-- 11 files changed, 28 insertions(+), 22 deletions(-) diff --git a/oneflow/core/job/resource_desc.cpp b/oneflow/core/job/resource_desc.cpp index 8137530592b..361d922b1ec 100644 --- a/oneflow/core/job/resource_desc.cpp +++ b/oneflow/core/job/resource_desc.cpp @@ -74,7 +74,14 @@ bool ResourceDesc::nccl_use_compute_stream() const { #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 return resource_.nccl_use_compute_stream(); #else - return true; // TODO: find a batter way for multi devices + return false; +#endif + +// TODO: find a batter way for multi devices +#if defined(WITH_NPU) + return true; +#else + return false; #endif } diff --git a/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp b/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp index 563170f15fd..d02f6316092 100644 --- a/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp +++ b/oneflow/core/job_rewriter/insert_nccl_logical_op_pass.cpp @@ -16,7 +16,7 @@ limitations under the License. #include "oneflow/core/auto_parallel/auto_memory.h" #include "oneflow/core/common/util.h" #include "oneflow/core/job/nd_sbp_util.h" -// #ifdef WITH_CUDA +#if defined(WITH_CUDA) || defined(WITH_NPU) #include "oneflow/core/framework/framework.h" #include "oneflow/core/framework/nd_sbp.h" #include "oneflow/core/framework/instructions_builder.h" @@ -146,7 +146,6 @@ void FindAllConnectedSubgraphForGpuExecOrder(std::vector> CHECK(visited.insert(seed_node).second); const ParallelDesc& seed_parallel_desc = seed_node->parallel_desc(); // NOTE(chengcheng): ONLY consider GPU op and parallel num > 1. - // if (seed_parallel_desc.device_type() != DeviceType::kCUDA) { continue; } if (seed_parallel_desc.device_type() == DeviceType::kCPU) { continue; } if (seed_parallel_desc.parallel_num() <= 1) { continue; } // NOTE(chengcheng): using fastest time shape for merge acc into bw subgraph. @@ -884,4 +883,4 @@ REGISTER_JOB_PASS("InsertNcclLogicalOpPass", InsertNcclLogicalOpPass); } // namespace oneflow -// #endif // WITH_CUDA +#endif // WITH_CUDA || WITH_NPU diff --git a/oneflow/core/job_rewriter/job_completer.cpp b/oneflow/core/job_rewriter/job_completer.cpp index c3b638d9f22..ae061f8d83e 100644 --- a/oneflow/core/job_rewriter/job_completer.cpp +++ b/oneflow/core/job_rewriter/job_completer.cpp @@ -153,7 +153,7 @@ Maybe JobCompleter::Complete(Job* job) { compile_tc->Count("[GraphCompile]" + job_name + " SystemOpFillJobNamePass", 1, true); JUST(JobPass4Name("DumpBlobParallelConfPass")(job, &job_pass_ctx)); compile_tc->Count("[GraphCompile]" + job_name + " DumpBlobParallelConfPass", 1, true); - // #ifdef WITH_CUDA +#if defined(WITH_CUDA) || defined(WITH_NPU) if (Singleton::Get()->nccl_use_compute_stream()) { // NOTE(chengcheng): this pass need as last pass for insert correct op with nccl boxing. JUST(JobPass4Name("InsertNcclLogicalOpPass")(job, &job_pass_ctx)); @@ -169,7 +169,7 @@ Maybe JobCompleter::Complete(Job* job) { JUST(JobPass4Name("DumpBlobParallelConfPass")(job, &job_pass_ctx)); compile_tc->Count("[GraphCompile]" + job_name + " DumpBlobParallelConfPass", 1, true); } - // #endif // WITH_CUDA +#endif // WITH_CUDA || WITH_NPU JUST(JobPass4Name("LogicalChainPass")(job, &job_pass_ctx)); JUST(JobPass4Name("DumpBlobParallelConfPass")(job, &job_pass_ctx)); diff --git a/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp b/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp index 90d5bd81908..33392b17a1c 100644 --- a/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp +++ b/oneflow/core/job_rewriter/nccl_logical_chain_strict_order_pass.cpp @@ -13,7 +13,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. */ -// #ifdef WITH_CUDA +#if defined(WITH_CUDA) || defined(WITH_NPU) #include "oneflow/core/auto_parallel/auto_memory.h" #include "oneflow/core/job/nd_sbp_util.h" #include "oneflow/core/framework/framework.h" @@ -210,4 +210,4 @@ REGISTER_JOB_PASS("NcclLogicalChainStrictOrderPass", NcclLogicalChainStrictOrder } // namespace oneflow -// #endif // WITH_CUDA +#endif // WITH_CUDA || WITH_NPU diff --git a/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp b/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp index 96faefdebd4..05e255493d3 100644 --- a/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp +++ b/oneflow/core/job_rewriter/nccl_logical_op_fusion_pass.cpp @@ -13,7 +13,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. */ -// #ifdef WITH_CUDA +#if defined(WITH_CUDA) || defined(WITH_NPU) #include "oneflow/core/auto_parallel/auto_memory.h" #include "oneflow/core/job/nd_sbp_util.h" #include "oneflow/core/framework/framework.h" @@ -293,4 +293,4 @@ REGISTER_JOB_PASS("NcclLogicalOpFusionPass", NcclLogicalOpFusionPass); } // namespace oneflow -// #endif // WITH_CUDA +#endif // WITH_CUDA || WITH_NPU diff --git a/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp b/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp index a4e4bb886fe..af81336ea70 100644 --- a/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp +++ b/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp @@ -22,7 +22,7 @@ limitations under the License. #include "oneflow/core/operator/nccl_send_recv_boxing_op_util.h" #include "oneflow/user/kernels/collective_communication/include/all_to_all.h" -// #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) namespace oneflow { @@ -254,4 +254,4 @@ REGISTER_SYSTEM_OP_KERNEL_UNIFIED_CCL_COMM_INIT(OperatorConf::kNcclSendRecvBoxin } // namespace oneflow -// #endif // WITH_CUDA && NCCL_VERSION_CODE > 2700 +#endif // WITH_CUDA || WITH_NPU diff --git a/oneflow/user/kernels/eager_nccl_s2s_kernel.cu b/oneflow/user/kernels/eager_nccl_s2s_kernel.cu index 9485a83a6b5..4b446445ccc 100644 --- a/oneflow/user/kernels/eager_nccl_s2s_kernel.cu +++ b/oneflow/user/kernels/eager_nccl_s2s_kernel.cu @@ -23,7 +23,7 @@ limitations under the License. #include "oneflow/core/ep/cuda/cuda_stream.h" #include "oneflow/user/kernels/collective_communication/include/all_to_all.h" -// #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) namespace oneflow { @@ -189,4 +189,4 @@ REGISTER_CUDA_EAGER_CCL_S2S_KERNEL(double) REGISTER_CUDA_EAGER_CCL_S2S_KERNEL(float16) } // namespace oneflow -// #endif // WITH_CUDA && NCCL_VERSION_CODE > 2700 +#endif // WITH_CUDA || WITH_NPU diff --git a/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp b/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp index 64e1f1bd3be..dded34c8eb2 100644 --- a/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp +++ b/oneflow/user/kernels/nccl_logical_2d_sbp_kernels.cpp @@ -25,7 +25,7 @@ limitations under the License. #include "oneflow/user/kernels/collective_communication/include/all_gather.h" #include "oneflow/user/kernels/collective_communication/include/all_to_all.h" -// #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) namespace oneflow { @@ -554,4 +554,4 @@ REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT("_nccl_logical_2D_same_dim1_all_reduc } // namespace oneflow -// #endif // WITH_CUDA && NCCL_VERSION_CODE > 2700 +#endif // WITH_CUDA || WITH_NPU diff --git a/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp b/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp index 93c72a14c67..aba3546a3d8 100644 --- a/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp +++ b/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp @@ -28,7 +28,7 @@ limitations under the License. #include "collective_communication/include/all_to_all.h" #include "collective_communication/include/reduce_scatter.h" -// #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) namespace oneflow { @@ -711,4 +711,4 @@ REGISTER_USER_KERNEL("_nccl_logical_fusion") } // namespace oneflow -// #endif // WITH_CUDA && NCCL_VERSION_CODE > 2700 +#endif // WITH_CUDA || WITH_NPU diff --git a/oneflow/user/kernels/nccl_logical_kernels.cpp b/oneflow/user/kernels/nccl_logical_kernels.cpp index b4e519744ec..dde7be74c28 100644 --- a/oneflow/user/kernels/nccl_logical_kernels.cpp +++ b/oneflow/user/kernels/nccl_logical_kernels.cpp @@ -28,7 +28,7 @@ limitations under the License. #include "oneflow/user/kernels/collective_communication/include/broadcast.h" #include "oneflow/user/kernels/collective_communication/include/reduce.h" -// #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) namespace oneflow { @@ -640,4 +640,4 @@ REGISTER_USER_KERNEL_UNIFIED_CCL_COMM_INIT("_nccl_logical_s2s"); } // namespace oneflow -// #endif // WITH_CUDA && NCCL_VERSION_CODE > 2700 +#endif // WITH_CUDA || WITH_NPU diff --git a/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp b/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp index 88ff9894070..dadbd0f9096 100644 --- a/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp +++ b/oneflow/user/kernels/nccl_logical_send_recv_kernel.cpp @@ -30,7 +30,7 @@ limitations under the License. #include "oneflow/core/operator/nccl_send_recv_boxing_op_util.h" #include "oneflow/user/kernels/collective_communication/include/all_to_all.h" -// #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) namespace oneflow { @@ -293,4 +293,4 @@ REGISTER_USER_KERNEL("_nccl_logical_send_recv") } // namespace oneflow -// #endif // WITH_CUDA && NCCL_VERSION_CODE > 2700 +#endif // WITH_CUDA || WITH_NPU From efd7b1c9717c94fcd4ce614ef5379641496b1240 Mon Sep 17 00:00:00 2001 From: luyang Date: Mon, 17 Mar 2025 07:50:45 +0000 Subject: [PATCH 05/16] refine nccl_use_compute_stream --- oneflow/core/job/resource_desc.cpp | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/oneflow/core/job/resource_desc.cpp b/oneflow/core/job/resource_desc.cpp index 361d922b1ec..9a625b088ef 100644 --- a/oneflow/core/job/resource_desc.cpp +++ b/oneflow/core/job/resource_desc.cpp @@ -73,12 +73,7 @@ CollectiveBoxingConf ResourceDesc::collective_boxing_conf() const { bool ResourceDesc::nccl_use_compute_stream() const { #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 return resource_.nccl_use_compute_stream(); -#else - return false; -#endif - -// TODO: find a batter way for multi devices -#if defined(WITH_NPU) +#elif defined(WITH_NPU) // TODO: find a batter way for multi devices return true; #else return false; From 135d7de3f2f7e4e39e3276156a8481b38f20b8b0 Mon Sep 17 00:00:00 2001 From: zly Date: Wed, 26 Mar 2025 13:28:12 +0800 Subject: [PATCH 06/16] fix test case(test_graph_sparse_softmax_cross_entropy.py) --- oneflow/user/kernels/eager_ccl_kernel.cpp | 44 ++++++++++--------- oneflow/user/kernels/eager_nccl_s2s_kernel.cu | 34 +++++++------- 2 files changed, 41 insertions(+), 37 deletions(-) diff --git a/oneflow/user/kernels/eager_ccl_kernel.cpp b/oneflow/user/kernels/eager_ccl_kernel.cpp index 7edddf695f1..bc17a34cd46 100644 --- a/oneflow/user/kernels/eager_ccl_kernel.cpp +++ b/oneflow/user/kernels/eager_ccl_kernel.cpp @@ -308,8 +308,8 @@ class EagerCclTouchKernel final : public user_op::OpKernel { private: using user_op::OpKernel::Compute; void Compute(user_op::KernelComputeContext* ctx, user_op::OpKernelState*, - const user_op::OpKernelCache* cache) const override{ - // Do nothing. + const user_op::OpKernelCache* cache) const override { + // Do nothing. }; bool AlwaysComputeWhenAllOutputsEmpty() const override { return true; } }; @@ -321,10 +321,10 @@ REGISTER_USER_KERNEL("eager_ccl_touch") namespace { -class EagerCclS2SOpKernelCache final : public user_op::OpKernelCache { +class EagerCclS2SCpuOpKernelCache final : public user_op::OpKernelCache { public: - explicit EagerCclS2SOpKernelCache(user_op::KernelCacheContext* ctx) { Init(ctx); } - ~EagerCclS2SOpKernelCache() override = default; + explicit EagerCclS2SCpuOpKernelCache(user_op::KernelCacheContext* ctx) { Init(ctx); } + ~EagerCclS2SCpuOpKernelCache() override = default; Symbol parallel_desc() const { return parallel_desc_; } @@ -339,7 +339,7 @@ class EagerCclS2SOpKernelCache final : public user_op::OpKernelCache { Symbol parallel_desc_; }; -size_t InferEagerCclS2SKernelTmpBufferSize(user_op::InferContext* ctx) { +size_t InferEagerCclS2SCpuKernelTmpBufferSize(user_op::InferContext* ctx) { const user_op::TensorDesc& in_tensor = ctx->InputTensorDesc("in", 0); size_t tensor_byte_size = in_tensor.shape().elem_cnt() * GetSizeOfDataType(in_tensor.data_type()); // NOTE(hanbinbin): Set tmp_buffer_size to twice tensor_byte_size because the @@ -364,24 +364,24 @@ static constexpr auto* GroupP2PPair = DECORATE(&RawGroupP2PPair, ThreadLocal); } // namespace template -class EagerCclS2SKernel final : public user_op::OpKernel { +class EagerCclS2SCPUKernel final : public user_op::OpKernel { public: - EagerCclS2SKernel() = default; - ~EagerCclS2SKernel() override = default; + EagerCclS2SCPUKernel() = default; + ~EagerCclS2SCPUKernel() override = default; void InitOpKernelCacheWithFlags( user_op::KernelCacheContext* ctx, int8_t flag, std::shared_ptr* cache_ptr) const override { // NOTE(jianhao): the cache only depends on parallel_conf, and the kernel is singleton // once parallel_conf is determined, so only init the cache at the first time. - if (*cache_ptr == nullptr) { *cache_ptr = std::make_shared(ctx); } + if (*cache_ptr == nullptr) { *cache_ptr = std::make_shared(ctx); } } private: using user_op::OpKernel::Compute; void Compute(user_op::KernelComputeContext* ctx, user_op::OpKernelState*, const user_op::OpKernelCache* cache) const override { - auto* kernel_cache = dynamic_cast(cache); + auto* kernel_cache = dynamic_cast(cache); CHECK(kernel_cache != nullptr); // NOTE(hanbinbin): Compute logic copy from _nccl_logical_s2s const user_op::Tensor* in = ctx->Tensor4ArgNameAndIndex("in", 0); @@ -490,20 +490,22 @@ class EagerCclS2SKernel final : public user_op::OpKernel { bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; } }; -#define REGISTER_EAGER_CCL_S2S_KERNEL(dtype) \ +#define REGISTER_EAGER_CCL_S2S_CPU_KERNEL(dtype) \ REGISTER_USER_KERNEL("eager_ccl_s2s") \ - .SetCreateFn>() \ + .SetCreateFn>() \ .SetIsMatchedHob(!(user_op::HobDeviceType() == DeviceType::kCUDA) \ && HobIsSendAndRecvRegistered() \ && (user_op::HobDataType("in", 0) == GetDataType::value) \ && (user_op::HobDataType("out", 0) == GetDataType::value)) \ - .SetInferTmpSizeFn(InferEagerCclS2SKernelTmpBufferSize); - -REGISTER_EAGER_CCL_S2S_KERNEL(int8_t) -REGISTER_EAGER_CCL_S2S_KERNEL(int32_t) -REGISTER_EAGER_CCL_S2S_KERNEL(int64_t) -REGISTER_EAGER_CCL_S2S_KERNEL(bool) -REGISTER_EAGER_CCL_S2S_KERNEL(float) -REGISTER_EAGER_CCL_S2S_KERNEL(double) + .SetInferTmpSizeFn(InferEagerCclS2SCpuKernelTmpBufferSize); + +REGISTER_EAGER_CCL_S2S_CPU_KERNEL(int8_t) +REGISTER_EAGER_CCL_S2S_CPU_KERNEL(int32_t) +REGISTER_EAGER_CCL_S2S_CPU_KERNEL(int64_t) +REGISTER_EAGER_CCL_S2S_CPU_KERNEL(bool) +REGISTER_EAGER_CCL_S2S_CPU_KERNEL(float) +REGISTER_EAGER_CCL_S2S_CPU_KERNEL(double) + +#undef REGISTER_EAGER_CCL_S2S_KERNEL } // namespace oneflow diff --git a/oneflow/user/kernels/eager_nccl_s2s_kernel.cu b/oneflow/user/kernels/eager_nccl_s2s_kernel.cu index 4b446445ccc..93070a08020 100644 --- a/oneflow/user/kernels/eager_nccl_s2s_kernel.cu +++ b/oneflow/user/kernels/eager_nccl_s2s_kernel.cu @@ -29,10 +29,10 @@ namespace oneflow { namespace { -class EagerCclOpKernelCache final : public user_op::OpKernelCache { +class EagerCclS2SOpKernelCache final : public user_op::OpKernelCache { public: - explicit EagerCclOpKernelCache(user_op::KernelCacheContext* ctx) { Init(ctx); } - ~EagerCclOpKernelCache() override = default; + explicit EagerCclS2SOpKernelCache(user_op::KernelCacheContext* ctx) { Init(ctx); } + ~EagerCclS2SOpKernelCache() override = default; Symbol parallel_desc() const { return parallel_desc_; } const ccl::CclComm& ccl_comm() const { return ccl_comm_; } @@ -60,11 +60,11 @@ size_t InferEagerCclS2SKernelTmpBufferSize(user_op::InferContext* ctx) { return tensor_byte_size * 2; } -void InitEagerCclOpKernelCache(user_op::KernelCacheContext* ctx, - std::shared_ptr* cache_ptr) { +void InitEagerCclS2SOpKernelCache(user_op::KernelCacheContext* ctx, + std::shared_ptr* cache_ptr) { // NOTE(jianhao): the cache only depends on parallel_conf, and the kernel is singleton // once parallel_conf is determined, so only init the cache at the first time. - if (*cache_ptr == nullptr) { *cache_ptr = std::make_shared(ctx); } + if (*cache_ptr == nullptr) { *cache_ptr = std::make_shared(ctx); } } } // namespace @@ -77,14 +77,14 @@ class EagerCclS2SKernel final : public user_op::OpKernel { void InitOpKernelCacheWithFlags( user_op::KernelCacheContext* ctx, int8_t flag, std::shared_ptr* cache_ptr) const override { - InitEagerCclOpKernelCache(ctx, cache_ptr); + InitEagerCclS2SOpKernelCache(ctx, cache_ptr); } private: using user_op::OpKernel::Compute; void Compute(user_op::KernelComputeContext* ctx, user_op::OpKernelState*, const user_op::OpKernelCache* cache) const override { - auto* kernel_cache = dynamic_cast(cache); + auto* kernel_cache = dynamic_cast(cache); CHECK(kernel_cache != nullptr); // NOTE(hanbinbin): Compute logic copy from _nccl_logical_s2s const user_op::Tensor* in = ctx->Tensor4ArgNameAndIndex("in", 0); @@ -172,7 +172,7 @@ class EagerCclS2SKernel final : public user_op::OpKernel { bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; } }; -#define REGISTER_CUDA_EAGER_CCL_S2S_KERNEL(dtype) \ +#define REGISTER_EAGER_CCL_S2S_KERNEL(dtype) \ REGISTER_USER_KERNEL("eager_ccl_s2s") \ .SetCreateFn>() \ .SetIsMatchedHob((user_op::HobDeviceType() == DeviceType::kCUDA) \ @@ -180,13 +180,15 @@ class EagerCclS2SKernel final : public user_op::OpKernel { && (user_op::HobDataType("out", 0) == GetDataType::value)) \ .SetInferTmpSizeFn(InferEagerCclS2SKernelTmpBufferSize); -REGISTER_CUDA_EAGER_CCL_S2S_KERNEL(int8_t) -REGISTER_CUDA_EAGER_CCL_S2S_KERNEL(int32_t) -REGISTER_CUDA_EAGER_CCL_S2S_KERNEL(int64_t) -REGISTER_CUDA_EAGER_CCL_S2S_KERNEL(bool) -REGISTER_CUDA_EAGER_CCL_S2S_KERNEL(float) -REGISTER_CUDA_EAGER_CCL_S2S_KERNEL(double) -REGISTER_CUDA_EAGER_CCL_S2S_KERNEL(float16) +REGISTER_EAGER_CCL_S2S_KERNEL(int8_t) +REGISTER_EAGER_CCL_S2S_KERNEL(int32_t) +REGISTER_EAGER_CCL_S2S_KERNEL(int64_t) +REGISTER_EAGER_CCL_S2S_KERNEL(bool) +REGISTER_EAGER_CCL_S2S_KERNEL(float) +REGISTER_EAGER_CCL_S2S_KERNEL(double) +REGISTER_EAGER_CCL_S2S_KERNEL(float16) +#undef REGISTER_EAGER_CCL_S2S_KERNEL + } // namespace oneflow #endif // WITH_CUDA || WITH_NPU From 635ef601902119db9c3a99fe7cf48eb9243a000b Mon Sep 17 00:00:00 2001 From: oneflow-ci-bot Date: Wed, 26 Mar 2025 05:29:40 +0000 Subject: [PATCH 07/16] auto format by CI --- oneflow/user/kernels/eager_ccl_kernel.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/oneflow/user/kernels/eager_ccl_kernel.cpp b/oneflow/user/kernels/eager_ccl_kernel.cpp index bc17a34cd46..9c568a75522 100644 --- a/oneflow/user/kernels/eager_ccl_kernel.cpp +++ b/oneflow/user/kernels/eager_ccl_kernel.cpp @@ -308,8 +308,8 @@ class EagerCclTouchKernel final : public user_op::OpKernel { private: using user_op::OpKernel::Compute; void Compute(user_op::KernelComputeContext* ctx, user_op::OpKernelState*, - const user_op::OpKernelCache* cache) const override { - // Do nothing. + const user_op::OpKernelCache* cache) const override{ + // Do nothing. }; bool AlwaysComputeWhenAllOutputsEmpty() const override { return true; } }; From ed30090af391f2cf21f1c088baf6c0625e78c4f6 Mon Sep 17 00:00:00 2001 From: zly Date: Wed, 26 Mar 2025 13:36:06 +0800 Subject: [PATCH 08/16] refine micro WITH_NPU --- oneflow/api/python/flags.cpp | 8 ++++++++ .../core/auto_parallel/boxing_collector.cpp | 8 ++++---- oneflow/core/common/device_type.h | 4 ++++ ...llective_boxing_sub_task_graph_builder.cpp | 2 ++ ...erarchical_sub_task_graph_builder_impl.cpp | 19 ++++++++----------- oneflow/user/kernels/stateful_opkernel.cpp | 2 +- 6 files changed, 27 insertions(+), 16 deletions(-) diff --git a/oneflow/api/python/flags.cpp b/oneflow/api/python/flags.cpp index add729a1da8..1f4f70a3f85 100644 --- a/oneflow/api/python/flags.cpp +++ b/oneflow/api/python/flags.cpp @@ -29,6 +29,14 @@ ONEFLOW_API_PYBIND11_MODULE("flags", m) { #endif // WITH_CUDA }); + m.def("with_npu", []() { +#ifdef WITH_NPU + return true; +#else + return false; +#endif // WITH_NPU + }); + m.def("cuda_version", []() { #ifdef WITH_CUDA return CUDA_VERSION; diff --git a/oneflow/core/auto_parallel/boxing_collector.cpp b/oneflow/core/auto_parallel/boxing_collector.cpp index 28a4805a69b..230a07e2ab3 100644 --- a/oneflow/core/auto_parallel/boxing_collector.cpp +++ b/oneflow/core/auto_parallel/boxing_collector.cpp @@ -581,14 +581,14 @@ Maybe BoxingCollector::AskSbpCombination(const NdSbp& sbp_producer, const return Maybe::Ok(); } -#ifdef WITH_CUDA + #if defined(WITH_CUDA) || defined(WITH_NPU) // Use a general basic communication if no P in the consumer if (((Singleton::Get()->nccl_use_compute_stream() && producer_parallel_desc == consumer_parallel_desc) || enable_general_basic_communication) && (!NdSbpHasPartialParallel(sbp_consumer)) - && producer_parallel_desc.device_type() == DeviceType::kCUDA - && consumer_parallel_desc.device_type() == DeviceType::kCUDA) { + && producer_parallel_desc.device_type() == consumer_parallel_desc.device_type() + && producer_parallel_desc.device_type() != DeviceType::kCPU) { if (NdSbpHasPartialParallel(sbp_producer) && NdSbpHasBroadcastParallel(sbp_consumer)) { // (?, P, ?)->(Si, Sj)->(?, B, ?), two-step transfer // Directly applying general basic communication would have O(n^2) time complexity for P->B @@ -600,7 +600,7 @@ Maybe BoxingCollector::AskSbpCombination(const NdSbp& sbp_producer, const // Otherwise, one-step transfer return Maybe::Ok(); } -#endif // WITH_CUDA +#endif // WITH_CUDA || WITH_NPU if (JUST(ComputeLazyCopyCostBetweenNdSbp(sbp_producer, sbp_consumer, logical_blob_desc, producer_parallel_desc, consumer_parallel_desc, diff --git a/oneflow/core/common/device_type.h b/oneflow/core/common/device_type.h index 10042d4b51f..3e2d8b65648 100644 --- a/oneflow/core/common/device_type.h +++ b/oneflow/core/common/device_type.h @@ -38,6 +38,10 @@ std::string PrintGeneratorAvailableDevices(); #define DEVICE_TYPE_SEQ \ OF_PP_MAKE_TUPLE_SEQ(DeviceType::kCPU) \ OF_PP_MAKE_TUPLE_SEQ(DeviceType::kCUDA) +#elif defined(WITH_NPU) +#define DEVICE_TYPE_SEQ \ + OF_PP_MAKE_TUPLE_SEQ(DeviceType::kCPU) \ + OF_PP_MAKE_TUPLE_SEQ(DeviceType::kNPU) #else #define DEVICE_TYPE_SEQ OF_PP_MAKE_TUPLE_SEQ(DeviceType::kCPU) #endif diff --git a/oneflow/core/graph/boxing/collective_boxing_sub_task_graph_builder.cpp b/oneflow/core/graph/boxing/collective_boxing_sub_task_graph_builder.cpp index 56afc7ac5d9..4e5742460e0 100644 --- a/oneflow/core/graph/boxing/collective_boxing_sub_task_graph_builder.cpp +++ b/oneflow/core/graph/boxing/collective_boxing_sub_task_graph_builder.cpp @@ -35,6 +35,8 @@ CollectiveBoxingSubTskGphBuilder::CollectiveBoxingSubTskGphBuilder() { if (collective_boxing_conf.nccl_enable_all_to_all()) { #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 builders.emplace_back(new CclAll2AllSubTskGphBuilder(DeviceType::kCUDA)); +#elif defined(WITH_NPU) + builders.emplace_back(new CclAll2AllSubTskGphBuilder(DeviceType::kNPU)); #else LOG(WARNING) << "nccl_enable_all_to_all is unavailable unless NCCL_VERSION > 2.7.0"; #endif diff --git a/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp b/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp index c5e7181de8e..4db16f8bf44 100644 --- a/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp +++ b/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp @@ -99,12 +99,9 @@ class NDNcclSendRecvBoxingSubTskGphBuilder final : public HierarchicalSubTskGphB const LogicalBlobId& lbi, const BlobDesc& logical_blob_desc, const NdSbp& in_nd_sbp, const NdSbp& out_nd_sbp, const Shape& time_shape) const override { - if (in_parallel_desc.device_type() != DeviceType::kCPU - && out_parallel_desc.device_type() != DeviceType::kCPU + if (in_parallel_desc.device_type()==out_parallel_desc.device_type() && in_parallel_desc.device_type() != DeviceType::kCPU && !NdSbpHasPartialParallel(out_nd_sbp)) { - // TODO: (zhaoluyang) use WITH_DEVICE or another way to support multi devices(e.g. - // cuda/npu/xpu....) - // #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700 + #if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) ParallelConf merged_parallel_conf; MergeParallelConf(in_parallel_desc.parallel_conf(), out_parallel_desc.parallel_conf(), &merged_parallel_conf); @@ -138,9 +135,9 @@ class NDNcclSendRecvBoxingSubTskGphBuilder final : public HierarchicalSubTskGphB if (has_output) { sorted_out_tasks->push_back(node); } } return BuildSubTskGphBuilderStatus("NDNcclSendRecvBoxingSubTskGphBuilder", ""); - // #else - // return Error::BoxingNotSupportedError() << "No CUDA or low NCCL version"; - // #endif + #else + return Error::BoxingNotSupportedError() << "No Device or low NCCL version"; + #endif } else { return Error::BoxingNotSupportedError() << "Partial SBP in the consumer or not running on CUDA"; @@ -268,9 +265,9 @@ Maybe DispatchHierarchicalSubTskGphBuilder::Build( &reduced_out_nd_sbp, logical_blob_desc.shape()); const auto& in_hierarchy = reduced_in_parallel_desc.hierarchy(); const auto& out_hierarchy = reduced_out_parallel_desc.hierarchy(); - if ((in_hierarchy->NumAxes() > 2 || out_hierarchy->NumAxes() > 2) - && reduced_in_parallel_desc.device_type() != DeviceType::kCPU - && reduced_out_parallel_desc.device_type() != DeviceType::kCPU) { + if ((in_hierarchy->NumAxes() > 2 || out_hierarchy->NumAxes() > 2) + && reduced_in_parallel_desc.device_type() == reduced_out_parallel_desc.device_type() + && reduced_in_parallel_desc.device_type() != DeviceType::kCPU) { return impl_->nd_nccl_send_recv_boxing_sub_tsk_gph_builder_->Build( ctx, sorted_in_tasks, sorted_out_tasks, sorted_ctrl_tasks, reduced_in_parallel_desc, reduced_out_parallel_desc, lbi, logical_blob_desc, reduced_in_nd_sbp, reduced_out_nd_sbp, diff --git a/oneflow/user/kernels/stateful_opkernel.cpp b/oneflow/user/kernels/stateful_opkernel.cpp index 641bbe0b3b1..685a10cbf5f 100644 --- a/oneflow/user/kernels/stateful_opkernel.cpp +++ b/oneflow/user/kernels/stateful_opkernel.cpp @@ -975,7 +975,7 @@ void StatefulOpKernel::Compute(eager::CallContext* call_ctx, ep::Stream* stream, OF_PROFILER_RANGE_GUARD("Compute"); auto er_guard = CHECK_JUST(profiler::EventRecorder::CreateKernelEventRecorder( op_type_name(), -#if defined(WITH_CUDA) +#if defined(WITH_CUDA) || defined(WITH_NPU) [compute_ctx]() -> int64_t { const auto CalMemorySize = [compute_ctx](const one::ArgVec& args) -> int64_t { const auto Func = [compute_ctx](int64_t mem_size, const auto& pair) { From 9a96086a90bfacda61b4b6d05218e84ec5cfff1d Mon Sep 17 00:00:00 2001 From: oneflow-ci-bot Date: Wed, 26 Mar 2025 05:39:49 +0000 Subject: [PATCH 09/16] auto format by CI --- oneflow/core/auto_parallel/boxing_collector.cpp | 2 +- .../hierarchical_sub_task_graph_builder_impl.cpp | 15 ++++++++------- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/oneflow/core/auto_parallel/boxing_collector.cpp b/oneflow/core/auto_parallel/boxing_collector.cpp index 230a07e2ab3..fca8f5f734c 100644 --- a/oneflow/core/auto_parallel/boxing_collector.cpp +++ b/oneflow/core/auto_parallel/boxing_collector.cpp @@ -581,7 +581,7 @@ Maybe BoxingCollector::AskSbpCombination(const NdSbp& sbp_producer, const return Maybe::Ok(); } - #if defined(WITH_CUDA) || defined(WITH_NPU) +#if defined(WITH_CUDA) || defined(WITH_NPU) // Use a general basic communication if no P in the consumer if (((Singleton::Get()->nccl_use_compute_stream() && producer_parallel_desc == consumer_parallel_desc) diff --git a/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp b/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp index 4db16f8bf44..ddd7d879ca5 100644 --- a/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp +++ b/oneflow/core/graph/boxing/hierarchical_sub_task_graph_builder_impl.cpp @@ -99,9 +99,10 @@ class NDNcclSendRecvBoxingSubTskGphBuilder final : public HierarchicalSubTskGphB const LogicalBlobId& lbi, const BlobDesc& logical_blob_desc, const NdSbp& in_nd_sbp, const NdSbp& out_nd_sbp, const Shape& time_shape) const override { - if (in_parallel_desc.device_type()==out_parallel_desc.device_type() && in_parallel_desc.device_type() != DeviceType::kCPU + if (in_parallel_desc.device_type() == out_parallel_desc.device_type() + && in_parallel_desc.device_type() != DeviceType::kCPU && !NdSbpHasPartialParallel(out_nd_sbp)) { - #if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) +#if (defined(WITH_CUDA) && (NCCL_VERSION_CODE > 2700)) || defined(WITH_NPU) ParallelConf merged_parallel_conf; MergeParallelConf(in_parallel_desc.parallel_conf(), out_parallel_desc.parallel_conf(), &merged_parallel_conf); @@ -135,9 +136,9 @@ class NDNcclSendRecvBoxingSubTskGphBuilder final : public HierarchicalSubTskGphB if (has_output) { sorted_out_tasks->push_back(node); } } return BuildSubTskGphBuilderStatus("NDNcclSendRecvBoxingSubTskGphBuilder", ""); - #else - return Error::BoxingNotSupportedError() << "No Device or low NCCL version"; - #endif +#else + return Error::BoxingNotSupportedError() << "No Device or low NCCL version"; +#endif } else { return Error::BoxingNotSupportedError() << "Partial SBP in the consumer or not running on CUDA"; @@ -265,8 +266,8 @@ Maybe DispatchHierarchicalSubTskGphBuilder::Build( &reduced_out_nd_sbp, logical_blob_desc.shape()); const auto& in_hierarchy = reduced_in_parallel_desc.hierarchy(); const auto& out_hierarchy = reduced_out_parallel_desc.hierarchy(); - if ((in_hierarchy->NumAxes() > 2 || out_hierarchy->NumAxes() > 2) - && reduced_in_parallel_desc.device_type() == reduced_out_parallel_desc.device_type() + if ((in_hierarchy->NumAxes() > 2 || out_hierarchy->NumAxes() > 2) + && reduced_in_parallel_desc.device_type() == reduced_out_parallel_desc.device_type() && reduced_in_parallel_desc.device_type() != DeviceType::kCPU) { return impl_->nd_nccl_send_recv_boxing_sub_tsk_gph_builder_->Build( ctx, sorted_in_tasks, sorted_out_tasks, sorted_ctrl_tasks, reduced_in_parallel_desc, From e7c002c333af2afb81a440eab673175ddc355b53 Mon Sep 17 00:00:00 2001 From: zhaoluyang Date: Thu, 27 Mar 2025 12:03:52 +0800 Subject: [PATCH 10/16] revert changes --- oneflow/core/common/device_type.h | 4 ---- oneflow/user/kernels/stateful_opkernel.cpp | 2 +- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/oneflow/core/common/device_type.h b/oneflow/core/common/device_type.h index 3e2d8b65648..10042d4b51f 100644 --- a/oneflow/core/common/device_type.h +++ b/oneflow/core/common/device_type.h @@ -38,10 +38,6 @@ std::string PrintGeneratorAvailableDevices(); #define DEVICE_TYPE_SEQ \ OF_PP_MAKE_TUPLE_SEQ(DeviceType::kCPU) \ OF_PP_MAKE_TUPLE_SEQ(DeviceType::kCUDA) -#elif defined(WITH_NPU) -#define DEVICE_TYPE_SEQ \ - OF_PP_MAKE_TUPLE_SEQ(DeviceType::kCPU) \ - OF_PP_MAKE_TUPLE_SEQ(DeviceType::kNPU) #else #define DEVICE_TYPE_SEQ OF_PP_MAKE_TUPLE_SEQ(DeviceType::kCPU) #endif diff --git a/oneflow/user/kernels/stateful_opkernel.cpp b/oneflow/user/kernels/stateful_opkernel.cpp index 685a10cbf5f..641bbe0b3b1 100644 --- a/oneflow/user/kernels/stateful_opkernel.cpp +++ b/oneflow/user/kernels/stateful_opkernel.cpp @@ -975,7 +975,7 @@ void StatefulOpKernel::Compute(eager::CallContext* call_ctx, ep::Stream* stream, OF_PROFILER_RANGE_GUARD("Compute"); auto er_guard = CHECK_JUST(profiler::EventRecorder::CreateKernelEventRecorder( op_type_name(), -#if defined(WITH_CUDA) || defined(WITH_NPU) +#if defined(WITH_CUDA) [compute_ctx]() -> int64_t { const auto CalMemorySize = [compute_ctx](const one::ArgVec& args) -> int64_t { const auto Func = [compute_ctx](int64_t mem_size, const auto& pair) { From 3fdbf48f7bf1e30db3479d18e10d8d7ea0ff9594 Mon Sep 17 00:00:00 2001 From: zhaoluyang Date: Mon, 7 Apr 2025 14:43:17 +0800 Subject: [PATCH 11/16] debug nll loss --- oneflow/core/autograd/gradient_funcs/nll.cpp | 12 +++--- oneflow/core/functional/functional_api.yaml | 2 +- oneflow/core/functional/impl/nn_functor.cpp | 41 +++++++++++++------ .../core/functional/impl/nn_grad_functor.cpp | 9 ++-- oneflow/ir/include/OneFlow/OneFlowUserOps.td | 9 ++-- oneflow/user/ops/nll_op.cpp | 13 ++++-- 6 files changed, 59 insertions(+), 27 deletions(-) diff --git a/oneflow/core/autograd/gradient_funcs/nll.cpp b/oneflow/core/autograd/gradient_funcs/nll.cpp index 76a946dd4b0..42d34d29df6 100644 --- a/oneflow/core/autograd/gradient_funcs/nll.cpp +++ b/oneflow/core/autograd/gradient_funcs/nll.cpp @@ -24,6 +24,7 @@ namespace one { struct NLLCaptureState : public AutoGradCaptureState { bool requires_grad = false; int64_t ignore_index = -100; + std::string reduction = "none"; }; class NLLGradFunction : public OpExprGradFunction { @@ -53,6 +54,7 @@ Maybe NLLGradFunction::Capture(NLLCaptureState* ctx, const TensorTuple& in ComposedAttrMap composed_attrs(attrs, base_attrs_); ctx->ignore_index = JUST(composed_attrs.GetAttr("ignore_index")); + ctx->reduction = JUST(composed_attrs.GetAttr("reduction")); ctx->SaveTensorForBackward(input); // input ctx->SaveTensorForBackward(JUST(VectorAt(inputs, 1))); // target if (inputs.size() == 3) { @@ -65,7 +67,7 @@ Maybe NLLGradFunction::Apply(const NLLCaptureState* ctx, const TensorTuple TensorTuple* in_grads) const { if (!ctx->requires_grad) { return Maybe::Ok(); } - CHECK_EQ_OR_RETURN(out_grads.size(), 2); // NOLINT(maybe-need-error-msg) + // CHECK_EQ_OR_RETURN(out_grads.size(), 2); // NOLINT(maybe-need-error-msg) CHECK_GE_OR_RETURN(ctx->SavedTensors().size(), 2) << Error::RuntimeError() << "The number of saved tensors is expected to be greater than or equal to 2, but got " @@ -77,13 +79,13 @@ Maybe NLLGradFunction::Apply(const NLLCaptureState* ctx, const TensorTuple in_grads->resize(ctx->SavedTensors().size()); if (ctx->SavedTensors().size() == 2) { - JUST(VectorAt(*in_grads, 0)) = - JUST(functional::NLLGrad(out_grad, input, target, NullOpt, ctx->ignore_index)); + JUST(VectorAt(*in_grads, 0)) = JUST( + functional::NLLGrad(out_grad, input, target, NullOpt, ctx->ignore_index, ctx->reduction)); } else { // has weight auto weight = JUST(VectorAt(ctx->SavedTensors(), 2)); - JUST(VectorAt(*in_grads, 0)) = - JUST(functional::NLLGrad(out_grad, input, target, weight, ctx->ignore_index)); + JUST(VectorAt(*in_grads, 0)) = JUST( + functional::NLLGrad(out_grad, input, target, weight, ctx->ignore_index, ctx->reduction)); } return Maybe::Ok(); diff --git a/oneflow/core/functional/functional_api.yaml b/oneflow/core/functional/functional_api.yaml index 5c6d8148ac1..05d5ab77900 100644 --- a/oneflow/core/functional/functional_api.yaml +++ b/oneflow/core/functional/functional_api.yaml @@ -1355,7 +1355,7 @@ bind_python: True - name: "nll_grad" - signature: "Tensor(Tensor out_grad, Tensor input, Tensor target, Tensor weight=None, Int64 ignore_index) => NLLGrad" + signature: 'Tensor(Tensor out_grad, Tensor input, Tensor target, Tensor weight=None, Int64 ignore_index, String reduction="none") => NLLGrad' bind_python: False - name: "binary_cross_entropy_loss" diff --git a/oneflow/core/functional/impl/nn_functor.cpp b/oneflow/core/functional/impl/nn_functor.cpp index db638fca99d..5722fe45d12 100644 --- a/oneflow/core/functional/impl/nn_functor.cpp +++ b/oneflow/core/functional/impl/nn_functor.cpp @@ -930,9 +930,9 @@ class SkipLayerNormFunctor { std::tuple(has_skip, has_gamma, has_beta, has_bias), op_expr)); } // has_bias - } // has_beta - } // has_gamma - } // has_skip + } // has_beta + } // has_gamma + } // has_skip } Maybe operator()(const std::shared_ptr& x, @@ -1170,8 +1170,8 @@ class SkipRMSNormFunctor { ops_.insert(std::pair, std::shared_ptr>( std::tuple(has_weight, has_skip, has_bias), op_expr)); } // has_bias - } // has_skip - } // has_weight + } // has_skip + } // has_weight } Maybe operator()(const std::shared_ptr& x, @@ -1477,7 +1477,7 @@ class MaxUnpoolNDFunctor { .Input("x") .Input("indices") .Output("y") - .Build())){}; + .Build())) {}; Maybe operator()(const std::shared_ptr& x, const std::shared_ptr& indices, const std::vector& kernel_size, @@ -1819,6 +1819,7 @@ class NLLLossFunctor { .Input("target") .Output("output") .Output("out_weight") + .Output("reduced_out") .Build()); op_weight_ = CHECK_JUST(one::OpBuilder("nll") @@ -1827,6 +1828,7 @@ class NLLLossFunctor { .Input("weight") .Output("output") .Output("out_weight") + .Output("reduced_out") .Build()); } @@ -1875,8 +1877,10 @@ class NLLLossFunctor { target_ = target; } - auto& attrs = THREAD_CACHED_MUTABLE_ATTR_MAP("ignore_index"); - attrs.SetAllAttrs(ignore_index); + // auto& attrs = THREAD_CACHED_MUTABLE_ATTR_MAP("ignore_index"); + // attrs.SetAllAttrs(ignore_index); + auto& attrs = THREAD_CACHED_MUTABLE_ATTR_MAP("ignore_index", "reduction"); + attrs.SetAllAttrs(ignore_index, reduction); std::shared_ptr nll_result; if (weight) { @@ -1890,9 +1894,14 @@ class NLLLossFunctor { if (K > 2) { output = JUST(functional::Reshape(output, *target_shape)); } if (reduction == "none") { return output; } - +#ifdef WITH_NPU + // npu device + if (!input->is_cpu()) { + auto reduced_out = JUST(VectorAt(*nll_result, 2)); + if (reduction == "sum" || reduction == "mean") { return reduced_out; } + } +#endif // WITH_NPU auto sum = JUST(functional::ReduceSum(output, {}, false, NullOpt)); - if (reduction == "sum") { return sum; } auto total_weight = @@ -1915,6 +1924,7 @@ class CrossEntropyFunctor { .Input("target") .Output("output") .Output("out_weight") + .Output("reduced_out") .Build()); op_nll_weight_ = CHECK_JUST(one::OpBuilder("nll") @@ -1923,6 +1933,7 @@ class CrossEntropyFunctor { .Input("weight") .Output("output") .Output("out_weight") + .Output("reduced_out") .Build()); } Maybe operator()(const std::shared_ptr& input, @@ -1959,8 +1970,8 @@ class CrossEntropyFunctor { const auto target_ = JUST(functional::Flatten(target, 0, target->shape()->NumAxes() - 1)); - auto& attrs = THREAD_CACHED_MUTABLE_ATTR_MAP("ignore_index"); - attrs.SetAllAttrs(ignore_index); + auto& attrs = THREAD_CACHED_MUTABLE_ATTR_MAP("ignore_index", "reduction"); + attrs.SetAllAttrs(ignore_index, reduction); std::shared_ptr nll_result; if (weight) { @@ -1974,6 +1985,12 @@ class CrossEntropyFunctor { output = JUST(functional::Reshape(output, *target_shape)); if (reduction == "none") { return output; } +#ifdef WITH_NPU + if (!input->is_cpu()) { + auto reduced_out = JUST(VectorAt(*nll_result, 2)); + if (reduction == "sum" || reduction == "mean") { return reduced_out; } + } +#endif // WITH_NPU auto sum = JUST(functional::ReduceSum(output, {}, false, NullOpt)); if (reduction == "sum") { return sum; } diff --git a/oneflow/core/functional/impl/nn_grad_functor.cpp b/oneflow/core/functional/impl/nn_grad_functor.cpp index e833c98d2ae..9e1d2eea809 100644 --- a/oneflow/core/functional/impl/nn_grad_functor.cpp +++ b/oneflow/core/functional/impl/nn_grad_functor.cpp @@ -468,9 +468,12 @@ class NLLGradFunctor { Maybe operator()(const std::shared_ptr& out_grad, const std::shared_ptr& input, const std::shared_ptr& target, - const Optional& weight, const int64_t ignore_index) const { - auto& attrs = THREAD_CACHED_MUTABLE_ATTR_MAP("ignore_index"); - attrs.SetAllAttrs(ignore_index); + const Optional& weight, const int64_t ignore_index, + const std::string& reduction) const { + // auto& attrs = THREAD_CACHED_MUTABLE_ATTR_MAP("ignore_index"); + // attrs.SetAllAttrs(ignore_index); + auto& attrs = THREAD_CACHED_MUTABLE_ATTR_MAP("ignore_index", "reduction"); + attrs.SetAllAttrs(ignore_index, reduction); if (weight) { return OpInterpUtil::Dispatch( diff --git a/oneflow/ir/include/OneFlow/OneFlowUserOps.td b/oneflow/ir/include/OneFlow/OneFlowUserOps.td index 7532cfe441e..1162e4fe58b 100644 --- a/oneflow/ir/include/OneFlow/OneFlowUserOps.td +++ b/oneflow/ir/include/OneFlow/OneFlowUserOps.td @@ -6106,10 +6106,12 @@ def OneFlow_NLLOp : OneFlow_BaseOp<"nll", [NoMemoryEffect, DeclareOpInterfaceMet ); let output = (outs OneFlow_Tensor:$output, - OneFlow_Tensor:$out_weight + OneFlow_Tensor:$out_weight, + OneFlow_Tensor:$reduced_out ); let attrs = (ins - DefaultValuedAttr:$ignore_index + DefaultValuedAttr:$ignore_index, + DefaultValuedAttr:$reduction ); let has_data_type_infer_fn = 1; let has_logical_tensor_desc_infer_fn = 1; @@ -6128,7 +6130,8 @@ def OneFlow_NLLGradOp : OneFlow_BaseOp<"nll_grad", [NoMemoryEffect, DeclareOpInt OneFlow_Tensor:$in_grad ); let attrs = (ins - DefaultValuedAttr:$ignore_index + DefaultValuedAttr:$ignore_index, + DefaultValuedAttr:$reduction ); let has_data_type_infer_fn = 1; let has_logical_tensor_desc_infer_fn = 1; diff --git a/oneflow/user/ops/nll_op.cpp b/oneflow/user/ops/nll_op.cpp index 7f9035539b7..305d714174b 100644 --- a/oneflow/user/ops/nll_op.cpp +++ b/oneflow/user/ops/nll_op.cpp @@ -31,6 +31,7 @@ namespace oneflow { ctx->SetOutputDType("output", 0, input_dtype); ctx->SetOutputDType("out_weight", 0, input_dtype); + ctx->SetOutputDType("reduced_out", 0, input_dtype); return Maybe::Ok(); } @@ -69,6 +70,10 @@ namespace oneflow { out_weight_desc->set_is_dynamic(is_dynamic); out_weight_desc->set_shape(Shape({N})); + user_op::TensorDesc* reduced_out_desc = ctx->MutOutputTensorDesc("reduced_out", 0); + reduced_out_desc->set_is_dynamic(false); + reduced_out_desc->set_shape(Shape({})); + return Maybe::Ok(); } @@ -78,7 +83,8 @@ namespace oneflow { .Split(user_op::OpArg("input", 0), 0) .Split(user_op::OpArg("target", 0), 0) .Split(user_op::OpArg("output", 0), 0) - .Split(user_op::OpArg("out_weight", 0), 0); + .Split(user_op::OpArg("out_weight", 0), 0) + .Broadcast(user_op::OpArg("reduced_out", 0)); if (ctx->user_op_conf().has_input("weight", 0)) { builder1.Broadcast(user_op::OpArg("weight", 0)); } @@ -89,8 +95,9 @@ namespace oneflow { auto builder2 = ctx->NewBuilder() .Split(user_op::OpArg("input", 0), shape.NumAxes() - 1) .Broadcast(user_op::OpArg("target", 0)) - .PartialSum(user_op::OpArg("output", 0)) - .PartialSum(user_op::OpArg("out_weight", 0)); + .Broadcast(user_op::OpArg("output", 0)) + .Broadcast(user_op::OpArg("out_weight", 0)) + .Broadcast(user_op::OpArg("reduced_out", 0)); if (ctx->user_op_conf().has_input("weight", 0)) { builder2.Split(user_op::OpArg("weight", 0), 0); } From 17391cdab79849313f2f5fc33a4e3638edc84a2b Mon Sep 17 00:00:00 2001 From: zhaoluyang Date: Wed, 9 Apr 2025 22:42:48 +0800 Subject: [PATCH 12/16] debug nll grad --- oneflow/core/autograd/gradient_funcs/nll.cpp | 21 ++++++++++++++----- oneflow/core/functional/functional_api.yaml | 2 +- .../core/functional/impl/nn_grad_functor.cpp | 9 ++++++-- oneflow/ir/include/OneFlow/OneFlowUserOps.td | 1 + oneflow/user/ops/nll_op.cpp | 2 ++ 5 files changed, 27 insertions(+), 8 deletions(-) diff --git a/oneflow/core/autograd/gradient_funcs/nll.cpp b/oneflow/core/autograd/gradient_funcs/nll.cpp index 42d34d29df6..734e748fe83 100644 --- a/oneflow/core/autograd/gradient_funcs/nll.cpp +++ b/oneflow/core/autograd/gradient_funcs/nll.cpp @@ -72,20 +72,31 @@ Maybe NLLGradFunction::Apply(const NLLCaptureState* ctx, const TensorTuple << Error::RuntimeError() << "The number of saved tensors is expected to be greater than or equal to 2, but got " << ctx->SavedTensors().size(); - const auto& out_grad = out_grads[0]; + const auto& out_grad = out_grads[0]; // for reduction="none" + const auto& reduced_out_grad = out_grads[2]; // for reduction="mean"/"sum" const auto& input = ctx->SavedTensors()[0]; const auto& target = ctx->SavedTensors()[1]; in_grads->resize(ctx->SavedTensors().size()); if (ctx->SavedTensors().size() == 2) { - JUST(VectorAt(*in_grads, 0)) = JUST( - functional::NLLGrad(out_grad, input, target, NullOpt, ctx->ignore_index, ctx->reduction)); + if (ctx->reduction == "none") { + JUST(VectorAt(*in_grads, 0)) = JUST(functional::NLLGrad( + out_grad, reduced_out_grad, input, target, NullOpt, ctx->ignore_index, ctx->reduction)); + } else { + JUST(VectorAt(*in_grads, 0)) = JUST(functional::NLLGrad( + out_grad, reduced_out_grad, input, target, NullOpt, ctx->ignore_index, ctx->reduction)); + } } else { // has weight auto weight = JUST(VectorAt(ctx->SavedTensors(), 2)); - JUST(VectorAt(*in_grads, 0)) = JUST( - functional::NLLGrad(out_grad, input, target, weight, ctx->ignore_index, ctx->reduction)); + if (ctx->reduction == "none") { + JUST(VectorAt(*in_grads, 0)) = JUST(functional::NLLGrad( + out_grad, reduced_out_grad, input, target, weight, ctx->ignore_index, ctx->reduction)); + } else { + JUST(VectorAt(*in_grads, 0)) = JUST(functional::NLLGrad( + out_grad, reduced_out_grad, input, target, weight, ctx->ignore_index, ctx->reduction)); + } } return Maybe::Ok(); diff --git a/oneflow/core/functional/functional_api.yaml b/oneflow/core/functional/functional_api.yaml index 05d5ab77900..8ab7f3d5aa7 100644 --- a/oneflow/core/functional/functional_api.yaml +++ b/oneflow/core/functional/functional_api.yaml @@ -1355,7 +1355,7 @@ bind_python: True - name: "nll_grad" - signature: 'Tensor(Tensor out_grad, Tensor input, Tensor target, Tensor weight=None, Int64 ignore_index, String reduction="none") => NLLGrad' + signature: 'Tensor(Tensor out_grad, Tensor reduced_out_grad, Tensor input, Tensor target, Tensor weight=None, Int64 ignore_index, String reduction="none") => NLLGrad' bind_python: False - name: "binary_cross_entropy_loss" diff --git a/oneflow/core/functional/impl/nn_grad_functor.cpp b/oneflow/core/functional/impl/nn_grad_functor.cpp index 9e1d2eea809..6cc68af005f 100644 --- a/oneflow/core/functional/impl/nn_grad_functor.cpp +++ b/oneflow/core/functional/impl/nn_grad_functor.cpp @@ -451,6 +451,7 @@ class NLLGradFunctor { NLLGradFunctor() { op_ = CHECK_JUST(one::OpBuilder("nll_grad") .Input("out_grad") + .Input("reduced_out_grad") .Input("input") .Input("target") .Output("in_grad") @@ -458,6 +459,7 @@ class NLLGradFunctor { op_weight_ = CHECK_JUST(one::OpBuilder("nll_grad") .Input("out_grad") + .Input("reduced_out_grad") .Input("input") .Input("target") .Input("weight") @@ -466,6 +468,7 @@ class NLLGradFunctor { } Maybe operator()(const std::shared_ptr& out_grad, + const std::shared_ptr& reduced_out_grad, const std::shared_ptr& input, const std::shared_ptr& target, const Optional& weight, const int64_t ignore_index, @@ -477,9 +480,11 @@ class NLLGradFunctor { if (weight) { return OpInterpUtil::Dispatch( - *op_weight_, {out_grad, input, target, JUST(JUST(weight)->detach())}, attrs); + *op_weight_, {out_grad, reduced_out_grad, input, target, JUST(JUST(weight)->detach())}, + attrs); } else { - return OpInterpUtil::Dispatch(*op_, {out_grad, input, target}, attrs); + return OpInterpUtil::Dispatch(*op_, {out_grad, reduced_out_grad, input, target}, + attrs); } } diff --git a/oneflow/ir/include/OneFlow/OneFlowUserOps.td b/oneflow/ir/include/OneFlow/OneFlowUserOps.td index 1162e4fe58b..2f4cc6ada66 100644 --- a/oneflow/ir/include/OneFlow/OneFlowUserOps.td +++ b/oneflow/ir/include/OneFlow/OneFlowUserOps.td @@ -6122,6 +6122,7 @@ def OneFlow_NLLOp : OneFlow_BaseOp<"nll", [NoMemoryEffect, DeclareOpInterfaceMet def OneFlow_NLLGradOp : OneFlow_BaseOp<"nll_grad", [NoMemoryEffect, DeclareOpInterfaceMethods]> { let input = (ins OneFlow_Tensor:$out_grad, + OneFlow_Tensor:$reduced_out_grad, OneFlow_Tensor:$input, OneFlow_Tensor:$target, Optional:$weight diff --git a/oneflow/user/ops/nll_op.cpp b/oneflow/user/ops/nll_op.cpp index 305d714174b..f3c93616bf8 100644 --- a/oneflow/user/ops/nll_op.cpp +++ b/oneflow/user/ops/nll_op.cpp @@ -179,6 +179,7 @@ namespace oneflow { .Split(user_op::OpArg("input", 0), 0) .Split(user_op::OpArg("target", 0), 0) .Split(user_op::OpArg("out_grad", 0), 0) + .Broadcast(user_op::OpArg("reduced_out_grad", 0)) .Split(user_op::OpArg("in_grad", 0), 0); if (ctx->user_op_conf().has_input("weight", 0)) { builder1.Broadcast(user_op::OpArg("weight", 0)); @@ -191,6 +192,7 @@ namespace oneflow { .Split(user_op::OpArg("input", 0), shape.NumAxes() - 1) .Broadcast(user_op::OpArg("target", 0)) .Broadcast(user_op::OpArg("out_grad", 0)) + .Broadcast(user_op::OpArg("reduced_out_grad", 0)) .Split(user_op::OpArg("in_grad", 0), shape.NumAxes() - 1); if (ctx->user_op_conf().has_input("weight", 0)) { builder2.Split(user_op::OpArg("weight", 0), 0); From dc3bc33a396facd69af1e00cd768892f168100ca Mon Sep 17 00:00:00 2001 From: zhaoluyang Date: Thu, 8 May 2025 18:30:37 +0800 Subject: [PATCH 13/16] refactor nll/nll_grad --- oneflow/core/autograd/gradient_funcs/nll.cpp | 34 ++++++++++++------- oneflow/core/functional/functional_api.yaml | 2 +- oneflow/core/functional/impl/nn_functor.cpp | 2 ++ .../core/functional/impl/nn_grad_functor.cpp | 10 ++++-- oneflow/ir/include/OneFlow/OneFlowUserOps.td | 4 ++- oneflow/user/ops/nll_op.cpp | 13 +++++-- 6 files changed, 45 insertions(+), 20 deletions(-) diff --git a/oneflow/core/autograd/gradient_funcs/nll.cpp b/oneflow/core/autograd/gradient_funcs/nll.cpp index 734e748fe83..d27a1959525 100644 --- a/oneflow/core/autograd/gradient_funcs/nll.cpp +++ b/oneflow/core/autograd/gradient_funcs/nll.cpp @@ -60,6 +60,7 @@ Maybe NLLGradFunction::Capture(NLLCaptureState* ctx, const TensorTuple& in if (inputs.size() == 3) { ctx->SaveTensorForBackward(inputs[2]); // weight } + ctx->SaveTensorForBackward(outputs[3]); // totol_weight return Maybe::Ok(); } @@ -67,11 +68,13 @@ Maybe NLLGradFunction::Apply(const NLLCaptureState* ctx, const TensorTuple TensorTuple* in_grads) const { if (!ctx->requires_grad) { return Maybe::Ok(); } - // CHECK_EQ_OR_RETURN(out_grads.size(), 2); // NOLINT(maybe-need-error-msg) + CHECK_EQ_OR_RETURN(out_grads.size(), 4); // NOLINT(maybe-need-error-msg) CHECK_GE_OR_RETURN(ctx->SavedTensors().size(), 2) << Error::RuntimeError() << "The number of saved tensors is expected to be greater than or equal to 2, but got " << ctx->SavedTensors().size(); + + // outputs: output, out_weight, reduced_out, total_weight const auto& out_grad = out_grads[0]; // for reduction="none" const auto& reduced_out_grad = out_grads[2]; // for reduction="mean"/"sum" const auto& input = ctx->SavedTensors()[0]; @@ -79,23 +82,28 @@ Maybe NLLGradFunction::Apply(const NLLCaptureState* ctx, const TensorTuple in_grads->resize(ctx->SavedTensors().size()); - if (ctx->SavedTensors().size() == 2) { + if (ctx->SavedTensors().size() == 3) { // no weight + const auto& total_weight = ctx->SavedTensors()[2]; if (ctx->reduction == "none") { - JUST(VectorAt(*in_grads, 0)) = JUST(functional::NLLGrad( - out_grad, reduced_out_grad, input, target, NullOpt, ctx->ignore_index, ctx->reduction)); + JUST(VectorAt(*in_grads, 0)) = + JUST(functional::NLLGrad(out_grad, reduced_out_grad, total_weight, input, target, NullOpt, + ctx->ignore_index, ctx->reduction)); } else { - JUST(VectorAt(*in_grads, 0)) = JUST(functional::NLLGrad( - out_grad, reduced_out_grad, input, target, NullOpt, ctx->ignore_index, ctx->reduction)); + JUST(VectorAt(*in_grads, 0)) = + JUST(functional::NLLGrad(out_grad, reduced_out_grad, total_weight, input, target, NullOpt, + ctx->ignore_index, ctx->reduction)); } - } else { - // has weight - auto weight = JUST(VectorAt(ctx->SavedTensors(), 2)); + } else if (ctx->SavedTensors().size() == 4) { // has weight + const auto& weight = ctx->SavedTensors()[2]; + const auto& total_weight = ctx->SavedTensors()[3]; if (ctx->reduction == "none") { - JUST(VectorAt(*in_grads, 0)) = JUST(functional::NLLGrad( - out_grad, reduced_out_grad, input, target, weight, ctx->ignore_index, ctx->reduction)); + JUST(VectorAt(*in_grads, 0)) = + JUST(functional::NLLGrad(out_grad, reduced_out_grad, total_weight, input, target, weight, + ctx->ignore_index, ctx->reduction)); } else { - JUST(VectorAt(*in_grads, 0)) = JUST(functional::NLLGrad( - out_grad, reduced_out_grad, input, target, weight, ctx->ignore_index, ctx->reduction)); + JUST(VectorAt(*in_grads, 0)) = + JUST(functional::NLLGrad(out_grad, reduced_out_grad, total_weight, input, target, weight, + ctx->ignore_index, ctx->reduction)); } } diff --git a/oneflow/core/functional/functional_api.yaml b/oneflow/core/functional/functional_api.yaml index 8ab7f3d5aa7..a5f453f397b 100644 --- a/oneflow/core/functional/functional_api.yaml +++ b/oneflow/core/functional/functional_api.yaml @@ -1355,7 +1355,7 @@ bind_python: True - name: "nll_grad" - signature: 'Tensor(Tensor out_grad, Tensor reduced_out_grad, Tensor input, Tensor target, Tensor weight=None, Int64 ignore_index, String reduction="none") => NLLGrad' + signature: 'Tensor(Tensor out_grad, Tensor reduced_out_grad, Tensor total_weight, Tensor input, Tensor target, Tensor weight=None, Int64 ignore_index, String reduction="none") => NLLGrad' bind_python: False - name: "binary_cross_entropy_loss" diff --git a/oneflow/core/functional/impl/nn_functor.cpp b/oneflow/core/functional/impl/nn_functor.cpp index be8f0053e0d..7a5e0ffdd2c 100644 --- a/oneflow/core/functional/impl/nn_functor.cpp +++ b/oneflow/core/functional/impl/nn_functor.cpp @@ -1820,6 +1820,7 @@ class NLLLossFunctor { .Output("output") .Output("out_weight") .Output("reduced_out") + .Output("total_weight") .Build()); op_weight_ = CHECK_JUST(one::OpBuilder("nll") @@ -1829,6 +1830,7 @@ class NLLLossFunctor { .Output("output") .Output("out_weight") .Output("reduced_out") + .Output("total_weight") .Build()); } diff --git a/oneflow/core/functional/impl/nn_grad_functor.cpp b/oneflow/core/functional/impl/nn_grad_functor.cpp index 6cc68af005f..ebf06a022df 100644 --- a/oneflow/core/functional/impl/nn_grad_functor.cpp +++ b/oneflow/core/functional/impl/nn_grad_functor.cpp @@ -452,6 +452,7 @@ class NLLGradFunctor { op_ = CHECK_JUST(one::OpBuilder("nll_grad") .Input("out_grad") .Input("reduced_out_grad") + .Input("total_weight") .Input("input") .Input("target") .Output("in_grad") @@ -460,6 +461,7 @@ class NLLGradFunctor { op_weight_ = CHECK_JUST(one::OpBuilder("nll_grad") .Input("out_grad") .Input("reduced_out_grad") + .Input("total_weight") .Input("input") .Input("target") .Input("weight") @@ -469,6 +471,7 @@ class NLLGradFunctor { Maybe operator()(const std::shared_ptr& out_grad, const std::shared_ptr& reduced_out_grad, + const std::shared_ptr& total_weight, const std::shared_ptr& input, const std::shared_ptr& target, const Optional& weight, const int64_t ignore_index, @@ -480,11 +483,12 @@ class NLLGradFunctor { if (weight) { return OpInterpUtil::Dispatch( - *op_weight_, {out_grad, reduced_out_grad, input, target, JUST(JUST(weight)->detach())}, + *op_weight_, + {out_grad, reduced_out_grad, total_weight, input, target, JUST(JUST(weight)->detach())}, attrs); } else { - return OpInterpUtil::Dispatch(*op_, {out_grad, reduced_out_grad, input, target}, - attrs); + return OpInterpUtil::Dispatch( + *op_, {out_grad, reduced_out_grad, total_weight, input, target}, attrs); } } diff --git a/oneflow/ir/include/OneFlow/OneFlowUserOps.td b/oneflow/ir/include/OneFlow/OneFlowUserOps.td index 2f4cc6ada66..f364235dc9b 100644 --- a/oneflow/ir/include/OneFlow/OneFlowUserOps.td +++ b/oneflow/ir/include/OneFlow/OneFlowUserOps.td @@ -6107,7 +6107,8 @@ def OneFlow_NLLOp : OneFlow_BaseOp<"nll", [NoMemoryEffect, DeclareOpInterfaceMet let output = (outs OneFlow_Tensor:$output, OneFlow_Tensor:$out_weight, - OneFlow_Tensor:$reduced_out + OneFlow_Tensor:$reduced_out, + OneFlow_Tensor:$total_weight ); let attrs = (ins DefaultValuedAttr:$ignore_index, @@ -6125,6 +6126,7 @@ def OneFlow_NLLGradOp : OneFlow_BaseOp<"nll_grad", [NoMemoryEffect, DeclareOpInt OneFlow_Tensor:$reduced_out_grad, OneFlow_Tensor:$input, OneFlow_Tensor:$target, + OneFlow_Tensor:$total_weight, Optional:$weight ); let output = (outs diff --git a/oneflow/user/ops/nll_op.cpp b/oneflow/user/ops/nll_op.cpp index f3c93616bf8..63c3d8d55ff 100644 --- a/oneflow/user/ops/nll_op.cpp +++ b/oneflow/user/ops/nll_op.cpp @@ -32,6 +32,7 @@ namespace oneflow { ctx->SetOutputDType("output", 0, input_dtype); ctx->SetOutputDType("out_weight", 0, input_dtype); ctx->SetOutputDType("reduced_out", 0, input_dtype); + ctx->SetOutputDType("total_weight", 0, input_dtype); return Maybe::Ok(); } @@ -74,6 +75,10 @@ namespace oneflow { reduced_out_desc->set_is_dynamic(false); reduced_out_desc->set_shape(Shape({})); + user_op::TensorDesc* total_weight_desc = ctx->MutOutputTensorDesc("total_weight", 0); + total_weight_desc->set_is_dynamic(false); + total_weight_desc->set_shape(Shape({})); + return Maybe::Ok(); } @@ -84,7 +89,8 @@ namespace oneflow { .Split(user_op::OpArg("target", 0), 0) .Split(user_op::OpArg("output", 0), 0) .Split(user_op::OpArg("out_weight", 0), 0) - .Broadcast(user_op::OpArg("reduced_out", 0)); + .Broadcast(user_op::OpArg("reduced_out", 0)) + .Broadcast(user_op::OpArg("total_weight", 0)); if (ctx->user_op_conf().has_input("weight", 0)) { builder1.Broadcast(user_op::OpArg("weight", 0)); } @@ -97,7 +103,8 @@ namespace oneflow { .Broadcast(user_op::OpArg("target", 0)) .Broadcast(user_op::OpArg("output", 0)) .Broadcast(user_op::OpArg("out_weight", 0)) - .Broadcast(user_op::OpArg("reduced_out", 0)); + .Broadcast(user_op::OpArg("reduced_out", 0)) + .Broadcast(user_op::OpArg("total_weight", 0)); if (ctx->user_op_conf().has_input("weight", 0)) { builder2.Split(user_op::OpArg("weight", 0), 0); } @@ -180,6 +187,7 @@ namespace oneflow { .Split(user_op::OpArg("target", 0), 0) .Split(user_op::OpArg("out_grad", 0), 0) .Broadcast(user_op::OpArg("reduced_out_grad", 0)) + .Broadcast(user_op::OpArg("total_weight", 0)) .Split(user_op::OpArg("in_grad", 0), 0); if (ctx->user_op_conf().has_input("weight", 0)) { builder1.Broadcast(user_op::OpArg("weight", 0)); @@ -193,6 +201,7 @@ namespace oneflow { .Broadcast(user_op::OpArg("target", 0)) .Broadcast(user_op::OpArg("out_grad", 0)) .Broadcast(user_op::OpArg("reduced_out_grad", 0)) + .Broadcast(user_op::OpArg("total_weight", 0)) .Split(user_op::OpArg("in_grad", 0), shape.NumAxes() - 1); if (ctx->user_op_conf().has_input("weight", 0)) { builder2.Split(user_op::OpArg("weight", 0), 0); From 22d387cafac5491f09f49b0d053fe8212eedc30b Mon Sep 17 00:00:00 2001 From: zhaoluyang Date: Mon, 12 May 2025 16:58:08 +0800 Subject: [PATCH 14/16] fix --- oneflow/core/autograd/gradient_funcs/nll.cpp | 26 +++++-------------- oneflow/core/functional/functional_api.yaml | 2 +- oneflow/core/functional/impl/nn_functor.cpp | 3 +++ .../core/functional/impl/nn_grad_functor.cpp | 11 ++++---- 4 files changed, 16 insertions(+), 26 deletions(-) diff --git a/oneflow/core/autograd/gradient_funcs/nll.cpp b/oneflow/core/autograd/gradient_funcs/nll.cpp index d27a1959525..713b5190b53 100644 --- a/oneflow/core/autograd/gradient_funcs/nll.cpp +++ b/oneflow/core/autograd/gradient_funcs/nll.cpp @@ -60,7 +60,7 @@ Maybe NLLGradFunction::Capture(NLLCaptureState* ctx, const TensorTuple& in if (inputs.size() == 3) { ctx->SaveTensorForBackward(inputs[2]); // weight } - ctx->SaveTensorForBackward(outputs[3]); // totol_weight + ctx->SaveTensorForBackward(outputs[3]); // total_weight return Maybe::Ok(); } @@ -84,27 +84,15 @@ Maybe NLLGradFunction::Apply(const NLLCaptureState* ctx, const TensorTuple if (ctx->SavedTensors().size() == 3) { // no weight const auto& total_weight = ctx->SavedTensors()[2]; - if (ctx->reduction == "none") { - JUST(VectorAt(*in_grads, 0)) = - JUST(functional::NLLGrad(out_grad, reduced_out_grad, total_weight, input, target, NullOpt, - ctx->ignore_index, ctx->reduction)); - } else { - JUST(VectorAt(*in_grads, 0)) = - JUST(functional::NLLGrad(out_grad, reduced_out_grad, total_weight, input, target, NullOpt, - ctx->ignore_index, ctx->reduction)); - } + JUST(VectorAt(*in_grads, 0)) = + JUST(functional::NLLGrad(out_grad, reduced_out_grad, input, target, total_weight, NullOpt, + ctx->ignore_index, ctx->reduction)); } else if (ctx->SavedTensors().size() == 4) { // has weight const auto& weight = ctx->SavedTensors()[2]; const auto& total_weight = ctx->SavedTensors()[3]; - if (ctx->reduction == "none") { - JUST(VectorAt(*in_grads, 0)) = - JUST(functional::NLLGrad(out_grad, reduced_out_grad, total_weight, input, target, weight, - ctx->ignore_index, ctx->reduction)); - } else { - JUST(VectorAt(*in_grads, 0)) = - JUST(functional::NLLGrad(out_grad, reduced_out_grad, total_weight, input, target, weight, - ctx->ignore_index, ctx->reduction)); - } + JUST(VectorAt(*in_grads, 0)) = + JUST(functional::NLLGrad(out_grad, reduced_out_grad, input, target, total_weight, weight, + ctx->ignore_index, ctx->reduction)); } return Maybe::Ok(); diff --git a/oneflow/core/functional/functional_api.yaml b/oneflow/core/functional/functional_api.yaml index a5f453f397b..ac454fd8ba7 100644 --- a/oneflow/core/functional/functional_api.yaml +++ b/oneflow/core/functional/functional_api.yaml @@ -1355,7 +1355,7 @@ bind_python: True - name: "nll_grad" - signature: 'Tensor(Tensor out_grad, Tensor reduced_out_grad, Tensor total_weight, Tensor input, Tensor target, Tensor weight=None, Int64 ignore_index, String reduction="none") => NLLGrad' + signature: 'Tensor(Tensor out_grad, Tensor reduced_out_grad, Tensor input, Tensor target, Tensor total_weight, Tensor weight=None, Int64 ignore_index, String reduction="none") => NLLGrad' bind_python: False - name: "binary_cross_entropy_loss" diff --git a/oneflow/core/functional/impl/nn_functor.cpp b/oneflow/core/functional/impl/nn_functor.cpp index 7a5e0ffdd2c..61cbdceff95 100644 --- a/oneflow/core/functional/impl/nn_functor.cpp +++ b/oneflow/core/functional/impl/nn_functor.cpp @@ -1927,6 +1927,7 @@ class CrossEntropyFunctor { .Output("output") .Output("out_weight") .Output("reduced_out") + .Output("total_weight") .Build()); op_nll_weight_ = CHECK_JUST(one::OpBuilder("nll") @@ -1936,6 +1937,7 @@ class CrossEntropyFunctor { .Output("output") .Output("out_weight") .Output("reduced_out") + .Output("total_weight") .Build()); } Maybe operator()(const std::shared_ptr& input, @@ -1990,6 +1992,7 @@ class CrossEntropyFunctor { #ifdef WITH_NPU if (!input->is_cpu()) { auto reduced_out = JUST(VectorAt(*nll_result, 2)); + // auto total_weight = JUST(VectorAt(*nll_result, 3)); if (reduction == "sum" || reduction == "mean") { return reduced_out; } } #endif // WITH_NPU diff --git a/oneflow/core/functional/impl/nn_grad_functor.cpp b/oneflow/core/functional/impl/nn_grad_functor.cpp index ebf06a022df..02dd4bba427 100644 --- a/oneflow/core/functional/impl/nn_grad_functor.cpp +++ b/oneflow/core/functional/impl/nn_grad_functor.cpp @@ -452,18 +452,18 @@ class NLLGradFunctor { op_ = CHECK_JUST(one::OpBuilder("nll_grad") .Input("out_grad") .Input("reduced_out_grad") - .Input("total_weight") .Input("input") .Input("target") + .Input("total_weight") .Output("in_grad") .Build()); op_weight_ = CHECK_JUST(one::OpBuilder("nll_grad") .Input("out_grad") .Input("reduced_out_grad") - .Input("total_weight") .Input("input") .Input("target") + .Input("total_weight") .Input("weight") .Output("in_grad") .Build()); @@ -471,24 +471,23 @@ class NLLGradFunctor { Maybe operator()(const std::shared_ptr& out_grad, const std::shared_ptr& reduced_out_grad, - const std::shared_ptr& total_weight, const std::shared_ptr& input, const std::shared_ptr& target, + const std::shared_ptr& total_weight, const Optional& weight, const int64_t ignore_index, const std::string& reduction) const { // auto& attrs = THREAD_CACHED_MUTABLE_ATTR_MAP("ignore_index"); // attrs.SetAllAttrs(ignore_index); auto& attrs = THREAD_CACHED_MUTABLE_ATTR_MAP("ignore_index", "reduction"); attrs.SetAllAttrs(ignore_index, reduction); - if (weight) { return OpInterpUtil::Dispatch( *op_weight_, - {out_grad, reduced_out_grad, total_weight, input, target, JUST(JUST(weight)->detach())}, + {out_grad, reduced_out_grad, input, target, total_weight, JUST(JUST(weight)->detach())}, attrs); } else { return OpInterpUtil::Dispatch( - *op_, {out_grad, reduced_out_grad, total_weight, input, target}, attrs); + *op_, {out_grad, reduced_out_grad, input, target, total_weight}, attrs); } } From 3152ef4a4f9df2e96867d94ed999f7c95adeb51f Mon Sep 17 00:00:00 2001 From: ShawnXuan Date: Wed, 14 May 2025 19:14:52 +0800 Subject: [PATCH 15/16] update --- oneflow/core/functional/impl/array_functor.cpp | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/oneflow/core/functional/impl/array_functor.cpp b/oneflow/core/functional/impl/array_functor.cpp index aef7ef62a3b..4f8f6a0dd78 100644 --- a/oneflow/core/functional/impl/array_functor.cpp +++ b/oneflow/core/functional/impl/array_functor.cpp @@ -588,7 +588,21 @@ class ArgWhereFunctor { const Symbol& dtype) const { auto& attrs = THREAD_CACHED_MUTABLE_ATTR_MAP("dtype"); attrs.SetAllAttrs(dtype->data_type()); +#ifdef WITH_NPU + std::string device_type = JUST(x->parallel_desc())->device_tag(); + if (device_type == DeviceType::kNPU) { + auto cpu_tensor = JUST(functional::ToDevice(x, "cpu")); + auto result = JUST(OpInterpUtil::Dispatch(*op_, {cpu_tensor}, attrs)); + for (int i = 0; i < result->size(); ++i) { + (*result)[i] = JUST(functional::ToDevice((*result)[i], "npu")); + } + return result; + } else { + return OpInterpUtil::Dispatch(*op_, {x}, attrs); + } +#else return OpInterpUtil::Dispatch(*op_, {x}, attrs); +#endif // WITH_NPU } private: From 01044a86fe734f9900b295499cc46d8c2e962739 Mon Sep 17 00:00:00 2001 From: zhaoluyang Date: Thu, 15 May 2025 11:20:57 +0800 Subject: [PATCH 16/16] refine --- oneflow/core/functional/impl/array_functor.cpp | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/oneflow/core/functional/impl/array_functor.cpp b/oneflow/core/functional/impl/array_functor.cpp index 4f8f6a0dd78..9c1a4b89cee 100644 --- a/oneflow/core/functional/impl/array_functor.cpp +++ b/oneflow/core/functional/impl/array_functor.cpp @@ -589,16 +589,22 @@ class ArgWhereFunctor { auto& attrs = THREAD_CACHED_MUTABLE_ATTR_MAP("dtype"); attrs.SetAllAttrs(dtype->data_type()); #ifdef WITH_NPU - std::string device_type = JUST(x->parallel_desc())->device_tag(); + auto device_type = DeviceType::kCPU; + if (x->is_global()) { + device_type = JUST(x->parallel_desc())->device_type(); + } else { + device_type = JUST(x->device())->enum_type(); + } if (device_type == DeviceType::kNPU) { - auto cpu_tensor = JUST(functional::ToDevice(x, "cpu")); + // NOTE: use cpu argwhere when device="npu" + auto cpu_tensor = JUST(one::functional::To(x, "cpu")); auto result = JUST(OpInterpUtil::Dispatch(*op_, {cpu_tensor}, attrs)); for (int i = 0; i < result->size(); ++i) { - (*result)[i] = JUST(functional::ToDevice((*result)[i], "npu")); + (*result)[i] = JUST(one::functional::To((*result)[i], "npu")); } return result; } else { - return OpInterpUtil::Dispatch(*op_, {x}, attrs); + return OpInterpUtil::Dispatch(*op_, {x}, attrs); } #else return OpInterpUtil::Dispatch(*op_, {x}, attrs);