THUDM / THUDM/slime

时常出现GPU显存分配失败导致的训练中断

Open
#1,053 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
8.5k
Forks
1.3k
Avg merge
5h 36m
Merged PRs (30d)
22

Description

个人环境:8 * H100(80G),用的slime官方镜像
复现on policy distillation用的脚本:

#!/bin/bash

set -ex

Start the teacher model server

TEACHER_IP="127.0.0.1" # Use localhost here, you can change it to your IP
TEACHER_PORT=13141
LOG_FILE="./slime/output/log/sglang_$(head /dev/urandom | tr -dc A-Za-z0-9 | head -c 6).log"
echo $LOG_FILE

Launch the teacher model server in the background/ huggingface格式

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python3 -m sglang.launch_server
--model-path ./huggingface_model/Qwen3-4B-Instruct-2507
--host 0.0.0.0
--port $TEACHER_PORT
--tp 8
--chunked-prefill-size 4096
--mem-fraction-static 0.2
> "$LOG_FILE" 2>&1 &

echo "Starting teacher model server..."

Wait for the teacher model server to be ready

until curl -sf http://$TEACHER_IP:$TEACHER_PORT/health_generate > /dev/null; do
echo "Waiting for the teacher model server to start..."
tail -n 10 "$LOG_FILE"
sleep 5
done

echo "Teacher model server is up and running at $TEACHER_IP:$TEACHER_PORT."
sleep 10

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

student config

source "./slime/scripts/models/qwen3-0.6B.sh"

CKPT_ARGS=(
--hf-checkpoint ./huggingface_model/Qwen3-0.6B
--ref-load ./megatron_models/Qwen3-0.6B
--load ./megatron_models/Qwen3-0.6B
--save ./slime/output/Qwen3_4B_to_0.6B
--save-interval 20
)

ROLLOUT_ARGS=(
--prompt-data ./huggingface_datasets/dapo-math-17k/dapo-math-17k.jsonl
--input-key prompt
--apply-chat-template
--rollout-shuffle
--num-rollout 2000
--rollout-batch-size 8
--n-samples-per-prompt 2
--global-batch-size 16
--rollout-max-response-len 16384
--rollout-temperature 0.8
--balance-data
)

RM_ARGS=(
--custom-rm-path on_policy_distillation.reward_func
--custom-reward-post-process-path on_policy_distillation.post_process_rewards
--rm-url http://$TEACHER_IP:$TEACHER_PORT/generate
)

EVAL_ARGS=(
)

PERF_ARGS=(
--tensor-model-parallel-size 2
--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 1

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

GRPO_ARGS=(
--advantage-estimator on_policy_distillation
--use-kl-loss
--kl-loss-coef 0.00
--kl-loss-type low_var_kl
--entropy-coef 0.00
)

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=(
)

student sglang

SGLANG_ARGS=(
--rollout-num-gpus-per-engine 1
--sglang-mem-fraction-static 0.3
)

MISC_ARGS=(
--attention-dropout 0.0
--hidden-dropout 0.0
--accumulate-allreduce-grads-in-fp32
--attention-softmax-in-fp32
--attention-backend flash
)

launch the master node of ray in container

export RAY_DISABLE_DOCKER_CPU_WARNING=1
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
echo "master node start!"

ray job submit --address="http://$MASTER_ADDR:8265"
--runtime-env-json='{
"env_vars": {
"PYTHONPATH": "../Megatron-LM/",
"CUDA_DEVICE_MAX_CONNECTIONS": "1"
}
}'
-- python3 train.py
--actor-num-nodes 1
--actor-num-gpus-per-node 8
--rollout-num-gpus 1
--colocate
${MODEL_ARGS[@]}
${CKPT_ARGS[@]}
${ROLLOUT_ARGS[@]}
${OPTIMIZER_ARGS[@]}
${GRPO_ARGS[@]}
${WANDB_ARGS[@]}
${PERF_ARGS[@]}
${EVAL_ARGS[@]}
${SGLANG_ARGS[@]}
${MISC_ARGS[@]}
${RM_ARGS[@]}

####clear after training
pkill -9 sglang
sleep 3
ray stop --force
pkill -9 ray
pkill -9 python
sleep 3
pkill -9 ray
pkill -9 python

使用上述脚本经常运行中出现下述:
torch.distributed.DistBackendError: NCCL error in: /pytorch/torch/csrc/distributed/c10d/NCCLUtils.cpp:77, unhandled cuda error (run with NCCL_DEBUG=INFO for details), NCCL version 2.27.5
ncclUnhandledCudaError: Call to CUDA function failed.
Last error:
Failed to CUDA calloc 10485760 bytes
mem_info={'gpu': '0', 'total_GB': 79.11, 'free_GB': 0.01, 'used_GB': 79.1, 'allocated_GB': 2.98, 'reserved_GB': 37.54}
mem_info={'gpu': '0', 'total_GB': 79.11, 'free_GB': 0.01, 'used_GB': 79.1, 'allocated_GB': 2.98, 'reserved_GB': 37.54}
看起来是cuda分配失败,有大佬咋回事嘛?上面哪个地方没配置正确吗

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start from the inline launch script and rerun the reproduction with NCCL_DEBUG=INFO as suggested by the error, while checking the generated sglang log and the reported GPU memory fields. Trace which launch settings coincide with the NCCL allocation failure, then verify the identified configuration can complete training without interruption.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
distributed-systems, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.