[Cherry pick to release/v0.5.15] Fix NVILA weight loading (#30400)
Co-authored-by: Brayden Zhong <brayden@radixark.ai>
This commit is contained in:
co-authored by
Brayden Zhong
parent
e2ea7aafad
commit
6875df3378
@@ -210,6 +210,10 @@ class NVILAForConditionalGeneration(nn.Module):
|
||||
if name.startswith("llm."):
|
||||
self.llm.load_weights([(name[len("llm.") :], loaded_weight)])
|
||||
else:
|
||||
if name not in params_dict and name.startswith(
|
||||
"vision_tower.vision_model."
|
||||
):
|
||||
name = "vision_tower." + name[len("vision_tower.vision_model.") :]
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(
|
||||
param, "weight_loader", weight_utils.default_weight_loader
|
||||
|
||||
@@ -168,6 +168,10 @@ class NVILALiteForConditionalGeneration(nn.Module):
|
||||
if name.startswith("llm."):
|
||||
self.llm.load_weights([(name[len("llm.") :], loaded_weight)])
|
||||
else:
|
||||
if name not in params_dict and name.startswith(
|
||||
"vision_tower.vision_model."
|
||||
):
|
||||
name = "vision_tower." + name[len("vision_tower.vision_model.") :]
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(
|
||||
param, "weight_loader", weight_utils.default_weight_loader
|
||||
|
||||
Reference in New Issue
Block a user