6ece16c4fe · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -0,0 +1,649 @@

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

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

3 +"""

4 +Create a reduced-layer version of the Maverick model for testing purposes.

5 +

6 +This script creates a new model with fewer layers by:

7 +1. Loading the original Maverick model configuration

8 +2. Creating a reduced configuration

9 +3. Generating compatible safetensors files with appropriate weights

10 +4. Creating the necessary index files for vLLM compatibility

11 +"""

12 +

13 +import json

14 +import shutil

15 +from pathlib import Path

16 +from typing import Any

17 +

18 +import pytest

19 +import torch

20 +from safetensors.torch import save_file

21 +from transformers import (AutoConfig, AutoProcessor, AutoTokenizer,

22 + GenerationConfig)

23 +

24 +from vllm import LLM, SamplingParams

25 +

26 +# Sample prompts for testing

27 +PROMPTS: list[str] = [

28 + "Hello, my name is",

29 + "The president of the United States is",

30 + "The capital of France is",

31 + "The future of AI is",

32 +]

33 +

34 +

35 +def run_maverick_serving(model: str):

36 + """Test Llama-4-Maverick model with vLLM LLM class using CLI equivalent

37 + options with reduced layers.

38 + """

39 +

40 + try:

41 + sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

42 +

43 + llm = LLM(

44 + model=model,

45 + max_model_len=2048,

46 + enforce_eager=True,

47 + tensor_parallel_size=8,

48 + enable_expert_parallel=True,

49 + trust_remote_code=True,

50 + gpu_memory_utilization=0.4,

51 + kv_cache_dtype="fp8",

52 + )

53 +

54 + outputs = llm.generate(PROMPTS, sampling_params)

55 +

56 + # Print the outputs

57 + print("\nGenerated Outputs:\n" + "-" * 60)

58 + for output in outputs:

59 + prompt = output.prompt

60 + generated_text = output.outputs[0].text

61 + print(f"Prompt: {prompt!r}")

62 + print(f"Output: {generated_text!r}")

63 + print("-" * 60)

64 +

65 + except Exception as e:

66 + print(f"Error initializing or running model: {e}")

67 + raise

68 +

69 +

70 +def create_reduced_maverick_model(

71 + original_model_name:

72 + str = "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",

73 + output_dir: str = "/tmp/reduced_maverick",

74 + text_layers: int = 4,

75 + num_experts: int = 4,

76 + vision_layers: int = 2,

77 + force_recreate: bool = False,

78 +) -> str:

79 + """

80 + Create a reduced-layer version of the Maverick model.

81 +

82 + Args:

83 + original_model_name: Name of the original Maverick model

84 + output_dir: Directory to save the reduced model

85 + text_layers: Number of text transformer layers

86 + num_experts: Number of experts per layer

87 + vision_layers: Number of vision transformer layers

88 + force_recreate: Whether to recreate if output_dir already exists

89 +

90 + Returns:

91 + Path to the created reduced model directory

92 + """

93 +

94 + print(

95 + f"Creating reduced Maverick model with {text_layers} text layers and "

96 + f"{vision_layers} vision layers...")

97 +

98 + # Create output directory

99 + output_path = Path(output_dir)

100 + if output_path.exists():

101 + if force_recreate:

102 + shutil.rmtree(output_path)

103 + else:

104 + print(f"Output directory {output_dir} already exists. "

105 + "Use --force-recreate to overwrite.")

106 + return str(output_path)

107 +

108 + output_path.mkdir(parents=True, exist_ok=True)

109 +

110 + try:

111 + print("Loading original model configuration...")

112 + original_config = AutoConfig.from_pretrained(original_model_name,

113 + trust_remote_code=True)

114 +

115 + print("Creating reduced configuration...")

116 + reduced_config = create_reduced_config(original_config, text_layers,

117 + num_experts, vision_layers)

118 +

119 + config_path = output_path / "config.json"

120 + with open(config_path, "w") as f:

121 + json.dump(reduced_config, f, indent=2)

122 + print(f"Saved reduced config to {config_path}")

123 +

124 + print("Copying tokenizer files...")

125 + copy_tokenizer_files(original_model_name, output_path)

126 +

127 + print("Creating reduced safetensors files...")

128 + create_reduced_safetensors(original_config, reduced_config,

129 + output_path)

130 +

131 + print("Creating preprocessor config...")

132 + create_preprocessor_config(original_config, output_path)

133 +

134 + try:

135 + gen_config = GenerationConfig.from_pretrained(original_model_name)

136 + gen_config.save_pretrained(output_path)

137 + print("Copied generation config")

138 + except Exception as e:

139 + print(f"Could not copy generation config: {e}")

140 +

141 + print(f"Successfully created reduced Maverick model at {output_path}")

142 + return str(output_path)

143 +

144 + except Exception as e:

145 + print(f"Error creating reduced model: {e}")

146 + # Clean up on failure

147 + if output_path.exists():

148 + shutil.rmtree(output_path)

149 + raise

150 +

151 +

152 +def create_reduced_config(original_config: Any, text_layers: int,

153 + num_experts: int,

154 + vision_layers: int) -> dict[str, Any]:

155 + """Create a reduced configuration based on the original."""

156 +

157 + # Convert config to dictionary

158 + config_dict = original_config.to_dict()

159 +

160 + # Reduce text layers

161 + if "text_config" in config_dict:

162 + original_text_layers = config_dict["text_config"]["num_hidden_layers"]

163 + config_dict["text_config"]["num_hidden_layers"] = text_layers

164 + print(

165 + f"Reduced text layers from {original_text_layers} to {text_layers}"

166 + )

167 +

168 + original_num_experts = config_dict["text_config"]["num_local_experts"]

169 + config_dict["text_config"]["num_local_experts"] = num_experts

170 + print(

171 + f"Reduced num experts from {original_num_experts} to {num_experts}"

172 + )

173 +

174 + hidden_dim_divisor = 4

175 +

176 + original_hidden_size = config_dict["text_config"]["hidden_size"]

177 + new_hidden_size = original_hidden_size // hidden_dim_divisor

178 + config_dict["text_config"]["hidden_size"] = new_hidden_size

179 + print(f"Reduced hidden size from {original_hidden_size} to "

180 + f"{new_hidden_size}")

181 +

182 + original_head_dim = config_dict["text_config"]["head_dim"]

183 + new_head_dim = original_head_dim // hidden_dim_divisor

184 + config_dict["text_config"]["head_dim"] = new_head_dim

185 + print(f"Reduced head dim from {original_head_dim} to {new_head_dim}")

186 +

187 + # Reduce vision layers

188 + if "vision_config" in config_dict:

189 + original_vision_layers = config_dict["vision_config"][

190 + "num_hidden_layers"]

191 + config_dict["vision_config"]["num_hidden_layers"] = vision_layers

192 + print(f"Reduced vision layers from {original_vision_layers} "

193 + f"to {vision_layers}")

194 +

195 + # Update model name to indicate it's a reduced version

196 + config_dict["_name_or_path"] = (

197 + f"reduced_maverick_{text_layers}t_{vision_layers}v")

198 +

199 + return config_dict

200 +

201 +

202 +def copy_tokenizer_files(original_model_name: str, output_path: Path) -> None:

203 + """Copy tokenizer files from the original model."""

204 +

205 + try:

206 + tokenizer = AutoTokenizer.from_pretrained(original_model_name,

207 + trust_remote_code=True)

208 + tokenizer.save_pretrained(output_path)

209 + print("Tokenizer files copied successfully")

210 + except Exception as e:

211 + print(f"Warning: Could not copy tokenizer files: {e}")

212 +

213 +

214 +def create_preprocessor_config(original_config: Any,

215 + output_path: Path) -> None:

216 + """Create preprocessor_config.json for multimodal model."""

217 +

218 + # Try to load the original preprocessor config

219 + try:

220 + processor = AutoProcessor.from_pretrained(

221 + original_config._name_or_path

222 + or "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",

223 + trust_remote_code=True,

224 + )

225 + processor.save_pretrained(output_path)

226 + print("Copied original preprocessor config")

227 + return

228 + except Exception as e:

229 + print(f"Could not copy original preprocessor config: {e}")

230 + raise

231 +

232 +

233 +def create_reduced_safetensors(original_config: Any, reduced_config: dict[str,

234 + Any],

235 + output_path: Path) -> None:

236 + """Create safetensors files with weights for the reduced model."""

237 +

238 + print("Generating synthetic weights for reduced model...")

239 +

240 + text_config = reduced_config["text_config"]

241 + vision_config = reduced_config["vision_config"]

242 +

243 + weights = {}

244 +

245 + print("Creating text model weights...")

246 + weights.update(create_text_model_weights(text_config))

247 +

248 + print("Creating vision model weights...")

249 + weights.update(create_vision_model_weights(vision_config))

250 +

251 + print("Creating shared model weights...")

252 + weights.update(create_shared_weights(text_config, vision_config))

253 +

254 + print("Saving weights to safetensors files...")

255 + save_weights_to_safetensors(weights, output_path)

256 +

257 +

258 +def create_text_model_weights(

259 + text_config: dict[str, Any]) -> dict[str, torch.Tensor]:

260 + """Create synthetic weights for the text model with MoE structure."""

261 +

262 + weights = {}

263 +

264 + vocab_size = text_config["vocab_size"]

265 + hidden_size = text_config["hidden_size"]

266 + intermediate_size = text_config["intermediate_size"]

267 + intermediate_size_mlp = text_config["intermediate_size_mlp"]

268 + num_layers = text_config["num_hidden_layers"]

269 + num_attention_heads = text_config["num_attention_heads"]

270 + num_key_value_heads = text_config.get("num_key_value_heads",

271 + num_attention_heads)

272 +

273 + # MoE specific parameters

274 + num_experts = text_config.get("num_local_experts")

275 + assert (num_experts

276 + is not None), "num_local_experts must be specified for MoE"

277 +

278 + head_dim = hidden_size // num_attention_heads

279 +

280 + # Embedding layers

281 + weights["language_model.model.embed_tokens.weight"] = torch.randn(

282 + vocab_size, hidden_size, dtype=torch.float16)

283 +

284 + # Transformer layers

285 + for layer_idx in range(num_layers):

286 + layer_prefix = f"language_model.model.layers.{layer_idx}"

287 + print(f"Creating weights for layer {layer_prefix}...")

288 +

289 + # Self-attention weights (separate q, k, v projections)

290 + weights[f"{layer_prefix}.self_attn.q_proj.weight"] = torch.randn(

291 + hidden_size, num_attention_heads * head_dim, dtype=torch.bfloat16)

292 + weights[f"{layer_prefix}.self_attn.k_proj.weight"] = torch.randn(

293 + hidden_size, num_key_value_heads * head_dim, dtype=torch.bfloat16)

294 + weights[f"{layer_prefix}.self_attn.v_proj.weight"] = torch.randn(

295 + num_key_value_heads * head_dim, hidden_size, dtype=torch.bfloat16)

296 + weights[f"{layer_prefix}.self_attn.o_proj.weight"] = torch.randn(

297 + hidden_size, num_attention_heads * head_dim, dtype=torch.bfloat16)

298 + print("Self-attention weights created.")

299 +

300 + # Feed-forward weights - MoE pattern based on interleave_moe_layer_step

301 + # For interleave_moe_layer_step=2: layers 1,3,5,... are MoE, layers

302 + # 0,2,4,... are dense

303 + interleave_step = text_config.get("interleave_moe_layer_step", 1)

304 + is_moe_layer = (interleave_step > 0

305 + and (layer_idx + 1) % interleave_step == 0)

306 +

307 + if is_moe_layer:

308 + # MoE layer structure

309 + # 1. Router weights

310 + weights[

311 + f"{layer_prefix}.feed_forward.router.weight"] = torch.randn(

312 + num_experts, hidden_size, dtype=torch.float16)

313 +

314 + # 2. Individual expert weights (not fused)

315 + for expert_idx in range(num_experts):

316 + expert_prefix = (

317 + f"{layer_prefix}.feed_forward.experts.{expert_idx}")

318 +

319 + weights[f"{expert_prefix}.gate_proj.weight"] = torch.randn(

320 + intermediate_size, hidden_size, dtype=torch.bfloat16)

321 + weights[f"{expert_prefix}.up_proj.weight"] = torch.randn(

322 + intermediate_size, hidden_size, dtype=torch.bfloat16)

323 + weights[f"{expert_prefix}.down_proj.weight"] = torch.randn(

324 + hidden_size, intermediate_size, dtype=torch.bfloat16)

325 +

326 + # Expert weight scales (FP8 quantization)

327 + weights[

328 + f"{expert_prefix}.gate_proj.weight_scale"] = torch.ones(

329 + intermediate_size, 1, dtype=torch.bfloat16)

330 + weights[f"{expert_prefix}.up_proj.weight_scale"] = torch.ones(

331 + intermediate_size, 1, dtype=torch.bfloat16)

332 + weights[

333 + f"{expert_prefix}.down_proj.weight_scale"] = torch.ones(

334 + hidden_size, 1, dtype=torch.bfloat16)

335 +

336 + # 3. Shared expert weights

337 + shared_expert_prefix = f"{layer_prefix}.feed_forward.shared_expert"

338 + weights[f"{shared_expert_prefix}.gate_proj.weight"] = torch.randn(

339 + intermediate_size, hidden_size, dtype=torch.bfloat16)

340 + weights[f"{shared_expert_prefix}.up_proj.weight"] = torch.randn(

341 + intermediate_size, hidden_size, dtype=torch.bfloat16)

342 + weights[f"{shared_expert_prefix}.down_proj.weight"] = torch.randn(

343 + hidden_size, intermediate_size, dtype=torch.bfloat16)

344 + print(f"MoE feed-forward weights created for layer {layer_idx}.")

345 + else:

346 + # Dense layer structure

347 + weights[f"{layer_prefix}.feed_forward.gate_proj.weight"] = (

348 + torch.randn(intermediate_size_mlp,

349 + hidden_size,

350 + dtype=torch.bfloat16))

351 + weights[f"{layer_prefix}.feed_forward.up_proj.weight"] = (

352 + torch.randn(intermediate_size_mlp,

353 + hidden_size,

354 + dtype=torch.bfloat16))

355 + weights[f"{layer_prefix}.feed_forward.down_proj.weight"] = (

356 + torch.randn(hidden_size,

357 + intermediate_size_mlp,

358 + dtype=torch.bfloat16))

359 + print(f"Dense feed-forward weights created for layer {layer_idx}.")

360 +

361 + # Layer norms

362 + weights[f"{layer_prefix}.input_layernorm.weight"] = torch.ones(

363 + hidden_size, dtype=torch.bfloat16)

364 + weights[

365 + f"{layer_prefix}.post_attention_layernorm.weight"] = torch.ones(

366 + hidden_size, dtype=torch.bfloat16)

367 + print("Layer norms created.")

368 +

369 + # Final layer norm and output projection

370 + weights["language_model.model.norm.weight"] = torch.ones(

371 + hidden_size, dtype=torch.bfloat16)

372 + weights["language_model.lm_head.weight"] = torch.randn(

373 + vocab_size, hidden_size, dtype=torch.bfloat16)

374 +

375 + return weights

376 +

377 +

378 +def create_vision_model_weights(

379 + vision_config: dict[str, Any]) -> dict[str, torch.Tensor]:

380 + """Create synthetic weights for the vision model."""

381 +

382 + weights = {}

383 +

384 + hidden_size = vision_config["hidden_size"]

385 + intermediate_size = vision_config["intermediate_size"]

386 + num_layers = vision_config["num_hidden_layers"]

387 +

388 + # Vision transformer layers

389 + for layer_idx in range(num_layers):

390 + layer_prefix = f"vision_model.model.layers.{layer_idx}"

391 +

392 + weights[f"{layer_prefix}.self_attn.q_proj.weight"] = torch.randn(

393 + hidden_size, hidden_size, dtype=torch.bfloat16)

394 + weights[f"{layer_prefix}.self_attn.q_proj.bias"] = torch.zeros(

395 + hidden_size, dtype=torch.bfloat16)

396 + weights[f"{layer_prefix}.self_attn.k_proj.weight"] = torch.randn(

397 + hidden_size, hidden_size, dtype=torch.bfloat16)

398 + weights[f"{layer_prefix}.self_attn.k_proj.bias"] = torch.zeros(

399 + hidden_size, dtype=torch.bfloat16)

400 + weights[f"{layer_prefix}.self_attn.v_proj.weight"] = torch.randn(

401 + hidden_size, hidden_size, dtype=torch.bfloat16)

402 + weights[f"{layer_prefix}.self_attn.v_proj.bias"] = torch.zeros(

403 + hidden_size, dtype=torch.bfloat16)

404 + weights[f"{layer_prefix}.self_attn.o_proj.weight"] = torch.randn(

405 + hidden_size, hidden_size, dtype=torch.bfloat16)

406 + weights[f"{layer_prefix}.self_attn.o_proj.bias"] = torch.zeros(

407 + hidden_size, dtype=torch.bfloat16)

408 +

409 + weights[f"{layer_prefix}.mlp.fc1.weight"] = torch.randn(

410 + intermediate_size, hidden_size, dtype=torch.bfloat16)

411 + weights[f"{layer_prefix}.mlp.fc1.bias"] = torch.zeros(

412 + intermediate_size, dtype=torch.bfloat16)

413 + weights[f"{layer_prefix}.mlp.fc2.weight"] = torch.randn(

414 + hidden_size, intermediate_size, dtype=torch.bfloat16)

415 + weights[f"{layer_prefix}.mlp.fc2.bias"] = torch.zeros(

416 + hidden_size, dtype=torch.bfloat16)

417 +

418 + weights[f"{layer_prefix}.input_layernorm.weight"] = torch.ones(

419 + hidden_size, dtype=torch.bfloat16)

420 + weights[f"{layer_prefix}.input_layernorm.bias"] = torch.zeros(

421 + hidden_size, dtype=torch.bfloat16)

422 + weights[

423 + f"{layer_prefix}.post_attention_layernorm.weight"] = torch.ones(

424 + hidden_size, dtype=torch.bfloat16)

425 + weights[f"{layer_prefix}.post_attention_layernorm.bias"] = torch.zeros(

426 + hidden_size, dtype=torch.bfloat16)

427 +

428 + return weights

429 +

430 +

431 +def create_shared_weights(

432 + text_config: dict[str, Any],

433 + vision_config: dict[str, Any]) -> dict[str, torch.Tensor]:

434 + """Create weights for shared components (vision-language connector)"""

435 +

436 + weights = {}

437 +

438 + text_hidden_size = text_config["hidden_size"]

439 + projector_input_dim = vision_config["projector_input_dim"]

440 +

441 + # Vision-language connector (projects vision features to text space)

442 + weights["multi_modal_projector.linear_1.weight"] = torch.randn(

443 + text_hidden_size, projector_input_dim, dtype=torch.bfloat16)

444 +

445 + return weights

446 +

447 +

448 +def save_weights_to_safetensors(weights: dict[str, torch.Tensor],

449 + output_path: Path) -> None:

450 + """Save weights to safetensors files and create index."""

451 +

452 + # Determine how to shard the weights

453 + max_shard_size = 5 * 1024 * 1024 * 1024 # 5GB per shard

454 +

455 + # Calculate sizes and create shards

456 + shards = []

457 + current_shard: dict[str, torch.Tensor] = {}

458 + current_size = 0

459 +

460 + for name, tensor in weights.items():

461 + tensor_size = tensor.numel() * tensor.element_size()

462 +

463 + if current_size + tensor_size > max_shard_size and current_shard:

464 + shards.append(current_shard)

465 + current_shard = {}

466 + current_size = 0

467 +

468 + current_shard[name] = tensor

469 + current_size += tensor_size

470 +

471 + if current_shard:

472 + shards.append(current_shard)

473 +

474 + # Save shards and create index

475 + weight_map = {}

476 +

477 + if len(shards) == 1:

478 + # Single file

479 + filename = "model.safetensors"

480 + save_file(shards[0], output_path / filename)

481 + weight_map = {name: filename for name in shards[0]}

482 + print(f"Saved weights to single file: {filename}")

483 + else:

484 + # Multiple shards

485 + for i, shard in enumerate(shards):

486 + filename = f"model-{i+1:05d}-of-{len(shards):05d}.safetensors"

487 + save_file(shard, output_path / filename)

488 + for name in shard:

489 + weight_map[name] = filename

490 + print(f"Saved shard {i+1}/{len(shards)}: {filename}")

491 +

492 + # Create index file

493 + index_data = {

494 + "metadata": {

495 + "total_size":

496 + sum(tensor.numel() * tensor.element_size()

497 + for tensor in weights.values())

498 + },

499 + "weight_map": weight_map,

500 + }

501 +

502 + index_path = output_path / "model.safetensors.index.json"

503 + with open(index_path, "w") as f:

504 + json.dump(index_data, f, indent=2)

505 +

506 + print(f"Created index file: {index_path}")

507 + print(f"Total model size: "

508 + f"{index_data['metadata']['total_size'] / (1024**3):.2f} GB")

509 +

510 +

511 +def run_reduced_model(model_path: str,

512 + should_profile: bool = False,

513 + **kwargs) -> None:

514 + """Test the created reduced model with vLLM."""

515 +

516 + print(f"\nTesting reduced model at {model_path}...")

517 +

518 + llm = LLM(

519 + model=model_path,

520 + trust_remote_code=True,

521 + max_model_len=512, # Small context for testing

522 + gpu_memory_utilization=0.3, # Conservative memory usage

523 + **kwargs,

524 + )

525 +

526 + sampling_params = SamplingParams(temperature=0.8,

527 + top_p=0.95,

528 + max_tokens=50)

529 +

530 + if should_profile:

531 + llm.start_profile()

532 + outputs = llm.generate(PROMPTS, sampling_params)

533 + if should_profile:

534 + llm.stop_profile()

535 +

536 + print("Test generation successful!")

537 + for output in outputs:

538 + print(f"Prompt: {output.prompt}")

539 + print(f"Output: "

540 + f"{output.outputs[0].text}")

541 + print("-" * 40)

542 +

543 +

544 +@pytest.mark.parametrize(

545 + "original_model_name,text_layers,num_experts,vision_layers,",

546 + [("meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", 4, 4, 2)])

547 +@pytest.mark.parametrize("enforce_eager", [True, False])

548 +@pytest.mark.parametrize("tp,ep", [(2, True)])

549 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")

550 +def test_dummy_maverick(

551 + original_model_name: str,

552 + text_layers: int,

553 + num_experts: int,

554 + vision_layers: int,

555 + enforce_eager: bool,

556 + tp: int,

557 + ep: bool,

558 + output_dir: str = "/tmp/reduced_maverick",

559 + force_recreate: bool = True,

560 + profile: bool = False,

561 +) -> None:

562 + model_path = create_reduced_maverick_model(

563 + original_model_name=original_model_name,

564 + output_dir=output_dir,

565 + text_layers=text_layers,

566 + num_experts=num_experts,

567 + vision_layers=vision_layers,

568 + force_recreate=force_recreate,

569 + )

570 +

571 + print(f"\nReduced model created successfully at: {model_path}")

572 +

573 + run_reduced_model(model_path=model_path,

574 + should_profile=profile,

575 + enforce_eager=enforce_eager,

576 + tensor_parallel_size=tp,

577 + enable_expert_parallel=ep)

578 +

579 +

580 +def main():

581 + """Main function to create and test the reduced model."""

582 +

583 + import argparse

584 +

585 + parser = argparse.ArgumentParser(

586 + description="Create a reduced-layer Maverick model")

587 + parser.add_argument(

588 + "--output-dir",

589 + default="/tmp/reduced_maverick",

590 + help="Output directory for the reduced model",

591 + )

592 + parser.add_argument(

593 + "--text-layers",

594 + type=int,

595 + default=4,

596 + help="Number of text transformer layers",

597 + )

598 + parser.add_argument("--num-experts",

599 + type=int,

600 + default=4,

601 + help="Number of experts")

602 + parser.add_argument(

603 + "--vision-layers",

604 + type=int,

605 + default=2,

606 + help="Number of vision transformer layers",

607 + )

608 + parser.add_argument(

609 + "--force-recreate",

610 + action="store_true",

611 + help="Force recreation if output directory exists",

612 + )

613 + parser.add_argument("--test",

614 + action="store_true",

615 + help="Test the created model with vLLM")

616 + parser.add_argument("--profile",

617 + action="store_true",

618 + help="Profile the created model with vLLM")

619 + parser.add_argument(

620 + "--test-original",

621 + action="store_true",

622 + help="Test the original model with vLLM",

623 + )

624 + parser.add_argument(

625 + "--original-model",

626 + default="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",

627 + help="Original model name to base the reduction on",

628 + )

629 +

630 + args = parser.parse_args()

631 +

632 + if args.test:

633 + test_dummy_maverick(original_model_name=args.original_model,

634 + output_dir=args.output_dir,

635 + text_layers=args.text_layers,

636 + num_experts=args.num_experts,

637 + vision_layers=args.vision_layers,

638 + force_recreate=args.force_recreate,

639 + tp=2,

640 + ep=True,

641 + enforce_eager=True,

642 + profile=args.profile)

643 +

644 + if args.test_original:

645 + run_maverick_serving(args.original_model)

646 +

647 +

648 +if __name__ == "__main__":

649 + exit(main())

[Misc] Add dummy maverick test (#21199)