diff --git a/KEY_CONVERSIONS.md b/KEY_CONVERSIONS.md new file mode 100644 index 0000000000..d72036c5aa --- /dev/null +++ b/KEY_CONVERSIONS.md @@ -0,0 +1,151 @@ +# SafeTensors Key Name Conversions + +This document describes the key name transformations applied by `convert_safetensors_keys.py`. + +## Transformation Rules + +### Top-Level Keys +| Original (Quantized) | Converted (BF16 Format) | Note | +|---------------------|------------------------|------| +| `head.weight` | `lm_head.weight` | LM head weight | +| `model.embed_tokens.weight` | `embed.weight` | Embedding layer | +| `model.norm.*` | `norm.*` | Model normalization | +| `model.layers.X.*` | `layers.X.*` | Strip `model.` prefix | + +### Layer Normalization +| Original (Quantized) | Converted (BF16 Format) | Note | +|---------------------|------------------------|------| +| `layers.X.input_layernorm.weight` | `layers.X.attn_norm.weight` | Attention pre-norm | +| `layers.X.post_attention_layernorm.weight` | `layers.X.ffn_norm.weight` | FFN pre-norm | +| `layers.X.self_attn.*` | `layers.X.attn.*` | Self-attention → attn | +| `layers.X.attn.q_a_norm.weight` | `layers.X.attn.q_norm.weight` | Query normalization | +| `layers.X.attn.compressor.kv_norm.weight` | `layers.X.attn.compressor.norm.weight` | Compressor norm | +| `layers.X.attn.compressor.indexer.kv_norm.weight` | `layers.X.attn.indexer.compressor.norm.weight` | Indexer compressor norm | + +### Attention Projections +| Original (Quantized) | Converted (BF16 Format) | Note | +|---------------------|------------------------|------| +| `layers.X.attn.sinks` | `layers.X.attn.attn_sink` | Attention sink | +| `layers.X.attn.kv_proj.*` | `layers.X.attn.wkv.*` | KV projection | +| `layers.X.attn.o_a_proj.*` | `layers.X.attn.wo_a.*` | Output A projection | +| `layers.X.attn.o_b_proj.*` | `layers.X.attn.wo_b.*` | Output B projection | +| `layers.X.attn.q_a_proj.*` | `layers.X.attn.wq_a.*` | Query A projection | +| `layers.X.attn.q_b_proj.*` | `layers.X.attn.wq_b.*` | Query B projection | + +### Compressor Projections +| Original (Quantized) | Converted (BF16 Format) | Note | +|---------------------|------------------------|------| +| `layers.X.attn.compressor.gate_proj.*` | `layers.X.attn.compressor.wgate.*` | Gate projection | +| `layers.X.attn.compressor.kv_proj.*` | `layers.X.attn.compressor.wkv.*` | KV projection | +| `layers.X.attn.compressor.position_bias` | `layers.X.attn.compressor.ape` | Positional encoding | + +### Indexer Compressor Projections +| Original (Quantized) | Converted (BF16 Format) | Note | +|---------------------|------------------------|------| +| `layers.X.attn.compressor.indexer.gate_proj.*` | `layers.X.attn.indexer.compressor.wgate.*` | Gate (reordered) | +| `layers.X.attn.compressor.indexer.kv_proj.*` | `layers.X.attn.indexer.compressor.wkv.*` | KV (reordered) | +| `layers.X.attn.compressor.indexer.position_bias` | `layers.X.attn.indexer.compressor.ape` | APE (reordered) | +| `layers.X.attn.compressor.indexer.q_b_proj.*` | `layers.X.attn.indexer.wq_b.*` | Query B (reordered) | +| `layers.X.attn.compressor.indexer.scorer.weights_proj.*` | `layers.X.attn.indexer.weights_proj.*` | Weights (flattened) | + +### FFN/MLP Transformations +| Original (Quantized) | Converted (BF16 Format) | Note | +|---------------------|------------------------|------| +| `layers.X.mlp.*` | `layers.X.ffn.*` | MLP → FFN | +| `layers.X.ffn.experts.Y.gate_proj.*` | `layers.X.ffn.experts.Y.w1.*` | Expert gate projection | +| `layers.X.ffn.experts.Y.up_proj.*` | `layers.X.ffn.experts.Y.w3.*` | Expert up projection | +| `layers.X.ffn.experts.Y.down_proj.*` | `layers.X.ffn.experts.Y.w2.*` | Expert down projection | +| `layers.X.ffn.shared_experts.gate_proj.*` | `layers.X.ffn.shared_experts.w1.*` | Shared gate projection | +| `layers.X.ffn.shared_experts.up_proj.*` | `layers.X.ffn.shared_experts.w3.*` | Shared up projection | +| `layers.X.ffn.shared_experts.down_proj.*` | `layers.X.ffn.shared_experts.w2.*` | Shared down projection | + +### Hadamard Compression (HC) Keys +| Original (Quantized) | Converted (BF16 Format) | Note | +|---------------------|------------------------|------| +| `model.hc_head.hc_base` | `hc_head_base` | HC head base (flatten) | +| `model.hc_head.hc_fn` | `hc_head_fn` | HC head function | +| `model.hc_head.hc_scale` | `hc_head_scale` | HC head scale | +| `layers.X.attn_hc.base` | `layers.X.hc_attn_base` | Attn HC base (reorder) | +| `layers.X.attn_hc.fn` | `layers.X.hc_attn_fn` | Attn HC function | +| `layers.X.attn_hc.scale` | `layers.X.hc_attn_scale` | Attn HC scale | +| `layers.X.ffn_hc.base` | `layers.X.hc_ffn_base` | FFN HC base (reorder) | +| `layers.X.ffn_hc.fn` | `layers.X.hc_ffn_fn` | FFN HC function | +| `layers.X.ffn_hc.scale` | `layers.X.hc_ffn_scale` | FFN HC scale | + +## Examples + +### Expert Layer Keys +``` +# Before +model.layers.0.mlp.experts.0.gate_proj.weight_packed +model.layers.0.mlp.experts.0.up_proj.weight_packed +model.layers.0.mlp.experts.0.down_proj.weight_packed + +# After +layers.0.ffn.experts.0.w1.weight_packed +layers.0.ffn.experts.0.w3.weight_packed +layers.0.ffn.experts.0.w2.weight_packed +``` + +### HC (Hadamard Compression) Keys +HC keys are restructured to match BF16 format: + +``` +# Before +model.hc_head.hc_base +model.hc_head.hc_fn +model.hc_head.hc_scale +model.layers.0.attn_hc.base +model.layers.0.attn_hc.fn +model.layers.0.attn_hc.scale +model.layers.0.ffn_hc.base +model.layers.0.ffn_hc.fn +model.layers.0.ffn_hc.scale + +# After +hc_head_base +hc_head_fn +hc_head_scale +layers.0.hc_attn_base +layers.0.hc_attn_fn +layers.0.hc_attn_scale +layers.0.hc_ffn_base +layers.0.hc_ffn_fn +layers.0.hc_ffn_scale +``` + +### Quantization-Specific Keys +Quantization metadata keys (e.g., `input_global_scale`, `weight_global_scale`, `weight_scale`) are preserved with their parent key transformation: + +``` +# Before +model.layers.0.mlp.experts.0.down_proj.input_global_scale +model.layers.0.mlp.experts.0.down_proj.weight_global_scale +model.layers.0.mlp.experts.0.down_proj.weight_scale + +# After +layers.0.ffn.experts.0.w2.input_global_scale +layers.0.ffn.experts.0.w2.weight_global_scale +layers.0.ffn.experts.0.w2.weight_scale +``` + +## Usage + +### Dry Run (Preview Changes) +```bash +python convert_safetensors_keys.py /path/to/model --dry-run +``` + +### Convert In-Place +```bash +python convert_safetensors_keys.py /path/to/model +``` + +**WARNING**: This modifies files in-place. While the script uses temporary files to avoid corruption, ensure you have enough disk space (approximately the size of the model directory). + +## Technical Details + +- Uses temporary files during conversion to avoid corruption +- Updates `model.safetensors.index.json` to reflect new key names +- Removes keys that don't exist in the BF16 model format +- Requires `safetensors` Python package diff --git a/convert_config.py b/convert_config.py new file mode 100755 index 0000000000..03faeb788f --- /dev/null +++ b/convert_config.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python3 +""" +Sync config.json fields from BF16 model to quantized model, +preserving the quantization_config in the target. +""" +import json +import argparse +from pathlib import Path + + +def sync_configs(source_path: str, target_path: str, dry_run: bool = False): + """ + Update target config with fields from source config, except quantization_config. + + Args: + source_path: Path to the source config.json (BF16 model) + target_path: Path to the target config.json (quantized model) + dry_run: If True, print changes without writing + """ + # Load both configs + print(f"Loading source config from: {source_path}") + with open(source_path, 'r') as f: + source_config = json.load(f) + + print(f"Loading target config from: {target_path}") + with open(target_path, 'r') as f: + target_config = json.load(f) + + # Preserve the quantization_config from target + quantization_config = target_config.get('quantization_config') + + if quantization_config is None: + print("WARNING: No quantization_config found in target!") + else: + print(f"Preserving quantization_config (size: {len(json.dumps(quantization_config))} chars)") + + # Track changes + added_keys = [] + modified_keys = [] + removed_keys = [] + + # Find keys that will be added or modified + for key, value in source_config.items(): + if key == 'quantization_config': + continue # Skip this key entirely + + if key not in target_config: + added_keys.append(key) + elif target_config[key] != value: + modified_keys.append(key) + + # Find keys that will be removed (exist in target but not source) + for key in target_config.keys(): + if key == 'quantization_config': + continue # Don't remove this + if key not in source_config: + removed_keys.append(key) + + # Print summary + print("\n=== Change Summary ===") + print(f"Keys to add: {len(added_keys)}") + if added_keys: + for key in added_keys: + print(f" + {key}") + + print(f"\nKeys to modify: {len(modified_keys)}") + if modified_keys: + for key in modified_keys: + print(f" ~ {key}") + print(f" Old: {target_config[key]}") + print(f" New: {source_config[key]}") + + print(f"\nKeys to remove: {len(removed_keys)}") + if removed_keys: + for key in removed_keys: + print(f" - {key}") + + if quantization_config: + print("\nPreserved key:") + print(" = quantization_config (unchanged)") + + # Create new config + new_config = source_config.copy() + + # Restore quantization_config + if quantization_config is not None: + new_config['quantization_config'] = quantization_config + + # Write or show result + if dry_run: + print("\n=== DRY RUN - No changes written ===") + print("\nTo apply changes, run without --dry-run") + else: + # Backup original + backup_path = f"{target_path}.backup" + print(f"\nBacking up original to: {backup_path}") + with open(backup_path, 'w') as f: + json.dump(target_config, f, indent=2) + + # Write new config + print(f"Writing updated config to: {target_path}") + with open(target_path, 'w') as f: + json.dump(new_config, f, indent=2) + + print("\n=== SUCCESS ===") + print(f"Config updated successfully!") + print(f"Backup saved to: {backup_path}") + + +def main(): + parser = argparse.ArgumentParser( + description="Sync config.json fields while preserving quantization_config" + ) + parser.add_argument( + '--source', + default='/mnt/nvme-data/engine/kylesayrs/DeepSeek-V4-Pro-BF16/config.json', + help='Path to source config.json (default: BF16 model)' + ) + parser.add_argument( + '--target', + default='/mnt/nvme-data/engine/kylesayrs/DeepSeek-V4-Pro-NVFP4-FP8-BLOCK/config.json', + help='Path to target config.json (default: quantized model)' + ) + parser.add_argument( + '--dry-run', + action='store_true', + help='Show what would change without modifying files' + ) + + args = parser.parse_args() + + # Verify files exist + if not Path(args.source).exists(): + print(f"ERROR: Source file not found: {args.source}") + return 1 + + if not Path(args.target).exists(): + print(f"ERROR: Target file not found: {args.target}") + return 1 + + sync_configs(args.source, args.target, args.dry_run) + return 0 + + +if __name__ == '__main__': + exit(main()) diff --git a/convert_safetensors_keys.py b/convert_safetensors_keys.py new file mode 100755 index 0000000000..410a90cf39 --- /dev/null +++ b/convert_safetensors_keys.py @@ -0,0 +1,292 @@ +#!/usr/bin/env python3 +""" +Convert safetensors key names to match BF16 model format. + +Key transformations: +1. head.weight -> lm_head.weight +2. model.embed_tokens.weight -> embed.weight +3. model.norm.* -> norm.* +4. model.layers.X. -> layers.X. +5. layers.X.input_layernorm.weight -> layers.X.attn_norm.weight +6. layers.X.post_attention_layernorm.weight -> layers.X.ffn_norm.weight +7. layers.X.self_attn.* -> layers.X.attn.* +8. layers.X.attn.q_a_norm.weight -> layers.X.attn.q_norm.weight +9. layers.X.attn.compressor.kv_norm.weight -> layers.X.attn.compressor.norm.weight +10. layers.X.attn.compressor.indexer.kv_norm.weight -> layers.X.attn.indexer.compressor.norm.weight +11. layers.X.attn.sinks -> layers.X.attn.attn_sink +12. layers.X.attn.kv_proj.* -> layers.X.attn.wkv.* +13. layers.X.attn.o_a_proj.* -> layers.X.attn.wo_a.* +14. layers.X.attn.o_b_proj.* -> layers.X.attn.wo_b.* +15. layers.X.attn.q_a_proj.* -> layers.X.attn.wq_a.* +16. layers.X.attn.q_b_proj.* -> layers.X.attn.wq_b.* +17. layers.X.attn.compressor.gate_proj.* -> layers.X.attn.compressor.wgate.* +18. layers.X.attn.compressor.kv_proj.* -> layers.X.attn.compressor.wkv.* +19. layers.X.attn.compressor.position_bias -> layers.X.attn.compressor.ape +20. layers.X.attn.compressor.indexer.gate_proj.* -> layers.X.attn.indexer.compressor.wgate.* +21. layers.X.attn.compressor.indexer.kv_proj.* -> layers.X.attn.indexer.compressor.wkv.* +22. layers.X.attn.compressor.indexer.position_bias -> layers.X.attn.indexer.compressor.ape +23. layers.X.attn.compressor.indexer.q_b_proj.* -> layers.X.attn.indexer.wq_b.* +24. layers.X.attn.compressor.indexer.scorer.weights_proj.* -> layers.X.attn.indexer.weights_proj.* +25. layers.X.mlp.* -> layers.X.ffn.* +26. layers.X.ffn.experts.Y.gate_proj.* -> layers.X.ffn.experts.Y.w1.* +27. layers.X.ffn.experts.Y.up_proj.* -> layers.X.ffn.experts.Y.w3.* +28. layers.X.ffn.experts.Y.down_proj.* -> layers.X.ffn.experts.Y.w2.* +29. layers.X.ffn.shared_experts.gate_proj.* -> layers.X.ffn.shared_experts.w1.* +30. layers.X.ffn.shared_experts.up_proj.* -> layers.X.ffn.shared_experts.w3.* +31. layers.X.ffn.shared_experts.down_proj.* -> layers.X.ffn.shared_experts.w2.* +32. model.hc_head.hc_base -> hc_head_base (flatten structure) +33. model.hc_head.hc_fn -> hc_head_fn +34. model.hc_head.hc_scale -> hc_head_scale +35. layers.X.attn_hc.base -> layers.X.hc_attn_base (reorder) +36. layers.X.attn_hc.fn -> layers.X.hc_attn_fn +37. layers.X.attn_hc.scale -> layers.X.hc_attn_scale +38. layers.X.ffn_hc.base -> layers.X.hc_ffn_base (reorder) +39. layers.X.ffn_hc.fn -> layers.X.hc_ffn_fn +40. layers.X.ffn_hc.scale -> layers.X.hc_ffn_scale + +This script modifies safetensors files in-place without creating backups. +""" + +import argparse +import json +import tempfile +from pathlib import Path + +from safetensors import safe_open +from safetensors.torch import save_file + + +def convert_key(key: str) -> str | None: + """Convert a key from quantized format to BF16 format. + + Returns None if the key should be removed. + """ + # lm_head -> head + if key.startswith("lm_head"): + return key.replace("lm_head", "head") + + # model.embed_tokens.weight -> embed.weight + if key == "model.embed_tokens.weight": + return "embed.weight" + + # model.norm -> norm + if key.startswith("model.norm"): + return key.replace("model.norm", "norm") + + # model.hc_head.hc_base -> hc_head_base (flatten the structure) + if key.startswith("model.hc_head."): + # model.hc_head.hc_base -> hc_head_base + # model.hc_head.hc_fn -> hc_head_fn + # model.hc_head.hc_scale -> hc_head_scale + return key.replace("model.hc_head.hc_", "hc_head_") + + # Handle layers (strip 'model.' prefix if present) + if key.startswith("model.layers."): + key = key[6:] # Remove "model." + + # Process all layer keys (whether they started with model.layers or just layers) + if key.startswith("layers."): + # input_layernorm -> attn_norm + if ".input_layernorm." in key: + key = key.replace(".input_layernorm.", ".attn_norm.") + + # post_attention_layernorm -> ffn_norm + if ".post_attention_layernorm." in key: + key = key.replace(".post_attention_layernorm.", ".ffn_norm.") + + # self_attn -> attn + if ".self_attn." in key: + key = key.replace(".self_attn.", ".attn.") + + # After self_attn->attn conversion, fix nested compressor/indexer structure + # compressor.indexer.kv_norm -> indexer.compressor.norm + if ".attn.compressor.indexer.kv_norm." in key: + key = key.replace(".attn.compressor.indexer.kv_norm.", ".attn.indexer.compressor.norm.") + # compressor.kv_norm -> compressor.norm (only if not already handled above) + elif ".attn.compressor.kv_norm." in key: + key = key.replace(".attn.compressor.kv_norm.", ".attn.compressor.norm.") + + # q_a_norm -> q_norm + if ".q_a_norm." in key: + key = key.replace(".q_a_norm.", ".q_norm.") + + # attn_hc.base -> hc_attn_base (reorder) + # layers.0.attn_hc.base -> layers.0.hc_attn_base + if ".attn_hc." in key: + key = key.replace(".attn_hc.", ".hc_attn_") + + # ffn_hc.base -> hc_ffn_base (reorder) + # layers.0.ffn_hc.base -> layers.0.hc_ffn_base + if ".ffn_hc." in key: + key = key.replace(".ffn_hc.", ".hc_ffn_") + + # mlp -> ffn (all MLP references become FFN) + if ".mlp." in key: + key = key.replace(".mlp.", ".ffn.") + + # After mlp->ffn conversion, handle projections + if ".ffn.experts." in key or ".ffn.shared_experts." in key: + # gate_proj -> w1, up_proj -> w3, down_proj -> w2 + key = key.replace(".gate_proj.", ".w1.") + key = key.replace(".up_proj.", ".w3.") + key = key.replace(".down_proj.", ".w2.") + + # Attention projection name conversions (proj suffix -> w prefix) + # Must be done carefully to avoid conflicts - process in specific order + + # First handle nested indexer compressor paths (most specific first) + if ".attn.compressor.indexer.scorer.weights_proj" in key: + key = key.replace(".attn.compressor.indexer.scorer.weights_proj", ".attn.indexer.weights_proj") + if ".attn.compressor.indexer.gate_proj" in key: + key = key.replace(".attn.compressor.indexer.gate_proj", ".attn.indexer.compressor.wgate") + if ".attn.compressor.indexer.kv_proj" in key: + key = key.replace(".attn.compressor.indexer.kv_proj", ".attn.indexer.compressor.wkv") + if ".attn.compressor.indexer.position_bias" in key: + key = key.replace(".attn.compressor.indexer.position_bias", ".attn.indexer.compressor.ape") + if ".attn.compressor.indexer.q_b_proj" in key: + key = key.replace(".attn.compressor.indexer.q_b_proj", ".attn.indexer.wq_b") + + # Then handle compressor paths + if ".attn.compressor.gate_proj" in key: + key = key.replace(".attn.compressor.gate_proj", ".attn.compressor.wgate") + if ".attn.compressor.kv_proj" in key: + key = key.replace(".attn.compressor.kv_proj", ".attn.compressor.wkv") + if ".attn.compressor.position_bias" in key: + key = key.replace(".attn.compressor.position_bias", ".attn.compressor.ape") + + # Finally handle top-level attention projections + if ".attn.sinks" in key: + key = key.replace(".attn.sinks", ".attn.attn_sink") + if ".attn.kv_proj" in key: + key = key.replace(".attn.kv_proj", ".attn.wkv") + if ".attn.o_a_proj" in key: + key = key.replace(".attn.o_a_proj", ".attn.wo_a") + if ".attn.o_b_proj" in key: + key = key.replace(".attn.o_b_proj", ".attn.wo_b") + if ".attn.q_a_proj" in key: + key = key.replace(".attn.q_a_proj", ".attn.wq_a") + if ".attn.q_b_proj" in key: + key = key.replace(".attn.q_b_proj", ".attn.wq_b") + + return key + + +def convert_file(input_path: Path, output_path: Path): + """Convert a single safetensors file.""" + print(f"Processing {input_path.name}...") + + tensors = {} + with safe_open(str(input_path), framework="pt") as f: + for old_key in f.keys(): + new_key = convert_key(old_key) + if new_key is not None: + tensors[new_key] = f.get_tensor(old_key) + + save_file(tensors, str(output_path)) + print(f" Converted {len(tensors)} tensors") + + +def update_index_file(index_path: Path): + """Update the model.safetensors.index.json file with new key names.""" + print(f"Processing {index_path.name}...") + + with open(index_path) as f: + index_data = json.load(f) + + if "weight_map" in index_data: + new_weight_map = {} + for old_key, shard_name in index_data["weight_map"].items(): + new_key = convert_key(old_key) + if new_key is not None: + new_weight_map[new_key] = shard_name + + index_data["weight_map"] = new_weight_map + + # Write to temp file first, then move to avoid corruption + temp_path = index_path.with_suffix(".json.tmp") + with open(temp_path, "w") as f: + json.dump(index_data, f, indent=2) + temp_path.replace(index_path) + + print(f" Updated {len(index_data['weight_map'])} entries") + + +def main(): + parser = argparse.ArgumentParser( + description="Convert safetensors key names in-place" + ) + parser.add_argument( + "model_dir", + type=Path, + help="Directory containing safetensors files to convert" + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Show what would be done without modifying files" + ) + args = parser.parse_args() + + model_dir = args.model_dir + if not model_dir.exists(): + print(f"Error: Directory {model_dir} does not exist") + return 1 + + # Find all safetensors files + safetensor_files = sorted(model_dir.glob("*.safetensors")) + if not safetensor_files: + print(f"No safetensors files found in {model_dir}") + return 1 + + if args.dry_run: + print("DRY RUN - No files will be modified") + print(f"\nFound {len(safetensor_files)} files to convert:") + for f in safetensor_files: + print(f" {f.name}") + + # Show sample key conversions + print("\nSample key conversions from first file:") + first_file = safetensor_files[0] + with safe_open(str(first_file), framework="pt") as f: + for i, old_key in enumerate(f.keys()): + if i >= 10: + break + new_key = convert_key(old_key) + if new_key is None: + print(f" {old_key} -> [REMOVED]") + else: + print(f" {old_key} -> {new_key}") + return 0 + + print(f"Converting {len(safetensor_files)} files in {model_dir}") + print("WARNING: This will modify files in-place!") + response = input("Continue? [y/N] ") + if response.lower() != 'y': + print("Aborted") + return 0 + + # Process each safetensors file + for shard_file in safetensor_files: + # Use a temp file to avoid corruption if the process fails + temp_path = shard_file.with_suffix(".safetensors.tmp") + try: + convert_file(shard_file, temp_path) + # Replace original with converted file + temp_path.replace(shard_file) + except Exception as e: + print(f"Error processing {shard_file.name}: {e}") + if temp_path.exists(): + temp_path.unlink() + raise + + # Update index file if it exists + index_path = model_dir / "model.safetensors.index.json" + if index_path.exists(): + update_index_file(index_path) + + print("\nConversion complete!") + return 0 + + +if __name__ == "__main__": + exit(main()) diff --git a/examples/quantizing_moe/deepseek_v4_pro_example.py b/examples/quantizing_moe/deepseek_v4_pro_example.py new file mode 100644 index 0000000000..7207610078 --- /dev/null +++ b/examples/quantizing_moe/deepseek_v4_pro_example.py @@ -0,0 +1,144 @@ +import torch +from compressed_tensors.distributed import init_dist +from compressed_tensors.quantization.quant_scheme import ( + FP8_BLOCK, + NVFP4, + QuantizationScheme, +) +from datasets import load_dataset +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.models.deepseek_v4.modeling_deepseek_v4 import ( + DeepseekV4PreTrainedModel, +) + +from llmcompressor import oneshot +from llmcompressor.datasets.utils import get_rank_partition +from llmcompressor.modifiers.quantization import QuantizationModifier +from llmcompressor.utils import load_context + +# Upstream BUG: norms should be loaded in float32, but usually aren't due to the base +# model having a quant_config which overrides this. Loading in float32 actually +# breaks the model definition (it expects bfloat16). Let's force load in bfloat16. +DeepseekV4PreTrainedModel._keep_in_fp32_modules_strict = set() + +# Select model and load it. +MODEL_ID = "/mnt/nvme-data/engine/kylesayrs/DeepSeek-V4-Pro-BF16" + +init_dist() +with load_context(): + model = AutoModelForCausalLM.from_pretrained( + MODEL_ID, + device_map="auto_offload", + max_memory={}, + offload_folder="/mnt/nvme-data/engine/kylesayrs/offload_folder", + ) + +# kluge for the way I saved the decompressed checkpoint +# mds = model.model.layers[-1].self_attn.wq_a._hf_hook.weights_map.dataset.index +# mds["model.hc_head.base"] = mds['model.hc_head.hc_base'] +# mds["model.hc_head.fn"] = mds['model.hc_head.hc_fn'] +# mds["model.hc_head.scale"] = mds['model.hc_head.hc_scale'] + +tokenizer = AutoTokenizer.from_pretrained("RedHatAI/DeepSeek-V4-Flash-BF16") + +# Select calibration dataset. +DATASET_ID = "HuggingFaceH4/ultrachat_200k" +DATASET_SPLIT = "train_sft" + +# Select number of samples. 512 samples is a good place to start. +# Increasing the number of samples can improve accuracy. +NUM_CALIBRATION_SAMPLES = 512 +MAX_SEQUENCE_LENGTH = 2048 + +# Load dataset and preprocess. +ds = load_dataset( + DATASET_ID, split=get_rank_partition(DATASET_SPLIT, NUM_CALIBRATION_SAMPLES) +) +ds = ds.shuffle(seed=42) + + +def preprocess(example): + # DeepSeek-V4 does not have a traditional chat template. + # Encode manually per https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash/tree/main/encoding + BOS = "<｜begin▁of▁sentence｜>" + EOS = "<｜end▁of▁sentence｜>" + text = BOS + for message in example["messages"]: + role = message["role"] + content = message["content"] + if role == "system": + text += content + elif role == "user": + text += f"<｜User｜>{content}" + elif role == "assistant": + text += f"<｜Assistant｜>{content}{EOS}" + + return {"text": text} + + +ds = ds.map(preprocess) + + +def tokenize(sample): + return tokenizer( + sample["text"], + padding=False, + max_length=MAX_SEQUENCE_LENGTH, + truncation=True, + add_special_tokens=False, + ) + + +ds = ds.map(tokenize, remove_columns=ds.column_names) + +# Configure the quantization algorithm to run. +# * quantize mlp/expert weights to NVFP4 +# * quantize attention projection weights to FP8_BLOCK +# model.model.layers.0.self_attn.q_a_proj +# +# wq_a | q_a_proj +# wq_b | q_b_proj +# wkv | kv_proj +# wo_a | o_a_proj +# wo_b | o_b_proj + +recipe = QuantizationModifier( + config_groups={ +# "attention": QuantizationScheme( +# targets=[ +# r"re:.*attn\.(q_a_proj|q_b_proj|kv_proj|o_a_proj|o_b_proj)$", +# r"re:.*attn\.compressor\.indexer\.q_b_proj$", +# r"re:.*shared_experts.*", +# ], +# **FP8_BLOCK, +# ), + "experts": QuantizationScheme( + targets=[ + r"re:.*mlp\.experts\..*(gate|up|down)_proj$", + ], + **NVFP4, + ), + }, + ignore=[], +) + +# Apply algorithms. +# due to the large size of DeepSeek-V4, we specify sequential targets such that +# only one block is loaded into GPU memory at a time + +oneshot( + model=model, + processor=tokenizer, + dataset=ds, + recipe=recipe, + batch_size=4, + shuffle_calibration_samples=False, +) + +# Save to disk compressed. +# SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4-FP8-BLOCK" +SAVE_DIR = "/mnt/nvme-data/engine/kylesayrs/DeepSeek-V4-Pro-NVFP4" +model.save_pretrained(SAVE_DIR, save_compressed=True) +tokenizer.save_pretrained(SAVE_DIR) + +torch.distributed.destroy_process_group() diff --git a/examples/quantizing_moe/deepseek_v4_pro_gptq.py b/examples/quantizing_moe/deepseek_v4_pro_gptq.py new file mode 100644 index 0000000000..212c0a3799 --- /dev/null +++ b/examples/quantizing_moe/deepseek_v4_pro_gptq.py @@ -0,0 +1,156 @@ +from compressed_tensors.distributed import init_dist +from compressed_tensors.quantization.quant_scheme import ( + FP8_BLOCK, + NVFP4, + QuantizationScheme, +) +from datasets import load_dataset +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.models.deepseek_v4.modeling_deepseek_v4 import ( + DeepseekV4PreTrainedModel, +) + +from llmcompressor import oneshot +from llmcompressor.datasets.utils import get_rank_partition +from llmcompressor.modifiers.gptq import GPTQModifier +from llmcompressor.utils import load_context + +# Upstream BUG: norms should be loaded in float32, but usually aren't due to the base +# model having a quant_config which overrides this. Loading in float32 actually +# breaks the model definition (it expects bfloat16). Let's force load in bfloat16. +DeepseekV4PreTrainedModel._keep_in_fp32_modules_strict = set() + +# Select model and load it. +# MODEL_ID = "RedHatAI/DeepSeek-V4-Flash-BF16" +MODEL_ID = "/mnt/nvme-data/engine/kylesayrs/DeepSeek-V4-Pro-BF16" +# MODEL_ID = "RedHatAI/DeepSeek-V4-Flash-BF16" +# MODEL_ID = "inference-optimization/DSV4-tiny-empty" + +init_dist() +with load_context(): + model = AutoModelForCausalLM.from_pretrained( + MODEL_ID, + device_map="auto_offload", + max_memory={}, + offload_folder="/mnt/nvme-data/engine/kylesayrs/offload_folder", + ) + +# kluge for the way I saved the decompressed checkpoint +# mds = model.model.layers[-1].self_attn.wq_a._hf_hook.weights_map.dataset.index +# mds["model.hc_head.base"] = mds['model.hc_head.hc_base'] +# mds["model.hc_head.fn"] = mds['model.hc_head.hc_fn'] +# mds["model.hc_head.scale"] = mds['model.hc_head.hc_scale'] + +tokenizer = AutoTokenizer.from_pretrained("RedHatAI/DeepSeek-V4-Flash-BF16") + +# Select calibration dataset. +DATASET_ID = "HuggingFaceH4/ultrachat_200k" +DATASET_SPLIT = "train_sft" + +# Select number of samples. 512 samples is a good place to start. +# Increasing the number of samples can improve accuracy. +NUM_CALIBRATION_SAMPLES = 512 +MAX_SEQUENCE_LENGTH = 2048 + +# Load dataset and preprocess. +ds = load_dataset( + DATASET_ID, split=get_rank_partition(DATASET_SPLIT, NUM_CALIBRATION_SAMPLES) +) +ds = ds.shuffle(seed=42) + + +def preprocess(example): + # DeepSeek-V4 does not have a traditional chat template. + # Encode manually per https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash/tree/main/encoding + BOS = "<｜begin▁of▁sentence｜>" + EOS = "<｜end▁of▁sentence｜>" + text = BOS + for message in example["messages"]: + role = message["role"] + content = message["content"] + if role == "system": + text += content + elif role == "user": + text += f"<｜User｜>{content}" + elif role == "assistant": + text += f"<｜Assistant｜>{content}{EOS}" + + return {"text": text} + + +ds = ds.map(preprocess) + + +def tokenize(sample): + return tokenizer( + sample["text"], + padding=False, + max_length=MAX_SEQUENCE_LENGTH, + truncation=True, + add_special_tokens=False, + ) + + +ds = ds.map(tokenize, remove_columns=ds.column_names) + +# Configure the quantization algorithm to run. +# * quantize mlp/expert weights to NVFP4 +# * quantize attention projection weights to FP8_BLOCK +# model.model.layers.0.self_attn.q_a_proj +# +# wq_a | q_a_proj +# wq_b | q_b_proj +# wkv | kv_proj +# wo_a | o_a_proj +# wo_b | o_b_proj + +recipe = GPTQModifier( + config_groups={ + "attention": QuantizationScheme( + targets=[ + r"re:.*attn\.(q_a_proj|q_b_proj|kv_proj|o_a_proj|o_b_proj)$", + r"re:.*attn\.compressor\.indexer\.q_b_proj$", + ], + **FP8_BLOCK, + ), + "experts": QuantizationScheme( + targets=[ + r"re:.*mlp\..*(gate|up|down)_proj$", + ], + **NVFP4, + ), + }, + ignore=[], +) + +import torch + +# Apply algorithms. +# due to the large size of DeepSeek-V4, we specify sequential targets such that +# only one block is loaded into GPU memory at a time + +torch.cuda.memory._record_memory_history() +try: + oneshot( + model=model, + processor=tokenizer, + dataset=ds, + recipe=recipe, + batch_size=4, + shuffle_calibration_samples=False, + sequential_targets=["DeepseekV4Attention", "ExpertMLP"], + sequential_targets_per_subgraph=(384 // 4 + 10), + propagate_error=True, # work around reliance on transformers cache + # something weird happens with the cache and propagation + ) +finally: + if torch.distributed.get_rank() == 0: + torch.cuda.memory._dump_snapshot("memory_snapshot5.pickle") + +# Save to disk compressed. +# SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4-FP8-BLOCK" +SAVE_DIR = "/mnt/nvme-data/engine/kylesayrs/DeepSeek-V4-Pro-NVFP4-FP8-BLOCK-GPTQ" +model.save_pretrained(SAVE_DIR, save_compressed=True) +tokenizer.save_pretrained(SAVE_DIR) + +torch.distributed.destroy_process_group() diff --git a/src/llmcompressor/modeling/moe/granitemoe.py b/src/llmcompressor/modeling/moe/granitemoe.py deleted file mode 100644 index 3cc264d0d5..0000000000 --- a/src/llmcompressor/modeling/moe/granitemoe.py +++ /dev/null @@ -1,84 +0,0 @@ -import torch -from compressed_tensors.offload import get_cache_init_kwargs, offload_module -from transformers.models.granitemoe.configuration_granitemoe import GraniteMoeConfig -from transformers.models.granitemoe.modeling_granitemoe import GraniteMoeParallelExperts - -from llmcompressor.modeling.moe.context import get_calibrate_all_experts_flag -from llmcompressor.modeling.moe.linear_experts import LinearExperts2D -from llmcompressor.utils.dev import skip_weights_initialize - - -class GraniteMoeLinearExperts(LinearExperts2D): - is_concatenated = False - is_transposed = False - has_bias = False - has_gate = False - - @classmethod - @torch.no_grad() - def from_experts_module( - cls, experts: "GraniteMoeParallelExperts", config: GraniteMoeConfig - ): - assert experts.num_experts == config.num_local_experts - - with skip_weights_initialize(): - self = cls( - experts.num_experts, experts.input_size, experts.output_size, config - ) - self.num_experts = experts.num_experts - - # TODO: experiment with copying views, not values - for i in range(experts.num_experts): - self[i].weight.copy_(experts.weight[i]) - - # copy offloading from original - offload_kwargs = get_cache_init_kwargs(experts) - for module in self.modules(): - offload_module(module, **offload_kwargs) - - return self - - def __init__( - self, - num_experts: int, - input_size: int, - output_size: int, - config: GraniteMoeConfig, - ) -> None: - self.num_experts = num_experts - self.input_size = input_size - self.output_size = output_size - - torch.nn.ModuleList.__init__( - self, - [ - torch.nn.Linear(input_size, output_size, bias=False, dtype=config.dtype) - for _ in range(num_experts) - ], - ) - - def forward(self, inputs: torch.Tensor, expert_size: list[int]): - """ - Forward pass of the GraniteMoeParallelExperts module. - Args: - inputs (Tensor): - Input tensor. - expert_size: - Expert size information. - Returns: - Tensor: Output tensor. - """ - output_list = [] - - for i in range(self.num_experts): - if get_calibrate_all_experts_flag(): - expert_out = self[i](inputs).split(expert_size, dim=0)[i] - else: - expert_out = self[i](inputs.split(expert_size, dim=0)[i]) - output_list.append(expert_out) - - return torch.cat(output_list, dim=0) - - -# register in registry -LinearExperts2D._registry[GraniteMoeParallelExperts] = GraniteMoeLinearExperts diff --git a/src/llmcompressor/modeling/moe/helpers.py b/src/llmcompressor/modeling/moe/helpers.py index 41c21f56db..1a1719e2a2 100644 --- a/src/llmcompressor/modeling/moe/helpers.py +++ b/src/llmcompressor/modeling/moe/helpers.py @@ -8,6 +8,7 @@ ) import torch +from compressed_tensors.offload import disable_onloading from loguru import logger from transformers import PreTrainedConfig @@ -36,14 +37,17 @@ class FusedExpertsProtocol(TorchModuleProtocol): @classmethod def __validate__(cls, object: object) -> bool: - return ( - isinstance(getattr(object, "down_proj", None), torch.nn.Parameter) - and ( - isinstance(getattr(object, "up_proj", None), torch.nn.Parameter) - or isinstance(getattr(object, "gate_up_proj", None), torch.nn.Parameter) + with disable_onloading(): + return ( + isinstance(getattr(object, "down_proj", None), torch.nn.Parameter) + and ( + isinstance(getattr(object, "up_proj", None), torch.nn.Parameter) + or isinstance( + getattr(object, "gate_up_proj", None), torch.nn.Parameter + ) + ) + and get_use_experts_implementation_args(object.__class__) is not None ) - and get_use_experts_implementation_args(object.__class__) is not None - ) def get_use_experts_implementation_args(experts_cls: type) -> dict[str, bool] | None: diff --git a/src/llmcompressor/modeling/moe/linear_experts.py b/src/llmcompressor/modeling/moe/linear_experts.py index e28a7cdd44..390aea8cd4 100644 --- a/src/llmcompressor/modeling/moe/linear_experts.py +++ b/src/llmcompressor/modeling/moe/linear_experts.py @@ -161,7 +161,6 @@ class LinearExperts2D(torch.nn.ModuleList): def get_registration( cls, key: type[torch.nn.Module], default: Any = None ) -> type["LinearExperts2D"]: - from .granitemoe import GraniteMoeLinearExperts # noqa: F401 from .llama4 import Llama4LinearExperts # noqa: F401 return cls._registry.get(key, default) diff --git a/src/llmcompressor/modifiers/gptq/gptq_quantize.py b/src/llmcompressor/modifiers/gptq/gptq_quantize.py index 1003fd247a..ea6169419a 100644 --- a/src/llmcompressor/modifiers/gptq/gptq_quantize.py +++ b/src/llmcompressor/modifiers/gptq/gptq_quantize.py @@ -39,7 +39,7 @@ def accumulate_hessian( match module: case torch.nn.Linear() | transformers.Conv1D(): - if len(inp.shape) == 3: + if len(inp.shape) > 2: inp = inp.reshape((-1, inp.shape[-1])) inp = inp.t() case torch.nn.Conv2d(): diff --git a/src/llmcompressor/modifiers/pruning/reap/utils.py b/src/llmcompressor/modifiers/pruning/reap/utils.py index c46baecc1c..d208dca68c 100644 --- a/src/llmcompressor/modifiers/pruning/reap/utils.py +++ b/src/llmcompressor/modifiers/pruning/reap/utils.py @@ -13,7 +13,6 @@ from loguru import logger from llmcompressor.modeling.moe.context import get_calibrate_all_experts_flag -from llmcompressor.modeling.moe.granitemoe import GraniteMoeLinearExperts from llmcompressor.modeling.moe.linear_experts import ExpertMLP, LinearExperts2D from llmcompressor.modeling.moe.llama4 import Llama4LinearExperts @@ -150,18 +149,12 @@ def get_moe_attrs(model: nn.Module, ignore: list[str]) -> MoeModelAttrs | None: experts = getattr(module, experts_attr) # REAP currently only supports LinearExperts2D experts, as they receive the # top_k indices and weights from the router in their forward pass. - # Granite and Llama4 experts diverge from this behavior, so they are - # unsupported for now. + # Llama4 experts diverge from this behavior, so it is unsupported for now. if not isinstance(experts, LinearExperts2D): logger.warning( f"Skipping layer {name}: experts module is not LinearExperts2D" ) continue - if isinstance(experts, GraniteMoeLinearExperts): - logger.warning( - f"Skipping unsupported GraniteMoeLinearExperts layer: {name}" - ) - continue if isinstance(experts, Llama4LinearExperts): logger.warning( f"Skipping unsupported Llama4LinearExperts layer: {name}" @@ -173,10 +166,9 @@ def get_moe_attrs(model: nn.Module, ignore: list[str]) -> MoeModelAttrs | None: raise ValueError( "Could not find any supported MoE layers with experts in " "LinearExperts2D format. Make sure the model has MoE layers " - "(excluding GraniteMoeLinearExperts and Llama4LinearExperts), " - "and that the name of its experts module is in EXPERTS_ATTRS " - "and it the name of its router module is in ROUTER_ATTRS in " - "reap/utils.py" + "(excluding Llama4LinearExperts), and that the name of its " + "experts module is in EXPERTS_ATTRS and it the name of its " + "router module is in ROUTER_ATTRS in reap/utils.py" ) logger.info( diff --git a/src/llmcompressor/observers/base.py b/src/llmcompressor/observers/base.py index 0c146dadd9..05bff162c1 100644 --- a/src/llmcompressor/observers/base.py +++ b/src/llmcompressor/observers/base.py @@ -97,6 +97,8 @@ def get_qparams(self) -> QParamsDict: fused_obs(fused_mod.weight) global_absmax = torch.max(global_absmax, -fused_obs.min_vals.min()) global_absmax = torch.max(global_absmax, fused_obs.max_vals.max()) + del fused_obs.min_vals + del fused_obs.max_vals global_scale = generate_gparam( -global_absmax.reshape(1), global_absmax.reshape(1) ) @@ -108,6 +110,9 @@ def get_qparams(self) -> QParamsDict: global_scale=global_scale, ) + del self.min_vals + del self.max_vals + return {"scale": scale, "zero_point": zero_point, "global_scale": global_scale} @torch.no_grad diff --git a/tests/llmcompressor/modeling/test_linearize.py b/tests/llmcompressor/modeling/test_linearize.py index ee69f6e499..4524efe658 100644 --- a/tests/llmcompressor/modeling/test_linearize.py +++ b/tests/llmcompressor/modeling/test_linearize.py @@ -10,7 +10,7 @@ from transformers.models.afmoe.configuration_afmoe import AfmoeConfig from transformers.models.afmoe.modeling_afmoe import AfmoeExperts from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config -from transformers.models.deepseek_v3.modeling_deepseek_v3 import DeepseekV3NaiveMoe +from transformers.models.deepseek_v3.modeling_deepseek_v3 import DeepseekV3Experts from transformers.models.deepseek_v4.configuration_deepseek_v4 import DeepseekV4Config from transformers.models.deepseek_v4.modeling_deepseek_v4 import ( DeepseekV4Experts, @@ -19,17 +19,17 @@ from transformers.models.gemma4.configuration_gemma4 import Gemma4TextConfig from transformers.models.gemma4.modeling_gemma4 import Gemma4TextExperts from transformers.models.glm4_moe.configuration_glm4_moe import Glm4MoeConfig -from transformers.models.glm4_moe.modeling_glm4_moe import Glm4MoeNaiveMoe +from transformers.models.glm4_moe.modeling_glm4_moe import Glm4MoeExperts from transformers.models.glm4_moe_lite.configuration_glm4_moe_lite import ( Glm4MoeLiteConfig, ) -from transformers.models.glm4_moe_lite.modeling_glm4_moe_lite import Glm4MoeLiteNaiveMoe +from transformers.models.glm4_moe_lite.modeling_glm4_moe_lite import Glm4MoeLiteExperts from transformers.models.glm_moe_dsa.configuration_glm_moe_dsa import GlmMoeDsaConfig -from transformers.models.glm_moe_dsa.modeling_glm_moe_dsa import GlmMoeDsaNaiveMoe +from transformers.models.glm_moe_dsa.modeling_glm_moe_dsa import GlmMoeDsaExperts from transformers.models.gpt_oss.configuration_gpt_oss import GptOssConfig from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts from transformers.models.granitemoe.configuration_granitemoe import GraniteMoeConfig -from transformers.models.granitemoe.modeling_granitemoe import GraniteMoeParallelExperts +from transformers.models.granitemoe.modeling_granitemoe import GraniteMoeExperts from transformers.models.hy_v3.configuration_hy_v3 import HYV3Config from transformers.models.hy_v3.modeling_hy_v3 import HYV3Experts from transformers.models.llama4.configuration_llama4 import ( @@ -167,27 +167,40 @@ def forward(self, *args, **kwargs): @pytest.mark.parametrize( "config_cls,experts_cls,kwargs", [ - (AfmoeConfig, AfmoeExperts, {}), + ( + AfmoeConfig, + AfmoeExperts, + {"hidden_size": 512, "intermediate_size": 1024, "num_experts": 16}, + ), ( DeepseekV3Config, - DeepseekV3NaiveMoe, - {"hidden_size": 512, "moe_intermediate_size": 1024}, + DeepseekV3Experts, + {"hidden_size": 512, "moe_intermediate_size": 1024, "n_routed_experts": 16}, + ), + ( + DeepseekV4Config, + DeepseekV4Experts, + {"hidden_size": 512, "moe_intermediate_size": 1024, "n_routed_experts": 16}, ), - (DeepseekV4Config, DeepseekV4Experts, {}), ( Gemma4TextConfig, Gemma4TextExperts, {"num_experts": 16, "top_k_experts": 4, "moe_intermediate_size": 2304}, ), - (Glm4MoeConfig, Glm4MoeNaiveMoe, {}), - (Glm4MoeLiteConfig, Glm4MoeLiteNaiveMoe, {}), - (GlmMoeDsaConfig, GlmMoeDsaNaiveMoe, {"hidden_size": 512}), + (Glm4MoeConfig, Glm4MoeExperts, {}), + (Glm4MoeLiteConfig, Glm4MoeLiteExperts, {}), + (GlmMoeDsaConfig, GlmMoeDsaExperts, {"hidden_size": 512}), + ( + GraniteMoeConfig, + GraniteMoeExperts, + {"hidden_size": 512, "intermediate_size": 1024, "num_local_experts": 4}, + ), (Qwen3_5MoeTextConfig, Qwen3_5MoeExperts, {}), (Qwen3MoeConfig, Qwen3MoeExperts, {}), (Qwen3NextConfig, Qwen3NextExperts, {}), (Qwen3VLMoeTextConfig, Qwen3VLMoeTextExperts, {}), (GptOssConfig, GptOssExperts, {}), - (HYV3Config, HYV3Experts, {}), + (HYV3Config, HYV3Experts, {"hidden_size": 512, "moe_intermediate_size": 1024}), ( NemotronHConfig, NemotronHExperts, @@ -230,33 +243,6 @@ def test_linearize_moe(config_cls, experts_cls, kwargs): assert torch.nn.functional.mse_loss(calib_outputs, true_outputs) < MODULE_MSE -def test_linearize_moe_granite(): - config = GraniteMoeConfig(hidden_size=512, intermediate_size=1024) - experts = GraniteMoeParallelExperts( - config.num_local_experts, config.hidden_size, config.intermediate_size - ) - init.normal_(experts.weight, mean=0.0, std=config.initializer_range) - - mock_model = DummyModel(experts, config) - linearize_moe(mock_model) - assert mock_model.module is not experts - - hidden_states = torch.randn(NUM_TEST_TOKENS, config.hidden_size, dtype=config.dtype) - expert_size = [ - (NUM_TEST_TOKENS // config.num_local_experts) - for _ in range(config.num_local_experts) - ] - expert_size[-1] += NUM_TEST_TOKENS % config.num_local_experts - true_outputs = experts(hidden_states, expert_size) - outputs = mock_model(hidden_states, expert_size) - with moe_calibration_context(): - calib_outputs = mock_model(hidden_states, expert_size) - - assert torch.any(true_outputs != 0), "Bad test setup, output is all zeros" - assert torch.nn.functional.mse_loss(outputs, true_outputs) < MODULE_MSE - assert torch.nn.functional.mse_loss(calib_outputs, true_outputs) < MODULE_MSE - - def test_linearize_moe_llama4(): text_config = Llama4TextConfig(hidden_size=512, intermediate_size=1024) config = Llama4Config(text_config=text_config)