fix(trainer-rank): keep retained checks paired with their layout

src/art/trainer_rank/_impl.py

@@ -3717,18 +3717,16 @@ def trusted(packed_tokens: int, signature: _MemorySignature) -> bool:
                         )
                         assert exact is not None
                         selected = exact
+                        # The minimum wave may execute cold after trust fails.
+                        # Keep its materialization paired with the retained check.
+                        layout_modes[width] = memory_minimal
                         if selected[0].fits and (
                             not profiled or trusted(selected[1], selected[3])
                         ):
-                            layout_modes[width] = memory_minimal
                             break
-                    else:
-                        # Nothing fit and trusted; keep the memory-minimal
-                        # pricing so the recorded failure is the monotone one.
-                        if not selected[0].fits:
-                            layout_modes.pop(width, None)
                 elif minimal_bound is not None:
                     selected = minimal_bound
+                    layout_modes[width] = True
                 if not selected[0].fits:
                     exact_failed_width = (
                         width

tests/unit/test_trainer_rank_weird_shapes.py

@@ -22,6 +22,7 @@
 )
 from art.trainer_rank._impl import (
     _CheckpointSlot,
+    _FlatForwardPlan,
     _flatten,
     _MemoryCheck,
     _MemoryProfile,
@@ -517,6 +518,104 @@ def test_dp_rank_forward_falls_back_to_memory_minimal_layout_before_refusing(
     assert rank.last_forward_telemetry()["selected_max_depth"] == 2
 
 
+@pytest.mark.parametrize(
+    "profile", ("absent", "tiny", "tiny_roomy", "trusted", "peer_untrusted")
+)
+def test_minimum_wave_materializes_the_layout_its_check_priced(
+    monkeypatch: pytest.MonkeyPatch, profile: str
+) -> None:
+    inputs = [_target_request(_tokens(*range(10_000, 10_040), tail)) for tail in (1, 2)]
+    rank = TrainerRank(_attention_runtime())
+    monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1))
+    unshared = rank._plan_flat_forward(inputs)
+    shared = rank._plan_flat_forward(inputs, memory_minimal=True)
+    assert (unshared.packed_tokens, shared.packed_tokens) == (82, 42)
+    if profile != "absent":
+        rank._memory_profiles[shared.signature] = _MemoryProfile(
+            bytes_per_token=0.0,
+            packed_tokens=1 if profile.startswith("tiny") else shared.packed_tokens,
+        )
+    if profile == "peer_untrusted":
+        # Local geometry is trusted, but another DP rank reports otherwise.
+        monkeypatch.setattr(rank, "_all_ranks_have_memory_profile", lambda **_: False)
+    shared_required = rank._memory_check(shared).estimated_required_bytes
+    unshared_required = rank._memory_check(unshared).estimated_required_bytes
+    assert shared_required < unshared_required
+    available = (
+        unshared_required
+        if profile == "tiny_roomy"
+        else (shared_required + unshared_required) // 2
+    )
+    monkeypatch.setattr(rank, "_available_memory_bytes", lambda: available)
+
+    candidate = rank._select_next_micro_batch([inputs], 0)
+
+    assert candidate.check.fits
+    assert isinstance(candidate.plan, _FlatForwardPlan)
+    assert candidate.plan.packed_tokens == shared.packed_tokens
+    assert rank._memory_check(candidate.plan) == candidate.check
+    assert candidate.cold_start == (profile != "trusted")
+
+
+@pytest.mark.parametrize("stage", ("bound", "exact"))
+def test_minimum_wave_rejection_keeps_its_layout_mode(
+    monkeypatch: pytest.MonkeyPatch, stage: str
+) -> None:
+    inputs = [_target_request(_tokens(*range(10_000, 10_040), tail)) for tail in (1, 2)]
+    rank = TrainerRank(_attention_runtime())
+    monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1))
+    shared = rank._plan_flat_forward(inputs, memory_minimal=True)
+    required = rank._memory_check(shared).estimated_required_bytes
+    available = iter(
+        (required - 1, required + 1, required - 1, required - 1)
+        if stage == "exact"
+        else (required - 1, required - 1)
+    )
+    monkeypatch.setattr(rank, "_available_memory_bytes", lambda: next(available))
+    modes: list[bool] = []
+    original_plan = rank._plan_flat_forward
+
+    def plan(requests, **kwargs):
+        modes.append(kwargs.get("memory_minimal", False))
+        return original_plan(requests, **kwargs)
+
+    original_error = RuntimeError("The split ladder owns the subsequent refusal")
+
+    def find(*args, **kwargs):
+        raise original_error
+
+    monkeypatch.setattr(rank, "_plan_flat_forward", plan)
+    monkeypatch.setattr(rank, "_find_admissible_forward", find)
+    with pytest.raises(RuntimeError) as caught:
+        rank._select_next_micro_batch([inputs], 0)
+    assert caught.value is original_error
+    assert modes == [True]
+    assert next(available, None) is None
+
+
+def test_minimum_wave_empty_dp_rank_keeps_collective_check_sequence(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    rank = TrainerRank(_attention_runtime())
+    monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (1, 2))
+    peer_required = iter((82, 42, 82, 42))
+    local_checks: list[tuple[int, bool]] = []
+
+    def check(required: int, *, sync_across_dp: bool = False) -> _MemoryCheck:
+        local_checks.append((required, sync_across_dp))
+        remote = next(peer_required)
+        return _MemoryCheck(remote, 60, remote <= 60)
+
+    monkeypatch.setattr(rank, "_memory_check_required", check)
+    candidate = rank._select_next_micro_batch([[_target_request(_tokens(1, 2))]], 0)
+
+    assert candidate.indices == ()
+    assert candidate.plan.packed_tokens == 0
+    assert candidate.check == _MemoryCheck(42, 60, True)
+    assert local_checks == [(0, True)] * 4
+    assert next(peer_required, None) is None
+
+
 def test_profiled_steady_state_keeps_the_wide_shared_wave(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None: