fixed issue with datatype

argonne-lcf · Feb 6, 2025 · a4372cd · a4372cd
1 parent e6a4697
commit a4372cd
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 7 deletions.
diff --git a/dlio_benchmark/checkpointing/base_checkpointing.py b/dlio_benchmark/checkpointing/base_checkpointing.py
@@ -34,7 +34,7 @@ def get_datatype_size(datatype):
     elif datatype == "fp64":
         return 8
     else:
-        raise Exception("Unsupported datatype")
+        raise Exception("Unsupported datatype {datatype}")
 
 class BaseCheckpointing(ABC):
 

diff --git a/dlio_benchmark/checkpointing/pytorch_checkpointing.py b/dlio_benchmark/checkpointing/pytorch_checkpointing.py
@@ -31,8 +31,10 @@ def get_torch_datatype(datatype):
         return torch.float16
     if datatype == "f64":
         return torch.float64
-    if datatype == "i8":
+    if datatype == "int8":
         return torch.int8
+    if datatype == "uint8":
+        return torch.uint8
     if datatype == "bf16": # bfloat16
         return torch.bfloat16
 
@@ -54,7 +56,7 @@ def __init__(self):
         super().__init__("pt")
 
     @dlp.log
-    def get_tensor(self, size, datatype="i8"):
+    def get_tensor(self, size, datatype="int8"):
         return torch.ones(size, dtype=get_torch_datatype(datatype))
 
     @dlp.log

diff --git a/dlio_benchmark/checkpointing/tf_checkpointing.py b/dlio_benchmark/checkpointing/tf_checkpointing.py
@@ -31,8 +31,10 @@ def get_tf_datatype(datatype):
         return tf.dtypes.float16
     if datatype == "f64":
         return tf.dtypes.float64
-    if datatype == "i8":
+    if datatype == "int8":
         return tf.dtypes.int8
+    if datatype == "uint8":
+        return tf.dtypes.uint8
     if datatype == "bf16": # bfloat16
         return tf.dtypes.bfloat16
 
@@ -54,7 +56,7 @@ def __init__(self):
         super().__init__("pb")
 
     @dlp.log
-    def get_tensor(self, size, datatype="i8"):
+    def get_tensor(self, size, datatype="int8"):
         return tf.random.uniform((size), maxval=100, dtype=get_tf_datatype(datatype))
 
     @dlp.log

diff --git a/dlio_benchmark/utils/config.py b/dlio_benchmark/utils/config.py
@@ -101,8 +101,8 @@ class ConfigArguments:
     epochs_between_evals: int = 1
     checkpoint_type: CheckpointLocationType = CheckpointLocationType.RANK_ZERO
     checkpoint_mechanism: CheckpointMechanismType = CheckpointMechanismType.NONE
-    checkpoint_model_datatype: str = "f16"
-    checkpoint_optimizer_datatype: str = "f32"
+    checkpoint_model_datatype: str = "fp16"
+    checkpoint_optimizer_datatype: str = "fp32"
     model_size: int = 10240
     vocab_size: int = 32000
     hidden_size: int = 2048