multistep: gradient checkpointing + tighter memory defaults

OOM in multi-turn rollouts because:
- Prompts grow each turn (turn 3 can hit 3-5k tokens with
prior_attempts folded into the chat user message)
- All B×G×turn_limit StepRecords sit in memory until the gradient
pass (24 records on default knobs)
- Three models co-resident on L40S (agent+LoRA, frozen target,
8-bit judge) leave only ~30 GB for activations + gradients

Fixes:
- --gradient-checkpointing default ON: ~80% activation memory
saved at ~30% extra compute. Critical for multi-step.
- --update-micro-batch 4 -> 2: half the activation memory per
backward
- --max-prompt-tokens 4096 -> 2048: drops the longest prior
turn first when chat prompt overflows
- --max-new-tokens 768 -> 384: half the per-turn generation cap
(bump back if thinking-mode answers truncate)
- hf_job_train_multistep.sh launcher exposes the new knobs as
env vars (MAX_PROMPT_TOKENS, MAX_NEW_TOKENS, UPDATE_MICRO_BATCH)
and passes them through

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

training/hf_job_train_multistep.sh

@@ -25,6 +25,11 @@ SFT_ADAPTER="${SFT_ADAPTER:-}"   # optional warmstart from a single-step adapter
 MAX_STEPS="${MAX_STEPS:-200}"
 NUM_GENS="${NUM_GENS:-4}"
 BATCH_SIZE="${BATCH_SIZE:-2}"
+# Memory-aware defaults — multi-turn prompts grow with prior_attempts
+# (turn 3 can hit 3-5k tokens), so tighter caps + grad checkpointing.
+MAX_PROMPT_TOKENS="${MAX_PROMPT_TOKENS:-2048}"
+MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-384}"
+UPDATE_MICRO_BATCH="${UPDATE_MICRO_BATCH:-2}"
 LR="${LR:-3e-6}"
 BETA="${BETA:-0.04}"
 TURN_LIMIT="${TURN_LIMIT:-3}"
@@ -88,6 +93,9 @@ python -u training/train_grpo_multistep.py \\
   --max-steps ${MAX_STEPS} \\
   --num-gens ${NUM_GENS} \\
   --batch-size ${BATCH_SIZE} \\
+  --max-prompt-tokens ${MAX_PROMPT_TOKENS} \\
+  --max-new-tokens ${MAX_NEW_TOKENS} \\
+  --update-micro-batch ${UPDATE_MICRO_BATCH} \\
   --lr ${LR} \\
   --beta ${BETA} \\
   --output-dir /app/outputs/grpo_multistep \\

training/train_grpo_multistep.py

@@ -235,11 +235,27 @@ def parse_args() -> argparse.Namespace:
     p.add_argument("--beta", type=float, default=0.04,
                    help="KL penalty vs frozen LoRA snapshot.")
     p.add_argument("--temperature", type=float, default=0.9)
-    p.add_argument("--max-new-tokens", type=int, default=768)
-    p.add_argument("--max-prompt-tokens", type=int, default=4096)
+    p.add_argument("--max-new-tokens", type=int, default=384,
+                   help="Per-turn agent generation cap. Trim from 768 to "
+                        "halve forward+backward memory. Bump back if "
+                        "thinking-mode answers get truncated.")
+    p.add_argument("--max-prompt-tokens", type=int, default=2048,
+                   help="Trim from 4096 — turn-3 prompts with "
+                        "prior_attempts can hit 3-5k tokens; truncating "
+                        "to 2k drops the longest prior turn first.")
     p.add_argument("--max-grad-norm", type=float, default=0.5)
-    p.add_argument("--update-micro-batch", type=int, default=4,
-                   help="Records per batched forward pass.")
+    p.add_argument("--update-micro-batch", type=int, default=2,
+                   help="Records per batched forward pass. 2 halves "
+                        "activation memory vs the default 4.")
+    p.add_argument("--gradient-checkpointing", action="store_true",
+                   default=True,
+                   help="Recompute forward activations during backward "
+                        "instead of caching. ~80%% activation memory "
+                        "saving at ~30%% extra compute. Default ON for "
+                        "multi-step because trajectory rollouts blow up "
+                        "activation memory.")
+    p.add_argument("--no-gradient-checkpointing",
+                   dest="gradient_checkpointing", action="store_false")
     p.add_argument("--save-every", type=int, default=50)
 
     # LoRA (used when --sft-adapter is not given — fresh LoRA init)
@@ -331,6 +347,26 @@ def main() -> None:
     n_tr = sum(p.numel() for p in model.parameters() if p.requires_grad)
     print(f"  trainable params: {n_tr:,}", flush=True)
 
+    # ---- Gradient checkpointing (default ON for multi-step) ----
+    # Saves ~80% activation memory at ~30% extra compute. Critical for
+    # multi-step because trajectory rollouts (B × G × turn_limit records)
+    # blow up activation memory during the backward pass.
+    if args.gradient_checkpointing:
+        # PEFT models need use_reentrant=False on modern PyTorch
+        try:
+            model.gradient_checkpointing_enable(
+                gradient_checkpointing_kwargs={"use_reentrant": False}
+            )
+        except TypeError:
+            # Older transformers/peft don't take the kwarg
+            model.gradient_checkpointing_enable()
+        # PEFT requires inputs to require grad when checkpointing the base
+        if hasattr(model, "enable_input_require_grads"):
+            model.enable_input_require_grads()
+        print("  gradient_checkpointing: ENABLED", flush=True)
+    else:
+        print("  gradient_checkpointing: disabled", flush=True)
+
     # ---- Snapshot trainable weights as the KL reference ----
     print("Snapshotting trainable weights as KL reference...", flush=True)
     ref_state: Dict[str, torch.Tensor] = {