yi_lu/laoyao_2b_moe
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add state-based training resume command

Browse Source
This commit is contained in:
yi_lu
2026-07-23 19:48:41 +08:00
parent 54ac00152a
commit 2cd0007685
10 changed files with 138 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
AGENTS.md
+10
View File
@@ -263,6 +263,16 @@ All long-running launchers use the fixed default container name
directory in `.runtime/active_training.txt`, so the stop command needs no
arguments. Do not use `docker kill --signal=SIGINT` for a routine stop.
Resume the last recorded run:
```bash
bash scripts/resume_training.sh
```
This reads `.runtime/active_training.txt` after an external interruption or
`.runtime/last_training.txt` after a graceful stop, validates the checkpoint
tracker, and invokes the original launcher.
Before stopping, check the latest checkpoint:
```bash
scripts/README.md
+23 -1
View File
@@ -18,6 +18,8 @@
iteration 完成后临时保存 checkpoint 并安全退出。
- `record_active_training.sh`: launcher 启动容器后原子记录当前 container、run 和
checkpoint 路径，供状态检查与无参数 graceful stop 使用。
- `resume_training.sh`: 读取最后一次训练状态，检查 checkpoint tracker，然后重新运行
原训练 launcher。
## g0050 下载与部署
@@ -117,7 +119,13 @@ laoyao-2b-pretraining
```
训练 launcher 确认容器启动后还会原子更新状态文件，记录 checkpoint 目录、run
目录和日志路径。即使状态文件缺失，停止脚本也会操作上述固定容器名。多实验并存时
目录、日志路径和 launcher，并保留：
```text
.runtime/last_training.txt
```
即使状态文件缺失，停止脚本也会操作上述固定容器名。多实验并存时
仍可显式覆盖：
```bash
@@ -132,3 +140,17 @@ bash scripts/graceful_stop_megatron_training.sh <container-name>
--graceful-exit-signal SIGTERM
--disable-graceful-exit-on-signal
```
停止后恢复最后一次训练：
```bash
bash scripts/resume_training.sh
```
该脚本优先读取 `active_training.txt`（适用于被外部强制终止的训练），否则读取
`last_training.txt`（适用于正常 graceful stop）。它会确认固定容器当前未运行、
checkpoint tracker 存在，再重新调用记录的 launcher。只检查、不启动：
```bash
RESUME_DRY_RUN=1 bash scripts/resume_training.sh
```
scripts/record_active_training.sh
+24 -14
View File
@@ -1,8 +1,8 @@
#!/usr/bin/env bash
set -euo pipefail
if [[ "$#" -ne 4 ]]; then
echo "usage: $0 <container-name> <checkpoint-dir> <run-dir> <log-file>" >&2
if [[ "$#" -ne 5 ]]; then
echo "usage: $0 <container-name> <checkpoint-dir> <run-dir> <log-file> <launcher-script>" >&2
exit 2
fi
@@ -10,9 +10,11 @@ CONTAINER_NAME="$1"
CHECKPOINT_DIR="$2"
RUN_DIR="$3"
LOG_FILE="$4"
LAUNCHER_SCRIPT="$5"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
TRAINING_STATE_FILE="${TRAINING_STATE_FILE:-${REPO_ROOT}/.runtime/active_training.txt}"
LAST_TRAINING_STATE_FILE="${LAST_TRAINING_STATE_FILE:-${REPO_ROOT}/.runtime/last_training.txt}"
STARTUP_WAIT_SECONDS="${STARTUP_WAIT_SECONDS:-30}"
deadline=$((SECONDS + STARTUP_WAIT_SECONDS))
@@ -24,17 +26,25 @@ while [[ "$(docker inspect -f '{{.State.Running}}' "$CONTAINER_NAME" 2>/dev/null
sleep 1
done
mkdir -p "$(dirname "$TRAINING_STATE_FILE")"
temporary_file="$(mktemp "${TRAINING_STATE_FILE}.tmp.XXXXXX")"
trap 'rm -f "$temporary_file"' EXIT
{
printf 'container_name=%s\n' "$CONTAINER_NAME"
printf 'checkpoint_dir=%s\n' "$CHECKPOINT_DIR"
printf 'run_dir=%s\n' "$RUN_DIR"
printf 'log_file=%s\n' "$LOG_FILE"
printf 'recorded_at=%s\n' "$(date --iso-8601=seconds)"
} >"$temporary_file"
mv "$temporary_file" "$TRAINING_STATE_FILE"
trap - EXIT
recorded_at="$(date --iso-8601=seconds)"
write_state_file() {
local state_file="$1"
local temporary_file
mkdir -p "$(dirname "$state_file")"
temporary_file="$(mktemp "${state_file}.tmp.XXXXXX")"
{
printf 'container_name=%s\n' "$CONTAINER_NAME"
printf 'checkpoint_dir=%s\n' "$CHECKPOINT_DIR"
printf 'run_dir=%s\n' "$RUN_DIR"
printf 'log_file=%s\n' "$LOG_FILE"
printf 'launcher_script=%s\n' "$LAUNCHER_SCRIPT"
printf 'recorded_at=%s\n' "$recorded_at"
} >"$temporary_file"
mv "$temporary_file" "$state_file"
}
write_state_file "$TRAINING_STATE_FILE"
write_state_file "$LAST_TRAINING_STATE_FILE"
echo "active_training_state=$TRAINING_STATE_FILE"
echo "last_training_state=$LAST_TRAINING_STATE_FILE"
scripts/resume_pretrain_8192_8gpu_mbs14.sh
+2 -1
View File
@@ -92,7 +92,8 @@ nohup docker run --rm --name "${CONTAINER_NAME}" \
--recompute-num-layers 1" \
> "${LOG_FILE}" 2>&1 &
bash "${SCRIPT_DIR}/record_active_training.sh" "$CONTAINER_NAME" "$CKPT_DIR" "$RUN_DIR" "$LOG_FILE"
bash "${SCRIPT_DIR}/record_active_training.sh" \
"$CONTAINER_NAME" "$CKPT_DIR" "$RUN_DIR" "$LOG_FILE" "${SCRIPT_DIR}/$(basename "${BASH_SOURCE[0]}")"
echo "Launched. Check with:"
echo " tail -f ${LOG_FILE}"
echo " docker ps"
scripts/resume_stage1_2_pretrain_8192_8gpu_mbs14.sh
+2 -1
View File
@@ -117,7 +117,8 @@ nohup docker run --rm --name "${CONTAINER_NAME}" \
> "${LOG_FILE}" 2>&1 &
bash "${SCRIPT_DIR}/record_active_training.sh" \
"$CONTAINER_NAME" "$HOST_CKPT_DIR" "$(dirname "$HOST_CKPT_DIR")" "$LOG_FILE"
"$CONTAINER_NAME" "$HOST_CKPT_DIR" "$(dirname "$HOST_CKPT_DIR")" "$LOG_FILE" \
"${SCRIPT_DIR}/$(basename "${BASH_SOURCE[0]}")"
echo "Launched. Check with:"
echo " tail -f ${LOG_FILE}"
echo " docker ps"
scripts/resume_training.sh
Executable
+69
View File
@@ -0,0 +1,69 @@
#!/usr/bin/env bash
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
ACTIVE_STATE_FILE="${TRAINING_STATE_FILE:-${REPO_ROOT}/.runtime/active_training.txt}"
LAST_STATE_FILE="${LAST_TRAINING_STATE_FILE:-${REPO_ROOT}/.runtime/last_training.txt}"
RESUME_DRY_RUN="${RESUME_DRY_RUN:-0}"
read_state_value() {
local state_file="$1"
local key="$2"
awk -v key="$key" '
index($0, key "=") == 1 {
sub(/^[^=]*=/, "")
print
exit
}
' "$state_file"
}
if [[ -f "$ACTIVE_STATE_FILE" ]]; then
STATE_FILE="$ACTIVE_STATE_FILE"
elif [[ -f "$LAST_STATE_FILE" ]]; then
STATE_FILE="$LAST_STATE_FILE"
else
echo "ERROR: no active or previous training state found." >&2
echo "checked: $ACTIVE_STATE_FILE" >&2
echo "checked: $LAST_STATE_FILE" >&2
exit 2
fi
CONTAINER_NAME="$(read_state_value "$STATE_FILE" container_name)"
CHECKPOINT_DIR="$(read_state_value "$STATE_FILE" checkpoint_dir)"
RUN_DIR="$(read_state_value "$STATE_FILE" run_dir)"
LOG_FILE="$(read_state_value "$STATE_FILE" log_file)"
LAUNCHER_SCRIPT="$(read_state_value "$STATE_FILE" launcher_script)"
CONTAINER_NAME="${CONTAINER_NAME:-laoyao-2b-pretraining}"
if [[ "$(docker inspect -f '{{.State.Running}}' "$CONTAINER_NAME" 2>/dev/null || true)" == "true" ]]; then
echo "ERROR: training container is already running: $CONTAINER_NAME" >&2
exit 3
fi
if [[ -z "$LAUNCHER_SCRIPT" || ! -f "$LAUNCHER_SCRIPT" ]]; then
echo "ERROR: recorded launcher script is missing: ${LAUNCHER_SCRIPT:-<empty>}" >&2
exit 4
fi
if [[ -z "$CHECKPOINT_DIR" || ! -f "${CHECKPOINT_DIR}/latest_checkpointed_iteration.txt" ]]; then
echo "ERROR: resumable checkpoint tracker is missing: ${CHECKPOINT_DIR:-<empty>}" >&2
exit 5
fi
latest_iteration="$(tr -d '[:space:]' < "${CHECKPOINT_DIR}/latest_checkpointed_iteration.txt")"
cat <<EOF
state_file=$STATE_FILE
container_name=$CONTAINER_NAME
launcher_script=$LAUNCHER_SCRIPT
run_dir=$RUN_DIR
checkpoint_dir=$CHECKPOINT_DIR
latest_checkpointed_iteration=$latest_iteration
log_file=$LOG_FILE
EOF
if [[ "$RESUME_DRY_RUN" == "1" ]]; then
echo "Dry run passed; training was not launched."
exit 0
fi
CONTAINER_NAME="$CONTAINER_NAME" bash "$LAUNCHER_SCRIPT"
scripts/train_stage2_1_plainqa_cpt_8192_8gpu_mbs14.sh
+2 -1
View File
@@ -154,5 +154,6 @@ nohup docker run --rm --name "$CONTAINER_NAME" \
--recompute-granularity full --recompute-method uniform --recompute-num-layers 1 \
$STAGE_TRANSITION_ARG" >"$LOG_FILE" 2>&1 &
bash "${SCRIPT_DIR}/record_active_training.sh" "$CONTAINER_NAME" "$CKPT_DIR" "$RUN_DIR" "$LOG_FILE"
bash "${SCRIPT_DIR}/record_active_training.sh" \
"$CONTAINER_NAME" "$CKPT_DIR" "$RUN_DIR" "$LOG_FILE" "${SCRIPT_DIR}/$(basename "${BASH_SOURCE[0]}")"
echo "Launched container=$CONTAINER_NAME log=$LOG_FILE"
scripts/train_stage2_1_pretrain_8192_8gpu_mbs14.sh
+2 -1
View File
@@ -141,5 +141,6 @@ nohup docker run --rm --name "$CONTAINER_NAME" \
--recompute-granularity full --recompute-method uniform --recompute-num-layers 1 \
$STAGE_TRANSITION_ARG" >"$LOG_FILE" 2>&1 &
bash "${SCRIPT_DIR}/record_active_training.sh" "$CONTAINER_NAME" "$CKPT_DIR" "$RUN_DIR" "$LOG_FILE"
bash "${SCRIPT_DIR}/record_active_training.sh" \
"$CONTAINER_NAME" "$CKPT_DIR" "$RUN_DIR" "$LOG_FILE" "${SCRIPT_DIR}/$(basename "${BASH_SOURCE[0]}")"
echo "Launched container=$CONTAINER_NAME log=$LOG_FILE"
scripts/train_stage2_2_plainqa_cpt_8192_8gpu_mbs14.sh
+2 -1
View File
@@ -147,5 +147,6 @@ nohup docker run --rm --name "$CONTAINER_NAME" \
--recompute-granularity full --recompute-method uniform --recompute-num-layers 1 \
$STAGE_TRANSITION_ARG" >"$LOG_FILE" 2>&1 &
bash "${SCRIPT_DIR}/record_active_training.sh" "$CONTAINER_NAME" "$CKPT_DIR" "$RUN_DIR" "$LOG_FILE"
bash "${SCRIPT_DIR}/record_active_training.sh" \
"$CONTAINER_NAME" "$CKPT_DIR" "$RUN_DIR" "$LOG_FILE" "${SCRIPT_DIR}/$(basename "${BASH_SOURCE[0]}")"
echo "Launched container=$CONTAINER_NAME log=$LOG_FILE"
scripts/train_stage2_2_pretrain_8192_8gpu_mbs14.sh
+2 -1
View File
@@ -147,5 +147,6 @@ nohup docker run --rm --name "$CONTAINER_NAME" \
--recompute-granularity full --recompute-method uniform --recompute-num-layers 1 \
$STAGE_TRANSITION_ARG" >"$LOG_FILE" 2>&1 &
bash "${SCRIPT_DIR}/record_active_training.sh" "$CONTAINER_NAME" "$CKPT_DIR" "$RUN_DIR" "$LOG_FILE"
bash "${SCRIPT_DIR}/record_active_training.sh" \
"$CONTAINER_NAME" "$CKPT_DIR" "$RUN_DIR" "$LOG_FILE" "${SCRIPT_DIR}/$(basename "${BASH_SOURCE[0]}")"
echo "Launched container=$CONTAINER_NAME log=$LOG_FILE"