Router auxiliary losses computed even when lambda = 0 #31

New Issue

2026-08-13T15:05:33+02:00

lars commented

2026-08-13 15:05:33 +02:00

Severity: low (wasted compute, no wrong results).

giant/training/trainers.py:618-622 computes balance_loss, classify_loss
and entropy_loss whenever a router exists, then adds each only if its lambda
is > 0. Every routed run therefore pays three extra router.gate() forward
passes per batch — including the default lambda_balance = 0.0 case, which is
exactly the configuration the failed 2026-07-22 benchmark ran.

Fix: guard the computation on the same > 0 condition that guards the
addition.

Migrated from issues.md (v0.3.0 branch review, 2026-08-13), Issue 6.

**Severity: low (wasted compute, no wrong results).** `giant/training/trainers.py:618-622` computes `balance_loss`, `classify_loss` and `entropy_loss` whenever a router exists, then adds each only if its lambda is `> 0`. Every routed run therefore pays three extra `router.gate()` forward passes per batch — including the default `lambda_balance = 0.0` case, which is exactly the configuration the failed 2026-07-22 benchmark ran. **Fix:** guard the computation on the same `> 0` condition that guards the addition. --- Migrated from `issues.md` (v0.3.0 branch review, 2026-08-13), Issue 6.

lars added the performance bug labels 2026-08-13 15:05:33 +02:00

lars commented

2026-08-13 16:18:25 +02:00

Fixed in f505fe7 on v0.3.0-stage2-autoregressive.

Guarded the three router auxiliary-loss computations in FlowDDPMStageTrainer._compute (balance_loss, classify_loss, entropy_loss) on the same lambda > 0 condition that already guarded their addition into the total loss, instead of computing all three unconditionally whenever a router exists. Confirmed WGANStageTrainer has no equivalent block, so this was the only call site. Added a test that mocks (wraps) the router's three loss methods and asserts zero calls at the default lambda=0 config and exactly one call each when lambdas are positive, so the guard is verified both ways. All four checks (pytest, ruff format, ruff check, ty check) pass clean.

Fixed in f505fe7 on v0.3.0-stage2-autoregressive. Guarded the three router auxiliary-loss computations in FlowDDPMStageTrainer._compute (balance_loss, classify_loss, entropy_loss) on the same lambda > 0 condition that already guarded their addition into the total loss, instead of computing all three unconditionally whenever a router exists. Confirmed WGANStageTrainer has no equivalent block, so this was the only call site. Added a test that mocks (wraps) the router's three loss methods and asserts zero calls at the default lambda=0 config and exactly one call each when lambdas are positive, so the guard is verified both ways. All four checks (pytest, ruff format, ruff check, ty check) pass clean.

lars closed this issue

2026-08-13 16:18:28 +02:00

lars referenced this issue from a commit

2026-08-13 16:21:01 +02:00

Skip router auxiliary loss compute when their lambda is 0 (gitea #31)

Sign in to join this conversation.

1 Participants

Notifications

Due Date

No due date set.

Dependencies

No dependencies set.

Reference: lars/giant#31