diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 4ae42575f..736bc951e 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1500,7 +1500,8 @@ def _moe_output_bytes_per_token( slot_ref: "LoRASlotRef | None" = None, ) -> int: """Known routed-expert working set, not a complete model/compiled bound.""" - if shape != ParallelShape(tp=1, cp=1): + # CP shards rows, not the per-token working set; EP dispatch is not modeled. + if (shape.tp, shape.ep, shape.etp) != (1, 1, 1): return 0 from megatron.core.extensions.transformer_engine import ( TEColumnParallelGroupedLinear, @@ -3461,7 +3462,9 @@ def _split_chunk_lower_cost( assert estimated is not None # rows are CPU copies physical_rows = self._physical_tokens(estimated) packed_tokens += physical_rows - group_rows.append((physical_rows, grad_enabled)) + # The most loaded CP rank holds at least an even share. + cp = max(1, self._topology_key()[2]) + group_rows.append((-(-physical_rows // cp), grad_enabled)) head_requests = tuple(requests[index] for index in group_indices) head_workspace_bytes = max( head_workspace_bytes, @@ -3536,7 +3539,7 @@ def _head_workspace_bytes(self, rows: int) -> int: rows <= 0 or self._padded_vocab_size is None or len(self.runtime.model) != 1 - or self._topology_key()[1:] != (1, 1, 1) + or self._topology_key()[1::2] != (1, 1) ): return 0 try: @@ -3776,9 +3779,21 @@ def _plan_head_workspace_bytes(self, plan: _FlatForwardPlan) -> int: return peak def _plan_group_rows(self, plan: _FlatForwardPlan) -> tuple[tuple[int, bool], ...]: + """Physical rows per group on the most loaded context-parallel rank.""" + topology = self._topology() if plan.signature.topology[2] > 1 else None return tuple( ( - self._physical_tokens(int(group.packed.tokens.numel())), + self._physical_tokens( + int(group.packed.tokens.numel()) + if topology is None + else max( + 1, + self._max_rank_model_tokens( + _pad_packed_batch(group.packed, multiple=int(topology.tp)), + topology=topology, + ), + ) + ), group.grad_enabled, ) for group in plan.groups @@ -3902,8 +3917,7 @@ def _checkpoint_memory_floor( or config.params_dtype is not torch.bfloat16 or self._param_dtype_size != 2 or next(self.runtime.model[0].parameters()).dtype is not torch.bfloat16 - or self._topology_key()[1:] != (1, 1, 1) - or _expert_parallel_shape(self.runtime.provider) != (1, 1) + or self._topology_key()[1::2] != (1, 1) or any( type(getattr(config, name, None)) is not type(value) or getattr(config, name) != value @@ -5106,25 +5120,32 @@ def fits(width: int) -> tuple[bool, bool]: width = normalize(width) result = estimate(width) if result is None: - # Estimator unavailable (device inputs): admit on the - # materialized plan, trying the cost-optimal layouts first and - # the memory-minimal layouts if those do not fit. + # Estimator unavailable (device inputs, or CP per-rank floors): + # admit on the materialized plan, trying the cost-optimal + # layouts first and the memory-minimal layouts if those do not + # fit or fall outside the profile's trust window. + def price(plan: _FlatForwardPlan) -> tuple[_MemoryCheck, bool, bool]: + check = self._memory_check( + plan, sync_across_dp=True, sync_planning_errors=True + ) + trusted = self._all_ranks_have_memory_profile( + packed_tokens=plan.packed_tokens, + signature=plan.signature, + ) + profiled = self._all_ranks_true( + plan.signature in self._memory_profiles + ) + return check, trusted, profiled + plan = materialize(width) - check = self._memory_check( - plan, sync_across_dp=True, sync_planning_errors=True - ) - if not check.fits and not layout_modes.get(width, False): + check, trusted, profiled = price(plan) + if ( + not check.fits or (profiled and not trusted) + ) and not layout_modes.get(width, False): layout_modes[width] = True plans.pop(width, None) plan = materialize(width) - check = self._memory_check( - plan, sync_across_dp=True, sync_planning_errors=True - ) - trusted = self._all_ranks_have_memory_profile( - packed_tokens=plan.packed_tokens, - signature=plan.signature, - ) - profiled = self._all_ranks_true(plan.signature in self._memory_profiles) + check, trusted, profiled = price(plan) else: check, trusted, profiled = result if width in plans: @@ -5197,6 +5218,19 @@ def candidate(width: int) -> _CandidateMicroBatch[ForwardInputsT]: first_estimate = estimate(min_width) if first_estimate is None or not (first_estimate[0].fits and first_estimate[1]): first = candidate(min_width) + if ( + first_estimate is None + and first.check.fits + and first.cold_start + and not layout_modes.get(min_width, False) + and self._all_ranks_true(first.plan.signature in self._memory_profiles) + ): + # Materialized pricing (DP-uniform): a profiled cost-optimal + # layout outside trust; full sharing may be trusted, as the + # estimator path would find. + layout_modes[min_width] = True + plans.pop(min_width, None) + first = candidate(min_width) if not first.check.fits: # The smallest wave cannot run unsplit: best effort is the # bounded split ladder. Each DP rank runs it on its own share, @@ -5715,6 +5749,7 @@ def _estimate_flat_forward( whose feasibility is monotone in width (valid for rejecting one). ``exact=True`` prices the planner's actual layouts (memoized by content) and is used only inside the band where those bounds disagree. + Under CP it returns None: per-rank floors need materialized layouts. """ if sync_planning_errors: @@ -5738,14 +5773,10 @@ def _estimate_flat_forward( # Pending saves require the actual bucket/replayed-tail geometry. # Existing unavailable handling materializes before admission. return None - if ( - self._topology_key()[2] > 1 - and self._recompute_granularity != "full" - and not self._geometry.moe_experts - and any(grad for (_, grad), _ in groups) - ): - # CP token ownership can be uneven. Use the existing exact-plan - # fallback; a global token count alone cannot price its peak. + if self._topology_key()[2] > 1: + # CP token ownership can be uneven and memory floors are priced + # per rank. Use the existing exact-plan fallback; a global token + # count alone cannot price its peak. return None packed_tokens = 0 head_workspace_bytes = 0 @@ -7724,7 +7755,12 @@ def _estimate_required_memory_bytes_from_values( static_compute = max( static_compute, *( - self._moe_workspace_bytes(packed_tokens, slot_ref=ref) + self._moe_workspace_bytes( + sum(rows for rows, _ in group_rows) + if signature.topology[2] > 1 and group_rows + else packed_tokens, + slot_ref=ref, + ) for ref in (slot_refs or (None,)) ), ) diff --git a/tests/unit/test_trainer_rank_admission_inputs.py b/tests/unit/test_trainer_rank_admission_inputs.py index 92b5dbde4..a79d198e1 100644 --- a/tests/unit/test_trainer_rank_admission_inputs.py +++ b/tests/unit/test_trainer_rank_admission_inputs.py @@ -93,7 +93,8 @@ def test_cp_gdn_segments_groups_and_retained_tokens_reach_exact_search(monkeypat assert _gdn_memory.model_shapes(rank) is None # This exercises the CP fallback. plan = rank._plan_flat_forward(requests) assert plan.grad_segment_count == 2 - assert rank._plan_group_rows(plan) == ((128, True), (80, False)) + # Rows on the most loaded CP rank (mocked at 3/4 of each group). + assert rank._plan_group_rows(plan) == ((96, True), (60, False)) assert rank._plan_retained_tokens(plan) == 156 for exact in (False, True): for memory_minimal in (False, True): @@ -108,7 +109,7 @@ def test_cp_gdn_segments_groups_and_retained_tokens_reach_exact_search(monkeypat requests, tuple(item.input_tokens for item in requests), checkpoint=Unset ) assert len(calls) == 1 - assert calls[0]["group_rows"] == ((128, True), (80, False)) + assert calls[0]["group_rows"] == ((32, True), (20, False)) # Even CP share. assert calls[0]["retained_tokens"] == 52 # Optimistic CP average only here. calls.clear() cost = rank._plan_cost(plan) diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index 7a5e7d952..c6e33ee9d 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -181,7 +181,7 @@ def test_actual_config_revalidated(field, value): assert r._checkpoint_memory_floor(((10, True),)) == (0, 0) -@pytest.mark.parametrize("axis", [1, 2, 3]) +@pytest.mark.parametrize("axis", [1, 3]) def test_topology_revalidated(axis): r = rank() topology = [1, 1, 1, 1] @@ -190,6 +190,76 @@ def test_topology_revalidated(axis): assert r._checkpoint_memory_floor(((10, True),)) == (0, 0) +@pytest.mark.parametrize("rows", [(10, True), (11, False)]) +@pytest.mark.parametrize("cp", [2, 4]) +def test_cp_floor_prices_rank_rows(cp, rows): + # Callers pass rows on the most loaded CP rank; the per-row floor matches CP1. + r = rank() + single = r._checkpoint_memory_floor((rows,)) + r._topology_key = lambda: (1, 1, cp, 1) + assert r._checkpoint_memory_floor((rows,)) == single != (0, 0) + + +@pytest.mark.parametrize("share", [lambda n: -(-n // 2), lambda n: n * 3 // 4]) +def test_cp_probe_defers_and_lower_bound_stays_below_exact(monkeypatch, share): + r = rank() + monkeypatch.setattr(r, "_topology_key", lambda: (1, 1, 2, 1)) + monkeypatch.setattr(r, "_topology", lambda: SimpleNamespace(cp=2, tp=1)) + # Even (tight) and uneven ownership of each group by the busiest CP rank. + monkeypatch.setattr( + r, "_max_rank_model_tokens", lambda batch, **_: share(batch.tokens.numel()) + ) + reqs = requests() + # Global counts would price the per-rank floors about cp times too high. + assert r._estimate_flat_forward(reqs) is None + plan = r._plan_flat_forward(reqs) + assert r._checkpoint_memory_floor(r._plan_group_rows(plan))[0] > 0 + exact = r._plan_cost(plan) + lower = r._split_chunk_lower_cost( + reqs, [q.input_tokens for q in reqs], checkpoint=Unset + ) + assert lower.required <= exact.required and lower.retained <= exact.retained + + +def test_cp_width_search_retries_full_sharing_inside_the_trust_window(monkeypatch): + r = rank() + r._dp_rank_and_size = lambda: (0, 1) + monkeypatch.setattr(r, "_topology_key", lambda: (1, 1, 2, 1)) + monkeypatch.setattr(r, "_topology", lambda: SimpleNamespace(cp=2, tp=1)) + monkeypatch.setattr( + r, "_max_rank_model_tokens", lambda batch, **_: batch.tokens.numel() * 3 // 4 + ) + r._available_memory_bytes = lambda: 1 << 50 + # Cost-optimal layouts stay unshared; memory-minimal layouts share fully. + monkeypatch.setattr( + r, + "_layout_anchor", + lambda *, memory_minimal=False: ( + "full_sharing" if memory_minimal else "no_sharing" + ), + ) + items = [ + ForwardInput( + input_tokens=torch.tensor([7, 100 + i]), hidden_states=True, no_grad=True + ) + for i in range(16) + ] + signature = r._plan_flat_forward(items[:1]).signature + r._memory_profiles[signature] = _MemoryProfile(bytes_per_token=1, packed_tokens=2) + # Unshared layouts leave the 8x window above width 8; full sharing + # (width + 1 packed rows) stays trusted through width 15. + selected = r._search_next_micro_batch(items, 0) + assert not isinstance(selected, _ForwardRefusal) + assert selected.stats_global_count == 15 + # The minimum wave retries too: one item's unshared layout (32 rows) is + # outside a 3-row profile's window; full sharing (17 rows) is inside it. + r._memory_profiles[signature] = _MemoryProfile(bytes_per_token=1, packed_tokens=3) + r._last_global_micro_batch_size = None + selected = r._search_next_micro_batch([items], 0) + assert not isinstance(selected, _ForwardRefusal) + assert selected.plan.packed_tokens == 17 and not selected.cold_start + + def test_dp_empty_and_local_count(): r = rank() r._topology_key = lambda: (3, 1, 1, 1) @@ -427,10 +497,3 @@ def test_no_grad_enclosure_config_guard(field, value): r = rank() setattr(r.runtime.model[0].decoder.config, field, value) assert r._checkpoint_memory_floor(((11, False),)) == (0, 0) - - -@pytest.mark.parametrize("cp", [2, 4]) -def test_no_grad_enclosure_keeps_cp_fallback(cp): - r = rank() - r._topology_key = lambda: (1, 1, cp, 1) - assert r._checkpoint_memory_floor(((11, False),)) == (0, 0) diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py index ccb1ed6a2..6b160e1c1 100644 --- a/tests/unit/test_trainer_rank_moe_memory.py +++ b/tests/unit/test_trainer_rank_moe_memory.py @@ -174,11 +174,12 @@ def test_unknown_fc2_input_keeps_previous_pair(layer, mode): @pytest.mark.parametrize("field", ["tp", "cp", "ep", "etp"]) def test_sharded_path_unchanged(layer, field): - assert ( - _moe_output_bytes_per_token( - [layer], replace(ParallelShape(tp=1, cp=1), **{field: 2}) - ) - == 0 + # CP shards rows, not the per-token working set; TP/EP/ETP are unmodeled. + single = _moe_output_bytes_per_token([layer], ParallelShape(tp=1, cp=1)) + sharded = replace(ParallelShape(tp=1, cp=1), **{field: 2}) + assert single > 0 + assert _moe_output_bytes_per_token([layer], sharded) == ( + single if field == "cp" else 0 ) diff --git a/tests/unit/test_trainer_rank_planning_status.py b/tests/unit/test_trainer_rank_planning_status.py index 8b290a4a2..b684aadce 100644 --- a/tests/unit/test_trainer_rank_planning_status.py +++ b/tests/unit/test_trainer_rank_planning_status.py @@ -112,7 +112,10 @@ def _worker(index: int, directory: Path) -> None: rank._dp_rank_and_size = lambda: (index, 2) rank._physical_tokens = lambda tokens: tokens rank._estimate_group_request_output_bytes = lambda requests: 0 - rank._memory_signature_from_requests = lambda *args, **kwargs: None + # Group rows read the plan's CP size; DP2/TP1/CP1/PP1 prices packed rows. + rank._memory_signature_from_requests = lambda *args, **kwargs: ( + SimpleNamespace(topology=(2, 1, 1, 1)) + ) rank._forward_item = lambda request: SimpleNamespace( input_ids=request.input_tokens, request=request ) diff --git a/tests/unit/test_trainer_rank_recompute_memory.py b/tests/unit/test_trainer_rank_recompute_memory.py index ea145e4d4..1d6eca3cb 100644 --- a/tests/unit/test_trainer_rank_recompute_memory.py +++ b/tests/unit/test_trainer_rank_recompute_memory.py @@ -421,27 +421,28 @@ def price(tokens=None, segments=plan.grad_segment_count): @pytest.mark.parametrize("kind", ("full", "no_grad", "moe")) -def test_cp_keeps_existing_full_no_grad_and_moe_costs(monkeypatch, kind): +def test_cp_prices_group_rows_at_the_most_loaded_rank(monkeypatch, kind): rank = _rank( "full" if kind == "full" else "selective", **({"num_moe_experts": 64} if kind == "moe" else {}), ) monkeypatch.setattr(rank, "_topology_key", lambda: (1, 1, 2, 1)) + monkeypatch.setattr(rank, "_topology", lambda: SimpleNamespace(cp=2, tp=1)) monkeypatch.setattr( - rank, - "_max_rank_model_tokens", - lambda *a, **kw: pytest.fail("unexpected CP plan"), + rank, "_max_rank_model_tokens", lambda batch, **_: batch.tokens.numel() - 1 ) plan = _plan(rank, no_grad=kind == "no_grad") + # Global counts cannot price per-rank floors: width probes defer to plans. + request = ForwardInput( + input_tokens=torch.tensor([1, 2]), hidden_states=True, no_grad=kind == "no_grad" + ) + assert rank._estimate_flat_forward([request]) is None + # Retention stays global for these kinds; floors use the most loaded rank. assert rank._plan_retained_tokens(plan) == plan.packed_tokens - assert rank._memory_check(plan).estimated_required_bytes == ( - rank._estimate_required_memory_bytes_from_values( - packed_tokens=plan.packed_tokens, - output_bytes=plan.output_bytes, - signature=plan.signature, - gdn_segments=plan.grad_segment_count, - ) + assert rank._plan_group_rows(plan) == tuple( + (int(group.packed.tokens.numel()) - 1, group.grad_enabled) + for group in plan.groups ) diff --git a/tests/unit/test_trainer_rank_weird_shapes.py b/tests/unit/test_trainer_rank_weird_shapes.py index 174e56499..51b0c313a 100644 --- a/tests/unit/test_trainer_rank_weird_shapes.py +++ b/tests/unit/test_trainer_rank_weird_shapes.py @@ -782,7 +782,8 @@ def test_adaptive_planner_globally_falls_back_when_one_rank_cannot_estimate( rank = TrainerRank(_runtime()) monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 2)) # Planning succeeds; only estimator availability and profile trust are false. - outcomes = iter((True, False, True, True, True, False)) + # The last outcome says no profile exists, so full sharing is not retried. + outcomes = iter((True, False, True, True, True, False, False)) monkeypatch.setattr(rank, "_all_ranks_true", lambda _local: next(outcomes)) plans = 0 original = rank._plan_flat_forward