Fix non-deterministic T5 calibration NaN on multi-GPU (#1636)

## Summary

Fixes known error failing nightly 2-gpu CI 4/5 times

- `tests/examples/llm_ptq/test_llm_ptq.py::test_ptq_t5` intermittently
fails on multi-GPU runners with `AssertionError: detected nan values in
amax. nan in original tensor: True` during FP8 calibration (in T5
encoder self-attention `self.o(attn_output)`).
- Root cause: `device_map="auto"` is memory-aware, so on a busy 2-GPU
box accelerate sometimes shards the tiny `t5-small` across both GPUs. T5
ties the encoder/decoder `shared` embeddings and relies on relative
position-bias buffers; splitting these across devices via naive
model-parallel hooks produces NaN activations (HF transformers
[#21093](https://github.com/huggingface/transformers/issues/21093)). The
placement varies with free memory at load time, which is why it failed
~4/5 runs but passed ~1/5 on the same machine.
- Fix: load T5 on a single device (`device_map=None`) in
`examples/llm_ptq/example_utils.py::get_model`, mirroring the existing
BART handling. The existing `model.to(device)` path then places it on a
single GPU, making calibration deterministic. `t5-small` is tiny, so
single-device placement is not a memory concern.

## Test plan

- Merge and see if nightly CI is no longer flaky
This commit is contained in:
Keval Morabia
2026-06-06 00:32:57 +05:30
committed by GitHub
parent 4eaf90a6ca
commit bcbe2b957e
+6
View File
@@ -598,6 +598,12 @@ def get_model(
# device_map "auto" and "cuda" triggers error regarding meta tensor from safetensors
device_map = None
if hf_config.model_type == "t5":
# device_map "auto" can naively shard T5's tied encoder/decoder embeddings and
# position-bias buffers across GPUs, which non-deterministically produces NaN
# activations during calibration on multi-GPU machines (see HF transformers #21093).
device_map = None
# Helper function to check if model has pack-quantized config
def has_pack_quantized_config(config):
# Check top-level quantization_config