File size: 7,549 Bytes
63c7b80 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 | """Smoke test the planner against the real frozen LeWM before any long run.
Section 5 of the brief: tiny batch, ~50 steps, H=3, T=1. Verify shapes, verify
both silent-failure invariants, verify one optimizer step reduces the loss, and
measure memory and it/s — ``H*T*n`` applications of ``f`` per sample is
expensive and the throughput here is what the time estimate is built from.
The unit tests cover the same invariants against a toy stand-in so they need no
checkpoint. This runs them against the actual world model, where the predictor
is a real 6-layer transformer and the rollout is where the memory goes.
"""
import argparse
import json
import sys
import time
from pathlib import Path
import torch
from torch.utils.data import DataLoader
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from lejepa_control.data import LatentGoalDataset, split_episodes # noqa: E402
from lejepa_control.world_model import load_lewm # noqa: E402
from lejepa_control_2.losses import DistanceScale, planner_loss # noqa: E402
from lejepa_control_2.scripts.train_planner import ( # noqa: E402
build_planner,
parse_args,
)
def check_invariants(planner, model, batch, scale, device, horizon, cycles):
"""The two failures that are silent, plus the activation count."""
planner.reset_call_counts()
out = planner(
model,
batch['context'],
batch['past_actions'],
batch['goal'],
horizon=horizon,
cycles=cycles,
)
B = batch['context'].size(0)
N = planner.num_context
shapes = {
'distances': (B, horizon),
'blocks': (B, horizon, planner.block_dim),
'raw': (B, horizon, planner.block_dim),
'contexts': (B, horizon, N, planner.latent_dim),
'answers': (B, horizon, planner.width),
'frames': (B, N + horizon, planner.latent_dim),
}
for key, want in shapes.items():
got = tuple(out[key].shape)
assert got == want, f'{key}: shape {got}, expected {want}'
assert torch.isfinite(out[key]).all(), f'{key} is not finite'
print(f' shapes ok ({len(shapes)} tensors, all finite)')
# invariant 1: the h-chain is differentiable all the way back to block 0
terminal = out['distance_seq'][-1]
assert terminal.requires_grad, 'terminal distance carries no gradient'
norms = []
for k in range(horizon):
g = torch.autograd.grad(
terminal.sum(), out['block_seq'][k], retain_graph=True,
allow_unused=True,
)[0]
assert g is not None and g.abs().sum() > 0, (
f'terminal distance does not reach block {k} — the h-chain is cut '
f'and the planner has collapsed to a greedy one-step policy'
)
norms.append(round(g.norm().item(), 6))
print(f' invariant 1 ok: d[-1] reaches every block, norms {norms}')
# invariant 3: the Change-1 detach schedule is actually in effect
expected = horizon * (planner.inner + 1)
unstaged = horizon * cycles * (planner.inner + 1)
assert planner.grad_calls == expected, (
f'grad-enabled f/g applications = {planner.grad_calls}, expected '
f'{expected} = H(n+1). Full backprop would be {unstaged}.'
)
print(
f' invariant 3 ok: {planner.grad_calls} grad-enabled applications '
f'= H(n+1); full backprop would be {unstaged}'
)
loss, _ = planner_loss(out, batch['goal_offset'], scale,
planner.action_embed)
loss.backward()
# invariant 2: the world model is frozen but transparent
leaked = [n for n, p in model.named_parameters() if p.grad is not None]
assert not leaked, f'frozen world model accumulated gradients: {leaked[:5]}'
assert planner.f.body.out.weight.grad.abs().sum() > 0, (
'no gradient reached f — the world model is opaque, not just frozen'
)
print(
f' invariant 2 ok: 0/{sum(1 for _ in model.parameters())} world-model '
f'params have grads, and f still trains'
)
planner.zero_grad(set_to_none=True)
def main():
argv = sys.argv[1:]
parser = argparse.ArgumentParser(add_help=False)
parser.add_argument('--smoke-steps', type=int, default=50)
smoke_args, rest = parser.parse_known_args(argv)
# tiny by default; anything on the command line still wins
defaults = ['--stage', 'A', '--batch-size', '8', '--pretrain-steps', '20']
args = parse_args(defaults + rest)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
torch.manual_seed(args.seed)
stats = json.loads((Path(args.latents) / 'stats.json').read_text())
model = load_lewm(device=device)
planner = build_planner(args, stats, stats['latent_dim'], device)
scale = DistanceScale().to(device)
horizon, cycles = 3, args.cycles
print(
f'smoke: H={horizon} T={cycles} n={args.inner} B={args.batch_size} '
f'on {device}, planner '
f'{sum(p.numel() for p in planner.parameters()) / 1e6:.2f}M params'
)
train_eps, _ = split_episodes(stats['n_episodes'])
dataset = LatentGoalDataset(
args.latents, episodes=train_eps, horizon=horizon, max_offset=3
)
loader = DataLoader(
dataset, batch_size=args.batch_size, shuffle=True, drop_last=True
)
raw = next(iter(loader))
batch = {k: v.to(device) for k, v in raw.items()}
scale.update(batch['context'], batch['goal'])
print('invariants:')
check_invariants(planner, model, batch, scale, device, horizon, cycles)
# does one optimizer step actually reduce the loss?
opt = torch.optim.AdamW(
planner.parameters(), lr=args.lr, weight_decay=args.weight_decay
)
if device == 'cuda':
torch.cuda.reset_peak_memory_stats()
losses = []
t0 = time.perf_counter()
for _ in range(smoke_args.smoke_steps):
scale.update(batch['context'], batch['goal'])
out = planner(
model, batch['context'], batch['past_actions'], batch['goal'],
horizon=horizon, cycles=cycles,
)
loss, metrics = planner_loss(
out, batch['goal_offset'], scale, planner.action_embed,
lambda_cycle=args.lambda_cycle,
lambda_support=args.lambda_support,
lambda_anchor=args.lambda_anchor,
lambda_sat=args.lambda_sat,
)
opt.zero_grad(set_to_none=True)
loss.backward()
torch.nn.utils.clip_grad_norm_(planner.parameters(), args.grad_clip)
opt.step()
losses.append(loss.item())
if device == 'cuda':
torch.cuda.synchronize()
elapsed = time.perf_counter() - t0
print(
f'overfit {smoke_args.smoke_steps} steps on a fixed batch: '
f'{losses[0]:.4f} -> {losses[-1]:.4f}'
)
assert losses[-1] < losses[0], 'loss did not fall on a fixed batch'
rate = smoke_args.smoke_steps / elapsed
peak = (
torch.cuda.max_memory_allocated() / 2**30 if device == 'cuda' else 0.0
)
print(
f'throughput {rate:.2f} it/s at B={args.batch_size} H={horizon} '
f'T={cycles}, peak VRAM {peak:.2f} GiB'
)
# what the real run costs: the driver's cost scales with H*T, and stages
# B-E run at B=32
scale_factor = (args.batch_size / 32) * (3 / horizon) * (1 / max(cycles, 1))
print(
f' extrapolated to B=32 H=3 T=3: '
f'~{rate * scale_factor * 3:.2f} it/s '
f'({20000 / max(rate * scale_factor * 3, 1e-6) / 60:.0f} min for 20k)'
)
print('SMOKE TEST PASSED')
if __name__ == '__main__':
main()
|