| import unittest | |
| from scripts.train_confidence_token_lazy_ddp import learning_rate_at_step | |
| class ConfidenceLearningRateTest(unittest.TestCase): | |
| def value(self, step: int) -> float: | |
| return learning_rate_at_step( | |
| step=step, | |
| total_steps=1440, | |
| schedule="warmup_cosine", | |
| peak_lr=1e-4, | |
| warmup_steps=36, | |
| warmup_start_lr=1e-5, | |
| min_lr=1e-6, | |
| ) | |
| def test_warmup_endpoints(self) -> None: | |
| self.assertAlmostEqual(self.value(0), 1e-5) | |
| self.assertAlmostEqual(self.value(35), 1e-4) | |
| def test_cosine_is_monotonic_and_reaches_floor(self) -> None: | |
| values = [self.value(step) for step in range(35, 1440)] | |
| self.assertTrue(all(left >= right for left, right in zip(values, values[1:]))) | |
| self.assertAlmostEqual(values[-1], 1e-6) | |
| def test_constant_schedule_is_unchanged(self) -> None: | |
| value = learning_rate_at_step( | |
| step=503, | |
| total_steps=1440, | |
| schedule="constant", | |
| peak_lr=3e-4, | |
| warmup_steps=0, | |
| warmup_start_lr=1e-5, | |
| min_lr=1e-6, | |
| ) | |
| self.assertEqual(value, 3e-4) | |
| if __name__ == "__main__": | |
| unittest.main() | |