Repository: huggingface/transformers
Pull request: Disable async loading when quantizing on the fly

# What does this PR do?

Fixes https://github.com/huggingface/transformers/issues/44387.

This PR disable async loading when we want to quantize the model. it is actually faster than doing a semaphore. If a quantizer happens to quantize faster than materializing the weights, we might consider adding adding a semaphore for that. 

For future reference: 

```
I tried to Materialize at most GLOBAL_WORKERS (4) full-precision tensors sit on GPU waiting to be quantized. As the main thread quantizes one and releases a permit, one worker loads the next. We can tune the number in the future it makes more sense in the future. 

This is problematic in the case where we use weight converters. This happens for some models like mixtral,
qwen2_moe, qwen3_vl_moe, lfm2_moe, ernie4_5_moe, solar_open, jamba, and others. For these, we don't use the semaphore and we just warn the user to set `HF_DEACTIVATE_ASYNC_LOAD` to not have memory issues. 

```

Results using `Qwen/Qwen2.5-7B-Instruct`

  ```python
┌───────────────────────────────┬───────┬───────────┬──────────┐
  │         Configuration         │ Time  │ Allocated │ Reserved │
  ├───────────────────────────────┼───────┼───────────┼──────────┤
  │ Baseline (no fix, async ON)   │ 30.5s │ 5.46 GB   │ 13.47 GB │
  ├───────────────────────────────┼───────┼───────────┼──────────┤
  │ Baseline (async OFF)          │ 29.6s │ 5.45 GB   │ 5.63 GB  │
  ├───────────────────────────────┼───────┼───────────┼──────────┤
  │ With semaphore fix (async ON) │ 30.3s │ 5.45 GB   │ 5.94 GB  │
  └───────────────────────────────┴───────┴───────────┴──────────┘
```

Results using `Qwen/Qwen3-30B-A3B` with FP8

 ```python
 ┌─────────┬────────────────┬──────────────────────────────────┐
  │         │     Async      │              Sync                │
  │         │   (default)    │   (HF_DEACTIVATE_ASYNC_LOAD=1)   │
  ├─────────┼────────────────┼──────────────────────────────────┤
  │ Time    │ 12.7s          │ 15.4s                            │
  ├─────────┼────────────────┼──────────────────────────────────┤
  │ Warning │ Yes            │ No                               │
  └─────────┴────────────────┴──────────────────────────────────┘
```

Changed files:
filename: src/transformers/core_model_loading.py
patch: @@ -791,7 +791,10 @@ def _materialize_copy(tensor: torch.Tensor, device=None, dtype=None) -> torch.Te
 
 
 def spawn_materialize(
-    thread_pool: ThreadPoolExecutor | None, tensor: torch.Tensor, device=None, dtype=None
+    thread_pool: ThreadPoolExecutor | None,
+    tensor: torch.Tensor,
+    device=None,
+    dtype=None,
 ) -> Future | Callable:
     """Materialize a tensor from file asynchronously if `thread_pool` is provided, or return a Callable that will
     load the tensor synchronously when called."""
@@ -1110,8 +1113,15 @@ def convert_and_load_state_dict_in_model(
     )
 
     # We use threading by default, if not explicitly deactivated via env variable. If we have to offload,
-    # we cannot use it either to control the memory as we are under memory constraints, so we need to be sequential
-    if is_env_variable_true("HF_DEACTIVATE_ASYNC_LOAD") or "disk" in device_map.values():
+    # we cannot use it either to control the memory as we are under memory constraints, so we need to be sequential.
+    # When doing on-the-fly quantization, we also use sync loading to avoid worker threads loading full-precision
+    # tensors to GPU faster than the main thread can quantize them, which would cause a large memory spike.
+    has_on_the_fly_quantization = hf_quantizer is not None and not hf_quantizer.pre_quantized
+    if (
+        is_env_variable_true("HF_DEACTIVATE_ASYNC_LOAD")
+        or "disk" in device_map.values()
+        or has_on_the_fly_quantization
+    ):
         thread_pool = None
     else:
         thread_pool = ThreadPoolExecutor(max_workers=GLOBAL_WORKERS)
@@ -1153,11 +1163,12 @@ def convert_and_load_state_dict_in_model(
                 source_pattern = original_key
 
             # 3. Handle dtype casting
-            if (
+            needs_quantization = (
                 hf_quantizer
                 and not hf_quantizer.pre_quantized
                 and hf_quantizer.param_needs_quantization(model, renamed_key)
-            ):
+            )
+            if needs_quantization:
                 mapping.quantization_operation = hf_quantizer.get_quantize_ops()
 
             _dtype = dtype