diff --git a/megatron/core/optimizer_param_scheduler.py b/megatron/core/optimizer_param_scheduler.py
index da7e0787..e515a802 100644
--- a/megatron/core/optimizer_param_scheduler.py
+++ b/megatron/core/optimizer_param_scheduler.py
@@ -53,6 +53,7 @@ class OptimizerParamScheduler:
         override_opt_param_scheduler: Optional[bool] = False,
         wsd_decay_steps: Optional[int] = None,
         lr_wsd_decay_style: Optional[str] = None,
+        stablelm2_scheduler_config=None,
     ) -> None:
 
         # Class values.
@@ -90,6 +91,15 @@ class OptimizerParamScheduler:
             assert not self.use_checkpoint_opt_param_scheduler, (
                 'both override and ' 'use-checkpoint are set.'
             )
+        
+        self.stablelm2_scheduler_config = stablelm2_scheduler_config
+        if self.stablelm2_scheduler_config is not None:
+          ## absolute samples
+          self.stablelm2_scheduler_config.rsqrt_samples += \
+              self.stablelm2_scheduler_config.cosine_samples
+          ## N of consine
+          if self.stablelm2_scheduler_config.cosine_period_samples == 0:
+            self.stablelm2_scheduler_config.cosine_period_samples = self.lr_decay_steps
 
         # Set the learning rate
         self.step(0)
@@ -150,6 +160,62 @@ class OptimizerParamScheduler:
             lr = max_lr * warmup_steps**0.5 / (num_steps**0.5)
             return max(min_lr, lr)
 
+        # stablelm2 scheduler of multiple stages
+        if self.stablelm2_scheduler_config is not None:
+            log_single_rank(logger, logging.INFO, f"> stablelm2_scheduler_config: {self.stablelm2_scheduler_config}")
+            if self.num_steps <= self.stablelm2_scheduler_config.cosine_samples:
+                ## cosine phase
+                # decay_ratio = float(self.num_steps) / float(self.lr_decay_steps)
+                # TODO
+                decay_ratio = float(self.num_steps) / float(self.stablelm2_scheduler_config.cosine_period_samples)
+                cosine_min_lr = self.stablelm2_scheduler_config.cosine_max_lr * 0.1
+                delta_lr = self.stablelm2_scheduler_config.cosine_max_lr - cosine_min_lr
+                coeff = 0.5 * (math.cos(2 * math.pi * decay_ratio) + 1.0)
+                self.stablelm2_scheduler_config.cosine_lr = cosine_min_lr + coeff * delta_lr
+                return self.stablelm2_scheduler_config.cosine_lr
+            elif self.num_steps <= self.stablelm2_scheduler_config.rsqrt_samples:
+                ## rsqrt phase
+                alpha = self.stablelm2_scheduler_config.alpha
+                beta = self.stablelm2_scheduler_config.beta
+                gbs = self.stablelm2_scheduler_config.global_batch_size * 1.0
+                self.stablelm2_scheduler_config.rsqrt_lr = alpha / ((self.num_steps / gbs + beta) ** 0.5)
+                return self.stablelm2_scheduler_config.rsqrt_lr
+            elif self.stablelm2_scheduler_config.decay_samples <= 0:
+                ## optional linear phase
+                decay_steps_ = self.lr_decay_steps - self.stablelm2_scheduler_config.rsqrt_samples
+                num_steps_ = self.num_steps - self.stablelm2_scheduler_config.rsqrt_samples
+                decay_ratio = float(num_steps_) / float(decay_steps_)
+                coeff = (1.0 - decay_ratio)
+                return coeff * self.stablelm2_scheduler_config.rsqrt_lr
+            else:
+                ## optional linear phase
+                valid_lr_decay_steps_ = min(
+                    self.lr_decay_steps,
+                    self.stablelm2_scheduler_config.rsqrt_samples + self.stablelm2_scheduler_config.decay_samples)
+                if self.num_steps <= valid_lr_decay_steps_:
+                    decay_steps_ = valid_lr_decay_steps_ - self.stablelm2_scheduler_config.rsqrt_samples
+                    num_steps_ = self.num_steps - self.stablelm2_scheduler_config.rsqrt_samples
+                    decay_ratio = float(num_steps_) / float(decay_steps_)
+                    coeff = (1.0 - decay_ratio)
+                    delta_lr = self.stablelm2_scheduler_config.rsqrt_lr - self.min_lr
+                    assert decay_ratio >= 0.0
+                    return coeff * delta_lr + self.min_lr
+                else:
+                    return self.min_lr
+
+        # Warmup-Stable-Decay(WSD)
+        if self.lr_decay_style == 'warmup-stable-decay':
+            W = self.lr_warmup_steps
+            S = round((self.lr_decay_steps - W) * 10. / 11.)
+            ## D is 10% of S.
+            T = self.lr_decay_steps - W - S
+            ## Warmup Phase, see above
+            ## Stable Phase
+            if self.num_steps < S:
+                return self.max_lr
+            else: # Decay Phase
+                return self.max_lr * 0.5 ** ((self.num_steps - S) / T)
+
         num_steps_ = self.num_steps - self.lr_warmup_steps
         decay_steps_ = self.lr_decay_steps - self.lr_warmup_steps
         decay_ratio = float(num_steps_) / float(decay_steps_)
