test_maverick.py 25.1 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Create a reduced-layer version of the Maverick model for testing purposes.

This script creates a new model with fewer layers by:
1. Loading the original Maverick model configuration
2. Creating a reduced configuration
3. Generating compatible safetensors files with appropriate weights
4. Creating the necessary index files for vLLM compatibility
"""

import json
import shutil
from pathlib import Path
from typing import Any

import pytest
import torch
from safetensors.torch import save_file
21
from transformers import AutoConfig, AutoProcessor, AutoTokenizer, GenerationConfig
22
23

from vllm import LLM, SamplingParams
24
from vllm.v1.executor.abstract import Executor
25
from vllm.v1.kv_cache_interface import ChunkedLocalAttentionSpec, FullAttentionSpec
26

27
28
from ....utils import multi_gpu_test

29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
# Sample prompts for testing
PROMPTS: list[str] = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]


def run_maverick_serving(model: str):
    """Test Llama-4-Maverick model with vLLM LLM class using CLI equivalent
    options with reduced layers.
    """

    try:
        sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

        llm = LLM(
            model=model,
            max_model_len=2048,
            enforce_eager=True,
            tensor_parallel_size=8,
            enable_expert_parallel=True,
            trust_remote_code=True,
            gpu_memory_utilization=0.4,
            kv_cache_dtype="fp8",
        )

        outputs = llm.generate(PROMPTS, sampling_params)

        # Print the outputs
        print("\nGenerated Outputs:\n" + "-" * 60)
        for output in outputs:
            prompt = output.prompt
            generated_text = output.outputs[0].text
            print(f"Prompt:    {prompt!r}")
            print(f"Output:    {generated_text!r}")
            print("-" * 60)

    except Exception as e:
        print(f"Error initializing or running model: {e}")
        raise


73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def get_rope_layers_config(model_path: str) -> list[int]:
    """
    Get the interleaved RoPE configuration from HuggingFace config

    Args:
        model_path: Path to the local directory containing the reduced
            Maverick model checkpoint

    Returns:
        List of 0 or 1 indicating whether each layer uses RoPE and local attn
        0 indicates that RoPE is not used while 1 indicates that RoPE is used.
    """
    config_path = Path(model_path) / "config.json"
    model_config = json.loads(config_path.read_text())
    text_config = model_config["text_config"]
    no_rope_layers = text_config["no_rope_layers"]
    print(f"Found no_rope_layers: {no_rope_layers}")
    return no_rope_layers


93
def create_reduced_maverick_model(
94
    original_model_name: str = "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
    output_dir: str = "/tmp/reduced_maverick",
    text_layers: int = 4,
    num_experts: int = 4,
    vision_layers: int = 2,
    force_recreate: bool = False,
) -> str:
    """
    Create a reduced-layer version of the Maverick model.

    Args:
        original_model_name: Name of the original Maverick model
        output_dir: Directory to save the reduced model
        text_layers: Number of text transformer layers
        num_experts: Number of experts per layer
        vision_layers: Number of vision transformer layers
        force_recreate: Whether to recreate if output_dir already exists

    Returns:
        Path to the created reduced model directory
    """

    print(
        f"Creating reduced Maverick model with {text_layers} text layers and "
118
119
        f"{vision_layers} vision layers..."
    )
120
121
122
123
124
125
126

    # Create output directory
    output_path = Path(output_dir)
    if output_path.exists():
        if force_recreate:
            shutil.rmtree(output_path)
        else:
127
128
129
130
            print(
                f"Output directory {output_dir} already exists. "
                "Use --force-recreate to overwrite."
            )
131
132
133
134
135
136
            return str(output_path)

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

    try:
        print("Loading original model configuration...")
137
138
139
        original_config = AutoConfig.from_pretrained(
            original_model_name, trust_remote_code=True
        )
140
        print("Creating reduced configuration...")
141
142
143
        reduced_config = create_reduced_config(
            original_config, text_layers, num_experts, vision_layers
        )
144
145
146
147
148
149
150
151
152
153

        config_path = output_path / "config.json"
        with open(config_path, "w") as f:
            json.dump(reduced_config, f, indent=2)
        print(f"Saved reduced config to {config_path}")

        print("Copying tokenizer files...")
        copy_tokenizer_files(original_model_name, output_path)

        print("Creating reduced safetensors files...")
154
        create_reduced_safetensors(original_config, reduced_config, output_path)
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176

        print("Creating preprocessor config...")
        create_preprocessor_config(original_config, output_path)

        try:
            gen_config = GenerationConfig.from_pretrained(original_model_name)
            gen_config.save_pretrained(output_path)
            print("Copied generation config")
        except Exception as e:
            print(f"Could not copy generation config: {e}")

        print(f"Successfully created reduced Maverick model at {output_path}")
        return str(output_path)

    except Exception as e:
        print(f"Error creating reduced model: {e}")
        # Clean up on failure
        if output_path.exists():
            shutil.rmtree(output_path)
        raise


177
178
179
def create_reduced_config(
    original_config: Any, text_layers: int, num_experts: int, vision_layers: int
) -> dict[str, Any]:
180
181
182
183
184
185
186
187
188
    """Create a reduced configuration based on the original."""

    # Convert config to dictionary
    config_dict = original_config.to_dict()

    # Reduce text layers
    if "text_config" in config_dict:
        original_text_layers = config_dict["text_config"]["num_hidden_layers"]
        config_dict["text_config"]["num_hidden_layers"] = text_layers
189
190
        original_layer_types = config_dict["text_config"]["layer_types"]
        config_dict["text_config"]["layer_types"] = original_layer_types[:text_layers]
191
        print(f"Reduced text layers from {original_text_layers} to {text_layers}")
192
193
194

        original_num_experts = config_dict["text_config"]["num_local_experts"]
        config_dict["text_config"]["num_local_experts"] = num_experts
195
        print(f"Reduced num experts from {original_num_experts} to {num_experts}")
196
197
198
199
200
201

        hidden_dim_divisor = 4

        original_hidden_size = config_dict["text_config"]["hidden_size"]
        new_hidden_size = original_hidden_size // hidden_dim_divisor
        config_dict["text_config"]["hidden_size"] = new_hidden_size
202
        print(f"Reduced hidden size from {original_hidden_size} to {new_hidden_size}")
203
204
205
206
207
208
209
210

        original_head_dim = config_dict["text_config"]["head_dim"]
        new_head_dim = original_head_dim // hidden_dim_divisor
        config_dict["text_config"]["head_dim"] = new_head_dim
        print(f"Reduced head dim from {original_head_dim} to {new_head_dim}")

    # Reduce vision layers
    if "vision_config" in config_dict:
211
        original_vision_layers = config_dict["vision_config"]["num_hidden_layers"]
212
        config_dict["vision_config"]["num_hidden_layers"] = vision_layers
213
        print(f"Reduced vision layers from {original_vision_layers} to {vision_layers}")
214
215

    # Update model name to indicate it's a reduced version
216
    config_dict["_name_or_path"] = f"reduced_maverick_{text_layers}t_{vision_layers}v"
217
218
219
220
221
222
223
224

    return config_dict


def copy_tokenizer_files(original_model_name: str, output_path: Path) -> None:
    """Copy tokenizer files from the original model."""

    try:
225
226
227
        tokenizer = AutoTokenizer.from_pretrained(
            original_model_name, trust_remote_code=True
        )
228
229
230
231
232
233
        tokenizer.save_pretrained(output_path)
        print("Tokenizer files copied successfully")
    except Exception as e:
        print(f"Warning: Could not copy tokenizer files: {e}")


234
def create_preprocessor_config(original_config: Any, output_path: Path) -> None:
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
    """Create preprocessor_config.json for multimodal model."""

    # Try to load the original preprocessor config
    try:
        processor = AutoProcessor.from_pretrained(
            original_config._name_or_path
            or "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
            trust_remote_code=True,
        )
        processor.save_pretrained(output_path)
        print("Copied original preprocessor config")
        return
    except Exception as e:
        print(f"Could not copy original preprocessor config: {e}")
        raise


252
253
254
def create_reduced_safetensors(
    original_config: Any, reduced_config: dict[str, Any], output_path: Path
) -> None:
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
    """Create safetensors files with weights for the reduced model."""

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

    text_config = reduced_config["text_config"]
    vision_config = reduced_config["vision_config"]

    weights = {}

    print("Creating text model weights...")
    weights.update(create_text_model_weights(text_config))

    print("Creating vision model weights...")
    weights.update(create_vision_model_weights(vision_config))

    print("Creating shared model weights...")
    weights.update(create_shared_weights(text_config, vision_config))

    print("Saving weights to safetensors files...")
    save_weights_to_safetensors(weights, output_path)


277
def create_text_model_weights(text_config: dict[str, Any]) -> dict[str, torch.Tensor]:
278
279
280
281
282
283
284
285
286
287
    """Create synthetic weights for the text model with MoE structure."""

    weights = {}

    vocab_size = text_config["vocab_size"]
    hidden_size = text_config["hidden_size"]
    intermediate_size = text_config["intermediate_size"]
    intermediate_size_mlp = text_config["intermediate_size_mlp"]
    num_layers = text_config["num_hidden_layers"]
    num_attention_heads = text_config["num_attention_heads"]
288
    num_key_value_heads = text_config.get("num_key_value_heads", num_attention_heads)
289
290
291

    # MoE specific parameters
    num_experts = text_config.get("num_local_experts")
292
    assert num_experts is not None, "num_local_experts must be specified for MoE"
293
294
295
296
297

    head_dim = hidden_size // num_attention_heads

    # Embedding layers
    weights["language_model.model.embed_tokens.weight"] = torch.randn(
298
299
        vocab_size, hidden_size, dtype=torch.float16
    )
300
301
302
303
304
305
306
307

    # Transformer layers
    for layer_idx in range(num_layers):
        layer_prefix = f"language_model.model.layers.{layer_idx}"
        print(f"Creating weights for layer {layer_prefix}...")

        # Self-attention weights (separate q, k, v projections)
        weights[f"{layer_prefix}.self_attn.q_proj.weight"] = torch.randn(
308
309
            hidden_size, num_attention_heads * head_dim, dtype=torch.bfloat16
        )
310
        weights[f"{layer_prefix}.self_attn.k_proj.weight"] = torch.randn(
311
312
            hidden_size, num_key_value_heads * head_dim, dtype=torch.bfloat16
        )
313
        weights[f"{layer_prefix}.self_attn.v_proj.weight"] = torch.randn(
314
315
            num_key_value_heads * head_dim, hidden_size, dtype=torch.bfloat16
        )
316
        weights[f"{layer_prefix}.self_attn.o_proj.weight"] = torch.randn(
317
318
            hidden_size, num_attention_heads * head_dim, dtype=torch.bfloat16
        )
319
320
321
322
323
324
        print("Self-attention weights created.")

        # Feed-forward weights - MoE pattern based on interleave_moe_layer_step
        # For interleave_moe_layer_step=2: layers 1,3,5,... are MoE, layers
        # 0,2,4,... are dense
        interleave_step = text_config.get("interleave_moe_layer_step", 1)
325
        is_moe_layer = interleave_step > 0 and (layer_idx + 1) % interleave_step == 0
326
327
328
329

        if is_moe_layer:
            # MoE layer structure
            # 1. Router weights
330
331
332
            weights[f"{layer_prefix}.feed_forward.router.weight"] = torch.randn(
                num_experts, hidden_size, dtype=torch.float16
            )
333
334
335

            # 2. Individual expert weights (not fused)
            for expert_idx in range(num_experts):
336
                expert_prefix = f"{layer_prefix}.feed_forward.experts.{expert_idx}"
337
338

                weights[f"{expert_prefix}.gate_proj.weight"] = torch.randn(
339
340
                    intermediate_size, hidden_size, dtype=torch.bfloat16
                )
341
                weights[f"{expert_prefix}.up_proj.weight"] = torch.randn(
342
343
                    intermediate_size, hidden_size, dtype=torch.bfloat16
                )
344
                weights[f"{expert_prefix}.down_proj.weight"] = torch.randn(
345
346
                    hidden_size, intermediate_size, dtype=torch.bfloat16
                )
347
348

                # Expert weight scales (FP8 quantization)
349
350
351
                weights[f"{expert_prefix}.gate_proj.weight_scale"] = torch.ones(
                    intermediate_size, 1, dtype=torch.bfloat16
                )
352
                weights[f"{expert_prefix}.up_proj.weight_scale"] = torch.ones(
353
354
355
356
357
                    intermediate_size, 1, dtype=torch.bfloat16
                )
                weights[f"{expert_prefix}.down_proj.weight_scale"] = torch.ones(
                    hidden_size, 1, dtype=torch.bfloat16
                )
358
359
360
361

            # 3. Shared expert weights
            shared_expert_prefix = f"{layer_prefix}.feed_forward.shared_expert"
            weights[f"{shared_expert_prefix}.gate_proj.weight"] = torch.randn(
362
363
                intermediate_size, hidden_size, dtype=torch.bfloat16
            )
364
            weights[f"{shared_expert_prefix}.up_proj.weight"] = torch.randn(
365
366
                intermediate_size, hidden_size, dtype=torch.bfloat16
            )
367
            weights[f"{shared_expert_prefix}.down_proj.weight"] = torch.randn(
368
369
                hidden_size, intermediate_size, dtype=torch.bfloat16
            )
370
371
372
            print(f"MoE feed-forward weights created for layer {layer_idx}.")
        else:
            # Dense layer structure
373
374
375
376
377
378
379
380
381
            weights[f"{layer_prefix}.feed_forward.gate_proj.weight"] = torch.randn(
                intermediate_size_mlp, hidden_size, dtype=torch.bfloat16
            )
            weights[f"{layer_prefix}.feed_forward.up_proj.weight"] = torch.randn(
                intermediate_size_mlp, hidden_size, dtype=torch.bfloat16
            )
            weights[f"{layer_prefix}.feed_forward.down_proj.weight"] = torch.randn(
                hidden_size, intermediate_size_mlp, dtype=torch.bfloat16
            )
382
383
384
385
            print(f"Dense feed-forward weights created for layer {layer_idx}.")

        # Layer norms
        weights[f"{layer_prefix}.input_layernorm.weight"] = torch.ones(
386
387
388
389
390
            hidden_size, dtype=torch.bfloat16
        )
        weights[f"{layer_prefix}.post_attention_layernorm.weight"] = torch.ones(
            hidden_size, dtype=torch.bfloat16
        )
391
392
393
394
        print("Layer norms created.")

    # Final layer norm and output projection
    weights["language_model.model.norm.weight"] = torch.ones(
395
396
        hidden_size, dtype=torch.bfloat16
    )
397
    weights["language_model.lm_head.weight"] = torch.randn(
398
399
        vocab_size, hidden_size, dtype=torch.bfloat16
    )
400
401
402
403
404

    return weights


def create_vision_model_weights(
405
406
    vision_config: dict[str, Any],
) -> dict[str, torch.Tensor]:
407
408
409
410
411
412
413
414
415
416
417
418
419
    """Create synthetic weights for the vision model."""

    weights = {}

    hidden_size = vision_config["hidden_size"]
    intermediate_size = vision_config["intermediate_size"]
    num_layers = vision_config["num_hidden_layers"]

    # Vision transformer layers
    for layer_idx in range(num_layers):
        layer_prefix = f"vision_model.model.layers.{layer_idx}"

        weights[f"{layer_prefix}.self_attn.q_proj.weight"] = torch.randn(
420
421
            hidden_size, hidden_size, dtype=torch.bfloat16
        )
422
        weights[f"{layer_prefix}.self_attn.q_proj.bias"] = torch.zeros(
423
424
            hidden_size, dtype=torch.bfloat16
        )
425
        weights[f"{layer_prefix}.self_attn.k_proj.weight"] = torch.randn(
426
427
            hidden_size, hidden_size, dtype=torch.bfloat16
        )
428
        weights[f"{layer_prefix}.self_attn.k_proj.bias"] = torch.zeros(
429
430
            hidden_size, dtype=torch.bfloat16
        )
431
        weights[f"{layer_prefix}.self_attn.v_proj.weight"] = torch.randn(
432
433
            hidden_size, hidden_size, dtype=torch.bfloat16
        )
434
        weights[f"{layer_prefix}.self_attn.v_proj.bias"] = torch.zeros(
435
436
            hidden_size, dtype=torch.bfloat16
        )
437
        weights[f"{layer_prefix}.self_attn.o_proj.weight"] = torch.randn(
438
439
            hidden_size, hidden_size, dtype=torch.bfloat16
        )
440
        weights[f"{layer_prefix}.self_attn.o_proj.bias"] = torch.zeros(
441
442
            hidden_size, dtype=torch.bfloat16
        )
443
444

        weights[f"{layer_prefix}.mlp.fc1.weight"] = torch.randn(
445
446
            intermediate_size, hidden_size, dtype=torch.bfloat16
        )
447
        weights[f"{layer_prefix}.mlp.fc1.bias"] = torch.zeros(
448
449
            intermediate_size, dtype=torch.bfloat16
        )
450
        weights[f"{layer_prefix}.mlp.fc2.weight"] = torch.randn(
451
452
            hidden_size, intermediate_size, dtype=torch.bfloat16
        )
453
        weights[f"{layer_prefix}.mlp.fc2.bias"] = torch.zeros(
454
455
            hidden_size, dtype=torch.bfloat16
        )
456
457

        weights[f"{layer_prefix}.input_layernorm.weight"] = torch.ones(
458
459
            hidden_size, dtype=torch.bfloat16
        )
460
        weights[f"{layer_prefix}.input_layernorm.bias"] = torch.zeros(
461
462
463
464
465
            hidden_size, dtype=torch.bfloat16
        )
        weights[f"{layer_prefix}.post_attention_layernorm.weight"] = torch.ones(
            hidden_size, dtype=torch.bfloat16
        )
466
        weights[f"{layer_prefix}.post_attention_layernorm.bias"] = torch.zeros(
467
468
            hidden_size, dtype=torch.bfloat16
        )
469
470
471
472
473

    return weights


def create_shared_weights(
474
475
    text_config: dict[str, Any], vision_config: dict[str, Any]
) -> dict[str, torch.Tensor]:
476
477
478
479
480
481
482
483
484
    """Create weights for shared components (vision-language connector)"""

    weights = {}

    text_hidden_size = text_config["hidden_size"]
    projector_input_dim = vision_config["projector_input_dim"]

    # Vision-language connector (projects vision features to text space)
    weights["multi_modal_projector.linear_1.weight"] = torch.randn(
485
486
        text_hidden_size, projector_input_dim, dtype=torch.bfloat16
    )
487
488
489
490

    return weights


491
492
493
def save_weights_to_safetensors(
    weights: dict[str, torch.Tensor], output_path: Path
) -> None:
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
    """Save weights to safetensors files and create index."""

    # Determine how to shard the weights
    max_shard_size = 5 * 1024 * 1024 * 1024  # 5GB per shard

    # Calculate sizes and create shards
    shards = []
    current_shard: dict[str, torch.Tensor] = {}
    current_size = 0

    for name, tensor in weights.items():
        tensor_size = tensor.numel() * tensor.element_size()

        if current_size + tensor_size > max_shard_size and current_shard:
            shards.append(current_shard)
            current_shard = {}
            current_size = 0

        current_shard[name] = tensor
        current_size += tensor_size

    if current_shard:
        shards.append(current_shard)

    # Save shards and create index
    weight_map = {}

    if len(shards) == 1:
        # Single file
        filename = "model.safetensors"
        save_file(shards[0], output_path / filename)
        weight_map = {name: filename for name in shards[0]}
        print(f"Saved weights to single file: {filename}")
    else:
        # Multiple shards
        for i, shard in enumerate(shards):
530
            filename = f"model-{i + 1:05d}-of-{len(shards):05d}.safetensors"
531
532
533
            save_file(shard, output_path / filename)
            for name in shard:
                weight_map[name] = filename
534
            print(f"Saved shard {i + 1}/{len(shards)}: {filename}")
535
536
537
538

    # Create index file
    index_data = {
        "metadata": {
539
540
541
            "total_size": sum(
                tensor.numel() * tensor.element_size() for tensor in weights.values()
            )
542
543
544
545
546
547
548
549
550
        },
        "weight_map": weight_map,
    }

    index_path = output_path / "model.safetensors.index.json"
    with open(index_path, "w") as f:
        json.dump(index_data, f, indent=2)

    print(f"Created index file: {index_path}")
551
552
553
    print(
        f"Total model size: {index_data['metadata']['total_size'] / (1024**3):.2f} GB"
    )
554
555


556
557
558
559
560
561
562
563
def check_attention_spec_interleaved_rope(
    llm: LLM,
    num_attention_layers: int,
    num_ranks: int,
    rope_layers: list[int],
):
    """Check that the attention spec is correct."""
    assert isinstance(llm.llm_engine.model_executor, Executor)
564
    kv_cache_specs_per_rank = llm.llm_engine.model_executor.get_kv_cache_specs()
565
566
567
568
569
570
571
572
573
    for rank in range(num_ranks):
        kv_cache_specs = kv_cache_specs_per_rank[rank]
        assert len(kv_cache_specs.keys()) == num_attention_layers
        for i in range(num_attention_layers):
            if rope_layers[i] == 0:
                expected_spec = FullAttentionSpec
            else:
                expected_spec = ChunkedLocalAttentionSpec
            assert isinstance(
574
575
576
                kv_cache_specs[f"language_model.model.layers.{i}.self_attn.attn"],
                expected_spec,
            )
577
578
579
580


def run_reduced_model(llm: LLM, should_profile: bool = False) -> None:
    """Test the created reduced model with vLLM."""
581
    sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=50)
582
583
584
585
586
587
588
589
590
591

    if should_profile:
        llm.start_profile()
    outputs = llm.generate(PROMPTS, sampling_params)
    if should_profile:
        llm.stop_profile()

    print("Test generation successful!")
    for output in outputs:
        print(f"Prompt: {output.prompt}")
592
        print(f"Output: {output.outputs[0].text}")
593
594
595
        print("-" * 40)


596
@multi_gpu_test(num_gpus=2)
597
598
@pytest.mark.parametrize(
    "original_model_name,text_layers,num_experts,vision_layers,",
599
600
    [("meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", 4, 4, 2)],
)
601
602
603
604
@pytest.mark.parametrize("enforce_eager", [True, False])
@pytest.mark.parametrize("tp,ep", [(2, True)])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_dummy_maverick(
605
    monkeypatch,
606
607
608
609
610
611
612
613
614
615
616
    original_model_name: str,
    text_layers: int,
    num_experts: int,
    vision_layers: int,
    enforce_eager: bool,
    tp: int,
    ep: bool,
    output_dir: str = "/tmp/reduced_maverick",
    force_recreate: bool = True,
    profile: bool = False,
) -> None:
617
618
619
    # Disable multiprocessing allows us to access model executor from LLM engine
    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")

620
621
622
623
624
625
626
627
628
629
630
    model_path = create_reduced_maverick_model(
        original_model_name=original_model_name,
        output_dir=output_dir,
        text_layers=text_layers,
        num_experts=num_experts,
        vision_layers=vision_layers,
        force_recreate=force_recreate,
    )

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

631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
    rope_layers = get_rope_layers_config(model_path)

    llm = LLM(
        model=model_path,
        trust_remote_code=True,
        max_model_len=512,  # Small context for testing
        gpu_memory_utilization=0.3,  # Conservative memory usage
        enforce_eager=enforce_eager,
        tensor_parallel_size=tp,
        enable_expert_parallel=ep,
    )

    check_attention_spec_interleaved_rope(
        llm,
        text_layers,
        tp,
        rope_layers,
    )

    print(f"\nTesting reduced model at {model_path}...")
    run_reduced_model(llm=llm, should_profile=profile)
652
653
654
655
656
657
658
659


def main():
    """Main function to create and test the reduced model."""

    import argparse

    parser = argparse.ArgumentParser(
660
661
        description="Create a reduced-layer Maverick model"
    )
662
663
664
665
666
667
668
669
670
671
672
    parser.add_argument(
        "--output-dir",
        default="/tmp/reduced_maverick",
        help="Output directory for the reduced model",
    )
    parser.add_argument(
        "--text-layers",
        type=int,
        default=4,
        help="Number of text transformer layers",
    )
673
    parser.add_argument("--num-experts", type=int, default=4, help="Number of experts")
674
675
676
677
678
679
680
681
682
683
684
    parser.add_argument(
        "--vision-layers",
        type=int,
        default=2,
        help="Number of vision transformer layers",
    )
    parser.add_argument(
        "--force-recreate",
        action="store_true",
        help="Force recreation if output directory exists",
    )
685
686
687
688
689
690
    parser.add_argument(
        "--test", action="store_true", help="Test the created model with vLLM"
    )
    parser.add_argument(
        "--profile", action="store_true", help="Profile the created model with vLLM"
    )
691
692
693
694
695
696
697
698
699
700
701
702
703
704
    parser.add_argument(
        "--test-original",
        action="store_true",
        help="Test the original model with vLLM",
    )
    parser.add_argument(
        "--original-model",
        default="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8",
        help="Original model name to base the reduction on",
    )

    args = parser.parse_args()

    if args.test:
705
706
707
708
709
710
711
712
713
714
715
716
        test_dummy_maverick(
            original_model_name=args.original_model,
            output_dir=args.output_dir,
            text_layers=args.text_layers,
            num_experts=args.num_experts,
            vision_layers=args.vision_layers,
            force_recreate=args.force_recreate,
            tp=2,
            ep=True,
            enforce_eager=True,
            profile=args.profile,
        )
717
718
719
720
721
722
723

    if args.test_original:
        run_maverick_serving(args.original_model)


if __name__ == "__main__":
    exit(main())