From 06ef2ffb736ec270334317c8c3156b949763068e Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 16 Jul 2026 13:47:47 -0600 Subject: [PATCH 01/10] fix: complete TrainerRank model runtime integration --- src/art/trainer_rank/__init__.py | 66 +++++++ .../megatron/lora/test_dynamic_lora_slots.py | 10 +- tests/unit/test_trainer_rank_validation.py | 183 +++++++++++++++++- 3 files changed, 256 insertions(+), 3 deletions(-) diff --git a/src/art/trainer_rank/__init__.py b/src/art/trainer_rank/__init__.py index ed3a05a8a..774826689 100644 --- a/src/art/trainer_rank/__init__.py +++ b/src/art/trainer_rank/__init__.py @@ -970,6 +970,11 @@ def _prepare_adapter_model( name: str, adapter_model: Mapping[str, torch.Tensor], ) -> dict[str, torch.Tensor]: + adapter_model = ( + self.runtime.model_support_handler.canonicalize_loaded_lora_state( + dict(adapter_model), self.runtime.model + ) + ) templates = self._local_lora_adapter_templates() keys = set(adapter_model) expected = set(templates) @@ -1153,6 +1158,9 @@ def _dynamic_optim_step( params: AdamParams, scale_grads: float, ) -> dict[str, float]: + self.runtime.model_support_handler.zero_internal_padding_grads( + self.runtime.model + ) selected = [] for name in checkpoint_names: self._guard_checkpoint_can_step(name) @@ -2335,9 +2343,14 @@ def _prepare_packed_forward( ) -> _PreparedPackedForward: topology = self._topology() batch = _pad_packed_batch(batch, multiple=int(topology.tp)) + self._configure_hybridep(batch, topology=topology) if int(topology.cp) > 1: return self._prepare_context_parallel_forward(batch, topology=topology) from art.megatron.prefix_tree_state import create_prefix_tree_state + from art.megatron.training.microbatches import ( + _art_flex_sliding_windows, + _gdn_planner_config_for_provider, + ) handler = self.runtime.model_support_handler provider = self.runtime.provider @@ -2348,9 +2361,13 @@ def _prepare_packed_forward( group_ids=batch.group_ids, parent_ids=batch.parent_ids, target_device=self.device, + input_pos=batch.position_ids, + sliding_windows=_art_flex_sliding_windows(provider), build_gdn_execution_spec=handler.build_gdn_execution_spec, + model_support_handler=handler, attention_head_dim=provider.kv_channels, attention_value_head_dim=provider.kv_channels, + gdn_planner_config=_gdn_planner_config_for_provider(provider, handler), ), packed_seq_params=None, positions_by_item=batch.positions_by_sequence, @@ -2364,6 +2381,55 @@ def _prepare_packed_forward( ), ) + def _configure_hybridep( + self, + batch: PrefixTreePack, + *, + topology: "ParallelTopology", + ) -> None: + from megatron.core import parallel_state as ps + + if int(ps.get_expert_model_parallel_world_size()) <= 1: + return + from art.megatron.train import ( + _ensure_hybridep_capacity, + _set_hybridep_token_count, + ) + + sequence_length = int(batch.tokens.shape[1]) + _ensure_hybridep_capacity( + self.runtime, + packed_sequence_length=sequence_length, + context_parallel_size=int(topology.cp), + ) + rows = sequence_length + if int(topology.cp) > 1: + from art.megatron.context_parallel.runtime import ( + context_parallel_rank_model_token_counts, + ) + from art.megatron.training.microbatches import ( + _context_parallel_config_for_provider, + _gdn_planner_config_for_provider, + ) + + handler = self.runtime.model_support_handler + rows = max( + context_parallel_rank_model_token_counts( + group_ids=batch.group_ids, + parent_ids=batch.parent_ids, + topology=topology, + config=_context_parallel_config_for_provider( + self.runtime.provider, self.device + ), + original_seq_len=sequence_length, + build_gdn_execution_spec=handler.build_gdn_execution_spec, + gdn_planner_config=_gdn_planner_config_for_provider( + self.runtime.provider, handler + ), + ) + ) + _set_hybridep_token_count(rows) + def _prepare_context_parallel_forward( self, batch: PrefixTreePack, diff --git a/tests/integration/megatron/lora/test_dynamic_lora_slots.py b/tests/integration/megatron/lora/test_dynamic_lora_slots.py index bc6ddd817..f30957204 100644 --- a/tests/integration/megatron/lora/test_dynamic_lora_slots.py +++ b/tests/integration/megatron/lora/test_dynamic_lora_slots.py @@ -388,7 +388,15 @@ def _assert_reload_replaces_slot_optimizer( def _trainer_for(lora: LoRA, device: torch.device) -> TrainerRank: trainer = TrainerRank.__new__(TrainerRank) - trainer.runtime = SimpleNamespace(model=[lora], optimizer=None) + trainer.runtime = SimpleNamespace( + model=[lora], + optimizer=None, + model_support_handler=SimpleNamespace( + canonicalize_loaded_lora_state=lambda state, _model: state, + zero_internal_padding_grads=lambda _model: None, + zero_internal_padding_params=lambda _model: None, + ), + ) trainer.device = device trainer._slot_stack = [] trainer._default_slot_ref = None diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py index bdd21ecb6..11926f32c 100644 --- a/tests/unit/test_trainer_rank_validation.py +++ b/tests/unit/test_trainer_rank_validation.py @@ -3,6 +3,7 @@ from collections.abc import Iterable from dataclasses import dataclass import gc +from importlib.util import find_spec import inspect from types import SimpleNamespace from typing import TYPE_CHECKING, Any, cast @@ -10,6 +11,7 @@ import pytest import torch +from art.megatron.prefix_tree_packing import prefix_tree_pack from art.trainer_rank import ( AdamParams, AdapterSelection, @@ -84,8 +86,18 @@ def _runtime( return SimpleNamespace( model=[model or torch.nn.Linear(1, 1)], optimizer=optimizer, - provider=SimpleNamespace(hidden_size=4, num_layers=1), - model_support_handler=SimpleNamespace(build_gdn_execution_spec=True), + provider=SimpleNamespace( + hidden_size=4, + num_layers=1, + kv_channels=2, + art_flex_sliding_windows=(16,), + ), + model_support_handler=SimpleNamespace( + build_gdn_execution_spec=True, + canonicalize_loaded_lora_state=lambda state, _model: state, + zero_internal_padding_grads=lambda _model: None, + zero_internal_padding_params=lambda _model: None, + ), ) # type: ignore @@ -200,6 +212,92 @@ def test_trainer_rank_accepts_shared_prefix_depth(depth: int) -> None: assert trainer.shared_prefix_max_depth == depth +@pytest.mark.skipif(find_spec("megatron") is None, reason="requires Megatron") +def test_cp1_packed_forward_uses_model_attention_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from art.megatron.context_parallel.types import ParallelTopology + + runtime = _runtime() + trainer = TrainerRank(runtime) + batch = prefix_tree_pack( + (torch.tensor([1, 2, 3]), torch.tensor([1, 2, 4])), max_depth=1 + ) + captured: dict[str, object] = {} + state = object() + + def create_state(**kwargs: object) -> object: + captured.update(kwargs) + return state + + monkeypatch.setattr(trainer, "_topology", lambda: ParallelTopology()) + monkeypatch.setattr(trainer, "_configure_hybridep", lambda *_args, **_kwargs: None) + monkeypatch.setattr( + "art.megatron.prefix_tree_state.create_prefix_tree_state", create_state + ) + monkeypatch.setattr( + "art.megatron.training.microbatches._gdn_planner_config_for_provider", + lambda provider, handler: "planner-config", + ) + + prepared = trainer._prepare_packed_forward(batch) + + assert prepared.attention_state is state + assert captured["model_support_handler"] is runtime.model_support_handler + assert captured["sliding_windows"] == (16,) + assert captured["gdn_planner_config"] == "planner-config" + torch.testing.assert_close( + cast(torch.Tensor, captured["input_pos"]), batch.position_ids + ) + + +@pytest.mark.skipif(find_spec("megatron") is None, reason="requires Megatron") +def test_hybridep_uses_maximum_cp_model_rows( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from art.megatron.context_parallel.types import ParallelTopology + + trainer = TrainerRank(_runtime()) + batch = prefix_tree_pack((torch.arange(9),), max_depth=0) + topology = ParallelTopology(cp=4) + calls: dict[str, object] = {} + + monkeypatch.setattr( + "megatron.core.parallel_state.get_expert_model_parallel_world_size", + lambda: 4, + ) + monkeypatch.setattr( + "art.megatron.train._ensure_hybridep_capacity", + lambda runtime, **kwargs: calls.update(capacity=kwargs), + ) + monkeypatch.setattr( + "art.megatron.train._set_hybridep_token_count", + lambda rows: calls.update(rows=rows), + ) + monkeypatch.setattr( + "art.megatron.context_parallel.runtime.context_parallel_rank_model_token_counts", + lambda **_kwargs: (8, 13, 9, 11), + ) + monkeypatch.setattr( + "art.megatron.training.microbatches._context_parallel_config_for_provider", + lambda *_: "cp-config", + ) + monkeypatch.setattr( + "art.megatron.training.microbatches._gdn_planner_config_for_provider", + lambda *_: "gdn-config", + ) + + trainer._configure_hybridep(batch, topology=topology) + + assert calls == { + "capacity": { + "packed_sequence_length": 9, + "context_parallel_size": 4, + }, + "rows": 13, + } + + def test_trainer_rank_adapter_stack_errors() -> None: trainer = TrainerRank(_runtime()) @@ -312,6 +410,50 @@ def test_load_checkpoint_slot_retains_config_and_uses_its_alpha( assert "student" not in trainer._checkpoint_slot_adapter_configs +@pytest.mark.skipif(find_spec("megatron") is None, reason="requires Megatron") +def test_slot_load_canonicalizes_only_incoming_adapter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[tuple[dict[str, torch.Tensor], object]] = [] + loaded_state: dict[str, torch.Tensor] = {} + runtime = _runtime() + runtime.model_support_handler.canonicalize_loaded_lora_state = lambda state, model: ( + calls.append((state, model)) + or {key: torch.zeros_like(value) for key, value in state.items()} + ) + runtime.model_support_handler.zero_internal_padding_params = lambda _model: ( + pytest.fail("slot load must not mutate unrelated slot parameters") + ) + trainer = TrainerRank(runtime) + monkeypatch.setattr( + trainer, + "_local_lora_adapter_templates", + lambda: {"weight": torch.empty(1)}, + ) + monkeypatch.setattr(trainer, "_guard_slot_can_load", lambda *_: None) + + def load_slot( + _model: object, + _ref: object, + adapter_model: dict[str, torch.Tensor], + **_kwargs: object, + ) -> int: + loaded_state.update(adapter_model) + return 1 + + monkeypatch.setattr( + "art.megatron.lora.load_lora_slot_into_model", + load_slot, + ) + + adapter = {"weight": torch.ones(1)} + trainer._load_slot("checkpoint", "student", adapter, trainable=True, alpha=None) + + assert calls == [(adapter, runtime.model)] + torch.testing.assert_close(loaded_state["weight"], torch.zeros(1)) + torch.testing.assert_close(adapter["weight"], torch.ones(1)) + + def test_checkpoint_slot_publish_requires_retained_adapter_config() -> None: trainer = TrainerRank(_runtime()) with pytest.raises(ValueError, match="Unknown checkpoint slot"): @@ -409,6 +551,43 @@ def test_optim_step_implicitly_steps_only_slots_with_grads( torch.testing.assert_close(untouched, before_untouched) +def test_dynamic_optimizer_zeroes_internal_padding_grads_before_step( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[str] = [] + param = torch.nn.Parameter(torch.tensor([1.0, 0.0])) + param.grad = torch.ones_like(param) + runtime = _runtime() + + def zero_padding_grads(_model: object) -> None: + calls.append("grads") + assert param.grad is not None + param.grad[-1] = 0.0 + + runtime.model_support_handler.zero_internal_padding_grads = zero_padding_grads + runtime.model_support_handler.zero_internal_padding_params = lambda _model: ( + pytest.fail("slot step must not mutate unrelated slot parameters") + ) + trainer = TrainerRank(runtime) + trainer._checkpoint_slot_params_by_name["student"] = (param,) + monkeypatch.setattr( + trainer, + "_reduce_dynamic_grads", + lambda params, **_kwargs: tuple(item.grad.float() for item in params), + ) + + trainer.optim_step( + params=AdamParams( + learning_rate=1e-2, + weight_decay=0.1, + grad_clip_norm=10.0, + ) + ) + + assert calls == ["grads"] + assert param[-1].item() == 0.0 + + def test_checkpoint_slot_optimizer_state_reproduces_exact_next_step( monkeypatch: pytest.MonkeyPatch, ) -> None: From 865c0c928430546732e93b21be2bfeb3b86c8692 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 16 Jul 2026 14:58:30 -0600 Subject: [PATCH 02/10] fix: harden TrainerRank retained graph state --- src/art/trainer_rank/__init__.py | 172 +++++++++++++++--- .../megatron/lora/test_dynamic_lora_slots.py | 50 +++++ tests/unit/test_trainer_rank_validation.py | 73 +++++++- 3 files changed, 266 insertions(+), 29 deletions(-) diff --git a/src/art/trainer_rank/__init__.py b/src/art/trainer_rank/__init__.py index 774826689..7dfb154ea 100644 --- a/src/art/trainer_rank/__init__.py +++ b/src/art/trainer_rank/__init__.py @@ -970,11 +970,6 @@ def _prepare_adapter_model( name: str, adapter_model: Mapping[str, torch.Tensor], ) -> dict[str, torch.Tensor]: - adapter_model = ( - self.runtime.model_support_handler.canonicalize_loaded_lora_state( - dict(adapter_model), self.runtime.model - ) - ) templates = self._local_lora_adapter_templates() keys = set(adapter_model) expected = set(templates) @@ -990,15 +985,19 @@ def _prepare_adapter_model( f"installed LoRA wrapper sites: {preview}{more}. Configure the " "Megatron runtime with matching LoRA target modules before loading." ) + local_state = { + key: tensor for key, tensor in adapter_model.items() if key in templates + } + adapter_model = ( + self.runtime.model_support_handler.canonicalize_loaded_lora_state( + local_state, self.runtime.model + ) + ) return { - key: ( - tensor.to( - device=templates[key].device, - dtype=templates[key].dtype, - non_blocking=True, - ) - if key in templates - else tensor + key: tensor.to( + device=templates[key].device, + dtype=templates[key].dtype, + non_blocking=True, ) for key, tensor in adapter_model.items() } @@ -1307,8 +1306,76 @@ def _restore_dynamic_optimizer( f"{name!r} has shape {tuple(value.shape)}, but the loaded " f"slot parameter has shape {tuple(param.shape)}." ) + self._zero_dynamic_optimizer_padding(name, dynamic) return dynamic + def _zero_dynamic_optimizer_padding( + self, + name: str, + dynamic: _DynamicOptimizer, + ) -> None: + masks = self._dynamic_optimizer_padding_masks(name) + with torch.no_grad(): + for param, mask in zip(dynamic.master_params, masks, strict=True): + param.masked_fill_(mask, 0) + for value in dynamic.optimizer.state.get(param, {}).values(): + if isinstance(value, torch.Tensor) and value.shape == param.shape: + value.masked_fill_(mask, 0) + + def _dynamic_optimizer_padding_masks(self, name: str) -> tuple[torch.Tensor, ...]: + params = self._checkpoint_slot_params_by_name[name] + masks = tuple(torch.zeros_like(param, dtype=torch.bool) for param in params) + param_indices = {id(param): index for index, param in enumerate(params)} + exported: dict[str, torch.Tensor] = {} + owners: dict[str, tuple[int, int | None]] = {} + ref = self._slot_ref("checkpoint", name) + + for chunk in self.runtime.model: + for module in chunk.modules(): + lora_params = getattr(module, "_lora_params", None) + expected_keys = getattr(module, "_expected_weight_keys", None) + if not callable(lora_params) or not callable(expected_keys): + continue + for suffix, param in lora_params(ref): + index = param_indices.get(id(param)) + if index is None: + continue + keys = expected_keys(str(suffix).removesuffix(".weight")) + if int(param.ndim) == 3: + if len(keys) != int(param.shape[0]): + raise TrainerRankSlotStateError( + f"Cannot map optimizer padding for checkpoint slot " + f"{name!r}: {len(keys)} adapter keys describe " + f"{int(param.shape[0])} local experts." + ) + for expert, key in enumerate(keys): + exported[str(key)] = torch.ones_like(param[expert].T) + owners[str(key)] = (index, expert) + else: + if len(keys) != 1: + raise TrainerRankSlotStateError( + f"Cannot map optimizer padding for checkpoint slot " + f"{name!r}: expected one adapter key, got {len(keys)}." + ) + key = str(keys[0]) + exported[key] = torch.ones_like(param.T) + owners[key] = (index, None) + + canonical = self.runtime.model_support_handler.canonicalize_loaded_lora_state( + exported, self.runtime.model + ) + for key, value in canonical.items(): + owner = owners.get(key) + if owner is None or not isinstance(value, torch.Tensor): + continue + index, expert = owner + mask = value.T == 0 + if expert is None: + masks[index].copy_(mask) + else: + masks[index][expert].copy_(mask) + return masks + def _reduce_dynamic_grads( self, params: Sequence[torch.nn.Parameter], @@ -1665,6 +1732,10 @@ def _execute_flat_plan(self, plan: _FlatForwardPlan) -> list[AnyForwardOutput]: outputs = [ ForwardOutput(None, None, None, None) for _ in range(plan.request_count) ] + if plan.groups: + self._configure_hybridep( + tuple(group.packed for group in plan.groups), topology=self._topology() + ) for group in plan.groups: from art.megatron.lora import use_lora_slot @@ -1681,7 +1752,9 @@ def _track_slot_graph_outputs( ref: "LoRASlotRef | None", outputs: Sequence[AnyForwardOutput], ) -> list[AnyForwardOutput]: - if ref is None or ref.name is None: + track_slot = ref is not None and ref.name is not None + track_hybridep = bool(getattr(self, "_hybridep_graph_tracking", False)) + if not track_slot and not track_hybridep: return list(outputs) marker: torch.Tensor | None = None @@ -1711,9 +1784,25 @@ def track(tensor: torch.Tensor | None) -> torch.Tensor | None: for output in outputs ] if marker is not None: - self._slot_graphs().setdefault(ref, []).append(weakref.ref(marker)) + marker_ref = weakref.ref(marker) + if track_slot: + self._slot_graphs().setdefault(ref, []).append(marker_ref) + if track_hybridep: + self._hybridep_graphs().append(marker_ref) return tracked_outputs + def _hybridep_graphs(self) -> list[weakref.ReferenceType[torch.Tensor]]: + graphs = getattr(self, "_pending_hybridep_graphs", None) + if graphs is None: + graphs = [] + self._pending_hybridep_graphs = graphs + return graphs + + def _has_live_hybridep_graphs(self) -> bool: + graphs = self._hybridep_graphs() + graphs[:] = [marker for marker in graphs if marker() is not None] + return bool(graphs) + def _slot_graphs( self, ) -> dict["LoRASlotRef", list[weakref.ReferenceType[torch.Tensor]]]: @@ -2343,7 +2432,6 @@ def _prepare_packed_forward( ) -> _PreparedPackedForward: topology = self._topology() batch = _pad_packed_batch(batch, multiple=int(topology.tp)) - self._configure_hybridep(batch, topology=topology) if int(topology.cp) > 1: return self._prepare_context_parallel_forward(batch, topology=topology) from art.megatron.prefix_tree_state import create_prefix_tree_state @@ -2383,26 +2471,68 @@ def _prepare_packed_forward( def _configure_hybridep( self, - batch: PrefixTreePack, + batches: Sequence[PrefixTreePack], *, topology: "ParallelTopology", ) -> None: from megatron.core import parallel_state as ps if int(ps.get_expert_model_parallel_world_size()) <= 1: + self._hybridep_graph_tracking = False + return + if not batches: return + from megatron.core.transformer.moe import fused_a2a + from art.megatron.train import ( _ensure_hybridep_capacity, + _hybridep_token_capacity, _set_hybridep_token_count, ) - sequence_length = int(batch.tokens.shape[1]) + padded = tuple( + _pad_packed_batch(batch, multiple=int(topology.tp)) for batch in batches + ) + sequence_length = max(int(batch.tokens.shape[1]) for batch in padded) + rows = max(self._hybridep_rows(batch, topology=topology) for batch in padded) + current = fused_a2a._hybrid_ep_buffer + live = self._has_live_hybridep_graphs() + required_capacity = _hybridep_token_capacity(sequence_length, int(topology.cp)) + if live and ( + current is None + or id(current) != getattr(self, "_hybridep_buffer_id", None) + or int(current.configurer.buffer_config.max_num_of_tokens_per_rank) + < required_capacity + ): + raise TrainerRankSlotStateError( + "Cannot grow or replace the HybridEP buffer while an earlier " + "TrainerRank forward still has a live backward graph. Finish " + "backward or release those outputs before forwarding a larger batch." + ) _ensure_hybridep_capacity( self.runtime, packed_sequence_length=sequence_length, context_parallel_size=int(topology.cp), ) - rows = sequence_length + current = fused_a2a._hybrid_ep_buffer + if current is None: + raise RuntimeError("HybridEP buffer was not initialized") + if live: + rows = max(rows, int(getattr(self, "_hybridep_rows_high_water", 0))) + _set_hybridep_token_count(rows) + self._hybridep_buffer_id = id(current) + self._hybridep_rows_high_water = rows + self._hybridep_graph_tracking = True + + def _hybridep_rows( + self, + batch: PrefixTreePack, + *, + topology: "ParallelTopology", + ) -> int: + sequence_length = int(batch.tokens.shape[1]) + if int(topology.cp) <= 1: + return sequence_length if int(topology.cp) > 1: from art.megatron.context_parallel.runtime import ( context_parallel_rank_model_token_counts, @@ -2413,7 +2543,7 @@ def _configure_hybridep( ) handler = self.runtime.model_support_handler - rows = max( + return max( context_parallel_rank_model_token_counts( group_ids=batch.group_ids, parent_ids=batch.parent_ids, @@ -2428,7 +2558,7 @@ def _configure_hybridep( ), ) ) - _set_hybridep_token_count(rows) + raise AssertionError("unreachable") def _prepare_context_parallel_forward( self, diff --git a/tests/integration/megatron/lora/test_dynamic_lora_slots.py b/tests/integration/megatron/lora/test_dynamic_lora_slots.py index f30957204..44cb6655c 100644 --- a/tests/integration/megatron/lora/test_dynamic_lora_slots.py +++ b/tests/integration/megatron/lora/test_dynamic_lora_slots.py @@ -5,6 +5,7 @@ from pathlib import Path import socket from types import SimpleNamespace +from typing import cast import pytest @@ -284,6 +285,55 @@ def _assert_distributed_optimizer_restore(device: torch.device) -> None: torch.testing.assert_close(actual, expected, atol=0, rtol=0) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required.") +def test_restored_dynamic_optimizer_canonicalizes_internal_padding() -> None: + with _single_rank_model_parallel(): + _assert_restored_dynamic_optimizer_canonicalizes_internal_padding() + + +def _assert_restored_dynamic_optimizer_canonicalizes_internal_padding() -> None: + device = torch.device("cuda") + ref = LoRASlotRef("checkpoint", "A") + adapter = _adapter("dense", rank=2, seed=17) + lora = LoRA("dense", 4, 5, 2, 32, torch.float32, device) + lora.load_lora_slot(ref, adapter, requires_grad=True) + trainer = _trainer_for(lora, device) + + def canonicalize( + state: dict[str, torch.Tensor], _model: object + ) -> dict[str, torch.Tensor]: + result = {key: value.clone() for key, value in state.items()} + for value in result.values(): + value[..., -1] = 0 + return result + + trainer.runtime.model_support_handler.canonicalize_loaded_lora_state = canonicalize + with use_lora_slot(ref): + lora(torch.randn(3, 4, device=device)).sum().backward() + trainer.optim_step( + params=AdamParams(learning_rate=1e-3, weight_decay=0.1, grad_clip_norm=0.0), + checkpoints=["A"], + ) + state = trainer.checkpoint_slot_optimizer_state("A") + assert state is not None + masks = trainer._dynamic_optimizer_padding_masks("A") + masters = cast(tuple[torch.Tensor, ...], state["master_params"]) + optimizer = cast(dict[str, object], state["optimizer"]) + optimizer_states = cast(dict[int, dict[str, object]], optimizer["state"]) + for index, (master, mask) in enumerate(zip(masters, masks, strict=True)): + master.masked_fill_(mask.cpu(), 5) + for value in optimizer_states[index].values(): + if isinstance(value, torch.Tensor) and value.shape == master.shape: + value.masked_fill_(mask.cpu(), 5) + + restored = trainer._restore_dynamic_optimizer("A", state) + for master, mask in zip(restored.master_params, masks, strict=True): + assert torch.count_nonzero(master[mask]) == 0 + for value in restored.optimizer.state[master].values(): + if isinstance(value, torch.Tensor) and value.shape == master.shape: + assert torch.count_nonzero(value[mask]) == 0 + + def _local_shard(full: torch.Tensor, rank: int, size: int) -> torch.Tensor: return full[:, rank * size : (rank + 1) * size].clone().requires_grad_() diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py index 11926f32c..69983e782 100644 --- a/tests/unit/test_trainer_rank_validation.py +++ b/tests/unit/test_trainer_rank_validation.py @@ -259,6 +259,7 @@ def test_hybridep_uses_maximum_cp_model_rows( trainer = TrainerRank(_runtime()) batch = prefix_tree_pack((torch.arange(9),), max_depth=0) + short_batch = prefix_tree_pack((torch.arange(5),), max_depth=0) topology = ParallelTopology(cp=4) calls: dict[str, object] = {} @@ -276,7 +277,12 @@ def test_hybridep_uses_maximum_cp_model_rows( ) monkeypatch.setattr( "art.megatron.context_parallel.runtime.context_parallel_rank_model_token_counts", - lambda **_kwargs: (8, 13, 9, 11), + lambda **kwargs: ( + 8, + int(cast(torch.Tensor, kwargs["group_ids"]).numel()) + 4, + 9, + 11, + ), ) monkeypatch.setattr( "art.megatron.training.microbatches._context_parallel_config_for_provider", @@ -287,7 +293,16 @@ def test_hybridep_uses_maximum_cp_model_rows( lambda *_: "gdn-config", ) - trainer._configure_hybridep(batch, topology=topology) + buffer = SimpleNamespace( + configurer=SimpleNamespace( + buffer_config=SimpleNamespace(max_num_of_tokens_per_rank=1024) + ) + ) + monkeypatch.setattr( + "megatron.core.transformer.moe.fused_a2a._hybrid_ep_buffer", buffer + ) + + trainer._configure_hybridep((batch, short_batch), topology=topology) assert calls == { "capacity": { @@ -298,6 +313,41 @@ def test_hybridep_uses_maximum_cp_model_rows( } +@pytest.mark.skipif(find_spec("megatron") is None, reason="requires Megatron") +def test_hybridep_rejects_buffer_growth_with_live_graph( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from art.megatron.context_parallel.types import ParallelTopology + + trainer = TrainerRank(_runtime()) + trainer._hybridep_graph_tracking = True + output = trainer._track_slot_graph_outputs( + None, + [ForwardOutput(torch.ones(1, requires_grad=True), None, None, None)], + )[0] + assert output.target_logprobs is not None + buffer = SimpleNamespace( + configurer=SimpleNamespace( + buffer_config=SimpleNamespace(max_num_of_tokens_per_rank=4) + ) + ) + trainer._hybridep_buffer_id = id(buffer) + trainer._hybridep_rows_high_water = 4 + monkeypatch.setattr( + "megatron.core.parallel_state.get_expert_model_parallel_world_size", + lambda: 4, + ) + monkeypatch.setattr( + "megatron.core.transformer.moe.fused_a2a._hybrid_ep_buffer", buffer + ) + + with pytest.raises(TrainerRankSlotStateError, match="live backward graph"): + trainer._configure_hybridep( + (prefix_tree_pack((torch.arange(9),), max_depth=0),), + topology=ParallelTopology(), + ) + + def test_trainer_rank_adapter_stack_errors() -> None: trainer = TrainerRank(_runtime()) @@ -411,7 +461,7 @@ def test_load_checkpoint_slot_retains_config_and_uses_its_alpha( @pytest.mark.skipif(find_spec("megatron") is None, reason="requires Megatron") -def test_slot_load_canonicalizes_only_incoming_adapter( +def test_slot_load_canonicalizes_only_local_incoming_adapter( monkeypatch: pytest.MonkeyPatch, ) -> None: calls: list[tuple[dict[str, torch.Tensor], object]] = [] @@ -426,10 +476,16 @@ def test_slot_load_canonicalizes_only_incoming_adapter( ) trainer = TrainerRank(runtime) monkeypatch.setattr( - trainer, - "_local_lora_adapter_templates", - lambda: {"weight": torch.empty(1)}, + trainer, "_local_lora_adapter_templates", lambda: {"weight": torch.empty(1)} ) + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 2) + + def gather_expected(values: list[set[str] | None], local: set[str]) -> None: + values[:] = [local, {"remote_weight"}] + + monkeypatch.setattr(torch.distributed, "all_gather_object", gather_expected) monkeypatch.setattr(trainer, "_guard_slot_can_load", lambda *_: None) def load_slot( @@ -446,11 +502,12 @@ def load_slot( load_slot, ) - adapter = {"weight": torch.ones(1)} + adapter = {"weight": torch.ones(1), "remote_weight": torch.ones(1)} trainer._load_slot("checkpoint", "student", adapter, trainable=True, alpha=None) - assert calls == [(adapter, runtime.model)] + assert calls == [({"weight": adapter["weight"]}, runtime.model)] torch.testing.assert_close(loaded_state["weight"], torch.zeros(1)) + assert "remote_weight" not in loaded_state torch.testing.assert_close(adapter["weight"], torch.ones(1)) From 04a3013cfaf50e70c90a0f7f5eca118e86c015e4 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 16 Jul 2026 15:01:41 -0600 Subject: [PATCH 03/10] test: cover unequal HybridEP slot graphs --- dev/trainer_rank_check.py | 32 ++++++++++++++++++++++---------- 1 file changed, 22 insertions(+), 10 deletions(-) diff --git a/dev/trainer_rank_check.py b/dev/trainer_rank_check.py index 54410613d..57a0b3c18 100644 --- a/dev/trainer_rank_check.py +++ b/dev/trainer_rank_check.py @@ -44,7 +44,7 @@ def main( depths: str = "0,1,2,3,4", performance_depth: int = 1, chunks: str = "17,512,8192", - workload: Literal["regular", "austin", "varied"] = "regular", + workload: Literal["regular", "austin", "varied", "unequal_slots"] = "regular", request: Literal["target", "multi", "topk", "logits", "hidden", "mixed"] = "target", families: int = 8, prefix_tokens: int = 128, @@ -397,15 +397,27 @@ def _performance( families, prefix_tokens, branches, completion_tokens = 30, 5000, 16, 100 rank = TrainerRank(runtime, shared_prefix_max_depth=depth, head_chunk_tokens=8192) slot_names = load_random_checkpoint_slots(runtime, rank, slots) - requests = _performance_requests( - request=request, - families=families, - prefix_tokens=prefix_tokens, - branches=branches, - completion_tokens=completion_tokens, - varied=workload == "varied", - slots=slot_names, - ) + if workload == "unequal_slots": + if len(slot_names) < 2: + raise ValueError("--workload unequal_slots requires --slots >= 2") + requests = [ + ForwardInput( + input_tokens=(tokens := _tokens(index * 1009, length)), + target_tokens=(tokens * 7 + 3) % 32_000, + checkpoint=slot_names[index], + ) + for index, length in enumerate((256, 64)) + ] + else: + requests = _performance_requests( + request=request, + families=families, + prefix_tokens=prefix_tokens, + branches=branches, + completion_tokens=completion_tokens, + varied=workload == "varied", + slots=slot_names, + ) dp_rank, dp_size = rank._dp_rank_and_size() plan = rank._plan_flat_forward(requests) assert workload != "austin" or plan.packed_tokens == 198_000 From 1b27ecfb91456280ebb907e8c51f5bb4c9a4dfd3 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 16 Jul 2026 15:03:19 -0600 Subject: [PATCH 04/10] perf: preserve per-group HybridEP extents --- src/art/trainer_rank/__init__.py | 56 +++++++++++++++------- tests/unit/test_trainer_rank_validation.py | 8 +--- 2 files changed, 40 insertions(+), 24 deletions(-) diff --git a/src/art/trainer_rank/__init__.py b/src/art/trainer_rank/__init__.py index 7dfb154ea..1d3b05e65 100644 --- a/src/art/trainer_rank/__init__.py +++ b/src/art/trainer_rank/__init__.py @@ -1732,19 +1732,32 @@ def _execute_flat_plan(self, plan: _FlatForwardPlan) -> list[AnyForwardOutput]: outputs = [ ForwardOutput(None, None, None, None) for _ in range(plan.request_count) ] - if plan.groups: + hybridep = ( self._configure_hybridep( tuple(group.packed for group in plan.groups), topology=self._topology() ) - for group in plan.groups: - from art.megatron.lora import use_lora_slot - - with use_lora_slot(group.slot_ref): - prepared = self._prepare_packed_forward(group.packed) - item_outputs = self._forward_packed(group.items, prepared) - item_outputs = self._track_slot_graph_outputs(group.slot_ref, item_outputs) - for index, output in zip(group.request_indices, item_outputs, strict=True): - outputs[index] = output + if plan.groups + else None + ) + try: + for group_index, group in enumerate(plan.groups): + from art.megatron.lora import use_lora_slot + + if hybridep is not None: + self._set_hybridep_rows(hybridep[0][group_index]) + with use_lora_slot(group.slot_ref): + prepared = self._prepare_packed_forward(group.packed) + item_outputs = self._forward_packed(group.items, prepared) + item_outputs = self._track_slot_graph_outputs( + group.slot_ref, item_outputs + ) + for index, output in zip( + group.request_indices, item_outputs, strict=True + ): + outputs[index] = output + finally: + if hybridep is not None: + self._set_hybridep_rows(hybridep[1]) return outputs def _track_slot_graph_outputs( @@ -2474,27 +2487,26 @@ def _configure_hybridep( batches: Sequence[PrefixTreePack], *, topology: "ParallelTopology", - ) -> None: + ) -> tuple[tuple[int, ...], int] | None: from megatron.core import parallel_state as ps if int(ps.get_expert_model_parallel_world_size()) <= 1: self._hybridep_graph_tracking = False - return + return None if not batches: - return + return None from megatron.core.transformer.moe import fused_a2a from art.megatron.train import ( _ensure_hybridep_capacity, _hybridep_token_capacity, - _set_hybridep_token_count, ) padded = tuple( _pad_packed_batch(batch, multiple=int(topology.tp)) for batch in batches ) sequence_length = max(int(batch.tokens.shape[1]) for batch in padded) - rows = max(self._hybridep_rows(batch, topology=topology) for batch in padded) + rows = tuple(self._hybridep_rows(batch, topology=topology) for batch in padded) current = fused_a2a._hybrid_ep_buffer live = self._has_live_hybridep_graphs() required_capacity = _hybridep_token_capacity(sequence_length, int(topology.cp)) @@ -2518,11 +2530,19 @@ def _configure_hybridep( if current is None: raise RuntimeError("HybridEP buffer was not initialized") if live: - rows = max(rows, int(getattr(self, "_hybridep_rows_high_water", 0))) - _set_hybridep_token_count(rows) + high_water = max(*rows, int(getattr(self, "_hybridep_rows_high_water", 0))) + else: + high_water = max(rows) self._hybridep_buffer_id = id(current) - self._hybridep_rows_high_water = rows + self._hybridep_rows_high_water = high_water self._hybridep_graph_tracking = True + return rows, high_water + + @staticmethod + def _set_hybridep_rows(rows: int) -> None: + from art.megatron.train import _set_hybridep_token_count + + _set_hybridep_token_count(rows) def _hybridep_rows( self, diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py index 69983e782..066117eb7 100644 --- a/tests/unit/test_trainer_rank_validation.py +++ b/tests/unit/test_trainer_rank_validation.py @@ -271,10 +271,6 @@ def test_hybridep_uses_maximum_cp_model_rows( "art.megatron.train._ensure_hybridep_capacity", lambda runtime, **kwargs: calls.update(capacity=kwargs), ) - monkeypatch.setattr( - "art.megatron.train._set_hybridep_token_count", - lambda rows: calls.update(rows=rows), - ) monkeypatch.setattr( "art.megatron.context_parallel.runtime.context_parallel_rank_model_token_counts", lambda **kwargs: ( @@ -302,15 +298,15 @@ def test_hybridep_uses_maximum_cp_model_rows( "megatron.core.transformer.moe.fused_a2a._hybrid_ep_buffer", buffer ) - trainer._configure_hybridep((batch, short_batch), topology=topology) + configured = trainer._configure_hybridep((batch, short_batch), topology=topology) assert calls == { "capacity": { "packed_sequence_length": 9, "context_parallel_size": 4, }, - "rows": 13, } + assert configured == ((13, 11), 13) @pytest.mark.skipif(find_spec("megatron") is None, reason="requires Megatron") From 3006f0a535ef87f2128afceda691e57ea37d52b1 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 16 Jul 2026 15:17:49 -0600 Subject: [PATCH 05/10] fix: validate dynamic slot state layouts --- src/art/trainer_rank/__init__.py | 19 +++++++++++ .../megatron/lora/test_dynamic_lora_slots.py | 33 +++++++++++++++---- 2 files changed, 46 insertions(+), 6 deletions(-) diff --git a/src/art/trainer_rank/__init__.py b/src/art/trainer_rank/__init__.py index 1d3b05e65..b22062bdf 100644 --- a/src/art/trainer_rank/__init__.py +++ b/src/art/trainer_rank/__init__.py @@ -550,6 +550,10 @@ def __init__( self._pending_slot_graphs: dict[ LoRASlotRef, list[weakref.ReferenceType[torch.Tensor]] ] = {} + self._pending_hybridep_graphs: list[weakref.ReferenceType[torch.Tensor]] = [] + self._hybridep_graph_tracking = False + self._hybridep_buffer_id: int | None = None + self._hybridep_rows_high_water = 0 self._memory_profiles: dict[_MemorySignature, _MemoryProfile] = {} self._last_global_micro_batch_size: int | None = None self.zero_grad() @@ -993,6 +997,11 @@ def _prepare_adapter_model( local_state, self.runtime.model ) ) + if set(adapter_model) != set(local_state): + raise TrainerRankSlotStateError( + "Model-specific LoRA canonicalization changed the adapter key set " + f"for {kind} slot {name!r}." + ) return { key: tensor.to( device=templates[key].device, @@ -1328,6 +1337,7 @@ def _dynamic_optimizer_padding_masks(self, name: str) -> tuple[torch.Tensor, ... param_indices = {id(param): index for index, param in enumerate(params)} exported: dict[str, torch.Tensor] = {} owners: dict[str, tuple[int, int | None]] = {} + mapped_indices: set[int] = set() ref = self._slot_ref("checkpoint", name) for chunk in self.runtime.model: @@ -1340,6 +1350,7 @@ def _dynamic_optimizer_padding_masks(self, name: str) -> tuple[torch.Tensor, ... index = param_indices.get(id(param)) if index is None: continue + mapped_indices.add(index) keys = expected_keys(str(suffix).removesuffix(".weight")) if int(param.ndim) == 3: if len(keys) != int(param.shape[0]): @@ -1361,6 +1372,14 @@ def _dynamic_optimizer_padding_masks(self, name: str) -> tuple[torch.Tensor, ... exported[key] = torch.ones_like(param.T) owners[key] = (index, None) + if mapped_indices and ( + missing := sorted(set(range(len(params))) - mapped_indices) + ): + raise TrainerRankSlotStateError( + f"Cannot map optimizer padding for checkpoint slot {name!r}: " + f"parameter indices {missing} do not belong to installed LoRA sites." + ) + canonical = self.runtime.model_support_handler.canonicalize_loaded_lora_state( exported, self.runtime.model ) diff --git a/tests/integration/megatron/lora/test_dynamic_lora_slots.py b/tests/integration/megatron/lora/test_dynamic_lora_slots.py index 44cb6655c..d2a9e2fb6 100644 --- a/tests/integration/megatron/lora/test_dynamic_lora_slots.py +++ b/tests/integration/megatron/lora/test_dynamic_lora_slots.py @@ -288,14 +288,35 @@ def _assert_distributed_optimizer_restore(device: torch.device) -> None: @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required.") def test_restored_dynamic_optimizer_canonicalizes_internal_padding() -> None: with _single_rank_model_parallel(): - _assert_restored_dynamic_optimizer_canonicalizes_internal_padding() + for num_local_experts in (1, 2): + _assert_restored_dynamic_optimizer_canonicalizes_internal_padding( + num_local_experts + ) -def _assert_restored_dynamic_optimizer_canonicalizes_internal_padding() -> None: +def _assert_restored_dynamic_optimizer_canonicalizes_internal_padding( + num_local_experts: int, +) -> None: device = torch.device("cuda") ref = LoRASlotRef("checkpoint", "A") - adapter = _adapter("dense", rank=2, seed=17) - lora = LoRA("dense", 4, 5, 2, 32, torch.float32, device) + prefix = "dense" if num_local_experts == 1 else "experts.{expert}" + adapter = { + key: value + for expert in range(num_local_experts) + for key, value in _adapter( + prefix.format(expert=expert), rank=2, seed=17 + expert + ).items() + } + lora = LoRA( + prefix, + 4, + 5, + 2, + 32, + torch.float32, + device, + num_local_experts=num_local_experts, + ) lora.load_lora_slot(ref, adapter, requires_grad=True) trainer = _trainer_for(lora, device) @@ -308,8 +329,8 @@ def canonicalize( return result trainer.runtime.model_support_handler.canonicalize_loaded_lora_state = canonicalize - with use_lora_slot(ref): - lora(torch.randn(3, 4, device=device)).sum().backward() + for param in trainer._checkpoint_slot_params_by_name["A"]: + param.grad = torch.ones_like(param) trainer.optim_step( params=AdamParams(learning_rate=1e-3, weight_decay=0.1, grad_clip_norm=0.0), checkpoints=["A"], From 0dc552309e21a5a3eac332527d4be6b1cf6ca855 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 16 Jul 2026 15:45:33 -0600 Subject: [PATCH 06/10] test: bound retained-slot GPU setup --- dev/trainer_rank_check.py | 7 ++++++- dev/trainer_rank_support.py | 15 ++++++++++++++- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/dev/trainer_rank_check.py b/dev/trainer_rank_check.py index 57a0b3c18..204146012 100644 --- a/dev/trainer_rank_check.py +++ b/dev/trainer_rank_check.py @@ -396,7 +396,12 @@ def _performance( if workload == "austin": families, prefix_tokens, branches, completion_tokens = 30, 5000, 16, 100 rank = TrainerRank(runtime, shared_prefix_max_depth=depth, head_chunk_tokens=8192) - slot_names = load_random_checkpoint_slots(runtime, rank, slots) + slot_names = load_random_checkpoint_slots( + runtime, + rank, + slots, + site_limit=1 if workload == "unequal_slots" else None, + ) if workload == "unequal_slots": if len(slot_names) < 2: raise ValueError("--workload unequal_slots requires --slots >= 2") diff --git a/dev/trainer_rank_support.py b/dev/trainer_rank_support.py index 6bc30ce9f..2d9e07423 100644 --- a/dev/trainer_rank_support.py +++ b/dev/trainer_rank_support.py @@ -12,6 +12,7 @@ def load_random_checkpoint_slots( count: int, *, lora_rank: int = 8, + site_limit: int | None = None, ) -> tuple[str, ...]: assert count >= 0, "slots must be >= 0" if count == 0: @@ -23,12 +24,24 @@ def load_random_checkpoint_slots( gathered, LoRAPublishPlanner(runtime.model).global_metadata({}) ) metadata = {meta.key: meta for values in gathered if values for meta in values} + selected = sorted(metadata.values(), key=lambda item: item.key) + if site_limit is not None: + pairs = [] + for meta in selected: + if ".lora_A." not in meta.key or ".experts." in meta.key: + continue + b_key = meta.key.replace(".lora_A.", ".lora_B.") + if b_meta := metadata.get(b_key): + pairs.append((meta, b_meta)) + selected = [meta for pair in pairs[:site_limit] for meta in pair] + if not selected: + raise RuntimeError("No replicated LoRA sites are available for the check") dtype = next(runtime.model[0].parameters()).dtype names = tuple(f"S{index}" for index in range(count)) for index, name in enumerate(names): generator = torch.Generator(device=rank.device).manual_seed(index + 1) adapter: dict[str, torch.Tensor] = {} - for meta in sorted(metadata.values(), key=lambda item: item.key): + for meta in selected: shape = list(meta.shape) if meta.manifest["sharded"]: axis = int(meta.manifest["export_shard_dim"]) From 78ec283a19f7592154080d08aece15ef69ce98bc Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 16 Jul 2026 15:49:26 -0600 Subject: [PATCH 07/10] test: trace unequal-slot GPU phases --- dev/trainer_rank_check.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/dev/trainer_rank_check.py b/dev/trainer_rank_check.py index 204146012..3981f5c36 100644 --- a/dev/trainer_rank_check.py +++ b/dev/trainer_rank_check.py @@ -77,6 +77,8 @@ def main( ), print_env=dist.get_rank() == 0, ) + if mode == "performance" and workload == "unequal_slots": + _trace_unequal_slots("runtime_ready") for chunk in runtime.model: chunk.eval() if mode == "correctness": @@ -396,12 +398,16 @@ def _performance( if workload == "austin": families, prefix_tokens, branches, completion_tokens = 30, 5000, 16, 100 rank = TrainerRank(runtime, shared_prefix_max_depth=depth, head_chunk_tokens=8192) + if workload == "unequal_slots": + _trace_unequal_slots("rank_ready") slot_names = load_random_checkpoint_slots( runtime, rank, slots, site_limit=1 if workload == "unequal_slots" else None, ) + if workload == "unequal_slots": + _trace_unequal_slots("slots_ready") if workload == "unequal_slots": if len(slot_names) < 2: raise ValueError("--workload unequal_slots requires --slots >= 2") @@ -425,9 +431,13 @@ def _performance( ) dp_rank, dp_size = rank._dp_rank_and_size() plan = rank._plan_flat_forward(requests) + if workload == "unequal_slots": + _trace_unequal_slots("plan_ready") assert workload != "austin" or plan.packed_tokens == 198_000 def step() -> list[MicroBatchStats]: + if workload == "unequal_slots": + _trace_unequal_slots("step_start") rank.zero_grad() stats: list[MicroBatchStats] = [] if adaptive: @@ -436,7 +446,11 @@ def step() -> list[MicroBatchStats]: stats.append(micro.stats) else: outputs = rank.dp_rank_forward(requests[dp_rank::dp_size]) + if workload == "unequal_slots": + _trace_unequal_slots("forward_ready") _output_loss(outputs).backward() + if workload == "unequal_slots": + _trace_unequal_slots("backward_ready") if optimizer_step: if not slot_names: raise ValueError("--optimizer-step requires --slots >= 1") @@ -478,6 +492,13 @@ def step() -> list[MicroBatchStats]: } +def _trace_unequal_slots(event: str) -> None: + print( + json.dumps({"event": event, "rank": dist.get_rank(), "time": time.time()}), + flush=True, + ) + + def _output_loss(outputs: Iterable[ForwardOutput]) -> torch.Tensor: terms: list[torch.Tensor] = [] for output in outputs: From 9bfe4ef8af8d89914b1febb300aa7a064c18edaf Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 16 Jul 2026 16:13:42 -0600 Subject: [PATCH 08/10] fix: reject unsafe TrainerRank DP EP topology --- src/art/trainer_rank/__init__.py | 24 +++++++++++++++++++++- tests/unit/test_trainer_rank_validation.py | 16 +++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/src/art/trainer_rank/__init__.py b/src/art/trainer_rank/__init__.py index b22062bdf..dfcee53f3 100644 --- a/src/art/trainer_rank/__init__.py +++ b/src/art/trainer_rank/__init__.py @@ -1751,6 +1751,7 @@ def _execute_flat_plan(self, plan: _FlatForwardPlan) -> list[AnyForwardOutput]: outputs = [ ForwardOutput(None, None, None, None) for _ in range(plan.request_count) ] + self._validate_hybridep_topology() hybridep = ( self._configure_hybridep( tuple(group.packed for group in plan.groups), topology=self._topology() @@ -2509,9 +2510,11 @@ def _configure_hybridep( ) -> tuple[tuple[int, ...], int] | None: from megatron.core import parallel_state as ps - if int(ps.get_expert_model_parallel_world_size()) <= 1: + expert_parallel_size = int(ps.get_expert_model_parallel_world_size()) + if expert_parallel_size <= 1: self._hybridep_graph_tracking = False return None + self._validate_hybridep_topology(topology) if not batches: return None from megatron.core.transformer.moe import fused_a2a @@ -2557,6 +2560,25 @@ def _configure_hybridep( self._hybridep_graph_tracking = True return rows, high_water + def _validate_hybridep_topology( + self, + topology: "ParallelTopology | None" = None, + ) -> None: + if topology is None: + configured_ep = int( + getattr(self.runtime.provider, "expert_model_parallel_size", 1) or 1 + ) + if configured_ep <= 1: + return + topology = self._topology() + if int(topology.dp) > 1: + raise NotImplementedError( + "TrainerRank does not support combining data parallelism with " + "expert parallelism because uneven DP inputs can desynchronize " + "HybridEP collectives. For MoE models, use DP=1 with CP and EP " + "set to the world size." + ) + @staticmethod def _set_hybridep_rows(rows: int) -> None: from art.megatron.train import _set_hybridep_token_count diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py index 066117eb7..754cd7777 100644 --- a/tests/unit/test_trainer_rank_validation.py +++ b/tests/unit/test_trainer_rank_validation.py @@ -251,6 +251,22 @@ def create_state(**kwargs: object) -> object: ) +@pytest.mark.skipif(find_spec("megatron") is None, reason="requires Megatron") +def test_hybridep_rejects_data_parallel_topology_before_empty_batch_return( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from art.megatron.context_parallel.types import ParallelTopology + + trainer = TrainerRank(_runtime()) + monkeypatch.setattr( + "megatron.core.parallel_state.get_expert_model_parallel_world_size", + lambda: 4, + ) + + with pytest.raises(NotImplementedError, match="DP=1"): + trainer._configure_hybridep((), topology=ParallelTopology(dp=4)) + + @pytest.mark.skipif(find_spec("megatron") is None, reason="requires Megatron") def test_hybridep_uses_maximum_cp_model_rows( monkeypatch: pytest.MonkeyPatch, From b1f51a674bdbf80f183bbae171f2290f1f301a2e Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 16 Jul 2026 16:20:00 -0600 Subject: [PATCH 09/10] test: cover TrainerRank DP EP guard entrypoint --- tests/unit/test_trainer_rank_validation.py | 23 +++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py index 754cd7777..415a2bb24 100644 --- a/tests/unit/test_trainer_rank_validation.py +++ b/tests/unit/test_trainer_rank_validation.py @@ -251,20 +251,21 @@ def create_state(**kwargs: object) -> object: ) -@pytest.mark.skipif(find_spec("megatron") is None, reason="requires Megatron") -def test_hybridep_rejects_data_parallel_topology_before_empty_batch_return( +@pytest.mark.parametrize("dp", [1, 4]) +def test_hybridep_validates_topology_for_empty_forward( monkeypatch: pytest.MonkeyPatch, + dp: int, ) -> None: - from art.megatron.context_parallel.types import ParallelTopology - - trainer = TrainerRank(_runtime()) - monkeypatch.setattr( - "megatron.core.parallel_state.get_expert_model_parallel_world_size", - lambda: 4, - ) + runtime = _runtime() + runtime.provider.expert_model_parallel_size = 4 + trainer = TrainerRank(runtime) + monkeypatch.setattr(trainer, "_topology", lambda: SimpleNamespace(dp=dp, cp=4)) - with pytest.raises(NotImplementedError, match="DP=1"): - trainer._configure_hybridep((), topology=ParallelTopology(dp=4)) + if dp > 1: + with pytest.raises(NotImplementedError, match="DP=1"): + trainer.dp_rank_forward([]) + else: + assert trainer.dp_rank_forward([]) == [] @pytest.mark.skipif(find_spec("megatron") is None, reason="requires Megatron") From d43cb5f691c6acb20055c8541ee3c1dd4d93fed3 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 16 Jul 2026 17:13:15 -0600 Subject: [PATCH 10/10] fix: avoid release REST outage in CI --- .github/workflows/prek.yml | 39 ++++----------- scripts/ci/github_release_assets.py | 78 +++++++++++++++++++++++++++++ 2 files changed, 89 insertions(+), 28 deletions(-) create mode 100644 scripts/ci/github_release_assets.py diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml index 88182576f..b5b649b4c 100644 --- a/.github/workflows/prek.yml +++ b/.github/workflows/prek.yml @@ -51,18 +51,9 @@ jobs: run: | fingerprint="${{ steps.fingerprint.outputs.fingerprint }}" part_prefix="${CI_UV_CACHE_ASSET_PREFIX}-${fingerprint}.tar.zst.part-" - release_api="https://api.github.com/repos/${GITHUB_REPOSITORY}/releases/tags/${CI_UV_CACHE_RELEASE_TAG}" - - release_json="$(curl -fsSL \ - -H "Authorization: Bearer ${GITHUB_TOKEN}" \ - -H "Accept: application/vnd.github+json" \ - "${release_api}" || true)" - - if [ -z "${release_json}" ]; then - echo "Cache release '${CI_UV_CACHE_RELEASE_TAG}' not found." - echo "cache-hit=false" >> "${GITHUB_OUTPUT}" - exit 0 - fi + release_json="$(python3 scripts/ci/github_release_assets.py \ + --repository "${GITHUB_REPOSITORY}" \ + --tag "${CI_UV_CACHE_RELEASE_TAG}")" hit="$(RELEASE_JSON="${release_json}" PART_PREFIX="${part_prefix}" python3 -c " import json, os, re @@ -73,7 +64,7 @@ jobs: int(m.group(1)) for a in payload.get('assets', []) for m in [pattern.match(a.get('name', ''))] - if m and a.get('id') is not None + if m and a.get('url') ) print('true' if parts and parts == list(range(len(parts))) else 'false') ")" @@ -147,22 +138,15 @@ jobs: env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} run: | - release_api="https://api.github.com/repos/${GITHUB_REPOSITORY}/releases/tags/${CI_UV_CACHE_RELEASE_TAG}" fingerprint="${{ needs.cache-status.outputs.fingerprint }}" part_prefix="${CI_UV_CACHE_ASSET_PREFIX}-${fingerprint}.tar.zst.part-" - release_json="$(curl -fsSL \ - -H "Authorization: Bearer ${GITHUB_TOKEN}" \ - -H "Accept: application/vnd.github+json" \ - "${release_api}" || true)" - - if [ -z "${release_json}" ]; then - echo "::error::Missing cache release '${CI_UV_CACHE_RELEASE_TAG}'." - exit 1 - fi + release_json="$(python3 scripts/ci/github_release_assets.py \ + --repository "${GITHUB_REPOSITORY}" \ + --tag "${CI_UV_CACHE_RELEASE_TAG}")" part_selection_file="/tmp/uv-cache-part-selection.txt" - if ! RELEASE_JSON="${release_json}" PART_PREFIX="${part_prefix}" python3 -c "import json, os, re, sys; payload=json.loads(os.environ['RELEASE_JSON']); part_prefix=os.environ['PART_PREFIX']; pattern=re.compile(r'^' + re.escape(part_prefix) + r'(\\d{3})$'); parts=[]; [parts.append((int(m.group(1)), int(a.get('id')), a.get('name'))) for a in payload.get('assets', []) for m in [pattern.match(a.get('name', ''))] if m and a.get('id') is not None]; parts.sort(key=lambda x: x[0]); indices=[p[0] for p in parts]; expected=list(range(len(parts))); print('\\n'.join(f'{asset_id} {name}' for _, asset_id, name in parts)) if parts and indices == expected else (_ for _ in ()).throw(SystemExit(2 if not parts else 3))" > "${part_selection_file}"; then + if ! RELEASE_JSON="${release_json}" PART_PREFIX="${part_prefix}" python3 -c "import json, os, re, sys; payload=json.loads(os.environ['RELEASE_JSON']); part_prefix=os.environ['PART_PREFIX']; pattern=re.compile(r'^' + re.escape(part_prefix) + r'(\\d{3})$'); parts=[]; [parts.append((int(m.group(1)), a.get('url'), a.get('name'))) for a in payload.get('assets', []) for m in [pattern.match(a.get('name', ''))] if m and a.get('url')]; parts.sort(key=lambda x: x[0]); indices=[p[0] for p in parts]; expected=list(range(len(parts))); print('\\n'.join(f'{url} {name}' for _, url, name in parts)) if parts and indices == expected else (_ for _ in ()).throw(SystemExit(2 if not parts else 3))" > "${part_selection_file}"; then echo "::error::No complete uv cache part set found for prefix '${part_prefix}'." exit 1 fi @@ -176,15 +160,14 @@ jobs: mkdir -p "${parts_dir}" awk -v d="${parts_dir}" '{print d "/" $2}' "${part_selection_file}" > "${part_paths_file}" - PARTS_DIR="${parts_dir}" GITHUB_TOKEN="${GITHUB_TOKEN}" GITHUB_REPOSITORY="${GITHUB_REPOSITORY}" \ + PARTS_DIR="${parts_dir}" GITHUB_TOKEN="${GITHUB_TOKEN}" \ xargs -n 2 -P 8 sh -c ' - asset_id="$1" + asset_url="$1" asset_name="$2" part_path="${PARTS_DIR}/${asset_name}" curl -fsSL -L \ -H "Authorization: Bearer ${GITHUB_TOKEN}" \ - -H "Accept: application/octet-stream" \ - "https://api.github.com/repos/${GITHUB_REPOSITORY}/releases/assets/${asset_id}" \ + "${asset_url}" \ -o "${part_path}" ' sh < "${part_selection_file}" diff --git a/scripts/ci/github_release_assets.py b/scripts/ci/github_release_assets.py new file mode 100644 index 000000000..5cb01099a --- /dev/null +++ b/scripts/ci/github_release_assets.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python3 +"""List GitHub release assets without the REST release endpoint.""" + +from __future__ import annotations + +import argparse +import json +import os +from urllib import request + +_QUERY = """ +query($owner: String!, $name: String!, $tag: String!, $cursor: String) { + repository(owner: $owner, name: $name) { + release(tagName: $tag) { + releaseAssets(first: 100, after: $cursor) { + nodes { name downloadUrl } + pageInfo { hasNextPage endCursor } + } + } + } +} +""" + + +def release_assets(repository: str, tag: str, token: str) -> list[dict[str, str]]: + owner, name = repository.split("/", 1) + cursor: str | None = None + assets: list[dict[str, str]] = [] + while True: + body = json.dumps( + { + "query": _QUERY, + "variables": { + "owner": owner, + "name": name, + "tag": tag, + "cursor": cursor, + }, + } + ).encode() + http_request = request.Request( + "https://api.github.com/graphql", + data=body, + headers={ + "Accept": "application/vnd.github+json", + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + }, + ) + with request.urlopen(http_request, timeout=30) as response: + payload = json.load(response) + if errors := payload.get("errors"): + raise RuntimeError(f"GitHub GraphQL error: {errors}") + release = payload["data"]["repository"]["release"] + if release is None: + return [] + page = release["releaseAssets"] + assets.extend( + {"name": node["name"], "url": node["downloadUrl"]} for node in page["nodes"] + ) + if not page["pageInfo"]["hasNextPage"]: + return assets + cursor = page["pageInfo"]["endCursor"] + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--repository", required=True) + parser.add_argument("--tag", required=True) + args = parser.parse_args() + token = os.environ.get("GITHUB_TOKEN") + if not token: + raise SystemExit("GITHUB_TOKEN is required") + print(json.dumps({"assets": release_assets(args.repository, args.tag, token)})) + + +if __name__ == "__main__": + main()