File size: 1,252 Bytes
ebad435 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | 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()
|