@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 273690a50a

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Core] Optimize LoRA weight loading (#25403)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>

Jee Jee Li committed Sep 23, 2025 at 18:19 UTC 273690a50ac2a5fa79fa7acc5077e49aa1af427e

10 files changed +83 -83

tests/lora/test_layers.py

+14 -12

@@ -164,8 +164,8 @@ def populate_loras(

164 weight=layer_weights,

165 generate_embeddings_tensor=generate_embeddings_tensor,

166 )

167 - sublora.lora_b = sublora.lora_b[:, (sublora_len *

168 - i):(sublora_len * (i + 1))]

167 + sublora.lora_b = sublora.lora_b[(sublora_len *

168 + i):(sublora_len * (i + 1)), :]

169 sublora.optimize()

170 subloras.append(sublora)

171

@@ -304,9 +304,9 @@ def test_embeddings(dist_init, num_loras, device, vocab_size, stage) -> None:

304 result = embedding(input_)

305 after_a = F.embedding(

306 input_,

307 - lora.lora_a,

307 + lora.lora_a.T,

308 )

309 - result += (after_a @ lora.lora_b)

309 + result += (after_a @ lora.lora_b.T)

310 expected_results.append(result)

311 expected_result = torch.cat(expected_results)

312

@@ -445,9 +445,9 @@ def test_embeddings_with_new_embeddings(dist_init, num_loras, device,

445 result = expanded_embedding(input_)

446 after_a = F.embedding(

447 original_input_,

448 - lora.lora_a,

448 + lora.lora_a.T,

449 )

450 - result += (after_a @ lora.lora_b)

450 + result += (after_a @ lora.lora_b.T)

451 expected_results.append(result)

452 expected_result = torch.cat(expected_results)

453

@@ -575,7 +575,7 @@ def test_lm_head_logits_processor(dist_init, num_loras, device, vocab_size,

575 lm_head=linear,

576 embedding_bias=None)

577 result[:, vocab_size + embeddings_tensor_len:] = float("-inf")

578 - result += input_ @ lora.lora_a @ lora.lora_b * lora.scaling

578 + result += input_ @ lora.lora_a.T @ lora.lora_b.T * lora.scaling

579 expected_results.append(result)

580 expected_result = torch.cat(expected_results)

581 logits_processor.org_vocab_size = vocab_size

@@ -692,9 +692,10 @@ def test_linear_replicated(

692

693 expected_results: list[torch.Tensor] = []

694 for input_, lora_id in zip(inputs, prompt_mapping):

695 +

696 lora = lora_dict[lora_id]

697 result = linear(input_)[0]

697 - result += input_ @ lora.lora_a @ lora.lora_b * lora.scaling

698 + result += input_ @ lora.lora_a.T @ lora.lora_b.T * lora.scaling

699 expected_results.append(result)

700 expected_result = torch.cat(expected_results)

701

@@ -817,7 +818,7 @@ def test_linear_parallel(dist_init, num_loras, orientation, fully_shard,

818 for input_, lora_id in zip(inputs, prompt_mapping):

819 lora = lora_dict[lora_id]

820 result = linear(input_)[0]

820 - result += input_ @ lora.lora_a @ lora.lora_b * lora.scaling

821 + result += input_ @ lora.lora_a.T @ lora.lora_b.T * lora.scaling

822 expected_results.append(result)

823 expected_result = torch.cat(expected_results)

824

@@ -965,9 +966,10 @@ def test_column_parallel_packed(dist_init, num_loras, repeats, fully_shard,

966 result = linear(input_)[0]

967 subloras = sublora_dict[lora_id]

968 for i, sublora in enumerate(subloras):

968 - result[:, sublora.lora_b.shape[1] * i:sublora.lora_b.shape[1] *

969 - (i + 1)] += (input_ @ sublora.lora_a @ sublora.lora_b *

970 - sublora.scaling)

969 + result[:, sublora.lora_b.shape[0] * i:sublora.lora_b.shape[0] *

970 + (i + 1)] += (

971 + input_ @ sublora.lora_a.T @ sublora.lora_b.T *

972 + sublora.scaling)

973 expected_results.append(result)

974 expected_result = torch.cat(expected_results)

975

tests/lora/test_lora_manager.py

+6 -6

@@ -63,9 +63,9 @@ def test_from_lora_tensors(sql_lora_files, device):

63 assert lora.lora_b is not None

64 assert lora.lora_a.device == torch.device(device)

65 assert lora.lora_b.device == torch.device(device)

66 - assert (lora.lora_a.shape[1] == lora.lora_b.shape[0]

66 + assert (lora.lora_a.shape[0] == lora.lora_b.shape[1]

67 ), f"{lora.lora_a.shape=}, {lora.lora_b.shape=}"

68 - assert lora.lora_a.shape[1] == 8

68 + assert lora.lora_a.shape[0] == 8

69 embeddings_module = next(

70 (k for k in EMBEDDING_MODULES if k in module_name), None)

71 if embeddings_module:

@@ -86,8 +86,8 @@ def create_lora(lora_id: int, model: nn.Module, sub_modules: list[str],

86 name,

87 8,

88 16,

89 - torch.rand([w.shape[1], 8], device=device),

90 - torch.rand([8, w.shape[0]], device=device),

89 + torch.rand([8, w.shape[1]], device=device),

90 + torch.rand([w.shape[0], 8], device=device),

91 )

92 return LoRAModel(lora_id, 8, loras)

@@ -109,8 +109,8 @@ def create_packed_lora(

109 replaced_module_name,

110 8,

111 16,

112 - torch.rand([w.shape[1], 8], device=device),

113 - torch.rand([8, w.shape[0] // len(replaced_module_names)],

112 + torch.rand([8, w.shape[1]], device=device),

113 + torch.rand([w.shape[0] // len(replaced_module_names), 8],

114 device=device),

115 )

116 return LoRAModel(lora_id, 8, loras)

tests/lora/utils.py

+4 -4

@@ -36,10 +36,10 @@ class DummyLoRAManager:

36 module_name,

37 rank=rank,

38 lora_alpha=1,

39 - lora_a=torch.rand([weight.shape[1], rank],

39 + lora_a=torch.rand([rank, weight.shape[1]],

40 dtype=weight.dtype,

41 device=self._device),

42 - lora_b=torch.rand([rank, weight.shape[0]],

42 + lora_b=torch.rand([weight.shape[0], rank],

43 dtype=weight.dtype,

44 device=self._device),

45 )

@@ -67,8 +67,8 @@ class DummyLoRAManager:

67 module_name,

68 rank=rank,

69 lora_alpha=1,

70 - lora_a=torch.rand([input_dim, rank], device="cuda"),

71 - lora_b=torch.rand([rank, output_dim], device="cuda"),

70 + lora_a=torch.rand([rank, input_dim], device="cuda"),

71 + lora_b=torch.rand([output_dim, input_dim], device="cuda"),

72 embeddings_tensor=embeddings_tensor,

73 )

74 self.set_module_lora(module_name, lora)

vllm/lora/layers/base_linear.py

+5 -5

@@ -121,18 +121,18 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):

121 lora_bias = self.slice_bias(lora_bias)

122

123 self.lora_a_stacked[0][index,

124 - 0, :lora_a.shape[1], :lora_a.shape[0]].copy_(

125 - lora_a.T, non_blocking=True)

124 + 0, :lora_a.shape[0], :lora_a.shape[1]].copy_(

125 + lora_a, non_blocking=True)

126 self.lora_b_stacked[0][index,

127 - 0, :lora_b.shape[1], :lora_b.shape[0]].copy_(

128 - lora_b.T, non_blocking=True)

127 + 0, :lora_b.shape[0], :lora_b.shape[1]].copy_(

128 + lora_b, non_blocking=True)

129 if lora_bias is not None:

130

131 self.lora_bias_stacked = cast(tuple[torch.Tensor, ...],

132 self.lora_bias_stacked)

133 assert len(self.lora_bias_stacked)

134 self.lora_bias_stacked[0][index, 0, :lora_bias.shape[0]].copy_(

135 - lora_bias.T, non_blocking=True)

135 + lora_bias, non_blocking=True)

136

137 def apply(self,

138 x: torch.Tensor,

vllm/lora/layers/column_parallel_linear.py

+33 -34

@@ -99,13 +99,13 @@ class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

99 if self.is_merged_col_linear:

100 tp_rank = get_tensor_model_parallel_rank()

101 shard_size = self.output_size // 2

102 - offset = lora_b.shape[-1] // 2

102 + offset = lora_b.shape[0] // 2

103

104 - left_weight = lora_b[:, tp_rank * shard_size:(tp_rank + 1) *

105 - shard_size]

106 - right_weight = lora_b[:, offset + tp_rank * shard_size:offset +

107 - (tp_rank + 1) * shard_size]

108 - lora_b = torch.cat([left_weight, right_weight], dim=1)

104 + left_weight = lora_b[tp_rank * shard_size:(tp_rank + 1) *

105 + shard_size, :]

106 + right_weight = lora_b[offset + tp_rank * shard_size:offset +

107 + (tp_rank + 1) * shard_size, :]

108 + lora_b = torch.cat([left_weight, right_weight], dim=0)

109 # Applicable to cases where the base_layer is

110 # ColumnParallelLinear.

111 else:

@@ -113,7 +113,7 @@ class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

113 shard_size = self.output_size

114 start_idx = tensor_model_parallel_rank * shard_size

115 end_idx = (tensor_model_parallel_rank + 1) * shard_size

116 - lora_b = lora_b[:, start_idx:end_idx]

116 + lora_b = lora_b[start_idx:end_idx, :]

117 return lora_b

118

119 def slice_bias(self, bias: torch.Tensor) -> torch.Tensor:

@@ -251,9 +251,8 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

251 for i, (shard_id, shard_size) in enumerate(

252 zip(self.output_ids, self.output_slices)):

253 if (lora_b_i := lora_b[i]) is not None:

254 - sliced_lora_b[i] = lora_b_i[:,

255 - shard_size * shard_id:shard_size *

256 - (shard_id + 1)]

254 + sliced_lora_b[i] = lora_b_i[shard_size * shard_id:shard_size *

255 + (shard_id + 1), :]

256 return sliced_lora_b

257

258 def slice_bias(

@@ -285,12 +284,12 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

284 for i in range(self.n_slices):

285 if (lora_a_i := lora_a[i]) is not None:

286 self.lora_a_stacked[i][

288 - index, 0, :lora_a_i.shape[1], :lora_a_i.shape[0]].copy_(

289 - lora_a_i.T, non_blocking=True)

287 + index, 0, :lora_a_i.shape[0], :lora_a_i.shape[1]].copy_(

288 + lora_a_i, non_blocking=True)

289 if (lora_b_i := lora_b[i]) is not None:

290 self.lora_b_stacked[i][

292 - index, 0, :lora_b_i.shape[1], :lora_b_i.shape[0]].copy_(

293 - lora_b_i.T, non_blocking=True)

291 + index, 0, :lora_b_i.shape[0], :lora_b_i.shape[1]].copy_(

292 + lora_b_i, non_blocking=True)

293

294 if lora_bias is not None:

295 self.lora_bias_stacked = cast(tuple[torch.Tensor, ...],

@@ -299,7 +298,7 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

298 if (lora_bias_i := lora_bias[i]) is not None:

299 self.lora_bias_stacked[i][index,

300 0, :lora_bias_i.shape[0]].copy_(

302 - lora_bias_i.T,

301 + lora_bias_i,

302 non_blocking=True)

303

304 @classmethod

@@ -345,18 +344,18 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

344 tp_rank = get_tensor_model_parallel_rank()

345 self.q_shard_id = tp_rank

346 self.kv_shard_id = tp_rank // self.base_layer.num_kv_head_replicas

348 - lora_b_q = lora_b[:, self.q_proj_shard_size *

347 + lora_b_q = lora_b[self.q_proj_shard_size *

348 self.q_shard_id:self.q_proj_shard_size *

350 - (self.q_shard_id + 1)]

349 + (self.q_shard_id + 1), :]

350 k_offset = self.q_proj_total_size

352 - lora_b_k = lora_b[:, k_offset +

351 + lora_b_k = lora_b[k_offset +

352 self.kv_proj_shard_size * self.kv_shard_id:k_offset +

354 - self.kv_proj_shard_size * (self.kv_shard_id + 1)]

353 + self.kv_proj_shard_size * (self.kv_shard_id + 1), :]

354 v_offset = k_offset + self.kv_proj_total_size

356 - lora_b_v = lora_b[:, v_offset +

355 + lora_b_v = lora_b[v_offset +

356 self.kv_proj_shard_size * self.kv_shard_id:v_offset +

358 - self.kv_proj_shard_size * (self.kv_shard_id + 1)]

359 - lora_b = torch.cat([lora_b_q, lora_b_k, lora_b_v], dim=1)

357 + self.kv_proj_shard_size * (self.kv_shard_id + 1), :]

358 + lora_b = torch.cat([lora_b_q, lora_b_k, lora_b_v], dim=0)

359 return lora_b

360

361 def slice_bias(self, bias: torch.Tensor) -> torch.Tensor:

@@ -465,7 +464,7 @@ class ColumnParallelLinearWithShardedLoRA(ColumnParallelLinearWithLoRA):

464 tp_rank = get_tensor_model_parallel_rank()

465 shard_size = self.lora_a_stacked[0].shape[2]

466 start_idx = tp_rank * shard_size

468 - lora_a = lora_a[:, start_idx:start_idx + shard_size]

467 + lora_a = lora_a[start_idx:start_idx + shard_size, :]

468 return lora_a

469

470 def apply(self,

@@ -508,10 +507,10 @@ class MergedColumnParallelLinearWithShardedLoRA(

507 output_shard_size = self.lora_a_stacked[0].shape[2]

508 output_start_idx = self.tp_rank * output_shard_size

509 lora_a = [

511 - lora_a[0][:, output_start_idx:output_start_idx +

512 - output_shard_size] if lora_a[0] is not None else None,

513 - lora_a[1][:, output_start_idx:output_start_idx +

514 - output_shard_size] if lora_a[1] is not None else None,

510 + lora_a[0][output_start_idx:output_start_idx +

511 + output_shard_size, :] if lora_a[0] is not None else None,

512 + lora_a[1][output_start_idx:output_start_idx +

513 + output_shard_size, :] if lora_a[1] is not None else None,

514 ]

515 return lora_a

516

@@ -551,7 +550,7 @@ class QKVParallelLinearWithShardedLoRA(QKVParallelLinearWithLoRA):

550 tp_rank = get_tensor_model_parallel_rank()

551 shard_size = self.lora_a_stacked[0].shape[2]

552 start_idx = tp_rank * shard_size

554 - lora_a = lora_a[:, start_idx:start_idx + shard_size]

553 + lora_a = lora_a[start_idx:start_idx + shard_size, :]

554 return lora_a

555

556 def apply(self,

@@ -589,12 +588,12 @@ class MergedQKVParallelLinearWithShardedLoRA(MergedQKVParallelLinearWithLoRA):

588 shard_size = [self.lora_a_stacked[i].shape[2] for i in range(3)]

589 start_idx = [self.tp_rank * shard_size[i] for i in range(3)]

590 lora_a = [

592 - lora_a[0][:, start_idx[0]:start_idx[0] +

593 - shard_size[0]] if lora_a[0] is not None else None,

594 - lora_a[1][:, start_idx[1]:start_idx[1] +

595 - shard_size[1]] if lora_a[1] is not None else None,

596 - lora_a[2][:, start_idx[2]:start_idx[2] +

597 - shard_size[2]] if lora_a[2] is not None else None,

591 + lora_a[0][start_idx[0]:start_idx[0] +

592 + shard_size[0], :] if lora_a[0] is not None else None,

593 + lora_a[1][start_idx[1]:start_idx[1] +

594 + shard_size[1], :] if lora_a[1] is not None else None,

595 + lora_a[2][start_idx[2]:start_idx[2] +

596 + shard_size[2], :] if lora_a[2] is not None else None,

597 ]

598 return lora_a

599

vllm/lora/layers/logits_processor.py

+4 -4

@@ -140,11 +140,11 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):

140 ):

141 self.reset_lora(index)

142 self.lora_a_stacked[index,

143 - 0, :lora_a.shape[1], :lora_a.shape[0]].copy_(

144 - lora_a.T, non_blocking=True)

143 + 0, :lora_a.shape[0], :lora_a.shape[1]].copy_(

144 + lora_a, non_blocking=True)

145 self.lora_b_stacked[index,

146 - 0, :lora_b.shape[1], :lora_b.shape[0]].copy_(

147 - lora_b.T, non_blocking=True)

146 + 0, :lora_b.shape[0], :lora_b.shape[1]].copy_(

147 + lora_b, non_blocking=True)

148 if embeddings_tensor is not None:

149 self.embeddings_tensors[

150 index,

vllm/lora/layers/row_parallel_linear.py

+2 -2

@@ -39,7 +39,7 @@ class RowParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

39 shard_size = self.input_size

40 start_idx = self.tp_rank * shard_size

41 end_idx = (self.tp_rank + 1) * shard_size

42 - lora_a = lora_a[start_idx:end_idx, :]

42 + lora_a = lora_a[:,start_idx:end_idx]

43 return lora_a

45 def slice_lora_b(self, lora_b: torch.Tensor) -> torch.Tensor:

@@ -122,7 +122,7 @@ class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):

122 shard_size = self.lora_b_stacked[0].shape[2]

123 start_idx = self.tp_rank * shard_size

124 end_idx = (self.tp_rank + 1) * shard_size

125 - lora_b = lora_b[:, start_idx:end_idx]

125 + lora_b = lora_b[ start_idx:end_idx,:]

126 return lora_b

127

128 def slice_bias(self, bias: torch.Tensor) -> torch.Tensor:

vllm/lora/layers/vocal_parallel_embedding.py

+6 -4

@@ -95,11 +95,13 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):

95 bias: Optional[torch.Tensor] = None,

96 ):

97 self.reset_lora(index)

98 - self.lora_a_stacked[index, :lora_a.shape[0], :lora_a.shape[1]].copy_(

99 - lora_a, non_blocking=True)

98 + # NOTE self.lora_a_stacked is row-major, and lora_a is col-major,

99 + # so we need transpose here

100 + self.lora_a_stacked[index, :lora_a.shape[1], :lora_a.shape[0]].copy_(

101 + lora_a.T, non_blocking=True)

102 self.lora_b_stacked[index,

101 - 0, :lora_b.shape[1], :lora_b.shape[0]].copy_(

102 - lora_b.T, non_blocking=True)

103 + 0, :lora_b.shape[0], :lora_b.shape[1]].copy_(

104 + lora_b, non_blocking=True)

105 if embeddings_tensor is not None:

106 self.embeddings_tensors[

107 index,

vllm/lora/lora_weights.py

+2 -2

@@ -86,11 +86,11 @@ class LoRALayerWeights:

86 embeddings_tensor_dim: Optional[int] = None,

87 bias_enabled: Optional[bool] = False) -> "LoRALayerWeights":

88 pin_memory = str(device) == "cpu" and is_pin_memory_available()

89 - lora_a = torch.zeros([input_dim, rank],

89 + lora_a = torch.zeros([rank, input_dim],

90 dtype=dtype,

91 device=device,

92 pin_memory=pin_memory)

93 - lora_b = torch.zeros([rank, output_dim],

93 + lora_b = torch.zeros([output_dim, rank],

94 dtype=dtype,

95 device=device,

96 pin_memory=pin_memory)

vllm/lora/models.py

+7 -10

@@ -152,30 +152,29 @@ class LoRAModel:

152 module_name, peft_helper, lora_embeddings_tensor)

153

154 if is_bias:

155 - loras[module_name].bias = tensor.to(device=device,

156 - dtype=dtype).t()

157 - bias = tensor.to(device=device, dtype=dtype).t()

155 + loras[module_name].bias = tensor.to(device=device, dtype=dtype)

156 + bias = tensor.to(device=device, dtype=dtype)

157 if pin_memory:

158 bias = bias.pin_memory()

159 loras[module_name].bias = bias

160 elif is_lora_a:

161 loras[module_name].lora_a = tensor.to(device=device,

163 - dtype=dtype).t()

162 + dtype=dtype)

163 if pin_memory:

164 loras[module_name].lora_a = loras[

165 module_name].lora_a.pin_memory()

166 else:

167 loras[module_name].lora_b = tensor.to(device=device,

169 - dtype=dtype).t()

168 + dtype=dtype)

169 assert embedding_padding_modules is not None

170 if any(name in module_name

171 for name in embedding_padding_modules

172 ) and target_embedding_padding is not None:

173 lora_b = loras[module_name].lora_b

175 - assert target_embedding_padding >= lora_b.shape[1]

176 - addition = target_embedding_padding - lora_b.shape[1]

174 + assert target_embedding_padding >= lora_b.shape[0]

175 + addition = target_embedding_padding - lora_b.shape[0]

176 loras[module_name].lora_b = torch.nn.functional.pad(

178 - lora_b, (0, addition))

177 + lora_b, (0, 0, 0, addition))

178 if pin_memory:

179 loras[module_name].lora_b = loras[

180 module_name].lora_b.pin_memory()

@@ -585,7 +584,6 @@ class LoRAModelManager:

584 "cpu",

585 bias_enabled=bias_enabled,

586 )

588 - lora.optimize()

587 else:

588 parts = module_name.split(".")

589 replacements = self.packed_modules_mapping[parts[-1]]

@@ -600,7 +598,6 @@ class LoRAModelManager:

598 "cpu",

599 bias_enabled=bias_enabled,

600 )

603 - lora.optimize()

601 subloras.append(lora)

602 lora = PackedLoRALayerWeights.pack(subloras)

603 model.loras[module_name] = lora