Mixture of Experts - #33
dan-smith-tech wants to merge 11 commits into
Conversation
Point the evenet submodule at https://github.com/dan-smith-tech/EveNet-Core (commit fd88624) which adds the MoE architecture (Gate, Expert, MoE classes, updated TransformerBlockModule, PETBody, EveNetModel, and network/loss/moe.py). The EveNetBackbone in model.py already passes all MoE hyperparameters so no further changes are needed here. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Add step 4 to set_peft_trainable for both ensemble modes: after the optional LayerNorm unfreezing, re-freeze all MoE-specific parameters (mlp.gate, mlp.routed_experts, mlp.shared_experts) in the PET transformer blocks. This makes the intent explicit and is robust to any future changes that might otherwise accidentally unfreeze gate routers or expert FFN weights during adapter-based fine-tuning. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
Pull request overview
This PR introduces Mixture-of-Experts (MoE) wiring into the Evenet-Lite backbone setup and adjusts PEFT training behavior to keep MoE weights frozen during adapter fine-tuning.
Changes:
- Pass MoE-related configuration fields from
pet_configinto the upstreamPETBodyconstructor. - Update PEFT parameter-freezing logic to explicitly keep MoE router/experts frozen even when LayerNorms are unfrozen.
- Change the
evenet/git submodule URL.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
evenet_lite/optim.py |
Ensures MoE gate/expert parameters remain frozen during PEFT fine-tuning. |
evenet_lite/model.py |
Threads MoE configuration options through to PETBody construction. |
.gitmodules |
Points the evenet submodule to a different remote URL. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| use_moe=pet_config.use_moe, | ||
| moe_base_num_experts=pet_config.moe_base_num_experts, | ||
| moe_base_select_top_k=pet_config.moe_base_select_top_k, | ||
| moe_num_shared_experts=pet_config.moe_num_shared_experts, | ||
| moe_expert_segmentation_factor=pet_config.moe_expert_segmentation_factor, |
There was a problem hiding this comment.
The default config shipped in evenet_lite/config/default_network_config.yaml does not define use_moe/moe_* under Body.PET, so accessing pet_config.use_moe etc will likely raise at runtime when constructing the default EvenetLiteClassifier model. To keep backward compatibility, either add these keys with defaults to the default YAML (and any other configs), or use safe fallbacks here (e.g., getattr(pet_config, "use_moe", False) and sensible defaults for the other MoE params).
| # 4) explicitly keep MoE params frozen (gate router and expert FFNs | ||
| # are pre-trained weights that must not be updated during adapter | ||
| # fine-tuning, even if LayerNorms were unselectively unfrozen above) | ||
| for name, p in m.backbone.named_parameters(): | ||
| if any(seg in name for seg in ("mlp.gate", "mlp.routed_experts", "mlp.shared_experts")): |
There was a problem hiding this comment.
This MoE-freezing block is duplicated in both the independent and non-ensemble branches. Consider extracting a small helper (e.g., _freeze_moe_params(backbone)) to avoid drift if the list of MoE parameter name patterns changes in the future.
| [submodule "evenet"] | ||
| path = evenet | ||
| url = git@github.com:EveNet-HEP/Core.git | ||
| url = https://github.com/dan-smith-tech/EveNet-Core.git |
There was a problem hiding this comment.
The submodule URL is being changed from the previous upstream (EveNet-HEP/Core.git) to what appears to be a personal fork (dan-smith-tech/EveNet-Core). This can impact reproducibility and access for other contributors/CI. If the intent is only to switch from SSH to HTTPS, consider keeping the upstream repo and just changing the protocol; if a fork is required for MoE support, it would be better to point to the canonical org/repo where those changes will live long-term (or document the rationale).
| url = https://github.com/dan-smith-tech/EveNet-Core.git | |
| url = https://github.com/EveNet-HEP/Core.git |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| use_moe: true | ||
| moe_base_num_experts: 16 |
There was a problem hiding this comment.
use_moe is enabled by default. If the intent of this PR is to add MoE parameters (not to change the default architecture), turning MoE on here will change behavior, resource usage, and potentially weight compatibility for anyone relying on the default config. Consider defaulting use_moe to false and letting users opt-in, or update the PR description to explicitly call out this default behavior change.
Pulls in MoE logging additions and ensures PETBody has use_moe support. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated 5 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # 4) explicitly keep MoE params frozen (gate router and expert FFNs | ||
| # are pre-trained weights that must not be updated during adapter | ||
| # fine-tuning, even if LayerNorms were unselectively unfrozen above) | ||
| for name, p in m.backbone.named_parameters(): | ||
| if any(seg in name for seg in ("mlp.gate", "mlp.routed_experts", "mlp.shared_experts")): | ||
| p.requires_grad_(False) |
There was a problem hiding this comment.
The MoE-freezing logic is duplicated in both the independent-ensemble and non-ensemble branches. Consider extracting this into a small helper (e.g., _freeze_moe_params(backbone)) to keep set_peft_trainable easier to maintain and reduce the chance the branches diverge in future edits.
| # 4) explicitly keep MoE params frozen (gate router and expert FFNs | ||
| # are pre-trained weights that must not be updated during adapter | ||
| # fine-tuning, even if LayerNorms were unselectively unfrozen above) | ||
| for name, p in m.backbone.named_parameters(): | ||
| if any(seg in name for seg in ("mlp.gate", "mlp.routed_experts", "mlp.shared_experts")): | ||
| p.requires_grad_(False) |
There was a problem hiding this comment.
This relies on substring-matching parameter names ("mlp.gate", "mlp.routed_experts", "mlp.shared_experts") to decide what to freeze. That’s brittle to upstream renames and can silently fail to freeze anything. Prefer freezing by module type (e.g., iterate modules and isinstance(..., MoE) or the router/expert submodules) or at least log/warn when no parameters matched while use_moe is enabled.
| use_moe: true | ||
| moe_base_num_experts: 16 | ||
| moe_expert_segmentation_factor: 1 | ||
| moe_base_select_top_k: 1 | ||
| moe_num_shared_experts: 1 | ||
| moe_scale_expert_dim: true | ||
| moe_alpha: 0.01 | ||
| moe_cz: 0.001 | ||
| moe_use_router_noise: true |
There was a problem hiding this comment.
default_network_config.yaml now enables MoE by default (use_moe: true). This is a behavior/perf change relative to previous defaults; the PR description only mentions adding parameters. If MoE is intended to be optional, consider defaulting use_moe to false and letting users opt in, or update the PR description/docs to clearly call out that MoE is now enabled by default.
| def on_epoch_end(self, trainer: "Trainer", epoch: int, metrics: Dict[str, float]) -> None: | ||
| if not trainer.is_rank_zero(): | ||
| return | ||
| if epoch % self.log_every_n_epochs != 0: | ||
| return | ||
| if getattr(trainer, "val_loader", None) is None: | ||
| return | ||
| try: | ||
| from evenet.network.layers.transformer import log_moe_expert_distribution | ||
| except ImportError: | ||
| logging.warning( | ||
| "MoEExpertDistributionCallback: could not import from evenet-core; skipping." | ||
| ) | ||
| return | ||
| model = trainer._unwrap_model() if hasattr(trainer, "_unwrap_model") else trainer.model | ||
| log_moe_expert_distribution(model, reset=False) # reset was done in on_epoch_start |
There was a problem hiding this comment.
In distributed training, this callback logs only on rank 0. If expert dispatch counts are accumulated per-rank (common under DDP), the printed distribution may reflect only rank 0’s shard rather than the global routing distribution. Consider reducing counts across ranks before logging (e.g., all-reduce inside log_moe_expert_distribution or by exposing counts and using trainer._all_reduce_tensor).
| [submodule "evenet"] | ||
| path = evenet | ||
| url = git@github.com:EveNet-HEP/Core.git | ||
| url = https://github.com/dan-smith-tech/EveNet-Core.git |
There was a problem hiding this comment.
The submodule URL was changed from the org SSH remote to a specific HTTPS fork. This can affect reproducibility and trust/supply-chain expectations (and is not mentioned in the PR description). If this is intentional, consider pointing to the canonical upstream repo (or documenting why a fork is required) and ensuring CI/dev workflows that rely on the submodule are updated accordingly.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 9 changed files in this pull request and generated 4 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # set peft | ||
| if getattr(self.config, "use_peft", True): | ||
| set_peft_trainable(self.model, train_layernorm=True) | ||
| if self.is_rank_zero(): |
There was a problem hiding this comment.
use_peft is checked with getattr(self.config, "use_peft", True) here, which defaults to PEFT-enabled even though TrainerConfig.use_peft defaults to False and other call sites use getattr(..., False). This can silently flip behavior for configs/objects missing the attribute. Consider using a consistent default (likely False) or removing getattr now that TrainerConfig always defines the field.
| metric_sum: Dict[str, float] = {"loss": 0.0, "accuracy": 0.0, "moe_l_aux": 0.0, "moe_cz_lz": 0.0} | ||
| metric_count: Dict[str, int] = {"loss": 0, "accuracy": 0, "moe_l_aux": 0, "moe_cz_lz": 0} | ||
| epoch_probs: List[torch.Tensor] = [] |
There was a problem hiding this comment.
MoE metrics (moe_l_aux, moe_cz_lz) are accumulated in metric_sum/metric_count, but only loss and accuracy are all-reduced across ranks. In DDP this will make the epoch-level MoE averages/logging rank-local (and inconsistent) while other metrics are global. Reduce MoE sums/counts as well (e.g., extend the tensors being reduced, or separately all-reduce those values).
|
|
||
| trainer = Trainer(model, feature_names, config, | ||
| callbacks=[MoEExpertDistributionCallback()]) | ||
| trainer.fit(train_dataset, val_dataset=val_dataset, epochs=50) |
There was a problem hiding this comment.
The docstring example uses trainer.fit(...), but Trainer in this repo exposes train(...) (no fit method). This example will fail if copied. Update the example to match the actual API (or add/alias a fit method if that’s intended).
| trainer.fit(train_dataset, val_dataset=val_dataset, epochs=50) | |
| trainer.train(train_dataset, val_dataset=val_dataset, epochs=50) |
| [submodule "evenet"] | ||
| path = evenet | ||
| url = git@github.com:EveNet-HEP/Core.git | ||
| url = https://github.com/dan-smith-tech/EveNet-Core.git | ||
| branch = main |
There was a problem hiding this comment.
This PR is described as adding MoE parameters, but it also repoints the evenet submodule URL (and pins a branch). That’s a significant supply-chain / reproducibility change and may be accidental. If the upstream repo moved, prefer the official organization URL and consider pinning to a specific commit/tag rather than a moving branch.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 9 changed files in this pull request and generated 5 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| metric_sum: Dict[str, float] = {"loss": 0.0, "accuracy": 0.0, "moe_l_aux": 0.0, "moe_cz_lz": 0.0} | ||
| metric_count: Dict[str, int] = {"loss": 0, "accuracy": 0, "moe_l_aux": 0, "moe_cz_lz": 0} |
There was a problem hiding this comment.
Because moe_l_aux/moe_cz_lz are always present in metric_sum, summarize_metrics() will always return these keys (typically 0.0), so they will be logged/printed even during validation or PEFT runs despite the intent to only report them for full fine-tuning. Consider only adding these accumulator keys when MoE is enabled and active (or remove them from the returned metrics when not applicable) to avoid noisy/ambiguous metrics.
| payload: Dict[str, Any] = { | ||
| "train/loss": loss, | ||
| "metrics/train_accuracy": accuracy, | ||
| "epoch": epoch + 1, | ||
| **self._optimizer_learning_rates(), | ||
| } | ||
| if not getattr(self.config, "use_peft", False): | ||
| if moe_l_aux != 0.0: | ||
| payload["train/moe_l_aux"] = moe_l_aux | ||
| if moe_cz_lz != 0.0: | ||
| payload["train/moe_cz_lz"] = moe_cz_lz | ||
| self.wandb_run.log(payload, step=step) |
There was a problem hiding this comment.
W&B logging mixes reduced_loss/reduced_accuracy (averaged across ranks) with moe_l_aux/moe_cz_lz taken only from the local rank-0 batch. This makes the step-level logs internally inconsistent under DDP. If these aux losses are meant to be comparable to the reduced loss, reduce/average the aux values across ranks before logging (or explicitly log them as rank-local).
| _moe_l_aux = _moe_l_aux + m.backbone.PET.moe_l_aux | ||
| _moe_cz_lz = _moe_cz_lz + m.backbone.PET.moe_cz_lz | ||
| if self.n_ensemble > 1: |
There was a problem hiding this comment.
The forward path reads m.backbone.PET.moe_l_aux / moe_cz_lz unconditionally. If MoE is disabled (use_moe=false) or the underlying PETBody implementation doesn’t define these attributes, this will raise AttributeError. Use getattr(..., default_tensor) (and/or gate this on the config) so the model remains usable without MoE support.
| use_moe=pet_config.use_moe, | ||
| moe_base_num_experts=pet_config.moe_base_num_experts, | ||
| moe_base_select_top_k=pet_config.moe_base_select_top_k, | ||
| moe_num_shared_experts=pet_config.moe_num_shared_experts, | ||
| moe_expert_segmentation_factor=pet_config.moe_expert_segmentation_factor, | ||
| moe_scale_expert_dim=pet_config.moe_scale_expert_dim, | ||
| moe_alpha=pet_config.moe_alpha, | ||
| moe_cz=pet_config.moe_cz, | ||
| moe_use_router_noise=pet_config.moe_use_router_noise, |
There was a problem hiding this comment.
The PETBody constructor is now passed many MoE-related fields via pet_config.. If a user supplies a custom network config that predates these keys, DotDict attribute access may fail at runtime. Consider using getattr(pet_config, , ) (or merging defaults into the loaded YAML) to preserve backward compatibility.
| use_moe=pet_config.use_moe, | |
| moe_base_num_experts=pet_config.moe_base_num_experts, | |
| moe_base_select_top_k=pet_config.moe_base_select_top_k, | |
| moe_num_shared_experts=pet_config.moe_num_shared_experts, | |
| moe_expert_segmentation_factor=pet_config.moe_expert_segmentation_factor, | |
| moe_scale_expert_dim=pet_config.moe_scale_expert_dim, | |
| moe_alpha=pet_config.moe_alpha, | |
| moe_cz=pet_config.moe_cz, | |
| moe_use_router_noise=pet_config.moe_use_router_noise, | |
| use_moe=getattr(pet_config, "use_moe", False), | |
| moe_base_num_experts=getattr(pet_config, "moe_base_num_experts", 0), | |
| moe_base_select_top_k=getattr(pet_config, "moe_base_select_top_k", 0), | |
| moe_num_shared_experts=getattr(pet_config, "moe_num_shared_experts", 0), | |
| moe_expert_segmentation_factor=getattr(pet_config, "moe_expert_segmentation_factor", 1), | |
| moe_scale_expert_dim=getattr(pet_config, "moe_scale_expert_dim", 1), | |
| moe_alpha=getattr(pet_config, "moe_alpha", 0.0), | |
| moe_cz=getattr(pet_config, "moe_cz", 0.0), | |
| moe_use_router_noise=getattr(pet_config, "moe_use_router_noise", False), |
| metric_sum: Dict[str, float] = {"loss": 0.0, "accuracy": 0.0, "moe_l_aux": 0.0, "moe_cz_lz": 0.0} | ||
| metric_count: Dict[str, int] = {"loss": 0, "accuracy": 0, "moe_l_aux": 0, "moe_cz_lz": 0} |
There was a problem hiding this comment.
The new MoE metrics (moe_l_aux/moe_cz_lz) are added to metric_sum/metric_count but are never all-reduced across ranks. In DDP this will produce rank-local values (and potentially inconsistent metrics across ranks) while loss/accuracy are reduced. Include these sums/counts in the same all-reduce path (or do a separate all-reduce) before computing epoch averages and returning metrics.
Add MoE parameters.