Update PyTorch setup for tf32 and change default WMT workload compute dtype to FLOAT32

Author: rka97

algoperf/pytorch_utils.py

@@ -20,6 +20,7 @@
 
 
 def pytorch_setup() -> Tuple[bool, int, torch.device, int]:
+  torch.set_float32_matmul_precision('high')
   use_pytorch_ddp = 'LOCAL_RANK' in os.environ
   rank = int(os.environ['LOCAL_RANK']) if use_pytorch_ddp else 0
   device = torch.device(f'cuda:{rank}' if torch.cuda.is_available() else 'cpu')

algoperf/workloads/wmt/workload.py

@@ -22,7 +22,7 @@ class BaseWmtWorkload(spec.Workload):
   """A WMT workload."""
 
   _vocab_size: int = 32000
-  _compute_dtype: spec.DTYPE = spec.DTYPE.BFLOAT16
+  _compute_dtype: spec.DTYPE = spec.DTYPE.FLOAT32
   _param_dtype: spec.DTYPE = spec.DTYPE.FLOAT32
 
   def __init__(self) -> None:

pyproject.toml

@@ -96,6 +96,7 @@ jax_core_deps = [
   "chex==0.1.86",
   "ml_dtypes==0.5.1",
   "protobuf==4.25.5",
+  "jmp",
 ]
 jax_cpu = [
   "jax==0.7.0",