File size: 3,196 Bytes
2622c40 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | #!/usr/bin/env bash
# Launch ONE of the pusht training runs with George's lift_new harness
# (port of dish_utensil_train with only the dataset and configs
# swapped). The sbatch wrappers run stage_pusht_local.sh on the training node
# first.
#
# Usage: run_one.sh <task_name> <main_process_port> [extra hydra overrides...]
# e.g.: NUM_GPUS=4 bash run_one.sh pusht_fastwam_scratch 29566
#
# task_name is one of the configs/task/*.yaml here; the output root and wandb
# run name are derived from it. wandb entity/project/mode are baked in (the
# checkout default points at a different entity, offline).
set -uo pipefail
TASK="${1:?task name required (e.g. pusht_fastwam_scratch)}"
PORT="${2:?main_process_port required (distinct per concurrent job)}"
shift 2
FASTWAM_ROOT=/shared_work/physical_intelligence/policies/Fast-WAM/fastwam
CONFIG_DIR=/shared_work/george/real_world_wam/pusht_train/configs
OUTPUT_ROOT=/shared_work/george/real_world_wam/runs/${TASK}
NUM_GPUS="${NUM_GPUS:-2}"
RUN_ID="$(date +%Y-%m-%d_%H-%M-%S)"
OUTPUT_DIR="${OUTPUT_ROOT}/${RUN_ID}"
ENV_PY=/shared_work/environments/miniconda3/envs/fastwam-robotwin-rui/bin
export PATH="${ENV_PY}:${PATH}"
# glibc malloc tuning — two constraints learned the hard way:
# (a) mmap churn: workers allocating/freeing ~16MB numpy frames above the
# default 128KB threshold produced tens of thousands of mmap calls/s
# (kernel page-allocator churn) → keep mid-size buffers pooled in arenas.
# (b) earlyoom on gpu-h200-103: the old blanket MALLOC_MMAP_THRESHOLD_=1GB
# let every dataloader worker hoard ~10GB of freed decode buffers
# (~1TB per 4-GPU job); two concurrent 4-GPU jobs overflowed the node's
# 2TB and earlyoom SIGTERM'd the fattest rank (killed jobs 1677 then
# 1676 on 2026-08-13, ~30 min into each overlap).
# 64MB threshold still pools the 16MB frame buffers (no churn) while rare
# big allocations go to mmap and are returned on free; ARENA_MAX caps
# per-process arena fragmentation; TRIM returns freed arena space to the OS.
export MALLOC_MMAP_THRESHOLD_=67108864
export MALLOC_ARENA_MAX=2
export MALLOC_TRIM_THRESHOLD_=134217728
# NB: do NOT set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True — amin's
# runs proved it makes DeepSpeed ZeRO-1 bf16 go loss=nan from ~step 10.
mkdir -p "${OUTPUT_DIR}"
# cwd must be the fastwam checkout: its configs and the Wan2.2 base weights
# (checkpoints/...) are referenced by relative paths. We only read from it;
# all outputs go to OUTPUT_DIR.
cd "${FASTWAM_ROOT}"
accelerate launch \
--config_file scripts/accelerate_configs/accelerate_zero1_ds.yaml \
--num_processes "${NUM_GPUS}" \
--main_process_port "${PORT}" \
scripts/train.py \
--config-dir "${CONFIG_DIR}" \
"task=${TASK}" \
"output_dir=${OUTPUT_DIR}" \
"wandb.name=${TASK}_${RUN_ID}" \
wandb.workspace=huaweiwam \
wandb.project=fastwam-realrobot \
"wandb.mode=${WANDB_MODE_OVERRIDE:-online}" \
hydra/job_logging=stdout \
"$@" 2>&1 | tee "${OUTPUT_DIR}/train_console.log"
# hydra/job_logging=stdout: Hydra's default file handler would try to write
# train.log into the read-only fastwam checkout (cwd) and crash at startup.
exit "${PIPESTATUS[0]}"
|