Libtorch doesn't mirror Pytorch behavior in Linear Layer due to cuBlas

username: M_T1
cooked: <p>I have as minimal python code as I could manage follows:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn
import torch.nn.functional as F

class AC(nn.Module):
    def __init__(self, state_dim, action_dim, action_std_init, device):
        super(AC, self).__init__()
        
        self.device = device

        # Actor network
        self.actor = nn.Sequential(
            nn.Linear(state_dim, 16),
            nn.ReLU(),
            nn.Linear(16, 32),
            nn.ReLU(),
            nn.Linear(32, action_dim)
        )
        
        # Critic network
        self.critic = nn.Sequential(
            nn.Linear(state_dim, 16),
            nn.ReLU(),
            nn.Linear(16, 32),
            nn.ReLU(),
            nn.Linear(32, 1)
        )

        # Action variance
        self.action_var = torch.full((action_dim,), action_std_init * action_std_init).to(device)

        self.to(device)

    def forward(self, state):
        raise NotImplementedError
    
    def evaluate(self, state, action):
        mu = self.actor(state)
        std = torch.sqrt(self.action_var.expand_as(mu))
        action_logprobs = -torch.pow(action - mu, 2) / (2 * self.action_var.expand_as(action)) - std.log() - torch.log(torch.sqrt(torch.tensor(2 * torch.pi)))
        action_logprobs = action_logprobs.sum(2)

        dist_entropy = 0.5 + 0.5 * torch.log(2 * torch.pi * self.action_var)
        dist_entropy = dist_entropy.sum()

        state_values = self.critic(state)

        return action_logprobs, state_values, dist_entropy

# Main code
if __name__ == "__main__":
    # Check if CUDA is available
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if not torch.cuda.is_available():
        print("CUDA is not available! Exiting...")
        exit()

    # Create a minimal AC model
    state_dim = 4
    action_dim = 2
    action_std_init = 0.1
    model = AC(state_dim, action_dim, action_std_init, device)
    num_envs = 4
    num_samples = 40

    # Create dummy input tensors with known shapes
    state2 = torch.rand((num_samples, num_envs, state_dim)).to(device)  
    action2 = torch.rand((num_samples, num_envs, action_dim)).to(device) 

    try:
        action_logprobs2, state_values2, dist_entropy2 = model.evaluate(state2, action2)
        
        print("Action logprobs:", action_logprobs2)
        print("State values:", state_values2)
        print("Dist entropy:", dist_entropy2)
    except Exception as e:
        print("Error during evaluation:", str(e))
</code></pre>
<p>It runs and seems to work as I would like it to.<br>
I mirror it in libtorch as follows:</p>
<pre data-code-wrap="c++"><code class="lang-c++">#include &lt;torch/torch.h&gt;
#include &lt;torch/script.h&gt;
#include &lt;cuda_runtime.h&gt;
#include &lt;iostream&gt;

struct AC : torch::nn::Module {
    torch::nn::Sequential actor, critic;
    torch::Tensor action_var;
    torch::Device device;

    // Constructor
    AC(long int state_dim, long int action_dim, double action_std_init, torch::Device device)
        : actor(register_module("actor", torch::nn::Sequential(
                                             torch::nn::Linear(state_dim, 16),
                                             torch::nn::ReLU(),
                                             torch::nn::Linear(16, 32),
                                             torch::nn::ReLU(),
                                             torch::nn::Linear(32, action_dim)))),
          critic(register_module("critic", torch::nn::Sequential(
                                               torch::nn::Linear(state_dim, 16),
                                               torch::nn::ReLU(),
                                               torch::nn::Linear(16, 32),
                                               torch::nn::ReLU(),
                                               torch::nn::Linear(32, 1)))),
          device(device),
          action_var(torch::full({action_dim}, action_std_init * action_std_init, torch::TensorOptions().device(device))) {
        this-&gt;to(device);
    }

    std::tuple&lt;torch::Tensor, torch::Tensor, torch::Tensor&gt; evaluate(const torch::Tensor&amp; state, const torch::Tensor&amp; action) {
        torch::Tensor mu = actor-&gt;forward(state);
        torch::Tensor std = torch::sqrt(action_var.expand_as(mu));
        torch::Tensor action_logprobs = -torch::pow(action - mu, 2) / (2 * action_var.expand_as(action)) - std.log() - torch::log(torch::sqrt(torch::tensor(2 * M_PI, torch::TensorOptions().device(device))));
        action_logprobs = action_logprobs.sum(2);

        torch::Tensor dist_entropy = 0.5 + 0.5 * torch::log(2 * M_PI * action_var);
        dist_entropy = dist_entropy.sum();

        torch::Tensor state_values = critic-&gt;forward(state);

        return std::make_tuple(action_logprobs, state_values, dist_entropy);
    }
};

int main() {
    // Check if CUDA is available
    if (!torch::cuda::is_available()) {
        std::cerr &lt;&lt; "CUDA is not available! Exiting..." &lt;&lt; std::endl;
        return -1;
    }

    // Set CUDA device
    int device_id = 0;
    cudaSetDevice(device_id);
    torch::Device device(torch::kCUDA, device_id);

    // Create a minimal AC3 model
    long int state_dim = 4;
    long int action_dim = 2;
    double action_std_init = 0.1;
    AC model(state_dim, action_dim, action_std_init, device);
    model.to(device);  

    long int num_envs = 4;
    long int num_samples = 40;

    // Create dummy input tensors with known shapes
    torch::Tensor state2 = torch::rand({num_samples, num_envs, state_dim}).to(device);  // Batch of states
    torch::Tensor action2 = torch::rand({num_samples, num_envs, action_dim}).to(device);  // Corresponding actions

    // Test the evaluate method
    try {
        auto [action_logprobs2, state_values2, dist_entropy2] = model.evaluate(state2, action2);

        std::cout &lt;&lt; "Action logprobs: " &lt;&lt; action_logprobs2 &lt;&lt; std::endl;
        std::cout &lt;&lt; "State values: " &lt;&lt; state_values2 &lt;&lt; std::endl;
        std::cout &lt;&lt; "Dist entropy: " &lt;&lt; dist_entropy2 &lt;&lt; std::endl;
    } catch (const c10::Error&amp; e) {
        std::cerr &lt;&lt; "Error during evaluation: " &lt;&lt; e.what() &lt;&lt; std::endl;
        return -1;
    }

    return 0;
}

</code></pre>
<p>I get the following error at the first line of the evaluate function (high in the callstack is the linear layer):</p>
<pre><code class="lang-auto">Error during evaluation: CUDA error: CUBLAS_STATUS_NOT_INITIALIZED when calling `cublasSgemm( handle, opa, opb, m, n, k, &amp;alpha, a, lda, b, ldb, &amp;beta, c, ldc)`
Exception raised from gemm&lt;float&gt; at ../aten/src/ATen/cuda/CUDABlas.cpp:427 (most recent call first):
</code></pre>
<p>However in running the non minimal example, tensor operations were working fine, and even using the same network with different dimension inputs and different function (shape:4,4) worked.</p>
<p>linking to cublas doesn’t seem to be the problem, I testing pure cublas calls in my build system and they worked fine.</p>
<p>Any idea on why I am seeing this discrepancy in libtorch and pytorch and how to fix it?</p>
<p>pytorch version: 2.1.0<br>
libtorch version: 2.1.0<br>
os: ubuntu 22.04 lts<br>
nvidia-driver: 550.54.15<br>
device: H100 80GB PCIE<br>
cuda version: 12.1</p>

username: ptrblck
cooked: <p>Check via <code>LD_DEBUG=libs</code> which cuBLAS libs are loaded and make sure they are the same.</p>

username: M_T1
cooked: <p>Found out I was linking wrong. Edited my build system and it worked.</p>