inference-optimization/ZAYA1-74B-preview-NVFP4

Summary

This is a quantized variant of Zyphra/ZAYA1-74B-preview with linears and experts quantized to NVFP4 using GPTQ.

Creation Process

1. Quantize the model to NVFP4 using LLM Compressor

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

# Run this example with `torchrun --nproc_per_node=N zaya1_nvfp4.py`
import torch
from compressed_tensors.offload import init_dist
from transformers import AutoModelForCausalLM, AutoTokenizer

from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import GPTQModifier
from llmcompressor.utils import load_context

# Load model
model_id = "Zyphra/ZAYA1-74B-preview"
init_dist()
with load_context():
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        device_map="auto_offload",
        max_memory={},
        offload_folder="offload_folder",
    )
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Configure the quantization algorithm and scheme.
# Ignore `lm_head` and the MoE router/gate layers (`model.layers.*.mlp.gate.*`).
recipe = GPTQModifier(
    targets="Linear", scheme="NVFP4", ignore=["lm_head", r"re:.*mlp\.gate\..*"]
)

# Apply quantization.
oneshot(
    model=model,
    processor=tokenizer,
    dataset="perfectblend",
    splits="train[:512]",
    recipe=recipe,
    max_seq_length=512,
    num_calibration_samples=256,
    batch_size=16,
)

# Save to disk in compressed-tensors format.
SAVE_DIR = model_id.rstrip("/").split("/")[-1] + "-NVFP4"
model.save_pretrained(SAVE_DIR)
tokenizer.save_pretrained(SAVE_DIR)

torch.distributed.destroy_process_group()

2. Once the model is saved, repack the 2d weights into 3d weights to match the original checkpoint format

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from compressed_tensors.entrypoints.convert import (
    convert_checkpoint,
    MoEExpertPacker,
)

# Example model with unfused, linearized weights
#   model.layers.0.mlp.experts.0.gate_proj.weight
#   model.layers.0.mlp.experts.0.up_proj.weight
#   model.layers.0.mlp.experts.0.down_proj.weight
MODEL_ID = "inference-optimization/ZAYA1-74B-preview-NVFP4-linear"
SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-3d"

# Convert 2d linearized weights into 3d packed weights
converter = MoEExpertPacker.from_pretrained(
    model_name_or_path=MODEL_ID,
    fuse_gate_up=True,
)

convert_checkpoint(
    model_stub=MODEL_ID,
    save_directory=SAVE_DIR,
    converter=converter,
    max_workers=8,
)

Serving

Serving in vLLM is currently limited, so only HF loading with a patch is supported as of now:

from transformers import AutoModelForCausalLM, AutoTokenizer
from llmcompressor.utils import load_context

model_id = "inference-optimization/ZAYA1-74B-preview-NVFP4-linear"
with load_context():
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_id)

input_ids = tokenizer("According to all known laws", return_tensors="pt").input_ids.to(model.device)
output = model.generate(input_ids, max_new_tokens=20)
print(tokenizer.decode(output[0]))

diff --git a/src/llmcompressor/modeling/moe/conversion_mappings.py b/src/llmcompressor/modeling/moe/conversion_mappings.py
index 36f90c1f0..9c8d4c64d 100644
--- a/src/llmcompressor/modeling/moe/conversion_mappings.py
+++ b/src/llmcompressor/modeling/moe/conversion_mappings.py
@@ -164,6 +164,10 @@ ARCH_TO_IMPORT_PATHS: dict[str, tuple[str | list[str], str | list[str]]] = {
         "transformers.models.qwen3_vl_moe.configuration_qwen3_vl_moe.Qwen3VLMoeTextConfig",
         "transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe.Qwen3VLMoeTextExperts",
     ),
+    "zaya": (
+        "transformers.models.zaya.configuration_zaya.ZayaConfig",
+        "transformers.models.zaya.modeling_zaya.ZayaExperts",
+    ),
 }
 
 ARCH_TO_2D_MAPPINGS = {
@@ -248,6 +252,10 @@ ARCH_TO_2D_MAPPINGS = {
             ),
         ],
     ),
+    "zaya": (
+        [],
+        []
+    ),
 }