[Model] Lfm2Moe #26344
[Model] Lfm2Moe #26344
Conversation
Signed-off-by: Paul Pak <paulpak58@gmail.com>
Signed-off-by: Paul Pak <paulpak58@gmail.com>
Signed-off-by: Paul Pak <paulpak58@gmail.com>
Signed-off-by: Paul Pak <paulpak58@gmail.com>
Signed-off-by: Paul Pak <paulpak58@gmail.com>
Signed-off-by: Paul Pak <paulpak58@gmail.com>
Signed-off-by: Paul Pak <paulpak58@gmail.com>
mergify bot added documentation ![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for the Lfm2Moe model, a Mixture-of-Experts variant. The changes are comprehensive, including the model implementation, configuration, and updates to the model registry and documentation. The overall implementation is solid, but I've identified a few critical issues that need to be addressed. These include a potential AttributeError in the model configuration due to a missing attribute, a faulty assertion in the expert parallelism logic that could lead to crashes, and a minor copy-paste error in an error message. Addressing these points will improve the robustness and correctness of the new model support.
| num_physical_experts: int, | ||
| num_local_physical_experts: int, | ||
| ) -> None: | ||
| assert self.num_local_physical_experts == num_local_physical_experts |
There was a problem hiding this comment.
The assertion self.num_local_physical_experts == num_local_physical_experts will fail if expert rebalancing is triggered by the EPLBScheduler. During rebalancing, the number of physical experts can change, leading to a change in num_local_physical_experts. This assertion compares the old value stored in the model with the new value from the scheduler, which will cause a crash. This assertion should be removed to allow for dynamic updates.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
| for layer in self.model.layers: | ||
| if isinstance(layer.feed_forward, Lfm2MoeSparseMoeBlock): | ||
| moe = layer.feed_forward | ||
| moe.n_local_physical_experts = num_local_physical_experts | ||
| moe.n_physical_experts = num_physical_experts | ||
| moe.n_redundant_experts = self.num_redundant_experts |
There was a problem hiding this comment.
In Lfm2MoeForCausalLM.update_physical_experts_metadata, the loop assumes every element in self.model.layers exposes a feed_forward attribute. When pipeline parallelism is enabled, make_layers inserts PPMissingLayer objects for non‑local layers, so this method will raise AttributeError on those placeholders. Adding an isinstance(layer, PPMissingLayer) or hasattr(layer, "feed_forward") guard before accessing layer.feed_forward would avoid the crash.
Useful? React with 👍 / 👎.
| num_physical_experts: int, | ||
| num_local_physical_experts: int, | ||
| ) -> None: | ||
| assert self.num_local_physical_experts == num_local_physical_experts |
Signed-off-by: Paul Pak <paulpak58@gmail.com>
github-actions bot added the ready 
Signed-off-by: Paul Pak <paulpak58@gmail.com> Signed-off-by: Patrick von Platen <patrick.v.platen@gmail.com>
Signed-off-by: Paul Pak <paulpak58@gmail.com> Signed-off-by: Patrick von Platen <patrick.v.platen@gmail.com>
Signed-off-by: Paul Pak <paulpak58@gmail.com>
Signed-off-by: Paul Pak <paulpak58@gmail.com> Signed-off-by: xuebwang-amd <xuebwang@amd.com>
Signed-off-by: Paul Pak <paulpak58@gmail.com> Signed-off-by: Dhruvil Bhatt <bhattdbh@amazon.com>
3 participants
Purpose
This PR implements LFM2-8B-A1B, a hybrid Mixture-of-Experts architecture variant of LFM2. The LFM2 family is optimized for on-device inference by combining short‑range, input‑aware gated convolutions with grouped‑query attention (GQA). LFM2‑MoE keeps this fast backbone and introduces sparse MoE feed‑forward networks to add representational capacity without significantly increasing the active compute path.
HF model
Test Plan
Test Result
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.