@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5abb117901

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] Ensure LoRA linear respect the base_layer's tp_size and tp_rank (#25487)

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

Jee Jee Li committed Sep 24, 2025 at 02:19 UTC 5abb1179012268a587043dbc1e5afccb4e41c3a5

5 files changed +24 -41

vllm/lora/layers/base_linear.py

+3 -2

@@ -24,11 +24,12 @@ class BaseLinearLayerWithLoRA(BaseLayerWithLoRA):

24 super().__init__()

25 self.base_layer = base_layer

26 self.input_size = self.base_layer.input_size

27 + # Ensure tp_size and tp_rank consistency with the base_layer.

28 + self.tp_size = self.base_layer.tp_size

29 + self.tp_rank = self.base_layer.tp_rank

30 self.device = _get_lora_device(self.base_layer)

31 self.lora_bias_stacked: Optional[tuple[torch.Tensor, ...]] = None

29 -

32 self.output_slices: tuple[int, ...]

31 - self.tp_size: int

33 self.output_size: int

34 self.n_slices: int

vllm/lora/layers/column_parallel_linear.py

+14 -26

@@ -8,9 +8,7 @@ import torch.nn as nn

8 from transformers import PretrainedConfig

10 from vllm.config.lora import LoRAConfig

11 -from vllm.distributed import (get_tensor_model_parallel_rank,

12 - get_tensor_model_parallel_world_size,

13 - tensor_model_parallel_all_gather)

11 +from vllm.distributed import tensor_model_parallel_all_gather

12 from vllm.distributed.utils import divide

13 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

14 MergedColumnParallelLinear,

@@ -85,7 +83,6 @@ class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

83 # inconsistent when TP is greater than 1.

84 self.is_merged_col_linear = type(

85 base_layer) is MergedColumnParallelLinear

88 - self.tp_size = get_tensor_model_parallel_world_size()

86 self.output_size = self.base_layer.output_size_per_partition

87 # There is only one LoRA layer

88 self.n_slices = 1

@@ -97,22 +94,20 @@ class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

94 # Applicable to cases where the base_layer is

95 # MergedColumnParallelLinear.

96 if self.is_merged_col_linear:

100 - tp_rank = get_tensor_model_parallel_rank()

97 shard_size = self.output_size // 2

98 offset = lora_b.shape[0] // 2

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

100 + left_weight = lora_b[self.tp_rank * shard_size:(self.tp_rank + 1) *

101 shard_size, :]

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

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

102 + right_weight = lora_b[offset + self.tp_rank * shard_size:offset +

103 + (self.tp_rank + 1) * shard_size, :]

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

105 # Applicable to cases where the base_layer is

106 # ColumnParallelLinear.

107 else:

112 - tensor_model_parallel_rank = get_tensor_model_parallel_rank()

108 shard_size = self.output_size

114 - start_idx = tensor_model_parallel_rank * shard_size

115 - end_idx = (tensor_model_parallel_rank + 1) * shard_size

109 + start_idx = self.tp_rank * shard_size

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

111 lora_b = lora_b[start_idx:end_idx, :]

112 return lora_b

113

@@ -120,10 +115,9 @@ class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

115 # TODO: Fix the slicing logic of bias.

116 if bias is None:

117 return bias

123 - tensor_model_parallel_rank = get_tensor_model_parallel_rank()

118 shard_size = self.output_size

125 - start_idx = tensor_model_parallel_rank * shard_size

126 - end_idx = (tensor_model_parallel_rank + 1) * shard_size

119 + start_idx = self.tp_rank * shard_size

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

121 bias = bias[start_idx:end_idx]

122 return bias

123

@@ -144,7 +138,7 @@ class ColumnParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

138

139 # Matrix multiply.

140 output_parallel = self.apply(input_, bias)

147 - if self.base_layer.gather_output:

141 + if self.base_layer.gather_output and self.tp_size > 1:

142 # All-gather across the partitions.

143 output = tensor_model_parallel_all_gather(output_parallel)

144 else:

@@ -185,8 +179,6 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

179 QKVParallelLinear]) -> None:

180 super().__init__(base_layer)

181 # There are two LoRA layers

188 - self.tp_size = get_tensor_model_parallel_world_size()

189 - self.tp_rank = get_tensor_model_parallel_rank()

182 # the output_sizes in MergedColumnParallelLinear is not sharded by tp

183 # we need to divide it by the tp_size to get correct slices size

184 output_sizes = self.base_layer.output_sizes

@@ -341,9 +333,9 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):

333 self.n_slices = 1

334

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

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

336 +

337 + self.q_shard_id = self.tp_rank

338 + self.kv_shard_id = self.tp_rank // self.base_layer.num_kv_head_replicas

339 lora_b_q = lora_b[self.q_proj_shard_size *

340 self.q_shard_id:self.q_proj_shard_size *

341 (self.q_shard_id + 1), :]

@@ -397,8 +389,6 @@ class MergedQKVParallelLinearWithLoRA(MergedColumnParallelLinearWithLoRA):

389 super().__init__(base_layer)

390 # There are three LoRA layer.

391 self.n_slices = len(self.base_layer.output_sizes)

400 - self.tp_size = get_tensor_model_parallel_world_size()

401 - self.tp_rank = get_tensor_model_parallel_rank()

392

393 self.q_proj_shard_size = (self.base_layer.num_heads *

394 self.base_layer.head_size)

@@ -461,9 +451,8 @@ class ColumnParallelLinearWithShardedLoRA(ColumnParallelLinearWithLoRA):

451 # Therefore, the sharding of `lora_a` only needs to correspond with the

452 # gather operation.

453 def slice_lora_a(self, lora_a: torch.Tensor) -> torch.Tensor:

464 - tp_rank = get_tensor_model_parallel_rank()

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

466 - start_idx = tp_rank * shard_size

455 + start_idx = self.tp_rank * shard_size

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

457 return lora_a

458

@@ -547,9 +536,8 @@ class QKVParallelLinearWithShardedLoRA(QKVParallelLinearWithLoRA):

536 """

537

538 def slice_lora_a(self, lora_a: torch.Tensor) -> torch.Tensor:

550 - tp_rank = get_tensor_model_parallel_rank()

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

552 - start_idx = tp_rank * shard_size

540 + start_idx = self.tp_rank * shard_size

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

542 return lora_a

543

vllm/lora/layers/replicated_linear.py

-1

@@ -18,7 +18,6 @@ class ReplicatedLinearWithLoRA(BaseLinearLayerWithLoRA):

18 def __init__(self, base_layer: ReplicatedLinear) -> None:

19 super().__init__(base_layer, )

20 # To ensure interface compatibility, set to 1 always.

21 - self.tp_size = 1

21 self.output_size = self.base_layer.output_size

22 self.n_slices = 1

vllm/lora/layers/row_parallel_linear.py

+5 -10

@@ -8,9 +8,7 @@ import torch.nn as nn

8 from transformers import PretrainedConfig

10 from vllm.config.lora import LoRAConfig

11 -from vllm.distributed import (get_tensor_model_parallel_rank,

12 - get_tensor_model_parallel_world_size,

13 - split_tensor_along_last_dim,

11 +from vllm.distributed import (split_tensor_along_last_dim,

12 tensor_model_parallel_all_reduce)

13 # yapf: disable

14 from vllm.model_executor.layers.linear import RowParallelLinear

@@ -25,12 +23,9 @@ class RowParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

23 def __init__(self, base_layer: RowParallelLinear) -> None:

24 super().__init__(base_layer)

28 - self.tp_size = get_tensor_model_parallel_world_size()

26 # reset input_size

27 self.input_size = self.base_layer.input_size_per_partition

28 self.output_size = self.base_layer.output_size

32 -

33 - self.tp_rank = get_tensor_model_parallel_rank()

29 # There is only one LoRA layer.

30 self.n_slices = 1

@@ -68,12 +63,12 @@ class RowParallelLinearWithLoRA(BaseLinearLayerWithLoRA):

63 else:

64 # TODO: simplify code below

65 splitted_input = split_tensor_along_last_dim(

71 - input_, num_partitions=self.base_layer.tp_size)

66 + input_, num_partitions=self.tp_size)

67 input_parallel = splitted_input[self.tp_rank].contiguous()

69 # Matrix multiply.

70 output_parallel = self.apply(input_parallel)

76 - if self.base_layer.reduce_results and self.base_layer.tp_size > 1:

71 + if self.base_layer.reduce_results and self.tp_size > 1:

72 output_ = tensor_model_parallel_all_reduce(output_parallel)

73 else:

74 output_ = output_parallel

@@ -154,8 +149,8 @@ class RowParallelLinearWithShardedLoRA(RowParallelLinearWithLoRA):

149 buffer, x, self.lora_a_stacked, 1.0)

150 if not current_platform.can_update_inplace():

151 buffer = shrunk_buffer

157 -

158 - buffer = tensor_model_parallel_all_reduce(buffer)

152 + if self.tp_size>1:

153 + buffer = tensor_model_parallel_all_reduce(buffer)

154

155 # following S-LoRA, allows the fusing of all_gather and all_reduce

156 # by adding the column partitioned lora output to a slice of output

vllm/lora/lora_weights.py

+2 -2

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

49 @property

50 def input_dim(self) -> int:

51 - return self.lora_a.shape[0]

51 + return self.lora_a.shape[1]

53 @property

54 def output_dim(self) -> int:

55 - return self.lora_b.shape[1]

55 + return self.lora_b.shape[0]

57 @property

58 def is_packed(self) -> bool: