atakan Claude Sonnet 5 commited on
Commit
a22c85e
·
1 Parent(s): 72bc69e

fix: Explicitly .to('cuda') instead of device_map='auto' on ZeroGPU

Browse files

Strong 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>

Files changed (1) hide show
  1. controlai_agent/orchestrator.py +12 -4
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
- dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 
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 torch.cuda.is_available() else None
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: