Router-only knobs (lambda_balance/lambda_proc/lambda_entropy/
gumbel_tau_start/_end) and WGAN-only knobs (n_critic/gp_weight) were
being logged to wandb's top-level run config unconditionally, even for
runs where routing or WGAN mode is off, implying hyperparameters from
an inactive code path. Extract _wandb_run_config and only include each
group when its gate is actually true (router.enabled / mode=="wgan");
the full model_config (with its router sub-dict) is still always
logged in full, so no information is lost.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Trains the routed trunk's forward combination as a hard one-hot sample
(matching eval-time top-1 dispatch exactly) while keeping a smooth gradient
on the backward pass, targeting the train/eval mismatch identified as a
likely contributor to experts overlapping instead of partitioning in the
first energy-router rollout benchmark. Off by default (model.router.gumbel);
existing routed configs/checkpoints are unaffected.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
lars