feat: add direct io backend

[McPatate]

Experimental changes to file loading

Warning

sorry for the messy thought dump you're about to read

Still measuring performance benefits. Probably implemented the wrong approach (full file loading in host memory), the goal with this impl is to try saturating disk read bandwidth on an nvme.

Running pytest benches/test_pt.py -k test_pt_sf_load_direct yields poor performance atm, but it's hard to compare to say test_pt_sf_load_cpu given os page cache is populated when writing the file to disk, which is done right before the benchmark(load_file, ...) call.

I want to continue experimenting with this loading method, perhaps instead imitating what is done in fastsafetensors' nogds_file_reader where they have a bounce buffer allocated with cudaHostAlloc from which they run cudaMemcpy once they pread memory in said buffer.

I'd also like to try playing around with creating a safetensors-distributed package that'd be built on top of this infrastructure layer for fast loading of files in tensor parallel scenarios. We'll probably want to port over some of the code in transformers in that package.

I'm not convinced we want to imitate fastsafetensors fully, where they scatter and broadcast tensors from gpu memory leveraging nvlink (iiuc). But at the same time, I'm not sure either we can slice the file like we currently do in our distributed weight loader in transformers while reading file in the bounce buffer, we want reads to be aligned when reading with O_DIRECT and this will get messy when sending over bytes to device. Need to investigate some more 🤕

References

• fastsafetensors paper
• fastsafetensors repo

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[danieldk]

I might be mistaken, but I think Vec uses the alignment of the type, so you can do this more safely using something like:

let mut vec: Vec<[u8; ALIGN]> = Vec::with_capacity((file_len + ALIGN - 1) / ALIGN);
// Use vec as u8 rather than [u8; ALIGN].

[danieldk]

Question for understanding, since region_bytes = blocks_per_thread * block_size. Suppose we have 16 blocks in the file and 4 threads. Then the responsibility by thread is:

0 0 0 0 1 1 1 1 2 2 2 2 3 3 3 3

I am curious how this compares to patterns like:

0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3