| """Patch T5 specroute gen_scripts to add P100 GPU detection and BSZ.""" |
| import re, os |
|
|
| BASE = '/Users/nnminh322/Desktop/personal/Continual/improve_gainlora' |
|
|
| T5_SCRIPTS = [ |
| os.path.join(BASE, 'gen_script_superni_order1_t5_specroute.sh'), |
| os.path.join(BASE, 'gen_script_superni_order2_t5_specroute.sh'), |
| os.path.join(BASE, 'gen_script_long_order3_t5_specroute.sh'), |
| os.path.join(BASE, 'gen_script_long_order4_t5_specroute.sh'), |
| ] |
|
|
| GPU_OLD = ( |
| 'else\n' |
| ' GPU_MODE="a100"\n' |
| ' GPU_IDS="${1:-0}"\n' |
| ' FP16_FLAG=""\n' |
| ' echo "[GPU] Strategy: A100 (single GPU, fp32)"\n' |
| 'fi' |
| ) |
|
|
| GPU_NEW = ( |
| 'elif [ "$GPU_MEM" -gt 16000 ]; then\n' |
| ' GPU_MODE="p100"\n' |
| ' GPU_IDS="${1:-0}"\n' |
| ' FP16_FLAG="--gradient_checkpointing"\n' |
| ' echo "[GPU] Strategy: P100 16GB (fp32 + gradient_checkpointing)"\n' |
| 'else\n' |
| ' GPU_MODE="a100"\n' |
| ' GPU_IDS="${1:-0}"\n' |
| ' FP16_FLAG=""\n' |
| ' echo "[GPU] Strategy: A100 (single GPU, fp32)"\n' |
| 'fi' |
| ) |
|
|
| BSZ_PAT = re.compile( |
| r'(elif \[ "\$GPU_MODE" = "t4_1gpu" \]; then\n BSZ=\d+; GA=\d+; EVAL_BSZ=\d+\n)' |
| r'(else\n BSZ=\d+; GA=\d+; EVAL_BSZ=\d+\n)' |
| ) |
|
|
| def add_p100(m): |
| return ( |
| m.group(1) |
| + 'elif [ "$GPU_MODE" = "p100" ]; then\n BSZ=8; GA=4; EVAL_BSZ=4\n' |
| + m.group(2) |
| ) |
|
|
| for name in T5_SCRIPTS: |
| if not os.path.exists(name): |
| print(f'SKIP (not found): {name}') |
| continue |
| with open(name) as f: |
| c = f.read() |
| n_detect = c.count(GPU_OLD) |
| c = c.replace(GPU_OLD, GPU_NEW, 1) |
| n_bsz = len(BSZ_PAT.findall(c)) |
| c = BSZ_PAT.sub(add_p100, c) |
| with open(name, 'w') as f: |
| f.write(c) |
| print(f'{name}: gpu_detect={n_detect} bsz_blocks={n_bsz}') |
|
|
| print('Done.') |
|
|
