diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index de15bbdee..978173a91 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1326,6 +1326,14 @@ def __init__(self, runtime: TrainingRuntime) -> None: "therefore requires PP=1 with exactly one local model chunk; " f"got pp={pp_size}, chunks={len(runtime.model)}" ) + if getattr(runtime.provider, "recompute_granularity", None) == "selective": + raise TrainerRankRuntimeSupportError( + "TrainerRank memory planning does not support selective recompute; " + "its activation estimate assumes full recompute. Use " + "ART_MEGATRON_RECOMPUTE_GRANULARITY=full with " + "ART_MEGATRON_RECOMPUTE_METHOD=uniform and " + "ART_MEGATRON_RECOMPUTE_NUM_LAYERS=1." + ) # Tensor parallelism is admitted: the vocab-parallel head, sequence- # parallel gather, TP padding of packed batches and sharded LoRA # gradient reduction pre-date the planner, memory checks all-reduce diff --git a/tests/unit/test_trainer_rank_topology.py b/tests/unit/test_trainer_rank_topology.py index d42609796..e706b36a7 100644 --- a/tests/unit/test_trainer_rank_topology.py +++ b/tests/unit/test_trainer_rank_topology.py @@ -11,7 +11,7 @@ from __future__ import annotations from types import SimpleNamespace -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal import pytest import torch @@ -58,6 +58,23 @@ def test_trainer_rank_accepts_tensor_parallel_runtimes(tp: int) -> None: assert rank.last_forward_telemetry is not None +@pytest.mark.parametrize("tp", (1, 2, 4)) +@pytest.mark.parametrize("granularity", (None, "full", "selective")) +def test_trainer_rank_recompute_support( + tp: int, granularity: Literal["full", "selective"] | None +) -> None: + runtime = _runtime(tp=tp) + runtime.provider.recompute_granularity = granularity + if granularity == "selective": + with pytest.raises( + TrainerRankRuntimeSupportError, + match="selective recompute.*ART_MEGATRON_RECOMPUTE_GRANULARITY=full", + ): + TrainerRank(runtime) + else: + TrainerRank(runtime) + + def test_trainer_rank_still_refuses_pipeline_parallel_runtimes() -> None: with pytest.raises(TrainerRankRuntimeSupportError, match="PP=1"): TrainerRank(_runtime(pp=2))