google / google/paxml

Use bfloat16 for eval

Open
#66 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
560
Forks
72
PR merge metrics
No merged PRs in 30d

Description

I'm running paxml on an Intel Xeon CPU server using the paxml/main.py program. I'm trying to create a model that creates weights in bfloat16, and uses that datatype during eval. I modified the LmCloudSpmd2B configuration with the following lines:
```
MODEL_DTYPE = jnp.bfloat16
ICI_MESH_SHAPE = [1, 1, 1]
```

The training status output includes the following output.
```
model.dtype : type/jax.numpy/float32
model.fprop_dtype : dtype[bfloat16]
```
All of the other operator datatypes are float32. When I run that model with the --eval switch all of the computation is in float32. How can I direct paxml to use bfloat16?

Tom

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.