fix(trainer-rank): account for recomputed mixer memory and head stages

.github/workflows/prek.yml

@@ -233,9 +233,11 @@ jobs:
             tests/unit/test_trainer_rank_profile_warm.py \
             tests/unit/test_trainer_rank_tp_floor.py \
             tests/unit/test_trainer_rank_checkpoint_gradient_memory.py \
+            tests/unit/test_trainer_rank_adapter_gradient_memory.py \
             tests/unit/test_trainer_rank_slot_memory.py \
             tests/unit/test_trainer_rank_moe_memory.py \
             tests/unit/test_trainer_rank_head_memory.py \
+            tests/unit/test_trainer_rank_head_stage_memory.py \
             tests/unit/test_grouped_planner_replay.py \
             tests/unit/test_planner_runtime_fact_guards.py \
             tests/unit/test_planner_replay_owner_budget.py \
@@ -282,9 +284,11 @@ jobs:
             --ignore=tests/unit/test_trainer_rank_profile_warm.py \
             --ignore=tests/unit/test_trainer_rank_tp_floor.py \
             --ignore=tests/unit/test_trainer_rank_checkpoint_gradient_memory.py \
+            --ignore=tests/unit/test_trainer_rank_adapter_gradient_memory.py \
             --ignore=tests/unit/test_trainer_rank_slot_memory.py \
             --ignore=tests/unit/test_trainer_rank_moe_memory.py \
             --ignore=tests/unit/test_trainer_rank_head_memory.py \
+            --ignore=tests/unit/test_trainer_rank_head_stage_memory.py \
             --ignore=tests/unit/test_grouped_planner_replay.py \
             --ignore=tests/unit/test_planner_runtime_fact_guards.py \
             --ignore=tests/unit/test_planner_replay_owner_budget.py \

src/art/megatron/context_parallel/runtime.py

@@ -433,6 +433,45 @@ def context_parallel_rank_model_token_counts(
     )
 
 
+def context_parallel_model_token_total(
+    *,
+    group_ids: torch.Tensor,
+    parent_ids: torch.Tensor,
+    topology: ParallelTopology,
+    config: ContextParallelConfig,
+    original_seq_len: int,
+    build_gdn_execution_spec: bool,
+    gdn_planner_config: Any | None = None,
+) -> int:
+    """Return the CP group's model rows in its larger physical layout.
+
+    Dispatch runs at least one row on every rank; an empty rank's padding row
+    passes through the model too.
+    """
+    planning_key, bundle, _group_ids_cpu, _parent_ids_cpu = (
+        _get_or_build_planning_bundle(
+            group_ids=group_ids,
+            parent_ids=parent_ids,
+            topology=topology,
+            config=config,
+            original_seq_len=original_seq_len,
+            build_gdn_execution_spec=build_gdn_execution_spec,
+        )
+    )
+    total = sum(
+        max(1, count) for count in bundle.token_layout_index.token_counts_by_rank
+    )
+    if not build_gdn_execution_spec:
+        return total
+    decision = _plan_gdn_global_execution(
+        planning_key=planning_key,
+        bundle=bundle,
+        topology=topology,
+        gdn_planner_config=gdn_planner_config,
+    )
+    return max(total, sum(max(1, count) for count in decision.gdn_token_counts_by_rank))
+
+
 def _normalized_chunk_size(
     *,
     valid_tokens: int,