feat: support context parallel for FSDP backend

[Beichen-Ma]

Description

This PR implements context parallel support for FSDP backend with all-gather communication type. Motivated by #412, support non-zigzag CP slicing format for FSDP CP & all-gather KV backend.

Tests

• tests/test_qwen3_0.6B_megatron_fsdp_align.py passed when CP enabled.
• Test numerical correctness by compare:
  • CP=1 and CP=2 produce identical gradient norms
  • Megatron CP=2 and FSDP CP=2 produce identical gradient norms

Next steps

• rollout data logging for fsdp cp>1

[gemini-code-assist]

Summary of Changes

Hello @Beichen-Ma, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces comprehensive support for context parallelism within the FSDP training backend, utilizing an 'all-gather' communication pattern. The changes span data preparation, model input handling, gradient processing, and loss calculation, ensuring that packed sequences and their associated metadata are correctly distributed and processed across parallel ranks. This enhancement significantly broadens the applicability of context parallelism to FSDP-based models, improving efficiency for certain workloads.

Highlights

• FSDP Context Parallelism (CP) Support: Enables context parallelism for the FSDP backend, specifically implementing an 'all-gather' communication strategy.
• Non-Zigzag CP Slicing: Introduces a new, non-zigzag CP slicing format for FSDP CP and all-gather KV backend, addressing previous limitations.
• Gradient Norm Synchronization: Implements explicit synchronization of gradient norms across the dp_cp_group for cp_size > 1 during gradient clipping to ensure correctness.
• Packed Batch Handling: Adds new utility functions and modifies data loading to correctly handle packed sequences and log probabilities under the 'all-gather' CP scheme.
• Conditional Logging: Temporarily disables rollout data logging for cp_size > 1 due to an identified logging problem, with a TODO to resolve it in future work.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces support for context parallelism (CP) in the FSDP backend using an all-gather communication strategy. This is a significant feature that enables more efficient training of large models. The changes are extensive, touching data processing, loss calculation, and distributed training utilities. The implementation is well-structured, with the new logic gated by a cp_comm_type flag. The PR also includes a manual implementation for gradient clipping when CP is enabled.

My review focuses on a performance improvement in the gradient norm calculation and suggests improving consistency in the codebase for safer data handling. Overall, this is a great contribution.

[gemini-code-assist]

The current implementation for calculating the total gradient norm involves calling .item() inside a loop over model parameters. This can be inefficient as it causes a GPU-to-CPU synchronization for each parameter. It's better to perform the accumulation on the GPU.

Suggested change

	total_norm_sq += param_norm.item() ** 2
	total_norm_sq += param_norm ** 2

[gemini-code-assist]

The zip function is used with strict=False. However, in other parts of the codebase, such as miles/backends/training_utils/data.py at line 195, strict=True is used for similar operations. The iterables total_lengths and response_lengths are expected to have the same length. Using strict=True would enforce this and help catch potential bugs early if the lengths mismatch. This applies to other new uses of zip in this PR as well. It would be good to be consistent and use strict=True where applicable.

Suggested change

	for total_length, response_length in zip(total_lengths, response_lengths, strict=False):
	for total_length, response_length in zip(total_lengths, response_lengths, strict=True):