File size: 1,652 Bytes
7f5d56d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 | 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
|