Ok it looks this issue exists on both cpu and gpu:
cpu: https://app.circleci.com/pipelines/github/pytorch/xla/14953/workflows/946bc70e-a4ba-4e89-be01-0d7ea7f653cb/jobs/34704
gpu: https://app.circleci.com/pipelines/github/pytorch/xla/14953/workflows/946bc70e-a4ba-4e89-be01-0d7ea7f653cb/jobs/34703

I have a feeling that it might be because sigmoid(x) = 0.5+0.5*tanh(0.5*x) is faster.. let me double check.

Yes, the tanh implementation is slightly faster on GPU.
Using the following script:

x = torch.rand(1000000000,device=device)
xm.mark_step()
t0 = time.time()
for _ in range(100):
    for _ in range(100):
        y = torch.sigmoid(x)
    xm.mark_step()
t1 = time.time()
print(t1-t0)

I'm getting 1.2621409893035889 with tanh implementation (with clamp) and 1.301847219467163 with normal implementation.

If we want to keep the tanh implementation, one way is to wrap it with xla::Clamp(zero, half + half * xla::Tanh(half * input), one).

I talked with Blake. Speed was the main reason we used tanh and TPU does not have this numerical instability issue. He suggested us to lower sigmod using XlaOp Logistic(XlaOp operand); which will have different TPU and GPU implementation in the backend to handle the subtle difference in accelerators.

Updated and thanks for the info. I just realize xla::Logistic is equivalent to torch.sigmoid.