feat: add 3FS copier support

[ABNER-1]

feat: add 3FS copier support

Summary

Add support for DeepSeek 3FS USRBIO as a new copier backend for high-performance safetensors file loading. This enables fastsafetensors to leverage 3FS's zero-copy I/O and asynchronous read capabilities when loading model weights from 3FS-mounted filesystems.

resolve #55

Motivation

3FS (Fire-Flyer File System) is a high-performance distributed filesystem designed for AI/ML workloads. By integrating 3FS USRBIO as a copier backend, fastsafetensors can achieve significantly higher I/O throughput when model weights are stored on 3FS-mounted paths, complementing the existing GDS and NoGDS copier backends.

Changes

New Files

• fastsafetensors/copier/threefs.py: Core 3FS copier implementation

  • ThreeFSFileCopier: Implements CopierInterface using 3FS USRBIO for file reads
  • new_threefs_file_copier(): Factory function registered as "3fs" copier type via the copier registry
  • Uses fastsafetensor-3fs-reader package for the actual 3FS I/O operations

• fastsafetensors/threefs_loader.py: High-level loader APIs for 3FS

  • ThreeFSLoader: Single-file loader extending BaseSafeTensorsFileLoader
  • ParallelThreeFSLoader: Pipeline-parallel loader extending PipelineParallel for batch loading multiple safetensors files

Modified Files

• fastsafetensors/copier/gds.py & fastsafetensors/copier/nogds.py: Added **kwargs to new_gds_file_copier() and new_nogds_file_copier() for forward-compatible copier constructor signatures
• fastsafetensors/copier/__init__.py: Export ThreeFSFileCopier
• fastsafetensors/parallel_loader.py: Minor logging improvements (added num_keys to batch summary)
• pyproject.toml: Added threefs optional dependency group (fastsafetensor-3fs-reader>=0.1.0)

Usage

# Single file loading
from fastsafetensors.threefs_loader import ThreeFSLoader

loader = ThreeFSLoader(pg=None, device="cuda:0", mount_point="/mnt/3fs")
loader.add_filenames({0: ["/mnt/3fs/model.safetensors"]})
bufs = loader.copy_files_to_device()
tensor = bufs.get_tensor("weight")
loader.close()

# Parallel loading (iterate_weights interface)
from fastsafetensors.threefs_loader import ParallelThreeFSLoader

loader = ParallelThreeFSLoader(
    pg=None,
    hf_weights_files=["/mnt/3fs/model-00001.safetensors", "/mnt/3fs/model-00002.safetensors"],
    device="cuda:0",
)
for key, tensor in loader.iterate_weights():
    # Process each tensor
    pass
loader.close()

Dependencies

This feature depends on the fastsafetensor-3fs-reader package, which is declared as an optional dependency:

pip install fastsafetensors[threefs]

If the 3FS reader package is not installed, the 3FS copier will not be available but all other functionality remains unaffected.

Testing

Tests are provided under tests/threefs/, automatically skipped when fastsafetensor_3fs_reader is not available. To run locally with mock backend:

FASTSAFETENSORS_BACKEND=mock TEST_FASTSAFETENSORS_FRAMEWORK=pytorch \
    pytest tests/threefs/ -v

Tests also run in real 3FS environments without any special environment variable -- the skip condition checks whether fastsafetensor_3fs_reader is importable.

tests/threefs/test_threefs.py -- Unit tests for core 3FS components (uses gpt2 model data)

• test_threefs_copier_registered: Verifies ThreeFSFileCopier is registered in the copier registry
• test_ThreeFSFileCopier: Tests the copier submit_io + wait_io with tensor data verification
• test_ThreeFSLoader: Tests single-file loading via ThreeFSLoader
• test_ThreeFSLoader_multiple_files: Tests multi-file loading with full tensor correctness

tests/threefs/test_parallel_threefs.py -- Integration tests for ParallelThreeFSLoader (uses example safetensors files)

• test_parallel_single_file: Single-file iteration with tensor data correctness verification
• test_parallel_multiple_files: Multi-file iteration with shape, dtype, and value verification
• test_parallel_close_and_memory_release: Resource cleanup and bounce buffer release after full iteration
• test_parallel_close_without_iterate: Graceful cleanup when closing without iterating

All 8 tests pass (~1.7s).

[takeshi-yoshimura]

@ABNER-1
Thanks for contributing your excellent work!
However, I’m a bit busy this week, so I’ll review this after April 13. please be patient. thanks.

[takeshi-yoshimura]

I expected that 3fs exposes generic, low-level python interfaces, but this seems specific to fastsafetensors. Does 3fs have plans to expose such interfaces? If so, this repository should contain ThreeFSFileReader.

[ABNER-1]

The 3fs reader repo is here: https://github.com/ABNER-1/fastsafetensor_3fs_reader.

It mainly uses pyusrbio for the implementation and includes tools to find optimal 3fs download/loading parameters.
I'm not sure if these files belong in fastsafetensors—let's discuss.

The repo is currently under my personal account; we can also consider moving it to a more appropriate org later.

[ABNER-1]

Some files were modified by pre-commit hooks. I'll revert them manually.

[takeshi-yoshimura]

@ABNER-1
I want to make sure who maintains the code for fastsafetensors_3fs_reader. I expect some parts can be generic for other use-cases and they should be maintained by external parties such as 3fs maintainers, but codes specific to fastsafetensors should be maintained in this repo.

Let me spend some more time to learn fastsafetensors_3fs_reader. Then, I will propose how to deal with this issue.