feat(archon): enable PP > 1 for RL training with XCCL weight sync

[rchardx]

Description

Enable pipeline parallelism (PP > 1) to work correctly with RL training by fixing data broadcast and weight synchronization issues.

Key changes:

• Replace cp_tp mesh with pp_cp_tp mesh to include PP in context_and_model_parallel_group
• Add per-PP-stage weight update groups for XCCL weight sync
• Add DistributedLock to serialize NCCL group initialization
• Add PP + weight_tying validation (unsupported combination)
• Fix MegatronEngine missing engine_lock in init

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not
  work as expected)
• Documentation update
• Code refactoring (no functional changes)
• Performance improvement
• Test coverage improvement

Checklist

• I have read the Contributing Guide
• I have run formatting tools (pre-commit or manual)
• I have run relevant unit tests and they pass
• I have added tests for new functionality
• I have updated documentation if needed
• My branch is up to date with main
• This PR introduces breaking changes (if yes, fill out details below)
• If this PR changes documentation, I have built and previewed it locally with
  jb build docs
• No critical issues raised by AI reviewers (/gemini review)

[gemini-code-assist]

Summary of Changes

Hello @rchardx, 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 significantly enhances the distributed training capabilities by enabling and stabilizing pipeline parallelism (PP > 1) for RL training. It achieves this by fundamentally restructuring how parallel groups are defined and utilized, introducing fine-grained, pipeline-stage-aware weight synchronization, and implementing critical locking mechanisms to maintain data integrity. The changes ensure that complex distributed setups can operate correctly and efficiently, while also preventing known incompatible configurations.

Highlights

• Pipeline Parallelism (PP > 1) Enablement: This pull request enables robust pipeline parallelism (PP > 1) for Reinforcement Learning (RL) training, addressing previous limitations in distributed data handling and weight synchronization.
• Parallel Group Refactoring: The cp_tp (context and tensor parallel) mesh has been replaced with a pp_cp_tp mesh, integrating pipeline parallelism into the core context and model parallel group definitions for more comprehensive distributed operations.
• Per-PP-Stage Weight Synchronization: New per-pipeline-stage weight update groups have been introduced for XCCL (NVIDIA Collective Communications Library) weight synchronization, allowing for efficient and correct weight updates across different pipeline stages.
• DistributedLock for Synchronization: A DistributedLock mechanism has been added to serialize NCCL group initialization and weight update operations, preventing race conditions and ensuring data consistency in distributed environments.
• Validation for Weight Tying with PP: New validation logic prevents the unsupported combination of pipeline parallelism (PP > 1) and weight tying, as tok_embeddings and output layers cannot share weights when distributed across different GPUs.
• MegatronEngine Fix: A fix was implemented in MegatronEngine to correctly acquire and release an engine_lock during distributed weight update group initialization, resolving a missing synchronization issue.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

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 enables pipeline parallelism (PP > 1) for RL training with XCCL weight synchronization. The changes include updating process group meshes to incorporate the pipeline parallel dimension, creating per-PP-stage weight update groups, and adding a distributed lock to serialize NCCL group initialization. Additionally, a validation check for the unsupported combination of pipeline parallelism and weight tying has been added.

My review identified a critical issue in three separate locations where a distributed lock is acquired. The current implementation does not guarantee that the lock will be released if an exception occurs within the locked code block, which could lead to deadlocks. I've provided suggestions to use try...finally blocks to ensure the lock is always released, making the implementation more robust.

[Copilot]

Pull request overview

This pull request enables pipeline parallelism (PP > 1) to work correctly with RL training by fixing data broadcast and weight synchronization issues when using XCCL for weight updates.

Changes:

• Replaced cp_tp mesh with pp_cp_tp mesh to include PP dimension in the context and model parallel group, enabling proper data broadcast across PP stages
• Added per-PP-stage weight update groups with DistributedLock to serialize NCCL group initialization and prevent race conditions
• Added validation to prevent PP with weight_tying, which is an unsupported combination due to shared weights residing on different GPUs

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.

File	Description
areal/experimental/models/archon/parallel_dims.py	Updated mesh definition from cp_tp to pp_cp_tp to include PP dimension for proper context and model parallel grouping
areal/experimental/engine/archon_engine.py	Added PP rank tracking, per-stage weight update groups, DistributedLock for NCCL serialization, and PP+weight_tying validation
areal/engine/megatron_engine.py	Added DistributedLock usage for NCCL group initialization consistency with ArchonEngine

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.