52d0cb8458 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -7,11 +7,13 @@ import torch

7 import torch.nn as nn

8 import torch.nn.functional as F

9 from torch.nn import LayerNorm

10 -from transformers.modeling_utils import PreTrainedModel

10 from transformers.models.qwen2_vl import Qwen2VLProcessor

11

12 from vllm.attention.layer import check_upstream_fa_availability

13 from vllm.config import VllmConfig

14 +from vllm.distributed import utils as dist_utils

15 +from vllm.distributed.parallel_state import (

16 + get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size)

17 from vllm.model_executor.layers.activation import SiluAndMul

18 from vllm.model_executor.layers.layernorm import RMSNorm

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

@@ -19,10 +21,14 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,

21 QKVParallelLinear,

22 RowParallelLinear)

23 from vllm.model_executor.layers.quantization import QuantizationConfig

24 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

25 from vllm.model_executor.models.interfaces import (MultiModalEmbeddings,

26 + SupportsLoRA,

27 SupportsMultiModal,

28 SupportsPP)

29 +from vllm.model_executor.models.module_mapping import MultiModelKeys

30 from vllm.model_executor.models.qwen2 import Qwen2ForCausalLM

31 +from vllm.model_executor.models.qwen2_5_vl import Qwen2_5_VisionAttention

32 from vllm.model_executor.models.qwen2_vl import (Qwen2VLDummyInputsBuilder,

33 Qwen2VLMultiModalProcessor,

34 Qwen2VLProcessingInfo)

@@ -38,6 +44,8 @@ from vllm.sequence import IntermediateTensors

44 from vllm.transformers_utils.configs.dotsocr import (DotsOCRConfig,

45 DotsVisionConfig)

46

47 +from .vision import run_dp_sharded_mrope_vision_model

48 +

49 IMAGE_TOKEN = "<|imgpad|>"

50

51

@@ -181,6 +189,8 @@ class PatchMerger(nn.Module):

189 context_dim: int,

190 spatial_merge_size: int = 2,

191 pre_norm="layernorm",

192 + prefix: str = "",

193 + use_data_parallel: bool = False,

194 ) -> None:

195 super().__init__()

196 self.hidden_size = context_dim * (spatial_merge_size**2)

@@ -189,21 +199,21 @@ class PatchMerger(nn.Module):

199 self.ln_q = LayerNorm(context_dim, eps=1e-6)

200 elif self.pre_norm == "rmsnorm":

201 self.ln_q = RMSNorm(context_dim, eps=1e-6)

192 - else:

193 - print("no norm in patch merger")

202

203 self.mlp = nn.Sequential(

204 ColumnParallelLinear(self.hidden_size,

205 self.hidden_size,

206 bias=True,

207 return_bias=False,

200 - disable_tp=True),

208 + prefix=f"{prefix}.0",

209 + disable_tp=use_data_parallel),

210 nn.GELU(),

211 RowParallelLinear(self.hidden_size,

212 dim,

213 bias=True,

214 return_bias=False,

206 - disable_tp=True),

215 + prefix=f"{prefix}.2",

216 + disable_tp=use_data_parallel),

217 )

218

219 def forward(self, x: torch.Tensor) -> torch.Tensor:

@@ -223,38 +233,36 @@ class DotsVisionAttention(nn.Module):

233 bias: bool = True,

234 *,

235 quant_config: Optional[QuantizationConfig] = None,

226 - prefix: str = "") -> None:

236 + prefix: str = "",

237 + use_data_parallel: bool = False) -> None:

238 super().__init__()

228 - from vllm.distributed import (parallel_state,

229 - tensor_model_parallel_all_gather)

230 - from vllm.distributed import utils as dist_utils

239

240 self.embed_dim = dim

233 - self.num_heads = num_heads

234 - self.head_dim = dim // num_heads

235 - self.tp_size = parallel_state.get_tensor_model_parallel_world_size()

236 - self.tp_rank = parallel_state.get_tensor_model_parallel_rank()

237 - self.num_heads_per_partition = dist_utils.divide(

241 + self.tp_size = (1 if use_data_parallel else

242 + get_tensor_model_parallel_world_size())

243 + self.tp_rank = (0 if use_data_parallel else

244 + get_tensor_model_parallel_rank())

245 + self.hidden_size_per_attention_head = dist_utils.divide(dim, num_heads)

246 + self.num_attention_heads_per_partition = dist_utils.divide(

247 num_heads, self.tp_size)

239 -

248 # qkv/proj follow Qwen2-VL style; bias controlled by arg

241 - self.qkv = QKVParallelLinear(hidden_size=dim,

242 - head_size=dim // num_heads,

243 - total_num_heads=num_heads,

244 - bias=bias,

245 - quant_config=quant_config,

246 - prefix=f"{prefix}.qkv")

249 + self.qkv = QKVParallelLinear(

250 + hidden_size=dim,

251 + head_size=self.hidden_size_per_attention_head,

252 + total_num_heads=num_heads,

253 + bias=bias,

254 + quant_config=quant_config,

255 + prefix=f"{prefix}.qkv",

256 + disable_tp=use_data_parallel)

257 self.proj = RowParallelLinear(input_size=dim,

258 output_size=dim,

259 bias=bias,

260 quant_config=quant_config,

251 - prefix=f"{prefix}.proj")

252 - self._all_gather = tensor_model_parallel_all_gather

253 - self._split_last = dist_utils.split_tensor_along_last_dim

254 -

261 + prefix=f"{prefix}.proj",

262 + disable_tp=use_data_parallel)

263 # Select attention backend

256 - self.attn_backend = get_vit_attn_backend(self.head_dim,

257 - torch.get_default_dtype())

264 + self.attn_backend = get_vit_attn_backend(

265 + self.hidden_size_per_attention_head, torch.get_default_dtype())

266 self.use_upstream_fa = False

267 if self.attn_backend != _Backend.FLASH_ATTN and \

268 check_upstream_fa_availability(torch.get_default_dtype()):

@@ -270,19 +278,6 @@ class DotsVisionAttention(nn.Module):

278 _Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA

279 }

280

273 - def _split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:

274 - # qkv: [S, B, 3*dim]

275 - seq_len, bs, _ = qkv.shape

276 - if self.tp_size > 1:

277 - qkv = self._all_gather(qkv)

278 - q, k, v = qkv.chunk(3, dim=2)

279 - if self.tp_size > 1:

280 - q = self._split_last(q, num_partitions=self.tp_size)[self.tp_rank]

281 - k = self._split_last(k, num_partitions=self.tp_size)[self.tp_rank]

282 - v = self._split_last(v, num_partitions=self.tp_size)[self.tp_rank]

283 - new_shape = (seq_len, bs, self.num_heads_per_partition, self.head_dim)

284 - return (q.view(*new_shape), k.view(*new_shape), v.view(*new_shape))

285 -

281 def forward(

282 self,

283 hidden_states: torch.Tensor,

@@ -295,7 +290,7 @@ class DotsVisionAttention(nn.Module):

290 # [S, C] -> [S, B=1, C]

291 x = hidden_states.unsqueeze(1)

292 x, _ = self.qkv(x)

298 - q, k, v = self._split_qkv(x)

293 + q, k, v = Qwen2_5_VisionAttention.split_qkv(self, x)

294 bs = q.shape[1]

295 # [S,B,H,D] -> [B,S,H,D]

296 q = q.permute(1, 0, 2, 3).contiguous()

@@ -327,8 +322,9 @@ class DotsVisionAttention(nn.Module):

322 max_seqlen_k=max_seqlen,

323 dropout_p=0.0,

324 causal=False)

330 - context_layer = output.view(bs, -1, self.num_heads_per_partition,

331 - self.head_dim)

325 + context_layer = output.view(bs, -1,

326 + self.num_attention_heads_per_partition,

327 + self.hidden_size_per_attention_head)

328 elif self.attn_backend == _Backend.TORCH_SDPA:

329 outputs = []

330 for i in range(1, len(cu_seqlens)):

@@ -368,7 +364,8 @@ class DotsSwiGLUFFN(nn.Module):

364 config,

365 *,

366 quant_config: Optional[QuantizationConfig] = None,

371 - prefix: str = ""):

367 + prefix: str = "",

368 + use_data_parallel: bool = False):

369 super().__init__()

370 hidden_features = config.intermediate_size

371 in_features = config.embed_dim

@@ -380,13 +377,13 @@ class DotsSwiGLUFFN(nn.Module):

377 bias=bias,

378 quant_config=quant_config,

379 prefix=f"{prefix}.fc13",

383 - disable_tp=True)

380 + disable_tp=use_data_parallel)

381 self.fc2 = RowParallelLinear(hidden_features,

382 in_features,

383 bias=bias,

384 quant_config=quant_config,

385 prefix=f"{prefix}.fc2",

389 - disable_tp=True)

386 + disable_tp=use_data_parallel)

387 self.act_fn = SiluAndMul()

388

389 def forward(self, x: torch.Tensor) -> torch.Tensor:

@@ -397,28 +394,36 @@ class DotsSwiGLUFFN(nn.Module):

394

395 def load_weights(self, weights: Iterable[tuple[str,

396 torch.Tensor]]) -> set[str]:

400 - params = dict(self.named_parameters())

401 - loaded: set[str] = set()

402 - for name, w in weights:

403 - # Map fc1 -> fc13 (shard 0)

404 - if name.startswith("fc1."):

405 - tgt = name.replace("fc1.", "fc13.")

406 - if tgt in params:

407 - params[tgt].weight_loader(params[tgt], w, 0)

408 - loaded.add(tgt)

409 - continue

410 - # Map fc3 -> fc13 (shard 1)

411 - if name.startswith("fc3."):

412 - tgt = name.replace("fc3.", "fc13.")

413 - if tgt in params:

414 - params[tgt].weight_loader(params[tgt], w, 1)

415 - loaded.add(tgt)

416 - continue

417 - # Pass-through for fc2 and others

418 - if name in params:

419 - params[name].weight_loader(params[name], w)

420 - loaded.add(name)

421 - return loaded

397 + stacked_params_mapping = [

398 + ("fc13", "fc1", 0),

399 + ("fc13", "fc3", 1),

400 + ]

401 + params_dict = dict(self.named_parameters())

402 + loaded_params: set[str] = set()

403 + for name, loaded_weight in weights:

404 +

405 + for param_name, weight_name, shard_id in stacked_params_mapping:

406 + if weight_name not in name:

407 + continue

408 + name = name.replace(weight_name, param_name)

409 + # Skip loading extra bias for GPTQ models.

410 + if name.endswith(".bias") and name not in params_dict:

411 + continue

412 + param = params_dict[name]

413 + weight_loader = param.weight_loader

414 + weight_loader(param, loaded_weight, shard_id)

415 + break

416 + else:

417 + # Skip loading extra bias for GPTQ models.

418 + if name.endswith(".bias") and name not in params_dict:

419 + continue

420 +

421 + param = params_dict[name]

422 + weight_loader = getattr(param, "weight_loader",

423 + default_weight_loader)

424 + weight_loader(param, loaded_weight)

425 + loaded_params.add(name)

426 + return loaded_params

427

428

429 class DotsPatchEmbed(nn.Module):

@@ -463,25 +468,28 @@ class DotsViTPreprocessor(nn.Module):

468

469 class DotsVisionBlock(nn.Module):

470

466 - def __init__(self,

467 - config,

468 - *,

469 - quant_config: Optional[QuantizationConfig] = None,

470 - prefix: str = ""):

471 + def __init__(

472 + self,

473 + config,

474 + *,

475 + quant_config: Optional[QuantizationConfig] = None,

476 + prefix: str = "",

477 + use_data_parallel: bool = False,

478 + ):

479 super().__init__()

480

473 - self.attn = DotsVisionAttention(

474 - config,

475 - config.embed_dim,

476 - num_heads=config.num_attention_heads,

477 - bias=config.use_bias,

478 - quant_config=quant_config,

479 - prefix=f"{prefix}.attn",

480 - )

481 + self.attn = DotsVisionAttention(config,

482 + config.embed_dim,

483 + num_heads=config.num_attention_heads,

484 + bias=config.use_bias,

485 + quant_config=quant_config,

486 + prefix=f"{prefix}.attn",

487 + use_data_parallel=use_data_parallel)

488 self.norm1 = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)

489 self.mlp = DotsSwiGLUFFN(config,

490 quant_config=quant_config,

484 - prefix=f"{prefix}.mlp")

491 + prefix=f"{prefix}.mlp",

492 + use_data_parallel=use_data_parallel)

493 self.norm2 = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)

494

495 def forward(self,

@@ -502,7 +510,7 @@ class DotsVisionBlock(nn.Module):

510 return hidden_states

511

512

505 -class DotsVisionTransformer(PreTrainedModel):

513 +class DotsVisionTransformer(nn.Module):

514

515 def __init__(

516 self,

@@ -512,8 +520,9 @@ class DotsVisionTransformer(PreTrainedModel):

520 num_hidden_layers_override: Optional[int] = None,

521 require_post_norm: Optional[bool] = None,

522 prefix: str = "",

523 + use_data_parallel: bool = False,

524 ) -> None:

516 - super().__init__(config)

525 + super().__init__()

526 self.config = config

527 self.spatial_merge_size = config.spatial_merge_size

528

@@ -526,14 +535,15 @@ class DotsVisionTransformer(PreTrainedModel):

535 if self.attn_backend != _Backend.FLASH_ATTN and \

536 check_upstream_fa_availability(torch.get_default_dtype()):

537 self.attn_backend = _Backend.FLASH_ATTN

529 -

538 + self.out_hidden_size = config.hidden_size

539 # Keep blocks for compatibility with other vision towers

540 num_layers = (config.num_hidden_layers if num_hidden_layers_override

541 is None else num_hidden_layers_override)

542 self.blocks = nn.ModuleList([

543 DotsVisionBlock(config,

544 quant_config=quant_config,

536 - prefix=f"{prefix}.blocks.{i}")

545 + prefix=f"{prefix}.blocks.{i}",

546 + use_data_parallel=use_data_parallel)

547 for i in range(num_layers)

548 ])

549 if require_post_norm is None:

@@ -548,6 +558,7 @@ class DotsVisionTransformer(PreTrainedModel):

558 dim=config.hidden_size,

559 context_dim=config.embed_dim,

560 spatial_merge_size=config.spatial_merge_size,

561 + use_data_parallel=use_data_parallel,

562 )

563

564 @property

@@ -604,7 +615,11 @@ class DotsVisionTransformer(PreTrainedModel):

615 return max_seqlen, seqlens

616

617 def forward(self, hidden_states: torch.Tensor,

607 - grid_thw: torch.Tensor) -> torch.Tensor:

618 + grid_thw: list[list[int]]) -> torch.Tensor:

619 + # Convert grid_thw to tensor (always expecting list format now)

620 + grid_thw = torch.tensor(grid_thw,

621 + device=hidden_states.device,

622 + dtype=torch.long)

623 hidden_states = hidden_states.to(self.dtype)

624 hidden_states = self.patch_embed(hidden_states, grid_thw)

625

@@ -638,7 +653,8 @@ class DotsVisionTransformer(PreTrainedModel):

653 info=DotsOCRProcessingInfo,

654 dummy_inputs=DotsOCRDummyInputsBuilder,

655 )

641 -class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

656 +class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP,

657 + SupportsLoRA):

658 hf_to_vllm_mapper = WeightsMapper(

659 orig_to_new_substr={

660 ".attn.qkv_proj.": ".attn.qkv.",

@@ -650,6 +666,21 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

666 },

667 )

668

669 + packed_modules_mapping = {

670 + "qkv_proj": [

671 + "q_proj",

672 + "k_proj",

673 + "v_proj",

674 + ],

675 + "gate_up_proj": [

676 + "gate_proj",

677 + "up_proj",

678 + ],

679 + ".attn.qkv": [".attn.qkv"],

680 + "fc13": ["fc1", "fc3"],

681 + }

682 + supports_encoder_tp_data = True

683 +

684 @classmethod

685 def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

686 if modality.startswith("image"):

@@ -660,19 +691,18 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

691

692 self.config: DotsOCRConfig = vllm_config.model_config.hf_config

693 self.quant_config = vllm_config.quant_config

663 - self.multimodal_config = vllm_config.model_config.multimodal_config

664 -

694 + multimodal_config = vllm_config.model_config.multimodal_config

695 + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"

696 if isinstance(self.config.vision_config, dict):

697 vision_config = DotsVisionConfig(**self.config.vision_config)

698 self.config.vision_config = vision_config

699 else:

700 vision_config = self.config.vision_config

670 -

701 self.vision_tower = DotsVisionTransformer(

702 vision_config,

703 quant_config=self.quant_config,

704 prefix=maybe_prefix(prefix, "vision_tower"),

675 - )

705 + use_data_parallel=self.use_data_parallel)

706 self.language_model: Qwen2ForCausalLM = init_vllm_registered_model(

707 vllm_config=vllm_config,

708 hf_config=self.config,

@@ -744,8 +774,17 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

774 else:

775 pixel_values = image_input["pixel_values"].type(

776 self.vision_tower.dtype)

747 - image_embeds = self.vision_tower(

748 - pixel_values, grid_thw)[:, :self.config.hidden_size]

777 +

778 + if self.use_data_parallel:

779 + return run_dp_sharded_mrope_vision_model(

780 + self.vision_tower,

781 + pixel_values,

782 + grid_thw_list,

783 + rope_type="rope_3d",

784 + )

785 + else:

786 + image_embeds = self.vision_tower(

787 + pixel_values, grid_thw)[:, :self.config.hidden_size]

788

789 # Split concatenated embeddings for each image item.

790 merge_size = self.vision_tower.spatial_merge_size

@@ -822,3 +861,13 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

861 torch.Tensor]]) -> set[str]:

862 loader = AutoWeightsLoader(self)

863 return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

864 +

865 + def get_mm_mapping(self) -> MultiModelKeys:

866 + """

867 + Get the module prefix in multimodal models

868 + """

869 + return MultiModelKeys.from_string_field(

870 + language_model="language_model",

871 + connector="vision_tower.merger",

872 + tower_model="vision_tower.",

873 + )

[Model] Improve DotsOCRForCausalLM (#25466)