# 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 from transformers import (AutoConfig, AutoProcessor, AutoTokenizer, GenerationConfig) from vllm import LLM, SamplingParams from vllm.v1.executor.abstract import Executor from vllm.v1.kv_cache_interface import (ChunkedLocalAttentionSpec, FullAttentionSpec) from ....utils import multi_gpu_test # 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 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 def create_reduced_maverick_model( original_model_name: str = "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", 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 " f"{vision_layers} vision layers...") # Create output directory output_path = Path(output_dir) if output_path.exists(): if force_recreate: shutil.rmtree(output_path) else: print(f"Output directory {output_dir} already exists. " "Use --force-recreate to overwrite.") return str(output_path) output_path.mkdir(parents=True, exist_ok=True) try: print("Loading original model configuration...") original_config = AutoConfig.from_pretrained(original_model_name, trust_remote_code=True) print("Creating reduced configuration...") reduced_config = create_reduced_config(original_config, text_layers, num_experts, vision_layers) 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...") create_reduced_safetensors(original_config, reduced_config, output_path) 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 def create_reduced_config(original_config: Any, text_layers: int, num_experts: int, vision_layers: int) -> dict[str, Any]: """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 print( f"Reduced text layers from {original_text_layers} to {text_layers}" ) original_num_experts = config_dict["text_config"]["num_local_experts"] config_dict["text_config"]["num_local_experts"] = num_experts print( f"Reduced num experts from {original_num_experts} to {num_experts}" ) 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 print(f"Reduced hidden size from {original_hidden_size} to " f"{new_hidden_size}") 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: original_vision_layers = config_dict["vision_config"][ "num_hidden_layers"] config_dict["vision_config"]["num_hidden_layers"] = vision_layers print(f"Reduced vision layers from {original_vision_layers} " f"to {vision_layers}") # Update model name to indicate it's a reduced version config_dict["_name_or_path"] = ( f"reduced_maverick_{text_layers}t_{vision_layers}v") return config_dict def copy_tokenizer_files(original_model_name: str, output_path: Path) -> None: """Copy tokenizer files from the original model.""" try: tokenizer = AutoTokenizer.from_pretrained(original_model_name, trust_remote_code=True) tokenizer.save_pretrained(output_path) print("Tokenizer files copied successfully") except Exception as e: print(f"Warning: Could not copy tokenizer files: {e}") def create_preprocessor_config(original_config: Any, output_path: Path) -> None: """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 def create_reduced_safetensors(original_config: Any, reduced_config: dict[str, Any], output_path: Path) -> None: """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) def create_text_model_weights( text_config: dict[str, Any]) -> dict[str, torch.Tensor]: """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"] num_key_value_heads = text_config.get("num_key_value_heads", num_attention_heads) # MoE specific parameters num_experts = text_config.get("num_local_experts") assert (num_experts is not None), "num_local_experts must be specified for MoE" head_dim = hidden_size // num_attention_heads # Embedding layers weights["language_model.model.embed_tokens.weight"] = torch.randn( vocab_size, hidden_size, dtype=torch.float16) # 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( hidden_size, num_attention_heads * head_dim, dtype=torch.bfloat16) weights[f"{layer_prefix}.self_attn.k_proj.weight"] = torch.randn( hidden_size, num_key_value_heads * head_dim, dtype=torch.bfloat16) weights[f"{layer_prefix}.self_attn.v_proj.weight"] = torch.randn( num_key_value_heads * head_dim, hidden_size, dtype=torch.bfloat16) weights[f"{layer_prefix}.self_attn.o_proj.weight"] = torch.randn( hidden_size, num_attention_heads * head_dim, dtype=torch.bfloat16) 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) is_moe_layer = (interleave_step > 0 and (layer_idx + 1) % interleave_step == 0) if is_moe_layer: # MoE layer structure # 1. Router weights weights[ f"{layer_prefix}.feed_forward.router.weight"] = torch.randn( num_experts, hidden_size, dtype=torch.float16) # 2. Individual expert weights (not fused) for expert_idx in range(num_experts): expert_prefix = ( f"{layer_prefix}.feed_forward.experts.{expert_idx}") weights[f"{expert_prefix}.gate_proj.weight"] = torch.randn( intermediate_size, hidden_size, dtype=torch.bfloat16) weights[f"{expert_prefix}.up_proj.weight"] = torch.randn( intermediate_size, hidden_size, dtype=torch.bfloat16) weights[f"{expert_prefix}.down_proj.weight"] = torch.randn( hidden_size, intermediate_size, dtype=torch.bfloat16) # Expert weight scales (FP8 quantization) weights[ f"{expert_prefix}.gate_proj.weight_scale"] = torch.ones( intermediate_size, 1, dtype=torch.bfloat16) weights[f"{expert_prefix}.up_proj.weight_scale"] = torch.ones( intermediate_size, 1, dtype=torch.bfloat16) weights[ f"{expert_prefix}.down_proj.weight_scale"] = torch.ones( hidden_size, 1, dtype=torch.bfloat16) # 3. Shared expert weights shared_expert_prefix = f"{layer_prefix}.feed_forward.shared_expert" weights[f"{shared_expert_prefix}.gate_proj.weight"] = torch.randn( intermediate_size, hidden_size, dtype=torch.bfloat16) weights[f"{shared_expert_prefix}.up_proj.weight"] = torch.randn( intermediate_size, hidden_size, dtype=torch.bfloat16) weights[f"{shared_expert_prefix}.down_proj.weight"] = torch.randn( hidden_size, intermediate_size, dtype=torch.bfloat16) print(f"MoE feed-forward weights created for layer {layer_idx}.") else: # Dense layer structure 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)) print(f"Dense feed-forward weights created for layer {layer_idx}.") # Layer norms weights[f"{layer_prefix}.input_layernorm.weight"] = torch.ones( hidden_size, dtype=torch.bfloat16) weights[ f"{layer_prefix}.post_attention_layernorm.weight"] = torch.ones( hidden_size, dtype=torch.bfloat16) print("Layer norms created.") # Final layer norm and output projection weights["language_model.model.norm.weight"] = torch.ones( hidden_size, dtype=torch.bfloat16) weights["language_model.lm_head.weight"] = torch.randn( vocab_size, hidden_size, dtype=torch.bfloat16) return weights def create_vision_model_weights( vision_config: dict[str, Any]) -> dict[str, torch.Tensor]: """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( hidden_size, hidden_size, dtype=torch.bfloat16) weights[f"{layer_prefix}.self_attn.q_proj.bias"] = torch.zeros( hidden_size, dtype=torch.bfloat16) weights[f"{layer_prefix}.self_attn.k_proj.weight"] = torch.randn( hidden_size, hidden_size, dtype=torch.bfloat16) weights[f"{layer_prefix}.self_attn.k_proj.bias"] = torch.zeros( hidden_size, dtype=torch.bfloat16) weights[f"{layer_prefix}.self_attn.v_proj.weight"] = torch.randn( hidden_size, hidden_size, dtype=torch.bfloat16) weights[f"{layer_prefix}.self_attn.v_proj.bias"] = torch.zeros( hidden_size, dtype=torch.bfloat16) weights[f"{layer_prefix}.self_attn.o_proj.weight"] = torch.randn( hidden_size, hidden_size, dtype=torch.bfloat16) weights[f"{layer_prefix}.self_attn.o_proj.bias"] = torch.zeros( hidden_size, dtype=torch.bfloat16) weights[f"{layer_prefix}.mlp.fc1.weight"] = torch.randn( intermediate_size, hidden_size, dtype=torch.bfloat16) weights[f"{layer_prefix}.mlp.fc1.bias"] = torch.zeros( intermediate_size, dtype=torch.bfloat16) weights[f"{layer_prefix}.mlp.fc2.weight"] = torch.randn( hidden_size, intermediate_size, dtype=torch.bfloat16) weights[f"{layer_prefix}.mlp.fc2.bias"] = torch.zeros( hidden_size, dtype=torch.bfloat16) weights[f"{layer_prefix}.input_layernorm.weight"] = torch.ones( hidden_size, dtype=torch.bfloat16) weights[f"{layer_prefix}.input_layernorm.bias"] = torch.zeros( hidden_size, dtype=torch.bfloat16) weights[ f"{layer_prefix}.post_attention_layernorm.weight"] = torch.ones( hidden_size, dtype=torch.bfloat16) weights[f"{layer_prefix}.post_attention_layernorm.bias"] = torch.zeros( hidden_size, dtype=torch.bfloat16) return weights def create_shared_weights( text_config: dict[str, Any], vision_config: dict[str, Any]) -> dict[str, torch.Tensor]: """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( text_hidden_size, projector_input_dim, dtype=torch.bfloat16) return weights def save_weights_to_safetensors(weights: dict[str, torch.Tensor], output_path: Path) -> None: """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): filename = f"model-{i+1:05d}-of-{len(shards):05d}.safetensors" save_file(shard, output_path / filename) for name in shard: weight_map[name] = filename print(f"Saved shard {i+1}/{len(shards)}: {filename}") # Create index file index_data = { "metadata": { "total_size": sum(tensor.numel() * tensor.element_size() for tensor in weights.values()) }, "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}") print(f"Total model size: " f"{index_data['metadata']['total_size'] / (1024**3):.2f} GB") 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) kv_cache_specs_per_rank = llm.llm_engine.model_executor.get_kv_cache_specs( ) 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( kv_cache_specs[ f"language_model.model.layers.{i}.self_attn.attn"], expected_spec) def run_reduced_model(llm: LLM, should_profile: bool = False) -> None: """Test the created reduced model with vLLM.""" sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=50) 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}") print(f"Output: " f"{output.outputs[0].text}") print("-" * 40) @multi_gpu_test(num_gpus=2) @pytest.mark.parametrize( "original_model_name,text_layers,num_experts,vision_layers,", [("meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", 4, 4, 2)]) @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( monkeypatch, 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: # Disable multiprocessing allows us to access model executor from LLM engine monkeypatch.setenv("VLLM_USE_V1", "1") monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") 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}") 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) def main(): """Main function to create and test the reduced model.""" import argparse parser = argparse.ArgumentParser( description="Create a reduced-layer Maverick model") 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", ) parser.add_argument("--num-experts", type=int, default=4, help="Number of experts") 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", ) 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") 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: 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) if args.test_original: run_maverick_serving(args.original_model) if __name__ == "__main__": exit(main())