Segmentation Fault When Training

Hi all,

I'm getting a segmentation fault while using deepspeed. I ran my training script through gdb and got that it's coming from deepspeed/ops/csrc/adam/cpu_adam.cpp:441, which does the following:

param_4[0].data = SIMD_LOAD(_params + i);

I know this isn't super helpful, but I'm not sure how to debug it better. The stacktrace was all ??s.

FWIW I'm using Pytorch 1.8.1 with cuda 11.1. I'm also using deepspeed within Pytorch Lightning and I set up the optimizer in the most basic way I could:

    def configure_optimizers(self):
        # create the optimizer
        optimizer = DeepSpeedCPUAdam(self.parameters(), lr=1e-3)
        return optimizer

I also was testing this on a single GPU before deploying it to a server and running it on a multi-GPU node--not sure if that would make a difference. And the only other weird thing in my setup that I can think of is that I am using the RC for triton v1.0.0 because for whatever reason I could not build triton v0.2.3 (and the author of triton suggested using 1.0.0 because it doesn't require building from scratch)

Start at deepspeed/ops/csrc/adam/cpu_adam.cpp:441 and reproduce the segmentation fault with the provided DeepSpeedCPUAdam setup in PyTorch Lightning. Inspect the gdb output and compare behavior on the reported single-GPU and multi-GPU environments, including the PyTorch 1.8.1, CUDA 11.1, and Triton 1.0.0 versions. Done means the crash cause is isolated and a verified fix or precise compatibility diagnosis is documented.