[Perf][Diffusion] Optimize HSDP startup via Rank-0 shared weight loading and accelerated LoRA delta computation (#7005)

Signed-off-by: samithuang <285365963@qq.com>

vllm_omni/diffusion/lora/loader.py

@@ -43,6 +43,7 @@ def _prepare_lora_delta(
     lora_a_suffix: str = "lora_A.weight",
     lora_b_suffix: str = "lora_B.weight",
     lora_bias_suffix: str = "bias",
+    compute_device: torch.device | None = None,
 ):
     used_keys: set[str] = set()
     # stacked_params_mapping                       param_to_weight_names
@@ -70,6 +71,8 @@ def _prepare_lora_delta(
                 if lora_bias_key not in lora_state_dict:
                     return None, used_keys
                 delta = lora_state_dict[lora_bias_key]
+                if compute_device is not None and compute_device.type != "cpu":
+                    delta = delta.to(compute_device)
                 stacked_deltas.append(delta)
                 used_keys.add(lora_bias_key)
         else:
@@ -80,6 +83,9 @@ def _prepare_lora_delta(
                     return None, used_keys
                 a = lora_state_dict[lora_a_key]
                 b = lora_state_dict[lora_b_key]
+                if compute_device is not None and compute_device.type != "cpu":
+                    a = a.to(compute_device)
+                    b = b.to(compute_device)
                 delta = torch.matmul(b, a)
                 stacked_deltas.append(delta)
                 used_keys.add(lora_a_key)
@@ -94,7 +100,10 @@ def _prepare_lora_delta(
         if lora_bias_key not in lora_state_dict:
             return None, used_keys
         used_keys.add(lora_bias_key)
-        return lora_state_dict[lora_bias_key], used_keys
+        delta = lora_state_dict[lora_bias_key]
+        if compute_device is not None and compute_device.type != "cpu":
+            delta = delta.to(compute_device)
+        return delta, used_keys
 
     lora_a_key = f"{base_key}.{lora_a_suffix}"
     lora_b_key = f"{base_key}.{lora_b_suffix}"
@@ -102,9 +111,13 @@ def _prepare_lora_delta(
         return None, used_keys
     a = lora_state_dict[lora_a_key]
     b = lora_state_dict[lora_b_key]
+    if compute_device is not None and compute_device.type != "cpu":
+        a = a.to(compute_device)
+        b = b.to(compute_device)
+    delta = torch.matmul(b, a)
     used_keys.add(lora_a_key)
     used_keys.add(lora_b_key)
-    return torch.matmul(b, a), used_keys
+    return delta, used_keys
 
 
 def _load_lora_state_dict(
@@ -267,6 +280,7 @@ class LoraLoaderMixin:
             else:
                 continue
 
+            compute_dev = params.device if params.device.type != "cpu" else None
             delta, used_keys = _prepare_lora_delta(
                 state_dict,
                 base_key,
@@ -275,6 +289,7 @@ class LoraLoaderMixin:
                 lora_a_suffix,
                 lora_b_suffix,
                 lora_bias_suffix,
+                compute_device=compute_dev,
             )
             if delta is None:
                 continue
@@ -319,6 +334,7 @@ class LoraLoaderMixin:
             else:
                 continue
 
+            compute_dev = param.device if param.device.type != "cpu" else None
             delta, used_keys = _prepare_lora_delta(
                 state_dict,
                 base_key,
@@ -327,6 +343,7 @@ class LoraLoaderMixin:
                 lora_a_suffix,
                 lora_b_suffix,
                 lora_bias_suffix,
+                compute_device=compute_dev,
             )
             if delta is None:
                 continue