DeepEP internode dispatch times out with NVSHMEM 3.6.5; backport DeepEP #564 into the Automodel image

[janbernloehr]

Describe the bug

The DeepSeek-V3.2 LoRA benchmark fails during the first MoE dispatch of iteration 0 when using the shipped deepep dispatcher on 16 nodes with 8 H100 GPUs per node:

RuntimeError: DeepEP error: timeout (dispatch CPU)

A controlled experiment identifies an incompatibility between the DeepEP source revision used by Automodel's container build and the NVSHMEM version installed alongside it:

• DeepEP: 42144303752422ade37f24bca9e2dde12df70e09, reported by the installed package as deep_ep 1.2.1+4214430.
• NVSHMEM: nvidia-nvshmem-cu13==3.6.5; the loaded runtime prints NVSHMEM v3.6.5.
• At this DeepEP revision, ibgda_get_rc() still indexes the default RC queue-pair array using the layout from before NVSHMEM 3.5.
• Rebuilding exactly that DeepEP revision with the compatibility changes from deepseek-ai/DeepEP#564, while keeping the dispatcher, NVSHMEM version, model configuration, and launch environment unchanged, fixes the failure.

The unpatched control reproduced the iteration-zero timeout. The patched run completed all 30 configured iterations and exited successfully on every node. This is a dependency-build fix: changing the benchmark to hybridep is an available workaround, but is not necessary once the DeepEP compatibility patch is applied.

Relevant public source:

• Automodel Dockerfile at the recorded failing source revision.
• The DeepEP patch applied by that Dockerfile. It adjusts a CUDA include path and a Python dependency; it does not contain nano-v2-9b& 12b #564's RC queue-pair compatibility changes.
• The affected benchmark configuration.
• DeepEP's unpatched queue-pair selection.

Steps/Code to reproduce bug

1. Use the affected dependency combination

Use an Automodel environment with the source and dependencies above, including the example configuration from Automodel commit e4f0cc9cd01f0f8ba51184dc3ec09578309deda6. The environment must contain the compiled DeepEP extension with its NVSHMEM internode support enabled.

The public Dockerfile at that revision documents the affected DeepEP/NVSHMEM pairing. A floating image tag or a newer dependency checkout may select a different combination; record the exact source revisions and loaded libraries when reproducing.

Captured environment details:

Component	Value
Hardware	16 nodes × 8 NVIDIA H100 80 GB SXM GPUs; 128 GPUs total
Interconnect	InfiniBand between nodes; NVLink within each node
Architecture	Linux x86-64
Python	3.12
PyTorch	2.14.0a0 development build
CUDA package family	CUDA 13
DeepEP source	42144303752422ade37f24bca9e2dde12df70e09
Installed DeepEP distribution	deep_ep 1.2.1+4214430
NVSHMEM Python distribution	nvidia-nvshmem-cu13 3.6.5
NVSHMEM loaded runtime	NVSHMEM v3.6.5, printed during initialization
NCCL	2.30.7
Automodel entry point	nemo_automodel.recipes.llm.benchmark
Recipe class	BenchmarkingRecipeForNextTokenPrediction

This is a summary of captured evidence, not a complete torch.utils.collect_env report. No standalone environment-collection report was captured.

The installed package version and the library actually loaded are the relevant evidence. A container-level NVSHMEM_VERSION environment variable may describe a different preinstalled copy and should not be used alone to identify the runtime.

2. Keep the shipped benchmark settings

Use examples/llm_benchmark/deepseek/dsv32_lora.yaml without a dispatcher override. The settings relevant to this reproduction are:

Setting	Value
Model	deepseek-ai/DeepSeek-V3.2
Model initialization	NeMoAutoModelForCausalLM.from_config
Precision	BF16
Seed	1234
Distributed strategy	fsdp2
Tensor / context / pipeline / expert parallel sizes	1 / 1 / 4 / 32
Sequence parallelism	Disabled
Pipeline schedule	interleaved1f1b
Pipeline microbatch size	1
Layers per virtual stage	2
Activation checkpointing	Enabled
Expert implementation	gmm
Dispatcher	deepep
Attention / linear / RMSNorm backends	te / te / te
Fake balanced gate	Enabled
LoRA	Rank 32, alpha 32, dropout 0, all target modules, moe_rank_scaling: true
Dataset	MockIterableDataset
Sequence length	4096
Local / global batch size	4 / 512
Total iterations	30
Warmup / measured iterations	10 / 20
Process-group timeout	10 minutes
Checkpoint saving	Disabled

The workload uses synthetic data and initializes the model from its configuration. Make the referenced Hugging Face model configuration/tokenizer available through the normal cache or download mechanism.

3. Launch one process launcher per node

The following is an equivalent portable torchrun launch. Execute it once on each of the 16 allocated nodes from the Automodel checkout, using the same environment on all nodes.

Set NODE_RANK to that node's index, 0 through 15. Set MASTER_ADDR to a reachable address on the first node, MASTER_PORT to a free rendezvous port, and RUN_ID to the same unique value on every node.

torchrun \
  --nnodes=16 \
  --nproc-per-node=8 \
  --node-rank="${NODE_RANK:?Set NODE_RANK to 0 through 15}" \
  --rdzv-backend=c10d \
  --rdzv-endpoint="${MASTER_ADDR:?Set MASTER_ADDR}:${MASTER_PORT:?Set MASTER_PORT}" \
  --rdzv-id="${RUN_ID:?Set the same RUN_ID on all nodes}" \
  -m nemo_automodel.recipes.llm.benchmark \
  --config examples/llm_benchmark/deepseek/dsv32_lora.yaml \
  --benchmark.json_output_path /tmp/dsv32-lora-benchmark-summary.json

The model/communication setup completes, and the benchmark logs iteration 0. The first MoE dispatch then stalls and fails before iteration 1.

The 128-GPU benchmark is the tested reproducer. A smaller standalone DeepEP reproducer has not been validated for this report.

4. Observe the initial error

A shortened call chain, with package-relative paths, is:

nemo_automodel/recipes/llm/benchmark.py: run_benchmark
  -> train_ft.py: _forward_backward_step
  -> torch.distributed.pipelining: schedule / stage forward
  -> MoE / LoRA expert token dispatch
  -> nemo_automodel/components/moe/megatron/token_dispatcher.py: dispatch_all_to_all
  -> token_dispatcher.py: dispatch / fused dispatch
  -> deep_ep/buffer.py:379: dispatch
  -> deep_ep/buffer.py:584: internode_dispatch
RuntimeError: DeepEP error: timeout (dispatch CPU)

Immediately before the timeout, the receive-count diagnostics remain at their uninitialized values:

num_recv_tokens: -1
num_rdma_recv_tokens: -1
moe_recv_expert_counter[0]: -1
...

An illegal CUDA memory access subsequently appears during error handling/teardown. The first observed failure is the DeepEP dispatch timeout.

Expected behavior

With model.backend.dispatcher: deepep, the benchmark should complete the configured 30 iterations, write benchmark_summary.json, and exit successfully on every node.

The container build should produce a DeepEP extension compatible with the NVSHMEM device-state layout it uses. Selecting the shipped deepep dispatcher should not require changing the benchmark to a different transport backend.

Additional context

Why this particular compatibility patch matters

DeepEP contains its own IBGDA device-side queue-pair access code. Matching the NVSHMEM host and device library versions alone is insufficient if that code still assumes an older internal layout.

At the affected DeepEP revision, the RC lookup uses:

return &state->globalmem.rcs[
    pe * num_rc_per_pe * state->num_devices_initialized +
    id % (num_rc_per_pe * state->num_devices_initialized)];

NVSHMEM 3.5 changed how the default queue pairs are stored. For the newer device-state layout, #564 uses:

return &state->globalmem.rcs[
    pe + nvshmemi_device_state_d.npes * id];

The patch selects the indexing scheme at compile time according to the NVSHMEM device-state type and retains support for the older layout. It also ensures the appropriate extern declarations for NVSHMEM device symbols through the __CUDACC_RDC__ guard.

DeepEP#564 was merged upstream on January 20, 2026, but these changes are absent from the specific revision pinned by the affected Automodel build.

The incorrect queue-pair indexing is consistent with the observed lack of progress. DeepEP's CPU-side dispatch wait polls the receive counts until they become nonnegative. Here they remain -1, and DeepEP's own timeout expires. Increasing a scheduler wall-time limit does not repair the queue-pair selection.

Controlled validation

Both arms rebuilt DeepEP from the same source revision in the same pinned base environment. The only intentional difference was whether the #564 compatibility changes were applied before compiling the extension.

The following were held constant:

• DeepEP source revision and installed NVSHMEM 3.6.5.
• PyTorch and other dependencies.
• Original benchmark configuration, including dispatcher: deepep.
• 16 × 8 H100 GPUs, PP4, EP32, BF16, sequence length 4096, and local batch size 4.
• Launch configuration and network settings.
• The selected eight InfiniBand HCAs.
• NVSHMEM_IB_SL unset in both arms; no routing override was introduced.

Both builds checked that deep_ep_cpp still links against libnvshmem. The passing result was not obtained by compiling out NVSHMEM support.

Arm	DeepEP build	Outcome
Control	Recorded source revision, rebuilt without #564	Iteration-zero timeout (dispatch CPU); nonzero exit after approximately 204.5 seconds
Treatment	Same source revision, rebuilt with #564's compatibility changes	All 30 iterations completed; exit 0 from every one of the 16 node-level launchers; approximately 1241.2 seconds total

Patched-run observations:

Measurement	Result
Completed iterations	30 / 30
Warmup / measured iterations	10 / 20
Reported step time	37.114 seconds
Final reported loss	12.2530
Reported peak allocated memory	46.48 GB
Reported MFU	15.5268%, with the config's peak of 989 TFLOP/s
Setup time	35.43 seconds
Warmup time	414.97 seconds
Measured training time	742.27 seconds
Benchmark summary	Written successfully
Node-level exit status	0 on all 16 nodes

Every node's patched-run log was checked for the dispatch timeout, DeepEP errors, stuck receive-counter diagnostics, and illegal-memory-access messages; none were found.

The paired experiment demonstrates that the compatibility backport is sufficient for this workload with the existing network configuration. It does not require enabling adaptive routing, changing HCA selection, or increasing a timeout.

Proposed fix and PR scope

Apply #564's compatibility changes as part of Automodel's DeepEP build, for example by extending docker/common/deepep.patch or adding a dedicated compatibility patch applied by docker/Dockerfile.

Keeping the current DeepEP revision and NVSHMEM version while adding this backport is the smallest change supported by the experiment. Moving to another DeepEP revision that includes the fix is an alternative, but would introduce additional source changes that were not tested here.

Suggested acceptance criteria:

1. The DeepEP extension is rebuilt with the corrected queue-pair indexing and device-symbol declaration handling.
2. NVSHMEM internode support remains compiled in, and the intended NVSHMEM runtime is loaded.
3. The original dsv32_lora.yaml completes all 30 steps with dispatcher: deepep, publishes its summary, and tears down cleanly.
4. Existing HybridEP support remains available.
5. Broader backend and platform coverage follows the repository's normal validation; this report does not claim that coverage.

The tested compatibility backport is included below. It applies to the DeepEP source tree, not directly to the Automodel repository.

DeepEP compatibility backport used by the passing treatment

diff --git a/csrc/kernels/configs.cuh b/csrc/kernels/configs.cuh
index 61ec54d..0e2091c 100644
--- a/csrc/kernels/configs.cuh
+++ b/csrc/kernels/configs.cuh
@@ -26,6 +26,13 @@
 #define __CUDACC_RDC__ // NOLINT(*-reserved-identifier)
 #endif
 
+// Define __CUDACC_RDC__ to ensure proper extern declarations for NVSHMEM device symbols
+#ifndef DISABLE_NVSHMEM
+#ifndef __CUDACC_RDC__
+#define __CUDACC_RDC__ // NOLINT(*-reserved-identifier)
+#endif
+#endif
+
 // Remove Torch restrictions
 #ifdef __CUDA_NO_HALF_CONVERSIONS__
 #undef __CUDA_NO_HALF_CONVERSIONS__
diff --git a/csrc/kernels/ibgda_device.cuh b/csrc/kernels/ibgda_device.cuh
index 7179a50..54af86e 100644
--- a/csrc/kernels/ibgda_device.cuh
+++ b/csrc/kernels/ibgda_device.cuh
@@ -7,6 +7,8 @@
 //  - nvshmem/src/include/non_abi/device/pt-to-pt/ibgda_device.cuh
 #pragma once
 
+#include <type_traits>
+
 #include "configs.cuh"
 #include "exception.cuh"
 #include "utils.cuh"
@@ -73,11 +75,25 @@ nvshmemi_ibgda_device_state_t* ibgda_get_state() {
     return &nvshmemi_ibgda_device_state_d;
 }
 
+// NVSHMEM 3.5.19 changed the RC QP layout. Select the matching indexing
+// scheme at compile time while retaining compatibility with the v1 layout.
+template <typename StateType>
+__device__ static __forceinline__
+nvshmemi_ibgda_device_qp_t* ibgda_get_rc_impl(StateType* state, int pe, int id) {
+    const auto num_rc_per_pe = state->num_rc_per_pe;
+    if constexpr (std::is_same_v<StateType, nvshmemi_ibgda_device_state_v1>) {
+        return &state->globalmem.rcs[
+            pe * num_rc_per_pe * state->num_devices_initialized +
+            id % (num_rc_per_pe * state->num_devices_initialized)];
+    } else {
+        return &state->globalmem.rcs[pe + nvshmemi_device_state_d.npes * id];
+    }
+}
+
 __device__ static __forceinline__
 nvshmemi_ibgda_device_qp_t* ibgda_get_rc(int pe, int id) {
     auto state = ibgda_get_state();
-    const auto num_rc_per_pe = ibgda_get_state()->num_rc_per_pe;
-    return &state->globalmem.rcs[pe * num_rc_per_pe * state->num_devices_initialized + id % (num_rc_per_pe * state->num_devices_initialized)];
+    return ibgda_get_rc_impl(state, pe, id);
 }
 
 __device__ static __forceinline__

Workaround

Using --model.backend.dispatcher hybridep also completed this benchmark: 30/30 steps, approximately 36.462 seconds per step, and final reported loss 12.2530.

That is a useful workaround, but it exercises a different backend. The proposed dependency fix allows the existing deepep configuration to run successfully.

The patched DeepEP run and the HybridEP run reported the same final loss at the logged precision. This is a useful consistency observation, not proof of bitwise equivalence, convergence, or correctness across other inputs.

Related public reports

• DeepEP#564: the merged NVSHMEM compatibility fix validated here.
• DeepEP#553: the NVSHMEM 3.5.19 timeout referenced by nano-v2-9b& 12b #564.
• Automodel#2614: aligned NVSHMEM versions in the build, but also documented that conventional DeepEP still hung after version alignment and switched the affected example to HybridEP. The present result identifies a specific missing source compatibility change that version alignment alone does not supply.
• Automodel#3906 and #3910: a similar timeout in Step-3.5-Flash LoRA was addressed by changing dispatcher and parallelism settings. That workload has not been tested with this backport, so a shared root cause remains a hypothesis.
• Automodel#2327: a related DeepSeek-V3 pretraining dispatch timeout with a different parallel configuration. The matching error string alone does not establish the same root cause.