8ea8ba275e · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -1,500 +0,0 @@

1 -# SPDX-License-Identifier: Apache-2.0

2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -

4 -from collections.abc import Iterable

5 -

6 -import torch

7 -import torch.nn as nn

8 -from transformers import SwinConfig

9 -from transformers.models.swin.modeling_swin import SwinEmbeddings, SwinPatchMerging

10 -from transformers.models.swin.modeling_swin import SwinLayer as HFSwinLayer

11 -from transformers.pytorch_utils import meshgrid

12 -

13 -from vllm.model_executor.layers.activation import get_act_fn

14 -from vllm.model_executor.layers.linear import (

15 - ColumnParallelLinear,

16 - QKVParallelLinear,

17 - RowParallelLinear,

18 -)

19 -from vllm.model_executor.layers.quantization import QuantizationConfig

20 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader

21 -

22 -

23 -class SwinSelfAttention(nn.Module):

24 - def __init__(

25 - self,

26 - config: SwinConfig,

27 - dim: int,

28 - num_heads: int,

29 - window_size: int,

30 - quant_config: QuantizationConfig | None = None,

31 - prefix: str = "",

32 - ) -> None:

33 - super().__init__()

34 - if dim % num_heads != 0:

35 - raise ValueError(

36 - f"The hidden size ({dim}) is not a multiple of the number of "

37 - f"attention heads ({num_heads})"

38 - )

39 -

40 - self.num_attention_heads = num_heads

41 - self.attention_head_size = int(dim / num_heads)

42 - self.all_head_size = self.num_attention_heads * self.attention_head_size

43 - self.window_size = (

44 - window_size

45 - if isinstance(window_size, Iterable)

46 - else (window_size, window_size)

47 - )

48 - self.scale = self.attention_head_size**-0.5

49 -

50 - self.relative_position_bias_table = nn.Parameter(

51 - torch.zeros(

52 - (2 * self.window_size[0] - 1) * (2 * self.window_size[1] - 1), num_heads

53 - )

54 - )

55 -

56 - # get pair-wise relative position index for each token inside the window

57 - coords_h = torch.arange(self.window_size[0])

58 - coords_w = torch.arange(self.window_size[1])

59 - coords = torch.stack(meshgrid([coords_h, coords_w], indexing="ij"))

60 - coords_flatten = torch.flatten(coords, 1)

61 - relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]

62 - relative_coords = relative_coords.permute(1, 2, 0).contiguous()

63 - relative_coords[:, :, 0] += self.window_size[0] - 1

64 - relative_coords[:, :, 1] += self.window_size[1] - 1

65 - relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1

66 - relative_position_index = relative_coords.sum(-1)

67 -

68 - self.relative_position_index = nn.Parameter(

69 - relative_position_index, requires_grad=False

70 - )

71 -

72 - self.qkv = QKVParallelLinear(

73 - hidden_size=dim,

74 - head_size=self.attention_head_size,

75 - total_num_heads=self.num_attention_heads,

76 - bias=config.qkv_bias,

77 - quant_config=quant_config,

78 - prefix=f"{prefix}.qkv",

79 - )

80 -

81 - def transpose_for_scores(self, x):

82 - new_x_shape = x.size()[:-1] + (

83 - self.num_attention_heads,

84 - self.attention_head_size,

85 - )

86 - x = x.view(new_x_shape)

87 - return x.permute(0, 2, 1, 3)

88 -

89 - def _get_rel_pos_bias(self) -> torch.Tensor:

90 - relative_position_bias = self.relative_position_bias_table[

91 - self.relative_position_index.view(-1)

92 - ]

93 - relative_position_bias = relative_position_bias.view(

94 - self.window_size[0] * self.window_size[1],

95 - self.window_size[0] * self.window_size[1],

96 - -1,

97 - )

98 - relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()

99 - return relative_position_bias.unsqueeze(0)

100 -

101 - def forward(

102 - self,

103 - hidden_states: torch.Tensor,

104 - attention_mask: torch.FloatTensor | None = None,

105 - output_attentions: bool | None = False,

106 - ) -> tuple[torch.Tensor, ...]:

107 - batch_size, dim, num_channels = hidden_states.shape

108 -

109 - qkv_output, _ = self.qkv(hidden_states)

110 - query_layer, key_layer, value_layer = qkv_output.chunk(3, dim=-1)

111 -

112 - key_layer = self.transpose_for_scores(key_layer)

113 - value_layer = self.transpose_for_scores(value_layer)

114 - query_layer = self.transpose_for_scores(query_layer)

115 -

116 - attention_scores = self._get_rel_pos_bias()

117 - if attention_mask is not None:

118 - mask_shape = attention_mask.shape[0]

119 - attention_mask_expanded = attention_mask.view(

120 - 1, mask_shape, 1, dim, dim

121 - ).expand(

122 - batch_size // mask_shape, mask_shape, self.num_attention_heads, dim, dim

123 - )

124 - attention_scores = attention_scores + attention_mask_expanded.unsqueeze(

125 - 1

126 - ).unsqueeze(0)

127 - attention_scores = attention_scores.view(

128 - -1, self.num_attention_heads, dim, dim

129 - )

130 -

131 - context_layer = torch.nn.functional.scaled_dot_product_attention(

132 - query_layer,

133 - key_layer,

134 - value_layer,

135 - attn_mask=attention_scores,

136 - dropout_p=0.0,

137 - )

138 - attention_probs = None

139 -

140 - context_layer = context_layer.permute(0, 2, 1, 3).contiguous()

141 - new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)

142 - context_layer = context_layer.view(new_context_layer_shape)

143 -

144 - outputs = (

145 - (context_layer, attention_probs) if output_attentions else (context_layer,)

146 - )

147 -

148 - return outputs

149 -

150 -

151 -class SwinSelfOutput(nn.Module):

152 - def __init__(

153 - self,

154 - config: SwinConfig,

155 - dim: int,

156 - quant_config: QuantizationConfig | None = None,

157 - prefix: str = "",

158 - ) -> None:

159 - super().__init__()

160 - self.dense = RowParallelLinear(

161 - input_size=dim,

162 - output_size=dim,

163 - quant_config=quant_config,

164 - prefix=f"{prefix}.dense",

165 - )

166 -

167 - def forward(

168 - self, hidden_states: torch.Tensor, input_tensor: torch.Tensor

169 - ) -> torch.Tensor:

170 - hidden_states, _ = self.dense(hidden_states)

171 -

172 - return hidden_states

173 -

174 -

175 -class SwinAttention(nn.Module):

176 - def __init__(

177 - self,

178 - config: SwinConfig,

179 - dim: int,

180 - num_heads: int,

181 - window_size: int,

182 - quant_config: QuantizationConfig | None = None,

183 - prefix: str = "",

184 - ) -> None:

185 - super().__init__()

186 - self.self = SwinSelfAttention(

187 - config,

188 - dim,

189 - num_heads,

190 - window_size,

191 - quant_config=quant_config,

192 - prefix=f"{prefix}.self",

193 - )

194 - self.output = SwinSelfOutput(

195 - config, dim, quant_config=quant_config, prefix=f"{prefix}.output"

196 - )

197 - self.pruned_heads = set()

198 -

199 - def forward(

200 - self,

201 - hidden_states: torch.Tensor,

202 - attention_mask: torch.FloatTensor | None = None,

203 - output_attentions: bool | None = False,

204 - ) -> tuple[torch.Tensor]:

205 - self_outputs = self.self(hidden_states, attention_mask, output_attentions)

206 - attention_output = self.output(self_outputs[0], hidden_states)

207 - outputs = (attention_output,) + self_outputs[1:]

208 - return outputs

209 -

210 -

211 -class SwinIntermediate(nn.Module):

212 - def __init__(

213 - self,

214 - config: SwinConfig,

215 - dim: int,

216 - quant_config: QuantizationConfig | None = None,

217 - prefix: str = "",

218 - ) -> None:

219 - super().__init__()

220 - self.dense = ColumnParallelLinear(

221 - dim,

222 - int(config.mlp_ratio * dim),

223 - quant_config=quant_config,

224 - prefix=f"{prefix}.dense",

225 - )

226 - self.intermediate_act_fn = get_act_fn(config.hidden_act)

227 -

228 - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

229 - hidden_states, _ = self.dense(hidden_states)

230 - hidden_states = self.intermediate_act_fn(hidden_states)

231 - return hidden_states

232 -

233 -

234 -class SwinOutput(nn.Module):

235 - def __init__(

236 - self,

237 - config: SwinConfig,

238 - dim: int,

239 - quant_config: QuantizationConfig | None = None,

240 - prefix: str = "",

241 - ) -> None:

242 - super().__init__()

243 - self.dense = RowParallelLinear(

244 - int(config.mlp_ratio * dim),

245 - dim,

246 - quant_config=quant_config,

247 - prefix=f"{prefix}.dense",

248 - )

249 -

250 - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

251 - hidden_states, _ = self.dense(hidden_states)

252 - return hidden_states

253 -

254 -

255 -class SwinLayer(HFSwinLayer):

256 - def __init__(

257 - self,

258 - config: SwinConfig,

259 - dim: int,

260 - input_resolution: int,

261 - num_heads: int,

262 - drop_path_rate: float = 0.0,

263 - shift_size: int = 0,

264 - quant_config: QuantizationConfig | None = None,

265 - prefix: str = "",

266 - ) -> None:

267 - super().__init__(

268 - config=config,

269 - dim=dim,

270 - input_resolution=input_resolution,

271 - num_heads=num_heads,

272 - drop_path_rate=drop_path_rate,

273 - shift_size=shift_size,

274 - )

275 -

276 - self.attention = SwinAttention(

277 - config,

278 - dim,

279 - num_heads,

280 - window_size=self.window_size,

281 - quant_config=quant_config,

282 - prefix=f"{prefix}.attention",

283 - )

284 - self.intermediate = SwinIntermediate(

285 - config, dim, quant_config=quant_config, prefix=f"{prefix}.intermediate"

286 - )

287 - self.output = SwinOutput(

288 - config, dim, quant_config=quant_config, prefix=f"{prefix}.output"

289 - )

290 -

291 -

292 -class SwinStage(nn.Module):

293 - def __init__(

294 - self,

295 - config: SwinConfig,

296 - dim: int,

297 - input_resolution: int,

298 - depth: int,

299 - num_heads: int,

300 - drop_path: list[float],

301 - downsample: SwinPatchMerging | None = None,

302 - quant_config: QuantizationConfig | None = None,

303 - prefix: str = "",

304 - ) -> None:

305 - super().__init__()

306 - self.config = config

307 - self.dim = dim

308 - self.blocks = nn.ModuleList(

309 - [

310 - SwinLayer(

311 - config=config,

312 - dim=dim,

313 - input_resolution=input_resolution,

314 - num_heads=num_heads,

315 - drop_path_rate=drop_path[layer_idx],

316 - shift_size=0 if (layer_idx % 2 == 0) else config.window_size // 2,

317 - quant_config=quant_config,

318 - prefix=f"{prefix}.blocks.{layer_idx}",

319 - )

320 - for layer_idx in range(depth)

321 - ]

322 - )

323 -

324 - # patch merging layer

325 - if downsample is not None:

326 - self.downsample = downsample(

327 - input_resolution, dim=dim, norm_layer=nn.LayerNorm

328 - )

329 - else:

330 - self.downsample = None

331 -

332 - self.pointing = False

333 -

334 - def forward(

335 - self,

336 - hidden_states: torch.Tensor,

337 - input_dimensions: tuple[int, int],

338 - output_attentions: bool | None = False,

339 - always_partition: bool | None = False,

340 - ) -> tuple[torch.Tensor]:

341 - height, width = input_dimensions

342 - for i, layer_module in enumerate(self.blocks):

343 - layer_outputs = layer_module(

344 - hidden_states,

345 - input_dimensions,

346 - output_attentions,

347 - always_partition,

348 - )

349 -

350 - hidden_states = layer_outputs[0]

351 -

352 - hidden_states_before_downsampling = hidden_states

353 - if self.downsample is not None:

354 - height_downsampled, width_downsampled = (height + 1) // 2, (width + 1) // 2

355 - output_dimensions = (height, width, height_downsampled, width_downsampled)

356 - hidden_states = self.downsample(

357 - hidden_states_before_downsampling, input_dimensions

358 - )

359 - else:

360 - output_dimensions = (height, width, height, width)

361 -

362 - stage_outputs = (

363 - hidden_states,

364 - hidden_states_before_downsampling,

365 - output_dimensions,

366 - )

367 -

368 - if output_attentions:

369 - stage_outputs += layer_outputs[1:]

370 - return stage_outputs

371 -

372 -

373 -class SwinEncoder(nn.Module):

374 - def __init__(

375 - self,

376 - config: SwinConfig,

377 - grid_size: int,

378 - quant_config: QuantizationConfig | None = None,

379 - prefix: str = "",

380 - ) -> None:

381 - super().__init__()

382 - self.num_layers = len(config.depths)

383 - self.config = config

384 - dpr = [

385 - x.item()

386 - for x in torch.linspace(

387 - 0, config.drop_path_rate, sum(config.depths), device="cpu"

388 - )

389 - ]

390 - self.layers = nn.ModuleList(

391 - [

392 - SwinStage(

393 - config=config,

394 - dim=int(config.embed_dim * 2**layer_idx),

395 - input_resolution=(

396 - grid_size[0] // (2**layer_idx),

397 - grid_size[1] // (2**layer_idx),

398 - ),

399 - depth=config.depths[layer_idx],

400 - num_heads=config.num_heads[layer_idx],

401 - drop_path=dpr[

402 - sum(config.depths[:layer_idx]) : sum(

403 - config.depths[: layer_idx + 1]

404 - )

405 - ],

406 - downsample=SwinPatchMerging

407 - if (layer_idx < self.num_layers - 1)

408 - else None,

409 - quant_config=quant_config,

410 - prefix=f"{prefix}.layers.{layer_idx}",

411 - )

412 - for layer_idx in range(self.num_layers)

413 - ]

414 - )

415 -

416 - def forward(

417 - self,

418 - hidden_states: torch.Tensor,

419 - input_dimensions: tuple[int, int],

420 - output_attentions: bool | None = False,

421 - always_partition: bool | None = False,

422 - ) -> tuple[torch.Tensor]:

423 - for i, layer_module in enumerate(self.layers):

424 - layer_outputs = layer_module(

425 - hidden_states,

426 - input_dimensions,

427 - output_attentions,

428 - always_partition,

429 - )

430 -

431 - hidden_states = layer_outputs[0]

432 - output_dimensions = layer_outputs[2]

433 -

434 - input_dimensions = (output_dimensions[-2], output_dimensions[-1])

435 -

436 - return hidden_states

437 -

438 -

439 -class SwinModel(nn.Module):

440 - config_class: SwinConfig

441 -

442 - def __init__(

443 - self,

444 - config: SwinConfig,

445 - quant_config: QuantizationConfig | None = None,

446 - prefix: str = "",

447 - ) -> None:

448 - super().__init__()

449 - self.config = config

450 - self.num_layers = len(config.depths)

451 - self.num_features = int(config.embed_dim * 2 ** (self.num_layers - 1))

452 -

453 - self.embeddings = SwinEmbeddings(config)

454 - self.encoder = SwinEncoder(

455 - config,

456 - self.embeddings.patch_grid,

457 - quant_config=quant_config,

458 - prefix=f"{prefix}.encoder",

459 - )

460 -

461 - def forward(

462 - self,

463 - pixel_values: torch.FloatTensor | None = None,

464 - output_attentions: bool | None = None,

465 - ) -> tuple[torch.Tensor]:

466 - embedding_output, input_dimensions = self.embeddings(pixel_values)

467 -

468 - encoder_outputs = self.encoder(

469 - embedding_output,

470 - input_dimensions,

471 - output_attentions=output_attentions,

472 - )

473 -

474 - return encoder_outputs

475 -

476 - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

477 - stacked_params_mapping = [

478 - ("qkv", "query", "q"),

479 - ("qkv", "key", "k"),

480 - ("qkv", "value", "v"),

481 - ]

482 - params_dict = dict(self.named_parameters())

483 - loaded_params: set[str] = set()

484 -

485 - for name, loaded_weight in weights:

486 - for param_name, weight_name, shard_id in stacked_params_mapping:

487 - if weight_name not in name:

488 - continue

489 - name = name.replace(weight_name, param_name)

490 -

491 - param = params_dict[name]

492 - weight_loader = param.weight_loader

493 - weight_loader(param, loaded_weight, shard_id)

494 - break

495 - else:

496 - param = params_dict[name]

497 - weight_loader = getattr(param, "weight_loader", default_weight_loader)

498 - weight_loader(param, loaded_weight)

499 - loaded_params.add(name)

500 - return loaded_params

[V0 deprecation] Remove Swin model (#35821)