Add constrained DecMem STM6/topk27 Minecraft evaluation

Add a read-budget profile preserving native patch retrieval, weights, sampler and four-latent blocks. Native settings remain the default. The new eight-H200 seed42/43 launcher runs a36-frame preflight, then300clips with600observations and500predictions each. Runtime metadata records actual settings and logical/backing KVstorage; this is not an exact whole-frame or persistent-storage cap. CPU selection/cache/interface/CLI and shell checks passed; GPU checks run in the submitted preflight.

DecMem/evaluation/worldmem_adapter.py

@@ -45,7 +45,7 @@ def prepare_controls(case, target_frames):
 
 
 class DecMemAdapter:
-    def __init__(self, checkpoint, device="cuda"):
+    def __init__(self, checkpoint, device="cuda", budget=False):
         self.checkpoint = str(Path(checkpoint).resolve())
         # Upstream model and VAE loading use paths relative to the DecMem checkout.
         os.chdir(REPO)
@@ -62,10 +62,29 @@ class DecMemAdapter:
         self.config.num_frames = self.config.image_or_video_shape[1]
         self.config.prope_max_time = 60
         self.config.rope_scale_pattern = "rope"
+        if budget:
+            self.config.model_kwargs.stm_sliding_window = 6
+            self.config.model_kwargs.ltm_topk = 27
         self.pipeline = WorldMemCausalDiffusionInferencePipeline(self.config, device=self.device)
         checkpoint_state = torch.load(self.checkpoint, map_location="cpu", weights_only=True, mmap=True)
         self.pipeline.generator.load_state_dict(checkpoint_state["generator"], strict=True)
         self.pipeline.to(device=self.device, dtype=self.dtype).eval().requires_grad_(False)
+        model = self.pipeline.generator.model
+        window, topk = model.stm_sliding_window, model.ltm_topk
+        if budget and any((a.stm_sliding_window, a.ltm_topk) != (6, 27) for a in [model, *(b.self_attn for b in model.blocks)]):
+            raise ValueError("DecMem read budget was not applied to every attention layer")
+        self.memory = {
+            "profile": "stm6_topk27" if budget else "native", "use_ltm_attn": model.use_ltm_attn,
+            "stm_sliding_window": window, "ltm_topk": topk, "ltm_tile_size": list(model.ltm_tile_size),
+            "num_frame_per_block": self.pipeline.num_frame_per_block,
+            "cached_past_latents": window-self.pipeline.num_frame_per_block,
+            "tokens_per_latent": self.pipeline.frame_seq_length,
+            "dense_key_tokens_per_query": window*self.pipeline.frame_seq_length,
+            "fine_key_tokens_per_query_upper_bound": topk*math.prod(model.ltm_tile_size),
+            "scope": "dense local context and sparse fine-attention reads; not a whole-frame or storage cap",
+            "candidate_history": "all committed history retained; coarse ranking includes all history and current block",
+            "temporal_caveat": "4x temporal VAE and 4-latent joint generation; first noisy block has 3 observed latents",
+        }
         self.parameter_counts = {
             "generator_total": sum(p.numel() for p in self.pipeline.generator.parameters()),
             "generator_trainable": sum(p.numel() for p in self.pipeline.generator.parameters() if p.requires_grad),
@@ -109,6 +128,12 @@ class DecMemAdapter:
             return_latents=True,
             skip_vae_decode=True,
         )
+        # STM slices retain their backing allocations; report both sizes before releasing KV.
+        caches = [value for layer in self.pipeline.kv_cache for value in layer.values() if torch.is_tensor(value)]
+        storages = {(value.device, value.untyped_storage().data_ptr()): value.untyped_storage().nbytes() for value in caches}
+        cache_usage = {"final_logical_bytes": sum(value.numel()*value.element_size() for value in caches),
+                       "final_storage_bytes": sum(storages.values())}
+        del caches, storages
         # Denoising is finished. Release global KV before decoding the complete causal timeline.
         self.pipeline.kv_cache = None
         self.pipeline.crossattn_cache = None
@@ -125,7 +150,7 @@ class DecMemAdapter:
             "batch_size": 1,
             "cfg_batch_size": 2,
             "sampling": {"solver": "UniPC", "steps": 20, "guidance_scale": 5.0, "timestep_shift": 5.0, "source": "scripts/infer_example.sh"},
-            "memory": {"use_ltm_attn": True, "ltm_topk": 80, "ltm_tile_size": [1, 8, 8], "stm_sliding_window": 8, "num_frame_per_block": 4},
+            "memory": {**self.memory, **cache_usage},
             "history_frames": history,
             "history_source_indices": case["source_indices"][:history].tolist(),
             "history_left_padding": padding,

deployment/evaluate_decmem_budget_8h200.sh

@@ -0,0 +1,87 @@
+#!/usr/bin/env bash
+#SBATCH --job-name=decmem-minecraft-stm6-topk27-full
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=8
+#SBATCH --cpus-per-task=14
+#SBATCH --gres=gpu:8
+#SBATCH --time=48:00:00
+#SBATCH --account=berzelius-2026-196
+#SBATCH --array=42-43%1
+#SBATCH --chdir=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals/DecMem
+#SBATCH --output=slurm_logs/%x-%A_%a.out
+#SBATCH --error=slurm_logs/%x-%A_%a.out
+
+set -euo pipefail
+module load buildenv-gcccuda/12.9.1-gcc11
+
+ROOT=/proj/cvl/users/x_fahkh2/worldmem-baseline-evals
+ENV_ROOT=/proj/cvl/users/x_fahkh2/envs
+REPO=$ROOT/DecMem
+PYTHON=$ENV_ROOT/decmem-eval/bin/python
+METRIC_PYTHON=$ENV_ROOT/geometry-forcing-eval/bin/python
+WORLD_MEM_ROOT=$ROOT/references/worldmem
+DATA_ROOT=/proj/cvl/users/x_fahkh2/WorldMem_Repro/datasets/minecraft
+MANIFEST=$ROOT/deployment/manifests/minecraft_test_300.csv
+FID_MANIFEST=$ROOT/deployment/manifests/minecraft_test_fid50.csv
+SEED=$SLURM_ARRAY_TASK_ID
+JOB=$SLURM_ARRAY_JOB_ID
+RUN=decmem_stm6_topk27_minecraft_300cases_seed${SEED}_job${JOB}
+OUTPUT=$REPO/outputs/$RUN
+PREFLIGHT_OUTPUT=$REPO/outputs/${RUN}_preflight
+
+export TMPDIR=$ROOT/cache TRITON_CACHE_DIR=$ROOT/cache
+export HF_HOME=$ROOT/cache/huggingface TORCH_HOME=$ROOT/cache/torch
+export PIP_CACHE_DIR=$ROOT/cache/pip CUDA_HOME=$CUDA_ROOT
+export XDG_CACHE_HOME=$ROOT/cache MPLCONFIGDIR=$ROOT/cache/matplotlib
+export PYTHONNOUSERSITE=1 PYTHONUNBUFFERED=1
+export PYTHONPATH=$REPO
+export OMP_NUM_THREADS=4 MKL_NUM_THREADS=4
+export WANDB_MODE=disabled HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1
+mkdir -p "$TMPDIR" "$OUTPUT"
+cd "$REPO"
+
+"$PYTHON" "$ROOT/deployment/verify_node.py" --output "$OUTPUT"
+cp "$0" "$OUTPUT/launcher.sh"
+cp "$REPO/evaluation/worldmem_adapter.py" "$OUTPUT/worldmem_adapter.py"
+
+# Read-budget ablation: STM6/top-k27 keeps all candidate history and block size4.
+# Exercise three denoised blocks and the complete scoring path before the full run.
+srun --exclusive --exact --nodes=1 --ntasks=1 --gpus-per-task=1 --cpus-per-task=14 --cpu-bind=cores \
+    "$PYTHON" "$ROOT/shared/run_and_score.py" \
+    --method decmem --decmem-budget \
+    --workspace-root "$ROOT" \
+    --metric-python "$METRIC_PYTHON" \
+    --worldmem-root "$WORLD_MEM_ROOT" \
+    --manifest "$MANIFEST" \
+    --data-root "$DATA_ROOT" \
+    --output "$PREFLIGHT_OUTPUT" \
+    --seed "$SEED" \
+    --limit 1 --max-generated 36 --rank 0 --world-size 1
+
+srun --exclusive --exact --nodes=1 --ntasks=1 --gpus-per-task=1 --cpus-per-task=14 \
+    "$METRIC_PYTHON" "$ROOT/shared/aggregate.py" \
+    --output "$PREFLIGHT_OUTPUT" \
+    --manifest "$MANIFEST" \
+    --dataset minecraft \
+    --seed "$SEED" \
+    --expected-frames 36 --limit 1
+
+srun --ntasks=8 --gpus-per-task=1 --cpus-per-task=14 --cpu-bind=cores --kill-on-bad-exit=1 \
+    "$PYTHON" "$ROOT/shared/run_and_score.py" \
+    --method decmem --decmem-budget \
+    --workspace-root "$ROOT" \
+    --metric-python "$METRIC_PYTHON" \
+    --worldmem-root "$WORLD_MEM_ROOT" \
+    --manifest "$MANIFEST" \
+    --data-root "$DATA_ROOT" \
+    --output "$OUTPUT" \
+    --seed "$SEED" \
+    --limit 300 --max-generated 500 --world-size 8
+
+srun --exclusive --exact --nodes=1 --ntasks=1 --gpus-per-task=1 --cpus-per-task=14 \
+    "$METRIC_PYTHON" "$ROOT/shared/aggregate.py" \
+    --output "$OUTPUT" \
+    --manifest "$MANIFEST" \
+    --dataset minecraft \
+    --fid-manifest "$FID_MANIFEST" \
+    --seed "$SEED"

shared/run_and_score.py

@@ -15,7 +15,7 @@ import torch
 from protocol import case_key, case_seed, load_minecraft_case, load_re10k_case, save_case, select_cases
 
 
-def load_adapter(method, root):
+def load_adapter(method, root, decmem_budget=False):
     entries = {
         "live": ("LIVE", "scripts/evaluate_worldmem.py", "LivePredictor"),
         "geometry-forcing": ("GeometryForcing", "evaluation/worldmem_adapter.py", "GeometryForcingAdapter"),
@@ -35,7 +35,7 @@ def load_adapter(method, root):
     if method == "geometry-forcing":
         return constructor(root / "checkpoints/geometry-forcing/geometry_forcing_state_dict.ckpt")
     if method == "decmem":
-        return constructor(root / "checkpoints/decmem/decmem.pt")
+        return constructor(root / "checkpoints/decmem/decmem.pt", budget=decmem_budget)
     return constructor(root / "checkpoints/matrix-game2", mouse_scale=1.0)
 
 
@@ -56,6 +56,7 @@ def main():
     parser.add_argument("--metric-python", type=Path, required=True)
     parser.add_argument("--vae-checkpoint", type=Path)
     parser.add_argument("--keep-tensor", action="store_true")
+    parser.add_argument("--decmem-budget", action="store_true", help="DecMem STM6/topk27 read-budget ablation")
     parser.add_argument("--dry-run", action="store_true")
     args = parser.parse_args()
     root = args.workspace_root.resolve()
@@ -71,10 +72,11 @@ def main():
     rows = select_cases(args.manifest, args.scope, args.rank, args.world_size, args.limit)
     if args.dry_run:
         print(json.dumps({"method": args.method, "rank": args.rank, "cases": [case_key(row) for row in rows],
-                          "seed": args.seed, "generated_frames": args.max_generated}, indent=2))
+                          "seed": args.seed, "generated_frames": args.max_generated,
+                          "decmem_budget": args.decmem_budget}, indent=2))
         return
     torch.cuda.set_device(0)
-    adapter = load_adapter(args.method, root)
+    adapter = load_adapter(args.method, root, args.decmem_budget)
     args.output.mkdir(parents=True, exist_ok=True)
     metadata = getattr(adapter, "metadata", getattr(adapter, "provenance", {}))
     if args.method in {"decmem", "matrix-game2"}:
@@ -83,6 +85,14 @@ def main():
         metadata = {**json.loads(manifest_path.read_text()),
                     "instantiated_parameter_counts": adapter.parameter_counts,
                     "runtime_strict_checkpoint_load": "passed"}
+        if args.method == "decmem":
+            for key in ("verification", "full_ready", "blocker"):
+                metadata.pop(key, None)
+            metadata["benchmark_protocol"] = metadata.pop("protocol")
+            metadata["generated_frames_per_case"] = args.max_generated
+            metadata["model"].update({key: adapter.memory[key] for key in ("stm_sliding_window", "ltm_topk")})
+            metadata["memory"] = adapter.memory
+            metadata["profile"] = adapter.memory["profile"]
     (args.output / f"method_manifest_rank{args.rank}.json").write_text(json.dumps(metadata, indent=2, default=str) + "\n")
     for row in rows:
         case = (load_minecraft_case if dataset == "minecraft" else load_re10k_case)(row, args.data_root)