Update lfm2 moe to use correct bridge elements - #1670
Open
TensorCruncher wants to merge 1 commit into
Open
Conversation
Contributor
Author
|
Contributor
Author
|
Compatibility checks are failing because unit tests have not been updated. Will do that once I have clarity on the test failures for verify model |
Contributor
Author
|
From the lfm2 moe modelling file in transformers we can see that the gate is not a linear nn module. Could this be the reason for the first test failure? The test fails on blocks > 2, so the first two dense mlps don't error. class Lfm2MoeTopKRouter(nn.Module):
def __init__(self, config):
super().__init__()
self.top_k = config.num_experts_per_tok
self.num_experts = config.num_experts
self.norm_topk_prob = config.norm_topk_prob
self.hidden_dim = config.hidden_size
self.weight = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim))
self.routed_scaling_factor = config.routed_scaling_factor
self.use_expert_bias = config.use_expert_bias
def forward(self, hidden_states, expert_bias=None):
router_logits = F.linear(hidden_states, self.weight)
routing_weights = router_logits.sigmoid()
if self.use_expert_bias:
scores_for_routing = routing_weights + expert_bias
_, selected_experts = torch.topk(scores_for_routing, k=self.top_k, dim=-1)
routing_weights = torch.gather(routing_weights, dim=1, index=selected_experts).type_as(router_logits)
else:
routing_weights, selected_experts = torch.topk(routing_weights, k=self.top_k, dim=-1)
if self.norm_topk_prob:
routing_weights = routing_weights / (routing_weights.sum(dim=-1, keepdim=True) + 1e-6)
routing_weights = routing_weights * self.routed_scaling_factor
return router_logits, routing_weights, selected_experts
class Lfm2MoeSparseMoeBlock(nn.Module):
def __init__(self, config):
super().__init__()
self.experts = Lfm2MoeExperts(config)
self.gate = Lfm2MoeTopKRouter(config)
self.use_expert_bias = config.use_expert_bias
if self.use_expert_bias:
self.expert_bias = nn.Buffer(torch.zeros(config.num_experts, dtype=torch.float32))
def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
batch_size, sequence_length, hidden_dim = hidden_states.shape
hidden_states_reshaped = hidden_states.view(-1, hidden_dim)
expert_bias = self.expert_bias if self.use_expert_bias else None
_, routing_weights, selected_experts = self.gate(hidden_states_reshaped, expert_bias)
final_hidden_states = self.experts(hidden_states_reshaped, selected_experts, routing_weights)
return final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) |
jlarson4
reviewed
Aug 14, 2026
jlarson4
left a comment
jlarson4
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for taking this on @TensorCruncher! Decomposing LFM2-MoE into real components is an improvement over the current system and I was able to reproduce the forward parity result on my end. A couple comments below, including the answer to your question:
| "out": LinearBridge(name="out_proj"), | ||
| }, | ||
| ), | ||
| "mlp": MoEBridge( |
Collaborator
There was a problem hiding this comment.
To answer your question: Your diagnosis is right, the structure of the benchmarks system is creating an error message that is masking the real problem. Lfm2MoeTopKRouter takes (hidden_states, expert_bias=None) and returns a 3-tuple, but the deeper problem is that use_expert_bias=True and expert_bias lives on the parent Lfm2MoeSparseMoeBlock. Calling the router with only hidden_states raises unsupported operand type(s) for +: 'Tensor' and 'NoneType'. The benchmark then retries with hidden_states= as a keyword, and that second attempt is what generates the LinearBridge.forward() missing 1 required positional argument: 'input' error.
This should be fixable by subclassing the existing MoERouterBridge and overriding get_random_inputs() to return {"args": (hidden_states, expert_bias)} with expert_bias=torch.zeros(num_experts). Pass config=self.cfg into the submodule too, without it d_model silently falls back to 768 and you'll get a shape error instead.
| "mlp": MoEBridge( | ||
| name="feed_forward", | ||
| config=self.cfg, | ||
| submodules={"gate": LinearBridge(name="gate")}, |
Collaborator
There was a problem hiding this comment.
Because the dense layers get a GatedMLPBridge with a gate submodule while sparse layers use MoEBridge's gate for the router, blocks.0.mlp.gate.hook_out is w1 and blocks.2.mlp.gate.hook_out is the router in the same model. MoEBridge's DENSE_SUBMODULE_KEYS comment describes this per-layer semantic flip as the thing the dense_* keys exist to prevent.
There is an example of the shape you should use in laguna.py:78-88 / llada2_moe.py:110-121. Using this instead should resolve some of your errors. Let me know if you have any questions on this parts specifically.
| Wrapping the HF layer as a whole preserves correct execution while avoiding | ||
| unresolved standard attention/MLP aliases on layers that do not have them. | ||
| """ | ||
| def set_original_component(self, original_component: Any) -> None: |
Collaborator
There was a problem hiding this comment.
I recently added per-layer dense binding to MoEBridge in #1666, so the set_original_component override on lines 27-42 should be safe to delete, as long as you have all the latest commits from dev. Run a git merge if you need to pull in the latest changes.
| rope_parameters = getattr(cfg, "rope_parameters", None) or {} | ||
| rope_theta = rope_parameters.get("rope_theta") or getattr(cfg, "rope_theta", None) | ||
| if rope_theta is not None: | ||
| self.cfg.rotary_base = rope_theta |
Collaborator
There was a problem hiding this comment.
TransformerBridgeConfig doesn't recover rotary_base from rope_parameters, so removing that block leaves cfg.rotary_base = 10000 instead of the model's 5,000,000. Forward parity hides it because RoPE is delegated to HF's model.pos_emb, but anything reading cfg.rotary_base now gets a value that's off by 500×. This deletion also drops default_prepend_bos = False, which changes tokenization on every string-input path.
Can you restore the rope_parameters/rope_theta propagation and the default_prepend_bos = False line? The other deletions in that block are fine, eps, num_experts, experts_per_token, moe_intermediate_size, and layer_types all still reach cfg without the explicit copies. test_norm_and_rope_config and test_default_prepend_bos_is_false are currently failing because they caught these two issues, please fix the adapter rather than updating those two tests.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Updated Lfm2MoE to use correct bridge elements so as to expand hook coverage.
The verify model test fails two tests, both due to the mlp/moe components.
Lfm2 MoE uses dense mlp for first two layers and moe for the remaining. This seems to be causing issues in certain tests. Forward parity passes though.
I will attach results of verify model run.