Add support for context parallelism (round 2)

[epwalsh]

Replaces #174.

There are two main differences between this implementation and that in #174:

1. This relies on ring-flash-attn instead of PyTorch's built-in experimental API, which gives us more flexibility and support for intra-document masking with CP.
2. This allows us to decouple the FSDP process group from the CP process group.