Regarding the parameter server idea, I guess parameter server and linear are quite different concepts as a PS is a training paradigm while a linear layer is a component of a model.
It looks like you have a use case to shard a linear layer across multiple processes/potentially multiple nodes. We are planning to build out sharding primitives within PyTorch, please see this RFC: [RFC] Model Sharding for distributed training · Issue #55207 · pytorch/pytorch · GitHub 2 which would suit this use case once it is built out.
Curious, if the input is a sparse feature, would an nn.EmbeddingBag work better than a linear layer?