support Llama4 with non uniformed intermediate size across layers for… (#10047)
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user