Strange behavior of saving sharded trainstate in GCP.
- Dominant language
- Python
- Stars
- 535
- Forks
- 101
- Avg merge
- 1d 13h
- Merged PRs (30d)
- 23
Description
A user posted in the Flax discussions about an orbax discrepancy between different zones in GCE. Do different zones have different orbax versions?
==================================================================
# what happened
When I save my sharded state in `asia-northeast3-a` in GCE with orbax, the orbax create `/tmp/orbax_ckpt/0/_sharding` file which starts with
```
{"dropout_rng":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"data\", \"model\"], \"partition_spec\": []}","opt_state.0.0.count":"{\"sharding_type\": \"NamedSharding\", \"shape\": [2, 1], \"axis_names\": [\"data\", \"model\"], \"partition_spec\": []}",
...
```
My sharded state has "dropout_rng" state, so above file make sense.
However, when I run same script in other region like `asia-southeast1-b`, the orbax create `_sharding` file without proper layer names, for example,
```
{"ZHJvcG91dF9ybmc=":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"partition_spec\": []}","b3B0X3N0YXRlLjAuMC5jb3VudA==":"{\"sharding_type\": \"NamedSharding\", \"shape\": [1, 1], \"axis_names\": [\"data\", \"model\"], \"partition_spec\": []}",
...
```
# Theory
I doubt that this is related to OCDBT, because the only difference in between terminal outputs is `ocdbt is intitialized` in `asia-northeast3-a` but the other regions are not having this message
`type_handlers.py:223] OCDBT is initialized successfully.`.
I checked `tensorstore==0.1.51` in all region.
Anyone can help me please?
Thank you.
_Originally posted by @sw32-seo in https://github.com/google/flax/discussions/3538_
Contributor guide
Assessment
This issue has not been assessed yet.