diff --git a/transformer_engine/pytorch/csrc/extensions/ep.cpp b/transformer_engine/pytorch/csrc/extensions/ep.cpp index 8173df947e..118f14a01f 100644 --- a/transformer_engine/pytorch/csrc/extensions/ep.cpp +++ b/transformer_engine/pytorch/csrc/extensions/ep.cpp @@ -145,7 +145,13 @@ void ep_initialize(uintptr_t comm_ptr, const std::string& group_name, int64_t nu .max_token_dtype = static_cast(GetTransformerEngineDType(torch_dtype)), .zero_copy = zero_copy ? 1 : 0, }; - nvte_ep_initialize(static_cast(ep_comm), &cfg); + // Release the GIL only around the native init. It must stay held while pybind11 casts + // the ``max_token_dtype`` object above and destroys the by-value ``pybind11::object`` + // parameter on return; releasing it across those trips pybind11's dec_ref GIL assertion. + { + pybind11::gil_scoped_release nogil; + nvte_ep_initialize(static_cast(ep_comm), &cfg); + } g_zero_copy_enabled.store(zero_copy, std::memory_order_relaxed); g_ep_initialized = true; g_ep_group_name = group_name; @@ -366,8 +372,7 @@ void register_ep_bindings(pybind11::module_& m) { "Initialize the EP backend; borrows torch's NCCL comm pointed to by ``comm_ptr``.", py::arg("comm_ptr"), py::arg("group_name"), py::arg("num_experts"), py::arg("max_tokens_per_rank"), py::arg("max_recv_tokens_per_rank"), py::arg("hidden_dim"), - py::arg("max_num_sms") = 0, py::arg("max_token_dtype"), py::arg("zero_copy") = false, - py::call_guard()); + py::arg("max_num_sms") = 0, py::arg("max_token_dtype"), py::arg("zero_copy") = false); m.def("ep_finalize", &ep_finalize, "Tear down the EP backend. Idempotent.", py::call_guard()); m.def("ep_get_zero_copy", &ep_get_zero_copy, "Return the current EP zero-copy toggle state."); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 9c9ec36138..7e9d114be8 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -716,6 +716,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { int num_max_streams, int comm_cga_size, int gemm_priority, int comm_priority, int num_comm_sm, bool set_sm_margin, bool atomic_gemm, bool rs_overlap_first_gemm) { + // Release the GIL only around the native construction (blocking collectives) to avoid + // tripping pybind11's inc_ref/dec_ref GIL assertions. + py::gil_scoped_release nogil; if (use_cublasmp) { return std::make_shared(helper, helper->mylocal, tp_size, comm_type, buffer_shape, buffer_dtype, num_comm_sm, @@ -726,8 +729,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { comm_cga_size, gemm_priority, comm_priority, num_comm_sm, set_sm_margin, atomic_gemm, rs_overlap_first_gemm); }), - py::call_guard(), py::arg("buffer_shape"), - py::arg("buffer_dtype"), py::arg("helper"), py::arg("tp_size"), + py::arg("buffer_shape"), py::arg("buffer_dtype"), py::arg("helper"), py::arg("tp_size"), py::arg("use_cublasmp") = false, py::arg("comm_type") = transformer_engine::CommOverlapType::RS, py::arg("num_splits") = 4, py::arg("num_max_streams") = NVTE_COMM_OVERLAP_MAX_STREAMS, @@ -751,6 +753,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { int comm_cga_size, int gemm_priority, int comm_priority, int num_comm_sm, bool set_sm_margin, bool atomic_gemm, bool use_ce, bool aggregate, bool use_cublasmp) { + // Release the GIL only around the native construction (blocking collectives) to avoid + // tripping pybind11's inc_ref/dec_ref GIL assertions. + py::gil_scoped_release nogil; if (use_cublasmp) { return std::make_shared(helper, helper->mylocal, tp_size, comm_type, buffer_shape, buffer_dtype, num_comm_sm, @@ -761,12 +766,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { gemm_priority, comm_priority, num_comm_sm, set_sm_margin, atomic_gemm, use_ce, aggregate); }), - py::call_guard(), py::arg("buffer_shape"), - py::arg("buffer_dtype"), py::arg("helper"), py::arg("tp_size"), py::arg("comm_type"), - py::arg("num_max_streams") = NVTE_COMM_OVERLAP_MAX_STREAMS, py::arg("comm_cga_size") = 1, - py::arg("gemm_priority") = 0, py::arg("comm_priority") = 0, py::arg("num_comm_sm") = 1, - py::arg("set_sm_margin") = false, py::arg("atomic_gemm") = false, - py::arg("use_ce") = true, py::arg("aggregate") = false, py::arg("use_cublasmp") = false) + py::arg("buffer_shape"), py::arg("buffer_dtype"), py::arg("helper"), py::arg("tp_size"), + py::arg("comm_type"), py::arg("num_max_streams") = NVTE_COMM_OVERLAP_MAX_STREAMS, + py::arg("comm_cga_size") = 1, py::arg("gemm_priority") = 0, py::arg("comm_priority") = 0, + py::arg("num_comm_sm") = 1, py::arg("set_sm_margin") = false, + py::arg("atomic_gemm") = false, py::arg("use_ce") = true, py::arg("aggregate") = false, + py::arg("use_cublasmp") = false) .def("copy_into_buffer", static_cast( &CommOverlapP2P::copy_into_buffer),