Buckets:
| """Out-of-memory retry check: a run where one unchecked micro-batch fails mid-backward must end with the same adapter.""" | |
| import sys, os | |
| sys.path.insert(0, "/workspace/decision-reasoner") | |
| os.chdir("/workspace/decision-reasoner/runs/v2") | |
| import torch, train | |
| from safetensors.torch import load_file | |
| inject = sys.argv[1] == "inject" | |
| out = f"/root/bench_tmp/oomtest/{sys.argv[1]}" | |
| orig = train.micro_loss | |
| calls = [0] | |
| def wrapped(forward, lm_weight, mb, *a, **k): | |
| calls[0] += 1 | |
| loss, st = orig(forward, lm_weight, mb, *a, **k) | |
| if inject and calls[0] == 2: # second micro-batch: earlier grads exist, fail part-way back | |
| model_params = [p for n, p in MODEL[0].named_parameters() if p.requires_grad and "layers.0." in n] | |
| h = [] | |
| def boom(g): | |
| for x in h: x.remove() | |
| raise torch.cuda.OutOfMemoryError("injected") | |
| h.append(model_params[0].register_hook(boom)) | |
| return loss, st | |
| train.micro_loss = wrapped | |
| MODEL = [None] | |
| orig_load = train.load_model | |
| def load(*a, **k): | |
| MODEL[0] = orig_load(*a, **k); return MODEL[0] | |
| train.load_model = load | |
| sys.argv = ["train.py", "--model", "/workspace/decision-reasoner/checkpoints/qwen3.5-0.8b-text", "--mode", "joint", "--alpha", "0.5", | |
| "--lr", "2e-4", "--subset", "data/sweep_ids.txt", "--readout-pool", "all", "--min-correct", "2", | |
| "--micro-tokens", "8192", "--ckpt-above", "8192", "--max-steps", "2", "--no-probe", "--out", out] | |
| train.main() | |
| print("calls", calls[0]) | |
Xet Storage Details
- Size:
- 1.51 kB
- Xet hash:
- be5ef73375677d25dcbe75b1e4443154e7c3cc429c2dee939c4d4a11e019b688
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.