JonathanColetti commited on
Commit
68f4fd7
·
verified ·
1 Parent(s): 471e398

release the world-decode cache before the block loop

Browse files
Files changed (1) hide show
  1. app.py +8 -2
app.py CHANGED
@@ -265,8 +265,14 @@ def run(world: str, steer_to: str, steer_at: float, seconds: float, seed: int):
265
 
266
  t0 = time.perf_counter()
267
  engine.start(world=WORLDS[world], seed=int(seed))
268
- _probe("after-world") # world decoded and the K/V cache allocated: the state the
269
- # first FFN linear fails in
 
 
 
 
 
 
270
  frames, pending_steer, last_push = [], steer_idx is not None, 0.0
271
 
272
  # The worker runs in a thread and reports failure by setting status.state rather
 
265
 
266
  t0 = time.perf_counter()
267
  engine.start(world=WORLDS[world], seed=int(seed))
268
+ # Opening a world decodes all 81 of its pixel frames in one call, which is by far
269
+ # the largest allocation in a run and leaves the caching allocator holding blocks
270
+ # the block loop cannot reuse (reserved 30.5 G against 13.5 G allocated). Growing
271
+ # past that OOMs, and a ZeroGPU worker cannot even format the OOM message --
272
+ # PyTorch asks NVML which processes hold memory, NVML is not available there, and
273
+ # the run dies on an internal assert instead. Hand the cache back first.
274
+ torch.cuda.empty_cache()
275
+ _probe("after-world")
276
  frames, pending_steer, last_push = [], steer_idx is not None, 0.0
277
 
278
  # The worker runs in a thread and reports failure by setting status.state rather