Commit graph

2 commits

Author SHA1 Message Date
pre-commit-ci[bot]
ea0545170a [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-17 20:42:25 +00:00
Hakan Baysal
9ad9d23349 studio/save: reach the UUID/MIG fallback, release sharded models before quantize
Two review fixes on the multi-GPU export sharding:

1. UUID/MIG CUDA_VISIBLE_DEVICES masks resolve to no numeric ids, so the
   len(visible) > 1 gate skipped get_device_map entirely and large exports on
   those hosts still stacked onto GPU0. An empty id list now routes to
   get_device_map(None), whose visible-count fallback exists for exactly this
   case; a genuinely GPU-less host still resolves "sequential" and keeps the
   loader default.

2. The compressed (FP8/NVFP4) export freed GPU memory before its llm-compressor
   subprocess only for single-device models -- a plain .to("cpu") is invalid on
   an accelerate-dispatched model, so a multi-GPU-sharded checkpoint stayed
   resident on every GPU while the subprocess loaded a second copy. The release
   is factored into _offload_model_for_quantize_subprocess /
   _restore_model_after_quantize_subprocess: dispatched all-GPU shards get their
   accelerate hooks removed, move to CPU, and are re-dispatched over the
   recorded hf_device_map afterwards. Maps with cpu/disk targets (already
   offloading) and quantized models are left alone, as before.
2026-07-17 23:40:21 +03:00