Commit 61873ad

authored and

Orbax Authors

committed

#v1 #safetensors Refactor SafetensorsLayout to eliminate unperformant multi-host loading code. On a single-host, a single read operation is performed, followed by device_put to local devices.

PiperOrigin-RevId: 893168206

1 parent 09d2982 commit 61873adCopy full SHA for 61873ad

2 files changed

+160

-254

lines changed

checkpoint/orbax/checkpoint/experimental/v1/_src/layout
- safetensors_layout.py
- safetensors_layout_test.py

2 files changed

+160

-254

lines changed

Comments

(0)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Commit 61873ad

2 files changed

2 files changed

File tree

2 files changed

2 files changed

0 commit comments