[Bug] OOM when running Qwen3-4B with offpolicy retool

Bug Description

When running Retool's off-policy GRPO experiments with Qwen3-4B, an OOM (Out Of Memory) error occurs after 37 steps, which is quite strange. Why does the training crash midway? Theoretically, GPU memory usage per step should remain stable.

Scripts

#!/bin/bash

for rerun the task

pkill -9 sglang
sleep 3
ray stop --force
pkill -9 ray
pkill -9 python
sleep 3
pkill -9 ray
pkill -9 python

set -ex

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)"

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"

source "/root/slime/scripts/models/qwen3-4B.sh"

MODEL_ARGS=(
--swiglu
--num-layers 36
--hidden-size 2560
--ffn-hidden-size 9728
--num-attention-heads 32
--group-query-attention
--num-query-groups 8
--use-rotary-position-embeddings
--disable-bias-linear
--normalization "RMSNorm"
--norm-epsilon 1e-6
--rotary-base "${MODEL_ARGS_ROTARY_BASE:-1000000}"
--vocab-size 151936
--kv-channels 128
--qk-layernorm
--rotary-base 5000000
)

CKPT_ARGS=(
--hf-checkpoint /mnt/data/zdm/models/font-info/qwen3-4b-sft
--ref-load /mnt/data/zdm/models/font-info/qwen3-4b-sft_torch_dist

--load ${LOG_DIR}/checkpoints/${EXP_NAME}

--save ${LOG_DIR}/checkpoints/${EXP_NAME}
--rotary-base 5000000
--save-interval 20
)

ROLLOUT_ARGS=(
--prompt-data /mnt/data/zdm/data/tir/tir_train_data/dapo-math-17k.jsonl
--input-key prompt
--label-key label
--apply-chat-template
--rollout-shuffle
--reward-key score
--num-rollout 3000
--rollout-batch-size 32
--n-samples-per-prompt 8
--rollout-max-response-len 8192
--rollout-max-prompt-len 1024
--rollout-temperature 1

--global-batch-size 256
--balance-data
)

EVAL_ARGS=(

# --eval-interval 2000

--eval-prompt-data aime /mnt/data/zdm/data/tir/tir_val_data/aime-2024.jsonl

--n-samples-per-eval-prompt 16

--eval-max-response-len 16384

--eval-top-p 1

--skip_eval_before_train

)

PERF_ARGS=(
--tensor-model-parallel-size 4
--sequence-parallel
--pipeline-model-parallel-size 1
--context-parallel-size 1
--expert-model-parallel-size 1
--expert-tensor-parallel-size 1

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

--micro-batch-size 1

--use-dynamic-batch-size
--max-tokens-per-gpu 9216

)

GRPO_ARGS=(
--advantage-estimator grpo
--get-mismatch-metrics

--use-kl-loss

--kl-loss-coef 0.00
--kl-loss-type low_var_kl
--entropy-coef 0.00

--mis-po-token-level-clip 0.5

--mis-po-token-level-clip-high 2

--mis-po-trajectory-level-clip 0.996

--mis-po-trajectory-level-clip-high 1.001

--eps-clip 0.2
--eps-clip-high 0.28
)

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

WANDB_ARGS=(
--use-tensorboard
--tensorboard_dir ${LOG_DIR}/tensorboards/${EXP_NAME}/

--use-wandb

--wandb-project slime-dapo

--wandb-group qwen3-4B-test-multi-turn

--wandb-key ${WANDB_KEY}

)

SGLANG_ARGS=(
--rollout-num-gpus-per-engine 2
--sglang-mem-fraction-static 0.7
)

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 generate_with_retool.generate
--custom-rm-path generate_with_retool.reward_func
)

launch the master node of ray in container

export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265

Build the runtime environment JSON with proper variable substitution

RUNTIME_ENV_JSON="{
"env_vars": {
"PYTHONPATH": "/root/Megatron-LM/:${SCRIPT_DIR}:/root/slime",
"CUDA_DEVICE_MAX_CONNECTIONS": "1",
"NCCL_NVLS_ENABLE": "${HAS_NVLINK}",
"PYTORCH_ALLOC_CONF": "expandable_segments:True"
}
}"

ray job submit --address="http://127.0.0.1:8265"
--runtime-env-json="${RUNTIME_ENV_JSON}"
-- python3 train.py
--actor-num-nodes 1
--actor-num-gpus-per-node 8
--colocate
${MODEL_ARGS[@]}
${CKPT_ARGS[@]}
${ROLLOUT_ARGS[@]}
${OPTIMIZER_ARGS[@]}
${GRPO_ARGS[@]}
${WANDB_ARGS[@]}
${PERF_ARGS[@]}
${EVAL_ARGS[@]}
${SGLANG_ARGS[@]}
${MISC_ARGS[@]}
${CUSTOM_ARGS[@]}

Steps to Reproduce

Could you please check it for me?

Expected Behavior

exp results

train(args)

File "/mnt/math/zxn/code/slime-20260305/slime/train.py", line 81, 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 2981, in get
values, debugger_breakpoint = worker.get_objects(
^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/ray/_private/worker.py", line 1012, in get_objects
raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(OutOfMemoryError): ray::MegatronTrainRayActor.train() (pid=1076529, ip=dlc4e52btvkgc6cp-master-0, actor_id=50f7ee525ea4e0e86f8909b702000000, repr=<slime.backends.megatron_utils.actor.MegatronTrainRayActor object at 0x7ecf08b497c0>)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/math/zxn/code/slime-20260305/slime/slime/backends/megatron_utils/actor.py", line 368, in train
return self.train_actor(rollout_id, rollout_data)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/math/zxn/code/slime-20260305/slime/slime/backends/megatron_utils/actor.py", line 440, in train_actor
self.compute_log_prob(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/math/zxn/code/slime-20260305/slime/slime/backends/megatron_utils/actor.py", line 346, in compute_log_prob
return forward_only(
^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 120, in decorate_context
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/mnt/math/zxn/code/slime-20260305/slime/slime/backends/megatron_utils/model.py", line 263, in forward_only
forward_data_store += forward_backward_func(
^^^^^^^^^^^^^^^^^^^^^^
File "/root/Megatron-LM/megatron/core/pipeline_parallel/schedules.py", line 632, in forward_backward_no_pipelining
output_tensor, num_tokens = forward_step(
^^^^^^^^^^^^^
File "/root/Megatron-LM/megatron/core/pipeline_parallel/schedules.py", line 417, in forward_step
output_tensor, loss_func = forward_step_func(data_iterator, model)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/math/zxn/code/slime-20260305/slime/slime/backends/megatron_utils/model.py", line 227, in forward_step
output_tensor = model(
^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/Megatron-LM/megatron/core/distributed/data_parallel_base.py", line 22, in forward
return self.module(*inputs, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/Megatron-LM/megatron/core/transformer/module.py", line 462, in forward
outputs = float16_to_fp32(outputs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/Megatron-LM/megatron/core/transformer/module.py", line 376, in float16_to_fp32
return conversion_helper(val, float_conversion)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/Megatron-LM/megatron/core/transformer/module.py", line 335, in conversion_helper
return conversion(val)
^^^^^^^^^^^^^^^
File "/root/Megatron-LM/megatron/core/transformer/module.py", line 373, in float_conversion
val = val.float()
^^^^^^^^^^^
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 61.34 GiB. GPU 3 has a total capacity of 79.11 GiB of which 27.07 GiB is free. Process 520735 has 3.17 GiB memory in use. Process 530150 has 48.84 GiB memory in use. Of the allocated memory 44.61 GiB is allocated by PyTorch, and 410.55 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
(MegatronTrainRayActor pid=1076535) [torch_memory_saver.cpp] TorchMemorySaver::malloc return OOM since memory_margin_bytes=1073741824 (alloc)size=65867350016 free_bytes=28822863872 [repeated 3x across cluster]
2026-03-25 21:19:13,067 ERR cli.py:73 -- ---------------------------------------
2026-03-25 21:19:13,067 ERR cli.py:74 -- Job 'raysubmit_RGD9ngBZxVzVGx8c' failed
2026-03-25 21:19:13,068 ERR cli.py:75 -- ---------------------------------------
2026-03-25 21:19:13,068 INFO cli.py:88 -- Status message: Job entrypoint command failed with exit code 1, last available logs (truncated to 20,000 chars):
return conversion_helper(val, float_conversion)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/Megatron-LM/megatron/core/transformer/module.py", line 335, in conversion_helper
return conversion(val)
^^^^^^^^^^^^^^^
File "/root/Megatron-LM/megatron/core/transformer/module.py", line 373, in float_conversion
val = val.float()
^^^^^^^^^^^
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 61.34 GiB. GPU 3 has a total capacity of 79.11 GiB of which 27.07 GiB is free. Process 520735 has 3.17 GiB memory in use. Process 530150 has 48.84 GiB memory in use. Of the allocated memory 44.61 GiB is allocated by PyTorch, and 410.55 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
(MegatronTrainRayActor pid=1076535) [torch_memory_saver.cpp] TorchMemorySaver::malloc return OOM since memory_margin_bytes=1073741824 (alloc)size=65867350016 free_bytes=28822863872 [repeated 3x across cluster]

Actual Behavior

Environment

• slime version:
• Python version:
• PyTorch version:
• CUDA/ROCm version:
• GPU type and count:
• OS:
• SGLang version (if relevant):
• Megatron-LM version (if relevant):

Logs

Additional Context

Or other optimizer scripts

Pre-submission Checklist

• I have read the CONTRIBUTING.md and understand the collaboration scope.
• I have read the documentation and my issue is not addressed there.
• I have searched for existing issues and this is not a duplicate.
• I have provided a minimal, reproducible example.

Start with the provided reproduction script and traceback, then read slime/backends/megatron_utils/actor.py around train, train_actor, and compute_log_prob, followed by model.py around forward_only and forward_step. Compare the failing step with earlier steps and inspect the Megatron float16_to_fp32 path named in the trace. Done means identifying the cause of the increasing allocation and demonstrating that the run no longer OOMs during the reported workload.