Hi Davide,
I ended up implementing LSTM from scratch in Pytorch as it gave me greater flexibility to try various modifications and not mess with the Pytorch internals. It didn’t take much time though.
But once you have an implementation, you can still automatically run it in parallel/distributed mode using Pytorch’s built in modules. Then, to explicitly define hidden states at every call, have the LSTM call the init method like this:
    def init_states(self, hiddensize, distbatchsize):
        """ Initialize h and c LSTMCell states """
        h = torch.randn(distbatchsize, hiddensize, self.input_dim[0], self.input_dim[1]).cuda()
        c = torch.randn(distbatchsize, hiddensize, self.input_dim[0], self.input_dim[1]).cuda()
        states = (h,c)
        return states


In data parallel mode, there is a separate network on each GPU, so this function is called locally on each GPU which has its local batch size. In that case, the hidden states should then be initialized with that local batch size.
The “distbatchsize” here does exactly that.  It is the batch size  local to the GPU after the full batch size has been scattered by the distributed module. So if your full batchsize = 8 and num_gpus = 4, then distbatchsize = 8/4 = 2. Pytorch automatically computes this and splits your input tensors into that size and sends it to all GPUs (Check the all_scatter module in MPI for more details). So How do you find the distbatchsize? The snippet below is from the LSTM forward method. The first dimension is the local batch size = 2. If you were running this on only 1 GPU, the distbatchsize = 8 i.e. the global batch size.
    def forward(self, input_tensor):
        """ Forward Pass """

        # Find the local batch size
        distbatchsize = input_tensor.size(0)

       # continue with init_states and forward pass


Not a clean solution, but it serves me well. Let me know if you have any comments!