refactor(archon): extract runner and weight sync into separate modules

[rchardx]

Description

Extract forward/backward execution and weight synchronization logic from ArchonEngine into dedicated modules for better separation of concerns.

Key changes:

• Add ForwardBackwardRunner abstraction with SequentialRunner and PipelinedRunner implementations
• Extract weight sync logic to WeightSyncState and dedicated functions
• Simplify planner agent guidelines with clearer structure

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 focuses on improving the modularity and clarity of the ArchonEngine, a core component for distributed ML training. By extracting distinct functionalities like microbatch runners and weight synchronization into their own modules, the engine's codebase becomes more organized and easier to manage. Additionally, the initialization process is made more transparent, and the planner agent's usage guidelines are refined to ensure it's applied effectively.

Highlights

• ArchonEngine Modularization: The core ArchonEngine has been significantly refactored by extracting its runner logic (for sequential and pipelined microbatch execution) and weight synchronization mechanisms into dedicated, separate modules. This improves maintainability and reduces the main engine file's complexity.
• Structured ArchonEngine Initialization: The __init__ method of ArchonEngine has been reorganized with clear, structured sections for configuration, model setup, distributed/parallelism settings, pipeline parallelism, rollout integration, and runtime state, enhancing readability and understanding.
• Optimized Import Statements: Imports that are only necessary for type checking have been moved behind typing.TYPE_CHECKING blocks, which can help reduce startup overhead and potential circular import issues.
• Enhanced Planner Agent Guidance: The documentation for the planner agent has been updated to provide clearer guidance on when to use it, including specific scenarios to avoid, and a more detailed structure for the planning process and output format.

🧠 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.

Changelog

• .claude/agents/planner.md
  • Updated section title from 'When to Activate' to 'When to Use'.
  • Added a new section 'Do NOT use for:' to clarify inappropriate use cases for the planner agent.
  • Expanded the 'Planning Process' section with detailed guidance on clarifying requirements and systematic research steps.
  • Revised the 'Plan Output' format to include both 'Quick Path' for simple tasks and 'Full Plan' for complex tasks, along with section guidelines.
  • Removed 'AReaL-Specific Guidelines' and 'Adding New Task Types' sections, streamlining the document.
• areal/experimental/engine/archon_engine.py
  • Removed numerous imports, delegating them to new modules or placing them under TYPE_CHECKING for type hints only.
  • Introduced imports for PipelinedRunner, SequentialRunner from archon_runner.py and WeightSyncState, init_weight_update_group, update_weights_from_distributed, update_weights_from_disk from archon_weight_sync.py.
  • Refactored the __init__ method into distinct, commented sections for better organization and readability.
  • Replaced the internal _forward_backward_pipelined and _forward_backward_sequential methods with a single call to self.runner.run().
  • Replaced internal weight synchronization methods (_init_weight_update_from_distributed, _update_weights_from_distributed, _update_bucket_weights_from_distributed, _update_weights_from_disk) with calls to the newly extracted functions in archon_weight_sync.py.
  • Added new private methods _setup_parallelism, _apply_pipeline_parallelism, _apply_parallelism, _create_runner, _prepare_pipelined_mb_inputs, and _create_pipelined_loss_fn to encapsulate logic related to parallelism and runner creation.
  • Removed the _get_full_tensor method, as it has been moved to archon_weight_sync.py.
• areal/experimental/engine/archon_runner.py
  • New file created to house the ForwardBackwardRunner abstract base class.
  • Implemented SequentialRunner for handling microbatch execution in non-pipeline parallel scenarios.
  • Implemented PipelinedRunner for managing microbatch execution using the Schedule1F1B strategy in pipeline parallel setups.
• areal/experimental/engine/archon_weight_sync.py
  • New file created to centralize weight synchronization logic.
  • Defined WeightSyncState class to manage the state of weight update groups.
  • Implemented init_weight_update_group for initializing distributed process groups for weight synchronization.
  • Implemented _get_full_tensor to retrieve full tensors from parameters, handling DTensor and CPU offload.
  • Implemented update_weights_from_distributed for broadcasting model parameters across distributed ranks.
  • Implemented _update_bucket_weights to handle chunked broadcasting of weights.
  • Implemented update_weights_from_disk for disk-based weight synchronization as a fallback mechanism.

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 is an excellent refactoring of the ArchonEngine. By extracting the runner logic into archon_runner.py and the weight synchronization logic into archon_weight_sync.py, the main engine class has become significantly more modular and maintainable. The reorganization of the __init__ and create_process_group methods with structured sections is a great improvement for readability. Moving imports behind TYPE_CHECKING is also a good practice for reducing startup overhead.

My only concern is with the changes in planner.md, where a potentially useful section with specific guidelines has been removed. Please see the detailed comment.

Overall, the code changes are of high quality and significantly improve the codebase's structure.

[Copilot]

Pull request overview

This refactor modularizes the Archon training engine by extracting pipeline/forward-backward execution and weight synchronization functionality into dedicated modules, while tightening imports and updating the planner agent guidelines. The intent is structural only, preserving existing behavior but improving readability and maintainability.

Changes:

• Extracted SequentialRunner and PipelinedRunner into areal/experimental/engine/archon_runner.py and wired ArchonEngine’s forward_backward_batch through a unified ForwardBackwardRunner abstraction.
• Extracted XCCL and disk-based weight synchronization logic into areal/experimental/engine/archon_weight_sync.py, replacing the inline _init_weight_update_from_distributed / _update_weights_from_* methods and introducing WeightSyncState.
• Simplified ArchonEngine by reorganizing initialization, parallelism setup, and runner creation, moving some imports behind TYPE_CHECKING; updated .claude/agents/planner.md with clearer “when to use” guidance and planning templates.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
areal/experimental/engine/archon_weight_sync.py	New module encapsulating Archon’s weight sync mechanisms (XCCL group init, DTensor handling, broadcast buckets, and disk-based sync), factored out from ArchonEngine.
areal/experimental/engine/archon_runner.py	New module defining ForwardBackwardRunner plus concrete SequentialRunner and PipelinedRunner (Schedule1F1B) used by Archon for non-PP and PP execution, respectively.
areal/experimental/engine/archon_engine.py	Refactored to use the new runner and weight-sync helpers, reorganized initialization and parallelism configuration, added WeightSyncState integration, and reduced runtime imports via TYPE_CHECKING.
.claude/agents/planner.md	Updated planner agent documentation: renamed “When to Activate” to “When to Use”, added concrete “do not use for” cases, clarified requirements/clarification questions, and tightened the planning output templates and maintainer guidelines.

---

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

[garrett4wade]

LGTM