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

areal/experimental/engine/archon_checkpoint.py

@@ -121,6 +121,31 @@ def load_model_from_hf(engine: ArchonEngine, path: str) -> None:
     # Convert back to Archon format
     archon_state_dict = engine.state_dict_adapter.from_hf(hf_state_dict)
 
+    # Compute key differences for diagnostics
+    model_keys = set(state_dict.keys())
+    loaded_keys = set(archon_state_dict.keys())
+
+    missing_keys = model_keys - loaded_keys
+    unexpected_keys = loaded_keys - model_keys
+
+    # Filter known expected missing keys
+    expected_missing = set()
+    for key in list(missing_keys):
+        # rotary_emb is computed at runtime, not stored in checkpoint
+        if "rotary_emb" in key:
+            expected_missing.add(key)
+    missing_keys -= expected_missing
+
+    if dist.get_rank() == 0:
+        if missing_keys:
+            engine.logger.warning(
+                f"Unexpected missing keys in checkpoint: {missing_keys}"
+            )
+        if unexpected_keys:
+            engine.logger.warning(
+                f"Unexpected extra keys in checkpoint: {unexpected_keys}"
+            )
+
     # Load into FSDP model (same as DCPState.load_state_dict)
     set_model_state_dict(
         engine.model,
@@ -142,6 +167,8 @@ def save_to_dcp(engine: ArchonEngine, path: str, with_optim: bool) -> None:
     state_dict = {"dcp": dcp_state}
     dcp.save(state_dict, checkpoint_id=path)
 
+    dist.barrier(group=engine.cpu_group)
+
 
 def load_from_dcp(engine: ArchonEngine, path: str, with_optim: bool) -> None:
     """Load model (and optionally optimizer) from DCP format."""
@@ -152,6 +179,8 @@ def load_from_dcp(engine: ArchonEngine, path: str, with_optim: bool) -> None:
     state_dict = {"dcp": dcp_state}
     dcp.load(state_dict=state_dict, checkpoint_id=path)
 
+    dist.barrier(group=engine.cpu_group)
+
 
 def save_optimizer_state(engine: ArchonEngine, path: str) -> None:
     """Save optimizer state to disk (sharded by rank)."""

areal/experimental/engine/archon_engine.py

@@ -85,7 +85,6 @@
 from areal.utils.network import find_free_ports, gethostip
 from areal.utils.offload import torch_memory_saver
 from areal.utils.perf_tracer import trace_perf
-from areal.utils.save_load import get_state_dict_from_repo_id_or_path
 
 
 @dataclass
@@ -265,6 +264,8 @@ def initialize(self, addr: str | None, ft_spec: FinetuneSpec, *args, **kwargs):
             f"Applied parallelism in {time.perf_counter() - tik:.2f} seconds"
         )
 
+        self._materialize_and_load_weights()
+
         self._create_optimizer(ft_spec)
 
         self._initialized = True
@@ -817,12 +818,20 @@ def _create_device_model(self):
         self.tokenizer = load_hf_tokenizer(self.config.path)
 
         tik = time.perf_counter()
-        with torch.device(current_platform.device_type):
-            model = self._create_llm_model()
 
-        self.logger.info(f"Model creation time: {time.perf_counter() - tik:.2f}s")
+        # Meta device mode: create structure only, no memory allocation
+        # Parameters exist only as metadata until materialized after FSDP
+        with torch.device("meta"):
+            model = self._create_model_structure()
+        model = model.to(getattr(torch, self.config.dtype))
         self.model = model
 
+        self.logger.info(
+            f"Model structure created on meta device in "
+            f"{time.perf_counter() - tik:.2f}s"
+        )
+        self.get_device_stats().log("after create model structure")
+
     def _build_ac_config(self) -> ActivationCheckpointConfig | None:
         # First check if gradient checkpointing is enabled
         if not self.config.gradient_checkpointing:
@@ -851,56 +860,43 @@ def _build_ac_config(self) -> ActivationCheckpointConfig | None:
 
         return ac_config
 
-    def _create_llm_model(self) -> nn.Module:
+    def _create_model_structure(self) -> nn.Module:
+        """Create model structure on meta device without loading weights."""
         model_args = self.spec.model_args_class.from_hf_config(
             self.model_config,
             is_critic=self.config.is_critic,
             attn_type=self.config.archon.attn_type,
         )
         model = self.spec.model_class(model_args)
+        return model
 
-        if not self.config.init_from_scratch:
-            self._load_hf_weights_to_archon_model(model)
+    def _materialize_and_load_weights(self):
+        """Materialize meta tensors and load weights after FSDP parallelization."""
+        if self.config.archon.offload_params:
+            init_device = "cpu"
+            buffer_device = current_platform.device_type
         else:
-            model.init_weights()
+            init_device = current_platform.device_type
+            buffer_device = init_device
 
-        dtype = getattr(torch, self.config.dtype)
-        model = model.to(dtype)
-
-        return model
+        tik = time.perf_counter()
 
-    def _load_hf_weights_to_archon_model(self, model: nn.Module):
-        self.logger.info(f"Loading weights from {self.config.path}")
+        self.model.to_empty(device=init_device)
 
-        hf_state_dict = get_state_dict_from_repo_id_or_path(self.config.path)
+        if not self.config.init_from_scratch:
+            load_model_from_hf(self, self.config.path)
+        else:
+            with torch.no_grad():
+                self.model.init_weights()
 
-        adapter = self.spec.state_dict_adapter_class(self.model_config)
+        self.model.init_buffers(buffer_device=buffer_device)
 
-        archon_state_dict = adapter.from_hf(hf_state_dict)
+        dist.barrier(group=self.cpu_group)
 
-        missing_keys, unexpected_keys = model.load_state_dict(
-            archon_state_dict, strict=False
+        self.logger.info(
+            f"Materialized and loaded weights in {time.perf_counter() - tik:.2f}s"
         )
-
-        if missing_keys:
-            expected_missing = {"rope_cache"}
-            if not self.config.is_critic:
-                expected_missing.add("score.weight")
-            else:
-                expected_missing.add("output.weight")
-
-            actual_missing = [k for k in missing_keys if k not in expected_missing]
-            if actual_missing:
-                self.logger.warning(
-                    f"Missing keys when loading weights: {actual_missing}"
-                )
-
-        if unexpected_keys:
-            self.logger.warning(
-                f"Unexpected keys when loading weights: {unexpected_keys}"
-            )
-
-        self.logger.info("Successfully loaded weights into Archon model")
+        self.get_device_stats().log("after materialize and load weights")
 
     def _create_optimizer(self, ft_spec: FinetuneSpec):
         if self.optimizer_config is None:

areal/experimental/models/archon/base.py

@@ -117,13 +117,6 @@ def convert_single_to_hf(
 class BaseArchonModel(nn.Module, ABC):
     """Base class for Archon models."""
 
-    model_args: BaseModelArgs
-    layers: nn.ModuleDict
-    tok_embeddings: nn.Embedding | None
-    norm: nn.Module | None
-    output: nn.Linear | None
-    score: nn.Linear | None
-
     @abstractmethod
     def forward(
         self,
@@ -134,7 +127,14 @@ def forward(
     ) -> torch.Tensor: ...
 
     @abstractmethod
-    def init_weights(self, buffer_device: torch.device | None = None) -> None: ...
+    def init_weights(self) -> None:
+        """Initialize model parameters."""
+        ...
+
+    @abstractmethod
+    def init_buffers(self, buffer_device: torch.device | str) -> None:
+        """Initialize model buffers (e.g., rope_cache)."""
+        ...
 
 
 __all__ = [

areal/experimental/models/archon/moe/moe.py

@@ -104,21 +104,8 @@ def __init__(self, moe_args: MoEArgs, dim: int, hidden_dim: int):
         # Buffers for auxiliary-loss-free load balancing
         # expert_bias is used to adjust routing probabilities
         # tokens_per_expert tracks usage for bias updates
-        if self.load_balance_coeff is not None:
-            assert self.load_balance_coeff > 0.0
-            self.register_buffer(
-                "expert_bias",
-                torch.zeros(moe_args.num_experts, dtype=torch.float32),
-                persistent=True,
-            )
-        else:
-            self.expert_bias = None
-
-        self.register_buffer(
-            "tokens_per_expert",
-            torch.zeros(moe_args.num_experts, dtype=torch.float32),
-            persistent=False,
-        )
+        self.expert_bias: torch.Tensor | None
+        self.tokens_per_expert: torch.Tensor
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Forward pass through MoE layer.
@@ -214,24 +201,39 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         return out.view(bs, slen, dim)
 
-    def init_weights(self, init_std: float, buffer_device: torch.device):
-        """Initialize weights.
+    def init_weights(self, init_std: float):
+        """Initialize MoE parameters.
 
         Args:
             init_std: Standard deviation for output projections.
-            buffer_device: Device for buffers.
         """
         self.experts.init_weights(init_std)
         self.router.init_weights(init_std)
 
         if self.shared_experts is not None:
             self.shared_experts.init_weights(init_std)
 
-        # Reset buffers
+    def init_buffers(self, buffer_device: torch.device | str):
+        """Initialize MoE buffers (tokens_per_expert, expert_bias).
+
+        Args:
+            buffer_device: Device for buffers.
+        """
         with torch.device(buffer_device):
-            self.tokens_per_expert = torch.zeros(self.num_experts, dtype=torch.float32)
+            self.register_buffer(
+                "tokens_per_expert",
+                torch.zeros(self.num_experts, dtype=torch.float32),
+                persistent=False,
+            )
             if self.load_balance_coeff is not None:
-                self.expert_bias = torch.zeros(self.num_experts, dtype=torch.float32)
+                assert self.load_balance_coeff > 0.0
+                self.register_buffer(
+                    "expert_bias",
+                    torch.zeros(self.num_experts, dtype=torch.float32),
+                    persistent=True,
+                )
+            else:
+                self.expert_bias = None
 
 
 __all__ = ["MoE", "FeedForward"]

areal/experimental/models/archon/qwen2/model/model.py

@@ -10,6 +10,7 @@
     SDPAWrapper,
     VarlenAttentionWrapper,
 )
+from areal.experimental.models.archon.base import BaseArchonModel
 from areal.experimental.models.archon.qwen2.model.args import Qwen2ModelArgs
 from areal.experimental.models.archon.qwen2.model.rope import (
     apply_rotary_emb,
@@ -241,14 +242,14 @@ def forward(
         x = x + self.feed_forward(self.ffn_norm(x))
         return x
 
-    def init_weights(self, buffer_device: torch.device):
+    def init_weights(self):
         for norm in (self.attention_norm, self.ffn_norm):
             norm.reset_parameters()
         self.attention.init_weights(self.weight_init_std)
         self.feed_forward.init_weights(self.weight_init_std)
 
 
-class Qwen2Model(nn.Module):
+class Qwen2Model(BaseArchonModel):
     """Qwen2 transformer model."""
 
     def __init__(self, model_args: Qwen2ModelArgs):
@@ -286,17 +287,13 @@ def _precompute_rope_cache(self) -> torch.Tensor:
             self.model_args.rope_theta,
         )
 
-    def init_weights(self, buffer_device: torch.device | None = None):
-        buffer_device = buffer_device or self.rope_cache.device
-        with torch.device(buffer_device):
-            self.rope_cache = self._precompute_rope_cache()
-
+    def init_weights(self):
         if self.tok_embeddings is not None:
             nn.init.normal_(self.tok_embeddings.weight)
 
         for layer in self.layers.values():
             if layer is not None:
-                layer.init_weights(buffer_device)
+                layer.init_weights()
 
         if self.norm is not None:
             self.norm.reset_parameters()
@@ -322,6 +319,10 @@ def init_weights(self, buffer_device: torch.device | None = None):
                 b=cutoff_factor * final_out_std,
             )
 
+    def init_buffers(self, buffer_device: torch.device | str):
+        with torch.device(buffer_device):
+            self.rope_cache = self._precompute_rope_cache()
+
     def forward(
         self,
         tokens: torch.Tensor,