flozi00 commited on
Commit
cfca239
·
verified ·
1 Parent(s): 2d55b77

zero-gpu: log the parameter device and VRAM so we can prove the weights are on the GPU

Browse files
Files changed (2) hide show
  1. BUILD.json +2 -2
  2. traceforge/space/local_llm.py +6 -0
BUILD.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
- "built_at": "2026-10-29T15:52:38",
3
- "commit": "e9f5ce0",
4
  "files": 53,
5
  "with_store": false,
6
  "with_card": true
 
1
  {
2
+ "built_at": "2026-10-29T16:33:11",
3
+ "commit": "8f3c1ee",
4
  "files": 53,
5
  "with_store": false,
6
  "with_card": true
traceforge/space/local_llm.py CHANGED
@@ -353,6 +353,9 @@ def generate(messages: list[dict], *, max_tokens: int = 512,
353
  stats["gpu"] = torch.cuda.get_device_name(0)
354
  stats["vram_gb"] = round(
355
  torch.cuda.get_device_properties(0).total_memory / 1e9)
 
 
 
356
  except Exception:
357
  pass
358
  LAST_GEN.update(stats)
@@ -442,6 +445,9 @@ class _Handler(BaseHTTPRequestHandler):
442
  f"{stats.get('seconds')}s {stats.get('tok_s')} tok/s "
443
  f"capped={stats.get('capped')} think={stats.get('think')} "
444
  f"skipped={stats.get('think_skipped')} "
 
 
 
445
  f"head={stats.get('raw_head')!r}", flush=True)
446
  self._send(200, {
447
  "id": "local-0",
 
353
  stats["gpu"] = torch.cuda.get_device_name(0)
354
  stats["vram_gb"] = round(
355
  torch.cuda.get_device_properties(0).total_memory / 1e9)
356
+ stats["param_device"] = str(next(model.parameters()).device)
357
+ stats["vram_used_gb"] = round(torch.cuda.memory_allocated() / 1e9, 1)
358
+ stats["attn"] = getattr(model.config, "_attn_implementation", "")
359
  except Exception:
360
  pass
361
  LAST_GEN.update(stats)
 
445
  f"{stats.get('seconds')}s {stats.get('tok_s')} tok/s "
446
  f"capped={stats.get('capped')} think={stats.get('think')} "
447
  f"skipped={stats.get('think_skipped')} "
448
+ f"dev={stats.get('param_device')} "
449
+ f"vram={stats.get('vram_used_gb')}GB "
450
+ f"attn={stats.get('attn')} "
451
  f"head={stats.get('raw_head')!r}", flush=True)
452
  self._send(200, {
453
  "id": "local-0",