diff --git a/helm/modules/product_mice.py b/helm/modules/product_mice.py index e1a9c8c..0628716 100644 --- a/helm/modules/product_mice.py +++ b/helm/modules/product_mice.py @@ -13,7 +13,6 @@ bias-update .float() calls. No NEW manual fp32 promotions. import torch from torch import nn import torch.nn.functional as F -import torch.distributed as dist from helm.hypercore.manifolds.product_lorentz import ProductLorentz from helm.hypercore.nn.linear.product_lorentz_linear import ProductLorentzLinear @@ -149,15 +148,13 @@ class ProductLorentzMoE(nn.Module): self.dim = args.dim self.manifold = manifold self.n_factors = manifold.n_factors - world_size = dist.get_world_size() if dist.is_initialized() else 1 - if args.n_routed_experts % world_size: - raise AssertionError(f'n_routed_experts must be divisible by world_size={world_size}') if args.n_activated_experts != 2 or args.n_shared_experts != 1: raise NotImplementedError('Validated HELM-D product MiCE configuration requires top2 routed and one shared expert') if self.dim % self.n_factors: raise ValueError('Model width must divide evenly into Lorentz factors') self.n_routed_experts = args.n_routed_experts - self.n_local_experts = args.n_routed_experts # match supplied HELM code + # Ordinary DDP replicates all experts on every rank, regardless of world size. + self.n_local_experts = args.n_routed_experts self.n_activated_experts = args.n_activated_experts self.experts_start_idx = 0 self.experts_end_idx = self.n_routed_experts