mixture optimization batches

reports/git-phase2-phase3-submit.md

@@ -0,0 +1,60 @@
+# Git Phase 2/3 提交推送核验报告
+
+## 核验结论
+
+截至 2026-08-02，**尚未完成 push**。执行 `git fetch --prune` 后，本地与 upstream 已分叉：本地领先 3 个提交、落后 32 个提交。远端分支不包含本地 `HEAD`。
+
+## 分支与远端
+
+- 当前分支：`codex/mixture-optimization-batches`
+- Upstream：`origin/codex/mixture-optimization-batches`
+- Ahead/behind（`HEAD...@{upstream}`）：`3 32`
+- 远端包含本地 HEAD：否
+- Fetch：成功；`origin/codex/mixture-optimization-batches` 从 `31aa643` 更新至 `0818eff`
+
+## 尚未推送的提交
+
+| Commit | Subject |
+|---|---|
+| `9f581948a7061d180e182785cc1df2246e46e118` | `fix: clarify sparse routing diagnostics` |
+| `11f542396af46551a3897833c5b39b91d7956e1c` | `fix: harden mixture routing correctness` |
+| `31be313e379dc9c05a65a5ee9aea6830933459fd` | `fix: stabilize MoA layouts and MPS numerics` |
+
+## Push 状态
+
+- 本轮仅按要求执行最终核验，未执行 `git push`。
+- 当前普通 push 将因远端领先 32 个提交而成为 non-fast-forward；需要先安全整合远端变更。
+- 工作区同时存在未提交源码改动，因此未擅自 rebase、merge、stash、reset 或创建提交。
+
+## `git status -sb`
+
+```text
+## codex/mixture-optimization-batches...origin/codex/mixture-optimization-batches [ahead 3, behind 32]
+ M tests/test_config_drift_detector.py
+ M tests/test_latent_mixture.py
+ M tests/test_molora_routing_aware_merge.py
+ M tests/test_molora_sparse_dispatch.py
+ M tests/test_peft_adapters.py
+ M tests/test_placement_plan_schema.py
+ M tests/test_vpeft_lora_e2e.py
+ M tools/config_drift_detector.py
+ M ultralytics/cfg/default.yaml
+ M ultralytics/nn/modules/latent_mixture.py
+ M ultralytics/nn/modules/moe/__init__.py
+ M ultralytics/nn/modules/moe/protocol.py
+ M ultralytics/nn/peft/molora/layer.py
+ M ultralytics/utils/lora/api.py
+ M ultralytics/utils/lora/config.py
+ M ultralytics/vpeft/__init__.py
+ M ultralytics/vpeft/placement_plan.py
+?? .session_tmps/
+?? .workbuddy/
+?? agent/logs/
+?? output/
+?? outputs/
+?? tests/test_moe_facade_protocol.py
+```
+
+## 建议后续动作
+
+在保留未提交用户改动的前提下，先完成或隔离剩余逻辑批次，然后将远端 32 个提交安全整合到当前分支，解决潜在冲突并重跑相关测试；确认 ahead/behind 后再普通 push。禁止 force push。

scripts/bench_moe_mps.py

@@ -34,6 +34,18 @@
 ]
 
 
+def mps_device_info():
+    """Return optional MPS metadata without assuming backend APIs exist."""
+    backend = getattr(torch.backends, "mps", None)
+    if backend is None:
+        return None, None
+    get_name = getattr(backend, "get_name", None)
+    get_core_count = getattr(backend, "get_core_count", None)
+    name = get_name() if callable(get_name) else None
+    core_count = get_core_count() if callable(get_core_count) else None
+    return name, core_count
+
+
 def pick_device():
     if torch.backends.mps.is_available():
         return torch.device("mps")
@@ -84,8 +96,13 @@ def main():
     dev = pick_device()
     print(f"[bench] device = {dev}, torch = {torch.__version__}")
     if dev.type == "mps":
-        print(f"[bench] mps name = {torch.backends.mps.get_name()}, "
-              f"core count = {torch.backends.mps.get_core_count()}")
+        name, core_count = mps_device_info()
+        details = []
+        if name is not None:
+            details.append(f"name = {name}")
+        if core_count is not None:
+            details.append(f"core count = {core_count}")
+        print(f"[bench] mps {', '.join(details) if details else 'metadata unavailable'}")
     print(f"[bench] bs={args.bs} imgsz={args.imgsz} runs={args.runs} warmup={args.warmup} dtype={args.dtype}")
     print("-" * 90)
     print(f"{'config':<30}{'params(M)':>11}{'median':>9}{'mean':>9}{'min':>9}{'p95':>9}{'fps':>8}")

tests/test_config_drift_detector.py

@@ -170,6 +170,27 @@ def test_incompatible_model_argument_counts_are_reported(tmp_path):
     assert any(issue.code == "MODEL_ARGS_INCOMPATIBLE" and "Known" in issue.message for issue in issues)
 
 
+def test_check_all_includes_yolo26_master_configs(tmp_path):
+    detector = ConfigDriftDetector(tmp_path)
+    _write(tmp_path / detector.DEFAULT_CFG, "alpha: 1\n")
+    _write(tmp_path / detector.YOLO26_MASTER_MODELS / "yolo26-master-test.yaml", "alpha: 1\nalpha: 2\n")
+
+    issues = detector.check_all()
+
+    assert any(
+        issue.code == "YAML_DUPLICATE_KEY" and issue.path.name == "yolo26-master-test.yaml" for issue in issues
+    )
+
+
+def test_lora_rank_pattern_cli_mapping_round_trip():
+    from ultralytics.utils.lora.config import LoRAConfig
+
+    rank_pattern = {"model.0.conv": 4, "model.1.conv": 8}
+    config = LoRAConfig.from_args(lora_rank_pattern=rank_pattern)
+
+    assert config.rank_pattern == rank_pattern
+
+
 def test_repository_has_no_configuration_drift():
     issues = ConfigDriftDetector(ROOT).check_all()
 

tests/test_ddp_lifecycle_ema_nan.py

@@ -440,6 +440,26 @@ def test_checkpoint_restore_tolerates_missing_lazy_ema_buffer():
     )
 
 
+def test_checkpoint_restore_migrates_legacy_three_slot_ema_state():
+    t = object.__new__(BaseTrainer)
+    t.model = nn.Linear(1, 1)
+    t.model.register_buffer("_mixture_loss_ema_buf", torch.tensor([1.0, 0.1, 0.1, 0.1]))
+    t.ema = ModelEMA(t.model)
+    t.optimizer = torch.optim.SGD(t.model.parameters(), lr=0.01)
+    t.scaler = torch.amp.GradScaler("cuda", enabled=False)
+    old_ema = nn.Linear(1, 1)
+    old_ema.register_buffer("_mixture_loss_ema_buf", torch.tensor([2.0, 0.2, 0.3]))
+
+    t._load_checkpoint_state(
+        {"ema": old_ema, "optimizer": None, "scaler": None, "best_fitness": 0.0, "updates": 7}
+    )
+
+    expected = torch.tensor([2.0, 0.2, 0.3, 0.1])
+    assert torch.equal(t.model._mixture_loss_ema_buf, expected)
+    assert torch.equal(t.ema.ema._mixture_loss_ema_buf, expected)
+    assert t.ema.updates == 7
+
+
 def test_healthy_checkpoint_rejects_nonfinite_state_and_preserves_prior(tmp_path):
     t = object.__new__(BaseTrainer)
     t.healthy = tmp_path / "last_healthy.pt"

tests/test_latent_mixture.py

@@ -101,6 +101,13 @@ def test_latent_mixture_publishes_single_train_aux_and_snapshot():
     assert records[0][0] is module
     snapshot = module.routing_snapshot()
     assert snapshot["family"] == "latent"
+    assert snapshot["top_k"] == snapshot["training_top_k"] == 4
+    assert snapshot["inference_top_k"] == 4
+    assert snapshot["executed_experts"] == 4
+    assert snapshot["configured_top_k"] == 4
+    assert snapshot["active_experts_per_sample"].tolist() == [4, 4]
+    assert snapshot["batch_expert_union"] == 4
+    assert snapshot["kernel_calls"] == 4
     assert snapshot["noise_std"] == 0.0
     assert snapshot["router_init_std"] == 0.0
     assert snapshot["identity_cold_start"] is False
@@ -238,3 +245,55 @@ def test_latent_temperature_anneal_updates_latent_modules():
 
     assert updated == 1
     assert float(module.temperature) == pytest.approx(0.475)
+
+
+def test_latent_value_fusion_has_explicit_causal_semantics():
+    torch.manual_seed(7)
+    first = torch.randn(2, 8, 4, 4)
+    auxiliary = torch.randn(2, 8, 4, 4)
+    changed = auxiliary.flip(0)
+
+    legacy = LatentMixture([8, 8], 8, residual_init=0.0).eval()
+    fused = LatentMixture(
+        [8, 8], 8, residual_init=0.0, value_fusion_mode="weighted_sum", value_fusion_weights=[1.0, 1.0]
+    ).eval()
+    fused.load_state_dict(legacy.state_dict(), strict=False)
+
+    with patch.object(legacy.router, "forward", return_value=(torch.zeros(2, 4), torch.full((2, 4), 0.25))):
+        legacy_a = legacy([first, auxiliary])
+        legacy_b = legacy([first, changed])
+    with patch.object(fused.router, "forward", return_value=(torch.zeros(2, 4), torch.full((2, 4), 0.25))):
+        fused_a = fused([first, auxiliary])
+        fused_b = fused([first, changed])
+
+    assert torch.allclose(legacy_a, legacy_b)
+    assert not torch.allclose(fused_a, fused_b)
+    assert fused.routing_snapshot()["value_fusion_mode"] == "weighted_sum"
+
+
+def test_latent_value_fusion_ablations_are_reproducible():
+    torch.manual_seed(8)
+    xs = [torch.randn(2, 8, 4, 4), torch.randn(2, 8, 4, 4)]
+    shuffled = [xs[0], xs[1].flip(0)]
+    first_only = LatentMixture(
+        [8, 8], 8, residual_init=0.0, value_fusion_mode="weighted_sum", value_fusion_weights=[1.0, 0.0]
+    ).eval()
+    equal = LatentMixture(
+        [8, 8], 8, residual_init=0.0, value_fusion_mode="weighted_sum", value_fusion_weights=[1.0, 1.0]
+    ).eval()
+    equal.load_state_dict(first_only.state_dict(), strict=False)
+
+    assert torch.allclose(first_only(xs), first_only(shuffled))
+    assert not torch.allclose(equal(xs), equal(shuffled))
+    assert torch.allclose(first_only.value_fusion_weights, torch.tensor([1.0, 0.0]))
+
+
+def test_latent_value_fusion_is_opt_in_and_checkpoint_compatible():
+    default = LatentMixture([8, 8], 8)
+    restored = LatentMixture([8, 8], 8)
+    restored.load_state_dict(default.state_dict(), strict=True)
+
+    assert default.value_fusion_mode == "router_only"
+    assert "value_fusion_weights" not in default.state_dict()
+    with pytest.raises(ValueError, match="value_fusion_mode"):
+        LatentMixture([8, 8], 8, value_fusion_mode="unknown")

tests/test_mixture_numeric.py

@@ -5,9 +5,11 @@
 import pytest
 import torch
 
+from ultralytics.nn.modules import LatentRouter
 from ultralytics.nn.modules._numeric import (
     all_reduce_mean,
     clamp_min_for_dtype,
+    disabled_autocast,
     fp_clamp_floor,
     stable_normalize,
 )
@@ -53,6 +55,27 @@ def test_all_mixture_namespaces_share_canonical_all_reduce_mean():
     assert mot_all_reduce_mean is all_reduce_mean
 
 
+def test_disabled_autocast_falls_back_when_device_is_unsupported():
+    with patch("torch.amp.autocast_mode.is_autocast_available", return_value=False), patch(
+        "torch.autocast", side_effect=AssertionError("unsupported autocast must not be constructed")
+    ):
+        with disabled_autocast("mps"):
+            result = torch.ones(1) + 1
+
+    assert result.item() == 2
+
+
+def test_disabled_autocast_uses_supported_cpu_context():
+    with patch("torch.amp.autocast_mode.is_autocast_available", return_value=True), patch(
+        "torch.autocast", wraps=torch.autocast
+    ) as autocast:
+        with disabled_autocast("cpu"):
+            result = torch.ones(1) + 1
+
+    autocast.assert_called_once_with(device_type="cpu", enabled=False)
+    assert result.item() == 2
+
+
 def test_all_reduce_mean_keeps_global_value_and_local_gradient():
     local = torch.tensor([0.2, 0.8], requires_grad=True)
 
@@ -102,6 +125,19 @@ def test_mot_router_keeps_fp32_logits_and_activation_dtype_weights():
     assert torch.allclose(weights.float().sum(dim=1), torch.ones_like(weights[:, 0].float()), atol=2e-3)
 
 
+def test_latent_router_keeps_fp32_routing_contract_after_dtype_conversion():
+    router = LatentRouter(8, num_experts=3).eval().half()
+    tokens = torch.randn(2, 8, dtype=torch.bfloat16)
+
+    logits, probs = router(tokens)
+
+    assert logits.dtype == torch.float32
+    assert probs.dtype == torch.float32
+    assert {parameter.dtype for parameter in router.parameters()} == {torch.float32}
+    assert torch.isfinite(probs).all()
+    assert torch.allclose(probs.sum(dim=-1), torch.ones_like(probs[:, 0]), atol=1e-6)
+
+
 @pytest.mark.parametrize("router_cls", [DualStreamGateRouter, ZeroCostRouter])
 def test_gated_router_keeps_fp32_aux_statistics_and_activation_dtype_weights(router_cls):
     router = router_cls(8, num_experts=3, top_k=2).train().half()

tests/test_moa.py

@@ -16,6 +16,7 @@
     _moa_router_aux_loss,
     _window_flash_attn,
 )
+from ultralytics.nn.modules.moa.heads import _window_partition_2d, _window_unpartition_2d
 from ultralytics.nn.modules._numeric import fp_clamp_floor as _fp_min
 from ultralytics.nn.tasks import DetectionModel
 from ultralytics.utils.loss import _collect_moa_aux_loss
@@ -412,22 +413,46 @@ def test_regional_attn_head_non_divisible_dim_and_heads():
 
 
 def test_regional_attn_head_small_spatial_dims():
-    """_RegionalAttnHead gracefully handles H=1 or W=1 feature maps."""
+    """_RegionalAttnHead keeps canonical memory format and gradients for singleton axes."""
     torch.manual_seed(0)
-    cases = [
-        (_RegionalAttnHead(32, num_heads=4), torch.randn(1, 32, 1, 16)),
-        (_RegionalAttnHead(32, num_heads=4), torch.randn(1, 32, 16, 1)),
-        (_RegionalAttnHead(32, num_heads=4), torch.randn(1, 32, 1, 1)),
-    ]
-    for module, x in cases:
-        module.train()
+    for height, width in ((1, 16), (16, 1), (1, 1)):
+        module = _RegionalAttnHead(32, num_heads=4).train()
+        x = torch.randn(2, 32, height, width, requires_grad=True)
         out = module(x)
         assert out.shape == x.shape
+        assert out.is_contiguous()
         assert torch.isfinite(out).all()
-        out.mean().backward()
+        out.square().mean().backward()
+        assert x.grad is not None and torch.isfinite(x.grad).all()
         assert _has_grad(module)
 
 
+def test_window_partition_unpartition_round_trip():
+    """Window helpers preserve row-major [H,W] layout, including multiple heads."""
+    spatial = torch.arange(2 * 3 * 4 * 6 * 5).reshape(2, 3, 4, 6, 5)
+    windows = _window_partition_2d(spatial, 2)
+    restored = _window_unpartition_2d(windows, 2, 2, 3, 4, 6)
+    assert torch.equal(restored, spatial)
+
+
+def test_window_flash_attn_matches_spatial_reference():
+    """Window attention matches an explicit per-window numerical reference."""
+    torch.manual_seed(7)
+    B, nh, H, W, hd, win = 2, 2, 4, 6, 3, 2
+    q, k, v = (torch.randn(B, nh, H * W, hd) for _ in range(3))
+    scale = hd ** -0.5
+
+    actual = _window_flash_attn(q, k, v, scale, win, H, W)
+    expected = torch.empty_like(actual.reshape(B, nh, H, W, hd))
+    for row in range(0, H, win):
+        for col in range(0, W, win):
+            indices = [(r * W + c) for r in range(row, row + win) for c in range(col, col + win)]
+            reference = _flash_attn(q[:, :, indices], k[:, :, indices], v[:, :, indices], scale)
+            expected[:, :, row : row + win, col : col + win] = reference.reshape(B, nh, win, win, hd)
+
+    assert torch.allclose(actual, expected.reshape(B, nh, H * W, hd), atol=1e-6, rtol=1e-5)
+
+
 def test_regional_attn_head_invalid_pool_stride():
     """_RegionalAttnHead raises ValueError for pool_stride < 1."""
     with pytest.raises(ValueError, match="pool_stride"):

tests/test_moe_facade_protocol.py

@@ -0,0 +1,39 @@
+"""MoE facade tiers, compatibility exports, and snapshot protocol regression tests."""
+
+import pytest
+import torch
+
+import ultralytics.nn.modules.moe as facade
+from ultralytics.nn.modules.moe.protocol import normalize_routing_snapshot, routing_metrics
+from ultralytics.nn.modules.moe.schedule import usage_gini as schedule_gini
+from ultralytics.nn.modules.moe.scheduler import compute_gini
+
+
+def test_public_api_tiers_are_disjoint_and_resolvable():
+    tiers = (facade.STABLE_MOE_CLASSES, facade.EXPERIMENTAL_MOE_CLASSES, facade.LEGACY_MOE_CLASSES)
+    assert not (tiers[0] & tiers[1] or tiers[0] & tiers[2] or tiers[1] & tiers[2])
+    for name in set().union(*tiers):
+        assert name in facade.__all__
+        assert getattr(facade, name) is not None
+
+
+def test_legacy_checkpoint_aliases_remain_importable():
+    assert facade.A2C2fMoE.__name__ == "A2C2fMoE"
+    assert facade.ABlockMoE.__name__ == "ABlockMoE"
+    assert facade.is_legacy_moe("A2C2fMoE")
+
+
+def test_snapshot_adapter_accepts_legacy_keys_without_mutation():
+    source = {"usage": torch.tensor([0.75, 0.25]), "counts": [3, 1], "top_k": 1}
+    normalized = normalize_routing_snapshot(source, num_experts=2)
+    assert normalized["expert_usage"] == [0.75, 0.25]
+    assert normalized["topk_counts"] == [3.0, 1.0]
+    assert "expert_usage" not in source
+
+
+def test_scheduler_and_diagnostics_share_gini_definition():
+    usage = torch.tensor([1.0, 0.0, 0.0, 0.0])
+    metrics = routing_metrics({"expert_usage": usage}, num_experts=4)
+    assert metrics.gini == pytest.approx(0.75)
+    assert compute_gini(usage) == pytest.approx(metrics.gini)
+    assert schedule_gini(usage) == pytest.approx(metrics.gini)