NeuroSenko/Step-5-Preview-Int4-g128

Step-5-Preview — Int4 g128 (MoE experts only)

A W4A16 quantization of Step-5-Preview (StepFun, ~604B-parameter MoE), made from the BF16 re-upload rene98c/Step-5-Preview-BF16.

• Size: 314.9 GiB (vs. 1131 GiB for BF16)
• Format: compressed-tensors, pack-quantized, int4 symmetric, group size 128
• What is quantized: routed MoE experts only (gate_proj / up_proj / down_proj)
• What stays in BF16: attention, shared expert, dense MLP layers, router, norms, embeddings, lm_head, the sparse indexer, MTP layers and the vision tower

⚠️ No inference engine can run this checkpoint end-to-end yet. See Runtime status below. It is published so people can work on runtime support without downloading 1.1 TB of BF16 weights.

Recipe

Routed experts are 96.7% of the model (1089 of 1126 GiB), so only they are quantized. Everything else is 1.6% of the weights, and int4 on attention is where most of the quality loss usually comes from, so it is kept in BF16.

The method is RTN with a per-group clip search. No calibration data is used. For every group of 128 weights, the scale is picked from 17 clip ratios in [0.60, 1.00] to minimize the rounding MSE. The search runs in bf16, the same precision the scale is stored in, so quantization and dequantization agree exactly.

Relative error of a single expert matrix, ‖(W_q − W)·x‖ / ‖W·x‖ on random Gaussian inputs (averaged over 8 experts from 4 tensors across the depth):

Recipe	Error
int4 g128, max-abs scale	11.74%
int4 g128 + clip search (this repo)	10.11%
int4 g128, asymmetric	10.05%

That is roughly the floor of a 16-level grid, not a flaw of this particular recipe. It is per-matrix noise, not the error of the whole model: attention and the shared expert are exact, and the noise from 8 routed experts across 92 layers should be largely uncorrelated. Asymmetric quantization gives the same error for more metadata, so the symmetric format is kept.

Smaller groups barely help, because the error comes from the 4-bit grid itself: g32 with the same clip search gives 10.12% for +25 GiB. int8 experts would be nearly lossless (0.73%) but take ~590 GiB in total.

These are weight-space numbers. Perplexity and KL divergence against BF16 have not been measured yet, because that needs a working forward pass (see below).

Tensor layout

The original checkpoint stores experts as fused 3D tensors [352, out, in]. Here they are split per expert, using the names that vLLM's step3p5 loader expects:

model.layers.N.moe.experts.E.{gate,up,down}_proj.{weight_packed,weight_scale,weight_shape}

The source repo stores two different copies of the MTP layer 92 (in model-00024 and in model-mtp3-full-00001) and two identical copies of embed_tokens. Its index points to model-mtp3-full-00001 and model-00024 respectively. Here, the copies the index does not point to were removed, so every tensor exists exactly once and a loader that reads whole shards cannot pick the wrong one. The unindexed vit.safetensors and model-mtp3-00001.safetensors from the source repo are not included either.

config.json puts text_config at the top level with architectures: ["Step3p5ForCausalLM"] and adds quantization_config. The original multimodal config is kept unchanged as config.step5v-original.json.

Runtime status

• transformers: has no step* architectures. The original repo's auto_map points to a configuration_step_robotics.py that is not published.
• vLLM: the existing step3p5 implementation almost matches this model: config keys, tensor names, FusedMoE and W4A16 MoE kernels (Marlin / Triton) are all there. Two things are missing:
  1. the sparse indexer on the 23 full-attention layers (sparse_indexer_*, ssmax_s);
  2. the vision tower.

The sparse indexer uses topk = 512, and sliding_window is also 512. For sequences of up to 512 tokens, the indexer's token selection is therefore a no-op and every token is attended. This makes short-context testing and layer-wise calibration possible without implementing the selection itself.

The existing GGUF builds (e.g. vcruz305/StepFun-5-Preview-GGUF) already run the model this way in llama.cpp, with the full-attention layers computed densely. One caveat from that work: the indexer path also carries a constant softmax scale (ssmax_s, 0.08496 vs. 1/sqrt(192) = 0.07217), which a plain dense fallback does not apply. A faithful short-context fallback has to keep that scale.

Plans

A calibrated version (GPTQ, possibly AutoRound) may follow, together with measured perplexity / KLD against BF16.

License

Same as the original Step-5-Preview release.