support Llama4 with non uniformed intermediate size across layers for… (#10047)

This commit is contained in:
gongwei-130
2025-09-05 17:28:15 -07:00
committed by GitHub
parent 273b28344b
commit ab62b135c1
7 changed files with 123 additions and 13 deletions
+9
View File
@@ -423,6 +423,12 @@ class Llama4DecoderLayer(nn.Module):
return self.config.num_local_experts > 0
return (layer_id + 1) % self.config.interleave_moe_layer_step == 0
def get_intermediate_size(self) -> int:
if isinstance(self.feed_forward, Llama4MoE):
return self.config.intermediate_size
else:
return self.config.intermediate_size_mlp
def forward(
self,
positions: torch.Tensor,
@@ -540,6 +546,9 @@ class Llama4ForCausalLM(LlamaForCausalLM):
def get_input_embeddings(self):
return self.model.embed_tokens
def get_layers(self):
return self.model.layers
def _init_model(
self,
config: Llama4TextConfig,