| |
| |
| |
| |
| @@ -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 |
|
|
