diff --git a/llmc/schedulers.h b/llmc/schedulers.h
index 44e9562..9ddc570 100644
--- a/llmc/schedulers.h
+++ b/llmc/schedulers.h
@@ -59,6 +59,26 @@ float get_learning_rate_constant(LearningRateScheduler *scheduler, int step) {
     return scheduler->learning_rate;
 }
 
+// wsd schedule: warmup linearly, keep constant, last 20% decay using 1 - sqrt decay to final_frac (should be 0.0)
+// https://arxiv.org/abs/2405.18392
+float get_learning_rate_wsd(LearningRateScheduler *scheduler, int step) {
+    int decay_point = (int)(0.8f * scheduler->train_num_batches);
+    float max_lr = scheduler->learning_rate;
+    float lr = max_lr;
+    if (step < scheduler->warmup_iterations) {
+        float decay_ratio = ((float)(step + 1)) / scheduler->warmup_iterations;
+        lr = max_lr * decay_ratio;
+    } else if (step < decay_point) {
+        // noop, keep lr constant
+    } else {
+        float decay_ratio = ((float)(step - decay_point)) / (scheduler->train_num_batches - decay_point);
+        assert(0.0f <= decay_ratio && decay_ratio <= 1.0f);
+        float min_lr = max_lr * scheduler->final_learning_rate_frac;
+        return min_lr + (1.0f - sqrtf(decay_ratio)) * (max_lr - min_lr);
+    }
+    return lr;
+}
+
 // return the learning rate at a given step
 float get_learning_rate(LearningRateScheduler *scheduler, int step) {
     float step_learning_rate;
@@ -68,11 +88,13 @@ float get_learning_rate(LearningRateScheduler *scheduler, int step) {
         step_learning_rate = get_learning_rate_linear(scheduler, step);
     } else if (strcmp(scheduler->type, "constant") == 0) {
         step_learning_rate = get_learning_rate_constant(scheduler, step);
+    } else if (strcmp(scheduler->type, "wsd") == 0) {
+        step_learning_rate = get_learning_rate_wsd(scheduler, step);
     } else {
-        printf("Unknown learning rate scheduler type\n");
+        fprintf(stderr, "Unknown learning rate scheduler type: %s\n", scheduler->type);
         exit(EXIT_FAILURE);
     }
     return step_learning_rate;
 }
 
-#endif // SCHEDULERS_H
\ No newline at end of file
+#endif // SCHEDULERS_H
diff --git a/train_gpt2.cu b/train_gpt2.cu
index 961dbc3..e07b431 100644
--- a/train_gpt2.cu
+++ b/train_gpt2.cu
@@ -21,7 +21,7 @@ GPT-2 Transformer Neural Net training loop. See README.md for usage.
 #include "llmc/dataloader.h"
 // defines: manual_seed, normal_ (same as torch.manual_seed and torch.normal)
 #include "llmc/rand.h"
-// defines learning rate schedulers
+// defines: lr_scheduler_init, get_learning_rate
 #include "llmc/schedulers.h"
 // defines: sample_softmax, random_f32
 #include "llmc/sampler.h"