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.
