From b81eef43c80474178ac1c1400c830225a0066a8d Mon Sep 17 00:00:00 2001
From: Pepijn <pepijn@huggingface.co>
Date: Thu, 16 Apr 2026 19:49:12 +0200
Subject: [PATCH] fix(profiling): wall_x OOM and xvla rename_map

- wall_x: switch to SGD optimizer + explicit scheduler overrides.
  The 4B-param model casts to bf16 internally, but AdamW's exp_avg/
  exp_avg_sq states blow past the 22 GB GPU. Same fix we applied to
  pi0/pi05/pi0_fast.
- xvla: fix rename_map. Dataset (libero_plus) exposes front/wrist
  image keys; the model expects image/image2. Previous map was
  direction-reversed and left the batch without any recognized
  image feature.

Made-with: Cursor
---
 profiling/model_profiling_specs.json | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/profiling/model_profiling_specs.json b/profiling/model_profiling_specs.json
index f30337663..9353995df 100644
--- a/profiling/model_profiling_specs.json
+++ b/profiling/model_profiling_specs.json
@@ -152,7 +152,17 @@
       "--policy.prediction_mode=diffusion",
       "--policy.attn_implementation=eager",
       "--policy.device=cuda",
-      "--batch_size=1"
+      "--batch_size=1",
+      "--use_policy_training_preset=false",
+      "--optimizer.type=sgd",
+      "--optimizer.lr=1e-5",
+      "--optimizer.weight_decay=0",
+      "--optimizer.grad_clip_norm=1.0",
+      "--scheduler.type=cosine_decay_with_warmup",
+      "--scheduler.peak_lr=1e-5",
+      "--scheduler.decay_lr=1e-6",
+      "--scheduler.num_warmup_steps=0",
+      "--scheduler.num_decay_steps=12"
     ]
   },
   "xvla": {
@@ -165,7 +175,7 @@
       "--policy.empty_cameras=1",
       "--policy.device=cuda",
       "--batch_size=1",
-      "--rename_map={\"observation.images.image\": \"observation.images.camera1\", \"observation.images.image2\": \"observation.images.camera2\"}"
+      "--rename_map={\"observation.images.front\": \"observation.images.image\", \"observation.images.wrist\": \"observation.images.image2\"}"
     ]
   }
 }