Skip to content

GitLab

Unverified Commit 29ff8716 authored by Sylvain Gugger's avatar Sylvain Gugger Committed by GitHub
Browse files

Add a progress bar for large model loading (#20713)

parent 5f94855d
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
src/transformers/modeling_utils.py
View file @ 29ff8716
...@@ -2659,6 +2659,8 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix ...@@ -2659,6 +2659,8 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix
else: else:
disk_only_shard_files = [] disk_only_shard_files = []
if len(resolved_archive_file) > 1:
resolved_archive_file = logging.tqdm(resolved_archive_file, desc="Loading checkpoint shards")
for shard_file in resolved_archive_file: for shard_file in resolved_archive_file:
# Skip the load for shards that only contain disk-offloaded weights when using safetensors for the offload. # Skip the load for shards that only contain disk-offloaded weights when using safetensors for the offload.
if shard_file in disk_only_shard_files: if shard_file in disk_only_shard_files:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment