nils/jspace
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

warm-start flag (stacking arm), adaptive flags for Blocksworld

Browse Source 
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Nils
2026-07-14 13:02:46 +02:00
co-authored by Claude Fable 5
parent aa3e11fb20
commit b461970c9a
3 changed files with 15 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/train_merge_code.py
+6 -1
View File
@@ -51,6 +51,7 @@ ap.add_argument("--deepk", type=int, default=0,
ap.add_argument("--bptt", type=int, default=0,
help="truncated BPTT: grads only through last N iterations")
ap.add_argument("--lr", type=float, default=1e-3)
ap.add_argument("--warm", default=None, help="warm-start adapter checkpoint")
ARGS, _ = ap.parse_known_args()
SEED = ARGS.seed
LR = ARGS.lr
@@ -59,7 +60,8 @@ SUFFIX = ((f"_s{SEED}" if SEED else "")
+ (f"_a{ARGS.alpha}" if ARGS.alpha != 0.3 else "")
+ ("_ad" if ARGS.adaptive else "")
+ (f"_dk{ARGS.deepk}" if ARGS.deepk else "")
+ (f"_lr{ARGS.lr}" if ARGS.lr != 1e-3 else ""))
+ (f"_lr{ARGS.lr}" if ARGS.lr != 1e-3 else "")
+ ("_warm" if ARGS.warm else ""))
PAUSE_ID = 6 # <unused0>
@@ -122,6 +124,9 @@ def main():
adapter = AdaptiveMergeAdapter(d=d, alpha0=ARGS.alpha).cuda()
else:
adapter = MergeAdapter(d=d, alpha=ARGS.alpha).cuda()
if ARGS.warm:
adapter.load_state_dict(torch.load(ARGS.warm, map_location="cuda"))
print("warm-started from", ARGS.warm, flush=True)
global K_BUCKETS
if ARGS.deepk:
f = ARGS.deepk / 4