I see, for this use case, an alternative is to use torch.distributed.rpc to connect the parameter server with trainers, and then let the parameter server periodically flush checkpoints to the external storage. So that you don’t have to pay the checkpointing overhead in every iteration.
Some related resources:

Building HogWild! 1 PS using torch.distributed.rpc: https://pytorch.org/tutorials/intermediate/rpc_param_server_tutorial.html 7

Batch updating PS (requires v1.6+): https://github.com/pytorch/tutorials/blob/release/1.6/intermediate_source/rpc_async_execution.rst 4