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>
Extends default_out_dir_name's non-default-field convention to the
router's new gumbel combine-weight flag and its learnable-knob toggles,
so gumbel sweep configs (learn_centers on/off, learn_width,
learn_temperature) resolve to distinguishable checkpoint directory
names instead of colliding on the same r-<type><n> token.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Extends the 10-expert EnergyRouter + physical-conditioning benchmark
config with the new opt-in gumbel combine weights, isolating the
learn_centers/learn_temperature axis: none, centers only, and
centers+temperature.
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