feat(archon): add DCP-based HF checkpoint save/load with MoE expert support#849
Conversation
Summary of ChangesHello @rchardx, I'm Gemini Code Assist1! 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 enhancement to the checkpointing mechanism for Archon models, particularly focusing on Mixture of Experts (MoE) models. It leverages PyTorch's Distributed Checkpoint (DCP) infrastructure to enable distributed saving and loading of model weights, addressing out-of-memory (OOM) issues that arise when dealing with large models. The changes include a new state dict adapter for MoE models, fixes for torch.compile compatibility, and extensive testing to ensure the correctness and efficiency of the new checkpointing process. Highlights
🧠 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 AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request includes changes to the archon engine and model, specifically focusing on improving checkpointing and state dictionary adapter functionality for MoE models. It introduces a MoEStateDictAdapter for handling expert weight conversion, adds support for multi-file safetensors loading, and integrates Hugging Face storage writer/reader for saving and loading models. Review comments suggest improvements such as using masked_fill_ for in-place operations, adding checks for dictionary keys before iteration, raising more specific exception types, adding comments to explain code sections, and improving error messages for better debugging.
There was a problem hiding this comment.
Pull request overview
Adds DCP-backed HuggingFace checkpoint save/load for Archon, including DTensor-aware MoE expert weight conversion to avoid rank-0 full state_dict gathers (and associated OOMs).
Changes:
- Replaces HF checkpoint save/load in
ArchonLMEnginewithtorch.distributed.checkpoint+ HF storage reader/writer. - Introduces
MoEStateDictAdapterand updates Qwen3/Qwen2 adapters for DTensor-aware MoE expert conversion and multi-file safetensors index support. - Adds new distributed/integration tests for DTensor checkpoint roundtrips, multi-file checkpoints, and engine save/load correctness.
Reviewed changes
Copilot reviewed 19 out of 19 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| areal/experimental/engine/archon_engine.py | Switches HF save/load to DCP + HF storage writer/reader; supports multi-file consolidation. |
| areal/experimental/models/archon/base.py | Adds hf_assets_path + multi-file model.safetensors.index.json parsing and HF storage reader helper. |
| areal/experimental/models/archon/moe_state_dict_adapter.py | New DTensor-aware MoE expert split/concat utilities for distributed checkpoints. |
| areal/experimental/models/archon/qwen3/model/state_dict_adapter.py | Migrates to MoEStateDictAdapter and adds DTensor-aware MoE HF conversion paths. |
| areal/experimental/models/archon/qwen2/model/state_dict_adapter.py | Extends adapter init for hf_assets_path and removes forced DTensor full gathers in HF conversion. |
| areal/experimental/models/archon/moe/router.py | Replaces in-place scatter with out-of-place for compile/checkpointing compatibility. |
| areal/experimental/models/archon/attention.py | Replaces in-place masked_fill with out-of-place for compile/checkpointing compatibility. |
| areal/experimental/models/archon/activation_checkpoint.py | Docstring/comment cleanup. |
| areal/experimental/models/archon/init.py | Exports MoEStateDictAdapter. |
| areal/tests/experimental/archon/torchrun/dist_utils.py | Refactors/shortens docstrings; keeps shared distributed test utilities. |
| areal/tests/experimental/archon/torchrun/run_ep_tests.py | Adds DTensor checkpoint roundtrip tests, placement-type checking, and Qwen2 dense DTensor tests. |
| areal/tests/experimental/archon/torchrun/run_checkpoint_tests.py | New torchrun integration tests for engine HF save/load (dense + MoE). |
| areal/tests/experimental/archon/test_state_dict_adapter.py | Adds unit tests for hf_assets_path, multi-file support, and MoE adapter helpers. |
| areal/tests/experimental/archon/test_distributed_ep.py | Adds multi-GPU DTensor checkpoint roundtrip tests. |
| areal/tests/experimental/archon/test_checkpoint_e2e.py | New end-to-end tests for multi-file index parsing and engine HF save/load. |
| areal/tests/experimental/archon/torchrun/run_vs_fsdp.py | Updates test import to dist_utils. |
| areal/tests/experimental/archon/torchrun/run_qwen3_parallelize.py | Updates test import to dist_utils. |
| areal/tests/experimental/archon/torchrun/run_parallel_dims.py | Updates test import to dist_utils. |
| areal/tests/experimental/archon/torchrun/run_forward.py | Minor comment cleanup. |
Comments suppressed due to low confidence (1)
areal/experimental/models/archon/qwen3/model/state_dict_adapter.py:130
- to_hf() does not strip wrapper name segments (e.g. "._checkpoint_wrapped_module", "._orig_mod") before key conversion, even though convert_single_to_hf() explicitly handles them. Since the new engine save/load path calls adapter.to_hf() on model state dicts, compiled / AC-wrapped models may silently drop keys (hf_key=None) and produce incomplete checkpoints. Suggest normalizing keys inside to_hf() (and any other bulk conversion entrypoints) the same way convert_single_to_hf() does before calling _convert_key_to_hf / MoE splitting.
for key, value in state_dict.items():
# Skip output.weight when weight tying is enabled
# HF will reconstruct lm_head.weight from embed_tokens.weight on load
if self.enable_weight_tying and key == "output.weight":
continue
if "moe.experts.w" in key:
# Split 3D expert weight into list of 2D weights
if isinstance(value, DTensor):
hf_pairs = self._split_moe_experts_distributed(key, value)
else:
hf_pairs = self._split_moe_experts(key, value)
hf_state_dict.update(hf_pairs)
else:
# Regular key mapping
hf_key = self._convert_key_to_hf(key)
if hf_key is not None:
hf_state_dict[hf_key] = value
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
rchardx
force-pushed
the
rchardx/state_dict
branch
from
136a73e to
07a1bbd
Compare
rchardx
had a problem deploying
to
AReaL-unittests
GitHub Actions
Failure
rchardx
force-pushed
the
rchardx/state_dict
branch
from
07a1bbd to
10344b0
Compare
rchardx
added
safe-to-test
rchardx
had a problem deploying
to
AReaL-unittests
GitHub Actions
Failure
rchardx
force-pushed
the
rchardx/state_dict
branch
from
10344b0 to
144a9b2
Compare
rchardx
added
safe-to-test
rchardx
had a problem deploying
to
AReaL-unittests
GitHub Actions
Error
rchardx
force-pushed
the
rchardx/state_dict
branch
from
144a9b2 to
cbbf9d4
Compare
rchardx
added
safe-to-test
rchardx
temporarily deployed
to
AReaL-unittests
GitHub Actions
Inactive
| hf_writer = HuggingFaceStorageWriter( | ||
| path=sharded_dir, | ||
| save_distributed=True, | ||
| fqn_to_index_mapping=fqn_to_index_mapping, | ||
| enable_consolidation=False, | ||
| ) | ||
| dcp.save(hf_state_dict, storage_writer=hf_writer) | ||
|
|
||
| # NOTE: consolidate_safetensors_files_on_every_rank() has internal barrier | ||
| consolidate_safetensors_files_on_every_rank( | ||
| input_dir=sharded_dir, | ||
| output_dir=path, | ||
| fqn_to_index_mapping=fqn_to_index_mapping, | ||
| num_threads=8, | ||
| ) |
There was a problem hiding this comment.
What about the performance?
There was a problem hiding this comment.
Our benchmark on a 12.55 GB model across 8 GPUs shows that DCP provides a modest 1.10x speedup (11.18s vs 12.27s), but the real advantage is memory efficiency: DCP uses 23x less CPU memory (0.45 GB vs 10.52 GB). The baseline approach with full_state_dict=True gathers the entire model to each rank before saving, consuming ~10.5 GB extra CPU memory per rank—this becomes the bottleneck for larger models where OOM is inevitable. DCP keeps data sharded throughout the save process, using only ~1.57 GB per rank. The time breakdown reveals that DCP's get_state_dict is 7x faster (0.39s vs 2.72s) by avoiding all-gather communication, though the two-phase write (DCP save + consolidate) takes longer than single-rank sequential write. In summary, DCP's primary value is enabling checkpoint saves for large models that would otherwise OOM, with a secondary benefit of slightly better throughput at scale.
Benchmarking: https://gist.github.com/rchardx/bc558b7bf7cab969e4075a71b17ed1a9
torchrun --nproc_per_node=8benchmark_hf_checkpoint_save.py --model-size large --warmupCreating FSDP2 model: large
World size: 8
Estimated size: 12.55 GB
Per-GPU shard: 1.57 GB
Model created and wrapped with FSDP2
GPU memory after model init: 1.38 GB
Running warmup...
Running BASELINE benchmark...
Running DCP benchmark...
================================================================================
Model: large
Layers: 32, Hidden: 4096
Estimated size: 12.55 GB
World size: 8 GPUs
Per-GPU shard size: 1.57 GB
================================================================================
BASELINE (full_state_dict=True, rank 0 saves):
Time:
All-gather: 2.72s
Save (rank 0): 9.54s
Total: 12.27s
Throughput: 1.02 GB/s
Memory (rank 0):
GPU before: 1.38 GB
GPU peak (after gather): 1.63 GB
GPU delta: +0.24 GB
CPU before: 2.10 GB
CPU after gather: 12.62 GB
CPU delta: +10.52 GB
DCP (full_state_dict=False, parallel save + consolidate):
Time:
Get state dict: 0.39s
DCP save: 7.51s
Consolidate: 3.28s
Total: 11.18s
Throughput: 1.12 GB/s
Memory (rank 0):
GPU before: 1.38 GB
GPU peak (after get): 1.38 GB
GPU delta: +0.00 GB
CPU before: 2.10 GB
CPU after get: 2.55 GB
CPU delta: +0.45 GB
--------------------------------------------------------------------------------
COMPARISON:
--------------------------------------------------------------------------------
Time speedup (DCP vs Baseline): 1.10x
CPU memory ratio (Baseline / DCP): 23.29x
Baseline needs 10.52 GB extra CPU memory
DCP needs 0.45 GB extra CPU memory
Theoretical analysis:
Model size: 12.55 GB
Baseline gathers full model to each rank -> ~12.55 GB extra memory
DCP keeps sharded data -> ~1.57 GB per rank
================================================================================
| def _load_safetensors_index(self, hf_assets_path: str) -> None: | ||
| """Load model.safetensors.index.json to support multi-file checkpoint.""" | ||
| mapping_path = os.path.join(hf_assets_path, "model.safetensors.index.json") | ||
| single_file_path = os.path.join(hf_assets_path, "model.safetensors") | ||
|
|
||
| try: | ||
| with open(mapping_path) as f: | ||
| hf_safetensors_index = json.load(f) | ||
| except FileNotFoundError: | ||
| if not os.path.exists(single_file_path): | ||
| logger.warning( | ||
| f"model.safetensors.index.json not found at hf_assets_path: {mapping_path}. " | ||
| "Defaulting to saving a single safetensors file if checkpoint is saved in HF format" | ||
| ) | ||
| return | ||
|
|
||
| if hf_safetensors_index: | ||
| self.fqn_to_index_mapping = {} | ||
| for hf_key, raw_index in hf_safetensors_index["weight_map"].items(): | ||
| match = re.search(r"\d+", raw_index) | ||
| if match: | ||
| self.fqn_to_index_mapping[hf_key] = int(match.group(0)) | ||
|
|
||
| def get_hf_storage_reader( | ||
| self, path: str, from_quantized: bool = False | ||
| ) -> HuggingFaceStorageReader: | ||
| """Return HuggingFaceStorageReader to read HF checkpoint. | ||
|
|
||
| Args: | ||
| path: The path to read HF checkpoint from. | ||
| from_quantized: Whether loading from quantized checkpoint format. | ||
| Note: Loading from quantized format is not supported by default. | ||
|
|
||
| Returns: | ||
| HuggingFaceStorageReader instance for reading HF checkpoint. | ||
| """ | ||
| if from_quantized: | ||
| logger.warning( | ||
| "Loading from quantized checkpoint format is not supported for this model." | ||
| ) | ||
| return HuggingFaceStorageReader(path) |
There was a problem hiding this comment.
An abstract class should not implement any informative methods. It should only define the API.
If we indeed want to share some commen methods, we should wrap them in an indepenent class and use the object on-demand in subclasses. Prefer composition over inheritance.
There was a problem hiding this comment.
This pattern is known as the Template Method Pattern, where an ABC defines an algorithm skeleton with concrete implementations while deferring certain steps to subclasses via abstract methods. It is widely used in Python's standard library. For example, collections.abc.Mapping only requires subclasses to implement __getitem__, __iter__, and __len__, but provides concrete implementations of get(), keys(), items(), and values() for free. Similarly, collections.abc.Set requires __contains__, __iter__, and __len__, but provides __le__, __and__, isdisjoint(), and other set operations as concrete methods. The shared implementation in BaseStateDictAdapter follows the same pattern and is minimal and stable. If it grows more complex or different subclasses need different implementations, we should revisit this decision and extract an HFAssetsHelper class for composition.
rchardx
force-pushed
the
rchardx/state_dict
branch
3 times, most recently
from
5f19e3f to
3b6bf4b
Compare
…upport Replace manual safetensors handling with PyTorch DCP infrastructure for HuggingFace format checkpoints. Previously, save/load required gathering full state dict to rank 0, causing OOM on large MoE models. Now uses HuggingFaceStorageWriter/Reader for distributed operations. Key changes: - Add MoEStateDictAdapter base class with DTensor-aware expert weight conversion (3D grouped <-> 2D individual), supporting EP and ETP placements via local expert index calculation - Fix in-place tensor ops (masked_fill_, scatter_) for torch.compile and checkpointing compatibility - Large MoE checkpoint save/load without OOM (no full tensor gather) - Correct EP/ETP checkpoint handling with distributed expert weights - Incremental weight sync to inference engines via convert_single_to_hf() - Foundation for PP support: layer-indexed key mapping and extensible DTensor placement logic for per-stage conversion
rchardx
force-pushed
the
rchardx/state_dict
branch
from
3b6bf4b to
b5c5393
Compare
rchardx
temporarily deployed
to
AReaL-unittests
GitHub Actions
Inactive
garrett4wade
left a comment
garrett4wade
left a comment
There was a problem hiding this comment.
LGTM. Waiting for CI
…upport (areal-project#849) Replace manual safetensors handling with PyTorch DCP infrastructure for HuggingFace format checkpoints. Previously, save/load required gathering full state dict to rank 0, causing OOM on large MoE models. Now uses HuggingFaceStorageWriter/Reader for distributed operations. Key changes: - Add MoEStateDictAdapter base class with DTensor-aware expert weight conversion (3D grouped <-> 2D individual), supporting EP and ETP placements via local expert index calculation - Fix in-place tensor ops (masked_fill_, scatter_) for torch.compile and checkpointing compatibility - Large MoE checkpoint save/load without OOM (no full tensor gather) - Correct EP/ETP checkpoint handling with distributed expert weights - Incremental weight sync to inference engines via convert_single_to_hf() - Foundation for PP support: layer-indexed key mapping and extensible DTensor placement logic for per-stage conversion
…upport (areal-project#849) Replace manual safetensors handling with PyTorch DCP infrastructure for HuggingFace format checkpoints. Previously, save/load required gathering full state dict to rank 0, causing OOM on large MoE models. Now uses HuggingFaceStorageWriter/Reader for distributed operations. Key changes: - Add MoEStateDictAdapter base class with DTensor-aware expert weight conversion (3D grouped <-> 2D individual), supporting EP and ETP placements via local expert index calculation - Fix in-place tensor ops (masked_fill_, scatter_) for torch.compile and checkpointing compatibility - Large MoE checkpoint save/load without OOM (no full tensor gather) - Correct EP/ETP checkpoint handling with distributed expert weights - Incremental weight sync to inference engines via convert_single_to_hf() - Foundation for PP support: layer-indexed key mapping and extensible DTensor placement logic for per-stage conversion
Description
Replace manual safetensors handling with PyTorch DCP infrastructure for HuggingFace format checkpoints. Previously, save/load required gathering full state dict to rank 0, causing OOM on large MoE models. Now uses HuggingFaceStorageWriter/Reader for distributed operations.
Key changes:
Type of Change
work as expected)
Checklist
jb build docs/gemini review)