[quantization] Instability on `llama` quantization

[stamalakhov]

What

There is some instability in llama-based models quantization. E.g.
Running the same command

python tico/quantization/wrapq/examples/quantize_full_qmodel_with_gptq.py --model unsloth/Llama-3.2-3B-Instruct --max_seq_len 2048 --linear_weight_bits 4 --gptq_mse smse --nsamples_for_qcalibration 128 --device cuda --lm_head_weight_bits 4 --save  "ptq_checkpoint"   --no_spinquant  --eval_tasks="mmlu,hellaswag,piqa,truthfulqa" --decode_calibration_steps 8 --sensitivity_path sensitivities_for_unsloth_Llama-3.2-3B-Instruct_wikitext_128_42.pt

on two different gpus produced two dirreferent ppl's (11.86 vs 12.14).

Running mse:

python tico/quantization/wrapq/examples/quantize_full_qmodel_with_gptq.py --model unsloth/Llama-3.2-3B-Instruct --max_seq_len 2048 --linear_weight_bits 4 --gptq_mse mse --nsamples_for_qcalibration 128 --device cuda --lm_head_weight_bits 4 --save  "ptq_checkpoint"   --no_spinquant  --eval_tasks="mmlu,hellaswag,piqa,truthfulqa" --decode_calibration_steps 8 

also produced two different ppl's (12.56 vs 12.52).

Let's make sure that:

1. this is inevitable
2. and/or reduce discrepancy of results.

[mhs4670go]

I think you can narrow the points.

1. Check if there's the same problem with no_gptq and no_spinquant

2. Check if there's the same problem with 1B model that might use a single GPU. I'm not sure but two gpus could be the reason.

[stamalakhov]

I think you can narrow the points.

1. Check if there's the same problem with no_gptq and no_spinquant

2. Check if there's the same problem with 1B model that might use a single GPU. I'm not sure but two gpus could be the reason.

@mhs4670go
Thank you for your reply.

1. ok, spinquant crashes with out of memory on llama (so it's just GPTQ/PTQ issue)
2. it's single GPU. I mean running the same script with the same parameters on different GPUs produce different results.

[mhs4670go]

FYI, I got same PPL when I tried below options with a single GPU.

 --max_seq_len 2048 --linear_weight_bits "4" --embedding_weight_bits 8 \
--lm_head_weight_bits 8 --gptq_mse mse

[stamalakhov]

You mean 12.56? or same PPL as a result?

[mhs4670go]

I've got 11.80 with bbc64aa.

[stamalakhov]

I've got 11.80 with bbc64aa.

Ahh. Got it. It's lm_head to 8 bits.

[stamalakhov]

It's GPTQ , no_GPTQ option produces the same results on tested GPUs.

[mhs4670go]

How about trying the following experiments to narrow it down:

• Disable actorder and check if the results become stable.
• Disable MSE-based GPTQ and see if that removes the variance.

This should help identify whether the instability comes from column reordering or the qparam search.

[stamalakhov]

Seems like torch version was the main cause of ppl fluctuation (~5 % off ppl change on different GPUs). After equalizing torch versions difference in ppl dropped to ~1%. Please see #670 (comment) .

torch	quantized_PPL
torch_2_6	11.85
torch_2_7	12.20
torch_2_8	12.05
torch_2_9	12.05
torch_2_10	11.91

[mhs4670go]

I think the current experiments in this PR are useful for narrowing down the issue, but the results suggest that Hessian precision alone is probably not the root cause.

fix_stability==2 and fix_stability==3 produce identical PPLs on each GPU, but there is still noticeable variance across GPUs.

One thing I would recommend checking next is the device handling inside GPTQQuantizer.convert().

Right now the code uses:

device = next(model.parameters()).device

to move cached inputs before running each layer. However, with device_map="balanced" the model can be sharded across multiple GPUs, and each layer may actually live on a different device.

Instead of using the first model parameter device globally, it would be safer to resolve the execution device per-layer:

def _module_device(module: torch.nn.Module) -> torch.device:
    for p in module.parameters(recurse=True):
        return p.device
    for b in module.buffers(recurse=True):
        return b.device
    return torch.device("cpu")

and then:

layer_device = _module_device(layer)

cache_args_batch = move_to_device(cache_args_batch, layer_device)
cache_kwargs_batch = move_to_device(cache_kwargs_batch, layer_device)

outs = layer(*cache_args_batch, **cache_kwargs_batch)

The same probably applies to the model.model.norm and lm_head path in _quantize_lm_head().

Finally, my guess is that the remaining instability may come from differences in activation/device execution paths across GPUs, which later get amplified by actorder.

[stamalakhov]

to move cached inputs before running each layer. However, with device_map="balanced" the model can be sharded across multiple GPUs, and each layer may actually live on a different device.

@mhs4670go
Thank you. Actually results in #656 (comment) were produced with just single gpu, like "CUDA_VISIBLE_DEVICES=0 python ...".

In my personal opinion:
the main cause of instability is inversion of ill-conditioned matrices (cholesky_inverse...), so the smallest change in Hessians will cause large variance on weights. (Damp factor is just used for it, but seems like it's not enough). Hessians are ill-conditioned, because its just $\sum_{N}(xx^T)$, where $N=128$, so these matrices are uninvertable without dampening even if they typed as double.

I tried to save Hessians to pt file and then to load them to another gpu, and PPL was the same, i mean results on different GPUs were the same.

Moreover for Llama3_2_1B, which is 15 layers long, variance in PPL's is much smaller than for for 27 layers long 3B model (not presented).

[stamalakhov]

@mhs4670go

1. this issue is about running quantization on different GPU servers, with single GPU
2. some instability seems to be inevitable for small number of samples (<5000)
3. we can improve stability using inputs from original model ("--gptq_use_orig_model"). I will provide results later
4. we can dampen more aggressively to improve stability of inversion (for SmolLM 0.1 dampening factor stabilized ppl). I will provide results later.

[stamalakhov]

@mhs4670go
Please see #670 (comment).
Seems like option 3 from above stabilized model. For calibrating Hessians it uses input from original model, which prevents GPTQ divergence and helps to stabilize it.