103f1ec8d3 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -27,7 +27,8 @@ from vllm.model_executor.sampling_metadata import SamplingMetadata

27 from vllm.sequence import IntermediateTensors

28 from vllm.utils import cdiv

29

30 -from .utils import extract_layer_index, maybe_prefix

30 +from .utils import (AutoWeightsLoader, WeightsMapper, extract_layer_index,

31 + maybe_prefix)

32

33

34 class OAIAttention(nn.Module):

@@ -203,6 +204,7 @@ class GptOssModel(nn.Module):

204 super().__init__()

205 self.config = vllm_config.model_config.hf_config

206 self.quant_config = vllm_config.quant_config

207 + self.parallel_config = vllm_config.parallel_config

208 self.config.hidden_size = self.config.hidden_size

209 self.embedding = VocabParallelEmbedding(

210 self.config.vocab_size,

@@ -225,64 +227,26 @@ class GptOssModel(nn.Module):

227 x = self.norm(x)

228 return x

229

228 -

229 -class GptOssForCausalLM(nn.Module):

230 -

231 - def __init__(

232 - self,

233 - vllm_config: VllmConfig,

234 - prefix: str = "",

235 - ):

236 - super().__init__()

237 - self.vllm_config = vllm_config

238 - self.model_config = vllm_config.model_config.hf_config

239 - self.model = GptOssModel(

240 - vllm_config=vllm_config,

241 - prefix=maybe_prefix(prefix, "model"),

242 - )

243 - self.lm_head = ParallelLMHead(

244 - self.model_config.vocab_size,

245 - self.model_config.hidden_size,

246 - )

247 - self.logits_processor = LogitsProcessor(self.model_config.vocab_size)

248 -

249 - def forward(self,

250 - input_ids: torch.Tensor,

251 - positions: torch.Tensor,

252 - intermediate_tensors: Optional[IntermediateTensors] = None,

253 - inputs_embeds: Optional[torch.Tensor] = None) -> torch.Tensor:

254 - assert intermediate_tensors is None

255 - assert inputs_embeds is None

256 - return self.model(input_ids, positions)

257 -

258 - def compute_logits(self, hidden_states: torch.Tensor,

259 - sampling_metadata: SamplingMetadata) -> torch.Tensor:

260 - logits = self.logits_processor(self.lm_head, hidden_states,

261 - sampling_metadata)

262 - return logits

263 -

230 def _load_weights_mxfp4(

265 - self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

266 - rename_mapping = {

267 - "self_attn": "attn",

268 - "input_layernorm.weight": "attn.norm.weight",

269 - "post_attention_layernorm.weight": "mlp.norm.weight",

270 - "embed_tokens": "embedding",

271 - }

272 -

273 - def maybe_rename(name: str) -> str:

274 - for remap_name, new_name in rename_mapping.items():

275 - if remap_name in name:

276 - return name.replace(remap_name, new_name)

277 - return name

278 -

231 + self,

232 + ep_rank_end: int,

233 + ep_rank_start: int,

234 + heads_per_rank: int,

235 + head_start: int,

236 + weights: Iterable[tuple[str, torch.Tensor]],

237 + stacked_params_mapping: list[tuple[str, ...]],

238 + ) -> set[str]:

239 params_dict = dict(self.named_parameters())

240 loaded_params: set[str] = set()

241 +

242 mxfp4_block = 32

243 + use_ep = self.parallel_config.enable_expert_parallel

244 + num_experts = self.config.num_local_experts

245

246 tp_rank = get_tensor_model_parallel_rank()

247 tp_size = get_tensor_model_parallel_world_size()

285 - intermediate_size = self.model_config.intermediate_size

248 +

249 + intermediate_size = self.config.intermediate_size

250 intermediate_size_block = intermediate_size // mxfp4_block

251 per_rank_intermediate_size_block = cdiv(intermediate_size_block,

252 tp_size)

@@ -294,33 +258,12 @@ class GptOssForCausalLM(nn.Module):

258 tp_rank_end = min((tp_rank + 1) * per_rank_intermediate_size,

259 intermediate_size)

260

297 - # Attention heads per rank

298 - heads_per_rank = self.model_config.num_attention_heads // tp_size

299 - head_start = tp_rank * heads_per_rank

300 -

301 - use_ep = self.vllm_config.parallel_config.enable_expert_parallel

302 - ep_size = get_ep_group().world_size

303 - ep_rank = get_ep_group().rank

304 - num_experts = self.model_config.num_local_experts

305 - experts_per_rank = num_experts // ep_size

306 - ep_rank_start = ep_rank * experts_per_rank

307 - ep_rank_end = (ep_rank + 1) * experts_per_rank

308 -

261 for name, weight in weights:

262 # FIXME(woosuk): Remove this after testing.

263 weight = weight.cuda()

264

313 - if "gate_up_proj_blocks" in name:

314 - # Handle MLP gate and up projection weights

315 - new_name = name.replace("gate_up_proj_blocks", "w13_weight")

316 -

317 - # flat weight from (E, 2 * N, block_size, entry_per_block)

318 - # to (E, 2 * N, -1), shouldn't trigger copy for contiguous

319 - weight = weight.view(num_experts, 2 * intermediate_size,

320 - -1).contiguous()

321 -

322 - # Extract gate and up projection parts

323 - # since the weight is shuffled, we can slice directly

265 + if ".w13_weight_scale" in name:

266 + # Handle MLP gate and up projection weights scale

267 if use_ep:

268 narrow_weight = weight[ep_rank_start:ep_rank_end, ...]

269 else:

@@ -328,43 +271,44 @@ class GptOssForCausalLM(nn.Module):

271 2 * tp_rank_start:2 * tp_rank_end,

272 ...]

273

331 - param = params_dict[new_name]

274 + param = params_dict[name]

275 weight_loader = getattr(param, "weight_loader",

276 default_weight_loader)

277 weight_loader(param,

278 narrow_weight,

336 - weight_name=new_name,

279 + weight_name=name,

280 shard_id=None,

281 expert_id=None)

339 - loaded_params.add(new_name)

340 -

341 - elif "down_proj_blocks" in name:

282 + loaded_params.add(name)

283 + continue

284 + elif ".w2_weight_scale" in name:

285 # Handle MLP down projection weights

343 - new_name = name.replace("down_proj_blocks", "w2_weight")

344 - # same flatten here, but since 2 mx4 value are packed in 1

345 - # uint8, divide by 2

346 - weight = weight.view(num_experts, -1,

347 - intermediate_size // 2).contiguous()

286 if use_ep:

287 narrow_weight = weight[ep_rank_start:ep_rank_end, ...]

288 else:

351 - narrow_weight = weight[...,

352 - tp_rank_start // 2:tp_rank_end // 2]

289 + narrow_weight = weight[..., tp_rank_start //

290 + mxfp4_block:tp_rank_end //

291 + mxfp4_block]

292

354 - param = params_dict[new_name]

293 + param = params_dict[name]

294 weight_loader = getattr(param, "weight_loader",

295 default_weight_loader)

296 weight_loader(param,

297 narrow_weight,

359 - weight_name=new_name,

298 + weight_name=name,

299 shard_id=None,

300 expert_id=None)

362 - loaded_params.add(new_name)

301 + loaded_params.add(name)

302 + continue

303 + elif ".w13_weight" in name:

304 + # Handle MLP gate and up projection weights

305 + # flat weight from (E, 2 * N, block_size, entry_per_block)

306 + # to (E, 2 * N, -1), shouldn't trigger copy for contiguous

307 + weight = weight.view(num_experts, 2 * intermediate_size,

308 + -1).contiguous()

309

364 - elif "gate_up_proj_scales" in name:

365 - # Handle MLP gate and up projection weights scale

366 - new_name = name.replace("gate_up_proj_scales",

367 - "w13_weight_scale")

310 + # Extract gate and up projection parts

311 + # since the weight is shuffled, we can slice directly

312 if use_ep:

313 narrow_weight = weight[ep_rank_start:ep_rank_end, ...]

314 else:

@@ -372,39 +316,40 @@ class GptOssForCausalLM(nn.Module):

316 2 * tp_rank_start:2 * tp_rank_end,

317 ...]

318

375 - param = params_dict[new_name]

319 + param = params_dict[name]

320 weight_loader = getattr(param, "weight_loader",

321 default_weight_loader)

322 weight_loader(param,

323 narrow_weight,

380 - weight_name=new_name,

324 + weight_name=name,

325 shard_id=None,

326 expert_id=None)

383 - loaded_params.add(new_name)

384 -

385 - elif "down_proj_scales" in name:

327 + loaded_params.add(name)

328 + continue

329 + elif ".w2_weight" in name:

330 # Handle MLP down projection weights

387 - new_name = name.replace("down_proj_scales", "w2_weight_scale")

331 + # same flatten here, but since 2 mx4 value are packed in 1

332 + # uint8, divide by 2

333 + weight = weight.view(num_experts, -1,

334 + intermediate_size // 2).contiguous()

335 if use_ep:

336 narrow_weight = weight[ep_rank_start:ep_rank_end, ...]

337 else:

391 - narrow_weight = weight[..., tp_rank_start //

392 - mxfp4_block:tp_rank_end //

393 - mxfp4_block]

338 + narrow_weight = weight[...,

339 + tp_rank_start // 2:tp_rank_end // 2]

340

395 - param = params_dict[new_name]

341 + param = params_dict[name]

342 weight_loader = getattr(param, "weight_loader",

343 default_weight_loader)

344 weight_loader(param,

345 narrow_weight,

400 - weight_name=new_name,

346 + weight_name=name,

347 shard_id=None,

348 expert_id=None)

403 - loaded_params.add(new_name)

404 - elif "gate_up_proj_bias" in name:

349 + loaded_params.add(name)

350 + continue

351 + elif ".w13_bias" in name:

352 # Handle MLP gate and up projection biases

406 - new_name = name.replace("gate_up_proj_bias", "w13_bias")

407 -

353 # Extract gate and up projection bias parts

354 if use_ep:

355 narrow_weight = weight[ep_rank_start:ep_rank_end, ...]

@@ -412,20 +357,19 @@ class GptOssForCausalLM(nn.Module):

357 narrow_weight = weight[:,

358 2 * tp_rank_start:2 * tp_rank_end]

359

415 - param = params_dict[new_name]

360 + param = params_dict[name]

361 weight_loader = getattr(param, "weight_loader",

362 default_weight_loader)

363 weight_loader(param,

364 narrow_weight,

420 - weight_name=new_name,

365 + weight_name=name,

366 shard_id=None,

367 expert_id=None)

423 - loaded_params.add(new_name)

424 -

425 - elif "down_proj_bias" in name:

368 + loaded_params.add(name)

369 + continue

370 + elif ".w2_bias" in name:

371 # Handle MLP down projection bias

427 - new_name = name.replace("down_proj_bias", "w2_bias")

428 - param = params_dict[new_name]

372 + param = params_dict[name]

373 weight_loader = getattr(param, "weight_loader",

374 default_weight_loader)

375 if use_ep:

@@ -436,87 +380,69 @@ class GptOssForCausalLM(nn.Module):

380 weight.zero_()

381 weight_loader(param,

382 weight,

439 - weight_name=new_name,

383 + weight_name=name,

384 shard_id=None,

385 expert_id=None)

442 - loaded_params.add(new_name)

386 + loaded_params.add(name)

387 + continue

388 elif "sinks" in name:

389 # Handle attention sinks (distributed across ranks)

445 - name = name.replace("self_attn", "attn")

390 param = params_dict[name]

391 narrow_weight = weight.narrow(0, head_start, heads_per_rank)

392 param.data.copy_(narrow_weight)

393 loaded_params.add(name)

450 - elif "q_proj" in name or "k_proj" in name or "v_proj" in name:

451 - shard_id = ("q" if "q_proj" in name else

452 - "k" if "k_proj" in name else "v")

453 - name = name.replace("self_attn", "attn")

454 - param_name = name.replace(f"{shard_id}_proj", "qkv")

455 - param = params_dict[param_name]

456 - weight_loader = param.weight_loader

457 - weight_loader(param, weight, loaded_shard_id=shard_id)

458 - loaded_params.add(param_name)

394 + continue

395 + for param_name, weight_name, shard_id in stacked_params_mapping:

396 + if weight_name not in name:

397 + continue

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

399 + param = params_dict[name]

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

401 + default_weight_loader)

402 + if weight_loader == default_weight_loader:

403 + weight_loader(param, weight)

404 + else:

405 + weight_loader(param, weight, shard_id)

406 + break

407 else:

408 # Handle all other weights with potential renaming

461 - renamed_name = maybe_rename(name)

462 - if renamed_name not in params_dict:

409 + if name not in params_dict:

410 continue

464 - param = params_dict[renamed_name]

411 + param = params_dict[name]

412 weight_loader = getattr(param, "weight_loader",

413 default_weight_loader)

414 weight_loader(param, weight)

468 - loaded_params.add(renamed_name)

469 -

415 + loaded_params.add(name)

416 return loaded_params

417

418 def _load_weights_other(

473 - self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

474 - rename_mapping = {

475 - "self_attn": "attn",

476 - "input_layernorm.weight": "attn.norm.weight",

477 - "post_attention_layernorm.weight": "mlp.norm.weight",

478 - "embed_tokens": "embedding",

479 - }

480 -

481 - def maybe_rename(name: str) -> str:

482 - for remap_name, new_name in rename_mapping.items():

483 - if remap_name in name:

484 - return name.replace(remap_name, new_name)

485 - return name

486 -

419 + self,

420 + ep_rank_start: int,

421 + ep_rank_end: int,

422 + heads_per_rank: int,

423 + head_start: int,

424 + weights: Iterable[tuple[str, torch.Tensor]],

425 + stacked_params_mapping: list[tuple[str, ...]],

426 + ) -> set[str]:

427 params_dict = dict(self.named_parameters())

428 loaded_params: set[str] = set()

429

430 + use_ep = self.parallel_config.enable_expert_parallel

431 +

432 tp_rank = get_tensor_model_parallel_rank()

433 tp_size = get_tensor_model_parallel_world_size()

492 - intermediate_size = self.model_config.intermediate_size

434

435 + intermediate_size = self.config.intermediate_size

436 per_rank_intermediate_size = cdiv(intermediate_size, tp_size)

437 # Calculate common slicing bounds for current rank

438 tp_rank_start = tp_rank * per_rank_intermediate_size

439 tp_rank_end = min((tp_rank + 1) * per_rank_intermediate_size,

440 intermediate_size)

441

500 - # Attention heads per rank

501 - heads_per_rank = self.model_config.num_attention_heads // tp_size

502 - head_start = tp_rank * heads_per_rank

503 -

504 - use_ep = self.vllm_config.parallel_config.enable_expert_parallel

505 - ep_size = get_ep_group().world_size

506 - ep_rank = get_ep_group().rank

507 - num_experts = self.model_config.num_local_experts

508 - experts_per_rank = num_experts // ep_size

509 - ep_rank_start = ep_rank * experts_per_rank

510 - ep_rank_end = (ep_rank + 1) * experts_per_rank

511 -

442 for name, weight in weights:

513 - if ".experts.gate_up_proj" in name and "bias" not in name:

443 + if ".w13_weight" in name:

444 # Handle MLP gate and up projection weights

515 - new_name = name.replace(".experts.gate_up_proj",

516 - ".experts.w13_weight")

517 -

445 # Extract gate and up projection parts

519 - # since the weight is shuffled, we can slice directly

446 if use_ep:

447 narrow_weight = weight[ep_rank_start:ep_rank_end, ...]

448 else:

@@ -524,30 +450,25 @@ class GptOssForCausalLM(nn.Module):

450 2 * tp_rank_start:2 * tp_rank_end]

451

452 narrow_weight = narrow_weight.permute(0, 2, 1).contiguous()

527 - param = params_dict[new_name]

453 + param = params_dict[name]

454

455 param.copy_(narrow_weight)

530 - loaded_params.add(new_name)

531 -

532 - elif ".experts.down_proj" in name and "bias" not in name:

456 + loaded_params.add(name)

457 + continue

458 + elif ".w2_weight" in name:

459 # Handle MLP down projection weights

534 - new_name = name.replace(".experts.down_proj",

535 - ".experts.w2_weight")

536 -

460 if use_ep:

461 narrow_weight = weight[ep_rank_start:ep_rank_end, ...]

462 else:

463 narrow_weight = weight[:, tp_rank_start:tp_rank_end, :]

464 narrow_weight = narrow_weight.permute(0, 2, 1).contiguous()

542 - param = params_dict[new_name]

465 + param = params_dict[name]

466

467 param.copy_(narrow_weight)

545 - loaded_params.add(new_name)

546 -

547 - elif "gate_up_proj_bias" in name:

468 + loaded_params.add(name)

469 + continue

470 + elif ".w13_bias" in name:

471 # Handle MLP gate and up projection biases

549 - new_name = name.replace("gate_up_proj_bias", "w13_bias")

550 -

472 # Extract gate and up projection bias parts

473 if use_ep:

474 narrow_weight = weight[ep_rank_start:ep_rank_end, ...]

@@ -555,60 +476,155 @@ class GptOssForCausalLM(nn.Module):

476 narrow_weight = weight[:,

477 2 * tp_rank_start:2 * tp_rank_end]

478

558 - param = params_dict[new_name]

559 -

479 + param = params_dict[name]

480 param.copy_(narrow_weight)

561 - loaded_params.add(new_name)

562 -

563 - elif "down_proj_bias" in name:

481 + loaded_params.add(name)

482 + continue

483 + elif ".w2_bias" in name:

484 # Handle MLP down projection bias

565 - new_name = name.replace("down_proj_bias", "w2_bias")

566 -

485 if use_ep:

486 weight = weight[ep_rank_start:ep_rank_end, ...]

487 else:

488 # (only load on rank 0 to avoid duplication)

489 if tp_rank != 0:

490 weight.zero_()

573 - param = params_dict[new_name]

491 + param = params_dict[name]

492 param.copy_(weight)

575 - loaded_params.add(new_name)

493 + loaded_params.add(name)

494 + continue

495 elif "sinks" in name:

496 # Handle attention sinks (distributed across ranks)

578 - name = name.replace("self_attn", "attn")

497 param = params_dict[name]

498 narrow_weight = weight.narrow(0, head_start, heads_per_rank)

499 param.data.copy_(narrow_weight)

500 loaded_params.add(name)

583 - elif "q_proj" in name or "k_proj" in name or "v_proj" in name:

584 - shard_id = ("q" if "q_proj" in name else

585 - "k" if "k_proj" in name else "v")

586 - name = name.replace("self_attn", "attn")

587 - param_name = name.replace(f"{shard_id}_proj", "qkv")

588 - param = params_dict[param_name]

589 - weight_loader = param.weight_loader

590 - weight_loader(param, weight, loaded_shard_id=shard_id)

591 - loaded_params.add(param_name)

501 + continue

502 + for param_name, weight_name, shard_id in stacked_params_mapping:

503 + if weight_name not in name:

504 + continue

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

506 + param = params_dict[name]

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

508 + default_weight_loader)

509 + if weight_loader == default_weight_loader:

510 + weight_loader(param, weight)

511 + else:

512 + weight_loader(param, weight, shard_id)

513 + break

514 else:

515 # Handle all other weights with potential renaming

594 -

595 - renamed_name = maybe_rename(name)

596 - if renamed_name not in params_dict:

516 + if name not in params_dict:

517 continue

598 - param = params_dict[renamed_name]

518 + param = params_dict[name]

519 weight_loader = getattr(param, "weight_loader",

520 default_weight_loader)

521 weight_loader(param, weight)

602 - loaded_params.add(renamed_name)

603 -

522 + loaded_params.add(name)

523 return loaded_params

524

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

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

608 - quant_method = (self.model_config.quantization_config['quant_method']

609 - if hasattr(self.model_config, "quantization_config")

610 - else None)

527 + stacked_params_mapping = [

528 + # (param_name, shard_name, shard_id)

529 + (".qkv", ".q_proj", "q"),

530 + (".qkv", ".k_proj", "k"),

531 + (".qkv", ".v_proj", "v"),

532 + ]

533 +

534 + tp_rank = get_tensor_model_parallel_rank()

535 + tp_size = get_tensor_model_parallel_world_size()

536 +

537 + # Attention heads per rank

538 + heads_per_rank = self.config.num_attention_heads // tp_size

539 + head_start = tp_rank * heads_per_rank

540 +

541 + ep_size = get_ep_group().world_size

542 + ep_rank = get_ep_group().rank

543 + num_experts = self.config.num_local_experts

544 + experts_per_rank = num_experts // ep_size

545 + ep_rank_start = ep_rank * experts_per_rank

546 + ep_rank_end = (ep_rank + 1) * experts_per_rank

547 +

548 + quant_method = (self.config.quantization_config['quant_method'] if

549 + hasattr(self.config, "quantization_config") else None)

550 if quant_method == "mxfp4":

612 - return self._load_weights_mxfp4(weights)

551 + return self._load_weights_mxfp4(ep_rank_end, ep_rank_start,

552 + heads_per_rank, head_start,

553 + weights, stacked_params_mapping)

554 else:

614 - return self._load_weights_other(weights)

555 + return self._load_weights_other(ep_rank_end, ep_rank_start,

556 + heads_per_rank, head_start,

557 + weights, stacked_params_mapping)

558 +

559 +

560 +class GptOssForCausalLM(nn.Module):

561 + packed_modules_mapping = {"qkv": ["q_proj", "k_proj", "v_proj"]}

562 +

563 + hf_to_vllm_mapper = WeightsMapper(

564 + orig_to_new_substr={

565 + ".self_attn.": ".attn.",

566 + ".post_attention_layernorm.": ".mlp.norm.",

567 + },

568 + orig_to_new_suffix={

569 + ".embed_tokens.weight": ".embedding.weight",

570 + ".input_layernorm.weight": ".attn.norm.weight",

571 + ".post_attention_layernorm.weight": ".mlp.norm.weight",

572 +

573 + # MoE MXFP4 weights

574 + ".gate_up_proj_blocks": ".w13_weight",

575 + ".down_proj_blocks": ".w2_weight",

576 + ".gate_up_proj_scales": ".w13_weight_scale",

577 + ".down_proj_scales": ".w2_weight_scale",

578 +

579 + # MoE other weights

580 + ".gate_up_proj": ".w13_weight",

581 + ".down_proj": ".w2_weight",

582 +

583 + # MoE Bias

584 + ".gate_up_proj_bias": ".w13_bias",

585 + ".down_proj_bias": ".w2_bias",

586 + },

587 + )

588 +

589 + def __init__(

590 + self,

591 + vllm_config: VllmConfig,

592 + prefix: str = "",

593 + ):

594 + super().__init__()

595 + self.vllm_config = vllm_config

596 + self.config = vllm_config.model_config.hf_config

597 +

598 + self.model = GptOssModel(

599 + vllm_config=vllm_config,

600 + prefix=maybe_prefix(prefix, "model"),

601 + )

602 + self.lm_head = ParallelLMHead(

603 + self.config.vocab_size,

604 + self.config.hidden_size,

605 + )

606 + self.logits_processor = LogitsProcessor(self.config.vocab_size)

607 +

608 + def forward(self,

609 + input_ids: torch.Tensor,

610 + positions: torch.Tensor,

611 + intermediate_tensors: Optional[IntermediateTensors] = None,

612 + inputs_embeds: Optional[torch.Tensor] = None) -> torch.Tensor:

613 + assert intermediate_tensors is None

614 + assert inputs_embeds is None

615 + return self.model(input_ids, positions)

616 +

617 + def compute_logits(self, hidden_states: torch.Tensor,

618 + sampling_metadata: SamplingMetadata) -> torch.Tensor:

619 + logits = self.logits_processor(self.lm_head, hidden_states,

620 + sampling_metadata)

621 + return logits

622 +

623 + def load_weights(self, weights: Iterable[tuple[str,

624 + torch.Tensor]]) -> set[str]:

625 + loader = AutoWeightsLoader(

626 + self,

627 + skip_prefixes=(["lm_head."]

628 + if self.config.tie_word_embeddings else None),

629 + )

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

[Model] use autoWeightsLoader for gptoss (#22446)