fix(merge): serialize GPU compression SVD to stop async-race abort
The compression SVD (torch.svd_lowrank) can race asynchronously on some stacks (Blackwell/cu130), corrupting CUDA state and aborting a later kernel — a C++ abort the Python try/except can't catch. This is what crashed Krea 2 merges (two different-form LoRAs route through dense merge + compression, the only path that runs GPU compression SVD here). Make the per-group cuda.synchronize() permanent, but gate it tightly on (GPU compression active) so it only fires on compression merges and fast no-compression merges keep their fully-pipelined speed. Previously this sync only ran under the PROFILE_MERGE sentinel. No version bump. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
parent
b160b02c51
commit
d4ef3c75d2
+9
-1
@@ -7505,6 +7505,14 @@ class LoRAOptimizer(_LoRAMergeBase):
|
||||
_merge_prof_lock = threading.Lock() if _profiling_on else None
|
||||
_prof_cuda = (use_gpu and compute_device is not None
|
||||
and getattr(compute_device, "type", None) == "cuda")
|
||||
# A GPU compression SVD (torch.svd_lowrank) can race asynchronously on
|
||||
# some stacks (Blackwell/cu130), corrupting state and aborting a later
|
||||
# kernel — a C++ abort the Python try/except can't catch. When the
|
||||
# compression SVD runs on GPU, serialize each group so the SVD finishes
|
||||
# before the next group launches. Gated tightly (compression ON + GPU
|
||||
# SVD) so fast no-compression merges keep their pipelined speed.
|
||||
_compress_sync = (_prof_cuda and compress_rank > 0
|
||||
and resolved_svd_device is not None)
|
||||
|
||||
def _prof_t():
|
||||
"""Synced timestamp — without the sync, async CUDA kernels would
|
||||
@@ -7963,7 +7971,7 @@ class LoRAOptimizer(_LoRAMergeBase):
|
||||
logging.info(f"[merge-debug] processing {idx + 1}/{len(group_items)}: "
|
||||
f"{_dbg_tk} (n_loras={prefix_stats.get(label_prefix, {}).get('n_loras', '?')})")
|
||||
result = _merge_one_group(label_prefix, target_group)
|
||||
if _merge_prof is not None and _prof_cuda:
|
||||
if _compress_sync or (_merge_prof is not None and _prof_cuda):
|
||||
torch.cuda.synchronize(compute_device)
|
||||
if _sl_key is not None:
|
||||
_sl_store(_sl_key, result)
|
||||
|
||||
Reference in New Issue
Block a user