checkpointing/HTTPTransport: added streaming serialization and parallel transfer support

[d4l3k]

This makes HTTP recovery much much faster with a few key changes:

1. use RWLock so multiple readers can be fetching the state_dict at the same time
2. use _streaming_save/_streaming_load when available which is ~2x faster on average and avoids the 2x memory overhead (requires PT nightly)
3. support parallel transfers by using pytree to divy up the leaf values into chunks which can be sent via parallel HTTP requests

This doesn't change the behavior of Manager by default as it seems that using parallel chunks w/ torch.save/load can actually increase time significantly as PyTorchStreamReader holds the GIL during deserialization.

The optimal config is with streaming and chunking enabled

Test plan:

pytest

Testing with 12GB total and 1MB tensors

pytorch nightly

$ python torchft/checkpointing/http_transport.py --num-chunks 0
INFO:__main__:fetching checkpoint took 6.626614563167095s
$ python torchft/checkpointing/http_transport.py --num-chunks 10
INFO:__main__:fetching checkpoint took 3.0460726767778397s
$ python torchft/checkpointing/http_transport.py --num-chunks 0 --device cuda
INFO:__main__:fetching checkpoint took 6.147395346313715s
$ python torchft/checkpointing/http_transport.py --num-chunks 10 --device cuda
INFO:__main__:fetching checkpoint took 2.9234009198844433s

pytorch 2.6.0

$ python torchft/checkpointing/http_transport.py --num-chunks 0
INFO:__main__:fetching checkpoint took 17.019980508834124s
$ python torchft/checkpointing/http_transport.py --num-chunks 10
INFO:__main__:fetching checkpoint took 40.383272521197796s

[H-Huang]

Looks good, thanks for the change!

python torchft/checkpointing/http_transport.py --num-chunks 0 --device cuda
INFO:main:fetching checkpoint took 12.673462141305208s

How come --num-chunks 0 with cuda is slower than CPU?

[H-Huang]

nit: maybe we could move the benchmarking code to its own folder

[d4l3k]

Moved to a _bench.py file and added a small test to make sure benchmark doesn't regress

[H-Huang]

nit: could you use tree_map here?

[d4l3k]

tree_map does the same flatten+unflatten so I think I'll just keep it like this to avoid duplicate mapping

[d4l3k]

Fixed slowness w/ CUDA due to duplicate transfers to CPU