This commit is contained in:
parent
f115389dd0
commit
b2075d945b
|
|
@ -2,7 +2,7 @@
|
|||
set -euo pipefail
|
||||
|
||||
# ===== 可调参数 =====
|
||||
CKPT_ROOT="/home/test/checkpoints/q3-32b-ds4"
|
||||
CKPT_ROOT="/home/test/checkpoints/q3-32b-ds4" # 如果分片实际在 checkpoint-62/global_step62 下，就把这里改成 .../checkpoint-62
|
||||
TAG="global_step62"
|
||||
HOSTS=(tn01 tn02 tn03 tn04 tn05 tn06)
|
||||
OUT_DIR="${CKPT_ROOT}/merged-${TAG}"
|
||||
|
|
@ -11,64 +11,67 @@ SSH_OPTS="-o BatchMode=yes -o StrictHostKeyChecking=accept-new -o ConnectTimeout
|
|||
RSYNC_OPTS="-a --info=progress2 --human-readable --partial --inplace"
|
||||
# ====================
|
||||
|
||||
echo "== 预检查 SSH 与远端目录 =="
|
||||
echo "== 预检查 SSH 与（非必需）远端目录存在 =="
|
||||
for h in "${HOSTS[@]}"; do
|
||||
if ! ssh ${SSH_OPTS} "$h" "true" >/dev/null 2>&1; then
|
||||
echo "!! 无法免密 SSH 到 $h（检查 ~/.ssh/config/authorized_keys/防火墙）" >&2
|
||||
exit 1
|
||||
fi
|
||||
if ! ssh ${SSH_OPTS} "$h" "test -d '${CKPT_ROOT}/${TAG}'"; then
|
||||
echo "!! $h 上缺少目录 ${CKPT_ROOT}/${TAG}，确认训练是否在该机产生了分片" >&2
|
||||
fi
|
||||
ssh ${SSH_OPTS} "$h" "true" >/dev/null || { echo "!! 无法免密 SSH 到 $h"; exit 1; }
|
||||
# 目录不存在也不致命，后面会跳过
|
||||
done
|
||||
|
||||
echo "== 1/4 开始按节点同步分片（仅 ${TAG}，带进度）=="
|
||||
mkdir -p "${CKPT_ROOT}"
|
||||
echo "== 1/4 同步各节点的 ${TAG} 整个目录（带进度）=="
|
||||
mkdir -p "${CKPT_ROOT}/${TAG}"
|
||||
LOCAL_HOST="$(hostname -s || hostname)"
|
||||
for h in "${HOSTS[@]}"; do
|
||||
[[ "$h" == "$LOCAL_HOST" ]] && { echo " - 跳过本机 $h"; continue; }
|
||||
echo " - 从 $h 拉取 ${CKPT_ROOT}/${TAG}/mp_rank_*/"
|
||||
# 只拉取该 step 下的 mp_rank_* 目录，避免无关文件
|
||||
rsync ${RSYNC_OPTS} -e "ssh ${SSH_OPTS}" \
|
||||
--include="${TAG}/" --include="${TAG}/mp_rank_*/" --include="${TAG}/mp_rank_*/**" \
|
||||
--exclude="*" \
|
||||
"${h}:${CKPT_ROOT}/" "${CKPT_ROOT}/"
|
||||
if ssh ${SSH_OPTS} "$h" "test -d '${CKPT_ROOT}/${TAG}'"; then
|
||||
echo " - 从 $h 拉取 ${CKPT_ROOT}/${TAG}/"
|
||||
# 不做 include/exclude 过滤，避免漏掉不同命名风格的分片文件
|
||||
rsync ${RSYNC_OPTS} -e "ssh ${SSH_OPTS}" \
|
||||
"${h}:${CKPT_ROOT}/${TAG}/" "${CKPT_ROOT}/${TAG}/" || true
|
||||
else
|
||||
echo " - $h 无 ${CKPT_ROOT}/${TAG}，跳过"
|
||||
fi
|
||||
done
|
||||
|
||||
echo "== 2/4 校验是否凑齐分片目录 =="
|
||||
if [[ ! -d "${CKPT_ROOT}/${TAG}" ]]; then
|
||||
echo "!! 未发现 ${CKPT_ROOT}/${TAG}" >&2; exit 1
|
||||
fi
|
||||
MP_CNT=$(find "${CKPT_ROOT}/${TAG}" -maxdepth 1 -type d -name "mp_rank_*" | wc -l | tr -d ' ')
|
||||
echo " - 已发现 mp_rank_* 目录数：${MP_CNT}"
|
||||
if [[ "${MP_CNT}" -eq 0 ]]; then
|
||||
echo "!! 没有任何 mp_rank_* 分片，请检查同步" >&2; exit 1
|
||||
echo "== 2/4 校验是否有分片“文件”（不是目录）=="
|
||||
# 兼容两种常见命名：mp_rank_*_model_states.pt 与 *mp_rank*model_states.pt（含 pp 维度）
|
||||
CNT_A=$(ls -1 "${CKPT_ROOT}/${TAG}"/mp_rank_*_model_states.pt 2>/dev/null | wc -l | tr -d ' ' || true)
|
||||
CNT_B=$(ls -1 "${CKPT_ROOT}/${TAG}"/*mp_rank*model_states.pt 2>/dev/null | wc -l | tr -d ' ' || true)
|
||||
CNT=$(( CNT_A + CNT_B ))
|
||||
echo " - 发现 model_states 分片文件数：${CNT}"
|
||||
if [[ "${CNT}" -eq 0 ]]; then
|
||||
echo "!! 未检测到任何 *_model_states.pt；请在各机上 ls 看看 ${CKPT_ROOT}/${TAG} 的实际文件名，再调整匹配规则" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "== 3/4 合并为 safetensors 输出到：${OUT_DIR} =="
|
||||
python - <<PY
|
||||
import os
|
||||
mkdir -p "${OUT_DIR}"
|
||||
if [[ -f "${CKPT_ROOT}/zero_to_fp32.py" ]]; then
|
||||
# 优先使用与该 checkpoint 同版本的官方脚本
|
||||
python "${CKPT_ROOT}/zero_to_fp32.py" \
|
||||
"${CKPT_ROOT}" \
|
||||
"${OUT_DIR}" \
|
||||
--tag "${TAG}" \
|
||||
--safe_serialization \
|
||||
--max_shard_size "${MAX_SHARD_SIZE}"
|
||||
else
|
||||
# 退回 DeepSpeed API
|
||||
python - <<PY
|
||||
from deepspeed.utils.zero_to_fp32 import convert_zero_checkpoint_to_fp32_state_dict
|
||||
ckpt_dir = r"${CKPT_ROOT}"
|
||||
out_dir = r"${OUT_DIR}"
|
||||
tag = r"${TAG}"
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
convert_zero_checkpoint_to_fp32_state_dict(
|
||||
checkpoint_dir=ckpt_dir,
|
||||
output_dir=out_dir,
|
||||
tag=tag,
|
||||
checkpoint_dir=r"${CKPT_ROOT}",
|
||||
output_dir=r"${OUT_DIR}",
|
||||
tag=r"${TAG}",
|
||||
safe_serialization=True,
|
||||
max_shard_size="${MAX_SHARD_SIZE}",
|
||||
max_shard_size=r"${MAX_SHARD_SIZE}",
|
||||
)
|
||||
print("合并完成：", out_dir)
|
||||
print("合并完成：", r"${OUT_DIR}")
|
||||
PY
|
||||
fi
|
||||
|
||||
echo "== 3.1 拷贝 config/tokenizer 工件（如存在）=="
|
||||
pushd "${CKPT_ROOT}" >/dev/null
|
||||
for f in config.json generation_config.json tokenizer_config.json tokenizer.json merges.txt vocab.json special_tokens_map.json added_tokens.json; do
|
||||
[[ -f "$f" ]] && cp -n "$f" "${OUT_DIR}/"
|
||||
[[ -f "${CKPT_ROOT}/${f}" ]] && cp -n "${CKPT_ROOT}/${f}" "${OUT_DIR}/"
|
||||
done
|
||||
popd >/dev/null
|
||||
|
||||
echo "== 4/4 自检（索引与config）=="
|
||||
python - <<'PY'
|
||||
|
|
|
|||
Loading…
Reference in New Issue