feat(archon): use meta device + DCP for memory-efficient model init

[rchardx]

Description

Optimize Archon engine initialization to use meta device + FSDP + DCP loading pattern, reducing peak memory from N × model_size to ~1/N × model_size per rank.

Changes:

• Create model structure on meta device (no memory allocation)
• Apply FSDP/TP/EP parallelization to meta device model
• Materialize tensors with to_empty() after FSDP
• Load weights via DCP (each rank reads only its shard)
• Add init_buffers() to BaseArchonModel for buffer-only initialization
• Separate init_weights() (parameters) from init_buffers() (rope_cache)

Related Issue

Fixes #857

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 introduces a significant performance improvement to the Archon engine by refactoring its model initialization process. By leveraging PyTorch's meta device functionality and a distributed checkpoint loading pattern, the system can now construct and load large language models with drastically reduced peak memory usage, making it more scalable and efficient for distributed training environments.

Highlights

• Memory-Efficient Model Initialization: The Archon engine now utilizes a meta device and a specific loading pattern (FSDP + DCP) to significantly reduce peak memory consumption during model initialization, aiming for ~1/N × model_size per rank.
• Separation of Weight and Buffer Initialization: The BaseArchonModel interface has been updated to distinguish between init_weights (for model parameters) and init_buffers (for non-parameter tensors like rope_cache), allowing for more granular control over initialization.
• Staged Model Materialization: Model structure is first created on a meta device without memory allocation, then parallelized (FSDP/TP/EP), and finally, tensors are materialized and weights loaded via DCP, ensuring each rank only loads its specific shard.

🧠 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 introduces a significant performance improvement by adopting the meta device initialization pattern for the Archon engine. This is a great change that will substantially reduce peak memory usage during model loading. The refactoring to separate model structure creation, materialization, and weight/buffer initialization is well-executed.

My review includes a couple of suggestions for improving code clarity and diagnostics:

• Updating a docstring in the base model class to reflect the new separation of concerns between init_weights and init_buffers.
• A note on a minor regression in diagnostic logging for missing weights during model loading.

[Copilot]

Pull request overview

This PR optimizes Archon engine initialization to use meta device + FSDP + DCP loading, significantly reducing peak memory usage from N × model_size to ~1/N × model_size per rank during distributed model loading.

Changes:

• Refactored model initialization to create model structure on meta device first, then materialize after FSDP parallelization
• Split buffer initialization from weight initialization by adding init_buffers() method to BaseArchonModel
• Replaced direct weight loading with DCP-based distributed loading where each rank reads only its shard

Reviewed changes

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

File	Description
areal/experimental/engine/archon_engine.py	Refactored _create_device_model() to create meta device model structure; added _materialize_and_load_weights() to handle materialization and DCP loading after FSDP; removed unused import
areal/experimental/models/archon/base.py	Added abstract init_buffers() method to BaseArchonModel for buffer-only initialization
areal/experimental/models/archon/qwen2/model/model.py	Moved rope_cache initialization from init_weights() to new init_buffers() method
areal/experimental/models/archon/qwen3/model/model.py	Moved rope_cache initialization from init_weights() to new init_buffers() method

---

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