Fused Self Attention Produces Large Errors When Using torch.randn Instead of torch.rand

[Jerry2423]

Describe the bug

In sd_attention_torch.py, when q_tensor (same for k and v) is generated using torch.rand, the results of fused_self_attn_for_SD_small_head_size match the expected output from cpu_golden_attn. However, when generated using torch.randn, the computed results show a significant discrepancy.

Expected Behavior

   q_tensor = torch.randn((4096, 64), dtype=torch.float32).to(device=device)
  k_tensor = torch.randn((4096, 64), dtype=torch.float32).to(device=device)
  v_tensor = torch.randn((4096, 64), dtype=torch.float32).to(device=device)

  output_nki = fused_self_attn_for_SD_small_head_size(q_tensor, k_tensor, v_tensor)

  output_torch = cpu_golden_attn(q_tensor, k_tensor, v_tensor)

  allclose = torch.allclose(output_torch, output_nki, atol=1e-5, rtol=1e-3)

  if allclose:
    print("NKI and Torch match")
  else:
    print("NKI and Torch differ")

Expected output - "NKI and Torch match"

Current Behavior

NKI and Torch differ

Reproduction Steps

In sd_attention_torch.py, replace q_tensor = torch.rand((4096, 64), dtype=torch.float32).to(device=device) with q_tensor = torch.randn((4096, 64), dtype=torch.float32).to(device=device). Do the same thing with k_tensor, v_tensor

Regression Issue

• Select this option if this issue appears to be a regression.

Possible Solution

No response

Additional Information/Context

No response

neuronx-cc version used

aws_neuronx_venv_pytorch_2_5_nxd_inference

Framework(s) and their versions used (JAX, PyTorch, etc..)

No response

[aws-haoyuli]

Thanks @Jerry2423 for filing the issue. We will take a look and get back to you.

[aws-cph]

Thank you for your patience as we investigated the torch.randn functionality on Trainium1.
I wanted to provide you with a clear update on the current status.
Unfortunately, torch.randn is a known limitation on trn1 instances, and there is no available workaround at this time. We also want to note that the example you were originally referencing has since been archived in our documentation as this use case is no longer supported.