diff --git a/training/expert_parallel/README.md b/training/expert_parallel/README.md new file mode 100644 index 000000000..65d6f8ced --- /dev/null +++ b/training/expert_parallel/README.md @@ -0,0 +1,92 @@ +# AutoEP Training Example + +AutoEP (Auto Expert Parallelism) automatically partitions MoE expert weights across GPUs and uses AllToAll communication to route tokens to the correct experts. +This example offers a quick start for AutoEP in DeepSpeed. + +## Quick Start + +### Prerequisites + +- 2+ GPUs (the fast grouped GEMM path works only on Hopper and Blackwell GPUs) +- Dependencies: + - PyTorch `>= 2.9.1` + - **DeepSpeed** with AutoEP: AutoEP has not been merged into `main` yet. `requirements.txt` installs the **tip of [PR #7938](https://github.com/deepspeedai/DeepSpeed/pull/7938)**. Manual install: `pip install "git+https://github.com/deepspeedai/DeepSpeed.git@refs/pull/7938/head#egg=deepspeed"`. + - **`transformers` `>= 5.2`** (see `requirements.txt`) + - Qwen3.5 requires specific kernel dependencies. See [VERIFICATION.md](VERIFICATION.md#qwen35-kernel-requirements) for more details. + - See `requirements.txt` for other dependencies. + +### Run + +The following launches causal LM training with **AutoEP + ZeRO-1** on a randomly initialized model built from the original **Qwen3.5-MoE** Hugging Face text config. +`--num_layers` overrides only the layer count in the original model config, which is useful when testing with limited GPU resources. `--dataset_name` and `--dataset_percentage` choose the Hugging Face training dataset and the percentage of the train split to use. + +```bash +deepspeed --num_gpus 8 train.py \ + --mode autoep \ + --model qwen3_5_moe \ + --autoep_size 8 \ + --num_layers 8 \ + --dataset_name wikitext \ + --dataset_percentage 10.0 \ + --steps 1000 +``` + +For this `--mode autoep` / `--model qwen3_5_moe` run, `train.py` derives the following `expert_parallel` section: + +```json + "expert_parallel": { + "enabled": true, + "autoep_size": 8, + "preset_model": "qwen3_5_moe" + } +``` + +Here are the key options in the DeepSpeed config for AutoEP: + +- **`enabled`** — Turns on AutoEP. +- **`autoep_size`** — Expert-parallel size. It must be specified with `--autoep_size` in AutoEP mode and must divide both the GPU count and the model's expert count. The benchmark commands use `8` for Qwen3.5 and `4` for Llama4 and Mixtral. +- **`preset_model`** — DeepSpeed's structural AutoEP preset id. The example's public `--model` choices intentionally use the same ids when an AutoEP preset exists. + +This example exposes three public `--model` choices: `qwen3_5_moe`, `llama4`, and `mixtral`. These match the DeepSpeed `preset_model` ids used by AutoEP for the same structures. The underlying AutoEP PR also defines additional structural preset ids: `qwen3_moe`, `deepseek_v2`, and `deepseek_v3`; those are not exposed as `--model` choices in this example. + +## Performance Benchmark + +We benchmarked Qwen3.5 and Mixtral with 8 layers, and Llama4 with 7 layers. Each model was run under matching conditions for AutoEP and the ZeRO-3 leaf baseline: 8 H100 GPUs, sequence length 1024, micro batch size 1, gradient accumulation 4, 100 optimizer steps, and steps 50-99 measured. The table below reports the side-by-side comparison. Llama4 uses 7 layers because the 8-layer ZeRO-3 leaf baseline OOMed during backward. + +| Model | ZeRO-3 leaf | AutoEP (+ZeRO-1) | +| --- | --- | --- | +| Qwen3.5 MoE | 42,128.05 tok/s, 34.99 GB | 87,540.15 tok/s, 25.58 GB (`2.08x` throughput, `0.73x` memory vs ZeRO-3) | +| Llama4 (7 layers) | 19,144.07 tok/s, 56.95 GB | 60,178.91 tok/s, 60.08 GB (`3.14x` throughput, `1.06x` memory vs 7-layer ZeRO-3) | +| Mixtral 8x7B | 32,622.11 tok/s, 50.47 GB | 69,052.31 tok/s, 35.03 GB (`2.12x` throughput, `0.69x` memory vs ZeRO-3) | + +Qwen3.5 reproduction steps and reference loss curves for the AutoEP and ZeRO-3 leaf comparison are in [VERIFICATION.md](VERIFICATION.md). + + +## Important Constraints + +### `autoep_size` requirements + +- Must be `<= num_experts` +- Must evenly divide `num_experts` +- Must evenly divide `world_size` +- `autoep_size=1` bypasses EP communication entirely (degenerate case) + +### Grouped GEMM backend + +`torch._grouped_mm` is required for the default production path. With the default `expert_parallel.use_grouped_mm=true`, DeepSpeed fails fast if `torch._grouped_mm` is unavailable. Set `use_grouped_mm=false` in the DeepSpeed config only for functional/debug runs that intentionally use the sequential for-loop path. On A100 (SM80), verify availability and actual throughput since the Hopper fast path may not activate. + +### Qwen3.5 linear-attention kernels + +For `--model qwen3_5_moe`, the required linear-attention kernel dependencies and verification checks are documented in [VERIFICATION.md](VERIFICATION.md#qwen35-kernel-requirements). + +### bf16 requirement + +`bf16` is recommended. `fp16` is functionally correct but not optimized for the Hopper grouped-GEMM fast path used by `torch._grouped_mm`. + +### Optimizer wiring + +AutoEP runs must let DeepSpeed build the optimizer from the JSON config (no client optimizer). This ensures `configure_moe_param_groups()` is invoked to split expert parameters into expert-data-parallel reduction groups. + +### Load balancing status + +DeepSeek-style auxiliary-loss-free (expert-bias) load balancing is **not yet implemented** in AutoEP. diff --git a/training/expert_parallel/VERIFICATION.md b/training/expert_parallel/VERIFICATION.md new file mode 100644 index 000000000..1eada3df9 --- /dev/null +++ b/training/expert_parallel/VERIFICATION.md @@ -0,0 +1,125 @@ +# Qwen3.5 AutoEP Reproduction + +This document describes how to reproduce the Qwen3.5 AutoEP sample and compare it with the ZeRO-3 leaf baseline. It is intentionally environment-neutral: choose local output directories and caches that fit your machine or cluster. + +The commands below use: + +- model preset: `qwen3_5_moe` +- DeepSpeed AutoEP preset: `qwen3_5_moe` +- layers: `8` +- dataset: `wikitext`, `dataset_percentage=10.0` +- tokenizer: `Qwen/Qwen3-0.6B` +- sequence length: `1024` +- micro batch size: `1` +- gradient accumulation: `4` +- world size: `8` +- steps: `100`, with steps `50-99` treated as the post-warmup measurement window + +AutoEP uses the sample's built-in `--mode autoep` config with `--autoep_size 8`: bf16, AdamW, ZeRO stage 1, `expert_parallel.enabled=true`, `autoep_size=8`, `preset_model=qwen3_5_moe`. The baseline uses `--mode zero3_leaf`: bf16, AdamW, ZeRO stage 3, and the Qwen3.5 MoE block registered as a ZeRO leaf module. + +## Install + +Run from the repository root: + +```bash +cd training/expert_parallel +python -m pip install -r requirements.txt +export TOKENIZERS_PARALLELISM=false +mkdir -p runs/qwen35/{init,autoep,zero3_leaf,compare} +``` + +You may set `HF_HOME` and `HF_DATASETS_CACHE` if your environment requires explicit Hugging Face cache locations. + + +## Qwen3.5 Kernel Requirements + +For `--model qwen3_5_moe`, `flash-linear-attention`, `causal-conv1d`, `flash-attn`, and `tilelang` on H100/Triton `>= 3.4` are verification requirements, not optional accelerators. The verification should fail if `transformers.utils.import_utils.is_flash_linear_attention_available()` or `is_causal_conv1d_available()` is false, or if a runtime inspection shows `Qwen3_5MoeGatedDeltaNet` using `torch_causal_conv1d_update` or `torch_chunk_gated_delta_rule`. + +`flash-attn` is also required when full-attention layers are configured to use `attn_implementation="flash_attention_2"`. + +## Create A Shared Initialization + +Use the same randomly initialized weights for AutoEP and ZeRO-3 leaf when comparing loss curves. This makes the metric comparison easier to interpret. + +```bash +python utils/prepare_init_weights.py \ + --model qwen3_5_moe \ + --num_layers 8 \ + --seed 42 \ + --output runs/qwen35/init/qwen35_l8_seed42.safetensors +``` + +## Run AutoEP + +```bash +deepspeed --num_gpus 8 --master_port 29104 train.py \ + --mode autoep \ + --model qwen3_5_moe \ + --autoep_size 8 \ + --num_layers 8 \ + --steps 100 \ + --warmup_steps 50 \ + --log_interval 1 \ + --seq_len 1024 \ + --micro_batch_size 1 \ + --grad_accum 4 \ + --seed 42 \ + --dataset_name wikitext \ + --dataset_percentage 10.0 \ + --tokenizer_name Qwen/Qwen3-0.6B \ + --load_init_weights runs/qwen35/init/qwen35_l8_seed42.safetensors \ + --metrics_out runs/qwen35/autoep/metrics.csv +``` + +## Run ZeRO-3 Leaf Baseline + +```bash +deepspeed --num_gpus 8 --master_port 29105 train.py \ + --mode zero3_leaf \ + --model qwen3_5_moe \ + --num_layers 8 \ + --steps 100 \ + --warmup_steps 50 \ + --log_interval 1 \ + --seq_len 1024 \ + --micro_batch_size 1 \ + --grad_accum 4 \ + --seed 42 \ + --dataset_name wikitext \ + --dataset_percentage 10.0 \ + --tokenizer_name Qwen/Qwen3-0.6B \ + --load_init_weights runs/qwen35/init/qwen35_l8_seed42.safetensors \ + --metrics_out runs/qwen35/zero3_leaf/metrics.csv +``` + +## Compare Metrics + +`compare_metrics.py` compares the loss, throughput, and peak memory reported by the two metrics CSV files, then generates summary JSON plus plots. + +```bash +python utils/compare_metrics.py \ + --autoep_csv runs/qwen35/autoep/metrics.csv \ + --zero3_leaf_csv runs/qwen35/zero3_leaf/metrics.csv \ + --warmup_steps 50 \ + --out_dir runs/qwen35/compare \ + --out_json runs/qwen35/compare/summary.json +``` + +Useful outputs: + +- `runs/qwen35/compare/summary.json` +- `runs/qwen35/compare/loss_curve.png` +- `runs/qwen35/compare/peak_memory_bar.png` +- `runs/qwen35/compare/throughput_bar.png` + +Small numeric differences are expected because AutoEP and ZeRO-3 leaf use different distributed execution paths. The comparison is intended to confirm that loss behavior remains aligned while reporting throughput and memory differences under the same model, data, seed, and initialization. + +## Reference Loss Curves + +The repository includes reference Qwen3.5 loss-curve images from a longer AutoEP and ZeRO-3 leaf comparison. They are useful for checking the expected curve shape after reproducing the workflow above, but the commands above are the source of truth for a fresh run in your own environment. + +![Qwen3.5 CE loss curve](images/qwen35_aux_10k_ce_loss_curve.png) + +![Qwen3.5 total loss curve](images/qwen35_aux_10k_total_loss_curve.png) + +![Qwen3.5 aux loss curve](images/qwen35_aux_10k_aux_loss_curve.png) diff --git a/training/expert_parallel/data_utils.py b/training/expert_parallel/data_utils.py new file mode 100644 index 000000000..d50ee9eaa --- /dev/null +++ b/training/expert_parallel/data_utils.py @@ -0,0 +1,303 @@ +"""Hugging Face text data loading and MoE config helpers for the AutoEP example.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Iterator + +import torch +from datasets import DownloadConfig, load_dataset +from datasets.utils.logging import disable_progress_bar +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler +from transformers import AutoTokenizer + + +def build_model_config( + config_cls: type[Any], + num_hidden_layers: int | None = None, +) -> Any: + """Build a Hugging Face config with an optional layer-count override.""" + kwargs: dict[str, Any] = {} + if num_hidden_layers is not None: + kwargs["num_hidden_layers"] = num_hidden_layers + return config_cls(**kwargs) + + +@dataclass +class CausalLmBatch: + """One micro-batch for causal LM training (CPU tensors until the training loop moves them).""" + + input_ids: torch.Tensor # [micro_batch_size, seq_len], dtype=torch.long + attention_mask: torch.Tensor # [micro_batch_size, seq_len], dtype=torch.long + labels: torch.Tensor # [micro_batch_size, seq_len], dtype=torch.long + + +def get_tokenizer(model_name: str, *, trust_remote_code: bool = True) -> Any: + """Load tokenizer; pad token follows ds_verify_loss behavior.""" + tokenizer = AutoTokenizer.from_pretrained( + model_name, trust_remote_code=trust_remote_code + ) + if tokenizer.pad_token is None: + if tokenizer.eos_token is not None: + tokenizer.pad_token = tokenizer.eos_token + else: + tokenizer.pad_token = tokenizer.convert_ids_to_tokens(2) + return tokenizer + + +def validate_tokenizer_vocab_size( + tokenizer: Any, + tokenizer_name: str, + expected_vocab_size: int, +) -> dict[str, Any]: + """Validate that tokenizer ids fit inside the model embedding table.""" + tokenizer_len = len(tokenizer) + tokenizer_vocab_size = getattr(tokenizer, "vocab_size", None) + if tokenizer_len > expected_vocab_size: + raise ValueError( + f"Tokenizer {tokenizer_name!r} len(tokenizer)={tokenizer_len} " + f"(vocab_size={tokenizer_vocab_size}) exceeds model " + f"vocab_size={expected_vocab_size}. " + "Pick a tokenizer whose ids fit within the model config." + ) + return { + "tokenizer_name": tokenizer_name, + "tokenizer_len": tokenizer_len, + "tokenizer_vocab_size": tokenizer_vocab_size, + "model_vocab_size": expected_vocab_size, + "exact_vocab_match": tokenizer_len == expected_vocab_size, + } + + +def _hf_train_split_string(dataset_fraction: float) -> str: + """Map fraction in (0, 1] to datasets split slice (matches ds_verify_loss).""" + if dataset_fraction >= 1.0: + return "train" + percentage_int = int(dataset_fraction * 100) + return f"train[:{percentage_int}%]" + + +def load_hf_text_dataset_rows( + dataset_name: str, + dataset_fraction: float, + *, + is_main_process: bool, +) -> tuple[Any, str]: + """Load raw text rows from Hugging Face (same presets as ds_verify_loss).""" + if not is_main_process: + disable_progress_bar() + + split_str = _hf_train_split_string(dataset_fraction) + dl_cfg = DownloadConfig(disable_tqdm=True) + + if is_main_process: + print(f"Loading HF dataset: {dataset_name} split={split_str!r} ...") + + if dataset_name == "wikitext": + dataset = load_dataset( + "wikitext", + "wikitext-103-raw-v1", + split=split_str, + download_config=dl_cfg, + ) + text_column = "text" + elif dataset_name == "openwebtext": + if dataset_fraction >= 1.0: + split_str = "train[:1%]" + dataset = load_dataset( + "openwebtext", split=split_str, download_config=dl_cfg + ) + text_column = "text" + elif dataset_name == "c4": + if dataset_fraction >= 1.0: + split_str = "train[:0.1%]" + dataset = load_dataset( + "c4", "en", split=split_str, download_config=dl_cfg + ) + text_column = "text" + elif dataset_name == "ag_news": + dataset = load_dataset( + "ag_news", split=split_str, download_config=dl_cfg + ) + text_column = "text" + else: + try: + dataset = load_dataset( + dataset_name, split=split_str, download_config=dl_cfg + ) + if "text" in dataset.column_names: + text_column = "text" + elif "content" in dataset.column_names: + text_column = "content" + elif "body" in dataset.column_names: + text_column = "body" + else: + text_column = dataset.column_names[0] + if is_main_process: + print( + f"Warning: using column {text_column!r}; " + f"columns={dataset.column_names}" + ) + except Exception as e: + if is_main_process: + print(f"Error loading {dataset_name!r}: {e}; falling back to wikitext.") + dataset = load_dataset( + "wikitext", + "wikitext-103-raw-v1", + split=split_str, + download_config=dl_cfg, + ) + text_column = "text" + + if is_main_process: + print(f"HF dataset rows: {len(dataset)} (text column={text_column!r})") + return dataset, text_column + + +def tokenize_hf_dataset( + dataset: Any, + text_column: str, + tokenizer: Any, + seq_len: int, + *, + is_main_process: bool, +) -> Any: + """Tokenize text column to fixed length (padding=max_length), torch columns.""" + + def has_text(example: dict[str, Any]) -> bool: + value = example.get(text_column) + return isinstance(value, str) and bool(value.strip()) + + def tokenize_fn(examples: dict[str, list]) -> dict[str, list]: + return tokenizer( + examples[text_column], + padding="max_length", + max_length=seq_len, + truncation=True, + ) + + if is_main_process: + print("Filtering empty HF text rows...") + dataset = dataset.filter( + has_text, + num_proc=1, + keep_in_memory=True, + ) + if len(dataset) == 0: + raise ValueError( + "HF dataset has no non-empty text rows; pick another dataset." + ) + if is_main_process: + print(f"Non-empty HF dataset rows: {len(dataset)}.") + print("Tokenizing HF dataset...") + tokenized = dataset.map( + tokenize_fn, + batched=True, + num_proc=1, + remove_columns=dataset.column_names, + keep_in_memory=True, + ) + tokenized.set_format( + type="torch", columns=["input_ids", "attention_mask"] + ) + if is_main_process: + print(f"Tokenization complete: {len(tokenized)} rows.") + if len(tokenized) == 0: + raise ValueError( + "Tokenized HF dataset is empty; increase dataset_percentage or pick another dataset." + ) + return tokenized + + +class HFBatchGenerator: + """Infinite iterator over a DataLoader; returns CausalLmBatch on CPU.""" + + def __init__( + self, + dataloader: DataLoader, + sampler: DistributedSampler, + ) -> None: + self.dataloader = dataloader + self.sampler = sampler + self._epoch = 0 + self._iter: Iterator | None = None + + def _next_raw_batch(self) -> dict[str, torch.Tensor]: + if self._iter is None: + self.sampler.set_epoch(self._epoch) + self._iter = iter(self.dataloader) + try: + return next(self._iter) + except StopIteration: + self._epoch += 1 + self.sampler.set_epoch(self._epoch) + self._iter = iter(self.dataloader) + return next(self._iter) + + def get_batch(self, optimizer_step: int, accum_idx: int) -> CausalLmBatch: + del optimizer_step, accum_idx # sequential consumption (like an infinite epoch) + batch = self._next_raw_batch() + input_ids = batch["input_ids"] + attention_mask = batch["attention_mask"] + labels = input_ids.clone() + labels = labels.masked_fill(attention_mask == 0, -100) + return CausalLmBatch( + input_ids=input_ids, + attention_mask=attention_mask, + labels=labels, + ) + + +def build_hf_batch_generator( + *, + dataset_name: str, + dataset_percentage: float, + tokenizer_name: str, + expected_vocab_size: int, + seq_len: int, + micro_batch_size: int, + dp_world_size: int, + dp_rank: int, + seed: int, + rank: int, + hf_num_dataloader_workers: int = 0, +) -> HFBatchGenerator: + """Load HF text data, tokenize, and build a per-DP-rank batch generator.""" + is_main = rank == 0 + if dataset_percentage <= 0: + raise ValueError("dataset_percentage must be positive") + if dataset_percentage < 1.0: + raise ValueError( + "dataset_percentage must be at least 1.0 because Hugging Face " + "split slicing uses whole percentages." + ) + fraction = min(dataset_percentage / 100.0, 1.0) + + tokenizer = get_tokenizer(tokenizer_name, trust_remote_code=True) + validate_tokenizer_vocab_size(tokenizer, tokenizer_name, expected_vocab_size) + raw, text_col = load_hf_text_dataset_rows( + dataset_name, fraction, is_main_process=is_main + ) + tokenized = tokenize_hf_dataset( + raw, + text_col, + tokenizer, + seq_len, + is_main_process=is_main, + ) + sampler = DistributedSampler( + tokenized, + num_replicas=dp_world_size, + rank=dp_rank, + shuffle=True, + seed=seed, + ) + loader = DataLoader( + tokenized, + batch_size=micro_batch_size, + sampler=sampler, + num_workers=hf_num_dataloader_workers, + pin_memory=torch.cuda.is_available(), + ) + return HFBatchGenerator(loader, sampler) diff --git a/training/expert_parallel/images/loss_curve.png b/training/expert_parallel/images/loss_curve.png new file mode 100644 index 000000000..ab88ec7c0 Binary files /dev/null and b/training/expert_parallel/images/loss_curve.png differ diff --git a/training/expert_parallel/images/peak_memory_bar.png b/training/expert_parallel/images/peak_memory_bar.png new file mode 100644 index 000000000..2e33148f4 Binary files /dev/null and b/training/expert_parallel/images/peak_memory_bar.png differ diff --git a/training/expert_parallel/images/qwen35_aux_10k_aux_loss_curve.png b/training/expert_parallel/images/qwen35_aux_10k_aux_loss_curve.png new file mode 100644 index 000000000..7d9a139bd Binary files /dev/null and b/training/expert_parallel/images/qwen35_aux_10k_aux_loss_curve.png differ diff --git a/training/expert_parallel/images/qwen35_aux_10k_ce_loss_curve.png b/training/expert_parallel/images/qwen35_aux_10k_ce_loss_curve.png new file mode 100644 index 000000000..cfc019bf8 Binary files /dev/null and b/training/expert_parallel/images/qwen35_aux_10k_ce_loss_curve.png differ diff --git a/training/expert_parallel/images/qwen35_aux_10k_total_loss_curve.png b/training/expert_parallel/images/qwen35_aux_10k_total_loss_curve.png new file mode 100644 index 000000000..9462a69f3 Binary files /dev/null and b/training/expert_parallel/images/qwen35_aux_10k_total_loss_curve.png differ diff --git a/training/expert_parallel/images/throughput_bar.png b/training/expert_parallel/images/throughput_bar.png new file mode 100644 index 000000000..a87b0162b Binary files /dev/null and b/training/expert_parallel/images/throughput_bar.png differ diff --git a/training/expert_parallel/init_weights.py b/training/expert_parallel/init_weights.py new file mode 100644 index 000000000..08e04879f --- /dev/null +++ b/training/expert_parallel/init_weights.py @@ -0,0 +1,198 @@ +"""Init-weights artifact helpers for AutoEP parity workflows.""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.metadata +import json +import os +from datetime import datetime, timezone +from typing import Any + +import torch +from safetensors.torch import load_file, save_file + +CURRENT_INIT_SCHEMA_VERSION = 1 + + +def sha256_file(path: str) -> str: + """Return SHA-256 hex digest of file bytes.""" + h = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(1024 * 1024), b""): + h.update(chunk) + return h.hexdigest() + + +def build_config_fingerprint(model_config: object) -> str: + """Build a stable fingerprint from shape-defining model config fields.""" + num_experts = getattr(model_config, "num_local_experts", None) + if num_experts is None: + num_experts = getattr(model_config, "num_experts", None) + intermediate = getattr(model_config, "intermediate_size", None) + if intermediate is None: + intermediate = getattr(model_config, "moe_intermediate_size", None) + fields = { + "num_hidden_layers": getattr(model_config, "num_hidden_layers", None), + "hidden_size": getattr(model_config, "hidden_size"), + "intermediate_size": intermediate, + "num_attention_heads": getattr(model_config, "num_attention_heads"), + "num_key_value_heads": getattr(model_config, "num_key_value_heads"), + "num_experts": num_experts, + "num_experts_per_tok": getattr(model_config, "num_experts_per_tok"), + "vocab_size": getattr(model_config, "vocab_size"), + "max_position_embeddings": getattr(model_config, "max_position_embeddings"), + } + payload = json.dumps(fields, sort_keys=True).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _meta_path(weights_path: str) -> str: + stem, _ = os.path.splitext(weights_path) + return f"{stem}_meta.json" + + +def _ensure_safetensors_path(path: str) -> None: + if not path.endswith(".safetensors"): + raise ValueError( + f"Init weights path must end with '.safetensors': {path}" + ) + + +def _fsync_dir(path: str) -> None: + dir_fd = os.open(os.path.dirname(path) or ".", os.O_DIRECTORY) + try: + os.fsync(dir_fd) + finally: + os.close(dir_fd) + + +def _write_json_atomic(path: str, payload: dict[str, Any]) -> None: + tmp = path + ".tmp" + with open(tmp, "w") as f: + json.dump(payload, f, indent=2) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp, path) + _fsync_dir(path) + + +def _model_state_dict_cpu_fp32(model: torch.nn.Module) -> dict[str, torch.Tensor]: + state_dict = model.state_dict() + converted: dict[str, torch.Tensor] = {} + for key in sorted(state_dict.keys()): + tensor = state_dict[key].detach().cpu().contiguous() + if torch.is_floating_point(tensor): + tensor = tensor.to(torch.float32) + converted[key] = tensor + return converted + + +def _transformers_version() -> str: + try: + return importlib.metadata.version("transformers") + except Exception: + return "unknown" + + +def save_init_weights_artifact( + path: str, + model: torch.nn.Module, + *, + args: argparse.Namespace, + model_config: object, + rank: int, +) -> dict[str, object]: + """Save init weights safetensors + sidecar metadata and return context.""" + _ensure_safetensors_path(path) + + out_path = os.path.abspath(path) + meta_path = _meta_path(out_path) + os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True) + + state_dict = _model_state_dict_cpu_fp32(model) + tmp_weights = out_path + ".tmp" + save_file(state_dict, tmp_weights) + with open(tmp_weights, "rb") as f: + os.fsync(f.fileno()) + os.replace(tmp_weights, out_path) + _fsync_dir(out_path) + + weights_sha256 = sha256_file(out_path) + config_fingerprint = build_config_fingerprint(model_config) + metadata = { + "schema_version": CURRENT_INIT_SCHEMA_VERSION, + "seed": int(getattr(args, "seed")), + "num_layers": int(getattr(args, "num_layers")), + "model_class": model.__class__.__name__, + "config_fingerprint": config_fingerprint, + "torch_version": torch.__version__, + "transformers_version": _transformers_version(), + "created_at": datetime.now(timezone.utc).isoformat(), + "saved_by_rank": int(rank), + } + _write_json_atomic(meta_path, metadata) + + return { + "init_weights_path": out_path, + "init_weights_sha256": weights_sha256, + "init_weights_loaded": False, + "init_weights_schema_version": CURRENT_INIT_SCHEMA_VERSION, + } + + +def load_init_weights_artifact( + path: str, + model: torch.nn.Module, + *, + args: argparse.Namespace, + model_config: object, +) -> dict[str, object]: + """Load and validate init weights artifact, then return metadata context.""" + _ensure_safetensors_path(path) + in_path = os.path.abspath(path) + meta_path = _meta_path(in_path) + + if not os.path.isfile(in_path): + raise FileNotFoundError(f"Init weights file does not exist: {in_path}") + if not os.path.isfile(meta_path): + raise FileNotFoundError(f"Init weights metadata file does not exist: {meta_path}") + + with open(meta_path) as f: + metadata = json.load(f) + + schema_version = int(metadata.get("schema_version", 0)) + if schema_version < 1: + raise ValueError(f"Invalid init-weights schema version: {schema_version}") + if schema_version > CURRENT_INIT_SCHEMA_VERSION: + raise ValueError( + "Schema version " + f"{schema_version} is not supported. Please upgrade your tools." + ) + + expected_layers = int(getattr(args, "num_layers")) + artifact_layers = int(metadata.get("num_layers", -1)) + if artifact_layers != expected_layers: + raise ValueError( + "Init weights num_layers mismatch: " + f"artifact={artifact_layers}, run={expected_layers}" + ) + + expected_fingerprint = build_config_fingerprint(model_config) + artifact_fingerprint = metadata.get("config_fingerprint") + if artifact_fingerprint != expected_fingerprint: + raise ValueError( + "Init weights config fingerprint mismatch: " + "artifact does not match current model configuration." + ) + + state_dict = load_file(in_path, device="cpu") + model.load_state_dict(state_dict, strict=True) + + return { + "init_weights_path": in_path, + "init_weights_sha256": sha256_file(in_path), + "init_weights_loaded": True, + "init_weights_schema_version": schema_version, + } diff --git a/training/expert_parallel/metrics.py b/training/expert_parallel/metrics.py new file mode 100644 index 000000000..58dbeb96c --- /dev/null +++ b/training/expert_parallel/metrics.py @@ -0,0 +1,68 @@ +"""Minimal CSV metrics and distributed reduction helpers for the AutoEP example.""" + +from __future__ import annotations + +import csv +import os +from typing import Any + +import torch +import torch.distributed as dist + +METRICS_COLUMNS: list[str] = [ + "step", + "loss", + "iter_time_sec", + "global_tokens_per_sec", + "cuda_memory_allocated_bytes", + "cuda_peak_memory_allocated_bytes", + "cuda_peak_memory_reserved_bytes", +] + + +class MetricsLogger: + """Write per-step metrics to CSV on rank 0.""" + + def __init__(self, csv_path: str, rank: int) -> None: + self.csv_path = csv_path + self.rank = rank + self._file = None + self._writer = None + + def log_step(self, metrics: dict[str, Any]) -> None: + if self.rank != 0: + return + if self._file is None: + parent = os.path.dirname(self.csv_path) + if parent: + os.makedirs(parent, exist_ok=True) + self._file = open(self.csv_path, "w", newline="") + self._writer = csv.DictWriter(self._file, fieldnames=METRICS_COLUMNS) + self._writer.writeheader() + self._writer.writerow({key: metrics.get(key, "") for key in METRICS_COLUMNS}) + self._file.flush() + + def close(self) -> None: + if self._file is not None: + self._file.flush() + self._file.close() + self._file = None + self._writer = None + + +def reduce_loss( + loss_tensor: torch.Tensor, + dp_world_size: int, + group: dist.ProcessGroup | None = None, +) -> float: + """Return the data-parallel mean of a scalar loss tensor.""" + loss_clone = loss_tensor.clone().detach() + dist.all_reduce(loss_clone, op=dist.ReduceOp.SUM, group=group) + return (loss_clone / dp_world_size).item() + + +def reduce_max(value: float, group: dist.ProcessGroup | None = None) -> float: + """Return max(value) across ranks in ``group``.""" + tensor = torch.tensor([value], device=torch.cuda.current_device()) + dist.all_reduce(tensor, op=dist.ReduceOp.MAX, group=group) + return tensor.item() diff --git a/training/expert_parallel/requirements.txt b/training/expert_parallel/requirements.txt new file mode 100644 index 000000000..64b985090 --- /dev/null +++ b/training/expert_parallel/requirements.txt @@ -0,0 +1,14 @@ +# Qwen3.5 AutoEP requires Transformers with Qwen3_5 MoE text-backbone support. +transformers>=5.2 +datasets>=2.14.0 +# Mandatory for Qwen3.5 verification: the model's linear-attention layers must not use torch fallbacks. +flash-linear-attention>=0.2.2 +causal-conv1d +# Mandatory when Qwen3.5 full-attention layers use the FlashAttention2 implementation. +flash-attn +# Mandatory on H100 with Triton >= 3.4 for Qwen3.5 fast-path verification. +tilelang +# AutoEP: PR #7938 head on upstream (not on PyPI yet). Same commits as tohtana/add_autoep on tohtana/DeepSpeed. +deepspeed @ git+https://github.com/deepspeedai/DeepSpeed.git@refs/pull/7938/head +matplotlib>=3.5.0 +torch>=2.9.1 diff --git a/training/expert_parallel/train.py b/training/expert_parallel/train.py new file mode 100644 index 000000000..1d1c13aab --- /dev/null +++ b/training/expert_parallel/train.py @@ -0,0 +1,484 @@ +"""Compact MoE causal LM training example for AutoEP and ZeRO-3 leaf. + +Launch with DeepSpeed: + + deepspeed --num_gpus 8 train.py --mode autoep --autoep_size 8 + deepspeed --num_gpus 8 train.py --mode zero3_leaf +""" + +from __future__ import annotations + +import argparse +import logging +import math +import os +import random +import sys +import time +from typing import Any, NamedTuple + +import numpy as np +import torch +from transformers import ( + AutoModelForCausalLM, + Llama4ForCausalLM, + Llama4TextConfig, + MixtralConfig, + Qwen3_5MoeForCausalLM, + Qwen3_5MoeTextConfig, +) + +import deepspeed + +from data_utils import ( + build_hf_batch_generator, + build_model_config, + get_tokenizer, + validate_tokenizer_vocab_size, +) +from init_weights import load_init_weights_artifact +from metrics import MetricsLogger, reduce_loss, reduce_max + +logger = logging.getLogger(__name__) + + +MODEL_PRESETS: dict[str, dict[str, Any]] = { + "mixtral": { + "architecture": "mixtral", + "config_cls": MixtralConfig, + "display_name": "Mixtral 8x7B", + "default_tokenizer_name": "mistralai/Mixtral-8x7B-v0.1", + }, + "qwen3_5_moe": { + "architecture": "qwen3_5_moe", + "config_cls": Qwen3_5MoeTextConfig, + "display_name": "Qwen3.5 MoE", + "default_tokenizer_name": "Qwen/Qwen3-0.6B", + }, + "llama4": { + "architecture": "llama4", + "config_cls": Llama4TextConfig, + "display_name": "Llama4 Scout", + "default_tokenizer_name": "meta-llama/Llama-4-Scout-17B-16E", + }, +} + +DEEPSPEED_LEAF_MOE_BLOCK_CLASS = { + "llama4": "transformers.models.llama4.modeling_llama4.Llama4TextMoe", + "mixtral": "transformers.models.mixtral.modeling_mixtral.MixtralSparseMoeBlock", + "qwen3_5_moe": ( + "transformers.models.qwen3_5_moe.modeling_qwen3_5_moe." + "Qwen3_5MoeSparseMoeBlock" + ), +} + + +class ModelPreset(NamedTuple): + architecture: str + config_cls: type[Any] + display_name: str + num_layers_overridden: bool + + +class TrainingState(NamedTuple): + rank: int + dp_world_size: int + engine: Any + batch_gen: Any + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="AutoEP / ZeRO-3 leaf MoE training") + parser.add_argument("--mode", choices=["autoep", "zero3_leaf"], default="autoep") + parser.add_argument("--model", choices=sorted(MODEL_PRESETS), default="qwen3_5_moe") + parser.add_argument("--num_layers", type=int, default=None) + parser.add_argument("--autoep_size", type=int, default=None) + parser.add_argument("--steps", type=int, default=50) + parser.add_argument("--warmup_steps", type=int, default=5) + parser.add_argument("--log_interval", type=int, default=1) + parser.add_argument("--seq_len", type=int, default=128) + parser.add_argument("--micro_batch_size", type=int, default=2) + parser.add_argument("--grad_accum", type=int, default=1) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--dataset_name", default="wikitext") + parser.add_argument("--dataset_percentage", type=float, default=10.0) + parser.add_argument("--tokenizer_name", default=None) + parser.add_argument("--hf_num_dataloader_workers", type=int, default=0) + parser.add_argument( + "--load_init_weights", + type=str, + default=None, + help="Load a shared initialization artifact created by utils/prepare_init_weights.py.", + ) + parser.add_argument("--metrics_out", default=None) + parser.add_argument("--local_rank", type=int, default=-1) + args = parser.parse_args() + + if args.mode == "autoep" and args.autoep_size is None: + parser.error("--autoep_size is required in AutoEP mode.") + if args.load_init_weights is not None: + if not args.load_init_weights.endswith(".safetensors"): + parser.error("--load_init_weights path must end with '.safetensors'.") + if not os.path.isfile(args.load_init_weights): + parser.error(f"--load_init_weights file does not exist: {args.load_init_weights}") + + return args + + +def resolve_model_preset(args: argparse.Namespace) -> ModelPreset: + preset = MODEL_PRESETS[args.model] + architecture = preset["architecture"] + config_cls = preset["config_cls"] + num_layers_overridden = args.num_layers is not None + if args.num_layers is None: + original_config = build_model_config(config_cls, None) + args.num_layers = int(original_config.num_hidden_layers) + if args.tokenizer_name is None: + args.tokenizer_name = preset["default_tokenizer_name"] + return ModelPreset( + architecture, + config_cls, + preset["display_name"], + num_layers_overridden, + ) + + +def build_model(architecture: str, model_config: Any) -> torch.nn.Module: + if architecture == "mixtral": + return AutoModelForCausalLM.from_config(model_config) + if architecture == "qwen3_5_moe": + return Qwen3_5MoeForCausalLM(model_config) + if architecture == "llama4": + return Llama4ForCausalLM(model_config) + raise ValueError(f"Unsupported architecture: {architecture!r}") + + +def num_experts_for_config(architecture: str, model_config: Any) -> int: + if architecture in {"mixtral", "llama4"}: + return int(model_config.num_local_experts) + if architecture == "qwen3_5_moe": + return int(model_config.num_experts) + raise ValueError(f"Unsupported architecture: {architecture!r}") + + +def build_deepspeed_config( + mode: str, + architecture: str, + micro_batch_size: int, + grad_accum: int, + autoep_size: int | None, +) -> dict[str, Any]: + config: dict[str, Any] = { + "bf16": {"enabled": True}, + "optimizer": {"type": "AdamW", "params": {"lr": 1e-4}}, + "scheduler": { + "type": "WarmupCosineLR", + "params": { + "total_num_steps": 1000, + "warmup_min_ratio": 0, + "warmup_num_steps": 100, + "cos_min_ratio": 0.001, + "warmup_type": "linear", + }, + }, + "train_micro_batch_size_per_gpu": micro_batch_size, + "gradient_accumulation_steps": grad_accum, + "steps_per_print": 10, + } + if mode == "autoep": + config["zero_optimization"] = {"stage": 1} + config["expert_parallel"] = { + "enabled": True, + "autoep_size": autoep_size, + "preset_model": architecture, + } + else: + config["zero_optimization"] = { + "stage": 3, + "stage3_param_persistence_threshold": 1e5, + "leaf_module": {"classes": [DEEPSPEED_LEAF_MOE_BLOCK_CLASS[architecture]]}, + } + return config + + +def validate_autoep_args( + architecture: str, + autoep_size: int, + num_experts: int, + world_size: int, +) -> None: + valid_sizes = [ + size + for size in range(1, min(num_experts, world_size) + 1) + if num_experts % size == 0 and world_size % size == 0 + ] + if autoep_size not in valid_sizes: + raise ValueError( + f"Invalid autoep_size={autoep_size} for architecture={architecture!r}; " + f"num_experts={num_experts}, world_size={world_size}, " + f"valid sizes={valid_sizes}" + ) + + from deepspeed.module_inject.auto_ep_config import PRESET_MODELS + + preset_id = architecture + if preset_id not in PRESET_MODELS: + raise ValueError( + f"DeepSpeed does not provide AutoEP preset_model={preset_id!r}; " + f"available presets={sorted(PRESET_MODELS)}" + ) + + +def setup_distributed(args: argparse.Namespace) -> tuple[int, int]: + deepspeed.init_distributed() + rank = int(os.environ.get("RANK", 0)) + world_size = int(os.environ.get("WORLD_SIZE", 1)) + if torch.cuda.is_available(): + local_rank = int(os.environ.get("LOCAL_RANK", args.local_rank)) + if local_rank >= 0: + torch.cuda.set_device(local_rank) + return rank, world_size + + +def seed_everything(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + +def sync_cuda() -> None: + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def load_initial_weights( + path: str | None, + model: torch.nn.Module, + args: argparse.Namespace, + model_config: Any, + rank: int, +) -> None: + if path is None: + return + try: + context = load_init_weights_artifact(path, model, args=args, model_config=model_config) + except Exception as exc: + logger.error("Failed to load init weights artifact: %s", exc) + sys.exit(2) + if torch.distributed.is_initialized(): + torch.distributed.barrier() + if rank == 0: + logger.info( + "Loaded init weights artifact %s (sha256=%s)", + context["init_weights_path"], + context["init_weights_sha256"], + ) + + +def prepare_training(args: argparse.Namespace) -> TrainingState: + if args.metrics_out is None: + args.metrics_out = f"metrics_{args.mode}.csv" + + rank, world_size = setup_distributed(args) + logging.basicConfig( + level=logging.INFO if rank == 0 else logging.WARNING, + format=f"[rank {rank}] %(levelname)s: %(message)s", + ) + seed_everything(args.seed) + + preset = resolve_model_preset(args) + model_config = build_model_config(preset.config_cls, args.num_layers) + num_experts = num_experts_for_config(preset.architecture, model_config) + autoep_size = args.autoep_size if args.mode == "autoep" else None + + if args.mode == "autoep": + try: + validate_autoep_args(preset.architecture, autoep_size, num_experts, world_size) + except ValueError as exc: + logger.error("AutoEP preflight failed: %s", exc) + sys.exit(2) + + ds_config = build_deepspeed_config( + args.mode, + preset.architecture, + args.micro_batch_size, + args.grad_accum, + autoep_size, + ) + + try: + tokenizer = get_tokenizer(args.tokenizer_name, trust_remote_code=True) + tokenizer_info = validate_tokenizer_vocab_size( + tokenizer, + args.tokenizer_name, + model_config.vocab_size, + ) + except ValueError as exc: + logger.error("Tokenizer validation failed: %s", exc) + sys.exit(2) + + if rank == 0: + logger.info("Mode: %s", args.mode) + logger.info( + "Model: %s (%s), layers=%s%s, hidden=%s, experts=%s", + args.model, + preset.display_name, + args.num_layers, + " from --num_layers" if preset.num_layers_overridden else " original default", + model_config.hidden_size, + num_experts, + ) + logger.info( + "Tokenizer %s: len=%s, vocab_size=%s, model_vocab_size=%s", + args.tokenizer_name, + tokenizer_info["tokenizer_len"], + tokenizer_info["tokenizer_vocab_size"], + tokenizer_info["model_vocab_size"], + ) + logger.info( + "Seq len=%s, micro batch=%s, grad_accum=%s, steps=%s", + args.seq_len, + args.micro_batch_size, + args.grad_accum, + args.steps, + ) + + model = build_model(preset.architecture, model_config) + load_initial_weights(args.load_init_weights, model, args, model_config, rank) + + try: + engine, _, _, _ = deepspeed.initialize( + model=model, + config=ds_config, + model_parameters=model.parameters(), + ) + except Exception as exc: + logger.error("deepspeed.initialize() failed: %s", exc) + sys.exit(2) + + import deepspeed.comm as dist_comm + + dp_rank = dist_comm.get_rank(engine.data_parallel_group) + dp_world_size = engine.dp_world_size + batch_gen = build_hf_batch_generator( + dataset_name=args.dataset_name, + dataset_percentage=args.dataset_percentage, + tokenizer_name=args.tokenizer_name, + expected_vocab_size=model_config.vocab_size, + seq_len=args.seq_len, + micro_batch_size=args.micro_batch_size, + dp_world_size=dp_world_size, + dp_rank=dp_rank, + seed=args.seed, + rank=rank, + hf_num_dataloader_workers=args.hf_num_dataloader_workers, + ) + if torch.distributed.is_initialized(): + torch.distributed.barrier() + + return TrainingState( + rank=rank, + dp_world_size=dp_world_size, + engine=engine, + batch_gen=batch_gen, + ) + + +def train(args: argparse.Namespace, state: TrainingState) -> None: + metrics_logger = MetricsLogger(args.metrics_out, state.rank) + if state.rank == 0: + logger.info( + "Starting training for %s optimizer steps (warmup=%s).", + args.steps, + args.warmup_steps, + ) + + for step in range(args.steps): + sync_cuda() + step_start = time.time() + last_loss = None + + for accum_idx in range(args.grad_accum): + batch = state.batch_gen.get_batch(step, accum_idx) + outputs = state.engine( + input_ids=batch.input_ids.to(state.engine.device), + attention_mask=batch.attention_mask.to(state.engine.device), + labels=batch.labels.to(state.engine.device), + ) + loss = outputs.loss + last_loss = loss.detach().clone() + state.engine.backward(loss) + state.engine.step() + + sync_cuda() + iter_time = time.time() - step_start + reduced_loss = reduce_loss( + last_loss, + state.dp_world_size, + group=state.engine.data_parallel_group, + ) + if not math.isfinite(reduced_loss): + if state.rank == 0: + logger.error("Non-finite loss at step %s: loss=%s", step, reduced_loss) + sys.exit(3) + + if step == args.warmup_steps - 1 and torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + + if step >= args.warmup_steps and step % args.log_interval == 0: + max_iter_time = reduce_max(iter_time) + mem_allocated = torch.cuda.memory_allocated() + mem_peak_allocated = torch.cuda.max_memory_allocated() + mem_peak_reserved = torch.cuda.max_memory_reserved() + global_tokens_per_sec = ( + args.seq_len + * args.micro_batch_size + * args.grad_accum + * state.dp_world_size + / max_iter_time + if max_iter_time > 0 + else 0 + ) + + metrics_logger.log_step( + { + "step": step, + "loss": reduced_loss, + "iter_time_sec": max_iter_time, + "global_tokens_per_sec": global_tokens_per_sec, + "cuda_memory_allocated_bytes": mem_allocated, + "cuda_peak_memory_allocated_bytes": mem_peak_allocated, + "cuda_peak_memory_reserved_bytes": mem_peak_reserved, + } + ) + if state.rank == 0: + logger.info( + "Step %s: loss=%.6f, time=%.3fs, global_tps=%.0f, peak_mem=%.2f GiB", + step, + reduced_loss, + max_iter_time, + global_tokens_per_sec, + mem_peak_allocated / (1024**3), + ) + + metrics_logger.close() + if state.rank == 0: + logger.info("Metrics written to %s", args.metrics_out) + + +def main() -> None: + args = parse_args() + state = prepare_training(args) + train(args, state) + + +if __name__ == "__main__": + try: + main() + except SystemExit: + raise + except Exception as exc: + logging.error("Unhandled exception: %s", exc, exc_info=True) + sys.exit(1) diff --git a/training/expert_parallel/utils/compare_metrics.py b/training/expert_parallel/utils/compare_metrics.py new file mode 100644 index 000000000..5660a0deb --- /dev/null +++ b/training/expert_parallel/utils/compare_metrics.py @@ -0,0 +1,195 @@ +"""Compare AutoEP and ZeRO-3 leaf CSV metrics.""" + +from __future__ import annotations + +import argparse +import csv +import json +import os +from statistics import mean +from typing import Any + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt + + +BYTES_PER_GIB = 1024**3 + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Compare AutoEP and ZeRO-3 leaf metrics") + parser.add_argument("--autoep_csv", required=True) + parser.add_argument("--zero3_leaf_csv", required=True) + parser.add_argument("--out_dir", required=True) + parser.add_argument("--out_json", required=True) + parser.add_argument("--autoep_label", default="AutoEP + ZeRO-1") + parser.add_argument("--zero3_leaf_label", default="HF + ZeRO-3 leaf") + parser.add_argument("--warmup_steps", type=int, default=5) + return parser.parse_args() + + +def load_rows(path: str, warmup_steps: int) -> list[dict[str, str]]: + with open(path, newline="") as f: + rows = [row for row in csv.DictReader(f) if int(row["step"]) >= warmup_steps] + if not rows: + raise ValueError(f"No rows at or after warmup step {warmup_steps}: {path}") + return rows + + +def align_rows( + autoep_rows: list[dict[str, str]], + zero3_rows: list[dict[str, str]], +) -> tuple[list[int], list[dict[str, str]], list[dict[str, str]]]: + autoep_by_step = {int(row["step"]): row for row in autoep_rows} + zero3_by_step = {int(row["step"]): row for row in zero3_rows} + steps = sorted(set(autoep_by_step) & set(zero3_by_step)) + if not steps: + raise ValueError("AutoEP and ZeRO-3 metrics do not share any post-warmup steps") + return steps, [autoep_by_step[step] for step in steps], [zero3_by_step[step] for step in steps] + + +def avg(rows: list[dict[str, str]], column: str) -> float: + return mean(float(row[column]) for row in rows) + + +def max_int(rows: list[dict[str, str]], column: str) -> int: + return max(int(row[column]) for row in rows) + + +def write_json(path: str, payload: dict[str, Any]) -> None: + parent = os.path.dirname(path) + if parent: + os.makedirs(parent, exist_ok=True) + tmp = path + ".tmp" + with open(tmp, "w") as f: + json.dump(payload, f, indent=2) + f.write("\n") + os.replace(tmp, path) + + +def save_loss_curve( + steps: list[int], + autoep_rows: list[dict[str, str]], + zero3_rows: list[dict[str, str]], + autoep_label: str, + zero3_label: str, + out_dir: str, +) -> str: + path = os.path.join(out_dir, "loss_curve.png") + fig, ax = plt.subplots(figsize=(10, 6)) + ax.plot(steps, [float(row["loss"]) for row in autoep_rows], label=autoep_label) + ax.plot(steps, [float(row["loss"]) for row in zero3_rows], label=zero3_label) + ax.set_xlabel("Optimizer Step") + ax.set_ylabel("Loss") + ax.set_title("Loss Curve Comparison") + ax.grid(True, alpha=0.3) + ax.legend() + fig.savefig(path, dpi=150, bbox_inches="tight") + plt.close(fig) + return path + + +def save_bar_chart( + values: list[float], + labels: list[str], + ylabel: str, + title: str, + path: str, + value_format: str, +) -> str: + fig, ax = plt.subplots(figsize=(8, 5)) + bars = ax.bar(labels, values) + ax.set_ylabel(ylabel) + ax.set_title(title) + for bar in bars: + ax.text( + bar.get_x() + bar.get_width() / 2.0, + bar.get_height(), + value_format.format(bar.get_height()), + ha="center", + va="bottom", + ) + fig.savefig(path, dpi=150, bbox_inches="tight") + plt.close(fig) + return path + + +def main() -> None: + args = parse_args() + os.makedirs(args.out_dir, exist_ok=True) + + autoep_rows = load_rows(args.autoep_csv, args.warmup_steps) + zero3_rows = load_rows(args.zero3_leaf_csv, args.warmup_steps) + steps, autoep_aligned, zero3_aligned = align_rows(autoep_rows, zero3_rows) + + autoep_loss = avg(autoep_aligned, "loss") + zero3_loss = avg(zero3_aligned, "loss") + autoep_tps = avg(autoep_aligned, "global_tokens_per_sec") + zero3_tps = avg(zero3_aligned, "global_tokens_per_sec") + autoep_peak_mem = max_int(autoep_aligned, "cuda_peak_memory_allocated_bytes") + zero3_peak_mem = max_int(zero3_aligned, "cuda_peak_memory_allocated_bytes") + + labels = [args.autoep_label, args.zero3_leaf_label] + plots = { + "loss_curve": save_loss_curve( + steps, + autoep_aligned, + zero3_aligned, + args.autoep_label, + args.zero3_leaf_label, + args.out_dir, + ), + "peak_memory_bar": save_bar_chart( + [autoep_peak_mem / BYTES_PER_GIB, zero3_peak_mem / BYTES_PER_GIB], + labels, + "Peak Memory (GiB)", + "Peak GPU Memory Comparison", + os.path.join(args.out_dir, "peak_memory_bar.png"), + "{:.2f}", + ), + "throughput_bar": save_bar_chart( + [autoep_tps, zero3_tps], + labels, + "Tokens/sec", + "Average Throughput Comparison", + os.path.join(args.out_dir, "throughput_bar.png"), + "{:.0f}", + ), + } + + summary = { + "aligned_steps": len(steps), + "loss": { + "autoep_mean": autoep_loss, + "zero3_leaf_mean": zero3_loss, + "mean_abs_diff": abs(autoep_loss - zero3_loss), + }, + "throughput": { + "autoep_tokens_per_sec": autoep_tps, + "zero3_leaf_tokens_per_sec": zero3_tps, + "ratio": autoep_tps / zero3_tps if zero3_tps else None, + }, + "peak_memory": { + "autoep_bytes": autoep_peak_mem, + "zero3_leaf_bytes": zero3_peak_mem, + "autoep_gib": autoep_peak_mem / BYTES_PER_GIB, + "zero3_leaf_gib": zero3_peak_mem / BYTES_PER_GIB, + "ratio": autoep_peak_mem / zero3_peak_mem if zero3_peak_mem else None, + }, + "plots": plots, + } + write_json(args.out_json, summary) + + print("\n=== Comparison Summary ===") + print(f"Aligned steps: {len(steps)}") + print(f"Mean loss: AutoEP={autoep_loss}, ZeRO-3={zero3_loss}") + print(f"Mean abs diff (loss): {summary['loss']['mean_abs_diff']}") + print(f"Peak memory ratio (AutoEP / ZeRO-3): {summary['peak_memory']['ratio']}") + print(f"Throughput ratio (AutoEP / ZeRO-3): {summary['throughput']['ratio']}") + print(f"Summary written to: {args.out_json}") + + +if __name__ == "__main__": + main() diff --git a/training/expert_parallel/utils/prepare_init_weights.py b/training/expert_parallel/utils/prepare_init_weights.py new file mode 100644 index 000000000..26241ee6f --- /dev/null +++ b/training/expert_parallel/utils/prepare_init_weights.py @@ -0,0 +1,76 @@ +"""Create a shared initialization artifact for AutoEP comparison runs. + +Run this as a normal Python script, not through the DeepSpeed launcher: + + python utils/prepare_init_weights.py --model qwen3_5_moe --num_layers 8 \ + --output runs/qwen35/init/qwen35_l8_seed42.safetensors +""" + +from __future__ import annotations + +import argparse +import logging +import os +import random +import sys + +import numpy as np +import torch + +sys.path.insert(0, os.path.dirname(os.path.dirname(__file__))) + +from init_weights import save_init_weights_artifact +from train import MODEL_PRESETS, build_model, build_model_config + +logger = logging.getLogger(__name__) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Create shared init weights artifact") + parser.add_argument("--model", choices=sorted(MODEL_PRESETS), default="qwen3_5_moe") + parser.add_argument("--num_layers", type=int, default=None) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "--output", + required=True, + help="Output .safetensors path. A sidecar *_meta.json file is also written.", + ) + args = parser.parse_args() + if not args.output.endswith(".safetensors"): + parser.error("--output path must end with '.safetensors'.") + return args + + +def seed_everything(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + + +def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + args = parse_args() + seed_everything(args.seed) + + architecture = MODEL_PRESETS[args.model]["architecture"] + model_config = build_model_config( + MODEL_PRESETS[args.model]["config_cls"], + args.num_layers, + ) + if args.num_layers is None: + args.num_layers = int(model_config.num_hidden_layers) + + model = build_model(architecture, model_config) + context = save_init_weights_artifact( + args.output, + model, + args=args, + model_config=model_config, + rank=0, + ) + logger.info("Saved init weights to %s", context["init_weights_path"]) + logger.info("sha256=%s", context["init_weights_sha256"]) + + +if __name__ == "__main__": + main()