| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Create a reduced-layer version of the Maverick model for testing purposes. |
| 4 | |
| 5 | This script creates a new model with fewer layers by: |
| 6 | 1. Loading the original Maverick model configuration |
| 7 | 2. Creating a reduced configuration |
| 8 | 3. Generating compatible safetensors files with appropriate weights |
| 9 | 4. Creating the necessary index files for vLLM compatibility |
| 10 | """ |
| 11 | |
| 12 | import json |
| 13 | import shutil |
| 14 | from pathlib import Path |
| 15 | from typing import Any |
| 16 | |
| 17 | import pytest |
| 18 | import torch |
| 19 | from safetensors.torch import save_file |
| 20 | from transformers import AutoConfig, AutoProcessor, AutoTokenizer, GenerationConfig |
| 21 | |
| 22 | from vllm import LLM, SamplingParams |
| 23 | from vllm.v1.executor.abstract import Executor |
| 24 | from vllm.v1.kv_cache_interface import ChunkedLocalAttentionSpec, FullAttentionSpec |
| 25 | |
| 26 | from ....utils import multi_gpu_test |
| 27 | |
| 28 | # Sample prompts for testing |
| 29 | PROMPTS: list[str] = [ |
| 30 | "Hello, my name is", |
| 31 | "The president of the United States is", |
| 32 | "The capital of France is", |
| 33 | "The future of AI is", |
| 34 | ] |
| 35 | |
| 36 | |
| 37 | def run_maverick_serving(model: str): |
| 38 | """Test Llama-4-Maverick model with vLLM LLM class using CLI equivalent |
| 39 | options with reduced layers. |
| 40 | """ |
| 41 | try: |
| 42 | sampling_params = SamplingParams(temperature=0.8, top_p=0.95) |
| 43 | |
| 44 | llm = LLM( |
| 45 | model=model, |
| 46 | max_model_len=2048, |
| 47 | enforce_eager=True, |
| 48 | tensor_parallel_size=8, |
| 49 | enable_expert_parallel=True, |
| 50 | trust_remote_code=True, |
| 51 | gpu_memory_utilization=0.4, |
| 52 | kv_cache_dtype="fp8", |
| 53 | ) |
| 54 | |
| 55 | outputs = llm.generate(PROMPTS, sampling_params) |
| 56 | |
| 57 | # Print the outputs |
| 58 | print("\nGenerated Outputs:\n" + "-" * 60) |
| 59 | for output in outputs: |
| 60 | prompt = output.prompt |
| 61 | generated_text = output.outputs[0].text |
| 62 | print(f"Prompt: {prompt!r}") |
| 63 | print(f"Output: {generated_text!r}") |
| 64 | print("-" * 60) |
| 65 | |
| 66 | except Exception as e: |
| 67 | print(f"Error initializing or running model: {e}") |
| 68 | raise |
| 69 | |
| 70 | |
| 71 | def get_rope_layers_config(model_path: str) -> list[int]: |
| 72 | """Get the interleaved RoPE configuration from HuggingFace config. |
| 73 | |
| 74 | Args: |
| 75 | model_path: Path to the local directory containing the reduced |
| 76 | Maverick model checkpoint |
| 77 | |
| 78 | Returns: |
| 79 | List of 0 or 1 indicating whether each layer uses RoPE and local attn |
| 80 | 0 indicates that RoPE is not used while 1 indicates that RoPE is used. |
| 81 | |
| 82 | """ |
| 83 | config_path = Path(model_path) / "config.json" |
| 84 | model_config = json.loads(config_path.read_text()) |
| 85 | text_config = model_config["text_config"] |
| 86 | no_rope_layers = text_config["no_rope_layers"] |
| 87 | print(f"Found no_rope_layers: {no_rope_layers}") |
| 88 | return no_rope_layers |
| 89 | |
| 90 | |
| 91 | def create_reduced_maverick_model( |
| 92 | original_model_name: str = "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", |
| 93 | output_dir: str = "/tmp/reduced_maverick", |
| 94 | text_layers: int = 4, |
| 95 | num_experts: int = 4, |
| 96 | vision_layers: int = 2, |
| 97 | force_recreate: bool = False, |
| 98 | ) -> str: |
| 99 | """Create a reduced-layer version of the Maverick model. |
| 100 | |
| 101 | Args: |
| 102 | original_model_name: Name of the original Maverick model |
| 103 | output_dir: Directory to save the reduced model |
| 104 | text_layers: Number of text transformer layers |
| 105 | num_experts: Number of experts per layer |
| 106 | vision_layers: Number of vision transformer layers |
| 107 | force_recreate: Whether to recreate if output_dir already exists |
| 108 | |
| 109 | Returns: |
| 110 | Path to the created reduced model directory |
| 111 | |
| 112 | """ |
| 113 | print( |
| 114 | f"Creating reduced Maverick model with {text_layers} text layers and " |
| 115 | f"{vision_layers} vision layers..." |
| 116 | ) |
| 117 | |
| 118 | # Create output directory |
| 119 | output_path = Path(output_dir) |
| 120 | if output_path.exists(): |
| 121 | if force_recreate: |
| 122 | shutil.rmtree(output_path) |
| 123 | else: |
| 124 | print( |
| 125 | f"Output directory {output_dir} already exists. " |
| 126 | "Use --force-recreate to overwrite." |
| 127 | ) |
| 128 | return str(output_path) |
| 129 | |
| 130 | output_path.mkdir(parents=True, exist_ok=True) |
| 131 | |
| 132 | try: |
| 133 | print("Loading original model configuration...") |
| 134 | original_config = AutoConfig.from_pretrained( |
| 135 | original_model_name, trust_remote_code=True |
| 136 | ) |
| 137 | print("Creating reduced configuration...") |
| 138 | reduced_config = create_reduced_config( |
| 139 | original_config, text_layers, num_experts, vision_layers |
| 140 | ) |
| 141 | |
| 142 | config_path = output_path / "config.json" |
| 143 | with open(config_path, "w") as f: |
| 144 | json.dump(reduced_config, f, indent=2) |
| 145 | print(f"Saved reduced config to {config_path}") |
| 146 | |
| 147 | print("Copying tokenizer files...") |
| 148 | copy_tokenizer_files(original_model_name, output_path) |
| 149 | |
| 150 | print("Creating reduced safetensors files...") |
| 151 | create_reduced_safetensors(original_config, reduced_config, output_path) |
| 152 | |
| 153 | print("Creating preprocessor config...") |
| 154 | create_preprocessor_config(original_config, output_path) |
| 155 | |
| 156 | try: |
| 157 | gen_config = GenerationConfig.from_pretrained(original_model_name) |
| 158 | gen_config.save_pretrained(output_path) |
| 159 | print("Copied generation config") |
| 160 | except Exception as e: |
| 161 | print(f"Could not copy generation config: {e}") |
| 162 | |
| 163 | print(f"Successfully created reduced Maverick model at {output_path}") |
| 164 | return str(output_path) |
| 165 | |
| 166 | except Exception as e: |
| 167 | print(f"Error creating reduced model: {e}") |
| 168 | # Clean up on failure |
| 169 | if output_path.exists(): |
| 170 | shutil.rmtree(output_path) |
| 171 | raise |
| 172 | |
| 173 | |
| 174 | def create_reduced_config( |
| 175 | original_config: Any, text_layers: int, num_experts: int, vision_layers: int |
| 176 | ) -> dict[str, Any]: |
| 177 | """Create a reduced configuration based on the original.""" |
| 178 | # Convert config to dictionary |
| 179 | config_dict = original_config.to_dict() |
| 180 | |
| 181 | # Reduce text layers |
| 182 | if "text_config" in config_dict: |
| 183 | original_text_layers = config_dict["text_config"]["num_hidden_layers"] |
| 184 | config_dict["text_config"]["num_hidden_layers"] = text_layers |
| 185 | original_layer_types = config_dict["text_config"]["layer_types"] |
| 186 | config_dict["text_config"]["layer_types"] = original_layer_types[:text_layers] |
| 187 | print(f"Reduced text layers from {original_text_layers} to {text_layers}") |
| 188 | |
| 189 | original_num_experts = config_dict["text_config"]["num_local_experts"] |
| 190 | config_dict["text_config"]["num_local_experts"] = num_experts |
| 191 | print(f"Reduced num experts from {original_num_experts} to {num_experts}") |
| 192 | |
| 193 | hidden_dim_divisor = 4 |
| 194 | |
| 195 | original_hidden_size = config_dict["text_config"]["hidden_size"] |
| 196 | new_hidden_size = original_hidden_size // hidden_dim_divisor |
| 197 | config_dict["text_config"]["hidden_size"] = new_hidden_size |
| 198 | print(f"Reduced hidden size from {original_hidden_size} to {new_hidden_size}") |
| 199 | |
| 200 | original_head_dim = config_dict["text_config"]["head_dim"] |
| 201 | new_head_dim = original_head_dim // hidden_dim_divisor |
| 202 | config_dict["text_config"]["head_dim"] = new_head_dim |
| 203 | print(f"Reduced head dim from {original_head_dim} to {new_head_dim}") |
| 204 | |
| 205 | # Reduce vision layers |
| 206 | if "vision_config" in config_dict: |
| 207 | original_vision_layers = config_dict["vision_config"]["num_hidden_layers"] |
| 208 | config_dict["vision_config"]["num_hidden_layers"] = vision_layers |
| 209 | print(f"Reduced vision layers from {original_vision_layers} to {vision_layers}") |
| 210 | |
| 211 | # Update model name to indicate it's a reduced version |
| 212 | config_dict["_name_or_path"] = f"reduced_maverick_{text_layers}t_{vision_layers}v" |
| 213 | |
| 214 | return config_dict |
| 215 | |
| 216 | |
| 217 | def copy_tokenizer_files(original_model_name: str, output_path: Path) -> None: |
| 218 | """Copy tokenizer files from the original model.""" |
| 219 | try: |
| 220 | tokenizer = AutoTokenizer.from_pretrained( |
| 221 | original_model_name, trust_remote_code=True |
| 222 | ) |
| 223 | tokenizer.save_pretrained(output_path) |
| 224 | print("Tokenizer files copied successfully") |
| 225 | except Exception as e: |
| 226 | print(f"Warning: Could not copy tokenizer files: {e}") |
| 227 | |
| 228 | |
| 229 | def create_preprocessor_config(original_config: Any, output_path: Path) -> None: |
| 230 | """Create preprocessor_config.json for multimodal model.""" |
| 231 | # Try to load the original preprocessor config |
| 232 | try: |
| 233 | processor = AutoProcessor.from_pretrained( |
| 234 | original_config._name_or_path |
| 235 | or "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", |
| 236 | trust_remote_code=True, |
| 237 | ) |
| 238 | processor.save_pretrained(output_path) |
| 239 | print("Copied original preprocessor config") |
| 240 | return |
| 241 | except Exception as e: |
| 242 | print(f"Could not copy original preprocessor config: {e}") |
| 243 | raise |
| 244 | |
| 245 | |
| 246 | def create_reduced_safetensors( |
| 247 | original_config: Any, reduced_config: dict[str, Any], output_path: Path |
| 248 | ) -> None: |
| 249 | """Create safetensors files with weights for the reduced model.""" |
| 250 | print("Generating synthetic weights for reduced model...") |
| 251 | |
| 252 | text_config = reduced_config["text_config"] |
| 253 | vision_config = reduced_config["vision_config"] |
| 254 | |
| 255 | weights = {} |
| 256 | |
| 257 | print("Creating text model weights...") |
| 258 | weights.update(create_text_model_weights(text_config)) |
| 259 | |
| 260 | print("Creating vision model weights...") |
| 261 | weights.update(create_vision_model_weights(vision_config)) |
| 262 | |
| 263 | print("Creating shared model weights...") |
| 264 | weights.update(create_shared_weights(text_config, vision_config)) |
| 265 | |
| 266 | print("Saving weights to safetensors files...") |
| 267 | save_weights_to_safetensors(weights, output_path) |
| 268 | |
| 269 | |
| 270 | def create_text_model_weights(text_config: dict[str, Any]) -> dict[str, torch.Tensor]: |
| 271 | """Create synthetic weights for the text model with MoE structure.""" |
| 272 | weights = {} |
| 273 | |
| 274 | vocab_size = text_config["vocab_size"] |
| 275 | hidden_size = text_config["hidden_size"] |
| 276 | intermediate_size = text_config["intermediate_size"] |
| 277 | intermediate_size_mlp = text_config["intermediate_size_mlp"] |
| 278 | num_layers = text_config["num_hidden_layers"] |
| 279 | num_attention_heads = text_config["num_attention_heads"] |
| 280 | num_key_value_heads = text_config.get("num_key_value_heads", num_attention_heads) |
| 281 | |
| 282 | # MoE specific parameters |
| 283 | num_experts = text_config.get("num_local_experts") |
| 284 | assert num_experts is not None, "num_local_experts must be specified for MoE" |
| 285 | |
| 286 | head_dim = hidden_size // num_attention_heads |
| 287 | |
| 288 | # Embedding layers |
| 289 | weights["language_model.model.embed_tokens.weight"] = torch.randn( |
| 290 | vocab_size, hidden_size, dtype=torch.float16 |
| 291 | ) |
| 292 | |
| 293 | # Transformer layers |
| 294 | for layer_idx in range(num_layers): |
| 295 | layer_prefix = f"language_model.model.layers.{layer_idx}" |
| 296 | print(f"Creating weights for layer {layer_prefix}...") |
| 297 | |
| 298 | # Self-attention weights (separate q, k, v projections) |
| 299 | weights[f"{layer_prefix}.self_attn.q_proj.weight"] = torch.randn( |
| 300 | num_attention_heads * head_dim, hidden_size, dtype=torch.bfloat16 |
| 301 | ) |
| 302 | weights[f"{layer_prefix}.self_attn.k_proj.weight"] = torch.randn( |
| 303 | num_key_value_heads * head_dim, hidden_size, dtype=torch.bfloat16 |
| 304 | ) |
| 305 | weights[f"{layer_prefix}.self_attn.v_proj.weight"] = torch.randn( |
| 306 | num_key_value_heads * head_dim, hidden_size, dtype=torch.bfloat16 |
| 307 | ) |
| 308 | weights[f"{layer_prefix}.self_attn.o_proj.weight"] = torch.randn( |
| 309 | hidden_size, num_attention_heads * head_dim, dtype=torch.bfloat16 |
| 310 | ) |
| 311 | print("Self-attention weights created.") |
| 312 | |
| 313 | # Feed-forward weights - MoE pattern based on interleave_moe_layer_step |
| 314 | # For interleave_moe_layer_step=2: layers 1,3,5,... are MoE, layers |
| 315 | # 0,2,4,... are dense |
| 316 | interleave_step = text_config.get("interleave_moe_layer_step", 1) |
| 317 | is_moe_layer = interleave_step > 0 and (layer_idx + 1) % interleave_step == 0 |
| 318 | |
| 319 | if is_moe_layer: |
| 320 | # MoE layer structure |
| 321 | # 1. Router weights |
| 322 | weights[f"{layer_prefix}.feed_forward.router.weight"] = torch.randn( |
| 323 | num_experts, hidden_size, dtype=torch.float16 |
| 324 | ) |
| 325 | |
| 326 | # 2. Individual expert weights (not fused) |
| 327 | for expert_idx in range(num_experts): |
| 328 | expert_prefix = f"{layer_prefix}.feed_forward.experts.{expert_idx}" |
| 329 | |
| 330 | weights[f"{expert_prefix}.gate_proj.weight"] = torch.randn( |
| 331 | intermediate_size, hidden_size, dtype=torch.bfloat16 |
| 332 | ) |
| 333 | weights[f"{expert_prefix}.up_proj.weight"] = torch.randn( |
| 334 | intermediate_size, hidden_size, dtype=torch.bfloat16 |
| 335 | ) |
| 336 | weights[f"{expert_prefix}.down_proj.weight"] = torch.randn( |
| 337 | hidden_size, intermediate_size, dtype=torch.bfloat16 |
| 338 | ) |
| 339 | |
| 340 | # Expert weight scales (FP8 quantization) |
| 341 | weights[f"{expert_prefix}.gate_proj.weight_scale"] = torch.ones( |
| 342 | intermediate_size, 1, dtype=torch.bfloat16 |
| 343 | ) |
| 344 | weights[f"{expert_prefix}.up_proj.weight_scale"] = torch.ones( |
| 345 | intermediate_size, 1, dtype=torch.bfloat16 |
| 346 | ) |
| 347 | weights[f"{expert_prefix}.down_proj.weight_scale"] = torch.ones( |
| 348 | hidden_size, 1, dtype=torch.bfloat16 |
| 349 | ) |
| 350 | |
| 351 | # 3. Shared expert weights |
| 352 | shared_expert_prefix = f"{layer_prefix}.feed_forward.shared_expert" |
| 353 | weights[f"{shared_expert_prefix}.gate_proj.weight"] = torch.randn( |
| 354 | intermediate_size, hidden_size, dtype=torch.bfloat16 |
| 355 | ) |
| 356 | weights[f"{shared_expert_prefix}.up_proj.weight"] = torch.randn( |
| 357 | intermediate_size, hidden_size, dtype=torch.bfloat16 |
| 358 | ) |
| 359 | weights[f"{shared_expert_prefix}.down_proj.weight"] = torch.randn( |
| 360 | hidden_size, intermediate_size, dtype=torch.bfloat16 |
| 361 | ) |
| 362 | print(f"MoE feed-forward weights created for layer {layer_idx}.") |
| 363 | else: |
| 364 | # Dense layer structure |
| 365 | weights[f"{layer_prefix}.feed_forward.gate_proj.weight"] = torch.randn( |
| 366 | intermediate_size_mlp, hidden_size, dtype=torch.bfloat16 |
| 367 | ) |
| 368 | weights[f"{layer_prefix}.feed_forward.up_proj.weight"] = torch.randn( |
| 369 | intermediate_size_mlp, hidden_size, dtype=torch.bfloat16 |
| 370 | ) |
| 371 | weights[f"{layer_prefix}.feed_forward.down_proj.weight"] = torch.randn( |
| 372 | hidden_size, intermediate_size_mlp, dtype=torch.bfloat16 |
| 373 | ) |
| 374 | print(f"Dense feed-forward weights created for layer {layer_idx}.") |
| 375 | |
| 376 | # Layer norms |
| 377 | weights[f"{layer_prefix}.input_layernorm.weight"] = torch.ones( |
| 378 | hidden_size, dtype=torch.bfloat16 |
| 379 | ) |
| 380 | weights[f"{layer_prefix}.post_attention_layernorm.weight"] = torch.ones( |
| 381 | hidden_size, dtype=torch.bfloat16 |
| 382 | ) |
| 383 | print("Layer norms created.") |
| 384 | |
| 385 | # Final layer norm and output projection |
| 386 | weights["language_model.model.norm.weight"] = torch.ones( |
| 387 | hidden_size, dtype=torch.bfloat16 |
| 388 | ) |
| 389 | weights["language_model.lm_head.weight"] = torch.randn( |
| 390 | vocab_size, hidden_size, dtype=torch.bfloat16 |
| 391 | ) |
| 392 | |
| 393 | return weights |
| 394 | |
| 395 | |
| 396 | def create_vision_model_weights( |
| 397 | vision_config: dict[str, Any], |
| 398 | ) -> dict[str, torch.Tensor]: |
| 399 | """Create synthetic weights for the vision model.""" |
| 400 | weights = {} |
| 401 | |
| 402 | hidden_size = vision_config["hidden_size"] |
| 403 | intermediate_size = vision_config["intermediate_size"] |
| 404 | num_layers = vision_config["num_hidden_layers"] |
| 405 | |
| 406 | # Vision transformer layers |
| 407 | for layer_idx in range(num_layers): |
| 408 | layer_prefix = f"vision_model.model.layers.{layer_idx}" |
| 409 | |
| 410 | weights[f"{layer_prefix}.self_attn.q_proj.weight"] = torch.randn( |
| 411 | hidden_size, hidden_size, dtype=torch.bfloat16 |
| 412 | ) |
| 413 | weights[f"{layer_prefix}.self_attn.q_proj.bias"] = torch.zeros( |
| 414 | hidden_size, dtype=torch.bfloat16 |
| 415 | ) |
| 416 | weights[f"{layer_prefix}.self_attn.k_proj.weight"] = torch.randn( |
| 417 | hidden_size, hidden_size, dtype=torch.bfloat16 |
| 418 | ) |
| 419 | weights[f"{layer_prefix}.self_attn.k_proj.bias"] = torch.zeros( |
| 420 | hidden_size, dtype=torch.bfloat16 |
| 421 | ) |
| 422 | weights[f"{layer_prefix}.self_attn.v_proj.weight"] = torch.randn( |
| 423 | hidden_size, hidden_size, dtype=torch.bfloat16 |
| 424 | ) |
| 425 | weights[f"{layer_prefix}.self_attn.v_proj.bias"] = torch.zeros( |
| 426 | hidden_size, dtype=torch.bfloat16 |
| 427 | ) |
| 428 | weights[f"{layer_prefix}.self_attn.o_proj.weight"] = torch.randn( |
| 429 | hidden_size, hidden_size, dtype=torch.bfloat16 |
| 430 | ) |
| 431 | weights[f"{layer_prefix}.self_attn.o_proj.bias"] = torch.zeros( |
| 432 | hidden_size, dtype=torch.bfloat16 |
| 433 | ) |
| 434 | |
| 435 | weights[f"{layer_prefix}.mlp.fc1.weight"] = torch.randn( |
| 436 | intermediate_size, hidden_size, dtype=torch.bfloat16 |
| 437 | ) |
| 438 | weights[f"{layer_prefix}.mlp.fc1.bias"] = torch.zeros( |
| 439 | intermediate_size, dtype=torch.bfloat16 |
| 440 | ) |
| 441 | weights[f"{layer_prefix}.mlp.fc2.weight"] = torch.randn( |
| 442 | hidden_size, intermediate_size, dtype=torch.bfloat16 |
| 443 | ) |
| 444 | weights[f"{layer_prefix}.mlp.fc2.bias"] = torch.zeros( |
| 445 | hidden_size, dtype=torch.bfloat16 |
| 446 | ) |
| 447 | |
| 448 | weights[f"{layer_prefix}.input_layernorm.weight"] = torch.ones( |
| 449 | hidden_size, dtype=torch.bfloat16 |
| 450 | ) |
| 451 | weights[f"{layer_prefix}.input_layernorm.bias"] = torch.zeros( |
| 452 | hidden_size, dtype=torch.bfloat16 |
| 453 | ) |
| 454 | weights[f"{layer_prefix}.post_attention_layernorm.weight"] = torch.ones( |
| 455 | hidden_size, dtype=torch.bfloat16 |
| 456 | ) |
| 457 | weights[f"{layer_prefix}.post_attention_layernorm.bias"] = torch.zeros( |
| 458 | hidden_size, dtype=torch.bfloat16 |
| 459 | ) |
| 460 | |
| 461 | return weights |
| 462 | |
| 463 | |
| 464 | def create_shared_weights( |
| 465 | text_config: dict[str, Any], vision_config: dict[str, Any] |
| 466 | ) -> dict[str, torch.Tensor]: |
| 467 | """Create weights for shared components (vision-language connector).""" |
| 468 | weights = {} |
| 469 | |
| 470 | text_hidden_size = text_config["hidden_size"] |
| 471 | projector_input_dim = vision_config["projector_input_dim"] |
| 472 | |
| 473 | # Vision-language connector (projects vision features to text space) |
| 474 | weights["multi_modal_projector.linear_1.weight"] = torch.randn( |
| 475 | text_hidden_size, projector_input_dim, dtype=torch.bfloat16 |
| 476 | ) |
| 477 | |
| 478 | return weights |
| 479 | |
| 480 | |
| 481 | def save_weights_to_safetensors( |
| 482 | weights: dict[str, torch.Tensor], output_path: Path |
| 483 | ) -> None: |
| 484 | """Save weights to safetensors files and create index.""" |
| 485 | # Determine how to shard the weights |
| 486 | max_shard_size = 5 * 1024 * 1024 * 1024 # 5GB per shard |
| 487 | |
| 488 | # Calculate sizes and create shards |
| 489 | shards = [] |
| 490 | current_shard: dict[str, torch.Tensor] = {} |
| 491 | current_size = 0 |
| 492 | |
| 493 | for name, tensor in weights.items(): |
| 494 | tensor_size = tensor.numel() * tensor.element_size() |
| 495 | |
| 496 | if current_size + tensor_size > max_shard_size and current_shard: |
| 497 | shards.append(current_shard) |
| 498 | current_shard = {} |
| 499 | current_size = 0 |
| 500 | |
| 501 | current_shard[name] = tensor |
| 502 | current_size += tensor_size |
| 503 | |
| 504 | if current_shard: |
| 505 | shards.append(current_shard) |
| 506 | |
| 507 | # Save shards and create index |
| 508 | weight_map = {} |
| 509 | |
| 510 | if len(shards) == 1: |
| 511 | # Single file |
| 512 | filename = "model.safetensors" |
| 513 | save_file(shards[0], output_path / filename) |
| 514 | weight_map = {name: filename for name in shards[0]} |
| 515 | print(f"Saved weights to single file: {filename}") |
| 516 | else: |
| 517 | # Multiple shards |
| 518 | for i, shard in enumerate(shards): |
| 519 | filename = f"model-{i + 1:05d}-of-{len(shards):05d}.safetensors" |
| 520 | save_file(shard, output_path / filename) |
| 521 | for name in shard: |
| 522 | weight_map[name] = filename |
| 523 | print(f"Saved shard {i + 1}/{len(shards)}: {filename}") |
| 524 | |
| 525 | # Create index file |
| 526 | index_data = { |
| 527 | "metadata": { |
| 528 | "total_size": sum( |
| 529 | tensor.numel() * tensor.element_size() for tensor in weights.values() |
| 530 | ) |
| 531 | }, |
| 532 | "weight_map": weight_map, |
| 533 | } |
| 534 | |
| 535 | index_path = output_path / "model.safetensors.index.json" |
| 536 | with open(index_path, "w") as f: |
| 537 | json.dump(index_data, f, indent=2) |
| 538 | |
| 539 | print(f"Created index file: {index_path}") |
| 540 | print( |
| 541 | f"Total model size: {index_data['metadata']['total_size'] / (1024**3):.2f} GB" |
| 542 | ) |
| 543 | |
| 544 | |
| 545 | def check_attention_spec_interleaved_rope( |
| 546 | llm: LLM, |
| 547 | num_attention_layers: int, |
| 548 | num_ranks: int, |
| 549 | rope_layers: list[int], |
| 550 | ): |
| 551 | """Check that the attention spec is correct.""" |
| 552 | assert isinstance(llm.llm_engine.model_executor, Executor) |
| 553 | kv_cache_specs_per_rank = llm.llm_engine.model_executor.get_kv_cache_specs() |
| 554 | for rank in range(num_ranks): |
| 555 | kv_cache_specs = kv_cache_specs_per_rank[rank] |
| 556 | assert len(kv_cache_specs.keys()) == num_attention_layers |
| 557 | for i in range(num_attention_layers): |
| 558 | if rope_layers[i] == 0: |
| 559 | expected_spec = FullAttentionSpec |
| 560 | else: |
| 561 | expected_spec = ChunkedLocalAttentionSpec |
| 562 | assert isinstance( |
| 563 | kv_cache_specs[f"language_model.model.layers.{i}.self_attn.attn"], |
| 564 | expected_spec, |
| 565 | ) |
| 566 | |
| 567 | |
| 568 | def run_reduced_model(llm: LLM, should_profile: bool = False) -> None: |
| 569 | """Test the created reduced model with vLLM.""" |
| 570 | sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=50) |
| 571 | |
| 572 | if should_profile: |
| 573 | llm.start_profile() |
| 574 | outputs = llm.generate(PROMPTS, sampling_params) |
| 575 | if should_profile: |
| 576 | llm.stop_profile() |
| 577 | |
| 578 | print("Test generation successful!") |
| 579 | for output in outputs: |
| 580 | print(f"Prompt: {output.prompt}") |
| 581 | print(f"Output: {output.outputs[0].text}") |
| 582 | print("-" * 40) |
| 583 | |
| 584 | |
| 585 | @multi_gpu_test(num_gpus=2) |
| 586 | @pytest.mark.parametrize( |
| 587 | "original_model_name,text_layers,num_experts,vision_layers,", |
| 588 | [("meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", 4, 4, 2)], |
| 589 | ) |
| 590 | @pytest.mark.parametrize("enforce_eager", [True, False]) |
| 591 | @pytest.mark.parametrize("tp,ep", [(2, True)]) |
| 592 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") |
| 593 | def test_dummy_maverick( |
| 594 | monkeypatch, |
| 595 | original_model_name: str, |
| 596 | text_layers: int, |
| 597 | num_experts: int, |
| 598 | vision_layers: int, |
| 599 | enforce_eager: bool, |
| 600 | tp: int, |
| 601 | ep: bool, |
| 602 | output_dir: str = "/tmp/reduced_maverick", |
| 603 | force_recreate: bool = True, |
| 604 | profile: bool = False, |
| 605 | ) -> None: |
| 606 | # Disable multiprocessing allows us to access model executor from LLM engine |
| 607 | monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") |
| 608 | # The forked test process may inherit CUDA state from test collection. |
| 609 | monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") |
| 610 | |
| 611 | model_path = create_reduced_maverick_model( |
| 612 | original_model_name=original_model_name, |
| 613 | output_dir=output_dir, |
| 614 | text_layers=text_layers, |
| 615 | num_experts=num_experts, |
| 616 | vision_layers=vision_layers, |
| 617 | force_recreate=force_recreate, |
| 618 | ) |
| 619 | |
| 620 | print(f"\nReduced model created successfully at: {model_path}") |
| 621 | |
| 622 | rope_layers = get_rope_layers_config(model_path) |
| 623 | |
| 624 | llm = LLM( |
| 625 | model=model_path, |
| 626 | trust_remote_code=True, |
| 627 | max_model_len=512, # Small context for testing |
| 628 | gpu_memory_utilization=0.3, # Conservative memory usage |
| 629 | enforce_eager=enforce_eager, |
| 630 | tensor_parallel_size=tp, |
| 631 | enable_expert_parallel=ep, |
| 632 | ) |
| 633 | |
| 634 | check_attention_spec_interleaved_rope( |
| 635 | llm, |
| 636 | text_layers, |
| 637 | tp, |
| 638 | rope_layers, |
| 639 | ) |
| 640 | |
| 641 | print(f"\nTesting reduced model at {model_path}...") |
| 642 | run_reduced_model(llm=llm, should_profile=profile) |
| 643 | |
| 644 | |
| 645 | def main(): |
| 646 | """Main function to create and test the reduced model.""" |
| 647 | import argparse |
| 648 | |
| 649 | parser = argparse.ArgumentParser( |
| 650 | description="Create a reduced-layer Maverick model" |
| 651 | ) |
| 652 | parser.add_argument( |
| 653 | "--output-dir", |
| 654 | default="/tmp/reduced_maverick", |
| 655 | help="Output directory for the reduced model", |
| 656 | ) |
| 657 | parser.add_argument( |
| 658 | "--text-layers", |
| 659 | type=int, |
| 660 | default=4, |
| 661 | help="Number of text transformer layers", |
| 662 | ) |
| 663 | parser.add_argument("--num-experts", type=int, default=4, help="Number of experts") |
| 664 | parser.add_argument( |
| 665 | "--vision-layers", |
| 666 | type=int, |
| 667 | default=2, |
| 668 | help="Number of vision transformer layers", |
| 669 | ) |
| 670 | parser.add_argument( |
| 671 | "--force-recreate", |
| 672 | action="store_true", |
| 673 | help="Force recreation if output directory exists", |
| 674 | ) |
| 675 | parser.add_argument( |
| 676 | "--test", action="store_true", help="Test the created model with vLLM" |
| 677 | ) |
| 678 | parser.add_argument( |
| 679 | "--profile", action="store_true", help="Profile the created model with vLLM" |
| 680 | ) |
| 681 | parser.add_argument( |
| 682 | "--test-original", |
| 683 | action="store_true", |
| 684 | help="Test the original model with vLLM", |
| 685 | ) |
| 686 | parser.add_argument( |
| 687 | "--original-model", |
| 688 | default="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", |
| 689 | help="Original model name to base the reduction on", |
| 690 | ) |
| 691 | |
| 692 | args = parser.parse_args() |
| 693 | |
| 694 | if args.test: |
| 695 | test_dummy_maverick( |
| 696 | original_model_name=args.original_model, |
| 697 | output_dir=args.output_dir, |
| 698 | text_layers=args.text_layers, |
| 699 | num_experts=args.num_experts, |
| 700 | vision_layers=args.vision_layers, |
| 701 | force_recreate=args.force_recreate, |
| 702 | tp=2, |
| 703 | ep=True, |
| 704 | enforce_eager=True, |
| 705 | profile=args.profile, |
| 706 | ) |
| 707 | |
| 708 | if args.test_original: |
| 709 | run_maverick_serving(args.original_model) |
| 710 | |
| 711 | |
| 712 | if __name__ == "__main__": |
| 713 | exit(main()) |