Reject unsupported selective recompute in TrainerRank planning

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

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))