Repository: huggingface/transformers
Pull request: Reduce tqdm verbosity during model loading

# What does this PR do?

Fixes https://github.com/huggingface/transformers/issues/44303 - see also comments here https://github.com/huggingface/transformers/pull/44316#issuecomment-3984362089.
Supersedes https://github.com/huggingface/transformers/pull/44316

Changed files:
filename: src/transformers/core_model_loading.py
patch: @@ -33,8 +33,9 @@
 
 from .integrations.accelerate import get_device, offload_weight
 from .integrations.tensor_parallel import ALL_PARALLEL_STYLES
-from .utils import is_env_variable_true, logging
+from .utils import is_env_variable_true
 from .utils.loading_report import LoadStateDictInfo
+from .utils.logging import get_logger, tqdm
 
 
 _torch_distributed_available = torch.distributed.is_available()
@@ -45,7 +46,7 @@
     from .quantizers import HfQuantizer
 
 
-logger = logging.get_logger(__name__)
+logger = get_logger(__name__)
 
 
 def build_glob_alternation(
@@ -1203,43 +1204,38 @@ def convert_and_load_state_dict_in_model(
             loading_info.unexpected_keys.add(renamed_key)
 
     try:
-        total_entries = len(param_name_to_load)
-        with logging.tqdm(total=total_entries, desc="Loading weights") as pbar:
-            for first_param_name, mapping in param_name_to_load.items():
-                pbar.update(1)
-                pbar.set_postfix({"Materializing param": first_param_name})
-                pbar.refresh()
-                try:
-                    realized_value = mapping.convert(
-                        first_param_name,
-                        model=model,
-                        config=model.config,
-                        hf_quantizer=hf_quantizer,
-                        loading_info=loading_info,
-                    )
-                    for target_name, param in realized_value.items():
-                        param = param[0] if isinstance(param, list) else param
-                        param_device = get_device(device_map, target_name)
-                        # Offloading support
-                        if param_device == "disk" and (target_name not in model_buffers or offload_buffers):
-                            disk_offload_index = offload_and_maybe_resave_param(
-                                target_name, param, loading_info, disk_offload_folder, disk_offload_index, mapping
-                            )
-                        else:
-                            set_param_for_module(
-                                model,
-                                target_name,
-                                param,
-                                loading_info,
-                                mapping.distributed_operation,
-                                hf_quantizer,
-                            )
-
-                    # Cleanup all the tensors that were gathered before next iteration
-                    del realized_value
-
-                except SkipParameters:
-                    continue
+        for first_param_name, mapping in tqdm(param_name_to_load.items(), desc="Loading weights"):
+            try:
+                realized_value = mapping.convert(
+                    first_param_name,
+                    model=model,
+                    config=model.config,
+                    hf_quantizer=hf_quantizer,
+                    loading_info=loading_info,
+                )
+                for target_name, param in realized_value.items():
+                    param = param[0] if isinstance(param, list) else param
+                    param_device = get_device(device_map, target_name)
+                    # Offloading support
+                    if param_device == "disk" and (target_name not in model_buffers or offload_buffers):
+                        disk_offload_index = offload_and_maybe_resave_param(
+                            target_name, param, loading_info, disk_offload_folder, disk_offload_index, mapping
+                        )
+                    else:
+                        set_param_for_module(
+                            model,
+                            target_name,
+                            param,
+                            loading_info,
+                            mapping.distributed_operation,
+                            hf_quantizer,
+                        )
+
+                # Cleanup all the tensors that were gathered before next iteration
+                del realized_value
+
+            except SkipParameters:
+                continue
 
     # Close the pool, independently of whether the code was interrupted or finished successfully
     finally: