lars
b51eafcfa5
CI / Lint (ruff check) (push) Successful in 28s
CI / Format (ruff format) (push) Successful in 28s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 21s
CI / Lint (ruff check) (pull_request) Successful in 25s
CI / Format (ruff format) (pull_request) Successful in 30s
CI / Tests (push) Successful in 1m37s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 31s
CI / Tests (pull_request) Successful in 55s
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>
27 lines
835 B
Python
27 lines
835 B
Python
"""Tests for giant/train.py helpers."""
|
|
|
|
from giant.train import _gumbel_tau
|
|
|
|
|
|
def test_gumbel_tau_at_step_zero_is_start():
|
|
assert _gumbel_tau(0, 1000, 1.0, 0.1) == 1.0
|
|
|
|
|
|
def test_gumbel_tau_at_total_steps_is_end():
|
|
assert abs(_gumbel_tau(1000, 1000, 1.0, 0.1) - 0.1) < 1e-9
|
|
|
|
|
|
def test_gumbel_tau_interpolates_linearly_midway():
|
|
assert abs(_gumbel_tau(500, 1000, 1.0, 0.1) - 0.55) < 1e-9
|
|
|
|
|
|
def test_gumbel_tau_clamps_beyond_total_steps():
|
|
assert _gumbel_tau(5000, 1000, 1.0, 0.1) == _gumbel_tau(1000, 1000, 1.0, 0.1)
|
|
|
|
|
|
def test_gumbel_tau_handles_zero_total_steps():
|
|
# total_steps=0 is guarded to 1 internally: step=0 gives zero progress
|
|
# (still tau_start), any step>=1 immediately clamps to full progress.
|
|
assert _gumbel_tau(0, 0, 1.0, 0.1) == 1.0
|
|
assert abs(_gumbel_tau(1, 0, 1.0, 0.1) - 0.1) < 1e-9
|