diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 534fd9642a..dcd1363ecc 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -26,6 +26,8 @@ ZERO_COPY = os.environ.get("NVTE_EP_ZERO_COPY", "0") == "1" +EAGER = os.environ.get("NVTE_EP_EAGER", "0") == "1" +OVERFLOW = os.environ.get("NVTE_EP_OVERFLOW", "0") == "1" # Must come after the transformer_engine import so libtransformer_engine.so is loaded. import transformer_engine_torch as tex # noqa: F401 @@ -43,6 +45,18 @@ def _zero_copy_test_include(fn): return fn +def _eager_test_include(fn): + """Mark a test to run in the eager pass; others skip there.""" + fn._eager_test_include = True + return fn + + +def _overflow_test_include(fn): + """Mark a test to run in the overflow (drop-on-overflow) pass; others skip there.""" + fn._overflow_test_include = True + return fn + + class _StageToSymm(torch.autograd.Function): """Identity op that stages ``src`` into a symm-mem buffer; grad passes through. Lets a test feed a symm-mem-backed, autograd-tracked tensor into ep_combine. @@ -126,6 +140,10 @@ def _make_cfg() -> _Cfg: active = min(cfg.num_experts, T * cfg.ep_size * TOP_K) overconc = cfg.num_experts // active cfg.recv_capacity_per_rank = NUM_LOCAL_EXPERTS * max(T * cfg.ep_size * TOP_K, 16) * overconc * 2 + if OVERFLOW: + # Undersize recv capacity so identity routing overflows a rank's budget; + # HT requires capacity >= max_tokens_per_rank. + cfg.recv_capacity_per_rank = TOKENS_PER_RANK cfg.device = torch.device("cuda", torch.cuda.current_device()) return cfg @@ -147,6 +165,9 @@ def setUpClass(cls): recv_capacity_per_rank=cls.cfg.recv_capacity_per_rank, hidden_dim=HIDDEN_DIM, zero_copy=ZERO_COPY, + eager=EAGER, + max_num_topk=TOP_K, + drop_on_overflow=OVERFLOW, ) def setUp(self): @@ -155,6 +176,14 @@ def setUp(self): getattr(self, self._testMethodName), "_zero_copy_test_include", False ): self.skipTest("not exercised in zero-copy mode") + # Only the eager-capable tests run in the eager pass. + if EAGER and not getattr(getattr(self, self._testMethodName), "_eager_test_include", False): + self.skipTest("not exercised in eager mode") + # Only the overflow-capable tests run in the overflow pass. + if OVERFLOW and not getattr( + getattr(self, self._testMethodName), "_overflow_test_include", False + ): + self.skipTest("not exercised in overflow mode") def _make_buffer(self, alignment=0, top_k=TOP_K): return EpBuffer( @@ -184,12 +213,12 @@ def _stage_grad_symm(self, x, symm_buf=None): return _GradToSymm.apply(x, symm_buf) def _make_raw_recv(self, dtype=torch.bfloat16): - """Raw recv tensors + token_counts for the primitive tests.""" + """Raw recv tensors + tokens_per_expert for the primitive tests.""" rc = self.cfg.recv_capacity_per_rank return ( torch.empty(rc, HIDDEN_DIM, dtype=dtype, device=self.cfg.device), torch.empty(rc, dtype=torch.float32, device=self.cfg.device), - torch.empty(NUM_LOCAL_EXPERTS, dtype=torch.int32, device=self.cfg.device), + torch.empty(NUM_LOCAL_EXPERTS, dtype=torch.int64, device=self.cfg.device), ) @staticmethod @@ -205,17 +234,63 @@ def _moe_step(self, buffer, topk_idx, tokens, w): # Prepare + @_eager_test_include def test_primitive_prepare(self): buf = self._make_buffer() topk_idx, _toks, _w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) - token_counts = ep_prepare(buf, topk_idx) + tokens_per_expert = ep_prepare(buf, topk_idx) torch.cuda.synchronize() - self.assertEqual(token_counts.shape, (NUM_LOCAL_EXPERTS,)) - local = int(token_counts.sum().item()) + self.assertEqual(tokens_per_expert.shape, (NUM_LOCAL_EXPERTS,)) + local = int(tokens_per_expert.sum().item()) total = torch.tensor([local], dtype=torch.int64, device=self.cfg.device) dist.all_reduce(total, op=dist.ReduceOp.SUM, group=self.ep_group) self.assertEqual(int(total.item()), self.cfg.world_size * TOKENS_PER_RANK * TOP_K) + @_eager_test_include + def test_eager_recv_sizing(self): + """Eager mode sizes dispatch outputs to the exact per-step recv-token total.""" + if not EAGER: + self.skipTest("eager-only assertions") + buf = self._make_buffer() + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + recv_t, recv_w, tokens_per_expert = ep_dispatch(buf, tokens, topk_idx, w) + torch.cuda.synchronize() + # The per-step recv-token total is exposed on the buffer (int64 [1]). + self.assertEqual(buf.total_recv_tokens.dtype, torch.int64) + total = int(buf.total_recv_tokens.item()) + # recv outputs are sized to the recv total, not recv_capacity_per_rank. + self.assertEqual(recv_t.shape[0], total) + self.assertEqual(recv_w.shape[0], total) + # padded total is at least the unpadded per-expert sum and within capacity. + self.assertGreaterEqual(total, int(tokens_per_expert.sum().item())) + self.assertLessEqual(total, self.cfg.recv_capacity_per_rank) + + @_overflow_test_include + def test_overflow_drop(self): + """drop_on_overflow: recv past capacity is dropped and dispatch continues + instead of trapping; the pre-drop recv total exceeds recv_capacity.""" + if not OVERFLOW: + self.skipTest("overflow-only assertions") + buf = self._make_buffer() + topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size) + # Identity routing sends TOKENS_PER_RANK * TOP_K tokens to each rank, which + # overflows the deliberately undersized capacity. + expected_recv = TOKENS_PER_RANK * TOP_K + self.assertGreater(expected_recv, self.cfg.recv_capacity_per_rank) + # total_recv_tokens reports the true (pre-drop) recv total, counting the + # tokens that will be dropped; the per-expert counts exclude them and sum + # to the kept tokens (capped at recv_capacity_per_rank). + tokens_per_expert = ep_prepare(buf, topk_idx) + torch.cuda.synchronize() + self.assertEqual(int(buf.total_recv_tokens.item()), expected_recv) + self.assertEqual(int(tokens_per_expert.sum().item()), self.cfg.recv_capacity_per_rank) + # Dispatch drops overflowing tokens and completes (no trap); recv outputs + # stay capped at recv_capacity_per_rank. + recv_t, recv_w, _ = ep_dispatch(buf, tokens, topk_idx, w) + torch.cuda.synchronize() + self.assertEqual(recv_t.shape[0], self.cfg.recv_capacity_per_rank) + self.assertEqual(recv_w.shape[0], self.cfg.recv_capacity_per_rank) + # Identity round-trip via raw primitives def test_primitive_dispatch_combine_identity(self): @@ -330,6 +405,7 @@ def test_zero_copy_pool_auto_alloc(self): # Multi-iter stability + @_eager_test_include def test_dispatch_autograd_multiple_iterations(self): """5 fwd+bwd iters on the same EpBuffer must be bit-stable.""" buf = self._make_buffer() @@ -479,6 +555,7 @@ def zero_grads(): ) @_zero_copy_test_include + @_eager_test_include def test_combine_autograd(self): """ep_combine fwd+bwd; bwd grad target is the EpBuffer symm buffer (zc) or in-flight.""" buf = self._make_buffer() diff --git a/tests/pytorch/distributed/run_test_ep.sh b/tests/pytorch/distributed/run_test_ep.sh index 62c9a0207f..10e814eb80 100755 --- a/tests/pytorch/distributed/run_test_ep.sh +++ b/tests/pytorch/distributed/run_test_ep.sh @@ -45,11 +45,13 @@ RET=0 run_pass() { local label="$1" local zc="$2" + local eager="${3:-0}" + local overflow="${4:-0}" local log="stdout_ep_${label}.txt" echo "=== Running ${SCRIPT} [${label}] on ${NUM_RANKS} GPUs (timeout=${TEST_TIMEOUT_S}s) ===" # setsid + kill-after so SIGKILL takes down the whole process group, not just torchrun. - NVTE_EP_ZERO_COPY="${zc}" setsid timeout --foreground --kill-after=10 --signal=TERM \ - "${TEST_TIMEOUT_S}" \ + NVTE_EP_ZERO_COPY="${zc}" NVTE_EP_EAGER="${eager}" NVTE_EP_OVERFLOW="${overflow}" setsid timeout --foreground \ + --kill-after=10 --signal=TERM "${TEST_TIMEOUT_S}" \ torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_RANKS}" \ "${SCRIPT}" 2>&1 | tee "${log}" local rc=${PIPESTATUS[0]} @@ -68,5 +70,7 @@ run_pass() { run_pass "default" 0 run_pass "zero_copy" 1 +run_pass "eager" 0 1 +run_pass "overflow" 0 0 1 exit $RET diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index a82ec1c98d..e55705357d 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -94,8 +94,15 @@ void EPBackend::validate_config(const NVTEEpGroupConfig& config) { NVTE_CHECK(config.num_experts > 0, "num_experts must be positive, got ", config.num_experts); NVTE_CHECK(config.max_tokens_per_rank > 0, "max_tokens_per_rank must be positive, got ", config.max_tokens_per_rank); - NVTE_CHECK(config.max_recv_tokens_per_rank > 0, "max_recv_tokens_per_rank must be positive, got ", + // 0 selects eager mode (NCCL_EP_AUTO); any explicit budget must be positive. + NVTE_CHECK(config.max_recv_tokens_per_rank >= 0, + "max_recv_tokens_per_rank must be non-negative, got ", config.max_recv_tokens_per_rank); + NVTE_CHECK(!(config.zero_copy && config.max_recv_tokens_per_rank == 0), + "zero-copy and eager (max_recv_tokens_per_rank = 0) modes are mutually exclusive"); + NVTE_CHECK(!(config.drop_on_overflow && config.max_recv_tokens_per_rank == 0), + "drop_on_overflow (overflow drop) is not supported in eager mode " + "(max_recv_tokens_per_rank = 0)"); NVTE_CHECK(config.hidden_dim > 0, "hidden_dim must be positive, got ", config.hidden_dim); NVTE_CHECK(config.max_token_dtype >= 0 && config.max_token_dtype < kNVTENumTypes, "max_token_dtype out of range, got ", static_cast(config.max_token_dtype)); @@ -210,9 +217,14 @@ void EPBackend::init(ncclComm_t ep_comm, NVTEEpGroupConfig group_config) { cfg.max_num_sms = group_config.num_comm_sms > 0 ? static_cast(group_config.num_comm_sms) : NCCL_EP_AUTO; - // Must be > 0; NCCL EP errors out on 0. + // 0 = NCCL_EP_AUTO, which enables eager mode (recv buffers sized per routing). cfg.max_recv_tokens_per_rank = static_cast(group_config.max_recv_tokens_per_rank); cfg.zero_copy = group_config.zero_copy ? NCCL_EP_ZERO_COPY_ON : NCCL_EP_ZERO_COPY_OFF; + // Upper bound on per-token top-k; required for eager mode with the + // expert-major layout, where NCCL EP sizes internal buffers from it. + cfg.num_topk = static_cast(group_config.num_topk); + cfg.overflow_policy = + group_config.drop_on_overflow ? NCCL_EP_OVERFLOW_DROP : NCCL_EP_OVERFLOW_AUTO; NVTE_CHECK_NCCL(ncclEpCreateGroup(&ep_group_, ep_comm, &cfg)); @@ -320,10 +332,8 @@ size_t EPBackend::handle_mem_size(NVTEEpLayerConfig layer_cfg) { } void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx, - NVTETensor recv_tokens_per_expert, - NVTETensor /*total_recv_tokens_per_rank*/, NVTEEpLayerConfig layer_cfg, - cudaStream_t stream) { - // total_recv_tokens_per_rank is a reserved placeholder; not yet populated. + NVTETensor recv_tokens_per_expert, NVTETensor total_recv_tokens_per_rank, + NVTEEpLayerConfig layer_cfg, cudaStream_t stream) { NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); NVTE_CHECK(layer_cfg.top_k > 0, "top_k must be > 0, got ", layer_cfg.top_k); NVTE_CHECK(nvte_tensor_shape(topk_idx).ndim == 2, "topk_idx must be 2D [T, top_k]"); @@ -331,16 +341,24 @@ void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx, NVTEShape topk_idx_shape; ncclEpTensor_t nccl_topk_idx = make_nccl_ep_tensor(topk_idx, topk_idx_shape); - // ncclEpUpdateHandle writes per-expert counts via expert_counters. + // ncclEpUpdateHandle writes per-expert counts via expert_counters and, when + // provided, the scalar padded recv-slot total via recv_total_counter. NVTEShape recv_tokens_per_expert_shape; ncclEpTensor_t recv_tokens_per_expert_desc; if (recv_tokens_per_expert != nullptr) { recv_tokens_per_expert_desc = make_nccl_ep_tensor(recv_tokens_per_expert, recv_tokens_per_expert_shape); } + NVTEShape total_recv_shape; + ncclEpTensor_t total_recv_desc; + if (total_recv_tokens_per_rank != nullptr) { + total_recv_desc = make_nccl_ep_tensor(total_recv_tokens_per_rank, total_recv_shape); + } ncclEpLayoutInfo_t layout_info = NCCL_EP_LAYOUT_INFO_INIT; layout_info.expert_counters = (recv_tokens_per_expert != nullptr) ? &recv_tokens_per_expert_desc : nullptr; + layout_info.recv_total_counter = + (total_recv_tokens_per_rank != nullptr) ? &total_recv_desc : nullptr; std::lock_guard lock(mutex_); NVTE_CHECK(initialized_, "EPBackend not initialized"); diff --git a/transformer_engine/common/include/transformer_engine/ep.h b/transformer_engine/common/include/transformer_engine/ep.h index 224622fd41..152add336e 100644 --- a/transformer_engine/common/include/transformer_engine/ep.h +++ b/transformer_engine/common/include/transformer_engine/ep.h @@ -44,7 +44,9 @@ typedef struct { int num_experts; /*! Upper bound on tokens this rank sends per dispatch. */ int max_tokens_per_rank; - /*! Upper bound on tokens this rank receives per dispatch (must be > 0). */ + /*! Upper bound on tokens this rank receives per dispatch. 0 selects eager + * mode: the library derives its internal bound and the caller sizes recv + * buffers to the per-routing recv count. */ int max_recv_tokens_per_rank; /*! Token hidden dimension. */ int hidden_dim; @@ -58,6 +60,14 @@ typedef struct { * by NVTECommWindow handles and transfer in place (no staging copies); * 0 (default) = staged. */ int zero_copy; + /*! Upper bound on per-token top-k across the group's handles. Required for + * eager mode (max_recv_tokens_per_rank == 0) with the expert-major layout, + * where it sizes internal buffers; 0 = unset. */ + int num_topk; + /*! Recv overflow policy. When nonzero, tokens past max_recv_tokens_per_rank + * are dropped and dispatch continues; 0 (default) traps on overflow. + * Not supported in eager mode. */ + int drop_on_overflow; } NVTEEpGroupConfig; /*! \brief Per-layer configuration consumed by nvte_ep_handle_mem_size and @@ -121,13 +131,14 @@ size_t nvte_ep_handle_mem_size(const NVTEEpLayerConfig* layer_cfg); * AllGathers topk_idx across the EP group and stages per-expert offsets and * counts into handle_mem so the matching dispatch/combine/_bwd can run with * no further routing computation. Must precede every dispatch/combine/_bwd - * that uses this handle_mem. recv_tokens_per_expert becomes host-valid after a - * stream sync. + * that uses this handle_mem. recv_tokens_per_expert and total_recv_tokens_per_rank + * become host-valid after a stream sync. * * \param[in] handle_mem uint8 routing-state buffer. * \param[in] topk_idx [T, top_k] int64 routing indices. - * \param[out] recv_tokens_per_expert [num_local_experts] int32 counts. - * \param[out] total_recv_tokens_per_rank Reserved placeholder; may be null. Unused for now. + * \param[out] recv_tokens_per_expert [num_local_experts] int32/int64 counts. + * \param[out] total_recv_tokens_per_rank Optional [1] int32/int64 scalar: padded + * recv-slot total for this rank. May be null. * \param[in] layer_cfg Per-call layer configuration (struct_size set). * \param[in] stream CUDA stream. */ diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 6edfbdc00e..75068c08a1 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -676,7 +676,7 @@ void grouped_swizzle_for_gemm(py::handle &tensor, bool rowwise, bool columnwise) void ep_initialize(uintptr_t comm_ptr, const std::string &group_name, int64_t num_experts, int64_t max_tokens_per_rank, int64_t max_recv_tokens_per_rank, int64_t hidden_dim, int64_t max_num_sms, pybind11::object max_token_dtype, - bool zero_copy); + bool zero_copy, int64_t max_num_topk, bool drop_on_overflow); void ep_finalize(); @@ -686,8 +686,9 @@ bool ep_get_zero_copy(); // Returns the handle_mem byte size for the given layer config. int64_t ep_handle_mem_size(int64_t top_k, int64_t dispatch_output_per_expert_alignment); -void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor token_counts, int64_t top_k, - int64_t dispatch_output_per_expert_alignment); +void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens_per_expert, + int64_t top_k, int64_t dispatch_output_per_expert_alignment, + at::Tensor total_recv_tokens); void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights); diff --git a/transformer_engine/pytorch/csrc/extensions/ep.cpp b/transformer_engine/pytorch/csrc/extensions/ep.cpp index 118f14a01f..ab8571899a 100644 --- a/transformer_engine/pytorch/csrc/extensions/ep.cpp +++ b/transformer_engine/pytorch/csrc/extensions/ep.cpp @@ -125,7 +125,7 @@ bool ep_get_zero_copy() { return g_zero_copy_enabled.load(std::memory_order_rela void ep_initialize(uintptr_t comm_ptr, const std::string& group_name, int64_t num_experts, int64_t max_tokens_per_rank, int64_t max_recv_tokens_per_rank, int64_t hidden_dim, int64_t max_num_sms, pybind11::object max_token_dtype, - bool zero_copy) { + bool zero_copy, int64_t max_num_topk, bool drop_on_overflow) { NVTE_CHECK(!group_name.empty(), "group_name must be non-empty (used for symm-mem lookup)"); NVTE_CHECK(comm_ptr != 0, "comm_ptr must be non-null (torch NCCL host comm pointer)"); NVTE_CHECK(!g_ep_initialized, "ep_initialize called twice without ep_finalize"); @@ -144,6 +144,8 @@ void ep_initialize(uintptr_t comm_ptr, const std::string& group_name, int64_t nu .num_comm_sms = static_cast(max_num_sms), .max_token_dtype = static_cast(GetTransformerEngineDType(torch_dtype)), .zero_copy = zero_copy ? 1 : 0, + .num_topk = static_cast(max_num_topk), + .drop_on_overflow = drop_on_overflow ? 1 : 0, }; // 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`` @@ -186,24 +188,34 @@ int64_t ep_handle_mem_size(int64_t top_k, int64_t dispatch_output_per_expert_ali // ── Per-step ops ───────────────────────────────────────────────────────────── -void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor token_counts, int64_t top_k, - int64_t dispatch_output_per_expert_alignment) { +void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens_per_expert, + int64_t top_k, int64_t dispatch_output_per_expert_alignment, + at::Tensor total_recv_tokens) { auto stream = at::cuda::getCurrentCUDAStream().stream(); NVTE_CHECK(topk_idx.dim() >= 2, "topk_idx must be at least 2D [..., top_k]"); auto idx_dtype = check_topk_idx_dtype(topk_idx); + // NCCL EP requires all prepare int output counters to share a dtype. + NVTE_CHECK(tokens_per_expert.scalar_type() == at::kLong, "tokens_per_expert must be int64"); + NVTE_CHECK(total_recv_tokens.scalar_type() == at::kLong, "total_recv_tokens must be int64"); const size_t T_flat = topk_idx.numel() / topk_idx.size(-1); const size_t topk_n = static_cast(topk_idx.size(-1)); auto topk_idx_te = makeTransformerEngineTensor(topk_idx.data_ptr(), Shape{T_flat, topk_n}, idx_dtype); - auto token_counts_te = makeTransformerEngineTensor( - token_counts.data_ptr(), Shape{static_cast(token_counts.numel())}, DType::kInt32); + auto tokens_per_expert_te = makeTransformerEngineTensor( + tokens_per_expert.data_ptr(), Shape{static_cast(tokens_per_expert.numel())}, + DType::kInt64); auto handle_mem_te = makeTransformerEngineTensor( handle_mem.data_ptr(), Shape{static_cast(handle_mem.numel())}, DType::kByte); + // [1] int64 scalar recv-slot total; lets the caller size dispatch outputs + // (eager) or detect overflow past recv_capacity_per_rank (graph mode). + auto total_recv_tokens_te = makeTransformerEngineTensor( + total_recv_tokens.data_ptr(), Shape{static_cast(total_recv_tokens.numel())}, + DType::kInt64); auto layer_cfg = make_layer_cfg(top_k, dispatch_output_per_expert_alignment); - nvte_ep_prepare(handle_mem_te.data(), topk_idx_te.data(), token_counts_te.data(), - /*total_recv_tokens_per_rank=*/nullptr, &layer_cfg, stream); + nvte_ep_prepare(handle_mem_te.data(), topk_idx_te.data(), tokens_per_expert_te.data(), + total_recv_tokens_te.data(), &layer_cfg, stream); } void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens, @@ -372,14 +384,18 @@ 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::arg("max_num_sms") = 0, py::arg("max_token_dtype"), py::arg("zero_copy") = false, + py::arg("max_num_topk") = 0, py::arg("drop_on_overflow") = 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."); m.def("ep_handle_mem_size", &ep_handle_mem_size, "Return the handle_mem byte size for the given layer config.", py::arg("top_k"), py::arg("dispatch_output_per_expert_alignment") = 0); - m.def("ep_prepare", &ep_prepare, "EP prepare", py::call_guard()); + m.def("ep_prepare", &ep_prepare, "EP prepare", py::arg("handle_mem"), py::arg("topk_idx"), + py::arg("tokens_per_expert"), py::arg("top_k"), + py::arg("dispatch_output_per_expert_alignment"), py::arg("total_recv_tokens"), + py::call_guard()); m.def("ep_dispatch", &ep_dispatch, "EP dispatch", py::call_guard()); m.def("ep_combine", &ep_combine, "EP combine", py::call_guard()); m.def("ep_dispatch_bwd", &ep_dispatch_bwd, "EP dispatch backward", diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py index 0ae4805caa..444d29871f 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -67,14 +67,16 @@ def _check_nccl_runtime_version() -> None: _BOOTSTRAPPED = False _ATEXIT_REGISTERED = False -# EP group captured at bootstrap; EpBuffer uses it to allocate the symm-mem -# combine grad buffer in zero-copy mode. +# EP group captured at bootstrap; used by the zero-copy symm-mem pool allocator. _EP_GROUP: Optional[dist.ProcessGroup] = None +# Eager-mode toggle captured at bootstrap; ep_dispatch reads it to size the recv +# outputs from the per-step recv-token total instead of recv_capacity_per_rank. +_EAGER = False def _atexit_finalize() -> None: """Best-effort teardown at interpreter shutdown; swallows errors.""" - global _BOOTSTRAPPED, _EP_GROUP + global _BOOTSTRAPPED, _EP_GROUP, _EAGER if _BOOTSTRAPPED: try: tex.ep_finalize() @@ -85,6 +87,7 @@ def _atexit_finalize() -> None: finally: _BOOTSTRAPPED = False _EP_GROUP = None + _EAGER = False def ep_bootstrap( @@ -95,6 +98,9 @@ def ep_bootstrap( hidden_dim: int, max_num_sms: int = 0, zero_copy: bool = False, + eager: bool = False, + max_num_topk: int = 0, + drop_on_overflow: bool = False, max_token_dtype: torch.dtype = torch.bfloat16, ) -> None: """Initialize EP by borrowing ep_group's NCCL comm. Call once per process. @@ -105,12 +111,29 @@ def ep_bootstrap( ``zero_copy`` opts the EP group into the symm-mem zero-copy IO path; pass ``True`` only when payload tensors are allocated via ``symm_mem_alloc``. Defaults to ``False``. + + ``eager`` sizes the dispatch recv outputs from the per-step recv-token total + reported by ep_prepare, instead of the fixed ``recv_capacity_per_rank`` upper + bound. This requires a host sync each step, so it is not CUDA-graph + capturable. Mutually exclusive with ``zero_copy``. Defaults to ``False``. + + ``max_num_topk`` is the upper bound on per-token top-k; it sizes NCCL EP + internal buffers and is required (>= 1) for ``eager`` mode. + + ``drop_on_overflow`` drops tokens that exceed ``recv_capacity_per_rank`` + instead of trapping. Not supported in ``eager`` mode. Defaults to ``False``. """ - global _BOOTSTRAPPED, _ATEXIT_REGISTERED, _EP_GROUP + global _BOOTSTRAPPED, _ATEXIT_REGISTERED, _EP_GROUP, _EAGER if _BOOTSTRAPPED: raise RuntimeError("ep_bootstrap was already called in this process") if ep_group.size() < 2: raise ValueError(f"ep_bootstrap requires ep_group.size() >= 2 (got {ep_group.size()}).") + if zero_copy and eager: + raise ValueError("ep_bootstrap: zero_copy and eager modes are mutually exclusive") + if eager and max_num_topk < 1: + raise ValueError("ep_bootstrap: eager mode requires max_num_topk >= 1") + if eager and drop_on_overflow: + raise ValueError("ep_bootstrap: drop_on_overflow is not supported in eager mode") _check_nccl_runtime_version() if zero_copy: warnings.warn( @@ -128,14 +151,19 @@ def ep_bootstrap( str(ep_group.group_name), int(num_experts), int(max_tokens_per_rank), - int(recv_capacity_per_rank), + # Eager mode sizes recv buffers per routing, so the group uses the + # library-derived bound (0 = NCCL_EP_AUTO) instead of a fixed budget. + 0 if eager else int(recv_capacity_per_rank), int(hidden_dim), int(max_num_sms), max_token_dtype, bool(zero_copy), + int(max_num_topk), + bool(drop_on_overflow), ) _BOOTSTRAPPED = True _EP_GROUP = ep_group + _EAGER = bool(eager) if not _ATEXIT_REGISTERED: atexit.register(_atexit_finalize) _ATEXIT_REGISTERED = True @@ -149,7 +177,7 @@ def ep_finalize() -> None: ``dist.destroy_process_group()``, since the borrowed NCCL comm becomes invalid once the PG is destroyed. """ - global _BOOTSTRAPPED, _EP_GROUP + global _BOOTSTRAPPED, _EP_GROUP, _EAGER if not _BOOTSTRAPPED: return try: @@ -157,6 +185,7 @@ def ep_finalize() -> None: finally: _BOOTSTRAPPED = False _EP_GROUP = None + _EAGER = False def is_symm_backed(t: torch.Tensor) -> bool: @@ -183,7 +212,7 @@ def is_symm_backed(t: torch.Tensor) -> bool: class EpBuffer: - """Per-microbatch EP layer state: handle_mem, token_counts, and shape/dtype config. + """Per-microbatch EP layer state: handle_mem, tokens_per_expert, and shape/dtype config. Use one EpBuffer per concurrently-in-flight call (e.g. per PP-1F1B microbatch). """ @@ -197,8 +226,11 @@ class EpBuffer: "num_local_experts", "payload_dtype", "device", - "token_counts", + "tokens_per_expert", "zero_copy", + "eager", + "total_recv_tokens", + "_host_total_recv_tokens", ) def __init__( @@ -226,12 +258,24 @@ def __init__( self.payload_dtype = payload_dtype self.device = device self.zero_copy = bool(tex.ep_get_zero_copy()) + self.eager = _EAGER size_bytes = tex.ep_handle_mem_size(self.top_k, self.alignment) self.handle_mem = torch.empty(int(size_bytes), dtype=torch.uint8, device=device) - self.token_counts = torch.empty(self.num_local_experts, dtype=torch.int32, device=device) + self.tokens_per_expert = torch.empty( + self.num_local_experts, dtype=torch.int64, device=device + ) # Persistent tensor; keep resident if activation CPU offloading is on. mark_not_offload(self.handle_mem) + # Per-step recv-token total (device int64 [1]), written by ep_prepare. + # The true pre-drop count: under drop_on_overflow it includes dropped + # tokens and may exceed recv_capacity_per_rank. Eager mode syncs it to + # size the recv outputs; graph mode reads it after replay to check overflow. + self.total_recv_tokens = torch.empty(1, dtype=torch.int64, device=device) + mark_not_offload(self.total_recv_tokens) + # Host mirror of total_recv_tokens; set by ep_prepare in eager mode (one + # D2H) to size the recv outputs. None otherwise. + self._host_total_recv_tokens: Optional[int] = None # torch.library custom ops (so they don't graph-break under torch.compile) @@ -241,17 +285,18 @@ def __init__( @torch.library.custom_op( f"{_LIB}::prepare", - mutates_args=("handle_mem", "token_counts"), + mutates_args=("handle_mem", "tokens_per_expert", "total_recv_tokens"), device_types="cuda", ) def _prepare_op( handle_mem: torch.Tensor, top_k: int, topk_idx: torch.Tensor, - token_counts: torch.Tensor, + tokens_per_expert: torch.Tensor, alignment: int, + total_recv_tokens: torch.Tensor, ) -> None: - tex.ep_prepare(handle_mem, topk_idx, token_counts, top_k, alignment) + tex.ep_prepare(handle_mem, topk_idx, tokens_per_expert, top_k, alignment, total_recv_tokens) @_prepare_op.register_fake @@ -341,13 +386,26 @@ def _(*_args, **_kw): def ep_prepare(buffer: "EpBuffer", topk_idx: torch.Tensor) -> torch.Tensor: """AllGather the routing map; fills ``buffer.handle_mem`` and returns - ``buffer.token_counts`` (int32, shape [num_local_experts]). topk_idx must + ``buffer.tokens_per_expert`` (int64, shape [num_local_experts]). topk_idx must be int32 or int64. + + Always fills ``buffer.total_recv_tokens`` (device int64 [1]) with the + per-step recv-token total. In eager mode its host value is cached in + ``buffer._host_total_recv_tokens`` (one D2H) to size the recv outputs; + otherwise it stays device-side for the caller to read after a graph replay + and compare against ``recv_capacity_per_rank`` to detect overflow. """ torch.ops.transformer_engine_ep.prepare( - buffer.handle_mem, buffer.top_k, topk_idx, buffer.token_counts, buffer.alignment + buffer.handle_mem, + buffer.top_k, + topk_idx, + buffer.tokens_per_expert, + buffer.alignment, + buffer.total_recv_tokens, ) - return buffer.token_counts + if buffer.eager: + buffer._host_total_recv_tokens = int(buffer.total_recv_tokens.item()) + return buffer.tokens_per_expert def _ep_dispatch_raw( @@ -383,15 +441,20 @@ def forward( # type: ignore[override] alignment: int, recv_tokens: torch.Tensor, recv_topk_weights: torch.Tensor, - token_counts: torch.Tensor, + tokens_per_expert: torch.Tensor, + total_recv_tokens: torch.Tensor, topk_idx: torch.Tensor, tokens: torch.Tensor, topk_weights: torch.Tensor, + skip_prepare: bool = False, ): - """Prepare + dispatch fwd.""" - torch.ops.transformer_engine_ep.prepare( - handle_mem, top_k, topk_idx, token_counts, alignment - ) + """Prepare + dispatch fwd. In eager mode the caller runs prepare first + (to size recv outputs from the recv-token total), so ``skip_prepare`` + avoids re-running the routing AllGather here.""" + if not skip_prepare: + torch.ops.transformer_engine_ep.prepare( + handle_mem, top_k, topk_idx, tokens_per_expert, alignment, total_recv_tokens + ) torch.ops.transformer_engine_ep.dispatch( handle_mem, topk_idx, @@ -409,13 +472,13 @@ def forward( # type: ignore[override] ctx.top_k = topk_weights.shape[-1] ctx.recv_capacity = recv_tokens.shape[0] ctx.hidden_dim = tokens.shape[-1] - ctx.mark_non_differentiable(token_counts) + ctx.mark_non_differentiable(tokens_per_expert) # Detach so the long-lived buffers aren't tracked as differentiable outputs; # autograd re-attaches grad_fn pointing back at this Function. - return recv_tokens.detach(), recv_topk_weights.detach(), token_counts + return recv_tokens.detach(), recv_topk_weights.detach(), tokens_per_expert @staticmethod - def backward(ctx, g_recv_tokens, g_recv_topk_weights, _g_token_counts): # type: ignore[override] + def backward(ctx, g_recv_tokens, g_recv_topk_weights, _g_tokens_per_expert): # type: ignore[override] """Dispatch bwd; normalizes grad-input layout, otherwise passes through.""" (handle_mem,) = ctx.saved_tensors device = handle_mem.device @@ -440,10 +503,12 @@ def backward(ctx, g_recv_tokens, g_recv_topk_weights, _g_token_counts): # type: None, # alignment None, # recv_tokens None, # recv_topk_weights - None, # token_counts + None, # tokens_per_expert + None, # total_recv_tokens None, # topk_idx grad_tokens.view(ctx.tokens_shape), grad_topk_weights.view(ctx.topk_weights_shape), + None, # skip_prepare ) @@ -539,8 +604,10 @@ def ep_dispatch( ``recv_tokens`` / ``recv_topk_weights`` are the dispatch recv outputs: pass caller-owned buffers (mcore-managed mode; in zero-copy they must be symm-mem-backed) or leave them None to allocate on - the fly (zero-copy: symm-mem pool; normal: plain). Returns (recv_tokens, recv_topk_weights, - token_counts); token_counts is non-diff. + the fly (zero-copy: symm-mem pool; normal: plain). In eager mode the recv outputs are sized to + this step's recv-token total and must not be caller-supplied. Returns (recv_tokens, + recv_topk_weights, tokens_per_expert); tokens_per_expert is non-diff. See ``buffer.total_recv_tokens`` for + the per-step recv total. """ _require_bf16("tokens", tokens) if topk_weights.dtype is not torch.float32: @@ -548,27 +615,41 @@ def ep_dispatch( f"topk_weights must be float32; got dtype={topk_weights.dtype}. " "Cast with topk_weights.float() before calling." ) + skip_prepare = False + if buffer.eager: + if recv_tokens is not None or recv_topk_weights is not None: + raise ValueError( + "eager mode sizes the recv outputs from the per-step recv-token total " + "and cannot use caller-supplied recv_tokens / recv_topk_weights" + ) + # Prepare first to learn this step's recv total, then size the recv + # outputs to it; skip re-running prepare inside the autograd forward. + ep_prepare(buffer, topk_idx) + rows = buffer._host_total_recv_tokens + skip_prepare = True + else: + rows = buffer.recv_capacity_per_rank if recv_tokens is None: recv_tokens = _alloc_io( - (buffer.recv_capacity_per_rank, buffer.hidden_dim), + (rows, buffer.hidden_dim), buffer.payload_dtype, buffer.device, buffer.zero_copy, ) if recv_topk_weights is None: - recv_topk_weights = _alloc_io( - (buffer.recv_capacity_per_rank,), torch.float32, buffer.device, buffer.zero_copy - ) + recv_topk_weights = _alloc_io((rows,), torch.float32, buffer.device, buffer.zero_copy) return _EpDispatch.apply( buffer.handle_mem, buffer.top_k, buffer.alignment, recv_tokens, recv_topk_weights, - buffer.token_counts, + buffer.tokens_per_expert, + buffer.total_recv_tokens, topk_idx, tokens, topk_weights, + skip_prepare, ) @@ -584,10 +665,15 @@ def ep_combine( ``expert_out`` is the combine input (always caller-supplied; in zero-copy it must be symm-mem- backed). ``grad_out`` is the backward's grad target: pass a caller-owned buffer (mcore-managed mode) or leave it None to allocate on the fly in the backward (zero-copy: symm-mem pool; normal: - plain). Result shape is (num_local_tokens, buffer.hidden_dim); defaults to - buffer.max_tokens_per_rank rows. + plain). In eager mode the grad target is sized per step and must not be caller-supplied. Result + shape is (num_local_tokens, buffer.hidden_dim); defaults to buffer.max_tokens_per_rank rows. """ _require_bf16("expert_out", expert_out) + if buffer.eager and grad_out is not None: + raise ValueError( + "eager mode sizes the combine grad target per step and cannot use a " + "caller-supplied grad_out" + ) if num_local_tokens is None: num_local_tokens = buffer.max_tokens_per_rank return _EpCombine.apply(