diff --git a/tests/unit/runtime/half_precision/test_bf16.py b/tests/unit/runtime/half_precision/test_bf16.py
index 740fa30641a1..3f551fb0fd4a 100644
--- a/tests/unit/runtime/half_precision/test_bf16.py
+++ b/tests/unit/runtime/half_precision/test_bf16.py
@@ -27,6 +27,7 @@ def test(self, zero_stage=2, use_cpu_offload=False):
             pytest.skip("cpu-adam is not compatible")
 
         config_dict = {
+            "train_micro_batch_size_per_gpu": 1,
             "steps_per_print": 1,
             "optimizer": {
                 "type": "Adam",
@@ -87,7 +88,7 @@ def test(self, zero_stage=2, use_cpu_offload=False):
             pytest.skip("cpu-adam is not compatible")
 
         config_dict = {
-            "train_batch_size": 4,
+            "train_micro_batch_size_per_gpu": 4,
             "steps_per_print": 1,
             "fp16": {
                 "enabled": False,
@@ -180,7 +181,7 @@ def test(self, optimizer_constructor, zero_stage=2):
             )
 
         config_dict = {
-            "train_batch_size": 2,
+            "train_micro_batch_size_per_gpu": 2,
             "steps_per_print": 1,
             "fp16": {
                 "enabled": False
@@ -209,7 +210,7 @@ def test(self):
             )
 
         config_dict = {
-            "train_batch_size": 2,
+            "train_micro_batch_size_per_gpu": 2,
             "steps_per_print": 1,
             "optimizer": {
                 "type": "Adam",
@@ -258,7 +259,7 @@ def test(self, stage=2):
             )
 
         config_dict = {
-            "train_batch_size": 1,
+            "train_micro_batch_size_per_gpu": 1,
             "steps_per_print": 1,
             "fp16": {
                 "enabled": False
@@ -301,7 +302,7 @@ def test(self, comp_type, comm_type):
         type_str = {torch.float16: "fp16", torch.bfloat16: "bfp16"}
 
         config_dict = {
-            "train_batch_size": 2,
+            "train_micro_batch_size_per_gpu": 2,
             "steps_per_print": 1,
             "fp16": {
                 "enabled": comp_type == torch.float16