Issue for ring-llama/test.ipynb

[fayejf]

Thanks for the great work!

https://github.com/cuda-mode/ring-attention/blob/main/ring-llama/test.ipynb

I can load the model with LlamaRingFlashAttention and move to the device but I've seen

RuntimeError: Default process group has not been initialized, please make sure to call init_process_group.

when I run y = model.generate

What did I miss? Thanks in advance!

[fayejf]

Also I was seeing

RuntimeError: Sizes of tensors must match except in dimension 2. Expected size 32 but got size 4 for tensor number 1 in the list.

with another model llama based model "01-ai/Yi-6B-200K"with LlamaRingFlashAttention
it was fine with LlamaFlashAttention2