[Minor] Add comment on skipping rope caches (#2004)

This commit is contained in:
Woosuk Kwon 2023-12-10 10:04:12 -08:00 committed by GitHub
parent 1aa1361510
commit 24cde76a15
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -322,9 +322,10 @@ class LlamaForCausalLM(nn.Module):
model_name_or_path, cache_dir, load_format, revision): model_name_or_path, cache_dir, load_format, revision):
if "rotary_emb.inv_freq" in name: if "rotary_emb.inv_freq" in name:
continue continue
if "rotary_emb.cos_cached" in name: if ("rotary_emb.cos_cached" in name
continue or "rotary_emb.sin_cached" in name):
if "rotary_emb.sin_cached" in name: # Models trained using ColossalAI may include these tensors in
# the checkpoint. Skip them.
continue continue
for (param_name, weight_name, shard_id) in stacked_params_mapping: for (param_name, weight_name, shard_id) in stacked_params_mapping:
if weight_name not in name: if weight_name not in name: