Add MoE load balancing loss to distillation

Author: JamesDeng42

src/maxtext/models/models.py

@@ -505,6 +505,8 @@ def __call__(
       mutable_collections.append("intermediates")
     if self.config.distill_beta > 0.0 and "intermediates" not in mutable_collections:
       mutable_collections.append("intermediates")
+    if self.config.load_balance_loss_weight > 0.0 and "intermediates" not in mutable_collections:
+      mutable_collections.append("intermediates")
 
     if self.config.pure_nnx_decoder:
       logits, hidden_state, kv_caches = self.decoder(

src/maxtext/models/qwen3.py

@@ -1055,7 +1055,7 @@ def __call__(
     # We sow the load balancing loss so it can be collected and added to the total loss
     # during training.
     if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None:
-      self.sow("intermediates", "moe_lb_loss", load_balance_loss)
+      self.moe_lb_loss = nnx.Intermediate(load_balance_loss)
 
     # Final residual connection (after the MoE block)
     layer_output = residual + mlp_output
@@ -1299,7 +1299,7 @@ def __call__(
     mlp_lnx, load_balance_loss, _ = self.moe_block(hidden_states)
     mlp_lnx = nn.with_logical_constraint(mlp_lnx, self.activation_axis_names)
     if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None:
-      self.sow("intermediates", "moe_lb_loss", load_balance_loss)
+      self.moe_lb_loss = nnx.Intermediate(load_balance_loss)
 
     layer_output = intermediate_inputs + mlp_lnx
     layer_output = nn.with_logical_constraint(layer_output, self.activation_axis_names)

src/maxtext/trainers/post_train/distillation/distillation_utils.py

@@ -52,6 +52,8 @@ class DistillationForwardOutput:
   logits: jax.Array
   #: out_projection_activations
   out_projection_activations: jax.Array | None = None
+  #: moe load balance loss
+  moe_lb_loss: jax.Array | None = None
 
 
 @flax.struct.dataclass(frozen=True)
@@ -373,13 +375,19 @@ def compute_loss(
 
     total_loss = base_logit_loss + feature_loss
 
+    moe_lb_loss = jnp.array(0.0)
+    if student_output.moe_lb_loss is not None:
+      moe_lb_loss = student_output.moe_lb_loss
+      total_loss += moe_lb_loss
+
     # 4. Return Loss AND Metrics
     metrics = {
         "distill/soft_loss": soft_loss,
         "distill/hard_loss": hard_loss,
         "distill/kl_div": jnp.mean(kl_div, where=mean_mask),
         "distill/teacher_loss": teacher_hard_loss,
         "distill/out_proj_feature_loss": feature_loss,
+        "distill/moe_lb_loss": moe_lb_loss,
         "distill/total_loss": total_loss,
         "distill/temperature": self.temperature,
         "distill/alpha": self.alpha,

src/maxtext/trainers/post_train/distillation/train_distill.py

@@ -150,8 +150,14 @@ def model_forward_fn(
     if config.distill_beta > 0.0:
       out_projection_activations = maxtext_utils.get_intermediate_value(model, "out_projection_activations", clear=True)
 
+    moe_lb_loss = None
+    if config.num_experts > 1 and config.load_balance_loss_weight > 0.0:
+      total_moe_lb_loss = maxtext_utils.get_intermediate_value(model, "moe_lb_loss", clear=False)
+      if total_moe_lb_loss is not None:
+        moe_lb_loss = jnp.mean(jnp.array(total_moe_lb_loss))
+
     retval = distillation_utils.DistillationForwardOutput(
-        logits=logits, out_projection_activations=out_projection_activations
+        logits=logits, out_projection_activations=out_projection_activations, moe_lb_loss=moe_lb_loss
     )
     return retval
 

src/maxtext/utils/maxtext_utils.py

@@ -1081,6 +1081,29 @@ def get_intermediate_value(model, nested_key, default=None, clear=False):
         intermediate_value = model.decoder.layers["self_attention"][nested_key].get_value()[-1]
         if clear:
           del model.decoder.layers["self_attention"][nested_key]
+    case "moe_lb_loss":
+      # pylint: disable=import-outside-toplevel
+      from flax import nnx
+
+      losses = []
+      for path, val in nnx.state(model).flat_state():
+        v = val.value if hasattr(val, "value") else val
+        if isinstance(v, dict) and "moe_lb_loss" in v:
+          losses.append(v["moe_lb_loss"][-1])
+          if clear:
+            v.pop("moe_lb_loss", None)
+        elif "moe_lb_loss" in path:
+          losses.append(v[-1] if isinstance(v, (list, tuple)) else v)
+          if clear:
+            curr = model
+            try:
+              for p in path[:-1]:
+                curr = getattr(curr, p)
+              delattr(curr, path[-1])
+            except AttributeError:
+              pass
+      if losses:
+        intermediate_value = losses
     case _:
       # Default case to handle any unknown nested keys
       raise ValueError(f"Incorrect nested_key: {nested_key}")

tests/post_training/unit/train_distill_test.py

@@ -437,6 +437,7 @@ def _test_monitored_strategy(self, *, sft_mode: bool, feature_loss_type: Literal
         "distill/kl_div",
         "distill/teacher_loss",
         "distill/out_proj_feature_loss",
+        "distill/moe_lb_loss",
         "distill/total_loss",
         "distill/temperature",
         "distill/alpha",