vLLM serving: patched deepseek_v4.py, disabled mega_moe, updated docs
- Add patches/deepseek_v4.py: patched vllm source file with modelopt NVFP4 weight name mappings (expert gate_proj→w1, mlp→ffn, self_attn→attn.mla_attn, compressor.kv_proj→wkv, etc.), E2M1 FP4→BF16 unpacking for stacked params, skip patterns for NVFP4 scale tensors on MergedColumnParallelLinear, and resilient loading for unknown params. - Update docker-compose.yml: copy patched deepseek_v4.py over original at container startup, remove --moe-backend=deep_gemm_mega_moe (no NVFP4 kernel). - Update patches/patch_vllm_weights.py: legacy runtime monkey-patch approach (doesn't work with worker processes), kept for reference. - Update README.md: added vLLM serving run history table (S1-S10), documented all open issues (MergedColumnParallelLinear+NVFP4, no mega_moe kernel, resilient loading), added vLLM-specific bug list and key notes. - Update scripts/serve_vllm.py: add WARN comment on mega_moe flag.
This commit is contained in:
127
README.md
127
README.md
@@ -1,10 +1,10 @@
|
||||
# DeepSeek V4 Pro → NVFP4 Quantization
|
||||
# DeepSeek V4 Pro → NVFP4 Quantization + vLLM Serving
|
||||
|
||||
Full NVFP4 quantization of DeepSeek V4 Pro on a single B200 node (8× B200, 2.7TB RAM, 13TB NVMe). **Result: 881GB NVFP4 (Run 11).**
|
||||
Full NVFP4 quantization of DeepSeek V4 Pro on a single B200 node (8× B200, 2.7TB RAM, 13TB NVMe). **Result: 881GB NVFP4 (Run 11).** Now working on vLLM serving of the quantized checkpoint.
|
||||
|
||||
**Cost:** ~$161/run at $23/hr (7 hours each). Don't waste runs.
|
||||
|
||||
## ✅ Final Result (Run 11)
|
||||
## ✅ Final Quantization Result (Run 11)
|
||||
|
||||
- **Output:** `/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4` — 881GB, 95 safetensors
|
||||
- **Config:** `nvfp4` full quantization, 128 calib samples, `kv_cache_qformat=fp8_cast`
|
||||
@@ -14,6 +14,79 @@ Full NVFP4 quantization of DeepSeek V4 Pro on a single B200 node (8× B200, 2.7T
|
||||
- **Calibrated state:** 721.4GB (insurance, can re-export with `--export-only`)
|
||||
- A few experts (11, 83, 100, 112, 254) had uncalibrated amax — weight-derived fallback used (normal for sparse MoE with 256 experts)
|
||||
|
||||
---
|
||||
|
||||
## 🔧 vLLM Serving (In Progress)
|
||||
|
||||
### Current Status: Debugging weight loading
|
||||
|
||||
The modelopt NVFP4 export and vllm have a chain of incompatibilities. We're progressively fixing them. The fundamental problem: **modelopt's NVFP4 quantization format and vllm's DeepSeek V4 serving code were never integrated.** NVIDIA's own published NVFP4 exports (DeepSeek-V3.2, MiniMax-M2.7) don't have these issues because they don't use MLA attention compression or 256-expert MoE — both of which create stacked/fused weight parameters that modelopt doesn't account for.
|
||||
|
||||
### Approach: Patched deepseek_v4.py + disabled mega_moe
|
||||
|
||||
Instead of runtime monkey-patching (which doesn't propagate to worker processes), we patch the vllm source file directly. The patched `deepseek_v4.py` is mounted into the Docker container and copied over the original before vllm starts.
|
||||
|
||||
We also disabled `--moe-backend=deep_gemm_mega_moe` because:
|
||||
1. The NVFP4 mega_moe kernel doesn't exist yet (NVIDIA hasn't built it)
|
||||
2. MegaMoE uses MXFP4 block scale format (32-col blocks), but modelopt exports NVFP4 (16-col blocks) — format mismatch
|
||||
3. MegaMoE doesn't register `weight_scale_2` or `input_scale` params, so those scales would be silently dropped
|
||||
|
||||
Instead, we use the standard FusedMoE path with `ModelOptNvFp4FusedMoE`, which natively supports NVFP4 expert weights.
|
||||
|
||||
### vLLM Serving Run History
|
||||
|
||||
| # | Date | Approach | Result | Root Cause | Fix/Next |
|
||||
|---|------|----------|--------|------------|----------|
|
||||
| S1 | May 10 09:34 | `patch_vllm_weights.py` runtime patch + mega_moe | ❌ `UnboundLocalError: name_mapped` | Expert weight names don't match any mapping → `name_mapped` never assigned | Add gate_proj→w1, up_proj→w3, down_proj→w2 mappings |
|
||||
| S2 | May 10 ~10:30 | Same, added expert rename regexes | ❌ Same error | `DeepseekV4ForCausalLM.hf_to_vllm_mapper` is a **class attribute** set at import time — patching the function doesn't update the cached mapper | Patch the class attribute directly |
|
||||
| S3 | May 10 ~11:00 | Patched class attr, but workers are separate processes | ❌ Same error in workers | Workers don't inherit in-memory patches — they fork before the patch runs | Patch the source file directly (`deepseek_v4.py`) |
|
||||
| S4 | May 10 ~11:30 | Direct source file patch + mega_moe | ❌ `KeyError: 'layers.0.mlp.experts.0.w2.weight'` | modelopt uses `mlp`, vllm uses `ffn` internally — missing `.mlp.` → `.ffn.` mapping | Add substr mapping |
|
||||
| S5 | May 10 ~12:00 | Added `mlp→ffn` mapping + mega_moe | ❌ `KeyError: 'fused_wkv_wgate.input_scale'` | Compressor fused params don't register `input_scale`/`weight_scale_2` | Add skip patterns for compressor/attention scale tensors |
|
||||
| S6 | May 10 ~12:30 | Added skip patterns + mega_moe | ❌ Shape mismatch: `w2_weight_scale (7168, 96) vs (7168, 192)` | NVFP4 uses 16-col block scales, mega_moe expects 32-col MXFP4 — format incompatibility | **Abandon mega_moe** — no NVFP4 mega_moe kernel exists |
|
||||
| S7 | May 10 ~13:00 | Disabled mega_moe, standard FusedMoE | ❌ `fused_wkv_wgate.weight` shape mismatch: param=(1024,7168) bf16, loaded=(512,3584) uint8 | `MergedColumnParallelLinear` creates weight as bf16 (not uint8), but modelopt exports NVFP4 packed uint8. `ModelOptNvFp4Config` only handles `Linear`, not `MergedColumnParallelLinear` | Unpack uint8→bf16 at load time |
|
||||
| S8 | May 10 ~13:30 | Added E2M1 unpacking for fused weights | ❌ `KeyError: 'fused_wkv_wgate.weight_scale'` | No `weight_scale` param registered for `MergedColumnParallelLinear` (same `ModelOptNvFp4Config` gap) | Skip all NVFP4 scale tensors for stacked/fused attention+compressor params |
|
||||
| S9 | May 10 ~14:00 | Added weight_scale skip patterns | ❌ `KeyError: 'compressor.kv_norm.weight'` | modelopt puts `kv_norm` under `compressor`, vllm has it at attention level (`attn.kv_norm`) | Add `compressor.kv_norm` → `kv_norm` mapping |
|
||||
| S10 | May 10 ~14:15 | Fixed kv_norm mapping | ❌ `KeyError: 'compressor.position_bias'` | modelopt exports params that don't exist in the vllm model | Make loading resilient to unknown params |
|
||||
|
||||
### Open Issues (as of May 10 ~16:00 UTC)
|
||||
|
||||
1. **MergedColumnParallelLinear + NVFP4 incompatibility** — The core problem. `ModelOptNvFp4Config.create_weights()` only handles `Linear` layers. For `MergedColumnParallelLinear` (used for `fused_wqa_wkv`, `fused_wkv_wgate`, `gate_up_proj`):
|
||||
- Weight param is created as `ModelWeightParameter` (bf16) instead of `PackedColumnParameter` (uint8)
|
||||
- `weight_scale`, `weight_scale_2`, `input_scale` params are never registered
|
||||
- `adjust_shard_indexes_for_packing` applies `packed_factor` to rows, but NVFP4 packs along columns
|
||||
- Current workaround: unpack uint8→bf16 at load time, skip scale tensors, let `process_weights_after_loading` re-quantize. This loses the calibration-optimized scales for attention/compressor/shared_expert weights.
|
||||
|
||||
2. **No NVFP4 mega_moe kernel** — We disabled mega_moe to avoid the format mismatch. Standard FusedMoE with `ModelOptNvFp4FusedMoE` works for expert weights, but loses the mega_moe performance optimization. When NVIDIA builds an NVFP4 mega_moe kernel, we can re-enable it.
|
||||
|
||||
3. **Resilient loading needed** — modelopt exports params (e.g., `compressor.position_bias`) that don't exist in the vllm model. Need to skip unknown params gracefully instead of crashing.
|
||||
|
||||
4. **Expert `weight_scale_2` handling with FusedMoE** — The standard FusedMoE path registers `w13_weight_scale_2` and `w2_weight_scale_2`, so expert global scales CAN be loaded. This works for experts. The issue is only with the stacked/fused attention params.
|
||||
|
||||
### What Each Patch Does
|
||||
|
||||
**`patches/deepseek_v4.py`** — Patched vllm source file, copied over the original at container startup. Contains:
|
||||
- **Regex mappings** (applied first by WeightsMapper):
|
||||
- Skip `weight_scale`, `weight_scale_2`, `input_scale` for compressor/attention fused params (no stacked param registered)
|
||||
- Skip `weight_scale`, `weight_scale_2`, `input_scale` for shared expert gate/up projections (stacked into `gate_up_proj`)
|
||||
- Expert projection rename: `gate_proj→w1`, `up_proj→w3`, `down_proj→w2` (only for `.experts.N.`, not `.shared_experts.`)
|
||||
- **Substr mappings** (applied after regex):
|
||||
- Attention: `self_attn→attn.mla_attn` with proper sub-projection names
|
||||
- `kv_norm` moved from compressor to attention level
|
||||
- `compressor.kv_proj→compressor.wkv`, `compressor.gate_proj→compressor.wgate`
|
||||
- `shared_experts.gate_proj→shared_experts.w1`, `shared_experts.up_proj→shared_experts.w3`
|
||||
- `.mlp.→.ffn.` (modelopt uses `mlp`, vllm uses `ffn`)
|
||||
- **E2M1 FP4→BF16 unpacking** for stacked params: When a uint8 packed NVFP4 weight is loaded into a bf16 param (MergedColumnParallelLinear), unpack using the E2M1 lookup table
|
||||
- **Resilient loading**: Skip unknown params that modelopt exports but vllm doesn't have
|
||||
|
||||
**`patches/patch_vllm_weights.py`** — Legacy runtime monkey-patch approach. Doesn't work because vllm workers are separate processes that don't inherit in-memory patches. Kept for reference.
|
||||
|
||||
**`docker-compose.yml`** — Docker Compose config:
|
||||
- Copies patched `deepseek_v4.py` before starting vllm
|
||||
- Removed `--moe-backend=deep_gemm_mega_moe` (no NVFP4 kernel exists)
|
||||
- All other vllm flags are critical for V4 (see `serve_vllm.py` for documentation)
|
||||
|
||||
---
|
||||
|
||||
## ⚠️ Model Config Patches (post-export)
|
||||
|
||||
modelopt 0.45.0.dev64's export produces configs that don't match what vllm expects at runtime. **NVIDIA's own published NVFP4 exports have the same gaps** — we compared against `nvidia/DeepSeek-V3.2-NVFP4` and `nvidia/MiniMax-M2.7-NVFP4` on HuggingFace. Neither includes `compress_ratios` or `scale_fmt` either. This is a modelopt ↔ vllm integration gap, not a problem with our quantization.
|
||||
@@ -94,7 +167,23 @@ python3 /root/nvidia-meeting/deepseek-v4-quant/scripts/quantize_nvfp4.py --valid
|
||||
|
||||
**Runtime:** Model loading ~50 min. Calibration ~5.5 hours. Export ~30-60 min. Total 7-8 hours.
|
||||
|
||||
## Run History
|
||||
### Step 3: Serve with vLLM
|
||||
|
||||
```bash
|
||||
cd /root/nvidia-meeting
|
||||
docker compose up -d
|
||||
```
|
||||
|
||||
Or without Docker:
|
||||
|
||||
```bash
|
||||
source /root/nvidia-meeting/venv/bin/activate
|
||||
python3 /root/nvidia-meeting/deepseek-v4-quant/scripts/serve_vllm.py
|
||||
```
|
||||
|
||||
**Note:** `serve_vllm.py` still references `--moe-backend=deep_gemm_mega_moe`. This needs to be removed when mega_moe support is ready. For now, use the Docker Compose setup which has it removed.
|
||||
|
||||
## Quantization Run History
|
||||
|
||||
| Run | Date | Commit | Calib | Result | Root Cause | Fix |
|
||||
|-----|------|--------|-------|--------|------------|-----|
|
||||
@@ -110,7 +199,7 @@ python3 /root/nvidia-meeting/deepseek-v4-quant/scripts/quantize_nvfp4.py --valid
|
||||
| 10 | May 9 ~15:30 | `5a72da7` | 128 | ❌ Export crash (calib ✅) | `get_weight_scaling_factor` reads stale GPU weight → `cudaErrorIllegalAddress` | Patch `_export_quantized_weight` to force weight to CPU at entry point |
|
||||
| 11 | May 9 ~22:50 | `07cd50e` | 128 | ✅ **SUCCESS** | — | 8 patches covering full export chain |
|
||||
|
||||
### Key Lessons
|
||||
### Key Lessons (Quantization)
|
||||
|
||||
**Run 2 — Stale GPU tensors:** `use_seq_device_map` shuffles layers through GPU for calibration. Quantizer amax tensors sit in VRAM for 5+ hours while CUDA's allocator churns memory. By export time, the GPU tensor metadata is valid but the underlying memory has been recycled — reading it triggers `cudaErrorIllegalAddress`. Fix: copy amax to CPU immediately after calibration.
|
||||
|
||||
@@ -133,6 +222,7 @@ python3 /root/nvidia-meeting/deepseek-v4-quant/scripts/quantize_nvfp4.py --valid
|
||||
- Don't skip the `__main__` post-parse conversions — `calib_size` must be int list, `dataset` must be list (Run 9)
|
||||
- Don't use shell script arg names (`--quant`, `--calib`, `--kv_cache_quant`, `--tp`) — use `hf_ptq.py` names (`--qformat`, `--calib_size`, `--kv_cache_qformat`, `--inference_tensor_parallel`)
|
||||
- Don't patch individual export functions one at a time — patch the entry point (`_export_quantized_weight`) so weight is on CPU for the entire chain (Run 10)
|
||||
- Don't use runtime monkey-patching for vllm serving — workers are separate processes that don't inherit patches. Patch the source file directly instead.
|
||||
|
||||
## Runtime Patches Applied by quantize_nvfp4.py
|
||||
|
||||
@@ -173,6 +263,24 @@ Patches 6-8 are belt-and-suspenders. Patch 4 is the one that matters — it move
|
||||
7. **Model loading OOM during expert weight conversion.** `AutoModelForCausalLM.from_pretrained` does `torch.cat` on GPU for expert `gate_up_proj` (31.5GB alloc), but only 25.9GB free with `device_map="sequential"`. Fixed by using modelopt's `get_model()` which sets `max_memory` per GPU before loading.
|
||||
8. **Export crash — stale GPU weight tensors in `get_weight_scaling_factor`.** Patches 1-3 only covered quantizer amax. The model weights themselves are also on stale GPU. `weight_scaling_factor_2.to(weight.device)` triggers `cudaErrorIllegalAddress`. Fixed by patching `_export_quantized_weight` to force weight to CPU at the entry point, covering the entire export chain.
|
||||
|
||||
### Bugs Found (V4 NVFP4 + vLLM serving)
|
||||
|
||||
1. **modelopt uses `mlp`, vllm uses `ffn`** — Module naming mismatch. Fixed with substr mapping.
|
||||
2. **modelopt uses `gate_proj`/`up_proj`/`down_proj`, vllm expects `w1`/`w3`/`w2`** — Expert weight naming mismatch. Fixed with regex mapping (only for `.experts.N.`, not `.shared_experts.`).
|
||||
3. **modelopt uses `self_attn` prefix, vllm uses `attn.mla_attn`** — Attention module naming. Fixed with substr mapping.
|
||||
4. **`kv_proj` maps to `wkv`, not `kv_proj`** — vllm stacks `wkv` + `wq_a` into `fused_wqa_wkv`. Fixed with substr mapping.
|
||||
5. **`compressor.kv_proj` → `compressor.wkv`** — Similar stacking for compressor. Fixed with substr mapping.
|
||||
6. **`compressor.kv_norm` → `attn.kv_norm`** — modelopt puts `kv_norm` under compressor, vllm has it at attention level. Fixed with substr mapping (must come before general compressor mapping).
|
||||
7. **`MergedColumnParallelLinear` + NVFP4 incompatibility** — `ModelOptNvFp4Config.create_weights()` only handles `Linear`, not `MergedColumnParallelLinear`. This causes:
|
||||
- Weight param created as bf16 instead of uint8 (PackedColumnParameter)
|
||||
- `weight_scale`/`weight_scale_2`/`input_scale` not registered for stacked params
|
||||
- `adjust_shard_indexes_for_packing` applies packed_factor to rows, but NVFP4 packs along columns
|
||||
- **Workaround:** Unpack uint8→bf16 at load time, skip scale tensors, rely on `process_weights_after_loading` re-quantization
|
||||
8. **No NVFP4 mega_moe kernel** — `DeepseekV4MegaMoEExperts` expects MXFP4 (32-col blocks), modelopt exports NVFP4 (16-col blocks). No kernel exists. **Abandoned mega_moe**, using standard FusedMoE instead.
|
||||
9. **`DeepseekV4ForCausalLM.hf_to_vllm_mapper` is a class attribute** — Runtime monkey-patching the factory function doesn't update the cached class attribute. Must patch the source file directly or update the class attribute explicitly.
|
||||
10. **vllm workers are separate processes** — In-memory monkey-patches don't propagate to workers. Must patch the source file directly.
|
||||
11. **modelopt exports params vllm doesn't have** — e.g., `compressor.position_bias`. Need resilient loading that skips unknown params.
|
||||
|
||||
## Dependencies (pinned versions)
|
||||
|
||||
- **nvidia-modelopt:** `0.45.0.dev64+g579fc6c31` (installed from git, not PyPI)
|
||||
@@ -191,6 +299,8 @@ The patches in `quantize_nvfp4.py` are for **modelopt 0.45.0.dev64** specificall
|
||||
- The amax snapshot (`v4_nvfp4_amax_snapshots.pt`) is ~50MB. Small, critical, cheap insurance.
|
||||
- The script calls `hf_main(args)` — the exact same entry point as the shell script. No pipeline divergence.
|
||||
- Must run from `/root/nvidia-meeting/modelopt-repo/examples/llm_ptq` (relative imports).
|
||||
- For vllm serving, the patched `deepseek_v4.py` must be mounted into the container — workers don't inherit in-memory patches.
|
||||
- We disabled `--moe-backend=deep_gemm_mega_moe` because no NVFP4 mega_moe kernel exists yet. Standard FusedMoE with `ModelOptNvFp4FusedMoE` handles expert weights correctly.
|
||||
|
||||
## File Layout
|
||||
|
||||
@@ -198,10 +308,15 @@ The patches in `quantize_nvfp4.py` are for **modelopt 0.45.0.dev64** specificall
|
||||
scripts/
|
||||
dequant_fp8_to_bf16.py — Step 1: FP8/FP4 → BF16 dequantization
|
||||
quantize_nvfp4.py — Step 2: NVFP4 quantization (patches + hf_main)
|
||||
serve_vllm.py — Step 3: vLLM serving (legacy, still has mega_moe flag)
|
||||
|
||||
patches/
|
||||
deepseek_v4.py — Patched vllm source file (copied over original at container startup)
|
||||
patch_vllm_weights.py — Legacy runtime monkey-patch (doesn't work with workers, kept for reference)
|
||||
quant_module_patched.py — (legacy) quant module patches
|
||||
patch_finegrained_fp8_blackwell.py — (legacy) FP8 kernel patches for Blackwell
|
||||
quant_module_patched.py — (legacy) quant module patches
|
||||
|
||||
docker-compose.yml — Docker Compose config for serving (uses patched deepseek_v4.py, no mega_moe)
|
||||
```
|
||||
|
||||
The `patches/` directory contains earlier approaches that modified modelopt source files directly. The current approach (`quantize_nvfp4.py`) uses runtime monkey-patching instead — no source files are modified.
|
||||
|
||||
@@ -6,7 +6,7 @@ services:
|
||||
- bash
|
||||
- -c
|
||||
- |
|
||||
python3 /patches/patch_vllm_weights.py
|
||||
cp /patches/deepseek_v4.py /usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/deepseek_v4.py
|
||||
exec vllm serve "$$@"
|
||||
- --
|
||||
environment:
|
||||
@@ -20,7 +20,6 @@ services:
|
||||
- --tensor-parallel-size=8
|
||||
- --compilation-config={"cudagraph_mode":"FULL_AND_PIECEWISE","custom_ops":["all"]}
|
||||
- --attention_config.use_fp4_indexer_cache=True
|
||||
- --moe-backend=deep_gemm_mega_moe
|
||||
- --tokenizer-mode=deepseek_v4
|
||||
- --tool-call-parser=deepseek_v4
|
||||
- --enable-auto-tool-choice
|
||||
@@ -42,5 +41,6 @@ services:
|
||||
stdin_open: true
|
||||
volumes:
|
||||
- /root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4:/model:ro
|
||||
- /root/nvidia-meeting/deepseek-v4-quant/patches/deepseek_v4.py:/patches/deepseek_v4.py:ro
|
||||
- /root/nvidia-meeting/deepseek-v4-quant/patches:/patches:ro
|
||||
network_mode: host
|
||||
|
||||
1719
patches/deepseek_v4.py
Normal file
1719
patches/deepseek_v4.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -2,11 +2,19 @@
|
||||
"""
|
||||
Patch vllm's DeepSeek V4 weight mapper to handle modelopt's NVFP4 export naming.
|
||||
|
||||
modelopt exports weights with `self_attn` prefix and other naming differences
|
||||
that vllm's _make_deepseek_v4_weights_mapper doesn't account for.
|
||||
modelopt exports weights with naming differences from what vllm's
|
||||
_make_deepseek_v4_weights_mapper + load_weights code expects:
|
||||
|
||||
This patch adds the missing substring mappings so modelopt-exported NVFP4
|
||||
checkpoints load correctly.
|
||||
1. Expert projections: modelopt uses gate_proj/up_proj/down_proj, vllm expects w1/w3/w2
|
||||
2. Shared expert projections: same gate_proj/up_proj naming, needs w1/w3 for stacking
|
||||
3. Compressor projections: kv_proj→wkv, gate_proj→wgate for fused stacking
|
||||
4. Attention projections: self_attn prefix, kv_proj→wkv for fused stacking, etc.
|
||||
5. Expert NVFP4 scales: weight_scale_2 and input_scale have no matching mega_moe params
|
||||
|
||||
CRITICAL: DeepseekV4ForCausalLM.hf_to_vllm_mapper is a CLASS attribute set at
|
||||
module import time. Simply patching _make_deepseek_v4_weights_mapper doesn't help
|
||||
because the class already cached the old mapper. We must also update the class
|
||||
attribute directly. Since expert_dtype=="fp4", __init__ doesn't recreate the mapper.
|
||||
|
||||
Drop into container as:
|
||||
python3 /patches/patch_vllm_weights.py
|
||||
@@ -14,50 +22,114 @@ Drop into container as:
|
||||
Or add to docker-compose.yml command before vllm serve.
|
||||
"""
|
||||
|
||||
import re
|
||||
import sys
|
||||
import importlib
|
||||
|
||||
|
||||
# Save original function BEFORE patching
|
||||
_original_make_mapper = None
|
||||
|
||||
|
||||
def make_patched_mapper(expert_dtype: str):
|
||||
"""Create a WeightsMapper with modelopt NVFP4 naming patches applied."""
|
||||
global _original_make_mapper
|
||||
# Use the saved original, not the (possibly patched) module attribute
|
||||
mapper = _original_make_mapper(expert_dtype)
|
||||
|
||||
# ══════════════════════════════════════════════════════════════════
|
||||
# Regex mappings (applied FIRST by WeightsMapper, before substr)
|
||||
# Order matters: skip patterns must come before rename patterns.
|
||||
# ══════════════════════════════════════════════════════════════════
|
||||
ordered_regexes = {}
|
||||
|
||||
# Skip expert NVFP4 scales that have no mega_moe params.
|
||||
# MUST come before gate_proj→w1 etc. because after renaming,
|
||||
# the key has "w1." not "gate_proj." and these patterns wouldn't match.
|
||||
#
|
||||
# modelopt's NVFP4 export includes weight_scale_2 (global scale) and
|
||||
# input_scale (activation scale) for each expert projection. But the
|
||||
# DeepseekV4MegaMoEExperts module only registers w13_weight_scale and
|
||||
# w2_weight_scale (E8M0 block scales) — no weight_scale_2 or input_scale.
|
||||
# Mapping to None tells WeightsMapper to skip these weights entirely.
|
||||
ordered_regexes[re.compile(r"\.experts\.\d+\.\w+_proj\.weight_scale_2$")] = None
|
||||
ordered_regexes[re.compile(r"\.experts\.\d+\.\w+_proj\.input_scale$")] = None
|
||||
|
||||
# Routed expert projections: gate_proj→w1, up_proj→w3, down_proj→w2
|
||||
# We use regex (not substr) to match ONLY .experts.N. — NOT .shared_experts.
|
||||
# Using substr ".down_proj." → ".w2." would also affect
|
||||
# shared_experts.down_proj, breaking shared expert loading
|
||||
# (vllm model uses down_proj, not w2, for shared experts).
|
||||
ordered_regexes[re.compile(r"(\.experts\.\d+\.)gate_proj\.")] = r"\1w1."
|
||||
ordered_regexes[re.compile(r"(\.experts\.\d+\.)up_proj\.")] = r"\1w3."
|
||||
ordered_regexes[re.compile(r"(\.experts\.\d+\.)down_proj\.")] = r"\1w2."
|
||||
|
||||
# Preserve any existing regex mappings from the original mapper
|
||||
if mapper.orig_to_new_regex:
|
||||
ordered_regexes.update(mapper.orig_to_new_regex)
|
||||
|
||||
mapper.orig_to_new_regex = ordered_regexes
|
||||
|
||||
# ══════════════════════════════════════════════════════════════════
|
||||
# Substr mappings (applied AFTER regex by WeightsMapper)
|
||||
# ══════════════════════════════════════════════════════════════════
|
||||
|
||||
# 1. Attention: self_attn → attn.mla_attn mappings
|
||||
# modelopt uses "self_attn" but vllm expects "attn" (mapped to "attn.mla_attn")
|
||||
mapper.orig_to_new_substr[".self_attn.q_a_proj."] = ".attn.mla_attn.wq_a."
|
||||
mapper.orig_to_new_substr[".self_attn.q_b_proj."] = ".attn.mla_attn.wq_b."
|
||||
mapper.orig_to_new_substr[".self_attn.q_a_norm."] = ".attn.mla_attn.q_norm."
|
||||
mapper.orig_to_new_substr[".self_attn.o_a_proj."] = ".attn.mla_attn.wo_a."
|
||||
mapper.orig_to_new_substr[".self_attn.o_b_proj."] = ".attn.mla_attn.wo_b."
|
||||
mapper.orig_to_new_substr[".self_attn.sinks"] = ".attn.mla_attn.attn_sink"
|
||||
|
||||
# CRITICAL: kv_proj must map to wkv (not kv_proj) because the stacking
|
||||
# code looks for "attn.wkv" to stack into fused_wqa_wkv.
|
||||
mapper.orig_to_new_substr[".self_attn.kv_proj."] = ".attn.mla_attn.wkv."
|
||||
mapper.orig_to_new_substr[".self_attn.kv_norm."] = ".attn.mla_attn.kv_norm."
|
||||
|
||||
# Compressor: self_attn.compressor → attn.mla_attn.compressor
|
||||
mapper.orig_to_new_substr[".self_attn.compressor."] = ".attn.mla_attn.compressor."
|
||||
|
||||
# Compressor projection renaming for stacking:
|
||||
# vllm stacks compressor.wkv + compressor.wgate → compressor.fused_wkv_wgate
|
||||
# modelopt exports as compressor.kv_proj and compressor.gate_proj
|
||||
mapper.orig_to_new_substr[".compressor.kv_proj."] = ".compressor.wkv."
|
||||
mapper.orig_to_new_substr[".compressor.gate_proj."] = ".compressor.wgate."
|
||||
|
||||
# 2. Shared expert projections: gate_proj→w1, up_proj→w3
|
||||
# vllm stacks shared_experts.w1 + shared_experts.w3 into
|
||||
# shared_experts.gate_up_proj. modelopt uses gate_proj/up_proj naming.
|
||||
# down_proj stays as-is (vllm model uses down_proj directly).
|
||||
mapper.orig_to_new_substr[".shared_experts.gate_proj."] = ".shared_experts.w1."
|
||||
mapper.orig_to_new_substr[".shared_experts.up_proj."] = ".shared_experts.w3."
|
||||
|
||||
return mapper
|
||||
|
||||
|
||||
def patch():
|
||||
global _original_make_mapper
|
||||
from vllm.model_executor.models import deepseek_v4
|
||||
|
||||
original_make_mapper = deepseek_v4._make_deepseek_v4_weights_mapper
|
||||
# 1. Save the original function BEFORE replacing it
|
||||
_original_make_mapper = deepseek_v4._make_deepseek_v4_weights_mapper
|
||||
|
||||
def patched_make_mapper(expert_dtype: str):
|
||||
mapper = original_make_mapper(expert_dtype)
|
||||
# 2. Patch the function so __init__ calls also get our mapper
|
||||
deepseek_v4._make_deepseek_v4_weights_mapper = make_patched_mapper
|
||||
print("✓ Patched _make_deepseek_v4_weights_mapper function")
|
||||
|
||||
# modelopt uses "self_attn" but vllm expects "attn" (which it then
|
||||
# maps to "attn.mla_attn" via the substr mapper)
|
||||
# We need: self_attn -> attn.mla_attn (skip the intermediate step)
|
||||
mapper.orig_to_new_substr[".self_attn.compressor."] = ".attn.mla_attn.compressor."
|
||||
mapper.orig_to_new_substr[".self_attn.kv_norm."] = ".attn.mla_attn.kv_norm."
|
||||
mapper.orig_to_new_substr[".self_attn.kv_proj."] = ".attn.mla_attn.kv_proj."
|
||||
mapper.orig_to_new_substr[".self_attn.o_a_proj."] = ".attn.mla_attn.wo_a."
|
||||
mapper.orig_to_new_substr[".self_attn.o_b_proj."] = ".attn.mla_attn.wo_b."
|
||||
mapper.orig_to_new_substr[".self_attn.q_a_proj."] = ".attn.mla_attn.wq_a."
|
||||
mapper.orig_to_new_substr[".self_attn.q_a_norm."] = ".attn.mla_attn.q_norm."
|
||||
mapper.orig_to_new_substr[".self_attn.q_b_proj."] = ".attn.mla_attn.wq_b."
|
||||
mapper.orig_to_new_substr[".self_attn.sinks"] = ".attn.mla_attn.attn_sink"
|
||||
# 3. CRITICAL: Also update the CLASS attribute directly.
|
||||
# DeepseekV4ForCausalLM.hf_to_vllm_mapper is set at class definition
|
||||
# time (module import). Our function patch above doesn't retroactively
|
||||
# update it. Since expert_dtype=="fp4", __init__ won't recreate it either.
|
||||
# We MUST update the class attribute directly.
|
||||
if hasattr(deepseek_v4, 'DeepseekV4ForCausalLM'):
|
||||
deepseek_v4.DeepseekV4ForCausalLM.hf_to_vllm_mapper = make_patched_mapper("fp4")
|
||||
print("✓ Updated DeepseekV4ForCausalLM.hf_to_vllm_mapper class attribute")
|
||||
else:
|
||||
print("⚠ DeepseekV4ForCausalLM not found (will be patched at import time)")
|
||||
|
||||
# modelopt names the indexer's sub-projects differently
|
||||
mapper.orig_to_new_substr[".self_attn.compressor.indexer.q_b_proj."] = ".attn.mla_attn.indexer.wq_b."
|
||||
mapper.orig_to_new_substr[".self_attn.compressor.indexer.kv_proj."] = ".attn.mla_attn.indexer.wkv."
|
||||
mapper.orig_to_new_substr[".self_attn.compressor.indexer.gate_proj."] = ".attn.mla_attn.indexer.gate."
|
||||
mapper.orig_to_new_substr[".self_attn.compressor.indexer.weights_proj."] = ".attn.mla_attn.indexer.wo_a."
|
||||
mapper.orig_to_new_substr[".self_attn.compressor.indexer.kv_norm."] = ".attn.mla_attn.indexer.kv_norm."
|
||||
mapper.orig_to_new_substr[".self_attn.compressor.indexer.position_bias"] = ".attn.mla_attn.indexer.position_bias"
|
||||
print("✓ All modelopt NVFP4 weight mapping patches applied")
|
||||
|
||||
# modelopt puts shared experts under mlp.shared_experts with correct names
|
||||
# but the mapper may try to rename .shared_experts. differently
|
||||
# Our model already has model.layers.N.mlp.shared_experts.down_proj etc.
|
||||
|
||||
# modelopt adds hc_head as a separate module (hc = hidden compression)
|
||||
# vllm doesn't have this in the mapper, but it should be handled by
|
||||
# the general weight loading if we don't filter it out
|
||||
|
||||
return mapper
|
||||
|
||||
deepseek_v4._make_deepseek_v4_weights_mapper = patched_make_mapper
|
||||
print("✓ Patched _make_deepseek_v4_weights_mapper for modelopt NVFP4 naming")
|
||||
|
||||
if __name__ == "__main__":
|
||||
patch()
|
||||
|
||||
@@ -82,7 +82,7 @@ cmd = [
|
||||
"--tensor-parallel-size", "8",
|
||||
"--compilation-config", '{"cudagraph_mode":"FULL_AND_PIECEWISE", "custom_ops":["all"]}',
|
||||
"--attention_config.use_fp4_indexer_cache=True",
|
||||
"--moe-backend", "deep_gemm_mega_moe",
|
||||
"--moe-backend", "deep_gemm_mega_moe", # WARN: No NVFP4 mega_moe kernel. Use docker-compose (omits this flag) instead.
|
||||
"--tokenizer-mode", "deepseek_v4",
|
||||
"--tool-call-parser", "deepseek_v4",
|
||||
"--enable-auto-tool-choice",
|
||||
|
||||
Reference in New Issue
Block a user