Repository navigation
Internal Change - #3473
Merged
Merged
Internal Change#3473
Conversation
Restructure the safetensors loader around the target sharding: per host, resolve the byte runs each local-device shard needs from the file blob, coalesce neighbouring runs per file, and serve each coalesced block as a ranged read. An over-read ratio caps block_size/needed_bytes so inner-dim (row-parallel TP) shardings don't amplify cross-host egress to N*file_size; leading-dim FSDP stays one contiguous read per host. Bound peak in-flight memory by the shared restore budget (MemoryOptions.read_concurrent_bytes) rather than a per-block cap: one limiter across every file's scheduler streams oversized blocks in budget-sized chunks, so host memory beyond the destination shard buffers is bounded regardless of block, shard, or file size. Expose the loader tuning knobs (max_over_read_ratio) on SafetensorsOptions and the load benchmark, add a TP ratio-sweep config, and extend the single- and multi-process layout tests to cover the sharded read paths.
…r file Replace the per-file scheduler and its cursor-callback contract with a pure read plan: byte runs carry their destination buffer views and every chunk's buffer copies are precomputed, so chunks are order-independent. Blocks larger than a fixed chunk size (128 MiB) are now fetched as parallel ranged reads under the same shared in-flight byte budget, instead of streaming sequentially at single-stream throughput -- the single-host whole-file case especially. Read all file headers (and length checks / stats) concurrently instead of one round trip at a time, and assemble each file's tensors as soon as its reads complete, overlapping host-to-device transfer with the remaining reads and releasing host buffers earlier.
The generator referenced the pre-rename singular module path and the fixture URIs predate the st_fixtures/ convention; the sweep failed at generator import. Also point the pair reference at mistral7b_tp.yaml.
The 128 MiB chunk size that splits coalesced blocks into concurrent ranged reads was a module constant; the right value is storage-backend dependent (request-rate quotas and per-request overhead vs parallel stream throughput). SafetensorsOptions.read_chunk_bytes overrides it, clamped to the in-flight budget as before, and the load benchmark gains a sweepable safetensors_read_chunk_mb knob to measure the tradeoff.
copybara-service
Bot
force-pushed
the
test_947822041
branch
15 times, most recently
from
July 14, 2026 20:15
fb97517 to
e2a5d1e
Compare
copybara-service
Bot
force-pushed
the
test_947822041
branch
from
July 14, 2026 20:37
e2a5d1e to
549d786
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Internal Change
FUTURE_COPYBARA_INTEGRATE_REVIEW=#3455