Spaces:
Running on Zero
Running on Zero
atakan Claude Sonnet 5 commited on
Commit ·
a22c85e
1
Parent(s): 72bc69e
fix: Explicitly .to('cuda') instead of device_map='auto' on ZeroGPU
Browse filesStrong suspicion for the ~5 minute hangs on a genuinely fast GPU
(RTX Pro 6000): device_map="auto" infers available VRAM at model-load
time, which on ZeroGPU happens outside any @spaces.GPU call -- no
physical GPU is attached to the process yet at that point, so
Accelerate may have been silently placing some or all layers on CPU
despite torch.cuda.is_available() reporting True. Replaced with a
plain load + explicit .to('cuda') and log the resulting device so this
is verifiable from the logs instead of inferred.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
controlai_agent/orchestrator.py
CHANGED
|
@@ -220,18 +220,26 @@ class ControlAIAgent:
|
|
| 220 |
hf_id = CONTROLAI_HF_REPO if "mlx" in str(model_path) or str(model_path).startswith("Qwen/") else model_path
|
| 221 |
print(f"Loading PyTorch model: {hf_id} (threads: {num_threads})...")
|
| 222 |
self.hf_tokenizer = AutoTokenizer.from_pretrained(hf_id, trust_remote_code=True)
|
| 223 |
-
|
|
|
|
| 224 |
# SDPA is a large, free speedup over the "eager" attention
|
| 225 |
# default -- meaningful for the long tool-schema prompt prefix.
|
| 226 |
-
attn_impl = "sdpa" if
|
| 227 |
self.model = AutoModelForCausalLM.from_pretrained(
|
| 228 |
hf_id,
|
| 229 |
torch_dtype=dtype,
|
| 230 |
low_cpu_mem_usage=True,
|
| 231 |
-
device_map="auto",
|
| 232 |
trust_remote_code=True,
|
| 233 |
attn_implementation=attn_impl,
|
| 234 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 235 |
# The fused ControlAI model already has the LoRA weights merged in;
|
| 236 |
# only apply a separate adapter when loading a plain base model.
|
| 237 |
if hf_id != CONTROLAI_HF_REPO and adapter_path and Path(adapter_path).exists():
|
|
@@ -241,7 +249,7 @@ class ControlAIAgent:
|
|
| 241 |
print(f"Loaded PEFT LoRA adapter from: {adapter_path}")
|
| 242 |
except Exception as exc:
|
| 243 |
print(f"Warning: Could not load LoRA adapter in PyTorch: {exc}")
|
| 244 |
-
print("PyTorch model loaded successfully.")
|
| 245 |
|
| 246 |
# Initialize local offline RAG index
|
| 247 |
try:
|
|
|
|
| 220 |
hf_id = CONTROLAI_HF_REPO if "mlx" in str(model_path) or str(model_path).startswith("Qwen/") else model_path
|
| 221 |
print(f"Loading PyTorch model: {hf_id} (threads: {num_threads})...")
|
| 222 |
self.hf_tokenizer = AutoTokenizer.from_pretrained(hf_id, trust_remote_code=True)
|
| 223 |
+
has_cuda = torch.cuda.is_available()
|
| 224 |
+
dtype = torch.bfloat16 if has_cuda else torch.float32
|
| 225 |
# SDPA is a large, free speedup over the "eager" attention
|
| 226 |
# default -- meaningful for the long tool-schema prompt prefix.
|
| 227 |
+
attn_impl = "sdpa" if has_cuda else None
|
| 228 |
self.model = AutoModelForCausalLM.from_pretrained(
|
| 229 |
hf_id,
|
| 230 |
torch_dtype=dtype,
|
| 231 |
low_cpu_mem_usage=True,
|
|
|
|
| 232 |
trust_remote_code=True,
|
| 233 |
attn_implementation=attn_impl,
|
| 234 |
)
|
| 235 |
+
if has_cuda:
|
| 236 |
+
# device_map="auto" infers available VRAM at load time; on
|
| 237 |
+
# ZeroGPU no physical GPU is attached to the process yet at
|
| 238 |
+
# this point (that only happens inside a @spaces.GPU call),
|
| 239 |
+
# so it can silently offload layers to CPU. An explicit
|
| 240 |
+
# move avoids that and any ambiguity about what actually
|
| 241 |
+
# ran where.
|
| 242 |
+
self.model = self.model.to("cuda")
|
| 243 |
# The fused ControlAI model already has the LoRA weights merged in;
|
| 244 |
# only apply a separate adapter when loading a plain base model.
|
| 245 |
if hf_id != CONTROLAI_HF_REPO and adapter_path and Path(adapter_path).exists():
|
|
|
|
| 249 |
print(f"Loaded PEFT LoRA adapter from: {adapter_path}")
|
| 250 |
except Exception as exc:
|
| 251 |
print(f"Warning: Could not load LoRA adapter in PyTorch: {exc}")
|
| 252 |
+
print(f"PyTorch model loaded successfully. Model device: {next(self.model.parameters()).device}")
|
| 253 |
|
| 254 |
# Initialize local offline RAG index
|
| 255 |
try:
|