Release test failing with DeepEP error

Release tests failing with DeepEP error off the main branch:

- llm_grpo_qwen3_5_35ba3b_dapo_4n8g_automodel
- llm_grpo_glm47_flash_4n8g_automodel

Error:

```
File "/opt/nemo-rl/3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/models/glm4_moe_lite/model.py", line 180, in forward
h = layer(
x=h,
...<3 lines>...
**attn_kwargs,
)
File "/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1885, in _call_impl
return inner()
File "/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1833, in inner
result = forward_call(*args, **kwargs)
File "/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2/lib/python3.13/site-packages/torch/distributed/algorithms/_checkpoint/checkpoint_wrapper.py", line 169, in forward
return self.checkpoint_fn( # type: ignore[misc]
~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^
self._checkpoint_wrapped_module, *args, **kwargs
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
)
^
File "/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2/lib/python3.13/site-packages/torch/_compile.py", line 54, in inner
return disable_fn(*args, **kwargs)
File "/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2/lib/python3.13/site-packages/torch/_dynamo/eval_frame.py", line 1263, in _fn
return fn(*args, **kwargs)
File "/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2/lib/python3.13/site-packages/torch/utils/checkpoint.py", line 512, in checkpoint
ret = function(*args, **kwargs)
File "/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
return forward_call(*args, **kwargs)
File "/opt/nemo-rl/3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/models/glm4_moe_lite/model.py", line 85, in forward
mlp_out = self._mlp(x=self.post_attention_layernorm(x), padding_mask=padding_mask)
File "/opt/nemo-rl/3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/models/glm4_moe_lite/model.py", line 94, in _mlp
return self.mlp(x, padding_mask)
~~~~~~~~^^^^^^^^^^^^^^^^^
File "/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
return forward_call(*args, **kwargs)
File "/opt/nemo-rl/3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/moe/layers.py", line 679, in forward
y = self.experts(x_latent, token_mask, weights, indices)
File "/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1790, in _call_impl
return forward_call(*args, **kwargs)
File "/opt/nemo-rl/3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/moe/experts.py", line 664, in forward
(permuted_local_hidden_states, tokens_per_expert, permuted_probs) = self.token_dispatcher.token_permutation2(
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
hidden_states=x,
^^^^^^^^^^^^^^^^
...<2 lines>...
token_indices=indices,
^^^^^^^^^^^^^^^^^^^^^^
)
^
File "/opt/nemo-rl/3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/moe/megatron/token_dispatcher.py", line 524, in token_permutation2
hidden_states, _ = self.dispatch_all_to_all(hidden_states, async_finish=False, allocate_on_comm_stream=False)
~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo-rl/3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/moe/megatron/token_dispatcher.py", line 469, in dispatch_all_to_all
self._comm_manager.dispatch(hidden_states, async_finish, allocate_on_comm_stream),
~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo-rl/3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/moe/megatron/token_dispatcher.py", line 177, in dispatch
) = fused_dispatch(
~~~~~~~~~~~~~~^
hidden_states,
^^^^^^^^^^^^^^
...<5 lines>...
allocate_on_comm_stream=allocate_on_comm_stream,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
)
^
File "/opt/nemo-rl/3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/moe/megatron/fused_a2a.py", line 272, in fused_dispatch
return FusedDispatch.apply(
~~~~~~~~~~~~~~~~~~~^
x.contiguous(),
^^^^^^^^^^^^^^^
...<5 lines>...
allocate_on_comm_stream,
^^^^^^^^^^^^^^^^^^^^^^^^
)
^
File "/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2/lib/python3.13/site-packages/torch/autograd/function.py", line 596, in apply
return super().apply(*args, **kwargs) # type: ignore[misc]
~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
File "/opt/nemo-rl/3rdparty/Automodel-workspace/Automodel/nemo_automodel/components/moe/megatron/fused_a2a.py", line 146, in forward
) = buffer.dispatch(
~~~~~~~~~~~~~~~^
x,
^^
...<8 lines>...
allocate_on_comm_stream=allocate_on_comm_stream,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
)
^
File "/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2/lib/python3.13/site-packages/deep_ep/buffer.py", line 376, in dispatch
return self.internode_dispatch(x, handle, num_tokens_per_rank, num_tokens_per_rdma_rank, is_token_in_rank,
~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
num_tokens_per_expert, topk_idx, topk_weights, expert_alignment, config, previous_event,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
async_finish, allocate_on_comm_stream)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/ray_venvs/nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2/lib/python3.13/site-packages/deep_ep/buffer.py", line 489, in internode_dispatch
recv_src_meta, send_rdma_head, send_nvl_head, event = self.runtime.internode_dispatch(
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
x, x_scales, topk_idx, topk_weights,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
num_tokens_per_rank, num_tokens_per_rdma_rank, is_token_in_rank, num_tokens_per_expert,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0, 0, None, None, None, None,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
expert_alignment, config, getattr(previous_event, 'event', None), async_finish, allocate_on_comm_stream)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: DeepEP error: timeout (dispatch CPU)

```