[Bug] `Unexpected result nan` on 64x H100, Qwen3 235B

Hi slime maintainers. I am running slime v0.1.0 for a stable version and encountering the same error multiple times.

Launch script:

#!/bin/bash

set -e

# will prevent ray from buffering stdout/stderr
export PYTHONBUFFERED=16

NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
if [ "$NVLINK_COUNT" -gt 0 ]; then
    HAS_NVLINK=1
else
    HAS_NVLINK=0
fi
echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
source "scripts/slime_models/Qwen3-235B-A22B-Thinking-2507_model_args.sh"

CKPT_ARGS=(
   --hf-checkpoint $SHARED_STORAGE_DATA/hf_models/Qwen3-235B-A22B-Thinking-2507
   --ref-load $SHARED_STORAGE_DATA/megatron_models/Qwen3-235B-A22B-Thinking-2507
   --load $SHARED_STORAGE_DATA/slime-ckpt/Qwen3-235B-A22B-Thinking-2507_slime_v1.5/
   --save $SHARED_STORAGE_DATA/slime-ckpt/Qwen3-235B-A22B-Thinking-2507_slime_v1.5/
   --save-interval 20
)

NODES=$TOTAL_NODES
PROMPT_LEN=4096
RESPONSE_LEN=$((65536 - PROMPT_LEN))
CONTEXT_LEN=$((PROMPT_LEN + RESPONSE_LEN))
GLOBAL_BATCH_SIZE=$((NODES * 64))
ROLLOUT_BATCH_SIZE=$((NODES * 8))

ROLLOUT_ARGS=(
   --prompt-data $SHARED_STORAGE_DATA/multi_pmc_qa/train.parquet
   --input-key conversation
   --label-key answer
   --rollout-shuffle
   --group-rm
   --num-rollout 3000
   --rollout-batch-size $ROLLOUT_BATCH_SIZE
   --n-samples-per-prompt 8
   --rollout-max-response-len $RESPONSE_LEN
   --rollout-temperature 0.8

   --global-batch-size $GLOBAL_BATCH_SIZE
   --balance-data
)

PERF_ARGS=(
   --tensor-model-parallel-size 4
   --sequence-parallel
   --pipeline-model-parallel-size 4
   --context-parallel-size 2
   --expert-model-parallel-size 16
   --expert-tensor-parallel-size 1
   --decoder-last-pipeline-num-layers 22

   --recompute-granularity full
   --recompute-method uniform
   --recompute-num-layers 1

   --use-dynamic-batch-size
   --max-tokens-per-gpu 16384
)

GRPO_ARGS=(
   --advantage-estimator grpo
   --use-kl-loss
   --kl-loss-coef 0.00
   --kl-loss-type low_var_kl
   --entropy-coef 0.00
   --eps-clip 0.2
   --eps-clip-high 0.28
   --use-tis
)

OPTIMIZER_ARGS=(
   --optimizer adam
   --lr 1e-6
   --lr-decay-style constant
   --weight-decay 0.1
   --adam-beta1 0.9
   --adam-beta2 0.98

   --optimizer-cpu-offload
   --overlap-cpu-optimizer-d2h-h2d
   --use-precision-aware-optimizer
)

WANDB_ARGS=(
   --use-wandb
   --wandb-team scientific_intelligence_nlp
   --wandb-project slime-training
   --wandb-group qwen3-30B-A3B-test
   --wandb-key $WANDB_API_KEY
   --wandb-host https://nn.wandb.io/
)

SGLANG_ARGS=(
   --rollout-num-gpus-per-engine 32
   --sglang-mem-fraction-static 0.7
   --sglang-enable-dp-attention
   --sglang-dp-size 4
   --sglang-ep-size 32
   --sglang-enable-dp-lm-head
   --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 256)
   --sglang-context-length $CONTEXT_LEN
)

MISC_ARGS=(
   # default dropout in megatron is 0.1
   --attention-dropout 0.0
   --hidden-dropout 0.0
   # should be good for model performance
   --accumulate-allreduce-grads-in-fp32
   --attention-softmax-in-fp32
   # need to comment this when using model with MLA
   --attention-backend flash
)

CUSTOM_ARGS=(
   --custom-generate-function-path deepresearch.generate.generate
   --custom-rm-path deepresearch.reward.reward_func
)

# Build the runtime environment JSON with proper variable substitution
RUNTIME_ENV_JSON="{
  \"env_vars\": {
    \"PYTHONPATH\": \"/root/Megatron-LM/\",
    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\"
  }
}"

ray job submit --address="http://$MASTER_ADDR:8265" \
   --runtime-env-json="${RUNTIME_ENV_JSON}" \
   -- python3 deepresearch/train.py \
   --actor-num-nodes $((NODES)) \
   --actor-num-gpus-per-node 8 \
   --rollout-num-gpus 64 \
   --colocate \
   --update-weight-buffer-size $(( 1024 * 1024 * 1024 * 4 )) \
   ${MODEL_ARGS[@]} \
   ${CKPT_ARGS[@]} \
   ${ROLLOUT_ARGS[@]} \
   ${OPTIMIZER_ARGS[@]} \
   ${GRPO_ARGS[@]} \
   ${DISTRIBUTED_ARGS[@]} \
   ${WANDB_ARGS[@]} \
   ${PERF_ARGS[@]} \
   ${SGLANG_ARGS[@]} \
   ${MISC_ARGS[@]} \
   ${CUSTOM_ARGS[@]}

Traceback:

Traceback (most recent call last):
  File "/workspace/deepresearch/train.py", line 89, in <module>
    train(args)
  File "/workspace/deepresearch/train.py", line 61, in train
    ray.get(actor_model.async_train(rollout_id, rollout_data_ref))
  File "/usr/local/lib/python3.12/dist-packages/ray/_private/auto_init_hook.py", line 22, in auto_init_wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/ray/_private/client_mode_hook.py", line 104, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/ray/_private/worker.py", line 2858, in get
    values, debugger_breakpoint = worker.get_objects(object_refs, timeout=timeout)
                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/ray/_private/worker.py", line 958, in get_objects
    raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(RuntimeError): ray::MegatronTrainRayActor.train() (pid=2526545, ip=172.27.6.172, actor_id=a94008d977b34ec82af4ae0602000000, repr=<slime.backends.megatron_utils.actor.MegatronTrainRayActor object at 0x152134dccc50>)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/slime/backends/megatron_utils/actor.py", line 258, in train
    train(
  File "/usr/local/lib/python3.12/dist-packages/slime/backends/megatron_utils/model.py", line 431, in train
    loss_dict, grad_norm = train_one_step(
                           ^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/slime/backends/megatron_utils/model.py", line 310, in train_one_step
    losses_reduced = forward_backward_func(
                     ^^^^^^^^^^^^^^^^^^^^^^
  File "/root/Megatron-LM/megatron/core/pipeline_parallel/schedules.py", line 2185, in forward_backward_pipelining_without_interleaving
    config.finalize_model_grads_func(
  File "/root/Megatron-LM/megatron/core/distributed/finalize_model_grads.py", line 422, in finalize_model_grads
    model_chunk.finish_grad_sync()
  File "/root/Megatron-LM/megatron/core/distributed/distributed_data_parallel.py", line 588, in finish_grad_sync
    bucket_group.finish_grad_sync()
  File "/root/Megatron-LM/megatron/core/distributed/param_and_grad_buffer.py", line 455, in finish_grad_sync
    self.start_grad_sync()
  File "/root/Megatron-LM/megatron/core/distributed/param_and_grad_buffer.py", line 333, in start_grad_sync
    self.check_grads(
  File "/root/Megatron-LM/megatron/core/distributed/param_and_grad_buffer.py", line 185, in check_grads
    rerun_state_machine.validate_result(
  File "/root/Megatron-LM/megatron/core/rerun_state_machine.py", line 505, in validate_result
    raise RuntimeError(full_message)
RuntimeError: Rank 50, node dgx172, device 2, iteration -1: Unexpected result nan (message='found NaN in local grad norm for bucket #0 in backward pass before data-parallel communication collective')

Start with deepresearch/train.py and trace the failing actor call into slime/backends/megatron_utils/actor.py and model.py. Then inspect Megatron-LM's param_and_grad_buffer.py and rerun_state_machine.py around the NaN grad-norm check, using the supplied 64x H100 Qwen3 launch configuration. Done means identifying a reproducible cause and validating a fix without the unexpected NaN result.