From c5ee6b353b96475e36b903b497c1a852a633d487 Mon Sep 17 00:00:00 2001 From: Han Qi Date: Sat, 14 Dec 2024 00:09:52 +0000 Subject: [PATCH] Raise error if weights are not loaded --- jetstream_pt/cli.py | 1 - jetstream_pt/fetch_models.py | 7 ++++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/jetstream_pt/cli.py b/jetstream_pt/cli.py index f5e7af2..5eca56b 100644 --- a/jetstream_pt/cli.py +++ b/jetstream_pt/cli.py @@ -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 = {} diff --git a/jetstream_pt/fetch_models.py b/jetstream_pt/fetch_models.py index 8c3b0a9..3bee8b1 100644 --- a/jetstream_pt/fetch_models.py +++ b/jetstream_pt/fetch_models.py @@ -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 @@ -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)