diff --git a/examples/pytorch/ep/bench/ep_bench.py b/examples/pytorch/ep/bench/ep_bench.py index 2b7a2c62e5..f80a57015c 100644 --- a/examples/pytorch/ep/bench/ep_bench.py +++ b/examples/pytorch/ep/bench/ep_bench.py @@ -181,8 +181,9 @@ def main(): ep_group, num_experts=E, max_tokens_per_rank=T, - recv_capacity_per_rank=recv_pr, hidden_dim=H, + num_topk=K, + recv_capacity_per_rank=recv_pr, max_num_sms=args.max_num_sms, ) @@ -207,8 +208,6 @@ def main(): recv_capacity_per_rank=recv_pr, hidden_dim=H, num_local_experts=num_local_experts, - dispatch_recv_tokens=caller_recv_tokens, - combine_grad_expert_out=caller_grad_expert_out, ) tokens = tokens_hbm @@ -235,8 +234,13 @@ def main(): eo_p = recv_tokens.detach().clone().requires_grad_(True) # Stand-in callables; the cuda-graph branch below swaps in graphed versions. - fwd_bwd_dispatch_fn = lambda x: ep_dispatch(buffer, x, topk_idx, topk_w)[0] # noqa: E731 - fwd_bwd_combine_fn = lambda expert_out: ep_combine(buffer, expert_out) # noqa: E731 + # caller_* are None unless the caller opted in, matching ep_dispatch/ep_combine defaults. + fwd_bwd_dispatch_fn = lambda x: ep_dispatch( # noqa: E731 + buffer, x, topk_idx, topk_w, recv_tokens=caller_recv_tokens + )[0] + fwd_bwd_combine_fn = lambda expert_out: ep_combine( # noqa: E731 + buffer, expert_out, grad_out=caller_grad_expert_out + ) def _dispatch_raw(): _ep_dispatch_raw(buffer, topk_idx, tokens, topk_w, recv_tokens, recv_w) @@ -246,7 +250,7 @@ def _combine_raw(): _ep_combine_raw(buffer, expert_out, out_buf) def _ep_dispatch_fwd(): - ep_dispatch(buffer, tokens.detach(), topk_idx, topk_w) + ep_dispatch(buffer, tokens.detach(), topk_idx, topk_w, recv_tokens=caller_recv_tokens) def _ep_dispatch_fwd_bwd(): tokens_p.grad = None @@ -288,11 +292,11 @@ def _ep_combine_fwd_bwd(): # Graph fwd+bwd of the autograd-wrapped ops via make_graphed_callables. class _DispatchMod(torch.nn.Module): def forward(self, x): - return ep_dispatch(buffer, x, topk_idx, topk_w)[0] + return ep_dispatch(buffer, x, topk_idx, topk_w, recv_tokens=caller_recv_tokens)[0] class _CombineMod(torch.nn.Module): def forward(self, expert_out): - return ep_combine(buffer, expert_out) + return ep_combine(buffer, expert_out, grad_out=caller_grad_expert_out) disp_mod = _DispatchMod().cuda() comb_mod = _CombineMod().cuda() diff --git a/examples/pytorch/ep/ep_moe.py b/examples/pytorch/ep/ep_moe.py index 149185f251..b47c334225 100644 --- a/examples/pytorch/ep/ep_moe.py +++ b/examples/pytorch/ep/ep_moe.py @@ -131,8 +131,9 @@ def main(): ep_group, num_experts=num_experts, max_tokens_per_rank=T, - recv_capacity_per_rank=recv_pr, hidden_dim=args.hidden, + num_topk=args.top_k, + recv_capacity_per_rank=recv_pr, ) try: _run_layer( @@ -182,15 +183,13 @@ def _run_layer(args, rank, world_size, ep_size, num_experts, num_local_experts, recv_capacity_per_rank=recv_pr, hidden_dim=args.hidden, num_local_experts=num_local_experts, - dispatch_recv_tokens=recv_tokens, - combine_grad_expert_out=grad_expert_out, ) - recv_t, recv_w_out, _tc = ep_dispatch(buffer, tokens, topk_idx, topk_w) + recv_t, recv_w_out, _tc = ep_dispatch(buffer, tokens, topk_idx, topk_w, recv_tokens=recv_tokens) expert_out = _batched_expert_linear(recv_t, kernels_local, num_local_experts) # Apply per-slot topk weighting before combine. expert_out = expert_out * recv_w_out.unsqueeze(-1).to(expert_out.dtype) - out = ep_combine(buffer, expert_out) + out = ep_combine(buffer, expert_out, grad_out=grad_expert_out) loss = 0.5 * (out.float() ** 2).sum() loss.backward() diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 534fd9642a..4778498b7d 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 @@ -144,9 +162,12 @@ def setUpClass(cls): cls.ep_group, num_experts=cls.cfg.num_experts, max_tokens_per_rank=TOKENS_PER_RANK, - recv_capacity_per_rank=cls.cfg.recv_capacity_per_rank, hidden_dim=HIDDEN_DIM, + num_topk=TOP_K, + # Omit recv_capacity_per_rank to select eager mode. + recv_capacity_per_rank=None if EAGER else cls.cfg.recv_capacity_per_rank, zero_copy=ZERO_COPY, + drop_on_overflow=OVERFLOW, ) def setUp(self): @@ -155,14 +176,22 @@ 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( top_k=top_k, max_tokens_per_rank=TOKENS_PER_RANK, - recv_capacity_per_rank=self.cfg.recv_capacity_per_rank, hidden_dim=HIDDEN_DIM, num_local_experts=NUM_LOCAL_EXPERTS, + recv_capacity_per_rank=None if EAGER else self.cfg.recv_capacity_per_rank, alignment=alignment, ) @@ -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..dd44773514 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,13 @@ 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; + // Per-token top-k; NCCL EP sizes internal buffers from it in eager mode. + 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 +331,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 +340,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..a4800944b5 100644 --- a/transformer_engine/common/include/transformer_engine/ep.h +++ b/transformer_engine/common/include/transformer_engine/ep.h @@ -44,7 +44,8 @@ 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 caller sizes recv buffers to the per-routing recv count. */ int max_recv_tokens_per_rank; /*! Token hidden dimension. */ int hidden_dim; @@ -58,6 +59,12 @@ typedef struct { * by NVTECommWindow handles and transfer in place (no staging copies); * 0 (default) = staged. */ int zero_copy; + /*! Per-token top-k; sizes NCCL EP internal buffers. Required in eager mode + * (max_recv_tokens_per_rank == 0); 0 = unset. */ + int num_topk; + /*! Recv overflow policy. Nonzero drops tokens past max_recv_tokens_per_rank + * and continues; 0 (default) traps. Not supported in eager mode. */ + int drop_on_overflow; } NVTEEpGroupConfig; /*! \brief Per-layer configuration consumed by nvte_ep_handle_mem_size and @@ -121,13 +128,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..95f02c64f0 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 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..9bf821721f 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 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(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("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..940812b6a4 100644 --- a/transformer_engine/pytorch/ep.py +++ b/transformer_engine/pytorch/ep.py @@ -67,14 +67,17 @@ 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 (set when recv_capacity_per_rank is +# omitted); ep_dispatch reads it to size the recv outputs from the per-step +# recv-token total instead of a fixed 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,16 +88,19 @@ def _atexit_finalize() -> None: finally: _BOOTSTRAPPED = False _EP_GROUP = None + _EAGER = False def ep_bootstrap( ep_group: dist.ProcessGroup, num_experts: int, max_tokens_per_rank: int, - recv_capacity_per_rank: int, hidden_dim: int, + num_topk: int, + recv_capacity_per_rank: Optional[int] = None, max_num_sms: int = 0, zero_copy: bool = False, + 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. @@ -102,15 +108,32 @@ def ep_bootstrap( max_token_dtype sets the widest token dtype this EP group will dispatch; it sizes NCCL EP staging buffers. + ``recv_capacity_per_rank`` bounds the tokens one rank receives per step and + sizes the recv outputs. Omit it (``None``) for eager mode, which sizes recv + outputs from the per-step recv total instead; eager needs a host sync each + step and is not CUDA-graph capturable. + ``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``. + Requires ``recv_capacity_per_rank``. + + ``num_topk`` is the per-token top-k; it sizes NCCL EP internal buffers. + + ``drop_on_overflow`` drops tokens exceeding ``recv_capacity_per_rank`` instead + of trapping. Requires ``recv_capacity_per_rank``. """ - global _BOOTSTRAPPED, _ATEXIT_REGISTERED, _EP_GROUP + global _BOOTSTRAPPED, _ATEXIT_REGISTERED, _EP_GROUP, _EAGER + eager = recv_capacity_per_rank is None 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 num_topk < 1: + raise ValueError(f"ep_bootstrap requires num_topk >= 1 (got {num_topk}).") + if zero_copy and eager: + raise ValueError("ep_bootstrap: zero_copy requires recv_capacity_per_rank") + if drop_on_overflow and eager: + raise ValueError("ep_bootstrap: drop_on_overflow requires recv_capacity_per_rank") _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 (recv_capacity_per_rank=None) sizes recv buffers per routing, + # so the group uses the library-derived bound (0 = NCCL_EP_AUTO). + int(recv_capacity_per_rank or 0), int(hidden_dim), int(max_num_sms), max_token_dtype, bool(zero_copy), + int(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,30 +226,43 @@ 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__( self, top_k: int, max_tokens_per_rank: int, - recv_capacity_per_rank: int, hidden_dim: int, num_local_experts: int, + recv_capacity_per_rank: Optional[int] = None, alignment: int = 0, payload_dtype: torch.dtype = torch.bfloat16, device: Optional[torch.device] = None, ) -> None: + if not _BOOTSTRAPPED: + raise RuntimeError("EpBuffer requires ep_bootstrap() to be called first.") if device is None: device = torch.device("cuda", torch.cuda.current_device()) alignment = int(alignment) if alignment > 1 and (alignment & (alignment - 1)) != 0: raise ValueError(f"alignment must be 0, 1, or a power of two (got {alignment}).") + self.eager = _EAGER + if not self.eager and recv_capacity_per_rank is None: + raise ValueError( + "EpBuffer requires recv_capacity_per_rank unless the EP group was " + "bootstrapped in eager mode (recv_capacity_per_rank omitted)." + ) self.top_k = int(top_k) self.alignment = alignment self.max_tokens_per_rank = int(max_tokens_per_rank) - self.recv_capacity_per_rank = int(recv_capacity_per_rank) + self.recv_capacity_per_rank = ( + None if recv_capacity_per_rank is None else int(recv_capacity_per_rank) + ) self.hidden_dim = int(hidden_dim) self.num_local_experts = int(num_local_experts) self.payload_dtype = payload_dtype @@ -229,9 +271,18 @@ def __init__( 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. + # Eager mode uses it to size the recv outputs; graph mode reads it after + # replay to detect overflow past recv_capacity_per_rank. + 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. + self._host_total_recv_tokens: Optional[int] = None # torch.library custom ops (so they don't graph-break under torch.compile) @@ -241,17 +292,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 +393,24 @@ 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. + + Also fills ``buffer.total_recv_tokens`` (device int64 [1]) with the per-step + recv total; eager mode mirrors it to the host to size the recv outputs, graph + mode reads it device-side 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( @@ -373,25 +436,19 @@ def _ep_combine_raw(buffer: "EpBuffer", expert_out: torch.Tensor, result: torch. class _EpDispatch(torch.autograd.Function): - """Autograd prepare+dispatch; bwd uses user-supplied grad inputs as-is.""" + """Autograd dispatch; caller runs prepare first. bwd uses user-supplied grad inputs as-is.""" @staticmethod def forward( # type: ignore[override] ctx, handle_mem: torch.Tensor, - top_k: int, - alignment: int, recv_tokens: torch.Tensor, recv_topk_weights: torch.Tensor, - token_counts: torch.Tensor, topk_idx: torch.Tensor, tokens: torch.Tensor, topk_weights: torch.Tensor, ): - """Prepare + dispatch fwd.""" - torch.ops.transformer_engine_ep.prepare( - handle_mem, top_k, topk_idx, token_counts, alignment - ) + """Dispatch fwd; prepare must have run into ``handle_mem`` beforehand.""" torch.ops.transformer_engine_ep.dispatch( handle_mem, topk_idx, @@ -407,15 +464,13 @@ def forward( # type: ignore[override] ctx.tokens_T_flat = tokens.numel() // tokens.shape[-1] ctx.topk_T_flat = topk_weights.numel() // topk_weights.shape[-1] 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) # 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() @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): # type: ignore[override] """Dispatch bwd; normalizes grad-input layout, otherwise passes through.""" (handle_mem,) = ctx.saved_tensors device = handle_mem.device @@ -436,11 +491,8 @@ def backward(ctx, g_recv_tokens, g_recv_topk_weights, _g_token_counts): # type: ) return ( None, # handle_mem - None, # top_k - None, # alignment None, # recv_tokens None, # recv_topk_weights - None, # token_counts None, # topk_idx grad_tokens.view(ctx.tokens_shape), grad_topk_weights.view(ctx.topk_weights_shape), @@ -539,8 +591,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,28 +602,33 @@ def ep_dispatch( f"topk_weights must be float32; got dtype={topk_weights.dtype}. " "Cast with topk_weights.float() before calling." ) + if buffer.eager and (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 (routing AllGather) up front so the recv outputs can be sized; in + # eager mode ep_prepare also host-syncs this step's recv-token total. + tokens_per_expert = ep_prepare(buffer, topk_idx) + rows = buffer._host_total_recv_tokens if buffer.eager else 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 - ) - return _EpDispatch.apply( + recv_topk_weights = _alloc_io((rows,), torch.float32, buffer.device, buffer.zero_copy) + recv_tokens, recv_topk_weights = _EpDispatch.apply( buffer.handle_mem, - buffer.top_k, - buffer.alignment, recv_tokens, recv_topk_weights, - buffer.token_counts, topk_idx, tokens, topk_weights, ) + return recv_tokens, recv_topk_weights, tokens_per_expert def ep_combine( @@ -584,10 +643,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(