| #!/usr/bin/env bash |
| |
| |
| |
| set -uo pipefail |
| cd "$(dirname "$0")" |
| ROOT="$(cd ../.. && pwd)" |
| export LD_LIBRARY_PATH="$ROOT/Tensor-Implementations/lib:$ROOT/Profiler/lib:${LD_LIBRARY_PATH:-}" |
| export OMPI_ALLOW_RUN_AS_ROOT=1 OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1 |
| export PATH=/usr/local/cuda/bin:$PATH |
| EXE=./build/bluscriptCP_exec |
| ARCH="CP_N_EMBD=1536 CP_N_LAYER=24 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=1" |
| CK=$PWD/checkpoints_680m/rank_0/varybatch.ckpt |
| probe(){ |
| local T=$1 S=$2 |
| local L="memprobe_${T}.log" |
| echo "=== probe T=$T scale=$S cp=2 ($((T/2)) tok/rank) ===" |
| env $ARCH CP_SIZE=2 CP_ATTN_MODE=ring CP_ATTN_FUSION=split CP_T=$T \ |
| YARN_SCALE=$S YARN_ORIG_MAXPOS=4096 CP_ROPE_THETA=500000 \ |
| CP_ZERO=1 CP_ZERO_V2=1 CP_BF16=1 CP_OPTIMIZER=muon \ |
| CP_WARMUP=2 CP_STABLE=0 CP_CKPT=0 CP_INIT_FROM="$CK" \ |
| CP_DATA_ROOT="$PWD/Data_Loader/Data/evalsrc" \ |
| CP_B=1 CP_GLOBAL_BATCH=$T CP_MAX_STEPS=3 \ |
| CUDA_VISIBLE_DEVICES=0,1 mpirun --allow-run-as-root -np 2 "$EXE" > "$L" 2>&1 |
| local rc=$? |
| if grep -qaiE 'out of memory|bad_alloc' "$L"; then |
| echo " OOM -> does NOT fit on 2x H100" |
| elif [ $rc -eq 0 ] && grep -qaE '^step +[0-9]+ \|' "$L"; then |
| echo " FITS: peak $(grep -aoE 'mem: [0-9]+MB' "$L" | tail -1) of 81559 MiB/card" |
| echo " step time $(grep -aoE 'dt: [0-9.]+ms' "$L" | tail -1)" |
| else |
| echo " FAILED rc=$rc"; grep -aA1 'terminate called' "$L" | grep -a 'what()' | head -1 |
| fi |
| } |
| probe 32768 8 |
| probe 65536 16 |
|
|
