diff --git a/areal/experimental/engine/archon_checkpoint.py b/areal/experimental/engine/archon_checkpoint.py new file mode 100644 index 0000000000..58eb3423c5 --- /dev/null +++ b/areal/experimental/engine/archon_checkpoint.py @@ -0,0 +1,179 @@ +from __future__ import annotations + +import os +import shutil +from typing import TYPE_CHECKING + +import torch +import torch.distributed as dist +import torch.distributed.checkpoint as dcp +from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, + set_model_state_dict, +) + +from areal.utils.fsdp.checkpoint import DCPState + +if TYPE_CHECKING: + from transformers import AutoProcessor, PreTrainedTokenizerFast + + from areal.experimental.engine.archon_engine import ArchonEngine + + +def save_model_to_hf( + engine: ArchonEngine, + path: str, + tokenizer: PreTrainedTokenizerFast | None, + processor: AutoProcessor | None = None, +) -> None: + """Save model in HuggingFace format using DCP infrastructure.""" + from torch.distributed.checkpoint import HuggingFaceStorageWriter + + if engine.model is None: + raise RuntimeError("Model not initialized") + if engine.state_dict_adapter is None: + raise RuntimeError("state_dict_adapter is required for HF format") + + engine.logger.info(f"Saving HF checkpoint to {path}") + os.makedirs(path, exist_ok=True) + + # Get distributed state dict + options = StateDictOptions(full_state_dict=False, cpu_offload=True) + state_dict = get_model_state_dict(engine.model, options=options) + + # Convert to HF format using adapter + hf_state_dict = engine.state_dict_adapter.to_hf(state_dict) + + fqn_to_index_mapping = engine.state_dict_adapter.fqn_to_index_mapping + + # NOTE: HuggingFaceStorageWriter always creates a sharded/ subdirectory when + # save_distributed=True. With enable_consolidation=True, it saves shards to + # path/sharded/, then consolidates to path/. The sharded/ directory is NOT + # automatically cleaned up by PyTorch, so we must remove it manually. + sharded_dir = os.path.join(path, "sharded") + + if fqn_to_index_mapping: + # Multi-file output: save to sharded/, then consolidate with all ranks + from torch.distributed.checkpoint._consolidate_hf_safetensors import ( + consolidate_safetensors_files_on_every_rank, + ) + + 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, + ) + else: + # Single-file output: auto consolidation + hf_writer = HuggingFaceStorageWriter( + path=path, + save_distributed=True, + enable_consolidation=True, + ) + dcp.save(hf_state_dict, storage_writer=hf_writer) + + # Clean up sharded/ directory after consolidation + if dist.get_rank() == 0 and os.path.exists(sharded_dir): + shutil.rmtree(sharded_dir) + + dist.barrier(group=engine.cpu_group) + + if dist.get_rank() == 0: + engine.model_config.save_pretrained(path) + if tokenizer is not None: + tokenizer.save_pretrained(path) + if processor is not None: + processor.save_pretrained(path) + dist.barrier(group=engine.cpu_group) + + +def load_model_from_hf(engine: ArchonEngine, path: str) -> None: + """Load model from HuggingFace format using DCP infrastructure.""" + if engine.model is None: + raise RuntimeError("Model not initialized") + if engine.state_dict_adapter is None: + raise RuntimeError("state_dict_adapter is required for HF format") + + engine.logger.info(f"Loading HF checkpoint from {path}") + + # Get model state dict structure (distributed) + options = StateDictOptions(full_state_dict=False, cpu_offload=True) + state_dict = get_model_state_dict(engine.model, options=options) + + # Convert to HF format to match checkpoint keys + hf_state_dict = engine.state_dict_adapter.to_hf(state_dict) + + # Load using DCP with HuggingFaceStorageReader + hf_reader = engine.state_dict_adapter.get_hf_storage_reader(path) + dcp.load(hf_state_dict, storage_reader=hf_reader) + + # Convert back to Archon format + archon_state_dict = engine.state_dict_adapter.from_hf(hf_state_dict) + + # Load into FSDP model (same as DCPState.load_state_dict) + set_model_state_dict( + engine.model, + model_state_dict=archon_state_dict, + options=StateDictOptions(strict=False), + ) + + dist.barrier(group=engine.cpu_group) + + +def save_to_dcp(engine: ArchonEngine, path: str, with_optim: bool) -> None: + """Save model (and optionally optimizer) using DCP format.""" + if engine.model is None: + raise RuntimeError("Model not initialized") + + os.makedirs(path, exist_ok=True) + + dcp_state = DCPState(engine.model, engine.optimizer if with_optim else None) + state_dict = {"dcp": dcp_state} + dcp.save(state_dict, checkpoint_id=path) + + +def load_from_dcp(engine: ArchonEngine, path: str, with_optim: bool) -> None: + """Load model (and optionally optimizer) from DCP format.""" + if engine.model is None: + raise RuntimeError("Model not initialized") + + dcp_state = DCPState(engine.model, engine.optimizer if with_optim else None) + state_dict = {"dcp": dcp_state} + dcp.load(state_dict=state_dict, checkpoint_id=path) + + +def save_optimizer_state(engine: ArchonEngine, path: str) -> None: + """Save optimizer state to disk (sharded by rank).""" + assert engine.optimizer is not None + assert dist.is_initialized() + rank = dist.get_rank() + shard_path = os.path.join( + path, f"optim_world_size_{engine.world_size}_rank_{rank}.pt" + ) + state_dict = engine.optimizer.state_dict() + torch.save(state_dict, shard_path) + dist.barrier(group=engine.cpu_group) + + +def load_optimizer_state(engine: ArchonEngine, path: str) -> None: + """Load optimizer state from disk (sharded by rank).""" + assert engine.optimizer is not None + assert dist.is_initialized() + rank = dist.get_rank() + shard_path = os.path.join( + path, f"optim_world_size_{engine.world_size}_rank_{rank}.pt" + ) + optimizer_state_dict = torch.load(shard_path, weights_only=False) + engine.optimizer.load_state_dict(optimizer_state_dict) + dist.barrier(group=engine.cpu_group) diff --git a/areal/experimental/engine/archon_engine.py b/areal/experimental/engine/archon_engine.py index 8c3e85ce8f..17c8669f18 100644 --- a/areal/experimental/engine/archon_engine.py +++ b/areal/experimental/engine/archon_engine.py @@ -12,13 +12,7 @@ import torch import torch.distributed as dist -import torch.distributed.checkpoint as dcp -from safetensors.torch import save_file from torch import nn -from torch.distributed.checkpoint.state_dict import ( - StateDictOptions, - get_model_state_dict, -) from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor import DTensor from torchdata.stateful_dataloader import StatefulDataLoader @@ -48,6 +42,14 @@ compute_total_loss_weight, reorder_and_pad_outputs, ) +from areal.experimental.engine.archon_checkpoint import ( + load_from_dcp, + load_model_from_hf, + load_optimizer_state, + save_model_to_hf, + save_optimizer_state, + save_to_dcp, +) from areal.experimental.models.archon import ( ArchonParallelDims, BaseStateDictAdapter, @@ -76,8 +78,7 @@ unsqueeze_mb_list, ) from areal.utils.distributed import init_custom_process_group, patch_dist_group_timeout -from areal.utils.fsdp import fsdp2_load_full_state_dict, get_cosine_schedule_with_warmup -from areal.utils.fsdp.checkpoint import DCPState +from areal.utils.fsdp import get_cosine_schedule_with_warmup from areal.utils.fsdp.grad import fsdp2_clip_grad_norm from areal.utils.functional import gather_logprobs, gather_logprobs_entropy from areal.utils.hf_utils import load_hf_tokenizer @@ -562,26 +563,26 @@ def update_weights(self, meta: WeightUpdateMeta): def save(self, meta: SaveLoadMeta): """Save model in HuggingFace or DCP format.""" if meta.weight_format == "hf": - self._save_model_to_hf(meta.path, meta.tokenizer, meta.processor) + save_model_to_hf(self, meta.path, meta.tokenizer, meta.processor) elif meta.weight_format == "dcp": - self._save_to_dcp(meta.path, meta.with_optim) + save_to_dcp(self, meta.path, meta.with_optim) else: raise ValueError(f"Unknown weight format {meta.weight_format}.") if meta.with_optim and meta.weight_format == "hf": - self._save_optimizer_state(meta.path) + save_optimizer_state(self, meta.path) def load(self, meta: SaveLoadMeta): """Load model from HuggingFace or DCP format.""" if meta.weight_format == "hf": - self._load_model_from_hf(meta.path) + load_model_from_hf(self, meta.path) elif meta.weight_format == "dcp": - self._load_from_dcp(meta.path, meta.with_optim) + load_from_dcp(self, meta.path, meta.with_optim) else: raise ValueError(f"Unknown weight format {meta.weight_format}.") if meta.with_optim and meta.weight_format == "hf": - self._load_optimizer_state(meta.path) + load_optimizer_state(self, meta.path) def offload(self) -> None: """Offload model memory to CPU using torch_memory_saver.""" @@ -642,7 +643,9 @@ def _validate_model_type(self) -> None: ) def _create_state_dict_adapter(self) -> BaseStateDictAdapter | None: - return self.spec.state_dict_adapter_class(self.model_config) + return self.spec.state_dict_adapter_class( + self.model_config, hf_assets_path=self.config.path + ) def _get_model_name_parameters(self) -> Iterator[tuple[str, nn.Parameter]]: return self.model.named_parameters() @@ -790,7 +793,7 @@ def _update_weights_from_disk(self, meta: WeightUpdateMeta): fut = self.rollout_engine.update_weights_from_disk(meta) assert meta.path is not None - self._save_model_to_hf(meta.path, self.tokenizer, None) + save_model_to_hf(self, meta.path, self.tokenizer, None) if dist.get_rank() == 0: update_name = names.update_weights_from_disk( @@ -807,95 +810,6 @@ def _update_weights_from_disk(self, meta: WeightUpdateMeta): current_platform.synchronize() dist.barrier(group=self.cpu_group) - def _save_model_to_hf( - self, - path: str, - tokenizer: PreTrainedTokenizerFast | None, - processor=None, - ): - """Save model in HuggingFace format.""" - if self.model is None: - raise RuntimeError("Model not initialized") - os.makedirs(path, exist_ok=True) - - options = StateDictOptions(full_state_dict=True, cpu_offload=True) - state_dict = get_model_state_dict(self.model, options=options) - if self.state_dict_adapter is not None: - state_dict = self.state_dict_adapter.to_hf(state_dict) - - if dist.get_rank() == 0: - os.makedirs(path, exist_ok=True) - model_path = os.path.join(path, "model.safetensors") - try: - save_file(state_dict, model_path) - except ImportError: - model_path = os.path.join(path, "pytorch_model.bin") - torch.save(state_dict, model_path) - - self.model_config.save_pretrained(path) - if tokenizer is not None: - tokenizer.save_pretrained(path) - if processor is not None: - processor.save_pretrained(path) - dist.barrier(group=self.cpu_group) - - def _load_model_from_hf(self, path: str): - """Load model from HuggingFace format.""" - if dist.get_rank() == 0: - full_state = get_state_dict_from_repo_id_or_path(path) - if self.state_dict_adapter is not None: - full_state = self.state_dict_adapter.from_hf(full_state) - else: - full_state = {} - - cpu_offload = self.config.archon.offload_params - fsdp2_load_full_state_dict( - self.model, - full_state, - cpu_offload, - tie_word_embeddings=self.model_config.tie_word_embeddings, - ) - - def _save_to_dcp(self, path: str, with_optim: bool): - if self.model is None: - raise RuntimeError("Model not initialized") - - os.makedirs(path, exist_ok=True) - - dcp_state = DCPState(self.model, self.optimizer if with_optim else None) - state_dict = {"dcp": dcp_state} - dcp.save(state_dict, checkpoint_id=path) - - def _load_from_dcp(self, path: str, with_optim: bool): - if self.model is None: - raise RuntimeError("Model not initialized") - - dcp_state = DCPState(self.model, self.optimizer if with_optim else None) - state_dict = {"dcp": dcp_state} - dcp.load(state_dict=state_dict, checkpoint_id=path) - - def _save_optimizer_state(self, path: str): - assert self.optimizer is not None - assert dist.is_initialized() - rank = dist.get_rank() - shard_path = os.path.join( - path, f"optim_world_size_{self.world_size}_rank_{rank}.pt" - ) - state_dict = self.optimizer.state_dict() - torch.save(state_dict, shard_path) - dist.barrier(group=self.cpu_group) - - def _load_optimizer_state(self, path: str): - assert self.optimizer is not None - assert dist.is_initialized() - rank = dist.get_rank() - shard_path = os.path.join( - path, f"optim_world_size_{self.world_size}_rank_{rank}.pt" - ) - optimizer_state_dict = torch.load(shard_path, weights_only=False) - self.optimizer.load_state_dict(optimizer_state_dict) - dist.barrier(group=self.cpu_group) - def _create_device_model(self): current_platform.set_device(int(os.environ["LOCAL_RANK"])) self.device = torch.device(int(os.environ["LOCAL_RANK"])) diff --git a/areal/experimental/models/archon/__init__.py b/areal/experimental/models/archon/__init__.py index 1ea3cbc5f2..a4f7ac38db 100644 --- a/areal/experimental/models/archon/__init__.py +++ b/areal/experimental/models/archon/__init__.py @@ -16,6 +16,10 @@ get_supported_model_types, is_supported_model, ) +from areal.experimental.models.archon.moe_weight_converter import ( + MoEConversionState, + MoEWeightConverter, +) from areal.experimental.models.archon.parallel_dims import ( ArchonParallelDims, ) @@ -25,6 +29,8 @@ "BaseStateDictAdapter", "ExpertParallel", "ExpertTensorParallel", + "MoEConversionState", + "MoEWeightConverter", "ModelSpec", "TensorParallel", "apply_expert_parallel", diff --git a/areal/experimental/models/archon/activation_checkpoint.py b/areal/experimental/models/archon/activation_checkpoint.py index fc6a2577db..833a127250 100644 --- a/areal/experimental/models/archon/activation_checkpoint.py +++ b/areal/experimental/models/archon/activation_checkpoint.py @@ -26,10 +26,7 @@ @dataclass class ActivationCheckpointConfig: - """Activation checkpointing configuration. - - Aligned with torchtitan.config.job_config.ActivationCheckpoint for compatibility. - """ + """Activation checkpointing configuration.""" mode: str = "selective" """Type of activation checkpointing to use: 'selective', 'full', 'memory_budget', 'none'""" @@ -250,8 +247,6 @@ def apply_ac( ) -> None: """Apply activation checkpointing to the model. - Aligned with torchtitan.distributed.activation_checkpoint.apply_ac - Args: model: The model to apply activation checkpointing to. ac_config: The activation checkpointing config. diff --git a/areal/experimental/models/archon/attention.py b/areal/experimental/models/archon/attention.py index 65673b28ec..dd29a44223 100644 --- a/areal/experimental/models/archon/attention.py +++ b/areal/experimental/models/archon/attention.py @@ -59,7 +59,7 @@ def create_block_causal_mask_2d( causal = positions.unsqueeze(0) <= positions.unsqueeze(1) mask = torch.full((seq_len, seq_len), float("-inf"), device=device, dtype=dtype) - mask.masked_fill_(same_seq & causal, 0.0) + mask = mask.masked_fill(same_seq & causal, 0.0) return mask diff --git a/areal/experimental/models/archon/base.py b/areal/experimental/models/archon/base.py index 6315183c0b..43daf98ff2 100644 --- a/areal/experimental/models/archon/base.py +++ b/areal/experimental/models/archon/base.py @@ -1,15 +1,23 @@ from __future__ import annotations +import json +import os +import re from abc import ABC, abstractmethod from dataclasses import dataclass from typing import TYPE_CHECKING, Any import torch import torch.nn as nn +from torch.distributed.checkpoint import HuggingFaceStorageReader + +from areal.utils import logging if TYPE_CHECKING: from transformers import PretrainedConfig +logger = logging.getLogger("ArchonModelBase") + @dataclass class BaseModelArgs(ABC): @@ -29,12 +37,70 @@ def from_hf_config( class BaseStateDictAdapter(ABC): - """Base class for HF <-> Archon state dict conversion.""" + """Base class for HF <-> Archon state dict conversion. + + Args: + model_config: HuggingFace model configuration + hf_assets_path: Path to HF assets folder containing tokenizer, model weights, etc. + If provided and contains model.safetensors.index.json, the index will be + parsed to build fqn_to_index_mapping for multi-file checkpoint support. + """ + + fqn_to_index_mapping: dict[str, int] | None - def __init__(self, model_config: PretrainedConfig): + def __init__( + self, model_config: PretrainedConfig, hf_assets_path: str | None = None + ): self.model_config = model_config self.from_hf_map: dict[str, str | None] = {} self.to_hf_map: dict[str, str] = {} + self.hf_assets_path = hf_assets_path + self.fqn_to_index_mapping = None + + if hf_assets_path: + self._load_safetensors_index(hf_assets_path) + + 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) @abstractmethod def from_hf(self, hf_state_dict: dict[str, Any]) -> dict[str, Any]: ... diff --git a/areal/experimental/models/archon/moe/router.py b/areal/experimental/models/archon/moe/router.py index c05626d916..86cb0cf93d 100644 --- a/areal/experimental/models/archon/moe/router.py +++ b/areal/experimental/models/archon/moe/router.py @@ -122,7 +122,7 @@ def _get_node_limited_routing_scores( # Create mask: False for selected groups (keep), True for others (mask) group_mask = torch.ones_like(group_scores, dtype=torch.bool) - group_mask.scatter_(1, group_idx, False) + group_mask = group_mask.scatter(1, group_idx, False) # Mask out experts from non-selected groups scores_for_choice = scores_grouped.masked_fill( diff --git a/areal/experimental/models/archon/moe_weight_converter.py b/areal/experimental/models/archon/moe_weight_converter.py new file mode 100644 index 0000000000..ca552dd751 --- /dev/null +++ b/areal/experimental/models/archon/moe_weight_converter.py @@ -0,0 +1,444 @@ +# Adapted from torchtitan: torchtitan/models/utils.py + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +import torch +from torch.distributed.tensor import DTensor +from torch.distributed.tensor.placement_types import Replicate, Shard, _StridedShard + +if TYPE_CHECKING: + from torch.distributed.device_mesh import DeviceMesh + + +@dataclass +class MoEConversionState: + """Explicit state container for DTensor conversion metadata. + + This dataclass holds the metadata needed for DTensor-aware MoE expert weight + conversion. The state is populated during to_hf (split) operations and consumed + during from_hf (concatenate) operations. + + Important: For DTensor checkpoints, the conversion methods have a stateful dependency: + `split_expert_weights_dtensor()` (called during to_hf) populates metadata that + `concatenate_expert_weights_dtensor()` (called during from_hf) requires. + Therefore, when working with DTensor expert weights, the same state instance + must be used for both to_hf() and from_hf() to ensure the metadata is available. + """ + + grouped_expert_weight_placements: dict[str, tuple] = field(default_factory=dict) + grouped_expert_weight_shape: dict[str, tuple[int, ...]] = field( + default_factory=dict + ) + local_experts_indices: dict[str, tuple[int, int]] = field(default_factory=dict) + + def clear(self) -> None: + """Clear all stored conversion metadata.""" + self.grouped_expert_weight_placements.clear() + self.grouped_expert_weight_shape.clear() + self.local_experts_indices.clear() + + +class MoEWeightConverter: + """DTensor-aware MoE expert weight converter. + + HF MoE models store experts as a module list each with 2D weights. In Archon, we + store experts as a 3D param with the first dimension being num_experts. The methods + in this class help convert 3D param into list of 2D params so that the checkpoint + can be loaded without incurring local memory overhead, and then concatenate + the results back to 3D param. + + This class provides: + - DTensor-aware methods for distributed checkpoint support (EP, EP+TP, EP+ETP) + - Offline conversion methods for non-distributed scenarios + + Note: This class is designed to be composed into StateDictAdapter classes rather + than inherited from. The conversion state is managed externally via MoEConversionState. + """ + + @staticmethod + def calculate_strided_shard_indices( + strided_shard_dim_degree: int, + strided_shard_dim_rank: int, + shard_dim_degree: int, + shard_dim_rank: int, + dim_size_to_split: int, + ) -> tuple[int, int]: + """Calculate start/end indices for [StridedShard(dim=i), Shard(dim=i)] placement. + + GPU Layout (strided_shard_rank, shard_rank): + + StridedShard Rank Shard rank + ┌─────────────────┐ + 0 │ GPU(0, 0) │ 0 + ────┼─────────────────┤ + 1 │ GPU(1, 0) │ + ────┼─────────────────┤ + 2 │ GPU(2, 0) │ + ──────┼─────────────────┼──── + 0 │ GPU(0, 1) │ 1 + ────┼─────────────────┤ + 1 │ GPU(1, 1) │ + ────┼─────────────────┤ + 2 │ GPU(2, 1) │ + └─────────────────┘ + + Calculates start_index from inner dimension (Shard) to outer dimension (StridedShard). + + Args: + strided_shard_dim_degree: Degree of the StridedShard mesh dimension + strided_shard_dim_rank: Rank in the StridedShard mesh dimension + shard_dim_degree: Degree of the Shard mesh dimension + shard_dim_rank: Rank in the Shard mesh dimension + dim_size_to_split: Total size of the dimension being split + + Returns: + Tuple of (start_index, end_index) for the local GPU + + Raises: + ValueError: If dimension cannot be evenly split + """ + block_size = dim_size_to_split // (strided_shard_dim_degree * shard_dim_degree) + + # Error out if can not evenly divide + if ( + block_size * (strided_shard_dim_degree * shard_dim_degree) + != dim_size_to_split + ): + raise ValueError( + f"Cannot evenly split dim_size {dim_size_to_split} with " + f"strided_shard_degree={strided_shard_dim_degree}, shard_degree={shard_dim_degree}" + ) + + start_index = block_size * ( + strided_shard_dim_degree * shard_dim_rank + strided_shard_dim_rank + ) + end_index = start_index + block_size + + return start_index, end_index + + @staticmethod + def calculate_indices_from_placements( + dim: int, + dim_size: int, + dtensor_placements: tuple, + device_mesh: DeviceMesh, + ) -> tuple[int | None, int | None]: + """Calculate local indices for a given dimension from DTensor placements. + + Handles various sharding strategies: + - 0 placements on dim: No split needed, returns (None, None) + - 1 placement (Shard): Simple division by shard degree + - 2 placements (StridedShard + Shard): Complex EP+ETP sharding + + Args: + dim: The dimension to calculate indices for + dim_size: Total size of the dimension + dtensor_placements: Tuple of DTensor placements + device_mesh: DeviceMesh for the DTensor + + Returns: + Tuple of (start_index, end_index) or (None, None) if no split on this dim + """ + mesh_names: list[str] = [] + dim_i_placements: list = [] + + # Find all device mesh dimensions that shard on dim-i + for i, name in enumerate(device_mesh.mesh_dim_names): + placement = dtensor_placements[i] + if isinstance(placement, (Shard, _StridedShard)): + if placement.dim == dim: + mesh_names.append(name) + dim_i_placements.append(placement) + elif isinstance(placement, Replicate): + # Replicate does not shard on any dimension, skip + pass + else: + raise ValueError( + f"Unexpected placement type: {type(placement).__name__} on " + f"mesh dim '{name}'. Expected Shard, _StridedShard, or Replicate." + ) + + # Calculate local indices based on sharding strategy + start_index, end_index = None, None + + if len(dim_i_placements) == 2: + # Handle StridedShard(i) + Shard(i) case (e.g., EP + ETP) + assert isinstance(dim_i_placements[0], _StridedShard), ( + "Expected StridedShard as first placement" + ) + + strided_shard_mesh = device_mesh[mesh_names[0]] + shard_mesh = device_mesh[mesh_names[1]] + + strided_degree = strided_shard_mesh.size() + strided_rank = strided_shard_mesh.get_local_rank() + shard_degree = shard_mesh.size() + shard_rank = shard_mesh.get_local_rank() + + start_index, end_index = MoEWeightConverter.calculate_strided_shard_indices( + strided_degree, strided_rank, shard_degree, shard_rank, dim_size + ) + + elif len(dim_i_placements) == 1: + # Handle single Shard(i) case (e.g., EP only) + assert not isinstance(dim_i_placements[0], _StridedShard), ( + "Expected regular Shard, not StridedShard" + ) + + shard_mesh = device_mesh[mesh_names[0]] + shard_degree = shard_mesh.size() + shard_rank = shard_mesh.get_local_rank() + + block_size = dim_size // shard_degree + if block_size * shard_degree != dim_size: + raise ValueError( + f"Dim {dim} size ({dim_size}) cannot be evenly divided by " + f"shard degree ({shard_degree})" + ) + + start_index = block_size * shard_rank + end_index = start_index + block_size + + elif len(dim_i_placements) == 0: + # No split on this dimension + return start_index, end_index + + else: + raise NotImplementedError( + f"Unsupported DTensor placements for GroupedExperts: " + f"{dtensor_placements} {dim_i_placements} {mesh_names}" + ) + + return start_index, end_index + + def split_expert_weights_dtensor( + self, + hf_abstract_key: str, + archon_abstract_key: str, + layer_id: str, + grouped_expert_weight: DTensor, + state: MoEConversionState, + ) -> dict[str, DTensor]: + """Split GroupedExperts weight into individual expert weights for distributed save. + + This method handles various sharding strategies for expert weights: + - FSDP + EP: StridedShard(0)Shard(0) or Shard(0) + - FSDP + ETP + EP: StridedShard(0)Shard(0)Shard(1/2) or StridedShard(1)Shard(0)Shard(1/2) + + Args: + hf_abstract_key: HuggingFace template key with {} placeholders for layer and expert IDs + archon_abstract_key: Archon template key with {} placeholder for layer ID + layer_id: Layer identifier (string) + grouped_expert_weight: DTensor containing all experts' weights (3D) + state: MoEConversionState to store metadata for from_hf reconstruction + + Returns: + Dictionary mapping individual expert HF keys to their DTensor weights (2D) + """ + device_mesh = grouped_expert_weight.device_mesh + dtensor_placements = grouped_expert_weight.placements + + # Use concrete key (with layer_id) to avoid collision between layers + archon_key = archon_abstract_key.format(layer_id) + + # Step 1: Store metadata for from_hf reconstruction + state.grouped_expert_weight_placements[archon_key] = dtensor_placements + state.grouped_expert_weight_shape[archon_key] = tuple( + grouped_expert_weight.shape + ) + + # Step 2: Calculate local expert indices from placements + num_experts = grouped_expert_weight.shape[0] + start_index, end_index = self.calculate_indices_from_placements( + dim=0, + dim_size=num_experts, + dtensor_placements=dtensor_placements, + device_mesh=device_mesh, + ) + assert start_index is not None and end_index is not None, ( + "Start index and end index cannot be None on dim-0!" + ) + + # Step 3: Store indices for from_hf reconstruction + state.local_experts_indices[archon_key] = (start_index, end_index) + + # Step 4: Create new placements for individual expert weights (2D) + # Expert dimension (dim-0) becomes Replicate, other dims shift down by 1 + new_placements = [] + for i, _name in enumerate(device_mesh.mesh_dim_names): + placement = dtensor_placements[i] + if hasattr(placement, "dim") and placement.dim == 0: + # Expert dimension is removed, convert to Replicate + new_placements.append(Replicate()) + elif isinstance(placement, Shard): + # Other Shard dimensions keep same dim (2D expert weight) + new_placements.append(Shard(placement.dim)) + elif isinstance(placement, _StridedShard): + # Keep strided shard with same parameters + new_placements.append( + _StridedShard(placement.dim, placement.split_factor) + ) + else: + # Replicate stays as Replicate + new_placements.append(placement) + + # Step 5: Extract local tensor and split into individual experts + local_grouped_weights = grouped_expert_weight._local_tensor + expected_local_experts = end_index - start_index + + if local_grouped_weights.shape[0] != expected_local_experts: + raise ValueError( + f"Local tensor shape mismatch: expected {expected_local_experts} experts, " + f"got {local_grouped_weights.shape[0]}" + ) + + local_expert_tensors: dict[str, DTensor] = {} + for expert_id in range(start_index, end_index): + hf_key = hf_abstract_key.format(layer_id, expert_id) + local_expert_index = expert_id - start_index + + # Extract individual expert weight and add temporary batch dimension + expert_weight = local_grouped_weights[local_expert_index, :, :].unsqueeze(0) + + # Create DTensor and remove batch dimension + expert_dtensor = DTensor.from_local( + expert_weight, device_mesh, tuple(new_placements), run_check=False + ).squeeze(0) + + local_expert_tensors[hf_key] = expert_dtensor + + return local_expert_tensors + + def concatenate_expert_weights_dtensor( + self, + expert_weights_by_layer: dict[str, dict[str, dict[int, Any]]], + archon_abstract_key: str, + layer_id: str, + device_mesh: DeviceMesh, + state: MoEConversionState, + ) -> DTensor | None: + """Concatenate individual expert weights back into GroupedExperts DTensor. + + Args: + expert_weights_by_layer: Dictionary tracking expert weights by layer, abstract key, and expert ID. + Structure: {layer_id: {abstract_key: {expert_id: tensor_weight}}} + archon_abstract_key: Archon template key with {} placeholder for layer ID + layer_id: Layer identifier (string) + device_mesh: DeviceMesh for the target GroupedExperts weight DTensor + state: MoEConversionState containing metadata from split operation + + Returns: + Concatenated GroupedExperts weight DTensor if all experts are available, otherwise None + """ + if layer_id not in expert_weights_by_layer: + return None + if archon_abstract_key not in expert_weights_by_layer[layer_id]: + return None + + # Use concrete key (with layer_id) to lookup metadata + archon_key = archon_abstract_key.format(layer_id) + + experts = expert_weights_by_layer[layer_id][archon_abstract_key] + expected_n_experts = ( + state.local_experts_indices[archon_key][1] + - state.local_experts_indices[archon_key][0] + ) + if len(experts) < expected_n_experts: + return None + + # Sort and stack expert weights + sorted_expert_ids = sorted(experts.keys()) + sorted_experts = [experts[i] for i in sorted_expert_ids] + + # Stack and get local tensor (input experts are DTensors) + stacked = torch.stack(sorted_experts, dim=0) + if isinstance(stacked, DTensor): + local_tensor = stacked._local_tensor + else: + local_tensor = stacked + + # Verify we have stored placements + assert ( + archon_key in state.grouped_expert_weight_placements + and archon_key in state.grouped_expert_weight_shape + ), ( + f"GroupedExperts weight metadata (placements, shape) for {archon_key} cannot be None!" + ) + + # Reconstruct DTensor with original 3D placements + stacked_dtensor = DTensor.from_local( + local_tensor, + device_mesh, + state.grouped_expert_weight_placements[archon_key], + run_check=False, + ) + + # Cleanup + del expert_weights_by_layer[layer_id][archon_abstract_key] + if not expert_weights_by_layer[layer_id]: + del expert_weights_by_layer[layer_id] + + return stacked_dtensor + + @staticmethod + def split_expert_weights_offline( + weight: torch.Tensor, n_experts: int + ) -> tuple[torch.Tensor, ...]: + """Split 3D expert weight into tuple of 2D weights. Used for offline conversion. + + NOTE: If used for online conversion with DTensors, torch.split() might incur + communication to gather the weight, causing OOM. + + Args: + weight: 3D tensor of shape (n_experts, out_dim, in_dim) + n_experts: Number of experts + + Returns: + Tuple of 2D tensors, one per expert + """ + split_weight = torch.split(weight, weight.shape[0] // n_experts, dim=0) + return split_weight + + @staticmethod + def concatenate_expert_weights_offline( + expert_weights_by_layer: dict[str, dict[str, dict[int, torch.Tensor]]], + abstract_key: str, + layer_num: str, + n_experts: int, + ) -> torch.Tensor | None: + """Concatenate 2D expert weights into 3D using torch.stack(). Used for offline conversion. + + Args: + expert_weights_by_layer: Dictionary tracking expert weights by layer, abstract key, and expert ID. + Structure: {layer_id: {abstract_key: {expert_id: tensor_weight}}} + abstract_key: Archon template key with {} placeholder for layer ID + layer_num: Layer identifier (string) + n_experts: Number of experts in the GroupedExperts module + + Returns: + Concatenated GroupedExperts weight if all experts are available, otherwise None + """ + if layer_num not in expert_weights_by_layer: + return None + if abstract_key not in expert_weights_by_layer[layer_num]: + return None + + experts = expert_weights_by_layer[layer_num][abstract_key] + if len(experts) < n_experts: + return None + + sorted_expert_ids = sorted(experts.keys()) + sorted_experts = [experts[i] for i in sorted_expert_ids] + stacked_tensor = torch.stack(sorted_experts, dim=0) + + del expert_weights_by_layer[layer_num][abstract_key] + if not expert_weights_by_layer[layer_num]: + del expert_weights_by_layer[layer_num] + + return stacked_tensor + + +__all__ = ["MoEConversionState", "MoEWeightConverter"] diff --git a/areal/experimental/models/archon/qwen2/model/state_dict_adapter.py b/areal/experimental/models/archon/qwen2/model/state_dict_adapter.py index b1189685b8..7021282e2b 100644 --- a/areal/experimental/models/archon/qwen2/model/state_dict_adapter.py +++ b/areal/experimental/models/archon/qwen2/model/state_dict_adapter.py @@ -6,7 +6,6 @@ from typing import TYPE_CHECKING, Any import torch -from torch.distributed.tensor import DTensor from areal.experimental.models.archon.base import BaseStateDictAdapter @@ -17,8 +16,10 @@ class Qwen2StateDictAdapter(BaseStateDictAdapter): """State dict adapter for Qwen2 models.""" - def __init__(self, model_config: PretrainedConfig): - super().__init__(model_config) + def __init__( + self, model_config: PretrainedConfig, hf_assets_path: str | None = None + ): + super().__init__(model_config, hf_assets_path) # HuggingFace -> Archon key mapping self.from_hf_map = { @@ -59,8 +60,6 @@ def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]: # Regular key mapping hf_key = self._convert_key_to_hf(key) if hf_key is not None: - if isinstance(value, DTensor): - value = value.full_tensor() hf_state_dict[hf_key] = value return hf_state_dict @@ -97,8 +96,6 @@ def convert_single_to_hf( hf_key = self._convert_key_to_hf(name) if hf_key is not None: - if isinstance(tensor, DTensor): - tensor = tensor.full_tensor() return [(hf_key, tensor)] return [] diff --git a/areal/experimental/models/archon/qwen3/model/state_dict_adapter.py b/areal/experimental/models/archon/qwen3/model/state_dict_adapter.py index a0f02a8013..e41631e949 100644 --- a/areal/experimental/models/archon/qwen3/model/state_dict_adapter.py +++ b/areal/experimental/models/archon/qwen3/model/state_dict_adapter.py @@ -9,6 +9,10 @@ from torch.distributed.tensor import DTensor from areal.experimental.models.archon.base import BaseStateDictAdapter +from areal.experimental.models.archon.moe_weight_converter import ( + MoEConversionState, + MoEWeightConverter, +) if TYPE_CHECKING: from transformers import PretrainedConfig @@ -20,11 +24,14 @@ class Qwen3StateDictAdapter(BaseStateDictAdapter): Handles: - Key name mapping between HF and Archon conventions - MoE expert weights: HF uses list of 2D weights, Archon uses 3D weights + - DTensor-aware distributed checkpoint support for MoE (EP, EP+TP, EP+ETP) - No weight permutation needed (unlike Llama3) """ - def __init__(self, model_config: PretrainedConfig): - super().__init__(model_config) + def __init__( + self, model_config: PretrainedConfig, hf_assets_path: str | None = None + ): + super().__init__(model_config, hf_assets_path) # HuggingFace -> Archon key mapping # Use {} as placeholder for layer numbers and expert ids @@ -70,31 +77,59 @@ def __init__(self, model_config: PretrainedConfig): num_experts = getattr(model_config, "num_experts", None) if num_experts is None: num_experts = getattr(model_config, "num_local_experts", None) - self.moe_enabled = num_experts is not None and num_experts > 1 + if num_experts is None: + moe_args = getattr(model_config, "moe_args", None) + if moe_args is not None: + num_experts = getattr(moe_args, "num_experts", None) + moe_enabled_flag = getattr(model_config, "moe_enabled", None) + self.moe_enabled = (num_experts is not None and num_experts > 1) or ( + moe_enabled_flag is True + ) if self.moe_enabled: + if num_experts is None: + raise ValueError( + "moe_enabled is True but num_experts could not be determined from " + "model_config. Expected one of: num_experts, num_local_experts, or " + "moe_args.num_experts" + ) self.num_experts = num_experts # Weight tying configuration self.enable_weight_tying = getattr(model_config, "tie_word_embeddings", False) + # HF abstract key templates for MoE expert weights + self._hf_expert_abstract_keys = { + "layers.{}.moe.experts.w1": "model.layers.{}.mlp.experts.{}.gate_proj.weight", + "layers.{}.moe.experts.w2": "model.layers.{}.mlp.experts.{}.down_proj.weight", + "layers.{}.moe.experts.w3": "model.layers.{}.mlp.experts.{}.up_proj.weight", + } + + # Composition: create MoE weight converter for DTensor-aware conversion + self._moe_converter = MoEWeightConverter() if self.moe_enabled else None + self._moe_state = MoEConversionState() if self.moe_enabled else None + def to_hf(self, state_dict: dict[str, Any]) -> dict[str, Any]: """Convert Archon state dict to HuggingFace format. Main transformations: 1. Key renaming: Archon names -> HF names 2. MoE: 3D weights (num_experts, out, in) -> list of 2D weights - 3. Weight tying: skip output.weight if enabled (shared with embeddings) + 3. Weight tying: skip output.weight if enabled (HF reconstructs from embeddings) """ hf_state_dict = {} 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 - hf_pairs = self._split_moe_experts(key, value) + 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 @@ -122,12 +157,45 @@ def from_hf(self, hf_state_dict: dict[str, Any]) -> dict[str, Any]: hf_state_dict["lm_head.weight"] = hf_state_dict["model.embed_tokens.weight"] state_dict = {} + # For DTensor path: {layer_id: {abstract_key: {expert_id: DTensor}}} + expert_weights_by_layer: dict[str, dict[str, dict[int, Any]]] = {} + # For offline path: {archon_key: (3D tensor, count)} expert_buffer: dict[str, tuple[torch.Tensor, int]] = {} for key, value in hf_state_dict.items(): - if ".mlp.experts." in key: - # Collect expert weights, merge later - self._collect_expert_weight(key, value, expert_buffer, state_dict) + if ".mlp.experts." in key and self.moe_enabled: + parsed = self._parse_expert_key(key) + if parsed is None: + continue + layer_id, expert_id, archon_abstract_key = parsed + + if isinstance(value, DTensor): + # DTensor path: collect and concatenate with distributed awareness + if layer_id not in expert_weights_by_layer: + expert_weights_by_layer[layer_id] = {} + if archon_abstract_key not in expert_weights_by_layer[layer_id]: + expert_weights_by_layer[layer_id][archon_abstract_key] = {} + + expert_weights_by_layer[layer_id][archon_abstract_key][ + expert_id + ] = value + + assert ( + self._moe_converter is not None and self._moe_state is not None + ), "MoE converter not initialized for MoE model" + result = self._moe_converter.concatenate_expert_weights_dtensor( + expert_weights_by_layer, + archon_abstract_key, + layer_id, + value.device_mesh, + self._moe_state, + ) + if result is not None: + archon_key = archon_abstract_key.format(layer_id) + state_dict[archon_key] = result + else: + # Offline path: use pre-allocated buffer + self._collect_expert_weight(key, value, expert_buffer, state_dict) else: # Regular key mapping archon_key = self._convert_key_from_hf(key) @@ -200,6 +268,49 @@ def _convert_key_from_hf(self, hf_key: str) -> str | None: return None + def _parse_expert_key(self, hf_key: str) -> tuple[str, int, str] | None: + """Parse HF expert key into (layer_id, expert_id, archon_abstract_key).""" + match = re.match( + r"model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight", + hf_key, + ) + if not match: + return None + + layer_id = match.group(1) + expert_id = int(match.group(2)) + proj_type = match.group(3) + + proj_map = { + "gate_proj": "layers.{}.moe.experts.w1", + "up_proj": "layers.{}.moe.experts.w3", + "down_proj": "layers.{}.moe.experts.w2", + } + archon_abstract_key = proj_map[proj_type] + + return layer_id, expert_id, archon_abstract_key + + def _split_moe_experts_distributed( + self, archon_key: str, weight_3d: DTensor + ) -> dict[str, DTensor]: + """Split 3D DTensor expert weight into 2D DTensors for distributed checkpoint.""" + match = re.search(r"layers\.(\d+)\.", archon_key) + if not match: + return {} + layer_id = match.group(1) + + archon_abstract_key = re.sub(r"layers\.\d+\.", "layers.{}.", archon_key) + hf_abstract_key = self._hf_expert_abstract_keys.get(archon_abstract_key) + if hf_abstract_key is None: + return {} + + assert self._moe_converter is not None and self._moe_state is not None, ( + "MoE converter not initialized for MoE model" + ) + return self._moe_converter.split_expert_weights_dtensor( + hf_abstract_key, archon_abstract_key, layer_id, weight_3d, self._moe_state + ) + def _split_moe_experts( self, archon_key: str, weight_3d: torch.Tensor ) -> dict[str, torch.Tensor]: @@ -263,20 +374,12 @@ def _collect_expert_weight( state_dict: Output state dict to write merged 3D weights """ # Parse key: extract layer_num, expert_num, and proj type - match = re.match( - r"model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight", - hf_key, - ) - if not match: + parsed = self._parse_expert_key(hf_key) + if parsed is None: return - layer_num = match.group(1) - expert_num = int(match.group(2)) - proj_type = match.group(3) - - # Map proj type to Archon weight name - proj_map = {"gate_proj": "w1", "up_proj": "w3", "down_proj": "w2"} - archon_key = f"layers.{layer_num}.moe.experts.{proj_map[proj_type]}" + layer_num, expert_num, archon_abstract_key = parsed + archon_key = archon_abstract_key.format(layer_num) # Pre-allocate 3D tensor on first expert, then fill in-place if archon_key not in buffer: diff --git a/areal/tests/experimental/archon/conftest.py b/areal/tests/experimental/archon/conftest.py new file mode 100644 index 0000000000..44fe2b8a4d --- /dev/null +++ b/areal/tests/experimental/archon/conftest.py @@ -0,0 +1,27 @@ +"""Pytest configuration for archon tests. + +NOTE: Archon engine requires PyTorch >= 2.9.1. +""" + +import pytest +import torch + +# Require PyTorch >= 2.9.1 for archon tests +_TORCH_VERSION = tuple(int(x) for x in torch.__version__.split("+")[0].split(".")[:3]) +_MIN_TORCH_VERSION = (2, 9, 1) + +if _TORCH_VERSION < _MIN_TORCH_VERSION: + collect_ignore_glob = ["test_*.py"] + + +def pytest_collection_modifyitems(config, items): + """Skip all archon tests if PyTorch version is too old.""" + if _TORCH_VERSION >= _MIN_TORCH_VERSION: + return + + skip_marker = pytest.mark.skip( + reason=f"Archon tests require PyTorch >= 2.9.1, but found {torch.__version__}" + ) + for item in items: + if "experimental/archon" in str(item.fspath): + item.add_marker(skip_marker) diff --git a/areal/tests/experimental/archon/test_activation_checkpoint.py b/areal/tests/experimental/archon/test_activation_checkpoint.py index 4ae5af9da8..8cc6c72c7b 100644 --- a/areal/tests/experimental/archon/test_activation_checkpoint.py +++ b/areal/tests/experimental/archon/test_activation_checkpoint.py @@ -186,8 +186,6 @@ def test_model_without_layers_raises(self): with pytest.raises(ValueError, match="must have a 'layers' attribute"): apply_ac(model, config) - # NOTE: Upgrading PyTorch will resolve this in the future. - @pytest.mark.slow def test_wrapped_model_forward_works(self): """Wrapped model should still produce correct forward output.""" model = DummyModel(num_layers=3, dim=4) @@ -206,8 +204,6 @@ def test_wrapped_model_forward_works(self): assert torch.allclose(ref_output, wrapped_output) - # NOTE: Upgrading PyTorch will resolve this in the future. - @pytest.mark.slow def test_wrapped_model_backward_works(self): """Wrapped model should support backward pass.""" model = DummyModel(num_layers=3, dim=4) diff --git a/areal/tests/experimental/archon/test_checkpoint_e2e.py b/areal/tests/experimental/archon/test_checkpoint_e2e.py new file mode 100644 index 0000000000..e6bdf59e94 --- /dev/null +++ b/areal/tests/experimental/archon/test_checkpoint_e2e.py @@ -0,0 +1,392 @@ +"""End-to-end checkpoint tests for State Dict Adapter functionality. + +Tests cover: +1. Multi-file safetensors loading (fqn_to_index_mapping) +2. ArchonEngine.save() and load() with HF format using DCP infrastructure + +These tests require: +- Real HF checkpoint files (Qwen3-30B-A3B MoE model) +- Multiple GPUs for distributed tests +- Sufficient GPU memory + +Run tests: + # Quick tests (single GPU, small model) + pytest areal/tests/experimental/archon/test_checkpoint_e2e.py -v -k "not moe" + + # Full MoE tests (multi-GPU, large model) + pytest areal/tests/experimental/archon/test_checkpoint_e2e.py -v -m multi_gpu +""" + +import json +import os +import subprocess + +import pytest +import torch + +from areal.platforms import current_platform +from areal.tests.experimental.archon.utils import ( + DENSE_MODEL_PATHS, + MOE_MODEL_PATHS, +) +from areal.utils.network import find_free_ports + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA not available" +) + + +# ============================================================================= +# Multi-file Safetensors Index Tests +# ============================================================================= + + +class TestMultiFileSafetensorsIndex: + """Tests for multi-file safetensors index loading.""" + + @pytest.fixture + def moe_model_path(self): + """Get path to Qwen3-30B-A3B MoE model.""" + path = MOE_MODEL_PATHS.get("qwen3_moe") + if path is None: + pytest.skip("Qwen3-30B-A3B model path not configured") + if not os.path.exists(path): + pytest.skip(f"Qwen3-30B-A3B model not found at {path}") + return path + + @pytest.fixture + def dense_model_path(self): + """Get path to Qwen3-0.6B dense model.""" + path = DENSE_MODEL_PATHS.get("qwen3") + if path is None: + pytest.skip("Qwen3-0.6B model path not configured") + if not os.path.exists(path): + pytest.skip(f"Qwen3-0.6B model not found at {path}") + return path + + def test_load_moe_safetensors_index(self, moe_model_path): + """Test loading multi-file safetensors index from real MoE checkpoint.""" + from transformers import AutoConfig + + from areal.experimental.models.archon.qwen3.model.state_dict_adapter import ( + Qwen3StateDictAdapter, + ) + + # Load config + config = AutoConfig.from_pretrained(moe_model_path, trust_remote_code=True) + + # Create adapter with hf_assets_path + adapter = Qwen3StateDictAdapter(config, hf_assets_path=moe_model_path) + + # Verify index was loaded + assert adapter.fqn_to_index_mapping is not None + assert len(adapter.fqn_to_index_mapping) > 0 + + # Verify index file structure + index_path = os.path.join(moe_model_path, "model.safetensors.index.json") + assert os.path.exists(index_path) + + with open(index_path) as f: + index_data = json.load(f) + + # Verify all keys in weight_map are in fqn_to_index_mapping + for hf_key in index_data["weight_map"]: + assert hf_key in adapter.fqn_to_index_mapping, ( + f"Missing key in fqn_to_index_mapping: {hf_key}" + ) + + # Print summary + print(f"\nLoaded index with {len(adapter.fqn_to_index_mapping)} weights") + print("Index maps to files 1-16 (16 safetensors files)") + + # Verify file indices are in expected range + indices = set(adapter.fqn_to_index_mapping.values()) + assert min(indices) == 1 + assert max(indices) == 16 # 16 files for Qwen3-30B-A3B + + def test_get_hf_storage_reader_moe(self, moe_model_path): + """Test creating HuggingFaceStorageReader for multi-file checkpoint.""" + from torch.distributed.checkpoint import HuggingFaceStorageReader + from transformers import AutoConfig + + from areal.experimental.models.archon.qwen3.model.state_dict_adapter import ( + Qwen3StateDictAdapter, + ) + + config = AutoConfig.from_pretrained(moe_model_path, trust_remote_code=True) + adapter = Qwen3StateDictAdapter(config, hf_assets_path=moe_model_path) + + # Get reader + reader = adapter.get_hf_storage_reader(moe_model_path) + assert isinstance(reader, HuggingFaceStorageReader) + + print(f"\nCreated HuggingFaceStorageReader for {moe_model_path}") + + def test_moe_config_parsing(self, moe_model_path): + """Test that MoE config is correctly parsed for expert count.""" + from transformers import AutoConfig + + from areal.experimental.models.archon.qwen3.model.state_dict_adapter import ( + Qwen3StateDictAdapter, + ) + + config = AutoConfig.from_pretrained(moe_model_path, trust_remote_code=True) + adapter = Qwen3StateDictAdapter(config, hf_assets_path=moe_model_path) + + # Qwen3-30B-A3B has 128 experts + assert adapter.moe_enabled is True + assert adapter.num_experts == 128 + + print(f"\nMoE config: {adapter.num_experts} experts") + + +# ============================================================================= +# Engine Integration Tests (Multi-GPU) +# ============================================================================= + + +def _run_engine_checkpoint_test( + n_gpus: int, test_type: str, model_path: str, output_file: str +): + """Run engine checkpoint test with torchrun.""" + port = find_free_ports(1)[0] + script_path = "areal/tests/experimental/archon/torchrun/run_checkpoint_tests.py" + + cmd = [ + "torchrun", + f"--nproc_per_node={n_gpus}", + "--nnodes=1", + "--master-addr=localhost", + f"--master_port={port}", + script_path, + f"--test_type={test_type}", + f"--model_path={model_path}", + f"--output={output_file}", + ] + + try: + subprocess.run(cmd, check=True, capture_output=True, text=True) + except subprocess.CalledProcessError as e: + pytest.fail(f"Test failed:\nSTDOUT: {e.stdout}\nSTDERR: {e.stderr}") + + # Verify result + with open(output_file) as f: + result = f.read().strip() + assert result == "Passed", f"Test failed: {result}" + + +@pytest.mark.multi_gpu +@pytest.mark.slow +class TestEngineCheckpointIntegration: + """End-to-end tests for ArchonEngine checkpoint methods. + + These tests require multiple GPUs and use real model checkpoints. + Tests use the unified save()/load() interface with HF format. + """ + + @pytest.fixture + def dense_model_path(self): + """Get path to Qwen3-0.6B dense model.""" + path = DENSE_MODEL_PATHS.get("qwen3") + if path is None: + pytest.skip("Qwen3-0.6B model path not configured") + if not os.path.exists(path): + pytest.skip(f"Qwen3-0.6B model not found at {path}") + return path + + @pytest.fixture + def moe_model_path(self): + """Get path to Qwen3-30B-A3B MoE model.""" + path = MOE_MODEL_PATHS.get("qwen3_moe") + if path is None: + pytest.skip("Qwen3-30B-A3B model path not configured") + if not os.path.exists(path): + pytest.skip(f"Qwen3-30B-A3B model not found at {path}") + return path + + def test_engine_save_hf_dense_2gpu(self, dense_model_path, tmp_path_factory): + """Test ArchonEngine.save() with weight_format="hf" on 2 GPUs.""" + if current_platform.device_count() < 2: + pytest.skip("This test requires 2 GPUs") + + output = tmp_path_factory.mktemp("test_output") / "result.out" + _run_engine_checkpoint_test( + n_gpus=2, + test_type="save_hf_dense", + model_path=dense_model_path, + output_file=str(output), + ) + + def test_engine_moe_checkpoint_4gpu(self, moe_model_path, tmp_path_factory): + """Test MoE checkpoint save/load on 4 GPUs with EP=4.""" + if current_platform.device_count() < 4: + pytest.skip("This test requires 4 GPUs") + + output = tmp_path_factory.mktemp("test_output") / "result.out" + _run_engine_checkpoint_test( + n_gpus=4, + test_type="moe_checkpoint", + model_path=moe_model_path, + output_file=str(output), + ) + + def test_engine_save_load_forward_match_2gpu( + self, dense_model_path, tmp_path_factory + ): + """Test forward output matches before save and after load (2 GPUs). + + This is the most critical test for checkpoint correctness: + 1. Initialize engine and run forward + 2. Save checkpoint + 3. Load checkpoint + 4. Run forward again + 5. Verify outputs match + """ + if current_platform.device_count() < 2: + pytest.skip("This test requires 2 GPUs") + + output = tmp_path_factory.mktemp("test_output") / "result.out" + _run_engine_checkpoint_test( + n_gpus=2, + test_type="save_load_forward_match", + model_path=dense_model_path, + output_file=str(output), + ) + + def test_engine_save_load_forward_match_1gpu( + self, dense_model_path, tmp_path_factory + ): + """Test forward output matches before save and after load (1 GPU). + + Single GPU version of forward match test for faster CI. + """ + if current_platform.device_count() < 1: + pytest.skip("This test requires at least 1 GPU") + + output = tmp_path_factory.mktemp("test_output") / "result.out" + _run_engine_checkpoint_test( + n_gpus=1, + test_type="save_load_forward_match", + model_path=dense_model_path, + output_file=str(output), + ) + + def test_engine_save_load_forward_match_with_compile_ac( + self, dense_model_path, tmp_path_factory + ): + """Test forward match with torch.compile + activation checkpointing. + + This test verifies that checkpoint save/load works correctly when the model + has wrapper prefixes in parameter names (e.g., _orig_mod from torch.compile, + _checkpoint_wrapped_module from activation checkpointing). + """ + if current_platform.device_count() < 1: + pytest.skip("This test requires at least 1 GPU") + + output = tmp_path_factory.mktemp("test_output") / "result.out" + _run_engine_checkpoint_test( + n_gpus=1, + test_type="save_load_forward_match_with_compile_ac", + model_path=dense_model_path, + output_file=str(output), + ) + + +# ============================================================================= +# Weight Comparison Tests +# ============================================================================= + + +class TestWeightComparisonAfterConversion: + """Tests that verify weight values are preserved after conversion.""" + + @pytest.fixture + def dense_model_path(self): + """Get path to Qwen3-0.6B dense model.""" + path = DENSE_MODEL_PATHS.get("qwen3") + if path is None: + pytest.skip("Qwen3-0.6B model path not configured") + if not os.path.exists(path): + pytest.skip(f"Qwen3-0.6B model not found at {path}") + return path + + @pytest.mark.slow + def test_adapter_roundtrip_preserves_weights(self, dense_model_path): + """Test that adapter roundtrip preserves exact weight values.""" + from safetensors.torch import load_file + from transformers import AutoConfig + + from areal.experimental.models.archon.qwen3.model.state_dict_adapter import ( + Qwen3StateDictAdapter, + ) + + config = AutoConfig.from_pretrained(dense_model_path, trust_remote_code=True) + adapter = Qwen3StateDictAdapter(config, hf_assets_path=dense_model_path) + + # Load original weights + safetensors_path = os.path.join(dense_model_path, "model.safetensors") + if os.path.exists(safetensors_path): + original_state = load_file(safetensors_path) + else: + # Multi-file checkpoint + index_path = os.path.join(dense_model_path, "model.safetensors.index.json") + with open(index_path) as f: + index = json.load(f) + + original_state = {} + loaded_files = set() + for key, filename in index["weight_map"].items(): + if filename not in loaded_files: + file_path = os.path.join(dense_model_path, filename) + file_state = load_file(file_path) + original_state.update(file_state) + loaded_files.add(filename) + + # Roundtrip: HF -> Archon -> HF + archon_state = adapter.from_hf(original_state) + roundtrip_state = adapter.to_hf(archon_state) + + # Check if weight tying is enabled + tie_word_embeddings = getattr(config, "tie_word_embeddings", False) + + # Compare weights + mismatches = [] + for key in original_state: + if "rotary_emb" in key: + continue # Skip rotary embeddings (not stored) + + # Skip lm_head.weight when tie_word_embeddings is True + # HF save_pretrained() doesn't save it either (reconstructed from embeddings) + if tie_word_embeddings and key == "lm_head.weight": + continue + + if key not in roundtrip_state: + mismatches.append(f"Missing key: {key}") + continue + + original = original_state[key] + roundtrip = roundtrip_state[key] + + if original.shape != roundtrip.shape: + mismatches.append( + f"Shape mismatch for {key}: {original.shape} vs {roundtrip.shape}" + ) + continue + + if not torch.allclose(original, roundtrip, rtol=1e-5, atol=1e-5): + max_diff = (original.float() - roundtrip.float()).abs().max().item() + mismatches.append(f"Value mismatch for {key}: max_diff={max_diff}") + + if mismatches: + for m in mismatches[:10]: # Show first 10 mismatches + print(m) + pytest.fail(f"Found {len(mismatches)} mismatches in roundtrip") + + skipped_msg = ( + " (lm_head.weight skipped due to tie_word_embeddings)" + if tie_word_embeddings + else "" + ) + print( + f"\nRoundtrip verified: {len(roundtrip_state)} weights preserved exactly{skipped_msg}" + ) diff --git a/areal/tests/experimental/archon/test_distributed_ep.py b/areal/tests/experimental/archon/test_distributed_ep.py index d0e6d78a9b..7330a5cbf3 100644 --- a/areal/tests/experimental/archon/test_distributed_ep.py +++ b/areal/tests/experimental/archon/test_distributed_ep.py @@ -165,3 +165,56 @@ def test_archon_ep_state_dict_update_2gpu(tmp_path_factory): pytest.skip("This test requires 2 GPUs") output = tmp_path_factory.mktemp("test_output") / "state_dict_update.out" _run_ep_test_with_torchrun(2, "state_dict_update", str(output)) + + +# ============================================================================= +# DTensor Checkpoint Roundtrip Tests +# ============================================================================= + + +@pytest.mark.multi_gpu +@pytest.mark.slow +def test_archon_ep_tp_dtensor_checkpoint_2gpu(tmp_path_factory): + """Test DTensor checkpoint roundtrip for EP+TP (ep=2, tp=2) on 2 GPUs. + + Tests MoEWeightConverter methods: + - split_expert_weights_dtensor(): 3D DTensor -> 2D DTensors + - concatenate_expert_weights_dtensor(): 2D DTensors -> 3D DTensor + + Verify: to_hf() -> from_hf() roundtrip preserves DTensor weights. + """ + if current_platform.device_count() < 2: + pytest.skip("This test requires 2 GPUs") + output = tmp_path_factory.mktemp("test_output") / "ep_tp_dtensor_checkpoint.out" + _run_ep_test_with_torchrun(2, "ep_tp_dtensor_checkpoint", str(output)) + + +@pytest.mark.multi_gpu +@pytest.mark.slow +def test_archon_ep_only_dtensor_checkpoint_2gpu(tmp_path_factory): + """Test DTensor checkpoint roundtrip for EP only (ep=2, tp=1) on 2 GPUs. + + Tests MoEWeightConverter methods with EP-only configuration. + + Verify: to_hf() -> from_hf() roundtrip preserves DTensor weights. + """ + if current_platform.device_count() < 2: + pytest.skip("This test requires 2 GPUs") + output = tmp_path_factory.mktemp("test_output") / "ep_only_dtensor_checkpoint.out" + _run_ep_test_with_torchrun(2, "ep_only_dtensor_checkpoint", str(output)) + + +@pytest.mark.multi_gpu +@pytest.mark.slow +def test_archon_etp_dtensor_checkpoint_4gpu(tmp_path_factory): + """Test DTensor checkpoint roundtrip for ETP (ep=2, tp=2, etp=2) on 4 GPUs. + + Tests MoEWeightConverter methods with ETP configuration + (StridedShard + Shard placement). + + Verify: to_hf() -> from_hf() roundtrip preserves DTensor weights. + """ + if current_platform.device_count() < 4: + pytest.skip("This test requires 4 GPUs") + output = tmp_path_factory.mktemp("test_output") / "etp_dtensor_checkpoint.out" + _run_ep_test_with_torchrun(4, "etp_dtensor_checkpoint", str(output)) diff --git a/areal/tests/experimental/archon/test_grouped_experts.py b/areal/tests/experimental/archon/test_grouped_experts.py index 51fca44c0f..5abbc938a5 100644 --- a/areal/tests/experimental/archon/test_grouped_experts.py +++ b/areal/tests/experimental/archon/test_grouped_experts.py @@ -13,10 +13,6 @@ _run_experts_for_loop, ) -# Marked as slow to exclude from CI pipeline. -# NOTE: Upgrading PyTorch will resolve this in the future. -pytestmark = pytest.mark.slow - class TestGroupedExpertsBasic: """Basic tests for GroupedExperts.""" diff --git a/areal/tests/experimental/archon/test_moe.py b/areal/tests/experimental/archon/test_moe.py index d304610b00..db8c713943 100644 --- a/areal/tests/experimental/archon/test_moe.py +++ b/areal/tests/experimental/archon/test_moe.py @@ -244,8 +244,6 @@ def test_expert_bias_buffer(self): class TestMoEGradients: """Tests for gradient flow.""" - # NOTE: Upgrading PyTorch will resolve this in the future. - @pytest.mark.slow def test_gradient_flow(self): """Test that gradients flow through MoE.""" dim = 32 @@ -265,8 +263,6 @@ def test_gradient_flow(self): assert moe.router.gate.weight.grad is not None assert moe.experts.w1.grad is not None - # NOTE: Upgrading PyTorch will resolve this in the future. - @pytest.mark.slow def test_gradient_all_experts(self): """Test that gradients reach all expert weights with enough tokens.""" dim = 32 diff --git a/areal/tests/experimental/archon/test_moe_utils.py b/areal/tests/experimental/archon/test_moe_utils.py index df2f17c09f..fc91c9406f 100644 --- a/areal/tests/experimental/archon/test_moe_utils.py +++ b/areal/tests/experimental/archon/test_moe_utils.py @@ -323,8 +323,6 @@ def test_full_moe_permutation_flow(self): assert merged.shape == tokens.shape - # NOTE: Upgrading PyTorch will resolve this in the future. - @pytest.mark.slow def test_gradient_flow(self): """Test that gradients flow correctly through permutation.""" num_tokens = 8 diff --git a/areal/tests/experimental/archon/test_qwen3_model_moe.py b/areal/tests/experimental/archon/test_qwen3_model_moe.py index a783ba2f60..4a15c8fcb4 100644 --- a/areal/tests/experimental/archon/test_qwen3_model_moe.py +++ b/areal/tests/experimental/archon/test_qwen3_model_moe.py @@ -277,8 +277,6 @@ def test_moe_block_forward(self, moe_args): @pytest.mark.skipif( not torch.cuda.is_available(), reason="CUDA required for MoE router histc" ) - # NOTE: Upgrading PyTorch will resolve this in the future. - @pytest.mark.slow def test_moe_block_gradient_flow(self, moe_args): """Test gradient flow through MoE TransformerBlock.""" block = TransformerBlock(layer_id=0, model_args=moe_args).cuda() @@ -478,8 +476,6 @@ def test_mixed_model_layer_structure(self, mixed_model_args): @pytest.mark.skipif( not torch.cuda.is_available(), reason="CUDA required for MoE router histc" ) - # NOTE: Upgrading PyTorch will resolve this in the future. - @pytest.mark.slow def test_moe_model_gradient_flow(self, moe_model_args): """Test gradient flow through MoE model.""" model = Qwen3Model(moe_model_args).cuda() @@ -516,8 +512,6 @@ def test_moe_model_gradient_flow(self, moe_model_args): @pytest.mark.skipif( not torch.cuda.is_available(), reason="CUDA required for MoE router histc" ) - # NOTE: Upgrading PyTorch will resolve this in the future. - @pytest.mark.slow def test_mixed_model_gradient_flow(self, mixed_model_args): """Test gradient flow through mixed MoE/dense model.""" model = Qwen3Model(mixed_model_args).cuda() diff --git a/areal/tests/experimental/archon/test_router.py b/areal/tests/experimental/archon/test_router.py index 53ea6a9936..108d8a4d72 100644 --- a/areal/tests/experimental/archon/test_router.py +++ b/areal/tests/experimental/archon/test_router.py @@ -353,8 +353,6 @@ def test_force_load_balance_top_k_2(self): class TestRouterGradients: """Tests for gradient flow through router.""" - # NOTE: Upgrading PyTorch will resolve this in the future. - @pytest.mark.slow def test_gradient_flow(self): """Test that gradients flow through router.""" dim = 64 diff --git a/areal/tests/experimental/archon/test_state_dict_adapter.py b/areal/tests/experimental/archon/test_state_dict_adapter.py index afcc390688..a210da8848 100644 --- a/areal/tests/experimental/archon/test_state_dict_adapter.py +++ b/areal/tests/experimental/archon/test_state_dict_adapter.py @@ -219,6 +219,41 @@ def test_convert_single_to_hf(self, adapter): assert result[0][0] == "model.layers.0.self_attn.q_proj.weight" assert torch.equal(result[0][1], tensor) + def test_qwen3_dense_no_moe_converter(self, adapter): + """Test that dense Qwen3 adapter does not initialize MoE converter.""" + assert adapter.moe_enabled is False + assert adapter._moe_converter is None + assert adapter._moe_state is None + + def test_qwen3_adapter_inherits_base_methods(self, adapter): + """Test that Qwen3StateDictAdapter inherits from BaseStateDictAdapter.""" + from areal.experimental.models.archon.base import BaseStateDictAdapter + + assert isinstance(adapter, BaseStateDictAdapter) + assert hasattr(adapter, "get_hf_storage_reader") + assert hasattr(adapter, "fqn_to_index_mapping") + + def test_qwen3_dense_with_weight_tying(self): + """Test Qwen3StateDictAdapter handles weight tying correctly for dense model.""" + + class MockQwen3DenseConfigTied: + model_type = "qwen3" + tie_word_embeddings = True + moe_enabled = False + + adapter = Qwen3StateDictAdapter(MockQwen3DenseConfigTied()) + + archon_state = { + "tok_embeddings.weight": torch.randn(1000, 64), + "output.weight": torch.randn(1000, 64), # Should be skipped + } + + hf_state = adapter.to_hf(archon_state) + + # output.weight should be skipped when weight tying is enabled + assert "lm_head.weight" not in hf_state + assert "model.embed_tokens.weight" in hf_state + class TestQwen3StateDictAdapterMoE: """Tests for Qwen3StateDictAdapter with MoE models.""" @@ -401,9 +436,6 @@ class MockConfig: return MockConfig() - # Marked as slow to exclude from CI pipeline. - # NOTE: Upgrading PyTorch will resolve this in the future. - @pytest.mark.slow @pytest.mark.skipif( not torch.cuda.is_available(), reason="MoE forward requires CUDA" ) @@ -749,3 +781,358 @@ def test_all_layer_indices_handled(self, model_configs): assert hf_key_back == hf_key, ( f"[{model_type}] Roundtrip failed for layer {layer_idx}" ) + + +# ============================================================================= +# Tests for HF Assets Path and Multi-File Checkpoint Support +# ============================================================================= + + +class TestHFAssetsPathSupport: + """Tests for hf_assets_path parameter and multi-file checkpoint support.""" + + @pytest.fixture + def mock_safetensors_index(self, tmp_path): + """Create a mock model.safetensors.index.json file.""" + index_data = { + "metadata": {"total_size": 12345678}, + "weight_map": { + "model.embed_tokens.weight": "model-00001-of-00019.safetensors", + "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00019.safetensors", + "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00019.safetensors", + "model.layers.0.self_attn.v_proj.weight": "model-00002-of-00019.safetensors", + "model.layers.0.self_attn.o_proj.weight": "model-00002-of-00019.safetensors", + "model.layers.0.mlp.gate_proj.weight": "model-00003-of-00019.safetensors", + "model.layers.0.mlp.up_proj.weight": "model-00003-of-00019.safetensors", + "model.layers.0.mlp.down_proj.weight": "model-00004-of-00019.safetensors", + "model.layers.10.self_attn.q_proj.weight": "model-00010-of-00019.safetensors", + "model.norm.weight": "model-00019-of-00019.safetensors", + "lm_head.weight": "model-00019-of-00019.safetensors", + }, + } + index_path = tmp_path / "model.safetensors.index.json" + import json + + with open(index_path, "w") as f: + json.dump(index_data, f) + return tmp_path + + def test_init_without_hf_assets_path(self): + """Test that adapter initializes correctly without hf_assets_path.""" + adapter = Qwen3StateDictAdapter(MockQwen3Config()) + + assert adapter.hf_assets_path is None + assert adapter.fqn_to_index_mapping is None + + def test_init_with_hf_assets_path(self, mock_safetensors_index): + """Test that adapter initializes correctly with hf_assets_path.""" + adapter = Qwen3StateDictAdapter( + MockQwen3Config(), hf_assets_path=str(mock_safetensors_index) + ) + + assert adapter.hf_assets_path == str(mock_safetensors_index) + assert adapter.fqn_to_index_mapping is not None + + def test_load_safetensors_index_mapping(self, mock_safetensors_index): + """Test that fqn_to_index_mapping is correctly populated.""" + adapter = Qwen3StateDictAdapter( + MockQwen3Config(), hf_assets_path=str(mock_safetensors_index) + ) + + mapping = adapter.fqn_to_index_mapping + assert mapping is not None + + # Verify index extraction from filenames + assert mapping["model.embed_tokens.weight"] == 1 + assert mapping["model.layers.0.self_attn.q_proj.weight"] == 1 + assert mapping["model.layers.0.self_attn.v_proj.weight"] == 2 + assert mapping["model.layers.0.mlp.gate_proj.weight"] == 3 + assert mapping["model.layers.0.mlp.down_proj.weight"] == 4 + assert mapping["model.layers.10.self_attn.q_proj.weight"] == 10 + assert mapping["model.norm.weight"] == 19 + assert mapping["lm_head.weight"] == 19 + + def test_load_safetensors_index_missing_file(self, tmp_path): + """Test that missing index file results in None mapping.""" + adapter = Qwen3StateDictAdapter(MockQwen3Config(), hf_assets_path=str(tmp_path)) + # When index file is missing, fqn_to_index_mapping should be None + assert adapter.fqn_to_index_mapping is None + + def test_get_hf_storage_reader(self, tmp_path): + """Test that get_hf_storage_reader returns correct type.""" + from torch.distributed.checkpoint import HuggingFaceStorageReader + + adapter = Qwen3StateDictAdapter(MockQwen3Config()) + reader = adapter.get_hf_storage_reader(str(tmp_path)) + + assert isinstance(reader, HuggingFaceStorageReader) + + def test_backward_compatibility_no_hf_assets_path(self): + """Test that existing code without hf_assets_path still works.""" + # This should work exactly as before + adapter = Qwen3StateDictAdapter(MockQwen3Config()) + + # All existing functionality should work + hf_key = "model.embed_tokens.weight" + archon_key = adapter._convert_key_from_hf(hf_key) + assert archon_key == "tok_embeddings.weight" + + hf_key_back = adapter._convert_key_to_hf(archon_key) + assert hf_key_back == hf_key + + def test_qwen2_adapter_with_hf_assets_path(self, mock_safetensors_index): + """Test that Qwen2StateDictAdapter also supports hf_assets_path.""" + from areal.experimental.models.archon.qwen2.model.state_dict_adapter import ( + Qwen2StateDictAdapter, + ) + + adapter = Qwen2StateDictAdapter( + MockQwen3Config(), hf_assets_path=str(mock_safetensors_index) + ) + + assert adapter.hf_assets_path == str(mock_safetensors_index) + assert adapter.fqn_to_index_mapping is not None + assert adapter.fqn_to_index_mapping["model.embed_tokens.weight"] == 1 + + +# ============================================================================= +# Tests for MoEWeightConverter Helper Methods +# ============================================================================= + + +class TestMoEWeightConverter: + """Tests for MoEWeightConverter methods.""" + + @pytest.fixture + def moe_converter(self): + from areal.experimental.models.archon.moe_weight_converter import ( + MoEWeightConverter, + ) + + return MoEWeightConverter + + def test_calculate_strided_shard_shard_indices_basic(self, moe_converter): + """Test calculate_strided_shard_indices with basic inputs.""" + # 8 experts split as [StridedShard(2), Shard(2)] + # Layout: + # StridedShard rank 0, Shard rank 0 -> experts 0,1 (block 0, position 0) + # StridedShard rank 1, Shard rank 0 -> experts 2,3 (block 1, position 0) + # StridedShard rank 0, Shard rank 1 -> experts 4,5 (block 0, position 1) + # StridedShard rank 1, Shard rank 1 -> experts 6,7 (block 1, position 1) + + # GPU (0,0): experts 0,1 + start, end = moe_converter.calculate_strided_shard_indices( + strided_shard_dim_degree=2, + strided_shard_dim_rank=0, + shard_dim_degree=2, + shard_dim_rank=0, + dim_size_to_split=8, + ) + assert start == 0 + assert end == 2 + + # GPU (1,0): experts 2,3 + start, end = moe_converter.calculate_strided_shard_indices( + strided_shard_dim_degree=2, + strided_shard_dim_rank=1, + shard_dim_degree=2, + shard_dim_rank=0, + dim_size_to_split=8, + ) + assert start == 2 + assert end == 4 + + # GPU (0,1): experts 4,5 + start, end = moe_converter.calculate_strided_shard_indices( + strided_shard_dim_degree=2, + strided_shard_dim_rank=0, + shard_dim_degree=2, + shard_dim_rank=1, + dim_size_to_split=8, + ) + assert start == 4 + assert end == 6 + + # GPU (1,1): experts 6,7 + start, end = moe_converter.calculate_strided_shard_indices( + strided_shard_dim_degree=2, + strided_shard_dim_rank=1, + shard_dim_degree=2, + shard_dim_rank=1, + dim_size_to_split=8, + ) + assert start == 6 + assert end == 8 + + def test_calculate_strided_shard_shard_indices_uneven_split_error( + self, moe_converter + ): + """Test that uneven split raises ValueError.""" + # 8 experts cannot be evenly split by strided_degree=3 and shard_degree=2 + # because 8 / (3 * 2) = 8/6 is not an integer + with pytest.raises(ValueError, match="Cannot evenly split"): + moe_converter.calculate_strided_shard_indices( + strided_shard_dim_degree=3, + strided_shard_dim_rank=0, + shard_dim_degree=2, + shard_dim_rank=0, + dim_size_to_split=8, + ) + + def test_calculate_strided_shard_shard_indices_single_shard(self, moe_converter): + """Test with shard_degree=1 (no additional sharding).""" + # 4 experts split only by StridedShard(4) + # GPU 0: experts 0 + # GPU 1: experts 1 + # GPU 2: experts 2 + # GPU 3: experts 3 + for rank in range(4): + start, end = moe_converter.calculate_strided_shard_indices( + strided_shard_dim_degree=4, + strided_shard_dim_rank=rank, + shard_dim_degree=1, + shard_dim_rank=0, + dim_size_to_split=4, + ) + assert start == rank + assert end == rank + 1 + + def test_calculate_strided_shard_shard_indices_edge_case_64_experts( + self, moe_converter + ): + """Test with 64 experts (common MoE configuration).""" + # 64 experts split as [StridedShard(4), Shard(4)] across 16 GPUs + # Block size = 64 / (4 * 4) = 4 experts per GPU + + # GPU (0,0): experts 0-3 + start, end = moe_converter.calculate_strided_shard_indices( + strided_shard_dim_degree=4, + strided_shard_dim_rank=0, + shard_dim_degree=4, + shard_dim_rank=0, + dim_size_to_split=64, + ) + assert start == 0 + assert end == 4 + + # GPU (3,3): experts 60-63 + start, end = moe_converter.calculate_strided_shard_indices( + strided_shard_dim_degree=4, + strided_shard_dim_rank=3, + shard_dim_degree=4, + shard_dim_rank=3, + dim_size_to_split=64, + ) + assert start == 60 + assert end == 64 + + +# ============================================================================= +# Tests for Qwen2StateDictAdapter Dense Model +# ============================================================================= + + +class TestQwen2StateDictAdapterDense: + """Tests for Qwen2StateDictAdapter with dense models.""" + + @pytest.fixture + def adapter(self): + from areal.experimental.models.archon.qwen2.model.state_dict_adapter import ( + Qwen2StateDictAdapter, + ) + + class MockQwen2Config: + model_type = "qwen2" + tie_word_embeddings = False + + return Qwen2StateDictAdapter(MockQwen2Config()) + + def test_to_hf_no_full_tensor_call(self, adapter): + """Test that to_hf() passes tensors through without modification.""" + # Create simple state dict + archon_state = { + "tok_embeddings.weight": torch.randn(1000, 64), + "layers.0.attention.wq.weight": torch.randn(64, 64), + "norm.weight": torch.randn(64), + "output.weight": torch.randn(1000, 64), + } + + hf_state = adapter.to_hf(archon_state) + + # Verify all keys are converted and values are identical (same tensor object) + assert "model.embed_tokens.weight" in hf_state + assert torch.equal( + hf_state["model.embed_tokens.weight"], archon_state["tok_embeddings.weight"] + ) + + def test_convert_single_to_hf_no_full_tensor_call(self, adapter): + """Test that convert_single_to_hf() returns tensor as-is.""" + name = "layers.0.attention.wq.weight" + tensor = torch.randn(64, 64) + + result = adapter.convert_single_to_hf(name, tensor) + + assert len(result) == 1 + hf_name, hf_tensor = result[0] + assert hf_name == "model.layers.0.self_attn.q_proj.weight" + # Should be the same tensor object (no copy or full_tensor) + assert hf_tensor is tensor + + def test_roundtrip_preserves_values(self, adapter): + """Test that roundtrip preserves all weight values.""" + archon_state = { + "tok_embeddings.weight": torch.randn(1000, 64), + "layers.0.attention.wq.weight": torch.randn(64, 64), + "layers.0.attention.wk.weight": torch.randn(16, 64), + "layers.0.attention.wv.weight": torch.randn(16, 64), + "layers.0.attention.wo.weight": torch.randn(64, 64), + "layers.0.feed_forward.w1.weight": torch.randn(128, 64), + "layers.0.feed_forward.w3.weight": torch.randn(128, 64), + "layers.0.feed_forward.w2.weight": torch.randn(64, 128), + "layers.0.attention_norm.weight": torch.randn(64), + "layers.0.ffn_norm.weight": torch.randn(64), + "norm.weight": torch.randn(64), + "output.weight": torch.randn(1000, 64), + } + + # Archon -> HF -> Archon + hf_state = adapter.to_hf(archon_state) + roundtrip_state = adapter.from_hf(hf_state) + + # Verify all values preserved + for key in archon_state: + assert key in roundtrip_state, f"Missing key: {key}" + assert torch.allclose(archon_state[key], roundtrip_state[key]), ( + f"Value mismatch for {key}" + ) + + def test_qwen2_adapter_inherits_base_methods(self, adapter): + """Test that Qwen2StateDictAdapter inherits from BaseStateDictAdapter.""" + from areal.experimental.models.archon.base import BaseStateDictAdapter + + assert isinstance(adapter, BaseStateDictAdapter) + assert hasattr(adapter, "get_hf_storage_reader") + assert hasattr(adapter, "fqn_to_index_mapping") + + def test_qwen2_adapter_with_weight_tying(self): + """Test Qwen2StateDictAdapter handles weight tying correctly.""" + from areal.experimental.models.archon.qwen2.model.state_dict_adapter import ( + Qwen2StateDictAdapter, + ) + + class MockQwen2ConfigTied: + model_type = "qwen2" + tie_word_embeddings = True + + adapter = Qwen2StateDictAdapter(MockQwen2ConfigTied()) + + archon_state = { + "tok_embeddings.weight": torch.randn(1000, 64), + "output.weight": torch.randn(1000, 64), # Should be skipped + } + + hf_state = adapter.to_hf(archon_state) + + # output.weight should be skipped when weight tying is enabled + assert "lm_head.weight" not in hf_state + assert "model.embed_tokens.weight" in hf_state diff --git a/areal/tests/experimental/archon/test_varlen_attention.py b/areal/tests/experimental/archon/test_varlen_attention.py index ca122fb376..916413bbed 100644 --- a/areal/tests/experimental/archon/test_varlen_attention.py +++ b/areal/tests/experimental/archon/test_varlen_attention.py @@ -70,8 +70,6 @@ def test_single_sequence(self): assert out.shape == q.shape assert not torch.isnan(out).any() - # NOTE: Upgrading PyTorch will resolve this in the future. - @pytest.mark.slow def test_gradient(self): """Test backward pass correctness.""" from areal.experimental.models.archon.varlen_attention import varlen_attn @@ -182,8 +180,6 @@ def test_wrapper_shape(self): assert out.shape == (batch, heads, seq_len, head_dim) assert not torch.isnan(out).any() - # NOTE: Upgrading PyTorch will resolve this in the future. - @pytest.mark.slow def test_wrapper_gradient(self): """Test wrapper backward pass.""" from areal.experimental.models.archon.varlen_attention import ( diff --git a/areal/tests/experimental/archon/torchrun/utils.py b/areal/tests/experimental/archon/torchrun/dist_utils.py similarity index 75% rename from areal/tests/experimental/archon/torchrun/utils.py rename to areal/tests/experimental/archon/torchrun/dist_utils.py index c5469efca6..802a7270bf 100644 --- a/areal/tests/experimental/archon/torchrun/utils.py +++ b/areal/tests/experimental/archon/torchrun/dist_utils.py @@ -10,7 +10,6 @@ def write_result(out: str, succ: bool) -> None: - """Write test result to file for pytest verification.""" with open(out, "w") as f: f.write("Passed" if succ else "Failed") @@ -28,7 +27,6 @@ def create_moe_model_args( max_seq_len: int = 8192, top_k: int = 2, ) -> Qwen3ModelArgs: - """Create model args for a small MoE model.""" return Qwen3ModelArgs( dim=dim, hidden_dim=hidden_dim, @@ -51,7 +49,7 @@ def create_moe_model_args( def gather_full_state_dict(model: torch.nn.Module) -> dict[str, torch.Tensor]: - """Gather all parameters including DTensors to full tensors.""" + """Gather all parameters (including DTensors) to full tensors.""" full_state = {} for name, param in model.named_parameters(): if isinstance(param, DTensor): @@ -75,14 +73,7 @@ def create_test_input( device: torch.device | str = "cuda", seed: int = 123, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: - """Create test input tensors. - - Returns: - tokens: (1, total_len) input token ids - positions: (1, total_len) position ids - cu_seqlens: (num_seqs + 1,) cumulative sequence lengths - seq_len_per_seq: sequence length per sequence - """ + """Return (tokens, positions, cu_seqlens, seq_len_per_seq).""" torch.manual_seed(seed) total_len = num_seqs * seq_len_per_seq @@ -102,7 +93,7 @@ def create_golden_model( device: torch.device | str, seed: int = 42, ) -> Qwen3Model: - """Create non-parallelized golden model for comparison.""" + """Create non-parallelized model for comparison.""" torch.manual_seed(seed) model = Qwen3Model(model_args) model.init_weights(device) @@ -117,24 +108,10 @@ def verify_outputs_match( atol: float = 1e-4, max_diff_threshold: float = 1e-2, ) -> tuple[bool, float, float]: - """Compare two outputs and return (success, max_diff, mean_diff). - - Args: - output1: First output tensor - output2: Second output tensor - rtol: Relative tolerance for allclose - atol: Absolute tolerance for allclose - max_diff_threshold: Maximum difference threshold for pass/fail - - Returns: - success: Whether outputs match within tolerances - max_diff: Maximum absolute difference - mean_diff: Mean absolute difference - """ + """Return (success, max_diff, mean_diff).""" max_diff = (output1 - output2).abs().max().item() mean_diff = (output1 - output2).abs().mean().item() - # Use both allclose and max_diff threshold allclose = torch.allclose(output1, output2, rtol=rtol, atol=atol) within_threshold = max_diff <= max_diff_threshold @@ -142,7 +119,6 @@ def verify_outputs_match( def print_rank0(msg: str) -> None: - """Print message only on rank 0.""" if dist.get_rank() == 0: print(msg) @@ -157,7 +133,6 @@ def create_dense_model_args( vocab_size: int = 1000, max_seq_len: int = 8192, ) -> Qwen3ModelArgs: - """Create model args for a small dense (non-MoE) model.""" return Qwen3ModelArgs( dim=dim, hidden_dim=hidden_dim, @@ -173,27 +148,16 @@ def create_dense_model_args( def validate_gradients(model: Qwen3Model) -> tuple[bool, list[str]]: - """Verify gradients flow through all key components. - - Args: - model: The model to check gradients for. - - Returns: - Tuple of (success, error_messages) - """ + """Return (success, error_messages).""" errors = [] - # Check embedding gradients if model.tok_embeddings.weight.grad is None: errors.append("tok_embeddings has no gradient") - # Check each layer for layer_id, layer in model.layers.items(): - # Attention if layer.attention.wq.weight.grad is None: errors.append(f"Layer {layer_id} attention.wq has no gradient") - # MoE or FFN if layer.moe is not None: if layer.moe.router.gate.weight.grad is None: errors.append(f"Layer {layer_id} MoE router has no gradient") @@ -203,7 +167,6 @@ def validate_gradients(model: Qwen3Model) -> tuple[bool, list[str]]: if layer.feed_forward.w1.weight.grad is None: errors.append(f"Layer {layer_id} FFN has no gradient") - # Check output layer if model.output is not None and model.output.weight.grad is None: errors.append("output projection has no gradient") @@ -211,14 +174,7 @@ def validate_gradients(model: Qwen3Model) -> tuple[bool, list[str]]: def validate_no_nan(output: torch.Tensor) -> bool: - """Check for NaN/Inf in output. - - Args: - output: Output tensor to check. - - Returns: - True if output is valid (no NaN/Inf), False otherwise. - """ + """Return True if output has no NaN/Inf.""" if torch.isnan(output).any(): return False if torch.isinf(output).any(): diff --git a/areal/tests/experimental/archon/torchrun/run_checkpoint_tests.py b/areal/tests/experimental/archon/torchrun/run_checkpoint_tests.py new file mode 100644 index 0000000000..77c1f97656 --- /dev/null +++ b/areal/tests/experimental/archon/torchrun/run_checkpoint_tests.py @@ -0,0 +1,659 @@ +#!/usr/bin/env python3 +"""Engine checkpoint integration tests. + +Tests for ArchonEngine save/load methods with HF format using DCP infrastructure. + +Run with: + torchrun --nproc_per_node=2 areal/tests/experimental/archon/torchrun/run_checkpoint_tests.py \ + --test_type=save_hf_dense --model_path=/path/to/model --output=/tmp/result.out + +Supported test types: + - save_hf_dense: Test save() with weight_format="hf" on dense model + - save_load_forward_match: Test save -> load -> forward output matches + - moe_checkpoint: Test MoE model checkpoint with EP +""" + +from __future__ import annotations + +import argparse +import os +import shutil +import tempfile + +import torch +import torch.distributed as dist + +from areal.api.alloc_mode import ParallelStrategy +from areal.api.cli_args import MicroBatchSpec, OptimizerConfig, TrainEngineConfig +from areal.api.io_struct import FinetuneSpec, SaveLoadMeta +from areal.experimental.engine.archon_engine import ArchonLMEngine +from areal.tests.experimental.archon.torchrun.dist_utils import ( + print_rank0, + write_result, +) + + +def create_engine_config(model_path: str) -> TrainEngineConfig: + """Create engine configuration for testing.""" + return TrainEngineConfig( + experiment_name="test_checkpoint", + trial_name="test", + path=model_path, + mb_spec=MicroBatchSpec(n_mbs=1), + optimizer=OptimizerConfig( + type="adam", + lr=1e-5, + weight_decay=0.01, + warmup_steps_proportion=0.0, + lr_scheduler_type="constant", + gradient_clipping=1.0, + ), + temperature=1.0, + ) + + +def test_save_hf_dense(model_path: str, output: str | None = None) -> bool: + """Test ArchonEngine.save() with weight_format="hf" on dense model. + + Steps: + 1. Initialize ArchonEngine with real HF checkpoint + 2. Call save(SaveLoadMeta(weight_format="hf")) to save in HF format + 3. Verify output contains safetensors files + 4. Verify config.json is copied + """ + rank = dist.get_rank() + world_size = dist.get_world_size() + + print_rank0("\n=== Test: save HF format (Dense) ===") + print_rank0(f"Model path: {model_path}") + print_rank0(f"World size: {world_size}") + + if rank == 0: + output_dir = tempfile.mkdtemp(prefix="hf_checkpoint_") + else: + output_dir = None + + output_dir_list = [output_dir] + dist.broadcast_object_list(output_dir_list, src=0) + output_dir = output_dir_list[0] + + success = True + + try: + config = create_engine_config(model_path) + engine = ArchonLMEngine(config) + + parallel_strategy = ParallelStrategy( + data_parallel_size=1, + tensor_parallel_size=world_size, + ) + ft_spec = FinetuneSpec(total_train_epochs=1, dataset_size=4, train_batch_size=4) + + engine.create_process_group(parallel_strategy=parallel_strategy) + engine.initialize(addr=None, ft_spec=ft_spec) + + print_rank0(f"Engine initialized with model: {engine.model.__class__.__name__}") + + # Save checkpoint using unified save() interface + save_meta = SaveLoadMeta( + path=output_dir, + weight_format="hf", + with_optim=False, + tokenizer=engine.tokenizer, + ) + engine.save(save_meta) + + dist.barrier() + + if rank == 0: + safetensors_files = list( + f for f in os.listdir(output_dir) if f.endswith(".safetensors") + ) + if len(safetensors_files) == 0: + print_rank0("ERROR: No safetensors files in output") + success = False + else: + print_rank0(f"Saved {len(safetensors_files)} safetensors files") + + config_path = os.path.join(output_dir, "config.json") + if not os.path.exists(config_path): + print_rank0("ERROR: config.json not found in output") + success = False + else: + print_rank0("config.json saved") + + engine.destroy() + + except Exception as e: + print_rank0(f"ERROR: {e}") + import traceback + + traceback.print_exc() + success = False + + finally: + dist.barrier() + if rank == 0 and output_dir and os.path.exists(output_dir): + shutil.rmtree(output_dir) + + if success: + print_rank0("save_hf_dense: PASSED") + else: + print_rank0("save_hf_dense: FAILED") + + if rank == 0 and output: + write_result(output, success) + + return success + + +def test_save_load_forward_match(model_path: str, output: str | None = None) -> bool: + """Test that forward output matches before save and after load. + + This is the most important test for checkpoint correctness: + 1. Initialize ArchonEngine with real HF checkpoint + 2. Run forward pass and record output + 3. Save checkpoint using save(weight_format="hf") + 4. Load checkpoint using load(weight_format="hf") + 5. Run forward pass again + 6. Verify outputs match exactly + + This test uses a small 0.6B model and runs on 1-2 GPUs. + """ + rank = dist.get_rank() + world_size = dist.get_world_size() + + print_rank0("\n=== Test: save_load_forward_match ===") + print_rank0(f"Model path: {model_path}") + print_rank0(f"World size: {world_size}") + + if rank == 0: + output_dir = tempfile.mkdtemp(prefix="forward_match_test_") + else: + output_dir = None + + output_dir_list = [output_dir] + dist.broadcast_object_list(output_dir_list, src=0) + output_dir = output_dir_list[0] + + success = True + + try: + config = create_engine_config(model_path) + engine = ArchonLMEngine(config) + + parallel_strategy = ParallelStrategy( + data_parallel_size=1, + tensor_parallel_size=world_size, + ) + ft_spec = FinetuneSpec(total_train_epochs=1, dataset_size=4, train_batch_size=4) + + engine.create_process_group(parallel_strategy=parallel_strategy) + engine.initialize(addr=None, ft_spec=ft_spec) + + print_rank0(f"Engine initialized with model: {engine.model.__class__.__name__}") + + torch.manual_seed(42) + batch_size = 4 + seq_len = 32 + vocab_size = engine.model_config.vocab_size + + input_ids = torch.randint( + 0, vocab_size, (1, batch_size * seq_len), device=engine.device + ) + positions = torch.arange(batch_size * seq_len, device=engine.device).unsqueeze( + 0 + ) + cu_seqlens = torch.tensor( + [i * seq_len for i in range(batch_size + 1)], + dtype=torch.int32, + device=engine.device, + ) + + engine.model.eval() + with torch.no_grad(): + output_before = engine.model( + input_ids, positions, cu_seqlens, max_seqlen=seq_len + ) + output_before = output_before.clone() + + print_rank0(f"Output before save: shape={output_before.shape}") + + save_meta = SaveLoadMeta( + path=output_dir, + weight_format="hf", + with_optim=False, + tokenizer=engine.tokenizer, + ) + engine.save(save_meta) + print_rank0(f"Checkpoint saved to {output_dir}") + + dist.barrier() + + load_meta = SaveLoadMeta( + path=output_dir, + weight_format="hf", + with_optim=False, + ) + engine.load(load_meta) + print_rank0("Checkpoint loaded") + + dist.barrier() + + engine.model.eval() + with torch.no_grad(): + output_after = engine.model( + input_ids, positions, cu_seqlens, max_seqlen=seq_len + ) + + print_rank0(f"Output after load: shape={output_after.shape}") + + max_diff = (output_before - output_after).abs().max().item() + mean_diff = (output_before - output_after).abs().mean().item() + allclose = torch.allclose(output_before, output_after, rtol=1e-4, atol=1e-4) + + print_rank0(f" Max diff: {max_diff:.6e}") + print_rank0(f" Mean diff: {mean_diff:.6e}") + print_rank0(f" Allclose (rtol=1e-4, atol=1e-4): {allclose}") + + if not allclose: + print_rank0("ERROR: Forward outputs do not match after load!") + # Allow small differences due to numerical precision + if max_diff > 1e-3: + success = False + print_rank0(f"ERROR: Max diff {max_diff} exceeds threshold 1e-3") + else: + print_rank0(f"WARNING: Max diff {max_diff} is small, treating as pass") + else: + print_rank0("Forward outputs match exactly!") + + engine.destroy() + + except Exception as e: + print_rank0(f"ERROR: {e}") + import traceback + + traceback.print_exc() + success = False + + finally: + dist.barrier() + if rank == 0 and output_dir and os.path.exists(output_dir): + shutil.rmtree(output_dir) + + if success: + print_rank0("save_load_forward_match: PASSED") + else: + print_rank0("save_load_forward_match: FAILED") + + if rank == 0 and output: + write_result(output, success) + + return success + + +def test_save_load_forward_match_with_compile_ac( + model_path: str, output: str | None = None +) -> bool: + """Test forward match with torch.compile and activation checkpointing enabled. + + This test verifies that checkpoint save/load works correctly when the model + has wrapper prefixes in parameter names (e.g., _orig_mod from torch.compile, + _checkpoint_wrapped_module from activation checkpointing). + + Steps: + 1. Initialize ArchonEngine with compile + AC enabled + 2. Run forward pass and record output + 3. Save checkpoint using save(weight_format="hf") + 4. Load checkpoint using load(weight_format="hf") + 5. Run forward pass again + 6. Verify outputs match exactly + """ + rank = dist.get_rank() + world_size = dist.get_world_size() + + print_rank0("\n=== Test: save_load_forward_match_with_compile_ac ===") + print_rank0(f"Model path: {model_path}") + print_rank0(f"World size: {world_size}") + + if rank == 0: + output_dir = tempfile.mkdtemp(prefix="forward_match_compile_ac_test_") + else: + output_dir = None + + output_dir_list = [output_dir] + dist.broadcast_object_list(output_dir_list, src=0) + output_dir = output_dir_list[0] + + success = True + + try: + config = create_engine_config(model_path) + # Enable compile and activation checkpointing + config.archon.enable_compile = True + config.gradient_checkpointing = True + config.archon.ac_mode = "full" + + engine = ArchonLMEngine(config) + + parallel_strategy = ParallelStrategy( + data_parallel_size=1, + tensor_parallel_size=world_size, + ) + ft_spec = FinetuneSpec(total_train_epochs=1, dataset_size=4, train_batch_size=4) + + engine.create_process_group(parallel_strategy=parallel_strategy) + engine.initialize(addr=None, ft_spec=ft_spec) + + print_rank0(f"Engine initialized with model: {engine.model.__class__.__name__}") + print_rank0(f" torch.compile enabled: {config.archon.enable_compile}") + print_rank0(f" Activation checkpointing: {config.archon.ac_mode}") + + # Verify wrapper prefixes exist in parameter names + param_names = [name for name, _ in engine.model.named_parameters()] + has_orig_mod = any("_orig_mod" in name for name in param_names) + has_checkpoint_wrapper = any( + "_checkpoint_wrapped_module" in name for name in param_names + ) + print_rank0(f" Has _orig_mod prefix: {has_orig_mod}") + print_rank0( + f" Has _checkpoint_wrapped_module prefix: {has_checkpoint_wrapper}" + ) + + torch.manual_seed(42) + batch_size = 4 + seq_len = 32 + vocab_size = engine.model_config.vocab_size + + input_ids = torch.randint( + 0, vocab_size, (1, batch_size * seq_len), device=engine.device + ) + positions = torch.arange(batch_size * seq_len, device=engine.device).unsqueeze( + 0 + ) + cu_seqlens = torch.tensor( + [i * seq_len for i in range(batch_size + 1)], + dtype=torch.int32, + device=engine.device, + ) + + engine.model.eval() + with torch.no_grad(): + output_before = engine.model( + input_ids, positions, cu_seqlens, max_seqlen=seq_len + ) + output_before = output_before.clone() + + print_rank0(f"Output before save: shape={output_before.shape}") + + save_meta = SaveLoadMeta( + path=output_dir, + weight_format="hf", + with_optim=False, + tokenizer=engine.tokenizer, + ) + engine.save(save_meta) + print_rank0(f"Checkpoint saved to {output_dir}") + + dist.barrier() + + load_meta = SaveLoadMeta( + path=output_dir, + weight_format="hf", + with_optim=False, + ) + engine.load(load_meta) + print_rank0("Checkpoint loaded") + + dist.barrier() + + engine.model.eval() + with torch.no_grad(): + output_after = engine.model( + input_ids, positions, cu_seqlens, max_seqlen=seq_len + ) + + print_rank0(f"Output after load: shape={output_after.shape}") + + max_diff = (output_before - output_after).abs().max().item() + mean_diff = (output_before - output_after).abs().mean().item() + allclose = torch.allclose(output_before, output_after, rtol=1e-4, atol=1e-4) + + print_rank0(f" Max diff: {max_diff:.6e}") + print_rank0(f" Mean diff: {mean_diff:.6e}") + print_rank0(f" Allclose (rtol=1e-4, atol=1e-4): {allclose}") + + if not allclose: + print_rank0("ERROR: Forward outputs do not match after load!") + # Allow small differences due to numerical precision + if max_diff > 1e-3: + success = False + print_rank0(f"ERROR: Max diff {max_diff} exceeds threshold 1e-3") + else: + print_rank0(f"WARNING: Max diff {max_diff} is small, treating as pass") + else: + print_rank0("Forward outputs match exactly!") + + engine.destroy() + + except Exception as e: + print_rank0(f"ERROR: {e}") + import traceback + + traceback.print_exc() + success = False + + finally: + dist.barrier() + if rank == 0 and output_dir and os.path.exists(output_dir): + shutil.rmtree(output_dir) + + if success: + print_rank0("save_load_forward_match_with_compile_ac: PASSED") + else: + print_rank0("save_load_forward_match_with_compile_ac: FAILED") + + if rank == 0 and output: + write_result(output, success) + + return success + + +def test_moe_checkpoint(model_path: str, output: str | None = None) -> bool: + """Test MoE model checkpoint with Expert Parallelism. + + This test requires 4 GPUs and tests: + 1. Loading MoE model with EP=4 + 2. Saving checkpoint without OOM (no full_tensor() in to_hf) + 3. Verifying saved checkpoint is valid + + Note: This test is memory-intensive due to MoE model size (30B params). + """ + rank = dist.get_rank() + world_size = dist.get_world_size() + + print_rank0("\n=== Test: moe_checkpoint ===") + print_rank0(f"Model path: {model_path}") + print_rank0(f"World size: {world_size}") + + if world_size < 4: + print_rank0("SKIP: MoE checkpoint test requires 4 GPUs") + if rank == 0 and output: + write_result(output, True) # Skip is not failure + return True + + if rank == 0: + output_dir = tempfile.mkdtemp(prefix="moe_checkpoint_") + else: + output_dir = None + + output_dir_list = [output_dir] + dist.broadcast_object_list(output_dir_list, src=0) + output_dir = output_dir_list[0] + + success = True + + try: + config = create_engine_config(model_path) + engine = ArchonLMEngine(config) + + parallel_strategy = ParallelStrategy( + data_parallel_size=1, + tensor_parallel_size=world_size, + ) + ft_spec = FinetuneSpec(total_train_epochs=1, dataset_size=4, train_batch_size=4) + + engine.create_process_group(parallel_strategy=parallel_strategy) + engine.initialize(addr=None, ft_spec=ft_spec) + + print_rank0("MoE Engine initialized") + print_rank0(f" Model: {engine.model.__class__.__name__}") + print_rank0( + f" Config experts: {getattr(engine.model_config, 'num_experts', 'N/A')}" + ) + + torch.cuda.synchronize() + mem_before = torch.cuda.memory_allocated() / 1024**3 + mem_before_max = torch.cuda.max_memory_allocated() / 1024**3 + print_rank0( + f" Memory before save: {mem_before:.2f} GB (max: {mem_before_max:.2f} GB)" + ) + + print_rank0("Saving MoE checkpoint...") + save_meta = SaveLoadMeta( + path=output_dir, + weight_format="hf", + with_optim=False, + tokenizer=engine.tokenizer, + ) + engine.save(save_meta) + + torch.cuda.synchronize() + mem_after = torch.cuda.memory_allocated() / 1024**3 + mem_after_max = torch.cuda.max_memory_allocated() / 1024**3 + print_rank0( + f" Memory after save: {mem_after:.2f} GB (max: {mem_after_max:.2f} GB)" + ) + + dist.barrier() + + if rank == 0: + safetensors_files = [ + f for f in os.listdir(output_dir) if f.endswith(".safetensors") + ] + if len(safetensors_files) == 0: + print_rank0("ERROR: No safetensors files in output") + success = False + else: + total_size = sum( + os.path.getsize(os.path.join(output_dir, f)) + for f in safetensors_files + ) + print_rank0( + f"Saved {len(safetensors_files)} files, " + f"total size: {total_size / 1024**3:.2f} GB" + ) + + engine.destroy() + + except torch.cuda.OutOfMemoryError as e: + print_rank0(f"ERROR: OOM during MoE checkpoint: {e}") + success = False + + except Exception as e: + print_rank0(f"ERROR: {e}") + import traceback + + traceback.print_exc() + success = False + + finally: + dist.barrier() + if rank == 0 and output_dir and os.path.exists(output_dir): + shutil.rmtree(output_dir) + + if success: + print_rank0("moe_checkpoint: PASSED") + else: + print_rank0("moe_checkpoint: FAILED") + + if rank == 0 and output: + write_result(output, success) + + return success + + +# ============================================================================= +# Main +# ============================================================================= + +TEST_REGISTRY = { + "save_hf_dense": test_save_hf_dense, + "save_load_forward_match": test_save_load_forward_match, + "save_load_forward_match_with_compile_ac": test_save_load_forward_match_with_compile_ac, + "moe_checkpoint": test_moe_checkpoint, +} + + +def main(): + parser = argparse.ArgumentParser(description="Engine Checkpoint Tests") + parser.add_argument( + "--test_type", + type=str, + required=True, + choices=list(TEST_REGISTRY.keys()), + help="Type of test to run", + ) + parser.add_argument( + "--model_path", + type=str, + required=True, + help="Path to HF model checkpoint", + ) + parser.add_argument( + "--output", + type=str, + default=None, + help="Output file for test result", + ) + args = parser.parse_args() + + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + torch.cuda.set_device(rank) + + print_rank0("=" * 60) + print_rank0(f"Running Checkpoint Test: {args.test_type}") + print_rank0("=" * 60) + + try: + test_fn = TEST_REGISTRY[args.test_type] + success = test_fn(args.model_path, args.output) + + dist.barrier() + + if success: + print_rank0(f"\n{'=' * 60}") + print_rank0(f"Checkpoint Test {args.test_type}: PASSED") + print_rank0("=" * 60) + else: + print_rank0(f"\n{'=' * 60}") + print_rank0(f"Checkpoint Test {args.test_type}: FAILED") + print_rank0("=" * 60) + + except Exception as e: + print(f"Rank {rank} failed with: {e}") + import traceback + + traceback.print_exc() + if rank == 0 and args.output: + write_result(args.output, False) + raise + + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/areal/tests/experimental/archon/torchrun/run_ep_tests.py b/areal/tests/experimental/archon/torchrun/run_ep_tests.py index e9d399d354..bd1168c91e 100644 --- a/areal/tests/experimental/archon/torchrun/run_ep_tests.py +++ b/areal/tests/experimental/archon/torchrun/run_ep_tests.py @@ -32,11 +32,14 @@ from areal.experimental.models.archon import ArchonParallelDims from areal.experimental.models.archon.qwen3.infra.parallelize import parallelize_qwen3 from areal.experimental.models.archon.qwen3.model.model import Qwen3Model +from areal.experimental.models.archon.qwen3.model.state_dict_adapter import ( + Qwen3StateDictAdapter, +) from areal.experimental.models.archon.ulysses import ( ulysses_gather_output, ulysses_slice_inputs, ) -from areal.tests.experimental.archon.torchrun.utils import ( +from areal.tests.experimental.archon.torchrun.dist_utils import ( create_golden_model, create_moe_model_args, create_test_input, @@ -52,7 +55,7 @@ def get_parallel_dims_ep_tp(world_size: int) -> ArchonParallelDims: - """Get parallel dims for EP+TP configuration (ep=ws, tp=ws).""" + """EP+TP config: ep=ws, tp=ws.""" return ArchonParallelDims( dp_shard=1, tp=world_size, @@ -64,11 +67,7 @@ def get_parallel_dims_ep_tp(world_size: int) -> ArchonParallelDims: def get_parallel_dims_ep_only(world_size: int) -> ArchonParallelDims: - """Get parallel dims for EP only configuration (ep=ws, tp=1). - - Note: dp_shard * tp * cp must equal world_size. - For EP only with 2 GPUs: dp_shard=2, tp=1, cp=1, ep=2 - """ + """EP only config: ep=ws, tp=1. dp_shard * tp * cp must equal world_size.""" return ArchonParallelDims( dp_shard=world_size, tp=1, @@ -80,42 +79,27 @@ def get_parallel_dims_ep_only(world_size: int) -> ArchonParallelDims: def get_parallel_dims_ep_cp(world_size: int) -> ArchonParallelDims: - """Get parallel dims for EP+CP configuration (ep=2, cp=2, 4 GPU). - - Note: dp_shard * tp * cp must equal world_size. - For EP+CP with 4 GPUs: dp_shard=2, tp=1, cp=2, ep=4 - """ + """EP+CP config: ep=4, cp=2, requires 4 GPUs.""" assert world_size == 4, "EP+CP requires 4 GPUs" return ArchonParallelDims( dp_shard=2, tp=1, cp=2, - ep=world_size, # ep=4 borrows from dp_shard * cp + ep=world_size, world_size=world_size, device_type="cuda", ) def get_parallel_dims_etp(world_size: int) -> ArchonParallelDims: - """Get parallel dims for ETP configuration. - - ETP (Expert Tensor Parallel) uses 2D sharding for expert weights. - EP borrows from dp_shard × cp only (not tp). - - Requires at least 4 GPUs. Extra GPUs go to dp_shard. - Config: tp=2, cp=1, ep=2, etp=2, dp_shard=world_size/4. - """ + """ETP config: tp=2, ep=2, etp=2, requires 4+ GPUs.""" assert world_size >= 4 and world_size % 4 == 0, ( "ETP requires at least 4 GPUs and world_size must be divisible by 4. " "Use tp_only_forward for 2 GPU TP-only MoE test." ) - - # ep borrows from dp_shard * cp = dp_shard (since cp=1) - # dp_shard_mod_ep = dp_shard * cp / ep = dp_shard / 2 - # dp_shard_in_ep = ep / cp = 2 dp_shard = world_size // 4 return ArchonParallelDims( - dp_shard=dp_shard * 2, # dp_shard must satisfy: dp_shard * tp * cp = world_size + dp_shard=dp_shard * 2, tp=2, cp=1, ep=2, @@ -126,10 +110,7 @@ def get_parallel_dims_etp(world_size: int) -> ArchonParallelDims: def get_parallel_dims_tp_only(world_size: int) -> ArchonParallelDims: - """Get parallel dims for TP-only MoE configuration (ep=1, tp=ws). - - This tests the TensorParallel class for MoE experts when EP is disabled. - """ + """TP-only MoE config: ep=1, tp=ws.""" return ArchonParallelDims( dp_shard=1, tp=world_size, @@ -147,7 +128,7 @@ def create_ep_model( device: torch.device, seed: int = 42, ) -> Qwen3Model: - """Create and parallelize a model with EP.""" + """Create and parallelize model with given parallel dims.""" torch.manual_seed(seed) model = Qwen3Model(model_args) model.init_weights(device) @@ -179,33 +160,18 @@ def test_forward( output: str | None = None, ep_model: Qwen3Model | None = None, ) -> tuple[bool, Qwen3Model]: - """Test forward numerical correctness for given parallel config. - - Verify: EP model output vs non-EP golden model output. - - Args: - parallel_dims: Pre-created ArchonParallelDims instance (reused across tests). - config_name: Name of the test configuration for logging. - output: Optional output file path for test result. - ep_model: Optional pre-created EP model to reuse. If None, creates a new one. - - Returns: - Tuple of (success, ep_model) where ep_model can be reused in subsequent tests. - """ + """Test forward correctness: EP model vs non-EP golden model. Returns (success, ep_model).""" rank = dist.get_rank() device = torch.device(f"cuda:{rank}") num_experts = 4 model_args = create_moe_model_args(num_experts=num_experts) - # Create golden model (non-parallelized) golden_model = create_golden_model(model_args, device, seed=42) - # Create EP model if not provided (reusing parallel_dims instance) if ep_model is None: ep_model = create_ep_model(model_args, parallel_dims, device, seed=42) - # Create test input (full sequence) tokens, positions, cu_seqlens, seq_len_per_seq = create_test_input( num_seqs=4, seq_len_per_seq=8, @@ -214,22 +180,16 @@ def test_forward( seed=123, ) - # Get CP info for input slicing and output gathering cp_group = parallel_dims.get_group("cp") cp_size = parallel_dims.cp cp_rank = parallel_dims.get_mesh("cp").get_local_rank() if cp_size > 1 else 0 - # Prepare inputs for EP model (slice if CP enabled) if cp_size > 1: - # Slice inputs for this CP rank inputs_dict = {"input_ids": tokens, "position_ids": positions} - # Create a dummy labels tensor for ulysses_slice_inputs labels = torch.zeros_like(tokens) sliced_inputs, _ = ulysses_slice_inputs(inputs_dict, labels, cp_rank, cp_size) tokens_ep = sliced_inputs["input_ids"] positions_ep = sliced_inputs["position_ids"] - # cu_seqlens should NOT be sliced - Ulysses gathers full sequence in attention - # so cu_seqlens refers to the full sequence positions cu_seqlens_ep = cu_seqlens seq_len_per_seq_ep = seq_len_per_seq else: @@ -238,17 +198,13 @@ def test_forward( cu_seqlens_ep = cu_seqlens seq_len_per_seq_ep = seq_len_per_seq - # Forward pass with torch.no_grad(): output_golden = golden_model(tokens, positions, cu_seqlens, seq_len_per_seq) output_ep = ep_model(tokens_ep, positions_ep, cu_seqlens_ep, seq_len_per_seq_ep) - # Gather EP output if CP enabled - # Model output has shape [1, seq_len/cp_size, dim], so use seq_dim=1 if cp_size > 1: output_ep = ulysses_gather_output(output_ep, cp_group, seq_dim=1) - # Verify outputs match success, max_diff, mean_diff = verify_outputs_match(output_golden, output_ep) print_rank0(f" {config_name} forward:") @@ -273,16 +229,7 @@ def test_output_consistency_across_ranks( config_name: str, ep_model: Qwen3Model | None = None, ) -> bool: - """Test that all ranks produce the same gathered output for the same input. - - For CP-enabled models, each CP rank processes different sequence slices, - so we compare the gathered (full) outputs across ranks. - - Args: - parallel_dims: Pre-created ArchonParallelDims instance (reused across tests). - config_name: Name of the test configuration for logging. - ep_model: Optional pre-created EP model to reuse. If None, creates a new one. - """ + """Test that all ranks produce the same gathered output.""" rank = dist.get_rank() world_size = dist.get_world_size() device = torch.device(f"cuda:{rank}") @@ -290,12 +237,10 @@ def test_output_consistency_across_ranks( num_experts = 4 model_args = create_moe_model_args(num_experts=num_experts) - # Reuse EP model if provided, otherwise create a new one if ep_model is None: ep_model = create_ep_model(model_args, parallel_dims, device, seed=42) model = ep_model - # Create test input (full sequence) tokens, positions, cu_seqlens, seq_len_per_seq = create_test_input( num_seqs=4, seq_len_per_seq=8, @@ -304,19 +249,16 @@ def test_output_consistency_across_ranks( seed=123, ) - # Get CP info for input slicing and output gathering cp_group = parallel_dims.get_group("cp") cp_size = parallel_dims.cp cp_rank = parallel_dims.get_mesh("cp").get_local_rank() if cp_size > 1 else 0 - # Prepare inputs for model (slice if CP enabled) if cp_size > 1: inputs_dict = {"input_ids": tokens, "position_ids": positions} labels = torch.zeros_like(tokens) sliced_inputs, _ = ulysses_slice_inputs(inputs_dict, labels, cp_rank, cp_size) tokens_model = sliced_inputs["input_ids"] positions_model = sliced_inputs["position_ids"] - # cu_seqlens should NOT be sliced - Ulysses gathers full sequence in attention cu_seqlens_model = cu_seqlens seq_len_per_seq_model = seq_len_per_seq else: @@ -325,25 +267,19 @@ def test_output_consistency_across_ranks( cu_seqlens_model = cu_seqlens seq_len_per_seq_model = seq_len_per_seq - # Synchronize all ranks before forward pass to ensure model state is consistent dist.barrier() - # Forward pass with torch.no_grad(): output = model( tokens_model, positions_model, cu_seqlens_model, seq_len_per_seq_model ) - # Gather output if CP enabled - # Model output has shape [1, seq_len/cp_size, dim], so use seq_dim=1 if cp_size > 1: output = ulysses_gather_output(output, cp_group, seq_dim=1) - # Gather outputs from all ranks output_list = [torch.zeros_like(output) for _ in range(world_size)] dist.all_gather(output_list, output) - # All outputs should be identical (after gathering CP outputs) success = True for i in range(1, world_size): if not torch.allclose(output_list[0], output_list[i], rtol=1e-5, atol=1e-5): @@ -372,18 +308,7 @@ def test_weight_sync( config_name: str, output: str | None = None, ) -> bool: - """Test weight gather, roundtrip, and cross-rank consistency. - - Verify: - 1. Gathered weights match original (before parallelization) - 2. Gathered weights are consistent across all ranks - 3. Loaded weights produce same output as original model - - Args: - parallel_dims: Pre-created ArchonParallelDims instance (reused across tests). - config_name: Name of the test configuration for logging. - output: Optional output file path for test result. - """ + """Test weight gather, roundtrip, and cross-rank consistency.""" rank = dist.get_rank() world_size = dist.get_world_size() device = torch.device(f"cuda:{rank}") @@ -391,14 +316,12 @@ def test_weight_sync( num_experts = 4 model_args = create_moe_model_args(num_experts=num_experts) - # Create model and save original weights before parallelization torch.manual_seed(42) model = Qwen3Model(model_args) model.init_weights(device) model = model.to(device) original_state = {k: v.clone() for k, v in model.state_dict().items()} - # Apply parallelization (reusing parallel_dims instance) parallelize_qwen3( model=model, parallel_dims=parallel_dims, @@ -411,12 +334,10 @@ def test_weight_sync( enable_compile=False, ) - # Gather sharded weights gathered_state = gather_full_state_dict(model) success = True - # Test 1: Verify expert weights match original print_rank0(f" {config_name} weight sync:") for name in original_state: if "experts" in name and name in gathered_state: @@ -435,7 +356,6 @@ def test_weight_sync( print_rank0(f" Weight mismatch for {name}: max_diff={max_diff}") success = False - # Test 2: Verify cross-rank consistency expert_weight_names = [ name for name in gathered_state.keys() if "experts.w1" in name ] @@ -453,13 +373,11 @@ def test_weight_sync( ) success = False - # Test 3: Verify roundtrip (load gathered weights to new model) torch.manual_seed(42) model_new = Qwen3Model(model_args) model_new.init_weights(device) model_new = model_new.to(device) - # Load gathered state loadable_state = {} model_new_state = model_new.state_dict() for name, tensor in gathered_state.items(): @@ -467,7 +385,6 @@ def test_weight_sync( loadable_state[name] = tensor model_new.load_state_dict(loadable_state, strict=False) - # Compare forward outputs tokens, positions, cu_seqlens, seq_len_per_seq = create_test_input( num_seqs=4, seq_len_per_seq=8, @@ -476,7 +393,6 @@ def test_weight_sync( seed=456, ) - # Create reference model golden_model = create_golden_model(model_args, device, seed=42) with torch.no_grad(): @@ -507,15 +423,7 @@ def test_weight_sync( def test_state_dict_update(output: str | None = None) -> bool: - """Test state dict correctness after optimizer step. - - Verify: - 1. Create EP model + optimizer - 2. Forward + backward + optimizer.step() - 3. Gather EP weights to full state_dict - 4. Create non-EP model, load gathered state_dict - 5. Two models with same input should produce same output - """ + """Test state dict correctness after optimizer step.""" rank = dist.get_rank() world_size = dist.get_world_size() device = torch.device(f"cuda:{rank}") @@ -523,14 +431,11 @@ def test_state_dict_update(output: str | None = None) -> bool: num_experts = 4 model_args = create_moe_model_args(num_experts=num_experts) - # Create EP model parallel_dims = get_parallel_dims_ep_tp(world_size) model = create_ep_model(model_args, parallel_dims, device, seed=42) - # Create optimizer optimizer = torch.optim.SGD(model.parameters(), lr=0.01) - # Training step tokens, positions, cu_seqlens, seq_len_per_seq = create_test_input( num_seqs=4, seq_len_per_seq=8, @@ -544,13 +449,9 @@ def test_state_dict_update(output: str | None = None) -> bool: loss.backward() optimizer.step() - # Gather updated weights gathered_state = gather_full_state_dict(model) - # Create new non-EP model and load gathered weights - torch.manual_seed( - 99 - ) # Different seed to ensure weights are loaded, not initialized + torch.manual_seed(99) model_new = Qwen3Model(model_args) model_new.init_weights(device) model_new = model_new.to(device) @@ -562,7 +463,6 @@ def test_state_dict_update(output: str | None = None) -> bool: loadable_state[name] = tensor model_new.load_state_dict(loadable_state, strict=False) - # Forward with new input tokens2, positions2, cu_seqlens2, seq_len_per_seq2 = create_test_input( num_seqs=4, seq_len_per_seq=8, @@ -572,12 +472,9 @@ def test_state_dict_update(output: str | None = None) -> bool: ) with torch.no_grad(): - # EP model forward output_ep = model(tokens2, positions2, cu_seqlens2, seq_len_per_seq2) - # Non-EP model forward output_new = model_new(tokens2, positions2, cu_seqlens2, seq_len_per_seq2) - # Verify outputs match success, max_diff, mean_diff = verify_outputs_match( output_ep, output_new, rtol=1e-4, atol=1e-4 ) @@ -709,6 +606,145 @@ def run_tp_only_forward(output: str | None = None) -> bool: return success +# ============================================================================= +# DTensor Checkpoint Roundtrip Tests +# ============================================================================= + + +def test_dtensor_checkpoint_roundtrip( + parallel_dims: ArchonParallelDims, + config_name: str, + output: str | None = None, +) -> bool: + """Test DTensor checkpoint roundtrip: to_hf() -> from_hf() preserves weights.""" + from torch.distributed.tensor import DTensor + + rank = dist.get_rank() + world_size = dist.get_world_size() + device = torch.device(f"cuda:{rank}") + + num_experts = 4 + model_args = create_moe_model_args(num_experts=num_experts) + + model = create_ep_model(model_args, parallel_dims, device, seed=42) + + original_dtensor_state = {} + for name, param in model.named_parameters(): + original_dtensor_state[name] = param + + adapter = Qwen3StateDictAdapter(model_args) + + success = True + print_rank0(f" {config_name} DTensor checkpoint roundtrip:") + + hf_state_dict = adapter.to_hf(original_dtensor_state) + + moe_2d_count = 0 + for key, value in hf_state_dict.items(): + if ".mlp.experts." in key: + moe_2d_count += 1 + if value.dim() != 2: + print_rank0(f" ERROR: {key} should be 2D, got {value.dim()}D") + success = False + + print_rank0(f" to_hf produced {moe_2d_count} 2D expert weights") + + reconstructed_state = adapter.from_hf(hf_state_dict) + + for name in original_dtensor_state: + if "moe.experts.w" in name: + original = original_dtensor_state[name] + if name not in reconstructed_state: + print_rank0(f" ERROR: {name} missing in reconstructed state") + success = False + continue + + reconstructed = reconstructed_state[name] + + if isinstance(original, DTensor): + original_local = original._local_tensor + else: + original_local = original + + if isinstance(reconstructed, DTensor): + reconstructed_local = reconstructed._local_tensor + else: + reconstructed_local = reconstructed + + if original_local.shape != reconstructed_local.shape: + print_rank0( + f" ERROR: Shape mismatch for {name}: " + f"{original_local.shape} vs {reconstructed_local.shape}" + ) + success = False + continue + + if not torch.allclose( + original_local, reconstructed_local, rtol=1e-5, atol=1e-5 + ): + max_diff = (original_local - reconstructed_local).abs().max().item() + print_rank0( + f" ERROR: Value mismatch for {name}: max_diff={max_diff}" + ) + success = False + + expert_keys = [k for k in reconstructed_state.keys() if "moe.experts.w1" in k] + for key in expert_keys[:1]: + value = reconstructed_state[key] + if isinstance(value, DTensor): + full_tensor = value.full_tensor() + else: + full_tensor = value + + tensor_list = [torch.zeros_like(full_tensor) for _ in range(world_size)] + dist.all_gather(tensor_list, full_tensor) + + for i in range(1, world_size): + if not torch.allclose(tensor_list[0], tensor_list[i], rtol=1e-5, atol=1e-5): + max_diff = (tensor_list[0] - tensor_list[i]).abs().max().item() + print_rank0( + f" ERROR: Cross-rank mismatch for {key} " + f"(rank 0 vs {i}): max_diff={max_diff}" + ) + success = False + + if success: + print_rank0(f" {config_name}_dtensor_checkpoint: PASSED") + else: + print_rank0(f" {config_name}_dtensor_checkpoint: FAILED") + + dist.barrier() + + if rank == 0 and output: + write_result(output, success) + + return success + + +def run_ep_tp_dtensor_checkpoint(output: str | None = None) -> bool: + """Run EP+TP DTensor checkpoint roundtrip test.""" + print_rank0("\n=== EP+TP DTensor Checkpoint Roundtrip Test ===") + world_size = dist.get_world_size() + parallel_dims = get_parallel_dims_ep_tp(world_size) + return test_dtensor_checkpoint_roundtrip(parallel_dims, "ep_tp", output) + + +def run_ep_only_dtensor_checkpoint(output: str | None = None) -> bool: + """Run EP only DTensor checkpoint roundtrip test.""" + print_rank0("\n=== EP Only DTensor Checkpoint Roundtrip Test ===") + world_size = dist.get_world_size() + parallel_dims = get_parallel_dims_ep_only(world_size) + return test_dtensor_checkpoint_roundtrip(parallel_dims, "ep_only", output) + + +def run_etp_dtensor_checkpoint(output: str | None = None) -> bool: + """Run ETP DTensor checkpoint roundtrip test (requires 4 GPUs).""" + print_rank0("\n=== ETP DTensor Checkpoint Roundtrip Test ===") + world_size = dist.get_world_size() + parallel_dims = get_parallel_dims_etp(world_size) + return test_dtensor_checkpoint_roundtrip(parallel_dims, "etp", output) + + # ============================================================================= # Main # ============================================================================= @@ -724,6 +760,10 @@ def run_tp_only_forward(output: str | None = None) -> bool: "etp_weight_sync": run_etp_weight_sync, "tp_only_forward": run_tp_only_forward, "state_dict_update": run_state_dict_update, + # DTensor checkpoint roundtrip tests + "ep_tp_dtensor_checkpoint": run_ep_tp_dtensor_checkpoint, + "ep_only_dtensor_checkpoint": run_ep_only_dtensor_checkpoint, + "etp_dtensor_checkpoint": run_etp_dtensor_checkpoint, } diff --git a/areal/tests/experimental/archon/torchrun/run_forward.py b/areal/tests/experimental/archon/torchrun/run_forward.py index 2e9fa2222b..d393a4962d 100644 --- a/areal/tests/experimental/archon/torchrun/run_forward.py +++ b/areal/tests/experimental/archon/torchrun/run_forward.py @@ -93,7 +93,7 @@ def make_engine(model_type: str, mb_spec: MicroBatchSpec): engine = ArchonEngine(config) - # Create parallel strategy (Phase 1: only dp_size) + # Create parallel strategy world_size = dist.get_world_size() if dist.is_initialized() else 1 parallel_strategy = ParallelStrategy(data_parallel_size=world_size) diff --git a/areal/tests/experimental/archon/torchrun/run_parallel_dims.py b/areal/tests/experimental/archon/torchrun/run_parallel_dims.py index d591626285..8956519af5 100644 --- a/areal/tests/experimental/archon/torchrun/run_parallel_dims.py +++ b/areal/tests/experimental/archon/torchrun/run_parallel_dims.py @@ -20,7 +20,10 @@ import torch.distributed as dist from areal.experimental.models.archon import ArchonParallelDims -from areal.tests.experimental.archon.torchrun.utils import print_rank0, write_result +from areal.tests.experimental.archon.torchrun.dist_utils import ( + print_rank0, + write_result, +) def test_ep_mesh_when_etp_disabled(output: str | None = None) -> bool: diff --git a/areal/tests/experimental/archon/torchrun/run_qwen3_parallelize.py b/areal/tests/experimental/archon/torchrun/run_qwen3_parallelize.py index a41ea4e71e..f3f90702c4 100644 --- a/areal/tests/experimental/archon/torchrun/run_qwen3_parallelize.py +++ b/areal/tests/experimental/archon/torchrun/run_qwen3_parallelize.py @@ -35,7 +35,7 @@ from areal.experimental.models.archon import ArchonParallelDims from areal.experimental.models.archon.qwen3.infra.parallelize import parallelize_qwen3 from areal.experimental.models.archon.qwen3.model.model import Qwen3Model -from areal.tests.experimental.archon.torchrun.utils import ( +from areal.tests.experimental.archon.torchrun.dist_utils import ( create_dense_model_args, create_moe_model_args, create_test_input, diff --git a/areal/tests/experimental/archon/torchrun/run_vs_fsdp.py b/areal/tests/experimental/archon/torchrun/run_vs_fsdp.py index 8019ea89b3..b103734642 100644 --- a/areal/tests/experimental/archon/torchrun/run_vs_fsdp.py +++ b/areal/tests/experimental/archon/torchrun/run_vs_fsdp.py @@ -20,7 +20,7 @@ from areal.engine.fsdp_engine import FSDPLMEngine from areal.experimental.engine.archon_engine import ArchonLMEngine from areal.platforms import current_platform -from areal.tests.experimental.archon.torchrun.utils import write_result +from areal.tests.experimental.archon.torchrun.dist_utils import write_result from areal.tests.utils import get_dataset_path, get_model_path from areal.utils.data import pad_sequences_to_tensors from areal.utils.hf_utils import load_hf_tokenizer