增加路由器参数

megatron/training/theoretical_memory_usage.py

@@ -59,6 +59,8 @@ def compute_weight_and_optimizer_memory(args, verbose=False):
             )
             # routed experts.
             + (2 * (args.ffn_hidden_size * args.hidden_size) * num_experts * gated_linear_multiplier)
+            # router
+            + args.hidden_size * num_experts
             # shared experts.
             + (2 * args.moe_shared_expert_intermediate_size * args.hidden_size)
             # Transformer layernorms.