From e86e70b33b5243eeb9fb03c6efaa2a103c8fbade Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Wed, 21 Oct 2026 11:12:37 +0000 Subject: [PATCH 1/3] Add AutoEP HF checkpoint to rollout precision check Load a truncated HF MoE checkpoint into an AutoEP training engine, train a few steps, and move the weights to rollout engines with a different EP size through a universal checkpoint and through AffineIR transfers. Check every step bitwise and report the time of each path. The README compares the weight flows. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi --- training/deepspeed_finetune_demo/README.md | 90 +++ .../autoep_rollout_e2e.py | 582 ++++++++++++++++++ .../run_autoep_rollout_e2e.sh | 36 ++ 3 files changed, 708 insertions(+) create mode 100644 training/deepspeed_finetune_demo/autoep_rollout_e2e.py create mode 100755 training/deepspeed_finetune_demo/run_autoep_rollout_e2e.sh diff --git a/training/deepspeed_finetune_demo/README.md b/training/deepspeed_finetune_demo/README.md index d5a167e9b..c85f682c9 100644 --- a/training/deepspeed_finetune_demo/README.md +++ b/training/deepspeed_finetune_demo/README.md @@ -156,6 +156,96 @@ maximum and mean absolute loss differences from the uninterrupted baseline. A successful run means both resume paths produced comparable losses; inspect the logs and differences when diagnosing a mismatch. +## AutoEP from HF checkpoint to rollout + +`autoep_rollout_e2e.py` follows weights from a Hugging Face checkpoint through AutoEP +training to a rollout engine with a different expert layout, and checks precision at +every step: + +1. Load the checkpoint, truncated to `NUM_LAYERS` (4 by default), into an AutoEP + training engine (`TRAIN_AUTOEP_SIZE`, default 4). Every weight must equal the + safetensors tensor bit for bit; logits are compared with plain transformers. +2. Train a few steps. +3. Move the trained weights to rollout engines (`ROLLOUT_AUTOEP_SIZE`, default 8) two ways: + - `ucp`: `save_checkpoint`, `ds_to_universal`, then `load_checkpoint` with `load_universal`; + - `affine`: AutoEP storage maps, `plan_transfer` and `execute_transfers` for the + experts, directly between devices; dense weights are copied on device (gathered + under ZeRO-3). +4. Generate greedily with `HybridEngineRollout` on the training engine and on both + rollout engines. + +### Weight flows + +The two paths of step 3 next to the existing in-place rollout: + +``` +Universal checkpoint (offline, through storage) + training engine save_checkpoint ZeRO shards ds_to_universal universal load_checkpoint rollout engine + EP=4, ZeRO-k ───────────────────────────> on disk ──────────────────────> checkpoint ──(load_universal)──> EP=8, NaN init + on disk + +AffineIR transfer (online, device to device) + training engine ── expert_storage_affine_maps() ──> source maps ─┐ + (live storage: EDP replicas, ZeRO-3 fragments) ├─ plan_transfer ─> TransferPlan ─ execute_transfers ─> rollout experts + rollout engine ── expert_storage_affine_maps() ──> target maps ─┘ (per tensor) (send/recv + local copy) + dense weights: copy_ on device (GatheredParameters under ZeRO-3) + +HybridEngineRollout (in place, reference) + training engine = rollout engine; sync_weights is a no-op, so the layout cannot change +``` + +| | Universal checkpoint | AffineIR transfer | HybridEngineRollout | +|------------------------|-----------------------------|-------------------------------------|-------------------------| +| Target layout | any EP size | any EP size or expert placement | same as training | +| Goes through storage | yes, full checkpoint twice | no | no | +| Source | saved ZeRO shards | live parameter storage | live parameters | +| Moves | all weights and, optionally, optimizer state | expert weights (dense copied separately) | nothing | + +Measured on 8x Intel Arc Pro B70, Qwen3-30B-A3B truncated to 4 layers, training at EP4 +with ZeRO-2, rollout at EP8: + +| Step | Universal checkpoint | AffineIR transfer | +|------------------------------------|----------------------|-------------------| +| save / plan | 6.8 s | 3.2 s | +| convert / execute | 108.9 s | 0.27 s | +| load / dense copy | 25.9 s | 0.01 s | +| rollout weights equal trained | bitwise (8.0e9 elements) | bitwise (8.0e9 elements) | + +AffineIR sent 4.2 GB between devices and copied 0.6 GB locally, in 12 rounds. + +### Checks + +Rollout engines start from NaN weights. The script exits with an error unless: + +- every training weight equals the safetensors tensor bit for bit after loading; +- every rollout weight equals the trained weight bit for bit (experts are matched by + expert id through plain collectives); +- the two paths' logits differ by no more than the forward's own run-to-run difference, + which is zero when the forward is deterministic. + +Logits and tokens against the training engine are reported but not required to match +bit for bit. The expert layout changes how tokens are grouped in the grouped GEMM, and +with it the bf16 accumulation order, so the training and rollout forwards round +differently even with identical weights. Set `ROLLOUT_AUTOEP_SIZE` equal to +`TRAIN_AUTOEP_SIZE` to remove this difference. + +Each stage writes `zero.json` with the checks, the time of every step (save, +convert, load; plan, execute, dense copy; generation) and the generated text. + +```bash +cd training/deepspeed_finetune_demo +MODEL_NAME=/path/to/Qwen3-30B-A3B ./run_autoep_rollout_e2e.sh +``` + +`NUM_GPUS`, `NUM_LAYERS`, `TRAIN_AUTOEP_SIZE`, `ROLLOUT_AUTOEP_SIZE`, `ZERO_STAGES` +(default `"2 3"`), and `OUTPUT_ROOT` override the defaults. The script needs a DeepSpeed +build with the AutoEP storage maps +([DeepSpeed#8800](https://github.com/deepspeedai/DeepSpeed/pull/8800)) and the AffineIR +transfer executor ([DeepSpeed#8809](https://github.com/deepspeedai/DeepSpeed/pull/8809)). +The UCP path writes a full training checkpoint and its universal form under +`OUTPUT_ROOT` (about 90 GB per stage for 4 layers of Qwen3-30B-A3B), removed after each +stage. + ## HumanEval results | Model | HumanEval (base) | HumanEval+ | diff --git a/training/deepspeed_finetune_demo/autoep_rollout_e2e.py b/training/deepspeed_finetune_demo/autoep_rollout_e2e.py new file mode 100644 index 000000000..bc4cb60f0 --- /dev/null +++ b/training/deepspeed_finetune_demo/autoep_rollout_e2e.py @@ -0,0 +1,582 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""AutoEP end to end: HF checkpoint -> ZeRO training -> rollout, through UCP or AffineIR. + +1. Load a Hugging Face MoE checkpoint, truncated to ``--num_layers``, into an AutoEP training + engine. Every weight is checked bit for bit against the safetensors files, and the logits + against plain transformers. +2. Train a few steps. +3. Move the trained weights into rollout engines with a different expert layout, two ways: + + - ``ucp``: ``save_checkpoint`` -> ``ds_to_universal`` -> ``load_checkpoint(load_universal)``. + - ``affine``: AutoEP storage maps -> ``plan_transfer`` -> ``execute_transfers`` for the experts, + on device and without files. Dense weights are copied on device (gathered under ZeRO-3). + + Rollout engines start from NaN weights, so any element a path fails to write shows up. +4. Generate greedily with ``HybridEngineRollout`` on the training engine and on every rollout engine. + +Oracles: the rollout weights equal the trained weights bit for bit (experts are matched by +expert id through plain collectives, independent of the affine code), and the two paths' logits +differ by no more than the forward's own run-to-run difference, which is zero when the forward is +deterministic. Logits and tokens against the training engine, whose expert layout differs, are +reported but not required to match bit for bit. +""" + +import argparse +import json +import math +import os +import platform +import shutil +import time +from types import SimpleNamespace + +import torch +from safetensors import safe_open +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + +import deepspeed +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator +from deepspeed.checkpoint.affine_transfer import PlanBudget, plan_transfer +from deepspeed.checkpoint.affine_transfer_executor import (StagingBudget, TransferEndpoints, TransferRequest, + execute_transfers) +from deepspeed.checkpoint.ds_to_universal import main as convert_to_universal +from deepspeed.module_inject.auto_ep_layer import AutoEPMoELayer +from deepspeed.runtime.rollout import HybridEngineRollout, RolloutRequest, SamplingConfig + +DTYPE = torch.bfloat16 +# AutoEP splits the HF fused gate/up projection into w1 (gate) and w3 (up). +HF_EXPERT_NAMES = {"w1": "gate_proj", "w2": "down_proj", "w3": "up_proj"} +PROMPTS = [ + "The capital of France is", + "def fibonacci(n):", + "Mixture-of-experts models route each token to", + "In 1969, the first humans", +] + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--model_name", required=True, help="local HF checkpoint directory") + parser.add_argument("--num_layers", type=int, default=4, help="truncate the model to this many layers") + parser.add_argument("--train_autoep_size", type=int, default=4) + parser.add_argument("--train_zero_stage", type=int, default=2, choices=(1, 2, 3)) + parser.add_argument("--rollout_autoep_size", type=int, default=8) + parser.add_argument("--paths", default="ucp,affine") + parser.add_argument("--train_steps", type=int, default=3) + parser.add_argument("--lr", type=float, default=1e-4) + parser.add_argument("--seq_len", type=int, default=128) + parser.add_argument("--max_new_tokens", type=int, default=32) + parser.add_argument("--budget_mb", type=float, default=64.0, help="executor staging bytes per peer") + parser.add_argument("--work_dir", required=True, help="checkpoint directory for the UCP path") + parser.add_argument("--output", default="autoep_rollout_e2e.json") + parser.add_argument("--local_rank", type=int, default=-1) + return parser.parse_args() + + +def device(): + return torch.device(get_accelerator().current_device_name()) + + +def synchronize(): + get_accelerator().synchronize() + + +def max_over_ranks(value): + tensor = torch.tensor([float(value)], dtype=torch.float64, device=device()) + dist.all_reduce(tensor, op=dist.ReduceOp.MAX) + return tensor.item() + + +def sum_over_ranks(value): + tensor = torch.tensor([float(value)], dtype=torch.float64, device=device()) + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + return int(tensor.item()) + + +class Timer: + """Wall time of a step, as the slowest rank saw it.""" + + def __init__(self): + self.seconds = {} + + def __call__(self, name): + timer = self + + class Scope: + + def __enter__(self): + dist.barrier() + synchronize() + self.started = time.perf_counter() + + def __exit__(self, *exc): + synchronize() + timer.seconds[name] = max_over_ranks(time.perf_counter() - self.started) + + return Scope() + + +# --------------------------------------------------------------------------- +# Engines +# --------------------------------------------------------------------------- + + +def truncated_config(args): + config = AutoConfig.from_pretrained(args.model_name) + config.num_hidden_layers = args.num_layers + return config + + +def ds_config(zero_stage, autoep_size, lr, load_universal=False): + config = { + "train_micro_batch_size_per_gpu": 1, + "gradient_accumulation_steps": 1, + "bf16": { + "enabled": True + }, + "zero_optimization": { + "stage": zero_stage + }, + "optimizer": { + "type": "AdamW", + "params": { + "lr": lr, + "torch_adam": True + } + }, + "expert_parallel": { + "enabled": True, + "autoep_size": autoep_size + }, + "zero_allow_untested_optimizer": True, + } + if load_universal: + config["checkpoint"] = {"load_universal": True} + return config + + +def initialize(model, config): + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) + return engine + + +def build_rollout_engine(args, load_universal): + """A rollout engine with every weight set to NaN, so that unwritten elements cannot pass.""" + model = AutoModelForCausalLM.from_config(truncated_config(args), dtype=DTYPE) + with torch.no_grad(): + for param in model.parameters(): + param.fill_(float("nan")) + # Universal checkpoints load through the ZeRO optimizer, so the rollout engine keeps one. Both + # paths use the same engine so that their layouts, and therefore their logits, are identical. + return initialize(model, ds_config(1, args.rollout_autoep_size, args.lr, load_universal)) + + +def autoep_layers(engine): + return {name: module for name, module in engine.module.named_modules() if isinstance(module, AutoEPMoELayer)} + + +def expert_param_ids(engine): + return {id(getattr(layer.experts, name)) for layer in autoep_layers(engine).values() for name in HF_EXPERT_NAMES} + + +def dense_params(engine): + expert_ids = expert_param_ids(engine) + return {name: param for name, param in engine.module.named_parameters() if id(param) not in expert_ids} + + +def full_parameter(param): + """The whole local parameter, gathered with ZeRO's own collective when it is partitioned.""" + if hasattr(param, "ds_id"): + with deepspeed.zero.GatheredParameters([param]): + return param.data.detach().clone() + return param.data.detach().clone() + + +def flat_storage(param, shard_numel): + # ZeRO-3 keeps the rank's partition in ds_tensor with trailing padding; the map excludes it. + if hasattr(param, "ds_id"): + return param.ds_tensor.data.view(-1)[:shard_numel] + return param.data.view(-1) + + +def local_expert_ids(layer): + # AutoEP packs EP rank r's experts as the contiguous range r * L .. (r + 1) * L - 1. The HF + # weight check below confirms this independently on the training engine. + count = layer.num_local_experts + return range(layer.ep_rank * count, (layer.ep_rank + 1) * count) + + +def gather_experts(layer, name): + """The layer's whole ``[num_experts, ...]`` tensor, gathered with plain collectives.""" + packed = full_parameter(getattr(layer.experts, name)).contiguous() + parts = [torch.empty_like(packed) for _ in range(layer.ep_size)] + dist.all_gather(parts, packed, group=layer.ep_group) + return torch.cat(parts, dim=0) + + +# --------------------------------------------------------------------------- +# Comparisons +# --------------------------------------------------------------------------- + + +def bit_mismatches(got, expected): + got = got.reshape(-1) + expected = expected.reshape(-1).to(got.device) + assert got.numel() == expected.numel() + return int((got.view(torch.int16) != expected.view(torch.int16)).sum().item()) + + +class WeightCheck: + """Counts of elements that differ bit for bit, summed over ranks.""" + + def __init__(self): + self.mismatched = 0 + self.checked = 0 + self.unmatched_names = [] + + def add(self, got, expected): + self.mismatched += bit_mismatches(got, expected) + self.checked += got.numel() + + def result(self): + return { + "mismatched_elements": sum_over_ranks(self.mismatched), + "checked_elements": sum_over_ranks(self.checked), + "unmatched_names": self.unmatched_names, + } + + +class HFCheckpoint: + + def __init__(self, path): + self.path = path + with open(os.path.join(path, "model.safetensors.index.json")) as stream: + self.weight_map = json.load(stream)["weight_map"] + self.files = {} + + def __contains__(self, key): + return key in self.weight_map + + def get(self, key): + file_name = self.weight_map[key] + if file_name not in self.files: + self.files[file_name] = safe_open(os.path.join(self.path, file_name), framework="pt", device="cpu") + return self.files[file_name].get_tensor(key) + + +def hf_dense_key(name): + return name.replace(".mlp.router.gate.", ".mlp.gate.") + + +def check_against_hf(engine, checkpoint): + """Every training weight equals the HF safetensors tensor of the same name or expert id.""" + check = WeightCheck() + for name, param in dense_params(engine).items(): + key = hf_dense_key(name) + if key not in checkpoint: + check.unmatched_names.append(name) + continue + check.add(full_parameter(param), checkpoint.get(key)) + for prefix, layer in autoep_layers(engine).items(): + for name, hf_name in HF_EXPERT_NAMES.items(): + local = full_parameter(getattr(layer.experts, name)) + for index, expert in enumerate(local_expert_ids(layer)): + check.add(local[index], checkpoint.get(f"{prefix}.experts.{expert}.{hf_name}.weight")) + return check.result() + + +def changed_fraction(engine, checkpoint): + """Fraction of trained weights that differ from the HF checkpoint, to show training moved them.""" + changed, total = 0, 0 + for name, param in dense_params(engine).items(): + value = full_parameter(param) + changed += bit_mismatches(value, checkpoint.get(hf_dense_key(name))) + total += value.numel() + return sum_over_ranks(changed) / max(sum_over_ranks(total), 1) + + +def check_rollout_weights(rollout, trained): + """Rollout weights equal the trained ones: dense by name, experts by expert id.""" + check = WeightCheck() + trained_dense = dense_params(trained) + for name, param in dense_params(rollout).items(): + check.add(full_parameter(param), full_parameter(trained_dense[name])) + trained_layers = autoep_layers(trained) + for prefix, layer in autoep_layers(rollout).items(): + for name in HF_EXPERT_NAMES: + reference = gather_experts(trained_layers[prefix], name) + local = full_parameter(getattr(layer.experts, name)) + expected = reference[local_expert_ids(layer).start:local_expert_ids(layer).stop] + check.add(local, expected) + return check.result() + + +# --------------------------------------------------------------------------- +# Weight movement +# --------------------------------------------------------------------------- + + +def sync_ucp(trained, args, timer): + tag = "trained" + save_dir = os.path.join(args.work_dir, "checkpoint") + with timer("ucp_save"): + trained.save_checkpoint(save_dir, tag=tag) + with timer("ucp_convert"): + if dist.get_rank() == 0: + convert_to_universal( + SimpleNamespace(input_folder=os.path.join(save_dir, tag), + output_folder=os.path.join(save_dir, f"{tag}_universal"), + num_extract_workers=4, + num_merge_workers=4, + keep_temp_folder=False, + strict=True, + inject_missing_state=False)) + rollout = build_rollout_engine(args, load_universal=True) + with timer("ucp_load"): + rollout.load_checkpoint(save_dir, + tag=f"{tag}_universal", + load_optimizer_states=False, + load_lr_scheduler_states=False) + return rollout + + +def sync_affine(trained, args, timer): + rollout = build_rollout_engine(args, load_universal=False) + world = dist.get_world_size() + rank = dist.get_rank() + endpoints = TransferEndpoints(source={r: r for r in range(world)}, target={r: r for r in range(world)}) + budget = StagingBudget(bytes_per_peer=int(args.budget_mb * (1 << 20))) + + trained_layers = autoep_layers(trained) + with timer("affine_plan"): + transfers = [] + for prefix, layer in autoep_layers(rollout).items(): + source_layer = trained_layers[prefix] + source_maps = source_layer.expert_storage_affine_maps() + target_maps = layer.expert_storage_affine_maps() + for name in HF_EXPERT_NAMES: + plan = plan_transfer(target_maps[name], source_maps[name], PlanBudget()) + source_param = getattr(source_layer.experts, name) + sources = {} + if rank in source_maps[name].shard_shapes: + numel = math.prod(source_maps[name].shard_shapes[rank]) + sources[rank] = flat_storage(source_param, numel) + targets = {} + if rank in target_maps[name].shard_shapes: + shape = target_maps[name].shard_shapes[rank] + targets[rank] = getattr(layer.experts, name).data.view(shape) + transfers.append(TransferRequest(f"{prefix}.{name}", plan, sources, targets)) + with timer("affine_execute"): + stats = execute_transfers(transfers, endpoints, DTYPE, budget) + with timer("affine_dense_copy"): + # Under ZeRO-0/1/2 every rank holds the whole dense weight, so this is a device copy. + # ZeRO-3 shards it; gathering here stands in for a ZeRO storage map. + trained_dense = dense_params(trained) + for name, param in dense_params(rollout).items(): + param.data.copy_(full_parameter(trained_dense[name])) + return rollout, { + "sent_bytes": sum_over_ranks(stats.sent_bytes), + "local_copy_bytes": sum_over_ranks(stats.local_copy_bytes), + "rounds": stats.rounds, + } + + +# --------------------------------------------------------------------------- +# Rollout +# --------------------------------------------------------------------------- + + +def prompts(tokenizer): + tokenizer.padding_side = "left" + batch = tokenizer(PROMPTS, return_tensors="pt", padding=True) + return batch["input_ids"].to(device()), batch["attention_mask"].to(device()) + + +def generate(engine, tokenizer, prompt_ids, prompt_mask, args, timer, name): + engine.eval() + rollout = HybridEngineRollout(engine, tokenizer) + with timer(f"{name}_generate"): + batch = rollout.generate(RolloutRequest(prompt_ids, prompt_mask), + SamplingConfig(max_new_tokens=args.max_new_tokens, temperature=0.0)) + return batch.input_ids + + +@torch.no_grad() +def logits_of(engine, input_ids, attention_mask): + engine.eval() + return engine.module(input_ids=input_ids, attention_mask=attention_mask).logits.float() + + +def all_ranks(value): + gathered = [None] * dist.get_world_size() + dist.all_gather_object(gathered, value) + return gathered + + +def compare_logits(got, expected, mask): + """Every rank runs the same prompts; each compares its own outputs, and results are combined.""" + valid = mask.bool() + difference = (got - expected).abs()[valid] + argmax_differs = (got.argmax(-1) != expected.argmax(-1))[valid] + by_rank = all_ranks({ + "bitwise_equal": bool(torch.equal(got, expected)), + "max_abs_diff": difference.max().item(), + "argmax_mismatched_positions": int(argmax_differs.sum().item()), + }) + return { + "bitwise_equal": all(entry["bitwise_equal"] for entry in by_rank), + "max_abs_diff": max(entry["max_abs_diff"] for entry in by_rank), + "reference_max_abs": expected[valid].abs().max().item(), + "argmax_mismatched_positions": sum(entry["argmax_mismatched_positions"] for entry in by_rank), + "positions": int(valid.sum().item()) * dist.get_world_size(), + "bitwise_equal_by_rank": [entry["bitwise_equal"] for entry in by_rank], + } + + +def compare_tokens(got, expected, prompt_len): + same = (got[:, prompt_len:] == expected[:, prompt_len:]) + first = [int((~row).nonzero()[0]) if not row.all() else None for row in same] + by_rank = all_ranks(first) + return { + "identical": all(divergence is None for rank_first in by_rank for divergence in rank_first), + "first_divergence_by_rank": by_rank, + } + + +# --------------------------------------------------------------------------- +# Driver +# --------------------------------------------------------------------------- + + +def train(engine, config, args): + engine.train() + generator = torch.Generator().manual_seed(4321 + dist.get_rank()) + losses = [] + for _ in range(args.train_steps): + tokens = torch.randint(0, config.vocab_size, (1, args.seq_len), generator=generator).to(device()) + loss = engine(input_ids=tokens, labels=tokens).loss + engine.backward(loss) + engine.step() + losses.append(loss.item()) + return losses + + +def manifest(args): + accelerator = get_accelerator() + return { + "args": vars(args), + "world_size": dist.get_world_size(), + "device_name": accelerator.get_device_name(0) if hasattr(accelerator, "get_device_name") else None, + "torch": torch.__version__, + "deepspeed": deepspeed.__version__, + "host": platform.node(), + } + + +def main(): + args = parse_args() + get_accelerator().set_device(int(os.environ.get("LOCAL_RANK", 0))) + deepspeed.init_distributed() + rank = dist.get_rank() + paths = args.paths.split(",") + if rank == 0: + shutil.rmtree(args.work_dir, ignore_errors=True) + os.makedirs(args.work_dir) + timer = Timer() + report = {"manifest": manifest(args), "oracles": {}} + oracles = report["oracles"] + + tokenizer = AutoTokenizer.from_pretrained(args.model_name) + prompt_ids, prompt_mask = prompts(tokenizer) + prompt_len = prompt_ids.shape[1] + checkpoint = HFCheckpoint(args.model_name) + config = truncated_config(args) + + # 1. HF checkpoint -> training engine. + with timer("hf_load"): + model = AutoModelForCausalLM.from_pretrained(args.model_name, config=config, dtype=DTYPE) + with torch.no_grad(): + model.to(device()).eval() + hf_logits = model(input_ids=prompt_ids, attention_mask=prompt_mask).logits.float() + trained = initialize(model, ds_config(args.train_zero_stage, args.train_autoep_size, args.lr)) + oracles["hf_weights"] = check_against_hf(trained, checkpoint) + oracles["hf_logits"] = compare_logits(logits_of(trained, prompt_ids, prompt_mask), hf_logits, prompt_mask) + del hf_logits + + # 2. Train. + with timer("train"): + report["losses"] = train(trained, config, args) + report["trained_dense_changed_fraction"] = changed_fraction(trained, checkpoint) + + # Shared-weight rollout on the training engine itself: the reference sequences. + reference_ids = generate(trained, tokenizer, prompt_ids, prompt_mask, args, timer, "trained") + reference_mask = (reference_ids != tokenizer.pad_token_id).long() + reference_mask[:, :prompt_len] = prompt_mask + trained_logits = logits_of(trained, reference_ids, reference_mask) + oracles["trained_logits_repeat"] = compare_logits(logits_of(trained, reference_ids, reference_mask), + trained_logits, reference_mask) + + # 3. Move the weights, 4. roll out. + rollouts = {} + for path in paths: + if path == "ucp": + rollout = sync_ucp(trained, args, timer) + elif path == "affine": + rollout, report["affine_transfer"] = sync_affine(trained, args, timer) + else: + raise ValueError(path) + result = {"weights_vs_trained": check_rollout_weights(rollout, trained)} + ids = generate(rollout, tokenizer, prompt_ids, prompt_mask, args, timer, path) + logits = logits_of(rollout, reference_ids, reference_mask) + # The same engine and inputs a second time: shows whether the forward itself is deterministic. + result["logits_repeat"] = compare_logits(logits_of(rollout, reference_ids, reference_mask), logits, + reference_mask) + result["tokens_vs_trained"] = compare_tokens(ids, reference_ids, prompt_len) + result["logits_vs_trained"] = compare_logits(logits, trained_logits, reference_mask) + rollouts[path] = (rollout, ids, logits) + oracles[path] = result + if len(rollouts) == 2: + (_, ids_a, logits_a), (_, ids_b, logits_b) = rollouts.values() + oracles["ucp_vs_affine"] = { + "tokens": compare_tokens(ids_a, ids_b, prompt_len), + "logits": compare_logits(logits_a, logits_b, reference_mask), + } + + report["seconds"] = timer.seconds + report["generations"] = {"trained": tokenizer.batch_decode(reference_ids[:, prompt_len:])} + for path, (_, ids, _) in rollouts.items(): + report["generations"][path] = tokenizer.batch_decode(ids[:, prompt_len:]) + + failures = [] + if oracles["hf_weights"]["mismatched_elements"] or oracles["hf_weights"]["unmatched_names"]: + failures.append("hf_weights") + if report["trained_dense_changed_fraction"] == 0: + failures.append("training did not change any weight") + for path in rollouts: + if oracles[path]["weights_vs_trained"]["mismatched_elements"]: + failures.append(f"{path} weights") + if "ucp_vs_affine" in oracles: + # Same layout and bitwise equal weights, so the two paths may differ only by the forward's + # own run-to-run noise: zero where the forward is deterministic. + repeats = [oracles["trained_logits_repeat"]] + [oracles[path]["logits_repeat"] for path in rollouts] + noise = max(repeat["max_abs_diff"] for repeat in repeats) + if oracles["ucp_vs_affine"]["logits"]["max_abs_diff"] > noise: + failures.append("ucp_vs_affine logits beyond run-to-run noise") + report["failures"] = failures + + if rank == 0: + print(json.dumps({"oracles": oracles, "seconds": timer.seconds, "failures": failures}, indent=2)) + with open(args.output, "w") as stream: + json.dump(report, stream, indent=2, default=str) + print(f"wrote {args.output}") + for rollout, _, _ in rollouts.values(): + rollout.destroy() + trained.destroy() + if failures: + raise SystemExit(f"failed: {failures}") + + +if __name__ == "__main__": + main() diff --git a/training/deepspeed_finetune_demo/run_autoep_rollout_e2e.sh b/training/deepspeed_finetune_demo/run_autoep_rollout_e2e.sh new file mode 100755 index 000000000..7926d77dd --- /dev/null +++ b/training/deepspeed_finetune_demo/run_autoep_rollout_e2e.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team + +set -euo pipefail + +# HF checkpoint -> AutoEP ZeRO training -> rollout at a new expert layout, once through a +# universal checkpoint (UCP) and once through AffineIR transfers, with bitwise oracles. +# +# MODEL_NAME must be a local HF checkpoint directory. Override NUM_GPUS, NUM_LAYERS, +# TRAIN_AUTOEP_SIZE, ROLLOUT_AUTOEP_SIZE, ZERO_STAGES, and OUTPUT_ROOT as needed. +# The UCP path writes a full training checkpoint and its universal form under OUTPUT_ROOT. + +NUM_GPUS="${NUM_GPUS:-8}" +NUM_LAYERS="${NUM_LAYERS:-4}" +TRAIN_AUTOEP_SIZE="${TRAIN_AUTOEP_SIZE:-4}" +ROLLOUT_AUTOEP_SIZE="${ROLLOUT_AUTOEP_SIZE:-8}" +ZERO_STAGES="${ZERO_STAGES:-2 3}" +MODEL_NAME="${MODEL_NAME:?set MODEL_NAME to a local HF checkpoint directory}" +OUTPUT_ROOT="${OUTPUT_ROOT:-${PWD}/autoep_rollout_e2e}" +EXAMPLE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +export TRANSFORMERS_OFFLINE="${TRANSFORMERS_OFFLINE:-1}" +mkdir -p "${OUTPUT_ROOT}" + +for stage in ${ZERO_STAGES}; do + torchrun --nproc_per_node "${NUM_GPUS}" "${EXAMPLE_DIR}/autoep_rollout_e2e.py" \ + --model_name "${MODEL_NAME}" \ + --num_layers "${NUM_LAYERS}" \ + --train_autoep_size "${TRAIN_AUTOEP_SIZE}" \ + --rollout_autoep_size "${ROLLOUT_AUTOEP_SIZE}" \ + --train_zero_stage "${stage}" \ + --work_dir "${OUTPUT_ROOT}/zero${stage}_checkpoint" \ + --output "${OUTPUT_ROOT}/zero${stage}.json" 2>&1 | tee "${OUTPUT_ROOT}/zero${stage}.log" + rm -rf "${OUTPUT_ROOT}/zero${stage}_checkpoint" +done From 91f80db63016b0ce3cd395bc52cceb863df51f54 Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Wed, 21 Oct 2026 11:37:06 +0000 Subject: [PATCH 2/3] AutoEP rollout check: run ZeRO-3 through AffineIR Release ZeRO-3 prefetched parameters after each forward so the later gathers do not find them in flight, report the executor's rounds as the maximum over ranks, and run only the AffineIR path for ZeRO-3, whose universal checkpoint the ZeRO-1 rollout engine cannot load. Add the ZeRO-3 results to the README. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi --- training/deepspeed_finetune_demo/README.md | 26 ++++++++++++------- .../autoep_rollout_e2e.py | 17 +++++++++--- .../run_autoep_rollout_e2e.sh | 7 +++++ 3 files changed, 36 insertions(+), 14 deletions(-) diff --git a/training/deepspeed_finetune_demo/README.md b/training/deepspeed_finetune_demo/README.md index c85f682c9..40b4c6082 100644 --- a/training/deepspeed_finetune_demo/README.md +++ b/training/deepspeed_finetune_demo/README.md @@ -201,17 +201,19 @@ HybridEngineRollout (in place, reference) | Source | saved ZeRO shards | live parameter storage | live parameters | | Moves | all weights and, optionally, optimizer state | expert weights (dense copied separately) | nothing | -Measured on 8x Intel Arc Pro B70, Qwen3-30B-A3B truncated to 4 layers, training at EP4 -with ZeRO-2, rollout at EP8: +Measured on 8x Intel Arc Pro B70, Qwen3-30B-A3B truncated to 4 layers, training at EP4, +rollout at EP8 (UCP is not run for ZeRO-3, see below): -| Step | Universal checkpoint | AffineIR transfer | -|------------------------------------|----------------------|-------------------| -| save / plan | 6.8 s | 3.2 s | -| convert / execute | 108.9 s | 0.27 s | -| load / dense copy | 25.9 s | 0.01 s | -| rollout weights equal trained | bitwise (8.0e9 elements) | bitwise (8.0e9 elements) | +| Step | UCP, ZeRO-2 | AffineIR, ZeRO-2 | AffineIR, ZeRO-3 | +|------------------------------------|--------------|------------------|------------------| +| save / plan | 7.1 s | 3.2 s | 0.05 s | +| convert / execute | 107.7 s | 0.27 s | 0.23 s | +| load / dense copy | 25.1 s | 0.01 s | 0.18 s | +| rollout weights equal trained | bitwise | bitwise | bitwise | -AffineIR sent 4.2 GB between devices and copied 0.6 GB locally, in 12 rounds. +Each check covers 8.0e9 elements. AffineIR sent 4.2 GB between devices and copied 0.6 GB +locally under ZeRO-2 (3.6 GB and 1.2 GB under ZeRO-3), in 12 rounds. Under ZeRO-3 the +dense copy gathers each weight. ### Checks @@ -242,7 +244,11 @@ MODEL_NAME=/path/to/Qwen3-30B-A3B ./run_autoep_rollout_e2e.sh build with the AutoEP storage maps ([DeepSpeed#8800](https://github.com/deepspeedai/DeepSpeed/pull/8800)) and the AffineIR transfer executor ([DeepSpeed#8809](https://github.com/deepspeedai/DeepSpeed/pull/8809)). -The UCP path writes a full training checkpoint and its universal form under +ZeRO-3 training with AutoEP also needs the fix that makes ZeRO-3 gather expert and dense +parameters in the same order on every rank; without it, training can hang from the third +step. A universal checkpoint converted from ZeRO-3 keeps ZeRO-3 model file names, which +the ZeRO-1 rollout engine cannot load, so the run script uses only the AffineIR path for +ZeRO-3. The UCP path writes a full training checkpoint and its universal form under `OUTPUT_ROOT` (about 90 GB per stage for 4 layers of Qwen3-30B-A3B), removed after each stage. diff --git a/training/deepspeed_finetune_demo/autoep_rollout_e2e.py b/training/deepspeed_finetune_demo/autoep_rollout_e2e.py index bc4cb60f0..c46ad8405 100644 --- a/training/deepspeed_finetune_demo/autoep_rollout_e2e.py +++ b/training/deepspeed_finetune_demo/autoep_rollout_e2e.py @@ -89,12 +89,16 @@ def max_over_ranks(value): return tensor.item() -def sum_over_ranks(value): +def reduce_over_ranks(value, op): tensor = torch.tensor([float(value)], dtype=torch.float64, device=device()) - dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + dist.all_reduce(tensor, op=op) return int(tensor.item()) +def sum_over_ranks(value): + return reduce_over_ranks(value, dist.ReduceOp.SUM) + + class Timer: """Wall time of a step, as the slowest rank saw it.""" @@ -379,7 +383,7 @@ def sync_affine(trained, args, timer): return rollout, { "sent_bytes": sum_over_ranks(stats.sent_bytes), "local_copy_bytes": sum_over_ranks(stats.local_copy_bytes), - "rounds": stats.rounds, + "rounds": reduce_over_ranks(stats.rounds, dist.ReduceOp.MAX), } @@ -400,13 +404,18 @@ def generate(engine, tokenizer, prompt_ids, prompt_mask, args, timer, name): with timer(f"{name}_generate"): batch = rollout.generate(RolloutRequest(prompt_ids, prompt_mask), SamplingConfig(max_new_tokens=args.max_new_tokens, temperature=0.0)) + # ZeRO-3 may leave prefetched parameters in flight after a forward; a later + # GatheredParameters on them would fail. + engine.empty_partition_cache() return batch.input_ids @torch.no_grad() def logits_of(engine, input_ids, attention_mask): engine.eval() - return engine.module(input_ids=input_ids, attention_mask=attention_mask).logits.float() + logits = engine.module(input_ids=input_ids, attention_mask=attention_mask).logits.float() + engine.empty_partition_cache() + return logits def all_ranks(value): diff --git a/training/deepspeed_finetune_demo/run_autoep_rollout_e2e.sh b/training/deepspeed_finetune_demo/run_autoep_rollout_e2e.sh index 7926d77dd..e0404a8f0 100755 --- a/training/deepspeed_finetune_demo/run_autoep_rollout_e2e.sh +++ b/training/deepspeed_finetune_demo/run_autoep_rollout_e2e.sh @@ -24,12 +24,19 @@ export TRANSFORMERS_OFFLINE="${TRANSFORMERS_OFFLINE:-1}" mkdir -p "${OUTPUT_ROOT}" for stage in ${ZERO_STAGES}; do + # A universal checkpoint converted from ZeRO-3 keeps ZeRO-3 model file names, which the + # ZeRO-1 rollout engine cannot load, so ZeRO-3 runs only the AffineIR path. + paths="ucp,affine" + if [ "${stage}" = "3" ]; then + paths="affine" + fi torchrun --nproc_per_node "${NUM_GPUS}" "${EXAMPLE_DIR}/autoep_rollout_e2e.py" \ --model_name "${MODEL_NAME}" \ --num_layers "${NUM_LAYERS}" \ --train_autoep_size "${TRAIN_AUTOEP_SIZE}" \ --rollout_autoep_size "${ROLLOUT_AUTOEP_SIZE}" \ --train_zero_stage "${stage}" \ + --paths "${paths}" \ --work_dir "${OUTPUT_ROOT}/zero${stage}_checkpoint" \ --output "${OUTPUT_ROOT}/zero${stage}.json" 2>&1 | tee "${OUTPUT_ROOT}/zero${stage}.log" rm -rf "${OUTPUT_ROOT}/zero${stage}_checkpoint" From 4adfb61e608de29547d2e40d00d2cd42410dc408 Mon Sep 17 00:00:00 2001 From: "Jin, Youzhi" Date: Wed, 21 Oct 2026 11:56:43 +0000 Subject: [PATCH 3/3] Link the ZeRO-3 gather-order fix (DeepSpeed#8876) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi --- training/deepspeed_finetune_demo/README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/training/deepspeed_finetune_demo/README.md b/training/deepspeed_finetune_demo/README.md index 40b4c6082..4fdc68d97 100644 --- a/training/deepspeed_finetune_demo/README.md +++ b/training/deepspeed_finetune_demo/README.md @@ -245,7 +245,8 @@ build with the AutoEP storage maps ([DeepSpeed#8800](https://github.com/deepspeedai/DeepSpeed/pull/8800)) and the AffineIR transfer executor ([DeepSpeed#8809](https://github.com/deepspeedai/DeepSpeed/pull/8809)). ZeRO-3 training with AutoEP also needs the fix that makes ZeRO-3 gather expert and dense -parameters in the same order on every rank; without it, training can hang from the third +parameters in the same order on every rank +([DeepSpeed#8876](https://github.com/deepspeedai/DeepSpeed/pull/8876)); without it, training can hang from the third step. A universal checkpoint converted from ZeRO-3 keeps ZeRO-3 model file names, which the ZeRO-1 rollout engine cannot load, so the run script uses only the AffineIR path for ZeRO-3. The UCP path writes a full training checkpoint and its universal form under