Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Raise error if weights are not loaded #206

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Raise error if weights are not loaded #206

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion jetstream_pt/cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,6 @@
flags.DEFINE_bool("enable_model_warmup", False, "enable model warmup")



def shard_weights(env, weights, weight_shardings):
"""Shard weights according to weight_shardings"""
sharded = {}
Expand Down
7 changes: 6 additions & 1 deletion jetstream_pt/fetch_models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,6 +159,10 @@ def _load_weights(directory):
for key in f.keys():
state_dict[key] = f.get_tensor(key).to(torch.bfloat16)
# Load the state_dict into the model
if not state_dict:
raise AssertionError(
f"Tried to load weights from {directory}, but couldn't find any."
)
return state_dict


Expand All @@ -177,7 +181,8 @@ def instantiate_model_from_repo_id(
"""Create model instance by hf model id.+"""
model_dir = _hf_dir(repo_id)
if not FLAGS.internal_use_random_weights and (
not os.path.exists(model_dir) or not os.listdir(model_dir)
not os.path.exists(model_dir)
or not glob.glob(os.path.join(model_dir, "*.safetensors"))
):
# no weights has been downloaded
_hf_download(repo_id, model_dir, FLAGS.hf_token)
Expand Down