From 6606c030d9d49c09f2e99dff8e518470d0efdd0a Mon Sep 17 00:00:00 2001 From: Dan Smith Date: Tue, 31 Mar 2026 17:15:30 +0100 Subject: [PATCH 1/9] copy files that are same --- network/body/embedding.py | 36 +++--- network/evenet_model.py | 27 ++-- network/layers/transformer.py | 223 ++++++++++++++++++++++++++++++++-- 3 files changed, 252 insertions(+), 34 deletions(-) diff --git a/network/body/embedding.py b/network/body/embedding.py index 83db476..4a8c541 100644 --- a/network/body/embedding.py +++ b/network/body/embedding.py @@ -8,7 +8,7 @@ from evenet.network.layers.transformer import TransformerBlockModule from evenet.network.layers.utils import RandomDrop import torch -from evenet.network.body.adapter import Adapter + class EmbeddingStack(nn.Module): def __init__(self, linear_block_type: str, @@ -263,7 +263,11 @@ class PETBody(nn.Module): def __init__( self, num_feat, num_keep, feature_drop, projection_dim, local, K, num_local, num_layers, num_heads, drop_probability, talking_head, layer_scale, - layer_scale_init, dropout, mode, use_adapter: bool = False + layer_scale_init, dropout, mode, use_moe: bool, + moe_base_num_experts: int, moe_base_select_top_k: int, + moe_num_shared_experts: int, moe_expert_segmentation_factor: int, + moe_scale_expert_dim: bool, moe_alpha: float, moe_cz: float, + moe_use_router_noise: bool ): super().__init__() self.num_keep = num_keep @@ -295,19 +299,16 @@ def __init__( self.transformer_blocks = nn.ModuleList([ TransformerBlockModule( projection_dim, num_heads, dropout, talking_head, layer_scale, layer_scale_init, - drop_probability + drop_probability, use_moe=use_moe, moe_base_num_experts=moe_base_num_experts, + moe_base_select_top_k=moe_base_select_top_k, moe_num_shared_experts=moe_num_shared_experts, + moe_expert_segmentation_factor=moe_expert_segmentation_factor, + moe_scale_expert_dim=moe_scale_expert_dim, + moe_alpha=moe_alpha, moe_cz=moe_cz, + moe_use_router_noise=moe_use_router_noise ) for _ in range(num_layers) ]) - self.use_adapter = use_adapter - if self.use_adapter: - self.adapters = nn.ModuleList([ - Adapter(projection_dim, bottleneck=16, dropout=dropout) - for _ in range(num_layers) - ]) - - def forward(self, input_features: Tensor, input_points: Tensor, @@ -350,16 +351,18 @@ def forward(self, encoded = local_features + encoded # Combine with original features skip_connection = encoded - for itransformer, transformer_block in enumerate(self.transformer_blocks): + moe_l_aux = encoded.new_zeros(()) + moe_cz_lz = encoded.new_zeros(()) + for transformer_block in self.transformer_blocks: encoded = transformer_block( x=encoded, mask=mask, attn_mask=attn_mask ) - if self.use_adapter: - encoded = self.adapters[itransformer](encoded) - encoded = encoded * mask.float() - + moe_l_aux += transformer_block.moe_l_aux.to(encoded.device) + moe_cz_lz += transformer_block.moe_cz_lz.to(encoded.device) + self.moe_l_aux = moe_l_aux + self.moe_cz_lz = moe_cz_lz return torch.add(encoded, skip_connection) @@ -440,4 +443,3 @@ def forward(self, x, time_mask, x_mask): x = (x + position_token) * x_mask.float() return x # (B, N, D) - diff --git a/network/evenet_model.py b/network/evenet_model.py index 016ab3f..2943d62 100644 --- a/network/evenet_model.py +++ b/network/evenet_model.py @@ -87,6 +87,9 @@ def __init__( self.global_input_dim: int = global_normalizer_info["norm_mask"].size()[-1] self.sequential_input_dim: int = input_normalizers_setting["SEQUENTIAL"]["norm_mask"].size()[-1] self.local_feature_indices = self.network_cfg.Body.PET.local_point_index + self.invisible_input_dim: int = len(normalization_dict["invisible_mean"]["Source"]) + self.invisible_padding: int = self.sequential_input_dim - self.invisible_input_dim + assert self.invisible_padding >= 0, f"Invisible Padding size {self.invisible_padding} is negative. " self.sequential_normalizer = Normalizer( norm_mask=input_normalizers_setting["SEQUENTIAL"]["norm_mask"].to(self.device), @@ -108,12 +111,7 @@ def __init__( norm_mask=torch.tensor([1], device=self.device, dtype=torch.bool) ) - self.invisible_padding: int = 0 if self.include_neutrino_generation: - self.invisible_input_dim: int = len(normalization_dict["invisible_mean"]["Source"]) - self.invisible_padding = self.sequential_input_dim - self.invisible_input_dim - assert self.invisible_padding >= 0, f"Invisible Padding size {self.invisible_padding} is negative. " - self.invisible_normalizer = Normalizer( mean=normalization_dict["invisible_mean"]["Source"].to(self.device), std=normalization_dict["invisible_std"]["Source"].to(self.device), @@ -155,6 +153,15 @@ def __init__( layer_scale_init=pet_config.layer_scale_init, dropout=pet_config.dropout, mode=pet_config.mode, + 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, ) # [2] Classification + Regression + Assignment Body @@ -376,7 +383,7 @@ def forward( num_point_cloud = x['num_sequential_vectors'].unsqueeze(-1) # (batch_size, 1) B, _, num_features = input_point_cloud.shape - if 'x_invisible' in x and self.include_neutrino_generation: + if 'x_invisible' in x: invisible_point_cloud = x['x_invisible'] pad_size = self.invisible_padding @@ -461,6 +468,8 @@ def forward( full_input_point_cloud = None full_global_conditions = None + moe_l_aux_total = torch.zeros((), device=input_point_cloud.device, dtype=input_point_cloud.dtype) + moe_cz_lz_total = torch.zeros((), device=input_point_cloud.device, dtype=input_point_cloud.dtype) for schedule_name, flag in schedules: if not flag: @@ -537,6 +546,8 @@ def forward( time=full_time, time_masking=time_masking ) + moe_l_aux_total += self.PET.moe_l_aux + moe_cz_lz_total += self.PET.moe_cz_lz if schedule_name == "deterministic" or schedule_name == "generation": ###################################### @@ -655,7 +666,9 @@ def forward( # "full_global_conditions": full_global_conditions, "alpha": alpha, "segmentation-mask": outputs.get("deterministic", {}).get("segmentation-out", {}).get("pred_masks", None), - "segmentation-aux": outputs.get("deterministic", {}).get("segmentation-out", {}).get("aux_outputs", None) + "segmentation-aux": outputs.get("deterministic", {}).get("segmentation-out", {}).get("aux_outputs", None), + "L_aux": moe_l_aux_total, + "cz_Lz": moe_cz_lz_total, } def predict_diffusion_vector( diff --git a/network/layers/transformer.py b/network/layers/transformer.py index b967084..2c75c29 100644 --- a/network/layers/transformer.py +++ b/network/layers/transformer.py @@ -1,6 +1,7 @@ import torch.nn as nn from torch import Tensor import torch +import torch.nn.functional as F from evenet.network.layers.utils import TalkingHeadAttention, StochasticDepth, LayerScale from evenet.network.layers.linear_block import GRUGate, GRUBlock @@ -8,9 +9,183 @@ from typing import Optional +class Gate(nn.Module): + def __init__( + self, + embed_dim: int, + num_experts: int, + select_top_k: int, + use_router_noise: bool + ) -> None: + super().__init__() + self.router = nn.Linear(embed_dim, num_experts, bias=False) + self.noise_router = nn.Linear(embed_dim, num_experts, bias=False) if use_router_noise else None + self.select_top_k = select_top_k + + def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]: + router_logits = self.router(x) + noisy_router_logits = router_logits + + # noisy top-k routing during training to keep exploration healthy + if self.training and self.noise_router is not None: + noise_std = F.softplus(self.noise_router(x)) + noisy_router_logits = noisy_router_logits + torch.randn_like(noisy_router_logits) * noise_std + + # get top-k expert scores per object + topk_logits, topk_indices = torch.topk(noisy_router_logits, self.select_top_k, dim=-1) + # probability distribution over selected experts only + topk_weights = torch.softmax(topk_logits, dim=-1) + + # dense gate weights over all experts from clean logits; used for losses/stats + dense_gate_weights = torch.softmax(router_logits, dim=-1) + + return router_logits, dense_gate_weights, topk_weights, topk_indices + + +class Expert(nn.Module): + def __init__(self, embed_dim: int, feedforward_dim: int, dropout: float) -> None: + super().__init__() + self.ffn = nn.Sequential( + nn.Linear(embed_dim, feedforward_dim), + # GELU activation maintained from original PET FFN + nn.GELU(approximate="none"), + nn.Dropout(dropout), + nn.Linear(feedforward_dim, embed_dim), + ) + + def forward(self, x: Tensor) -> Tensor: + return self.ffn(x) + + +class MoE(nn.Module): + def __init__( + self, + embed_dim: int, + feedforward_dim: int, + base_num_experts: int, + base_select_top_k: int, + num_shared_experts: int, + expert_segmentation_factor: int, + scale_expert_dim: bool, + alpha: float, + c_z: float, + use_router_noise: bool, + dropout: float + ) -> None: + super().__init__() + self.embed_dim = embed_dim + self.base_num_experts = base_num_experts + self.base_select_top_k = base_select_top_k + self.expert_segmentation_factor = expert_segmentation_factor + self.num_shared_experts = num_shared_experts + self.alpha = alpha + self.c_z = c_z + + total_experts = self.base_num_experts * self.expert_segmentation_factor + # num_experts is the total budget - routed experts fill the remainder after reserving shared slots + self.num_experts = total_experts - num_shared_experts + self.select_top_k = self.base_select_top_k * self.expert_segmentation_factor + # when scale_expert_dim is True divide each expert's hidden dim by select_top_k to keep per-token compute constant vs. a vanilla FFN + # note: even shared experts are being impacted by segmentation scaling of k + self.expert_hidden_dim = int(feedforward_dim / (self.select_top_k + self.num_shared_experts)) if scale_expert_dim else feedforward_dim + + self.gate = Gate(embed_dim, self.num_experts, self.select_top_k, use_router_noise=use_router_noise) + self.routed_experts = nn.ModuleList([ + Expert(embed_dim, self.expert_hidden_dim, dropout) + for _ in range(self.num_experts) + ]) + self.shared_experts = nn.ModuleList([ + Expert(embed_dim, self.expert_hidden_dim, dropout) + for _ in range(self.num_shared_experts) + ]) + + def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]: + original_shape = x.shape + # print(f"[MoE] forward pass received objects shape={tuple(original_shape)}") + # collapse batch/object axes to a 2D tensor so that each row corresponds to a single object to route to experts + x = x.reshape(-1, x.shape[-1]) + num_objects = x.shape[0] + + router_logits, dense_gate_weights, topk_weights, topk_indices = self.gate(x) + + routed_output = torch.zeros((num_objects, self.embed_dim), dtype=x.dtype, device=x.device) + objects_per_expert = torch.zeros(self.num_experts, dtype=torch.long, device=x.device) + + if num_objects > 0: + # get flat list of each object id repeated for each of its top-k experts - e.g., [0, 0, 1, 1, 2, 2, ...] + object_indices = torch.arange(num_objects, device=x.device).unsqueeze(1).expand(-1, self.select_top_k).reshape(-1) + # get flat list of which expert each object is assigned to - of size [num_objects * top_k] + expert_indices = topk_indices.reshape(-1) + # get flat list of corresponding expert weights for each object - of size [num_objects * top_k] + expert_weights = topk_weights.reshape(-1) + + # sort by expert index so that all objects for each expert are grouped together + order = torch.argsort(expert_indices) + # update to be in sorted order by expert index + object_indices = object_indices[order] + expert_indices = expert_indices[order] + expert_weights = expert_weights[order] + # count how many objects are assigned to each expert to know how to split the input tensor for each expert's forward pass + objects_per_expert = torch.bincount(expert_indices, minlength=self.num_experts) + + cursor = 0 + # iterate through each expert's assigned objects in order of expert index + for expert_id, count in enumerate(objects_per_expert.tolist()): + if count == 0: + continue + end = cursor + count + + # create minibatch if all objects assigned to the current expert + current_object_indices = object_indices[cursor:end] + current_inputs = x.index_select(0, current_object_indices) + + # forward pass through the current expert with created minibatch + current_outputs = self.routed_experts[expert_id](current_inputs) + # get the corresponding expert weights for the current expert's assigned objects + current_weights = expert_weights[cursor:end].unsqueeze(-1) + + # weight the expert outputs by the corresponding expert weights for each object, + # then add to the correct rows of the final output tensor using the object indices + routed_output.index_add_(0, current_object_indices, current_outputs * current_weights) + + cursor = end + + # as shared experts are not part of the routing decisions, + # run all objects through all shared experts and add to the final output + if self.num_shared_experts > 0: + shared_output = torch.zeros_like(routed_output) + for shared_expert in self.shared_experts: + shared_output = shared_output + shared_expert(x) + final_output = routed_output + shared_output + else: + final_output = routed_output + + # fi - proportion of objects assigned to each expert + denom = max(num_objects * self.select_top_k, 1) + dispatch_fraction = objects_per_expert.to(dtype=dense_gate_weights.dtype) / denom + # pi - average probability of each expert being selected across all objects + mean_router_prob = dense_gate_weights.mean(dim=0) if num_objects > 0 else torch.zeros_like(dispatch_fraction) + # l_aux is the sum across experts of fi * pi, scaled by alpha and num_experts + l_aux = self.alpha * self.num_experts * torch.sum(dispatch_fraction * mean_router_prob) + + if num_objects > 0: + # use clean router logits without noise (pre-softmax) + cz_lz = self.c_z * torch.mean(torch.logsumexp(router_logits, dim=-1).pow(2)) + else: + cz_lz = torch.zeros((), dtype=x.dtype, device=x.device) + + # convert back to original batch/object shape, with the MoE output in the last dimension + final_output = final_output.view(original_shape[0], original_shape[1], self.embed_dim) + + return final_output, l_aux, cz_lz + + class TransformerBlockModule(nn.Module): def __init__(self, projection_dim, num_heads, dropout, talking_head, layer_scale, layer_scale_init, - drop_probability): + drop_probability, use_moe: bool, moe_base_num_experts: int, moe_base_select_top_k: int, + moe_num_shared_experts: int, moe_expert_segmentation_factor: int, + moe_scale_expert_dim: bool, moe_alpha: float, moe_cz: float, + moe_use_router_noise: bool): super().__init__() self.projection_dim = projection_dim self.num_heads = num_heads @@ -18,6 +193,7 @@ def __init__(self, projection_dim, num_heads, dropout, talking_head, layer_scale self.talking_head = talking_head self.layer_scale_flag = layer_scale self.drop_probability = drop_probability + self.use_moe = use_moe self.norm1 = nn.LayerNorm(projection_dim) self.norm2 = nn.LayerNorm(projection_dim) @@ -27,13 +203,30 @@ def __init__(self, projection_dim, num_heads, dropout, talking_head, layer_scale else: self.attn = nn.MultiheadAttention(projection_dim, num_heads, dropout, batch_first=True) - self.mlp = nn.Sequential( - nn.Linear(projection_dim, 2 * projection_dim), - nn.GELU(approximate='none'), - nn.Dropout(dropout), - nn.Linear(2 * projection_dim, projection_dim), - ) + if not self.use_moe: + self.mlp = nn.Sequential( + nn.Linear(projection_dim, 2 * projection_dim), + nn.GELU(approximate="none"), + nn.Dropout(dropout), + nn.Linear(2 * projection_dim, projection_dim), + ) + else: + self.mlp = MoE( + embed_dim=projection_dim, + feedforward_dim=2 * projection_dim, + base_num_experts=moe_base_num_experts, + base_select_top_k=moe_base_select_top_k, + num_shared_experts=moe_num_shared_experts, + expert_segmentation_factor=moe_expert_segmentation_factor, + scale_expert_dim=moe_scale_expert_dim, + alpha=moe_alpha, + c_z=moe_cz, + use_router_noise=moe_use_router_noise, + dropout=dropout, + ) + self.moe_l_aux = torch.tensor(0.0) + self.moe_cz_lz = torch.tensor(0.0) self.drop_path = StochasticDepth(drop_probability) if layer_scale: @@ -42,6 +235,9 @@ def __init__(self, projection_dim, num_heads, dropout, talking_head, layer_scale def forward(self, x, mask, attn_mask=None): # TransformerBlock input shapes: x: torch.Size([B, P, 128]), mask: torch.Size([B, P, 1]) + self.moe_l_aux = x.new_zeros(()) + self.moe_cz_lz = x.new_zeros(()) + padding_mask = ~(mask.squeeze(2).bool()) if mask is not None else None # [batch_size, num_objects] if self.talking_head: @@ -70,11 +266,19 @@ def forward(self, x, mask, attn_mask=None): # Input updates: torch.Size([B, P, 128]), mask: torch.Size([B, P]) x2 = x + self.drop_path(self.layer_scale1(updates, mask)) x3 = self.norm2(x2) - x = x2 + self.drop_path(self.layer_scale2(self.mlp(x3), mask)) + if self.use_moe: + x4, self.moe_l_aux, self.moe_cz_lz = self.mlp(x3) + else: + x4 = self.mlp(x3) + x = x2 + self.drop_path(self.layer_scale2(x4, mask)) else: x2 = x + self.drop_path(updates) x3 = self.norm2(x2) - x = x2 + self.drop_path(self.mlp(x3)) + if self.use_moe: + x4, self.moe_l_aux, self.moe_cz_lz = self.mlp(x3) + else: + x4 = self.mlp(x3) + x = x2 + self.drop_path(x4) if mask is not None: x = x * mask @@ -395,4 +599,3 @@ def forward(self, tgt = tgt + self.dropout3(tgt2) tgt = self.norm3(tgt) return tgt - From a3b4c3d2a060f483a7c9f548736f23f1ddadbbce Mon Sep 17 00:00:00 2001 From: Dan Smith Date: Tue, 31 Mar 2026 17:45:16 +0100 Subject: [PATCH 2/9] add moe loss --- network/loss/moe.py | 40 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) create mode 100644 network/loss/moe.py diff --git a/network/loss/moe.py b/network/loss/moe.py new file mode 100644 index 0000000..d9ce242 --- /dev/null +++ b/network/loss/moe.py @@ -0,0 +1,40 @@ +"""MoE load-balancing losses. + +The auxiliary load-balancing loss (``L_aux``) and router z-loss (``cz_Lz``) are +pre-computed inside the model's forward pass (see ``MoE.forward`` in +``network/layers/transformer.py``) and aggregated by ``EveNetModel.shared_step``. + +This module provides the canonical helper for training engines to collect those +losses from the model output dict and add them to the overall loss with weight 1.0. +""" + +from typing import Optional, Tuple + +from torch import Tensor + + +def loss( + l_aux: Optional[Tensor], + cz_lz: Optional[Tensor], +) -> Tuple[Optional[Tensor], Optional[Tensor]]: + """Return the pre-computed MoE load-balancing losses unchanged. + + Both values are already scalar tensors produced during the model forward + pass: + + * ``l_aux`` – auxiliary load-balancing loss (Switch-Transformer style): + ``alpha * num_experts * Σ(fi * pi)`` + * ``cz_lz`` – router z-loss that penalises large logit magnitudes: + ``c_z * mean(logsumexp(router_logits, dim=-1)²)`` + + They should be summed into the total loss with weight 1.0 (no task-weight + scaling) by the calling engine. + + Args: + l_aux: Scalar auxiliary loss tensor, or ``None`` when MoE is disabled. + cz_lz: Scalar z-loss tensor, or ``None`` when MoE is disabled. + + Returns: + ``(l_aux, cz_lz)`` – the same tensors passed in, unchanged. + """ + return l_aux, cz_lz From 4a87ebdea8c6d0c819733def56d4f8e6ab697afa Mon Sep 17 00:00:00 2001 From: Dan Smith Date: Tue, 31 Mar 2026 18:02:50 +0100 Subject: [PATCH 3/9] Revert "add moe loss" This reverts commit a3b4c3d2a060f483a7c9f548736f23f1ddadbbce. --- network/loss/moe.py | 40 ---------------------------------------- 1 file changed, 40 deletions(-) delete mode 100644 network/loss/moe.py diff --git a/network/loss/moe.py b/network/loss/moe.py deleted file mode 100644 index d9ce242..0000000 --- a/network/loss/moe.py +++ /dev/null @@ -1,40 +0,0 @@ -"""MoE load-balancing losses. - -The auxiliary load-balancing loss (``L_aux``) and router z-loss (``cz_Lz``) are -pre-computed inside the model's forward pass (see ``MoE.forward`` in -``network/layers/transformer.py``) and aggregated by ``EveNetModel.shared_step``. - -This module provides the canonical helper for training engines to collect those -losses from the model output dict and add them to the overall loss with weight 1.0. -""" - -from typing import Optional, Tuple - -from torch import Tensor - - -def loss( - l_aux: Optional[Tensor], - cz_lz: Optional[Tensor], -) -> Tuple[Optional[Tensor], Optional[Tensor]]: - """Return the pre-computed MoE load-balancing losses unchanged. - - Both values are already scalar tensors produced during the model forward - pass: - - * ``l_aux`` – auxiliary load-balancing loss (Switch-Transformer style): - ``alpha * num_experts * Σ(fi * pi)`` - * ``cz_lz`` – router z-loss that penalises large logit magnitudes: - ``c_z * mean(logsumexp(router_logits, dim=-1)²)`` - - They should be summed into the total loss with weight 1.0 (no task-weight - scaling) by the calling engine. - - Args: - l_aux: Scalar auxiliary loss tensor, or ``None`` when MoE is disabled. - cz_lz: Scalar z-loss tensor, or ``None`` when MoE is disabled. - - Returns: - ``(l_aux, cz_lz)`` – the same tensors passed in, unchanged. - """ - return l_aux, cz_lz From dad72b651261122560d3baec9997635bba1f480c Mon Sep 17 00:00:00 2001 From: Dan Smith Date: Tue, 31 Mar 2026 18:02:52 +0100 Subject: [PATCH 4/9] Revert "copy files that are same" This reverts commit 6606c030d9d49c09f2e99dff8e518470d0efdd0a. --- network/body/embedding.py | 36 +++--- network/evenet_model.py | 27 ++-- network/layers/transformer.py | 223 ++-------------------------------- 3 files changed, 34 insertions(+), 252 deletions(-) diff --git a/network/body/embedding.py b/network/body/embedding.py index 4a8c541..83db476 100644 --- a/network/body/embedding.py +++ b/network/body/embedding.py @@ -8,7 +8,7 @@ from evenet.network.layers.transformer import TransformerBlockModule from evenet.network.layers.utils import RandomDrop import torch - +from evenet.network.body.adapter import Adapter class EmbeddingStack(nn.Module): def __init__(self, linear_block_type: str, @@ -263,11 +263,7 @@ class PETBody(nn.Module): def __init__( self, num_feat, num_keep, feature_drop, projection_dim, local, K, num_local, num_layers, num_heads, drop_probability, talking_head, layer_scale, - layer_scale_init, dropout, mode, use_moe: bool, - moe_base_num_experts: int, moe_base_select_top_k: int, - moe_num_shared_experts: int, moe_expert_segmentation_factor: int, - moe_scale_expert_dim: bool, moe_alpha: float, moe_cz: float, - moe_use_router_noise: bool + layer_scale_init, dropout, mode, use_adapter: bool = False ): super().__init__() self.num_keep = num_keep @@ -299,16 +295,19 @@ def __init__( self.transformer_blocks = nn.ModuleList([ TransformerBlockModule( projection_dim, num_heads, dropout, talking_head, layer_scale, layer_scale_init, - drop_probability, use_moe=use_moe, moe_base_num_experts=moe_base_num_experts, - moe_base_select_top_k=moe_base_select_top_k, moe_num_shared_experts=moe_num_shared_experts, - moe_expert_segmentation_factor=moe_expert_segmentation_factor, - moe_scale_expert_dim=moe_scale_expert_dim, - moe_alpha=moe_alpha, moe_cz=moe_cz, - moe_use_router_noise=moe_use_router_noise + drop_probability ) for _ in range(num_layers) ]) + self.use_adapter = use_adapter + if self.use_adapter: + self.adapters = nn.ModuleList([ + Adapter(projection_dim, bottleneck=16, dropout=dropout) + for _ in range(num_layers) + ]) + + def forward(self, input_features: Tensor, input_points: Tensor, @@ -351,18 +350,16 @@ def forward(self, encoded = local_features + encoded # Combine with original features skip_connection = encoded - moe_l_aux = encoded.new_zeros(()) - moe_cz_lz = encoded.new_zeros(()) - for transformer_block in self.transformer_blocks: + for itransformer, transformer_block in enumerate(self.transformer_blocks): encoded = transformer_block( x=encoded, mask=mask, attn_mask=attn_mask ) - moe_l_aux += transformer_block.moe_l_aux.to(encoded.device) - moe_cz_lz += transformer_block.moe_cz_lz.to(encoded.device) - self.moe_l_aux = moe_l_aux - self.moe_cz_lz = moe_cz_lz + if self.use_adapter: + encoded = self.adapters[itransformer](encoded) + encoded = encoded * mask.float() + return torch.add(encoded, skip_connection) @@ -443,3 +440,4 @@ def forward(self, x, time_mask, x_mask): x = (x + position_token) * x_mask.float() return x # (B, N, D) + diff --git a/network/evenet_model.py b/network/evenet_model.py index 2943d62..016ab3f 100644 --- a/network/evenet_model.py +++ b/network/evenet_model.py @@ -87,9 +87,6 @@ def __init__( self.global_input_dim: int = global_normalizer_info["norm_mask"].size()[-1] self.sequential_input_dim: int = input_normalizers_setting["SEQUENTIAL"]["norm_mask"].size()[-1] self.local_feature_indices = self.network_cfg.Body.PET.local_point_index - self.invisible_input_dim: int = len(normalization_dict["invisible_mean"]["Source"]) - self.invisible_padding: int = self.sequential_input_dim - self.invisible_input_dim - assert self.invisible_padding >= 0, f"Invisible Padding size {self.invisible_padding} is negative. " self.sequential_normalizer = Normalizer( norm_mask=input_normalizers_setting["SEQUENTIAL"]["norm_mask"].to(self.device), @@ -111,7 +108,12 @@ def __init__( norm_mask=torch.tensor([1], device=self.device, dtype=torch.bool) ) + self.invisible_padding: int = 0 if self.include_neutrino_generation: + self.invisible_input_dim: int = len(normalization_dict["invisible_mean"]["Source"]) + self.invisible_padding = self.sequential_input_dim - self.invisible_input_dim + assert self.invisible_padding >= 0, f"Invisible Padding size {self.invisible_padding} is negative. " + self.invisible_normalizer = Normalizer( mean=normalization_dict["invisible_mean"]["Source"].to(self.device), std=normalization_dict["invisible_std"]["Source"].to(self.device), @@ -153,15 +155,6 @@ def __init__( layer_scale_init=pet_config.layer_scale_init, dropout=pet_config.dropout, mode=pet_config.mode, - 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, ) # [2] Classification + Regression + Assignment Body @@ -383,7 +376,7 @@ def forward( num_point_cloud = x['num_sequential_vectors'].unsqueeze(-1) # (batch_size, 1) B, _, num_features = input_point_cloud.shape - if 'x_invisible' in x: + if 'x_invisible' in x and self.include_neutrino_generation: invisible_point_cloud = x['x_invisible'] pad_size = self.invisible_padding @@ -468,8 +461,6 @@ def forward( full_input_point_cloud = None full_global_conditions = None - moe_l_aux_total = torch.zeros((), device=input_point_cloud.device, dtype=input_point_cloud.dtype) - moe_cz_lz_total = torch.zeros((), device=input_point_cloud.device, dtype=input_point_cloud.dtype) for schedule_name, flag in schedules: if not flag: @@ -546,8 +537,6 @@ def forward( time=full_time, time_masking=time_masking ) - moe_l_aux_total += self.PET.moe_l_aux - moe_cz_lz_total += self.PET.moe_cz_lz if schedule_name == "deterministic" or schedule_name == "generation": ###################################### @@ -666,9 +655,7 @@ def forward( # "full_global_conditions": full_global_conditions, "alpha": alpha, "segmentation-mask": outputs.get("deterministic", {}).get("segmentation-out", {}).get("pred_masks", None), - "segmentation-aux": outputs.get("deterministic", {}).get("segmentation-out", {}).get("aux_outputs", None), - "L_aux": moe_l_aux_total, - "cz_Lz": moe_cz_lz_total, + "segmentation-aux": outputs.get("deterministic", {}).get("segmentation-out", {}).get("aux_outputs", None) } def predict_diffusion_vector( diff --git a/network/layers/transformer.py b/network/layers/transformer.py index 2c75c29..b967084 100644 --- a/network/layers/transformer.py +++ b/network/layers/transformer.py @@ -1,7 +1,6 @@ import torch.nn as nn from torch import Tensor import torch -import torch.nn.functional as F from evenet.network.layers.utils import TalkingHeadAttention, StochasticDepth, LayerScale from evenet.network.layers.linear_block import GRUGate, GRUBlock @@ -9,183 +8,9 @@ from typing import Optional -class Gate(nn.Module): - def __init__( - self, - embed_dim: int, - num_experts: int, - select_top_k: int, - use_router_noise: bool - ) -> None: - super().__init__() - self.router = nn.Linear(embed_dim, num_experts, bias=False) - self.noise_router = nn.Linear(embed_dim, num_experts, bias=False) if use_router_noise else None - self.select_top_k = select_top_k - - def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]: - router_logits = self.router(x) - noisy_router_logits = router_logits - - # noisy top-k routing during training to keep exploration healthy - if self.training and self.noise_router is not None: - noise_std = F.softplus(self.noise_router(x)) - noisy_router_logits = noisy_router_logits + torch.randn_like(noisy_router_logits) * noise_std - - # get top-k expert scores per object - topk_logits, topk_indices = torch.topk(noisy_router_logits, self.select_top_k, dim=-1) - # probability distribution over selected experts only - topk_weights = torch.softmax(topk_logits, dim=-1) - - # dense gate weights over all experts from clean logits; used for losses/stats - dense_gate_weights = torch.softmax(router_logits, dim=-1) - - return router_logits, dense_gate_weights, topk_weights, topk_indices - - -class Expert(nn.Module): - def __init__(self, embed_dim: int, feedforward_dim: int, dropout: float) -> None: - super().__init__() - self.ffn = nn.Sequential( - nn.Linear(embed_dim, feedforward_dim), - # GELU activation maintained from original PET FFN - nn.GELU(approximate="none"), - nn.Dropout(dropout), - nn.Linear(feedforward_dim, embed_dim), - ) - - def forward(self, x: Tensor) -> Tensor: - return self.ffn(x) - - -class MoE(nn.Module): - def __init__( - self, - embed_dim: int, - feedforward_dim: int, - base_num_experts: int, - base_select_top_k: int, - num_shared_experts: int, - expert_segmentation_factor: int, - scale_expert_dim: bool, - alpha: float, - c_z: float, - use_router_noise: bool, - dropout: float - ) -> None: - super().__init__() - self.embed_dim = embed_dim - self.base_num_experts = base_num_experts - self.base_select_top_k = base_select_top_k - self.expert_segmentation_factor = expert_segmentation_factor - self.num_shared_experts = num_shared_experts - self.alpha = alpha - self.c_z = c_z - - total_experts = self.base_num_experts * self.expert_segmentation_factor - # num_experts is the total budget - routed experts fill the remainder after reserving shared slots - self.num_experts = total_experts - num_shared_experts - self.select_top_k = self.base_select_top_k * self.expert_segmentation_factor - # when scale_expert_dim is True divide each expert's hidden dim by select_top_k to keep per-token compute constant vs. a vanilla FFN - # note: even shared experts are being impacted by segmentation scaling of k - self.expert_hidden_dim = int(feedforward_dim / (self.select_top_k + self.num_shared_experts)) if scale_expert_dim else feedforward_dim - - self.gate = Gate(embed_dim, self.num_experts, self.select_top_k, use_router_noise=use_router_noise) - self.routed_experts = nn.ModuleList([ - Expert(embed_dim, self.expert_hidden_dim, dropout) - for _ in range(self.num_experts) - ]) - self.shared_experts = nn.ModuleList([ - Expert(embed_dim, self.expert_hidden_dim, dropout) - for _ in range(self.num_shared_experts) - ]) - - def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]: - original_shape = x.shape - # print(f"[MoE] forward pass received objects shape={tuple(original_shape)}") - # collapse batch/object axes to a 2D tensor so that each row corresponds to a single object to route to experts - x = x.reshape(-1, x.shape[-1]) - num_objects = x.shape[0] - - router_logits, dense_gate_weights, topk_weights, topk_indices = self.gate(x) - - routed_output = torch.zeros((num_objects, self.embed_dim), dtype=x.dtype, device=x.device) - objects_per_expert = torch.zeros(self.num_experts, dtype=torch.long, device=x.device) - - if num_objects > 0: - # get flat list of each object id repeated for each of its top-k experts - e.g., [0, 0, 1, 1, 2, 2, ...] - object_indices = torch.arange(num_objects, device=x.device).unsqueeze(1).expand(-1, self.select_top_k).reshape(-1) - # get flat list of which expert each object is assigned to - of size [num_objects * top_k] - expert_indices = topk_indices.reshape(-1) - # get flat list of corresponding expert weights for each object - of size [num_objects * top_k] - expert_weights = topk_weights.reshape(-1) - - # sort by expert index so that all objects for each expert are grouped together - order = torch.argsort(expert_indices) - # update to be in sorted order by expert index - object_indices = object_indices[order] - expert_indices = expert_indices[order] - expert_weights = expert_weights[order] - # count how many objects are assigned to each expert to know how to split the input tensor for each expert's forward pass - objects_per_expert = torch.bincount(expert_indices, minlength=self.num_experts) - - cursor = 0 - # iterate through each expert's assigned objects in order of expert index - for expert_id, count in enumerate(objects_per_expert.tolist()): - if count == 0: - continue - end = cursor + count - - # create minibatch if all objects assigned to the current expert - current_object_indices = object_indices[cursor:end] - current_inputs = x.index_select(0, current_object_indices) - - # forward pass through the current expert with created minibatch - current_outputs = self.routed_experts[expert_id](current_inputs) - # get the corresponding expert weights for the current expert's assigned objects - current_weights = expert_weights[cursor:end].unsqueeze(-1) - - # weight the expert outputs by the corresponding expert weights for each object, - # then add to the correct rows of the final output tensor using the object indices - routed_output.index_add_(0, current_object_indices, current_outputs * current_weights) - - cursor = end - - # as shared experts are not part of the routing decisions, - # run all objects through all shared experts and add to the final output - if self.num_shared_experts > 0: - shared_output = torch.zeros_like(routed_output) - for shared_expert in self.shared_experts: - shared_output = shared_output + shared_expert(x) - final_output = routed_output + shared_output - else: - final_output = routed_output - - # fi - proportion of objects assigned to each expert - denom = max(num_objects * self.select_top_k, 1) - dispatch_fraction = objects_per_expert.to(dtype=dense_gate_weights.dtype) / denom - # pi - average probability of each expert being selected across all objects - mean_router_prob = dense_gate_weights.mean(dim=0) if num_objects > 0 else torch.zeros_like(dispatch_fraction) - # l_aux is the sum across experts of fi * pi, scaled by alpha and num_experts - l_aux = self.alpha * self.num_experts * torch.sum(dispatch_fraction * mean_router_prob) - - if num_objects > 0: - # use clean router logits without noise (pre-softmax) - cz_lz = self.c_z * torch.mean(torch.logsumexp(router_logits, dim=-1).pow(2)) - else: - cz_lz = torch.zeros((), dtype=x.dtype, device=x.device) - - # convert back to original batch/object shape, with the MoE output in the last dimension - final_output = final_output.view(original_shape[0], original_shape[1], self.embed_dim) - - return final_output, l_aux, cz_lz - - class TransformerBlockModule(nn.Module): def __init__(self, projection_dim, num_heads, dropout, talking_head, layer_scale, layer_scale_init, - drop_probability, use_moe: bool, moe_base_num_experts: int, moe_base_select_top_k: int, - moe_num_shared_experts: int, moe_expert_segmentation_factor: int, - moe_scale_expert_dim: bool, moe_alpha: float, moe_cz: float, - moe_use_router_noise: bool): + drop_probability): super().__init__() self.projection_dim = projection_dim self.num_heads = num_heads @@ -193,7 +18,6 @@ def __init__(self, projection_dim, num_heads, dropout, talking_head, layer_scale self.talking_head = talking_head self.layer_scale_flag = layer_scale self.drop_probability = drop_probability - self.use_moe = use_moe self.norm1 = nn.LayerNorm(projection_dim) self.norm2 = nn.LayerNorm(projection_dim) @@ -203,30 +27,13 @@ def __init__(self, projection_dim, num_heads, dropout, talking_head, layer_scale else: self.attn = nn.MultiheadAttention(projection_dim, num_heads, dropout, batch_first=True) - if not self.use_moe: - self.mlp = nn.Sequential( - nn.Linear(projection_dim, 2 * projection_dim), - nn.GELU(approximate="none"), - nn.Dropout(dropout), - nn.Linear(2 * projection_dim, projection_dim), - ) - else: - self.mlp = MoE( - embed_dim=projection_dim, - feedforward_dim=2 * projection_dim, - base_num_experts=moe_base_num_experts, - base_select_top_k=moe_base_select_top_k, - num_shared_experts=moe_num_shared_experts, - expert_segmentation_factor=moe_expert_segmentation_factor, - scale_expert_dim=moe_scale_expert_dim, - alpha=moe_alpha, - c_z=moe_cz, - use_router_noise=moe_use_router_noise, - dropout=dropout, - ) + self.mlp = nn.Sequential( + nn.Linear(projection_dim, 2 * projection_dim), + nn.GELU(approximate='none'), + nn.Dropout(dropout), + nn.Linear(2 * projection_dim, projection_dim), + ) - self.moe_l_aux = torch.tensor(0.0) - self.moe_cz_lz = torch.tensor(0.0) self.drop_path = StochasticDepth(drop_probability) if layer_scale: @@ -235,9 +42,6 @@ def __init__(self, projection_dim, num_heads, dropout, talking_head, layer_scale def forward(self, x, mask, attn_mask=None): # TransformerBlock input shapes: x: torch.Size([B, P, 128]), mask: torch.Size([B, P, 1]) - self.moe_l_aux = x.new_zeros(()) - self.moe_cz_lz = x.new_zeros(()) - padding_mask = ~(mask.squeeze(2).bool()) if mask is not None else None # [batch_size, num_objects] if self.talking_head: @@ -266,19 +70,11 @@ def forward(self, x, mask, attn_mask=None): # Input updates: torch.Size([B, P, 128]), mask: torch.Size([B, P]) x2 = x + self.drop_path(self.layer_scale1(updates, mask)) x3 = self.norm2(x2) - if self.use_moe: - x4, self.moe_l_aux, self.moe_cz_lz = self.mlp(x3) - else: - x4 = self.mlp(x3) - x = x2 + self.drop_path(self.layer_scale2(x4, mask)) + x = x2 + self.drop_path(self.layer_scale2(self.mlp(x3), mask)) else: x2 = x + self.drop_path(updates) x3 = self.norm2(x2) - if self.use_moe: - x4, self.moe_l_aux, self.moe_cz_lz = self.mlp(x3) - else: - x4 = self.mlp(x3) - x = x2 + self.drop_path(x4) + x = x2 + self.drop_path(self.mlp(x3)) if mask is not None: x = x * mask @@ -599,3 +395,4 @@ def forward(self, tgt = tgt + self.dropout3(tgt2) tgt = self.norm3(tgt) return tgt + From fd88624e89b5a24cff44a667ff089fb2a6ab4240 Mon Sep 17 00:00:00 2001 From: Dan Smith Date: Tue, 31 Mar 2026 18:17:18 +0100 Subject: [PATCH 5/9] Add MoE architecture to transformer, embedding, model, and loss - Add Gate, Expert, MoE classes to network/layers/transformer.py - Update TransformerBlockModule with optional MoE FFN (use_moe param), storing per-block moe_l_aux and moe_cz_lz after each forward pass - Update PETBody in network/body/embedding.py to accept all MoE hyperparameters alongside existing use_adapter, pass them to each TransformerBlockModule, and accumulate moe_l_aux/moe_cz_lz across transformer blocks in forward - Update EveNetModel in network/evenet_model.py to pass MoE params to PETBody, track moe_l_aux_total/moe_cz_lz_total in forward, and return them as 'L_aux'/'cz_Lz' in the output dict - Add network/loss/moe.py with a loss() helper that extracts L_aux and cz_Lz from the model output dict for integration into training loops Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- network/body/embedding.py | 20 ++- network/evenet_model.py | 17 ++- network/layers/transformer.py | 222 ++++++++++++++++++++++++++++++++-- network/loss/moe.py | 25 ++++ 4 files changed, 271 insertions(+), 13 deletions(-) create mode 100644 network/loss/moe.py diff --git a/network/body/embedding.py b/network/body/embedding.py index 83db476..5e618ba 100644 --- a/network/body/embedding.py +++ b/network/body/embedding.py @@ -263,7 +263,11 @@ class PETBody(nn.Module): def __init__( self, num_feat, num_keep, feature_drop, projection_dim, local, K, num_local, num_layers, num_heads, drop_probability, talking_head, layer_scale, - layer_scale_init, dropout, mode, use_adapter: bool = False + layer_scale_init, dropout, mode, use_adapter: bool = False, + use_moe: bool = False, moe_base_num_experts: int = 4, + moe_base_select_top_k: int = 2, moe_num_shared_experts: int = 0, + moe_expert_segmentation_factor: int = 1, moe_scale_expert_dim: bool = False, + moe_alpha: float = 0.01, moe_cz: float = 0.0, moe_use_router_noise: bool = False ): super().__init__() self.num_keep = num_keep @@ -295,7 +299,12 @@ def __init__( self.transformer_blocks = nn.ModuleList([ TransformerBlockModule( projection_dim, num_heads, dropout, talking_head, layer_scale, layer_scale_init, - drop_probability + drop_probability, use_moe=use_moe, moe_base_num_experts=moe_base_num_experts, + moe_base_select_top_k=moe_base_select_top_k, moe_num_shared_experts=moe_num_shared_experts, + moe_expert_segmentation_factor=moe_expert_segmentation_factor, + moe_scale_expert_dim=moe_scale_expert_dim, + moe_alpha=moe_alpha, moe_cz=moe_cz, + moe_use_router_noise=moe_use_router_noise ) for _ in range(num_layers) ]) @@ -350,16 +359,21 @@ def forward(self, encoded = local_features + encoded # Combine with original features skip_connection = encoded + moe_l_aux = encoded.new_zeros(()) + moe_cz_lz = encoded.new_zeros(()) for itransformer, transformer_block in enumerate(self.transformer_blocks): encoded = transformer_block( x=encoded, mask=mask, attn_mask=attn_mask ) + moe_l_aux += transformer_block.moe_l_aux.to(encoded.device) + moe_cz_lz += transformer_block.moe_cz_lz.to(encoded.device) if self.use_adapter: encoded = self.adapters[itransformer](encoded) encoded = encoded * mask.float() - + self.moe_l_aux = moe_l_aux + self.moe_cz_lz = moe_cz_lz return torch.add(encoded, skip_connection) diff --git a/network/evenet_model.py b/network/evenet_model.py index 016ab3f..565a6ac 100644 --- a/network/evenet_model.py +++ b/network/evenet_model.py @@ -155,6 +155,15 @@ def __init__( layer_scale_init=pet_config.layer_scale_init, dropout=pet_config.dropout, mode=pet_config.mode, + 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, ) # [2] Classification + Regression + Assignment Body @@ -461,6 +470,8 @@ def forward( full_input_point_cloud = None full_global_conditions = None + moe_l_aux_total = torch.zeros((), device=input_point_cloud.device, dtype=input_point_cloud.dtype) + moe_cz_lz_total = torch.zeros((), device=input_point_cloud.device, dtype=input_point_cloud.dtype) for schedule_name, flag in schedules: if not flag: @@ -537,6 +548,8 @@ def forward( time=full_time, time_masking=time_masking ) + moe_l_aux_total += self.PET.moe_l_aux + moe_cz_lz_total += self.PET.moe_cz_lz if schedule_name == "deterministic" or schedule_name == "generation": ###################################### @@ -655,7 +668,9 @@ def forward( # "full_global_conditions": full_global_conditions, "alpha": alpha, "segmentation-mask": outputs.get("deterministic", {}).get("segmentation-out", {}).get("pred_masks", None), - "segmentation-aux": outputs.get("deterministic", {}).get("segmentation-out", {}).get("aux_outputs", None) + "segmentation-aux": outputs.get("deterministic", {}).get("segmentation-out", {}).get("aux_outputs", None), + "L_aux": moe_l_aux_total, + "cz_Lz": moe_cz_lz_total, } def predict_diffusion_vector( diff --git a/network/layers/transformer.py b/network/layers/transformer.py index b967084..6517e13 100644 --- a/network/layers/transformer.py +++ b/network/layers/transformer.py @@ -1,6 +1,7 @@ import torch.nn as nn from torch import Tensor import torch +import torch.nn.functional as F from evenet.network.layers.utils import TalkingHeadAttention, StochasticDepth, LayerScale from evenet.network.layers.linear_block import GRUGate, GRUBlock @@ -8,9 +9,183 @@ from typing import Optional + +class Gate(nn.Module): + def __init__( + self, + embed_dim: int, + num_experts: int, + select_top_k: int, + use_router_noise: bool + ) -> None: + super().__init__() + self.router = nn.Linear(embed_dim, num_experts, bias=False) + self.noise_router = nn.Linear(embed_dim, num_experts, bias=False) if use_router_noise else None + self.select_top_k = select_top_k + + def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]: + router_logits = self.router(x) + noisy_router_logits = router_logits + + # noisy top-k routing during training to keep exploration healthy + if self.training and self.noise_router is not None: + noise_std = F.softplus(self.noise_router(x)) + noisy_router_logits = noisy_router_logits + torch.randn_like(noisy_router_logits) * noise_std + + # get top-k expert scores per object + topk_logits, topk_indices = torch.topk(noisy_router_logits, self.select_top_k, dim=-1) + # probability distribution over selected experts only + topk_weights = torch.softmax(topk_logits, dim=-1) + + # dense gate weights over all experts from clean logits; used for losses/stats + dense_gate_weights = torch.softmax(router_logits, dim=-1) + + return router_logits, dense_gate_weights, topk_weights, topk_indices + + +class Expert(nn.Module): + def __init__(self, embed_dim: int, feedforward_dim: int, dropout: float) -> None: + super().__init__() + self.ffn = nn.Sequential( + nn.Linear(embed_dim, feedforward_dim), + # GELU activation maintained from original PET FFN + nn.GELU(approximate="none"), + nn.Dropout(dropout), + nn.Linear(feedforward_dim, embed_dim), + ) + + def forward(self, x: Tensor) -> Tensor: + return self.ffn(x) + + +class MoE(nn.Module): + def __init__( + self, + embed_dim: int, + feedforward_dim: int, + base_num_experts: int, + base_select_top_k: int, + num_shared_experts: int, + expert_segmentation_factor: int, + scale_expert_dim: bool, + alpha: float, + c_z: float, + use_router_noise: bool, + dropout: float + ) -> None: + super().__init__() + self.embed_dim = embed_dim + self.base_num_experts = base_num_experts + self.base_select_top_k = base_select_top_k + self.expert_segmentation_factor = expert_segmentation_factor + self.num_shared_experts = num_shared_experts + self.alpha = alpha + self.c_z = c_z + + total_experts = self.base_num_experts * self.expert_segmentation_factor + # num_experts is the total budget - routed experts fill the remainder after reserving shared slots + self.num_experts = total_experts - num_shared_experts + self.select_top_k = self.base_select_top_k * self.expert_segmentation_factor + # when scale_expert_dim is True divide each expert's hidden dim by select_top_k to keep per-token compute constant vs. a vanilla FFN + # note: even shared experts are being impacted by segmentation scaling of k + self.expert_hidden_dim = int(feedforward_dim / (self.select_top_k + self.num_shared_experts)) if scale_expert_dim else feedforward_dim + + self.gate = Gate(embed_dim, self.num_experts, self.select_top_k, use_router_noise=use_router_noise) + self.routed_experts = nn.ModuleList([ + Expert(embed_dim, self.expert_hidden_dim, dropout) + for _ in range(self.num_experts) + ]) + self.shared_experts = nn.ModuleList([ + Expert(embed_dim, self.expert_hidden_dim, dropout) + for _ in range(self.num_shared_experts) + ]) + + def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]: + original_shape = x.shape + # collapse batch/object axes to a 2D tensor so that each row corresponds to a single object to route to experts + x = x.reshape(-1, x.shape[-1]) + num_objects = x.shape[0] + + router_logits, dense_gate_weights, topk_weights, topk_indices = self.gate(x) + + routed_output = torch.zeros((num_objects, self.embed_dim), dtype=x.dtype, device=x.device) + objects_per_expert = torch.zeros(self.num_experts, dtype=torch.long, device=x.device) + + if num_objects > 0: + # get flat list of each object id repeated for each of its top-k experts - e.g., [0, 0, 1, 1, 2, 2, ...] + object_indices = torch.arange(num_objects, device=x.device).unsqueeze(1).expand(-1, self.select_top_k).reshape(-1) + # get flat list of which expert each object is assigned to - of size [num_objects * top_k] + expert_indices = topk_indices.reshape(-1) + # get flat list of corresponding expert weights for each object - of size [num_objects * top_k] + expert_weights = topk_weights.reshape(-1) + + # sort by expert index so that all objects for each expert are grouped together + order = torch.argsort(expert_indices) + # update to be in sorted order by expert index + object_indices = object_indices[order] + expert_indices = expert_indices[order] + expert_weights = expert_weights[order] + # count how many objects are assigned to each expert to know how to split the input tensor for each expert's forward pass + objects_per_expert = torch.bincount(expert_indices, minlength=self.num_experts) + + cursor = 0 + # iterate through each expert's assigned objects in order of expert index + for expert_id, count in enumerate(objects_per_expert.tolist()): + if count == 0: + continue + end = cursor + count + + # create minibatch if all objects assigned to the current expert + current_object_indices = object_indices[cursor:end] + current_inputs = x.index_select(0, current_object_indices) + + # forward pass through the current expert with created minibatch + current_outputs = self.routed_experts[expert_id](current_inputs) + # get the corresponding expert weights for the current expert's assigned objects + current_weights = expert_weights[cursor:end].unsqueeze(-1) + + # weight the expert outputs by the corresponding expert weights for each object, + # then add to the correct rows of the final output tensor using the object indices + routed_output.index_add_(0, current_object_indices, current_outputs * current_weights) + + cursor = end + + # as shared experts are not part of the routing decisions, + # run all objects through all shared experts and add to the final output + if self.num_shared_experts > 0: + shared_output = torch.zeros_like(routed_output) + for shared_expert in self.shared_experts: + shared_output = shared_output + shared_expert(x) + final_output = routed_output + shared_output + else: + final_output = routed_output + + # fi - proportion of objects assigned to each expert + denom = max(num_objects * self.select_top_k, 1) + dispatch_fraction = objects_per_expert.to(dtype=dense_gate_weights.dtype) / denom + # pi - average probability of each expert being selected across all objects + mean_router_prob = dense_gate_weights.mean(dim=0) if num_objects > 0 else torch.zeros_like(dispatch_fraction) + # l_aux is the sum across experts of fi * pi, scaled by alpha and num_experts + l_aux = self.alpha * self.num_experts * torch.sum(dispatch_fraction * mean_router_prob) + + if num_objects > 0: + # use clean router logits without noise (pre-softmax) + cz_lz = self.c_z * torch.mean(torch.logsumexp(router_logits, dim=-1).pow(2)) + else: + cz_lz = torch.zeros((), dtype=x.dtype, device=x.device) + + # convert back to original batch/object shape, with the MoE output in the last dimension + final_output = final_output.view(original_shape[0], original_shape[1], self.embed_dim) + + return final_output, l_aux, cz_lz + + class TransformerBlockModule(nn.Module): def __init__(self, projection_dim, num_heads, dropout, talking_head, layer_scale, layer_scale_init, - drop_probability): + drop_probability, use_moe: bool = False, moe_base_num_experts: int = 4, + moe_base_select_top_k: int = 2, moe_num_shared_experts: int = 0, + moe_expert_segmentation_factor: int = 1, moe_scale_expert_dim: bool = False, + moe_alpha: float = 0.01, moe_cz: float = 0.0, moe_use_router_noise: bool = False): super().__init__() self.projection_dim = projection_dim self.num_heads = num_heads @@ -18,6 +193,7 @@ def __init__(self, projection_dim, num_heads, dropout, talking_head, layer_scale self.talking_head = talking_head self.layer_scale_flag = layer_scale self.drop_probability = drop_probability + self.use_moe = use_moe self.norm1 = nn.LayerNorm(projection_dim) self.norm2 = nn.LayerNorm(projection_dim) @@ -27,13 +203,30 @@ def __init__(self, projection_dim, num_heads, dropout, talking_head, layer_scale else: self.attn = nn.MultiheadAttention(projection_dim, num_heads, dropout, batch_first=True) - self.mlp = nn.Sequential( - nn.Linear(projection_dim, 2 * projection_dim), - nn.GELU(approximate='none'), - nn.Dropout(dropout), - nn.Linear(2 * projection_dim, projection_dim), - ) + if not self.use_moe: + self.mlp = nn.Sequential( + nn.Linear(projection_dim, 2 * projection_dim), + nn.GELU(approximate="none"), + nn.Dropout(dropout), + nn.Linear(2 * projection_dim, projection_dim), + ) + else: + self.mlp = MoE( + embed_dim=projection_dim, + feedforward_dim=2 * projection_dim, + base_num_experts=moe_base_num_experts, + base_select_top_k=moe_base_select_top_k, + num_shared_experts=moe_num_shared_experts, + expert_segmentation_factor=moe_expert_segmentation_factor, + scale_expert_dim=moe_scale_expert_dim, + alpha=moe_alpha, + c_z=moe_cz, + use_router_noise=moe_use_router_noise, + dropout=dropout, + ) + self.moe_l_aux = torch.tensor(0.0) + self.moe_cz_lz = torch.tensor(0.0) self.drop_path = StochasticDepth(drop_probability) if layer_scale: @@ -42,6 +235,9 @@ def __init__(self, projection_dim, num_heads, dropout, talking_head, layer_scale def forward(self, x, mask, attn_mask=None): # TransformerBlock input shapes: x: torch.Size([B, P, 128]), mask: torch.Size([B, P, 1]) + self.moe_l_aux = x.new_zeros(()) + self.moe_cz_lz = x.new_zeros(()) + padding_mask = ~(mask.squeeze(2).bool()) if mask is not None else None # [batch_size, num_objects] if self.talking_head: @@ -70,11 +266,19 @@ def forward(self, x, mask, attn_mask=None): # Input updates: torch.Size([B, P, 128]), mask: torch.Size([B, P]) x2 = x + self.drop_path(self.layer_scale1(updates, mask)) x3 = self.norm2(x2) - x = x2 + self.drop_path(self.layer_scale2(self.mlp(x3), mask)) + if self.use_moe: + x4, self.moe_l_aux, self.moe_cz_lz = self.mlp(x3) + else: + x4 = self.mlp(x3) + x = x2 + self.drop_path(self.layer_scale2(x4, mask)) else: x2 = x + self.drop_path(updates) x3 = self.norm2(x2) - x = x2 + self.drop_path(self.mlp(x3)) + if self.use_moe: + x4, self.moe_l_aux, self.moe_cz_lz = self.mlp(x3) + else: + x4 = self.mlp(x3) + x = x2 + self.drop_path(x4) if mask is not None: x = x * mask diff --git a/network/loss/moe.py b/network/loss/moe.py new file mode 100644 index 0000000..cfc8988 --- /dev/null +++ b/network/loss/moe.py @@ -0,0 +1,25 @@ +from torch import Tensor +from typing import Optional, Tuple, Dict, Any + + +def loss(model_output: Dict[str, Any]) -> Tuple[Optional[Tensor], Optional[Tensor]]: + """ + Extract MoE auxiliary losses from model output. + + The auxiliary load-balancing loss (l_aux) and z-loss (cz_lz) are computed + during the model forward pass inside each MoE transformer block and + accumulated onto the model output dict under the keys "L_aux" and "cz_Lz". + Callers should add both directly to their total loss without additional + scaling; the weights are already baked in via the alpha and c_z hyperparameters + configured on each MoE layer. + + Args: + model_output: dict returned by EveNetModel.forward() + + Returns: + l_aux: auxiliary load-balancing loss, or None when MoE is disabled + cz_lz: z-loss regularisation term, or None when MoE is disabled + """ + l_aux = model_output.get("L_aux", None) + cz_lz = model_output.get("cz_Lz", None) + return l_aux, cz_lz From eac4a9e420893c170f946d29cf210fd92585ee3c Mon Sep 17 00:00:00 2001 From: Dan Smith Date: Wed, 8 Apr 2026 01:39:21 +0100 Subject: [PATCH 6/9] add extensive moe logging --- network/evenet_model.py | 36 +++++++++++++++++++++++++++ network/layers/transformer.py | 47 +++++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+) diff --git a/network/evenet_model.py b/network/evenet_model.py index 565a6ac..068c0f6 100644 --- a/network/evenet_model.py +++ b/network/evenet_model.py @@ -21,6 +21,7 @@ from torch import Tensor, nn from typing import Dict, Optional, Any, Union import re +import logging class EveNetModel(nn.Module): @@ -326,6 +327,41 @@ def __init__( ("deterministic", self.include_classification or self.include_assignment or self.include_regression or self.include_segmentation), ] + self._log_backbone_setup() + + def _log_backbone_setup(self) -> None: + logger = logging.getLogger(__name__) + pet_cfg = self.network_cfg.Body.PET + + pretrain_path = getattr(getattr(self.options, "Training", None), "pretrain_model_load_path", None) + if pretrain_path: + logger.info(f"[Backbone] Pretrain path : {pretrain_path}") + else: + logger.warning("[Backbone] No pretrain_model_load_path set — training from scratch.") + + logger.info( + f"[Backbone] PET config : " + f"layers={pet_cfg.num_layers}, " + f"heads={pet_cfg.num_heads}, " + f"dim={pet_cfg.hidden_dim}, " + f"mode={pet_cfg.mode}" + ) + logger.info( + f"[Backbone] MoE enabled : {pet_cfg.use_moe}" + ) + if pet_cfg.use_moe: + logger.info( + f"[Backbone] MoE config : " + f"num_experts={pet_cfg.moe_base_num_experts}, " + f"top_k={pet_cfg.moe_base_select_top_k}, " + f"shared_experts={pet_cfg.moe_num_shared_experts}, " + f"seg_factor={pet_cfg.moe_expert_segmentation_factor}, " + f"scale_dim={pet_cfg.moe_scale_expert_dim}, " + f"alpha={pet_cfg.moe_alpha}, " + f"cz={pet_cfg.moe_cz}, " + f"router_noise={pet_cfg.moe_use_router_noise}" + ) + def forward( self, x: Dict[str, Tensor], time: Tensor, progressive_params: dict = None, diff --git a/network/layers/transformer.py b/network/layers/transformer.py index 6517e13..ba65863 100644 --- a/network/layers/transformer.py +++ b/network/layers/transformer.py @@ -1,3 +1,5 @@ +import logging + import torch.nn as nn from torch import Tensor import torch @@ -9,6 +11,8 @@ from typing import Optional +_moe_logger = logging.getLogger(__name__) + class Gate(nn.Module): def __init__( @@ -99,6 +103,7 @@ def __init__( Expert(embed_dim, self.expert_hidden_dim, dropout) for _ in range(self.num_shared_experts) ]) + self._forward_logged = False def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]: original_shape = x.shape @@ -108,6 +113,48 @@ def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]: router_logits, dense_gate_weights, topk_weights, topk_indices = self.gate(x) + if not self._forward_logged: + expert_w_in = self.routed_experts[0].ffn[0].weight.shape # (hidden, embed) + expert_w_out = self.routed_experts[0].ffn[3].weight.shape # (embed, hidden) + gate_w = self.gate.router.weight.shape # (num_experts, embed) + + shape_ok = ( + x.shape[-1] == self.embed_dim + and router_logits.shape == (num_objects, self.num_experts) + and topk_indices.shape == (num_objects, self.select_top_k) + and expert_w_in[0] == self.expert_hidden_dim + and expert_w_in[1] == self.embed_dim + and gate_w[0] == self.num_experts + ) + level = _moe_logger.info if shape_ok else _moe_logger.error + + level( + f"[MoE] First forward — input: {tuple(original_shape)} " + f"(flattened tokens: {num_objects})" + ) + level( + f"[MoE] gate weight : {tuple(gate_w)} " + f"→ router_logits: {tuple(router_logits.shape)} " + f"topk_indices: {tuple(topk_indices.shape)}" + ) + level( + f"[MoE] expert ffn[0] : {tuple(expert_w_in)} " + f"ffn[3]: {tuple(expert_w_out)} " + f"(expected hidden={self.expert_hidden_dim}, embed={self.embed_dim})" + ) + if self.num_shared_experts > 0: + sh_w_in = self.shared_experts[0].ffn[0].weight.shape + sh_w_out = self.shared_experts[0].ffn[3].weight.shape + level( + f"[MoE] shared ffn[0] : {tuple(sh_w_in)} " + f"ffn[3]: {tuple(sh_w_out)}" + ) + if not shape_ok: + _moe_logger.error("[MoE] ❌ Shape mismatch detected — check config vs checkpoint.") + else: + _moe_logger.info("[MoE] ✅ All shapes consistent.") + self._forward_logged = True + routed_output = torch.zeros((num_objects, self.embed_dim), dtype=x.dtype, device=x.device) objects_per_expert = torch.zeros(self.num_experts, dtype=torch.long, device=x.device) From ebf8e89f7d6ff16e10a3572e53d81ee366dfec9c Mon Sep 17 00:00:00 2001 From: Dan Smith Date: Thu, 9 Apr 2026 17:38:22 +0100 Subject: [PATCH 7/9] log expert usage distribution --- network/layers/transformer.py | 64 +++++++++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/network/layers/transformer.py b/network/layers/transformer.py index ba65863..54bca13 100644 --- a/network/layers/transformer.py +++ b/network/layers/transformer.py @@ -104,6 +104,11 @@ def __init__( for _ in range(self.num_shared_experts) ]) self._forward_logged = False + self.register_buffer('expert_dispatch_counts', torch.zeros(self.num_experts, dtype=torch.long)) + + def reset_expert_dispatch_counts(self) -> None: + """Reset the accumulated eval-time dispatch counts to zero.""" + self.expert_dispatch_counts.zero_() def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]: original_shape = x.shape @@ -175,6 +180,9 @@ def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]: # count how many objects are assigned to each expert to know how to split the input tensor for each expert's forward pass objects_per_expert = torch.bincount(expert_indices, minlength=self.num_experts) + if not self.training: + self.expert_dispatch_counts.add_(objects_per_expert.to(self.expert_dispatch_counts.device)) + cursor = 0 # iterate through each expert's assigned objects in order of expert index for expert_id, count in enumerate(objects_per_expert.tolist()): @@ -647,3 +655,59 @@ def forward(self, tgt = self.norm3(tgt) return tgt + +def log_moe_expert_distribution( + model: nn.Module, + reset: bool = True, + logger: Optional[logging.Logger] = None, +) -> None: + """Log the expert dispatch distribution for every MoE layer in *model*. + + Call this after running eval batches to see how evenly inputs were routed. + Counts are accumulated across all eval forward passes since the last reset. + + Args: + model: Any nn.Module that may contain MoE sub-modules. + reset: If True (default), zero the counts after logging so the next + eval epoch starts fresh. + logger: Logger to write to. Defaults to this module's logger. + + Example usage in an eval loop:: + + model.eval() + for batch in eval_loader: + with torch.no_grad(): + model(batch) + log_moe_expert_distribution(model) # prints distribution, resets counts + """ + log = (logger or _moe_logger).info + + moe_layers = [(name, m) for name, m in model.named_modules() if isinstance(m, MoE)] + if not moe_layers: + _moe_logger.warning("log_moe_expert_distribution: no MoE layers found in model.") + return + + for name, moe in moe_layers: + counts = moe.expert_dispatch_counts.cpu() + total = counts.sum().item() + + if total == 0: + log(f"[MoE dist] {name}: no eval data recorded (counts are all zero).") + if reset: + moe.reset_expert_dispatch_counts() + continue + + uniform = total / moe.num_experts + lines = [ + f"[MoE dist] {name} " + f"(total dispatches={total:,}, ideal per expert={uniform:,.1f})" + ] + for i, c in enumerate(counts.tolist()): + pct = 100.0 * c / total + bar = "█" * int(pct / 2) # each block ≈ 2 % + lines.append(f" expert {i:3d}: {c:9,d} ({pct:5.1f}%) {bar}") + log("\n".join(lines)) + + if reset: + moe.reset_expert_dispatch_counts() + From 96bffdd4d026e4303f107eb407dc696890257141 Mon Sep 17 00:00:00 2001 From: Dan Smith Date: Thu, 9 Apr 2026 17:39:11 +0100 Subject: [PATCH 8/9] logging --- utilities/logger.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/utilities/logger.py b/utilities/logger.py index 4c5ba3e..d9846b3 100644 --- a/utilities/logger.py +++ b/utilities/logger.py @@ -123,6 +123,6 @@ def setup_logging(log_level=logging.INFO, rank: int = 0, log_dir="logs"): format=f"%(asctime)s | %(name)s | %(levelname)s | %(message)s", handlers=[ logging.FileHandler(log_file, mode="a"), - # logging.StreamHandler() if rank == 0 else logging.NullHandler(), # Only rank 0 logs to stdout + logging.StreamHandler() if rank == 0 else logging.NullHandler(), ] ) From 2f66c19d61cd5f2c727ad466eb348624af3a2f36 Mon Sep 17 00:00:00 2001 From: Dan Smith Date: Fri, 10 Apr 2026 05:35:09 +0100 Subject: [PATCH 9/9] fix logging --- network/layers/transformer.py | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/network/layers/transformer.py b/network/layers/transformer.py index 54bca13..bd4372e 100644 --- a/network/layers/transformer.py +++ b/network/layers/transformer.py @@ -131,18 +131,16 @@ def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]: and expert_w_in[1] == self.embed_dim and gate_w[0] == self.num_experts ) - level = _moe_logger.info if shape_ok else _moe_logger.error - - level( + print( f"[MoE] First forward — input: {tuple(original_shape)} " f"(flattened tokens: {num_objects})" ) - level( + print( f"[MoE] gate weight : {tuple(gate_w)} " f"→ router_logits: {tuple(router_logits.shape)} " f"topk_indices: {tuple(topk_indices.shape)}" ) - level( + print( f"[MoE] expert ffn[0] : {tuple(expert_w_in)} " f"ffn[3]: {tuple(expert_w_out)} " f"(expected hidden={self.expert_hidden_dim}, embed={self.embed_dim})" @@ -150,14 +148,14 @@ def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]: if self.num_shared_experts > 0: sh_w_in = self.shared_experts[0].ffn[0].weight.shape sh_w_out = self.shared_experts[0].ffn[3].weight.shape - level( + print( f"[MoE] shared ffn[0] : {tuple(sh_w_in)} " f"ffn[3]: {tuple(sh_w_out)}" ) if not shape_ok: - _moe_logger.error("[MoE] ❌ Shape mismatch detected — check config vs checkpoint.") + print("[MoE] ❌ Shape mismatch detected — check config vs checkpoint.") else: - _moe_logger.info("[MoE] ✅ All shapes consistent.") + print("[MoE] ✅ All shapes consistent.") self._forward_logged = True routed_output = torch.zeros((num_objects, self.embed_dim), dtype=x.dtype, device=x.device) @@ -680,11 +678,12 @@ def log_moe_expert_distribution( model(batch) log_moe_expert_distribution(model) # prints distribution, resets counts """ - log = (logger or _moe_logger).info + _log = (logger.info if logger is not None else print) + _warn = (logger.warning if logger is not None else print) moe_layers = [(name, m) for name, m in model.named_modules() if isinstance(m, MoE)] if not moe_layers: - _moe_logger.warning("log_moe_expert_distribution: no MoE layers found in model.") + _warn("log_moe_expert_distribution: no MoE layers found in model.") return for name, moe in moe_layers: @@ -692,7 +691,7 @@ def log_moe_expert_distribution( total = counts.sum().item() if total == 0: - log(f"[MoE dist] {name}: no eval data recorded (counts are all zero).") + _log(f"[MoE dist] {name}: no eval data recorded (counts are all zero).") if reset: moe.reset_expert_dispatch_counts() continue @@ -706,7 +705,7 @@ def log_moe_expert_distribution( pct = 100.0 * c / total bar = "█" * int(pct / 2) # each block ≈ 2 % lines.append(f" expert {i:3d}: {c:9,d} ({pct:5.1f}%) {bar}") - log("\n".join(lines)) + _log("\n".join(lines)) if reset: moe.reset_expert_dispatch_counts()