Two cases the release helper skipped outright, both of which leave GPU memory
held while the compressed subprocess or the torchao device_map="auto" reload
allocates a second copy:
- Quantized models. ExportBackend.load_checkpoint loads 4-bit by DEFAULT, so the
common Studio export hit the is_loaded_in_4bit guard and kept a quantized shard
on every visible GPU. They are now attempted like any other model: transformers
refuses .to() for some bitsandbytes builds, but that refusal raises before
anything moves, so the existing recovery path restores the model and returns
None -- best-effort where the stack allows it, old behaviour where it does not.
- Maps that spill to CPU. Any non-GPU target disqualified the whole model even
though the GPU-mapped modules were still resident and are exactly what needs
reclaiming. A cpu spill is safe to move (those weights are already in host RAM)
and is now released; only disk/meta targets are still skipped, because
accelerate keeps those parameters off the model and moving would try to
materialize the whole checkpoint. An all-CPU map is skipped as a no-op.
The portable torchao FP8/INT8 export freed the in-memory model only when every
parameter sat on one device, then reloaded a second copy with
device_map="auto". A checkpoint loaded through the new multi-GPU export map is
accelerate-dispatched across several GPUs, so that single-device gate never
fired and the original stayed resident on every GPU during the reload -- an OOM
for exactly the models large enough to have needed the sharded load.
It now uses the same _offload_model_for_quantize_subprocess /
_restore_model_after_quantize_subprocess pair as the compressed export, which
removes the accelerate hooks, moves to CPU, and re-dispatches over the recorded
hf_device_map afterwards. Those helpers are extended to XPU as well, since
torchao also runs on Intel GPUs and the path they replace covered both.
Two review fixes on the multi-GPU export path:
1. The LoRA-merge save path budgeted every merged tensor against GPU0
(get_device_properties(0) + unqualified memory_allocated()). A merged tensor
lives on the GPU of its source layer, so for a model sharded across GPUs
(the device_map="balanced" this PR enables) GPU1+ could OOM as their weights
accumulated while only GPU0's headroom was checked. Budget against W's own
device via a per-device cache; single-GPU behavior is unchanged (W on GPU0).
2. _offload_model_for_quantize_subprocess removed the accelerate hooks and then
moved a dispatched model to CPU; if that move raised (host RAM too small for
the sharded checkpoint) the model was left hookless and half-moved, breaking
later exports in the same worker. It now re-dispatches (or, for the
single-device path, moves back) on a failed move before aborting the offload.
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.
