Commit 2fbebca

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: 893720776

1 parent 39262e1 commit 2fbebcaCopy full SHA for 2fbebca

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 2fbebca

2 files changed

2 files changed

File tree

2 files changed

2 files changed

0 commit comments