#!/usr/bin/env python3
"""Standalone Qwen3-MoE DDP auxiliary-scaling reproducer.
This file intentionally has no PartitionCheck import. It calls the frozen
Transformers Trainer and Qwen3-MoE entrypoints directly, writes raw FP64 vector
bytes plus a compact JSON receipt, and can compare two independently produced
runs without using the detector under evaluation.
"""
from __future__ import annotations
import argparse
import hashlib
import inspect
import json
import math
import os
import platform
import shutil
import sys
import tempfile
from pathlib import Path
from typing import Any, Final
CANDIDATE_ID: Final = "HQ03-QWEN-DDP-AUX-SCALE-V1"
TRANSFORMERS_VERSION: Final = "5.15.0"
CUDA_TORCH_VERSION: Final = "2.13.0+cu130"
CPU_TORCH_VERSION: Final = "2.13.0+cpu"
TRAINER_SOURCE_SHA256: Final = "6ea719d4f225c6fcbccfcd2137cdfb304426c93c4dd65983ae4186f6e8fb3d8d"
QWEN_SOURCE_SHA256: Final = "56d820671d810b68f31056605cec0c674994c8f962370194225911ac6a71a365"
ROUTER_AUX_COEFFICIENT: Final = 0.01
GLOBAL_TOKEN_MASS: Final = 64
LEARNING_RATE: Final = 0.01
MODES: Final = ("current", "ce_only", "aux_only", "component_aware_repair")
PARTITIONS: Final = (
"ddp1_full_global",
"ddp2_mirrored_local_distribution",
"ddp2_rank_skewed_boundary_diagnostic",
)
VECTOR_NAMES: Final = (
"total_gradient",
"router_gradient",
"nonrouter_gradient",
"parameter_delta",
)
class ReproError(RuntimeError):
"""Fail-closed standalone-reproducer error."""
def _canonical_json(value: object) -> bytes:
return (
json.dumps(
value,
allow_nan=False,
ensure_ascii=True,
separators=(",", ":"),
sort_keys=True,
).encode("utf-8")
+ b"\n"
)
def _sha256_bytes(value: bytes) -> str:
return hashlib.sha256(value).hexdigest()
def _sha256_file(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for block in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(block)
return digest.hexdigest()
def _write_new(path: Path, value: bytes) -> None:
if path.exists() or path.is_symlink():
raise ReproError(f"refusing to replace existing output: {path}")
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_name(f".{path.name}.tmp-{os.getpid()}")
temporary.write_bytes(value)
os.replace(temporary, path)
def _fixture(variant: int, device: Any) -> dict[str, Any]:
import torch
if variant < 1 or variant > 999:
raise ReproError("fixture variant must be in [1, 999]")
rows: list[Any] = []
for base_index in range(4):
length = 9
step = 3 + base_index
offset = 7 + base_index * 11 + variant * 5
sequence = (torch.arange(length, dtype=torch.long) * step + offset) % 125 + 3
rows.extend((sequence.clone(), sequence.clone()))
input_ids = torch.stack(rows).to(device=device)
return {
"attention_mask": torch.ones_like(input_ids),
"input_ids": input_ids,
"labels": input_ids.clone(),
}
def _partition_indices(partition: str, world_size: int, rank: int) -> list[int]:
if partition == "ddp1_full_global":
if world_size != 1 or rank != 0:
raise ReproError("DDP1 partition requires exactly one process")
return list(range(8))
if world_size != 2 or rank not in {0, 1}:
raise ReproError("DDP2 partition requires exactly two processes")
if partition == "ddp2_mirrored_local_distribution":
return ([0, 2, 4, 6], [1, 3, 5, 7])[rank]
if partition == "ddp2_rank_skewed_boundary_diagnostic":
return ([0, 1, 2, 3], [4, 5, 6, 7])[rank]
raise ReproError(f"unsupported partition: {partition}")
def _model(device: Any, coefficient: float) -> Any:
import torch
from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM
torch.manual_seed(17)
config = Qwen3MoeConfig(
attention_dropout=0.0,
bos_token_id=1,
eos_token_id=2,
head_dim=16,
hidden_size=64,
intermediate_size=128,
max_position_embeddings=128,
moe_intermediate_size=64,
num_attention_heads=4,
num_experts=4,
num_experts_per_tok=1,
num_hidden_layers=2,
num_key_value_heads=2,
output_router_logits=True,
pad_token_id=0,
router_aux_loss_coef=coefficient,
use_cache=False,
vocab_size=128,
)
return Qwen3MoeForCausalLM(config).to(device=device, dtype=torch.float32)
def _trainer(model: Any, *, use_cpu: bool) -> Any:
from transformers import Trainer, TrainingArguments
output_dir = tempfile.mkdtemp(prefix="hq03-standalone-trainer-")
arguments = TrainingArguments(
average_tokens_across_devices=True,
bf16=False,
dataloader_pin_memory=False,
disable_tqdm=True,
fp16=False,
gradient_accumulation_steps=1,
max_grad_norm=0.0,
output_dir=output_dir,
per_device_train_batch_size=8,
remove_unused_columns=False,
report_to="none",
seed=17,
tf32=False,
use_cpu=use_cpu,
)
trainer = Trainer(model=model, args=arguments)
if trainer.model_accepts_loss_kwargs is not True:
raise ReproError("frozen Qwen target no longer accepts num_items_in_batch")
if trainer.accelerator.mixed_precision != "no":
raise ReproError("standalone reproducer requires explicit FP32")
trainer.current_gradient_accumulation_steps = 1
trainer._hq03_temporary_output = output_dir
return trainer
def _flatten_parameters(named: list[tuple[str, Any]]) -> Any:
import torch
return torch.cat(
[
parameter.detach().reshape(-1).to(device="cpu", dtype=torch.float64)
for _, parameter in named
]
)
def _flatten_gradients(named: list[tuple[str, Any]], *, group: str) -> tuple[Any, int]:
import torch
parts = []
none_count = 0
for name, parameter in named:
is_router = ".mlp.gate.weight" in name
if group == "router" and not is_router:
continue
if group == "nonrouter" and is_router:
continue
if parameter.grad is None:
none_count += 1
parts.append(torch.zeros(parameter.numel(), dtype=torch.float64))
else:
parts.append(parameter.grad.detach().reshape(-1).to(device="cpu", dtype=torch.float64))
if not parts:
raise ReproError(f"empty gradient group: {group}")
return torch.cat(parts), none_count
def _mean_across_ranks(value: Any) -> Any:
import torch.distributed as dist
result = value.detach().clone().to(dtype=value.new_empty(()).float().dtype)
dist.all_reduce(result, op=dist.ReduceOp.SUM)
return result / dist.get_world_size()
def _tensor_bytes(value: Any) -> bytes:
tensor = value.detach().to(device="cpu", dtype=value.new_empty(()).double().dtype).contiguous()
return tensor.numpy().astype("<f8", copy=False).tobytes(order="C")
def _tensor_summary(value: Any) -> dict[str, object]:
import torch
tensor = value.detach().to(device="cpu", dtype=torch.float64).reshape(-1)
raw = _tensor_bytes(tensor)
return {
"l2": float(torch.linalg.vector_norm(tensor).item()),
"max_abs": float(torch.max(torch.abs(tensor)).item()) if tensor.numel() else 0.0,
"numel": int(tensor.numel()),
"sha256": _sha256_bytes(raw),
}
def _runtime(profile: str, expected_world_size: int) -> tuple[Any, Any, int, int, Any]:
import torch
import torch.distributed as dist
import transformers
world_size = int(os.environ.get("WORLD_SIZE", "0"))
rank = int(os.environ.get("RANK", "-1"))
local_rank = int(os.environ.get("LOCAL_RANK", "-1"))
if world_size != expected_world_size or rank not in range(world_size):
raise ReproError("torchrun world/rank identity differs from request")
if sys.version_info[:2] != (3, 11) or transformers.__version__ != TRANSFORMERS_VERSION:
raise ReproError("frozen Python/Transformers runtime differs")
if profile == "cuda-certified":
if (
torch.__version__ != CUDA_TORCH_VERSION
or torch.cuda.device_count() < world_size
or not dist.is_nccl_available()
):
raise ReproError("frozen CUDA/NCCL runtime differs")
backend = "nccl"
device = torch.device("cuda", local_rank)
torch.cuda.set_device(device)
torch.cuda.manual_seed_all(17)
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
elif profile == "cpu-smoke":
if torch.__version__ != CPU_TORCH_VERSION or not dist.is_gloo_available():
raise ReproError("frozen CPU/Gloo smoke runtime differs")
backend = "gloo"
device = torch.device("cpu")
else:
raise ReproError(f"unsupported runtime profile: {profile}")
torch.manual_seed(17)
torch.use_deterministic_algorithms(True)
dist.init_process_group(backend=backend)
return torch, dist, rank, local_rank, device
def _source_receipt() -> dict[str, object]:
import torch
import transformers
from transformers import Qwen3MoeForCausalLM, Trainer
trainer_path = Path(inspect.getsourcefile(Trainer) or "")
qwen_path = Path(inspect.getsourcefile(Qwen3MoeForCausalLM) or "")
if not trainer_path.is_file() or not qwen_path.is_file():
raise ReproError("cannot resolve installed target source files")
trainer_sha = _sha256_file(trainer_path)
qwen_sha = _sha256_file(qwen_path)
if trainer_sha != TRAINER_SOURCE_SHA256 or qwen_sha != QWEN_SOURCE_SHA256:
raise ReproError("installed Transformers target source differs from frozen receipts")
return {
"python": platform.python_version(),
"qwen_source_sha256": qwen_sha,
"torch": torch.__version__,
"trainer_source_sha256": trainer_sha,
"transformers": transformers.__version__,
}
def _run_worker(args: argparse.Namespace) -> int:
from torch.nn.parallel import DistributedDataParallel
expected_world_size = 1 if args.partition == "ddp1_full_global" else 2
torch, dist, rank, local_rank, device = _runtime(args.runtime_profile, expected_world_size)
trainer: Any | None = None
try:
source = _source_receipt()
full_batch = _fixture(args.fixture_variant, device)
indices = _partition_indices(args.partition, expected_world_size, rank)
selected = {name: value[indices] for name, value in full_batch.items()}
coefficient = 0.0 if args.mode == "ce_only" else ROUTER_AUX_COEFFICIENT
model = _model(device, coefficient)
named = list(model.named_parameters())
initial = _flatten_parameters(named)
initial_sha256 = _sha256_bytes(_tensor_bytes(initial))
model.zero_grad(set_to_none=True)
if device.type == "cuda":
executed = DistributedDataParallel(
model,
device_ids=[local_rank],
output_device=local_rank,
find_unused_parameters=True,
)
else:
executed = DistributedDataParallel(model, find_unused_parameters=True)
if args.mode in {"current", "ce_only"}:
trainer = _trainer(model, use_cpu=device.type == "cpu")
reported = trainer.training_step(
executed,
selected,
num_items_in_batch=GLOBAL_TOKEN_MASS,
)
with torch.no_grad():
traced, outputs = trainer.compute_loss(
model,
selected,
return_outputs=True,
num_items_in_batch=GLOBAL_TOKEN_MASS,
)
local_aux = expected_world_size * coefficient * outputs.aux_loss
local_ce = traced - local_aux
objective = _mean_across_ranks(reported)
ce_component = _mean_across_ranks(local_ce)
aux_component = _mean_across_ranks(local_aux)
else:
outputs = executed(
input_ids=selected["input_ids"],
attention_mask=selected["attention_mask"],
output_router_logits=True,
use_cache=False,
)
local_aux_unscaled = ROUTER_AUX_COEFFICIENT * outputs.aux_loss
if args.mode == "aux_only":
local_ce = torch.zeros_like(local_aux_unscaled)
local_aux = expected_world_size * local_aux_unscaled
loss = local_aux
elif args.mode == "component_aware_repair":
shifted = torch.nn.functional.pad(selected["labels"], (0, 1), value=-100)[
..., 1:
].contiguous()
ce_sum = torch.nn.functional.cross_entropy(
outputs.logits.float().reshape(-1, outputs.logits.shape[-1]),
shifted.reshape(-1),
ignore_index=-100,
reduction="sum",
)
local_ce = expected_world_size * ce_sum / GLOBAL_TOKEN_MASS
local_aux = local_aux_unscaled
loss = local_ce + local_aux
else:
raise ReproError(f"unsupported worker mode: {args.mode}")
loss.backward()
objective = _mean_across_ranks(loss)
ce_component = _mean_across_ranks(local_ce)
aux_component = _mean_across_ranks(local_aux)
total, total_none = _flatten_gradients(named, group="all")
router, router_none = _flatten_gradients(named, group="router")
nonrouter, nonrouter_none = _flatten_gradients(named, group="nonrouter")
optimizer = torch.optim.SGD(
model.parameters(),
lr=LEARNING_RATE,
momentum=0.0,
weight_decay=0.0,
)
optimizer.step()
updated = _flatten_parameters(named)
vectors = {
"nonrouter_gradient": nonrouter,
"parameter_delta": updated - initial,
"router_gradient": router,
"total_gradient": total,
}
local_hashes = {name: _tensor_summary(value)["sha256"] for name, value in vectors.items()}
rank_hashes: list[dict[str, object] | None] = [None for _ in range(expected_world_size)]
dist.all_gather_object(rank_hashes, local_hashes)
if any(row != rank_hashes[0] for row in rank_hashes):
raise ReproError("DDP ranks reconstructed different reduced vectors")
dist.barrier()
if rank != 0:
return 0
output_dir = Path(args.out)
if output_dir.exists() or output_dir.is_symlink():
raise ReproError("standalone worker output directory must be fresh")
output_dir.mkdir(parents=True)
vector_receipts: dict[str, object] = {}
for name, vector in vectors.items():
raw = _tensor_bytes(vector)
vector_path = output_dir / f"{name}.f64le.bin"
_write_new(vector_path, raw)
vector_receipts[name] = {
**_tensor_summary(vector),
"path": vector_path.name,
"size_bytes": len(raw),
}
result = {
"attempt": {
"arm_label": args.arm_label,
"attempt_id": args.attempt_id,
"block_id": args.block_id,
"scientific_seed": 17,
},
"authority": (
"EXPERIMENTAL_STANDALONE_REPRO"
if args.runtime_profile == "cuda-certified"
else "CPU_SMOKE_NOT_BUG_EVIDENCE"
),
"candidate_id": CANDIDATE_ID,
"fixture": {
"duplicate_payload_pairs": [[0, 1], [2, 3], [4, 5], [6, 7]],
"global_token_mass": GLOBAL_TOKEN_MASS,
"selected_indices_by_rank": (
[indices]
if expected_world_size == 1
else [
_partition_indices(args.partition, expected_world_size, item)
for item in range(expected_world_size)
]
),
"variant": args.fixture_variant,
},
"grad_is_none": {
"nonrouter": nonrouter_none,
"router": router_none,
"total": total_none,
},
"imports_partitioncheck": any(
name == "partitioncheck" or name.startswith("partitioncheck.")
for name in sys.modules
),
"initial_state_sha256": initial_sha256,
"mode": args.mode,
"objective_components": {
"auxiliary": float(aux_component.item()),
"ce": float(ce_component.item()),
"total": float(objective.item()),
},
"partition": args.partition,
"rank_vector_sha256": rank_hashes,
"runtime_profile": args.runtime_profile,
"schema_version": "hq03.standalone-qwen-ddp-aux-repro.v1",
"source_receipt": source,
"vectors": vector_receipts,
"world_size": expected_world_size,
}
if result["imports_partitioncheck"] is not False:
raise ReproError("standalone reproducer imported PartitionCheck")
unsigned = dict(result)
unsigned["result_sha256"] = _sha256_bytes(_canonical_json(result))
_write_new(output_dir / "result.json", _canonical_json(unsigned))
return 0
finally:
if trainer is not None:
temporary = getattr(trainer, "_hq03_temporary_output", None)
if isinstance(temporary, str):
shutil.rmtree(temporary, ignore_errors=True)
if dist.is_initialized():
dist.destroy_process_group()
def _load_result(path: Path) -> dict[str, Any]:
raw = path.read_bytes()
value = json.loads(raw)
if (
not isinstance(value, dict)
or value.get("schema_version") != "hq03.standalone-qwen-ddp-aux-repro.v1"
):
raise ReproError(f"invalid standalone result: {path}")
address = value.get("result_sha256")
unsigned = {name: item for name, item in value.items() if name != "result_sha256"}
if not isinstance(address, str) or _sha256_bytes(_canonical_json(unsigned)) != address:
raise ReproError(f"standalone result content address differs: {path}")
return value
def _load_vector(directory: Path, result: dict[str, Any], name: str) -> Any:
import torch
vectors = result.get("vectors")
if not isinstance(vectors, dict) or not isinstance(vectors.get(name), dict):
raise ReproError(f"missing vector receipt: {name}")
receipt = vectors[name]
path_value = receipt.get("path")
if not isinstance(path_value, str) or Path(path_value).name != path_value:
raise ReproError(f"unsafe vector path: {name}")
raw = (directory / path_value).read_bytes()
if _sha256_bytes(raw) != receipt.get("sha256") or len(raw) != receipt.get("size_bytes"):
raise ReproError(f"vector bytes differ from receipt: {name}")
if len(raw) % 8:
raise ReproError(f"vector byte length is malformed: {name}")
return torch.frombuffer(bytearray(raw), dtype=torch.float64).clone()
def _comparison(left: Any, right: Any) -> dict[str, float]:
import torch
if left.shape != right.shape or left.numel() == 0:
raise ReproError("comparison vectors have different or empty shapes")
difference = left - right
absolute_l2 = float(torch.linalg.vector_norm(difference).item())
denominator = max(
float(torch.linalg.vector_norm(left).item()),
float(torch.linalg.vector_norm(right).item()),
sys.float_info.min,
)
values = {
"absolute_l2": absolute_l2,
"max_absolute": float(torch.max(torch.abs(difference)).item()),
"relative_l2": absolute_l2 / denominator,
}
if not all(math.isfinite(value) for value in values.values()):
raise ReproError("comparison produced a non-finite metric")
return values
def _compare(args: argparse.Namespace) -> int:
left_dir = Path(args.left)
right_dir = Path(args.right)
left = _load_result(left_dir / "result.json")
right = _load_result(right_dir / "result.json")
if left["candidate_id"] != CANDIDATE_ID or right["candidate_id"] != CANDIDATE_ID:
raise ReproError("comparison candidate identity differs")
comparisons = {
name: _comparison(
_load_vector(left_dir, left, name),
_load_vector(right_dir, right, name),
)
for name in VECTOR_NAMES
}
result: dict[str, object] = {
"candidate_id": CANDIDATE_ID,
"comparisons": comparisons,
"left": {
"mode": left["mode"],
"partition": left["partition"],
"result_sha256": left["result_sha256"],
},
"right": {
"mode": right["mode"],
"partition": right["partition"],
"result_sha256": right["result_sha256"],
},
"schema_version": "hq03.standalone-comparison.v1",
}
result["comparison_sha256"] = _sha256_bytes(_canonical_json(result))
_write_new(Path(args.out), _canonical_json(result))
return 0
def _parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description=__doc__)
commands = parser.add_subparsers(dest="command", required=True)
worker = commands.add_parser("worker")
worker.add_argument("--runtime-profile", choices=("cpu-smoke", "cuda-certified"), required=True)
worker.add_argument("--mode", choices=MODES, required=True)
worker.add_argument("--partition", choices=PARTITIONS, required=True)
worker.add_argument("--fixture-variant", type=int, required=True)
worker.add_argument("--attempt-id", required=True)
worker.add_argument("--block-id", required=True)
worker.add_argument("--arm-label", choices=("A", "B", "C", "BENIGN"), required=True)
worker.add_argument("--out", type=Path, required=True)
compare = commands.add_parser("compare")
compare.add_argument("--left", type=Path, required=True)
compare.add_argument("--right", type=Path, required=True)
compare.add_argument("--out", type=Path, required=True)
return parser
def main() -> int:
args = _parser().parse_args()
if args.command == "worker":
return _run_worker(args)
if args.command == "compare":
return _compare(args)
raise AssertionError("unreachable standalone command")
if __name__ == "__main__":
raise SystemExit(main())
Summary
With
average_tokens_across_devices=True,Trainerscales the entire scalar returned by Qwen3-MoE to compensate for global token normalization, including the router auxiliary loss. DDP already averages that auxiliary term across ranks, so its effective coefficient grows with the number of processes.Moving the same logical batch from one rank to two therefore changes the router gradient, even with matching routing distributions and an unchanged
router_aux_loss_coef.System Info
Original GPU environment: Python 3.11, PyTorch 2.13.0+cu130, FP32, Transformers 5.15.0, commit 5eddc12e, Qwen3-MoE, DDP.
The results below are from the pinned version. Source inspection on 2026-09-10 found the same calculation in 4815a0a6; the full experiment has not been rerun on that commit.
Information
Tasks
examplesfolderReproduction
The standalone script below uses fixture variant 17: eight sequences of length nine, arranged as four identical pairs, giving 64 valid next-token targets. It creates a small, randomly initialized
Qwen3MoeForCausalLMwithoutput_router_logits=Trueandrouter_aux_loss_coef=0.01. It runs in FP32 withaverage_tokens_across_devices=True,gradient_accumulation_steps=1, no label smoothing or custom loss, and a globalnum_items_in_batch=64. The model accepts the loss kwargs, soTrainer.compute_lossenters the cross-device compensation branch.The test compares one rank processing the full batch with two ranks processing matched local routing distributions. Model initialization, global token count, auxiliary coefficient, and the optimizer update are fixed. Matching the routing distributions makes the rank-local and full-batch auxiliary statistics agree, so loss scaling is the only remaining difference.
0.43621003.39e-73.52e-7These are the original pinned-version GPU results, not a new run on current
main. The local source patch and the component-wise repair both close the gap; the behavior was reproduced on two hosts.Run the standalone reproduction (one GPU versus two GPUs)
Save the Python script in the next section as
qwen_ddp_aux_repro.py, then run the following in the environment specified above. No pretrained model or dataset download is needed. The script checks the pinned Trainer and Qwen source hashes and requires a fresh output directory for each run.currentruns the native Trainer path.component_aware_repaircomputes the CE and auxiliary components separately.ce_onlysets the auxiliary coefficient to zero. The final three lines print the router-gradient relative L2 differences.For a CPU-only smoke check, use PyTorch 2.13.0+cpu and replace
cuda-certifiedwithcpu-smoke; the GPU visibility settings can be omitted. That is a separate smoke check, not the GPU result reported in the table.Complete standalone script: qwen_ddp_aux_repro.py
Expected behavior
For the same logical batch with matched routing distributions, changing the DDP world size should preserve the effective
router_aux_loss_coefand produce matching router gradients within numerical tolerance. World-size compensation should apply only to the globally token-normalized cross-entropy term, while the router auxiliary term keeps its configured coefficient.Root cause
Qwen3-MoE returns
token_loss + router_aux_loss_coef * aux_lossas one scalar, andTrainer.compute_lossapplies the data-parallel scale to that scalar. Writingce_lossfor the local CE sum divided by the global token count andauxfor the weighted auxiliary term, the current computation isThe CE term needs this factor because DDP averages gradients. The rank-local auxiliary term should keep its configured coefficient after that average. Scaling both terms multiplies the auxiliary contribution by
world_size.Proposed fix
Apply the compensation to the token-normalized component only, then add the weighted auxiliary loss back:
The tested patch extracts
router_aux_loss_coef * outputs.aux_loss, subtracts it before scaling, and adds it back afterward. It fixes the coefficient error shown above and does not change the scope of Qwen's routing statistics.Local trainer.py patch for the pinned commit
AI assistance was used to prepare this report and reproducer. The measurements above are from the pinned GPU experiments; this edit does not claim a new current-main or GPU test run.