Upload scripts/memprobe.sh with huggingface_hub
Browse files- scripts/memprobe.sh +36 -0
scripts/memprobe.sh
ADDED
|
@@ -0,0 +1,36 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env bash
|
| 2 |
+
# Can 2x H100 TRAIN at 32k / 64k? Measure, do not extrapolate.
|
| 3 |
+
# Phase 1 measured 50,117 MB at T=16,384 cp=2. Everything above that has been a
|
| 4 |
+
# 2-point fit whose second point I can no longer verify.
|
| 5 |
+
set -uo pipefail
|
| 6 |
+
cd "$(dirname "$0")"
|
| 7 |
+
ROOT="$(cd ../.. && pwd)"
|
| 8 |
+
export LD_LIBRARY_PATH="$ROOT/Tensor-Implementations/lib:$ROOT/Profiler/lib:${LD_LIBRARY_PATH:-}"
|
| 9 |
+
export OMPI_ALLOW_RUN_AS_ROOT=1 OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1
|
| 10 |
+
export PATH=/usr/local/cuda/bin:$PATH
|
| 11 |
+
EXE=./build/bluscriptCP_exec
|
| 12 |
+
ARCH="CP_N_EMBD=1536 CP_N_LAYER=24 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=1"
|
| 13 |
+
CK=$PWD/checkpoints_680m/rank_0/varybatch.ckpt
|
| 14 |
+
probe(){ # T SCALE
|
| 15 |
+
local T=$1 S=$2
|
| 16 |
+
local L="memprobe_${T}.log"
|
| 17 |
+
echo "=== probe T=$T scale=$S cp=2 ($((T/2)) tok/rank) ==="
|
| 18 |
+
env $ARCH CP_SIZE=2 CP_ATTN_MODE=ring CP_ATTN_FUSION=split CP_T=$T \
|
| 19 |
+
YARN_SCALE=$S YARN_ORIG_MAXPOS=4096 CP_ROPE_THETA=500000 \
|
| 20 |
+
CP_ZERO=1 CP_ZERO_V2=1 CP_BF16=1 CP_OPTIMIZER=muon \
|
| 21 |
+
CP_WARMUP=2 CP_STABLE=0 CP_CKPT=0 CP_INIT_FROM="$CK" \
|
| 22 |
+
CP_DATA_ROOT="$PWD/Data_Loader/Data/evalsrc" \
|
| 23 |
+
CP_B=1 CP_GLOBAL_BATCH=$T CP_MAX_STEPS=3 \
|
| 24 |
+
CUDA_VISIBLE_DEVICES=0,1 mpirun --allow-run-as-root -np 2 "$EXE" > "$L" 2>&1
|
| 25 |
+
local rc=$?
|
| 26 |
+
if grep -qaiE 'out of memory|bad_alloc' "$L"; then
|
| 27 |
+
echo " OOM -> does NOT fit on 2x H100"
|
| 28 |
+
elif [ $rc -eq 0 ] && grep -qaE '^step +[0-9]+ \|' "$L"; then
|
| 29 |
+
echo " FITS: peak $(grep -aoE 'mem: [0-9]+MB' "$L" | tail -1) of 81559 MiB/card"
|
| 30 |
+
echo " step time $(grep -aoE 'dt: [0-9.]+ms' "$L" | tail -1)"
|
| 31 |
+
else
|
| 32 |
+
echo " FAILED rc=$rc"; grep -aA1 'terminate called' "$L" | grep -a 'what()' | head -1
|
| 33 |
+
fi
|
| 34 |
+
}
|
| 35 |
+
probe 32768 8
|
| 36 |
+
probe 65536 16
|