mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user