Instantiate, loading and saving different model with HYBRID_SHARD FSDP

Assume I am using HYBRID_SHARD FSDP and model are sharded in one node, is it possible to create different model and loading different model checkpoint for different node, and save different checkpoint per node.

This is under the scenario where the FSDP model does not engage in the model’s forward and backward passing.

coef = node_id
model = ModelClass(coef)
# how to construct FSDP
# how to load checkpoint based on node_id
# how to save FSDP checkpoint based on node_id