From 11e37d5ec233f9178be259688ebc6f5738bac3fc Mon Sep 17 00:00:00 2001 From: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com> Date: Sat, 3 Oct 2026 18:40:50 +0100 Subject: [PATCH] Keep balanced routing auxiliary-loss gradients finite Signed-off-by: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com> --- vmoe/nn/routing.py | 8 ++-- vmoe/nn/routing_test.py | 81 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 85 insertions(+), 4 deletions(-) diff --git a/vmoe/nn/routing.py b/vmoe/nn/routing.py index d36307f..740cfba 100644 --- a/vmoe/nn/routing.py +++ b/vmoe/nn/routing.py @@ -151,10 +151,10 @@ def _gshard_auxiliary_loss(cls, gates: Array) -> Array: def _importance_auxiliary_loss(cls, gates: Array) -> Array: axis = tuple(range(gates.ndim - 1)) # All except last. importance_per_expert = jnp.sum(gates, axis=axis) - std_importance_per_expert = jnp.std(importance_per_expert) + variance_importance_per_expert = jnp.var(importance_per_expert) mean_importance_per_expert = jnp.mean(importance_per_expert) - # Compute coefficient of variation (i.e. std/mean) squared. - return (std_importance_per_expert / mean_importance_per_expert)**2 + # Compute CV squared directly, avoiding the undefined std gradient at zero. + return variance_importance_per_expert / mean_importance_per_expert**2 @classmethod def _load_auxiliary_loss(cls, logits: Array, logits_noisy: Array, @@ -180,7 +180,7 @@ def _load_auxiliary_loss(cls, logits: Array, logits_noisy: Array, # We compute the average such probability for each expert over examples. p_mean = jnp.mean(p, axis=0) # Compute p_mean's coefficient of variation squared. - return (jnp.std(p_mean) / jnp.mean(p_mean))**2 + return jnp.var(p_mean) / jnp.mean(p_mean) ** 2 class NoisyTopItemsPerExpertRouter(nn.Module): diff --git a/vmoe/nn/routing_test.py b/vmoe/nn/routing_test.py index ad7d3f6..aba1e06 100644 --- a/vmoe/nn/routing_test.py +++ b/vmoe/nn/routing_test.py @@ -19,6 +19,8 @@ import chex import jax import jax.numpy as jnp +import numpy as np +import optax from vmoe.nn import routing @@ -163,5 +165,84 @@ def test_forward_not_deterministic(self): chex.assert_trees_all_equal_comparator(different_fn, error_msg_fn, y1, y2) +class BalancedAuxiliaryGradientTest(parameterized.TestCase): + """Balanced expert statistics have a smooth zero-variance loss.""" + + @parameterized.parameters(1, 2, 4) + def test_importance_value_gradient_and_hessian_at_balance(self, experts): + gates = jnp.full((4, experts), 1.0 / experts) + loss = routing.NoisyTopExpertsPerItemRouter._importance_auxiliary_loss + for function in (jax.value_and_grad(loss), jax.jit(jax.value_and_grad(loss))): + value, gradient = function(gates) + self.assertEqual(float(value), 0.0) + np.testing.assert_array_equal(gradient, np.zeros_like(gates)) + direction = jnp.arange(gates.size, dtype=gates.dtype).reshape(gates.shape) + _, tangent = jax.jvp(jax.grad(loss), (gates,), (direction,)) + self.assertTrue(np.all(np.isfinite(tangent))) + + @parameterized.product(experts=(1, 3), selected=(1,)) + def test_balanced_load_has_finite_zero_gradients(self, experts, selected): + logits = jnp.zeros((4, experts)) + + def loss(values): + return routing.NoisyTopExpertsPerItemRouter._load_auxiliary_loss( + values, logits, noise_std=0.3, num_selected_experts=selected + ) + + value, gradient = jax.jit(jax.value_and_grad(loss))(logits) + self.assertEqual(float(value), 0.0) + np.testing.assert_array_equal(gradient, np.zeros_like(logits)) + + def test_nonuniform_balanced_gates_and_unbalanced_reference(self): + loss = routing.NoisyTopExpertsPerItemRouter._importance_auxiliary_loss + balanced = jnp.array([[0.9, 0.1], [0.1, 0.9]]) + np.testing.assert_array_equal(jax.grad(loss)(balanced), np.zeros((2, 2))) + gates = jnp.array( + [[0.5, 0.4, 0.1], [0.3, 0.3, 0.4], [0.1, 0.2, 0.7], [0.8, 0.2, 0.0]] + ) + totals = np.asarray(gates).sum(0) + expected = np.mean((totals - totals.mean()) ** 2) / totals.mean() ** 2 + np.testing.assert_allclose(loss(gates), expected, rtol=1e-5) + analytical = jax.grad(loss)(gates) + eps = 1e-3 + for i in range(gates.shape[0]): + for j in range(gates.shape[1]): + plus = gates.at[i, j].add(eps) + minus = gates.at[i, j].add(-eps) + finite_difference = (loss(plus) - loss(minus)) / (2 * eps) + np.testing.assert_allclose( + analytical[i, j], finite_difference, rtol=1e-3, atol=1e-5 + ) + + @parameterized.parameters(True, False) + def test_real_router_balanced_initialization_allows_training_update( + self, deterministic + ): + layer = routing.NoisyTopExpertsPerItemRouter( + num_experts=3, + num_selected_experts=2, + deterministic=deterministic, + dispatcher={'name': 'einsum', 'batch_priority': False, 'capacity': 4}, + ) + inputs = jnp.ones((3, 4, 2)) + params = {'dense': {'kernel': jnp.zeros((2, 3))}} + key = jax.random.PRNGKey(13) + + def objective(p): + dispatcher, metrics = layer.apply({'params': p}, inputs, rngs={'gating': key}) + dispatched = dispatcher.dispatch(inputs) + combined = dispatcher.combine(dispatched * 2.0) + return jnp.mean(combined**2) + metrics['auxiliary_loss'].mean() + + value, grads = jax.jit(jax.value_and_grad(objective))(params) + self.assertTrue(np.isfinite(value)) + for gradient in jax.tree.leaves(grads): + self.assertTrue(np.all(np.isfinite(gradient))) + optimizer = optax.sgd(1e-3) + updates, _ = optimizer.update(grads, optimizer.init(params), params) + updated = optax.apply_updates(params, updates) + self.assertTrue(np.isfinite(objective(updated))) + + if __name__ == '__main__': absltest.main()