Make standalone runner Colab-safe (disable device_map by default)

run_standalone_inference.py

@@ -51,6 +51,11 @@ def parse_args() -> argparse.Namespace:
     p.add_argument("--max-new-tokens", type=int, default=96)
     p.add_argument("--seed", type=int, default=42)
     p.add_argument("--device", choices=["auto", "cuda", "cpu"], default="auto")
+    p.add_argument(
+        "--use-device-map",
+        action="store_true",
+        help="Use transformers/accelerate device_map='auto'. Disabled by default for maximum Colab compatibility.",
+    )
     p.add_argument("--prompts-file", default=None, help="Optional .txt file with one prompt per line")
     return p.parse_args()
 
@@ -203,9 +208,10 @@ def main() -> None:
         local,
         trust_remote_code=True,
         torch_dtype=torch_dtype,
-        device_map="auto" if device.type == "cuda" else None,
+        device_map="auto" if (device.type == "cuda" and args.use_device_map) else None,
+        low_cpu_mem_usage=bool(device.type == "cuda" and args.use_device_map),
     )
-    if device.type == "cpu":
+    if device.type == "cpu" or (device.type == "cuda" and not args.use_device_map):
         model = model.to(device)
     model.eval()
 
@@ -227,7 +233,7 @@ def main() -> None:
     max_prompt_tokens = max(1, context_window - args.max_new_tokens)
 
     print(f"repo={args.repo_id} revision={args.revision} snapshot_sha={sha}")
-    print(f"device={device} dtype={torch_dtype} prompt_style={prompt_style}")
+    print(f"device={device} dtype={torch_dtype} prompt_style={prompt_style} use_device_map={args.use_device_map}")
     print(f"context_window={context_window} max_prompt_tokens={max_prompt_tokens}")
 
     torch.manual_seed(args.seed)