diff --git a/mlperf_logging/compliance_checker/training_5.1.0/closed_llama31_405b.yaml b/mlperf_logging/compliance_checker/training_5.1.0/closed_llama31_405b.yaml
index c47fd87..90e2d45 100644
--- a/mlperf_logging/compliance_checker/training_5.1.0/closed_llama31_405b.yaml
+++ b/mlperf_logging/compliance_checker/training_5.1.0/closed_llama31_405b.yaml
@@ -24,12 +24,15 @@
     REQ:   EXACTLY_ONE
 
 - KEY:
-    NAME:  opt_learning_rate_decay_steps
+    NAME:  opt_learning_rate_warmup_steps
     REQ:   EXACTLY_ONE
+    POST: >
+        s['opt_learning_rate_warmup_steps'] = math.ceil(8000 * 1152 / s['global_batch_size'] )
 
 - KEY:
-    NAME:  opt_learning_rate_warmup_steps
+    NAME:  opt_learning_rate_decay_steps
     REQ:   EXACTLY_ONE
+    CHECK: " v['value'] == math.ceil(1_200_000 * 1152 / s['global_batch_size'] ) - s['opt_learning_rate_warmup_steps'] "
 
 - KEY:
     NAME:  opt_learning_rate_decay_schedule
diff --git a/mlperf_logging/compliance_checker/training_5.1.0/closed_llama31_8b.yaml b/mlperf_logging/compliance_checker/training_5.1.0/closed_llama31_8b.yaml
index 2aba722..38cb96f 100644
--- a/mlperf_logging/compliance_checker/training_5.1.0/closed_llama31_8b.yaml
+++ b/mlperf_logging/compliance_checker/training_5.1.0/closed_llama31_8b.yaml
@@ -23,13 +23,15 @@
     REQ:   EXACTLY_ONE
 
 - KEY:
-    NAME:  opt_learning_rate_decay_steps
+    NAME:  opt_learning_rate_warmup_steps
     REQ:   EXACTLY_ONE
-    CHECK: " v['value'] == 1200000 " 
+    POST: >
+        s['opt_learning_rate_warmup_steps'] = v['value']
 
 - KEY:
-    NAME:  opt_learning_rate_warmup_steps
+    NAME:  opt_learning_rate_decay_steps
     REQ:   EXACTLY_ONE
+    CHECK: " v['value'] == math.ceil(1_200_000 / s['global_batch_size'] ) - s['opt_learning_rate_warmup_steps'] "
 
 - KEY:
     NAME:  opt_learning_rate_decay_schedule