diff --git a/dev/trainer_rank.py b/dev/trainer_rank.py new file mode 100644 index 000000000..0dd8b24bc --- /dev/null +++ b/dev/trainer_rank.py @@ -0,0 +1,92 @@ +from itertools import islice +import os + +import torch +import torch.distributed as dist +from trainer_rank_support import load_random_checkpoint_slots +from transformers import AutoTokenizer +import typer + +from art.trainer_rank import AdamParams, ForwardInput, TrainerRank + + +def main( + model: str = "Qwen/Qwen3-0.6B", + samples: int = 16, + steps: int = 1, + lr: float = 5e-5, + layers: int = 2, + lora_rank: int = 8, + max_seq_length: int = 256, +) -> None: + os.environ.setdefault("ART_MEGATRON_TENSOR_MODEL_PARALLEL_SIZE", "1") + os.environ.setdefault("ART_MEGATRON_CONTEXT_PARALLEL_SIZE", "1") + os.environ.setdefault("ART_MEGATRON_PIPELINE_MODEL_PARALLEL_SIZE", "1") + + if not torch.cuda.is_available(): + raise RuntimeError("dev/trainer_rank.py requires CUDA") + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + dist.init_process_group(backend="nccl") + + try: + from datasets import load_dataset + + from art.megatron import train as megatron_train + + tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True) + inputs: list[ForwardInput[torch.Tensor, None, None, None]] = [] + rows = load_dataset("roneneldan/TinyStories", split="train", streaming=True) + for row in islice(rows, samples): + token_ids = tokenizer( + str(row["text"]), # type: ignore[index] + add_special_tokens=True, + truncation=True, + max_length=max_seq_length + 1, + return_tensors="pt", + )["input_ids"].reshape(-1) + inputs.append( + ForwardInput( + input_tokens=token_ids[:-1], + target_tokens=token_ids[1:], + ) + ) + + runtime = megatron_train.build_training_runtime( + model_identifier=model, + provider_configure=lambda provider: setattr(provider, "num_layers", layers), + print_env=dist.get_rank() == 0, + ) + rank = TrainerRank(runtime) + (slot,) = load_random_checkpoint_slots(runtime, rank, 1, lora_rank=lora_rank) + rank.set_checkpoint(slot) + + for step in range(steps): + loss_sum = torch.tensor(0.0, device=rank.device) + token_count = torch.tensor(0.0, device=rank.device) + for micro in rank.forward_micro_batches(inputs): + loss = torch.tensor(0.0, device=rank.device) + for output in micro.outputs: + assert output.target_logprobs is not None + loss = loss - output.target_logprobs.sum() + token_count += output.target_logprobs.numel() + loss.backward() + loss_sum += loss.detach() + + rank.dp_reduce(loss_sum) + rank.dp_reduce(token_count) + scale = 1.0 / max(float(token_count.item()), 1.0) + metrics = rank.optim_step( + params=AdamParams(learning_rate=lr), + scale_grads=scale, + ) + metrics["loss"] = float(loss_sum.item() * scale) + metrics["tokens"] = float(token_count.item()) + if dist.get_rank() == 0: + print(f"step={step} {metrics}", flush=True) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +if __name__ == "__main__": + typer.run(main) diff --git a/dev/trainer_rank_check.py b/dev/trainer_rank_check.py new file mode 100644 index 000000000..54410613d --- /dev/null +++ b/dev/trainer_rank_check.py @@ -0,0 +1,596 @@ +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +import json +import os +import statistics +import time +from typing import Any, Literal, cast + +import torch +import torch.distributed as dist +from trainer_rank_support import load_random_checkpoint_slots +import typer + +from art.megatron.prefix_tree_packing import prefix_tree_pack +from art.trainer_rank import ( + AdamParams, + ForwardInput, + ForwardOutput, + MicroBatchStats, + TopK, + TrainerRank, + Unset, +) + + +@dataclass(frozen=True) +class Diff: + mean_abs_pct: float = 0.0 + max_abs_diff: float = 0.0 + + def merge(self, other: Diff) -> Diff: + return Diff( + max(self.mean_abs_pct, other.mean_abs_pct), + max(self.max_abs_diff, other.max_abs_diff), + ) + + +def main( + mode: Literal["correctness", "performance"] = "correctness", + model: str = "Qwen/Qwen3-0.6B", + layers: int = 1, + depths: str = "0,1,2,3,4", + performance_depth: int = 1, + chunks: str = "17,512,8192", + workload: Literal["regular", "austin", "varied"] = "regular", + request: Literal["target", "multi", "topk", "logits", "hidden", "mixed"] = "target", + families: int = 8, + prefix_tokens: int = 128, + branches: int = 4, + completion_tokens: int = 32, + slots: int = 0, + adaptive: bool = False, + optimizer_step: bool = False, + warmup: int = 3, + repeat: int = 10, + output_jsonl: str = "", +) -> None: + os.environ.setdefault("ART_MEGATRON_TENSOR_MODEL_PARALLEL_SIZE", "1") + os.environ.setdefault("ART_MEGATRON_CONTEXT_PARALLEL_SIZE", "1") + os.environ.setdefault("ART_MEGATRON_PIPELINE_MODEL_PARALLEL_SIZE", "1") + if not torch.cuda.is_available(): + raise RuntimeError("dev/trainer_rank_check.py requires CUDA") + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + dist.init_process_group(backend="nccl") + try: + from art.megatron import train as megatron_train + + torch.manual_seed(1234) + runtime = megatron_train.build_training_runtime( + model_identifier=model, + provider_configure=( + (lambda provider: setattr(provider, "num_layers", layers)) + if layers > 0 + else None + ), + print_env=dist.get_rank() == 0, + ) + for chunk in runtime.model: + chunk.eval() + if mode == "correctness": + payload = _correctness( + runtime, + depths=_ints(depths), + chunks=_ints(chunks), + slots=slots, + ) + else: + payload = _performance( + runtime, + depth=performance_depth, + workload=workload, + request=request, + families=families, + prefix_tokens=prefix_tokens, + branches=branches, + completion_tokens=completion_tokens, + slots=slots, + adaptive=adaptive, + optimizer_step=optimizer_step, + warmup=warmup, + repeat=repeat, + ) + payload.update(_topology(), model=model, layers=layers, mode=mode) + if dist.get_rank() == 0: + line = json.dumps(payload, sort_keys=True) + print(line, flush=True) + if output_jsonl: + with open(output_jsonl, "a", encoding="utf-8") as output: + output.write(line + "\n") + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def _correctness( + runtime: Any, + *, + depths: tuple[int, ...], + chunks: tuple[int, ...], + slots: int, +) -> dict[str, object]: + assert depths and chunks, "depths and chunks must not be empty" + slot_names = load_random_checkpoint_slots(runtime, TrainerRank(runtime), slots) + requests = _correctness_requests(slot_names) + reference = _global_outputs( + TrainerRank( + runtime, + shared_prefix_max_depth=0, + head_chunk_tokens=max(chunks), + ), + requests, + ) + worst = Diff() + grad_worst = Diff() + rows: list[dict[str, object]] = [] + for depth in depths: + depth_reference: list[dict[str, object]] | None = None + for chunk_tokens in chunks: + rank = TrainerRank( + runtime, + shared_prefix_max_depth=depth, + head_chunk_tokens=chunk_tokens, + ) + outputs = _global_outputs(rank, requests) + if dist.get_rank() == 0: + assert reference is not None and outputs is not None + independent_diff = _compare_outputs( + outputs, + reference, + tolerance=5e-3, + ) + chunk_diff = ( + Diff() + if depth_reference is None + else _compare_outputs(outputs, depth_reference, tolerance=2e-5) + ) + depth_reference = outputs + worst = worst.merge(independent_diff).merge(chunk_diff) + rows.append( + { + "depth": depth, + "head_chunk_tokens": chunk_tokens, + "independent_mean_abs_pct": independent_diff.mean_abs_pct, + "chunk_mean_abs_pct": chunk_diff.mean_abs_pct, + } + ) + print(rows[-1], flush=True) + grad_diff = _head_backward_chunk_parity( + runtime, requests, depth=depth, chunks=chunks + ) + grad_worst = grad_worst.merge(grad_diff) + return { + "request_combinations": 16, + "slots": slots, + "rows": rows, + "mean_abs_pct": worst.mean_abs_pct, + "max_abs_diff": worst.max_abs_diff, + "head_backward_mean_abs_pct": grad_worst.mean_abs_pct, + "head_backward_max_abs_diff": grad_worst.max_abs_diff, + } + + +def _local_outputs( + rank: TrainerRank, + indexed_requests: Sequence[tuple[int, ForwardInput]], +) -> list[dict[str, object]]: + from art.megatron.lora import use_lora_slot + + requests = [request for _, request in indexed_requests] + plan = rank._plan_flat_forward(requests) + outputs: list[ForwardOutput] = [ + ForwardOutput(None, None, None, None) for _ in requests + ] + sources: list[torch.Tensor] = [torch.empty(0, dtype=torch.long) for _ in requests] + for group in plan.groups: + prepared = rank._prepare_packed_forward(group.packed) + with use_lora_slot(group.slot_ref): + group_outputs = rank._forward_packed(group.items, prepared) + for index, source, output in zip( + group.request_indices, + prepared.source_positions_by_item, + group_outputs, + strict=True, + ): + sources[index] = source + outputs[index] = output + return [ + _output_record(global_index, source, output) + for (global_index, _), source, output in zip( + indexed_requests, + sources, + outputs, + strict=True, + ) + ] + + +def _global_outputs( + rank: TrainerRank, + requests: Sequence[ForwardInput], +) -> list[dict[str, object]] | None: + from megatron.core import parallel_state as ps + + dp_rank = int(ps.get_data_parallel_rank()) + dp_size = int(ps.get_data_parallel_world_size()) + indexed = list(enumerate(requests))[dp_rank::dp_size] + local = _local_outputs(rank, indexed) + gathered: list[list[dict[str, object]] | None] = [None] * dist.get_world_size() + dist.all_gather_object(gathered, local) + if dist.get_rank() != 0: + return None + records = [ + record for rank_records in gathered if rank_records for record in rank_records + ] + return [ + _reconstruct(index, request, records) for index, request in enumerate(requests) + ] + + +def _output_record( + index: int, + source_positions: torch.Tensor, + output: ForwardOutput, +) -> dict[str, object]: + return { + "index": index, + "source": source_positions.detach().cpu(), + "target": _cpu(output.target_logprobs), + "topk_logprobs": _cpu(None if output.top_k is None else output.top_k.logprobs), + "topk_tokens": _cpu(None if output.top_k is None else output.top_k.tokens), + "logits": _cpu(output.logits), + "hidden": _cpu(output.hidden_states), + } + + +def _reconstruct( + index: int, + request: ForwardInput, + records: Sequence[dict[str, object]], +) -> dict[str, object]: + selected = [record for record in records if record["index"] == index] + return { + key: _reconstruct_tensor( + selected, + key, + length=int(request.input_tokens.numel()), + ) + for key in ("target", "topk_logprobs", "topk_tokens", "logits", "hidden") + } + + +def _reconstruct_tensor( + records: Sequence[dict[str, object]], + key: str, + *, + length: int, +) -> torch.Tensor | None: + values = [ + record[key] for record in records if isinstance(record[key], torch.Tensor) + ] + if not values: + return None + first = cast(torch.Tensor, values[0]) + output = torch.empty((length, *first.shape[1:]), dtype=first.dtype) + filled = torch.zeros(length, dtype=torch.bool) + for record in records: + value = record[key] + if not isinstance(value, torch.Tensor): + continue + source = cast(torch.Tensor, record["source"]) + duplicate = filled.index_select(0, source) + if bool(duplicate.any()): + torch.testing.assert_close( + output.index_select(0, source[duplicate]), + value[duplicate], + atol=2e-5, + rtol=2e-5, + ) + output[source] = value + filled[source] = True + assert bool(filled.all()), f"{key} reconstruction missed positions" + return output + + +def _compare_outputs( + actual: Sequence[dict[str, object]], + expected: Sequence[dict[str, object]], + *, + tolerance: float, +) -> Diff: + worst = Diff() + for actual_output, expected_output in zip(actual, expected, strict=True): + for key in actual_output: + actual_tensor = actual_output[key] + expected_tensor = expected_output[key] + if actual_tensor is None or expected_tensor is None: + if actual_tensor is not expected_tensor: + raise AssertionError(f"{key} None mismatch") + continue + assert isinstance(actual_tensor, torch.Tensor) + assert isinstance(expected_tensor, torch.Tensor) + if key == "topk_tokens": + if not torch.equal( + actual_tensor.sort(dim=-1).values, + expected_tensor.sort(dim=-1).values, + ): + raise AssertionError("top-k token sets differ") + continue + diff = _diff(actual_tensor, expected_tensor) + if diff.mean_abs_pct > tolerance: + raise AssertionError( + f"{key} mean_abs_pct={diff.mean_abs_pct} exceeds {tolerance}" + ) + worst = worst.merge(diff) + return worst + + +def _head_backward_chunk_parity( + runtime: Any, + requests: Sequence[ForwardInput], + *, + depth: int, + chunks: Sequence[int], +) -> Diff: + active = [ + request + for request in requests + if request.target_tokens is not None or request.top_k is not None + ] + rank = TrainerRank( + runtime, + shared_prefix_max_depth=depth, + head_chunk_tokens=chunks[0], + ) + items = [rank._forward_item(request) for request in active] + prepared = rank._prepare_packed_forward( + prefix_tree_pack( + (item.input_ids for item in items), + max_depth=depth, + ) + ) + with torch.no_grad(): + hidden = rank._gather_sequence_parallel_hidden(rank._decoder_hidden(prepared)) + gradients: list[torch.Tensor] = [] + for chunk_tokens in (chunks[0], chunks[-1]): + rank.head_chunk_tokens = chunk_tokens + candidate = hidden.detach().requires_grad_(True) + outputs = rank._project_head(items, prepared, candidate) + _output_loss(outputs).backward() + assert candidate.grad is not None + gradients.append(candidate.grad) + diff = _diff(gradients[0], gradients[1]) + if diff.mean_abs_pct > 2e-3: + raise AssertionError(f"head gradient mean_abs_pct={diff.mean_abs_pct}") + return diff + + +def _performance( + runtime: Any, + *, + depth: int, + workload: str, + request: str, + families: int, + prefix_tokens: int, + branches: int, + completion_tokens: int, + slots: int, + adaptive: bool, + optimizer_step: bool, + warmup: int, + repeat: int, +) -> dict[str, object]: + 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) + 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 + + def step() -> list[MicroBatchStats]: + rank.zero_grad() + stats: list[MicroBatchStats] = [] + if adaptive: + for micro in rank.forward_micro_batches(requests): + _output_loss(cast(Sequence[ForwardOutput], micro.outputs)).backward() + stats.append(micro.stats) + else: + outputs = rank.dp_rank_forward(requests[dp_rank::dp_size]) + _output_loss(outputs).backward() + if optimizer_step: + if not slot_names: + raise ValueError("--optimizer-step requires --slots >= 1") + rank.optim_step(params=AdamParams(learning_rate=1e-5)) + return stats + + for _ in range(warmup): + step() + times: list[float] = [] + all_stats: list[MicroBatchStats] = [] + torch.cuda.reset_peak_memory_stats() + for _ in range(repeat): + torch.cuda.synchronize() + started = time.perf_counter() + all_stats.extend(step()) + torch.cuda.synchronize() + times.append(time.perf_counter() - started) + median = statistics.median(times) + free, total = torch.cuda.mem_get_info() + return { + "depth": depth, + "workload": workload, + "request": request, + "adaptive": adaptive, + "optimizer_step": optimizer_step, + "slots": slots, + "warmup": warmup, + "repeat": repeat, + "packed_tokens": plan.packed_tokens, + "logical_tokens": plan.logical_tokens, + "median_s": median, + "packed_tok_s": plan.packed_tokens / median, + "logical_tok_s": plan.logical_tokens / median, + "peak_allocated_gb": torch.cuda.max_memory_allocated() / 1024**3, + "peak_reserved_gb": torch.cuda.max_memory_reserved() / 1024**3, + "device_used_gb": (total - free) / 1024**3, + "windows": [stat.global_count for stat in all_stats], + "rejected_candidates": sum(stat.rejected_candidates for stat in all_stats), + } + + +def _output_loss(outputs: Iterable[ForwardOutput]) -> torch.Tensor: + terms: list[torch.Tensor] = [] + for output in outputs: + if output.target_logprobs is not None: + terms.append(-output.target_logprobs.float().sum()) + if output.top_k is not None: + terms.append(-output.top_k.logprobs.float().sum()) + if output.logits is not None: + terms.append(output.logits.float().square().mean()) + if output.hidden_states is not None: + terms.append(output.hidden_states.float().square().mean()) + if not terms: + raise RuntimeError("request produced no differentiable outputs") + return torch.stack(terms).sum() + + +def _correctness_requests(slots: Sequence[str] = ()) -> list[ForwardInput]: + requests: list[ForwardInput] = [] + for mask in range(16): + tokens = torch.tensor( + [11, 12, 20 + mask // 8, 30 + mask // 4 % 2, 40 + mask // 2 % 2, 50 + mask] + ) + tokens = tokens.reshape(2, 3) if mask == 15 else tokens + labels: torch.Tensor | None = None + if mask & 1: + labels = (tokens * 7 + mask) % 1000 + if mask == 1: + labels = torch.stack((labels, (labels + 17) % 1000), dim=1) + labels[2, 1] = -100 + requests.append( + ForwardInput( + input_tokens=tokens, + target_tokens=labels, + top_k=3 if mask & 2 else None, + logits=bool(mask & 4), + hidden_states=bool(mask & 8), + checkpoint=slots[mask % len(slots)] if slots else Unset, + ) + ) + return requests + + +def _performance_requests( + *, + request: str, + families: int, + prefix_tokens: int, + branches: int, + completion_tokens: int, + varied: bool, + slots: Sequence[str], +) -> list[ForwardInput]: + requests: list[ForwardInput] = [] + for family in range(families): + family_base = family * 10_000_019 + prefix_len = prefix_tokens + ((family * 97) % 257 - 128 if varied else 0) + prefix = _tokens(family_base, max(1, prefix_len)) + family_branches = max(1, branches + ((family % 5) - 2 if varied else 0)) + for branch in range(family_branches): + completion_len = completion_tokens + ( + (branch * 17) % 33 - 16 if varied else 0 + ) + tokens = torch.cat( + ( + prefix, + _tokens(family_base + branch * 1009 + 17, max(1, completion_len)), + ) + ) + labels = (tokens * 7 + 3) % 32_000 + labels[: int(prefix.numel())] = -100 + if request == "multi": + labels = torch.stack( + tuple((labels + offset) % 32_000 for offset in range(4)), dim=1 + ) + labels[: int(prefix.numel())] = -100 + requests.append( + ForwardInput( + input_tokens=tokens, + target_tokens=labels + if request in {"target", "multi", "mixed"} + else None, + top_k=10 if request in {"topk", "mixed"} else None, + logits=request == "logits" + or request == "mixed" + and branch % 16 == 0, + hidden_states=request == "hidden" + or request == "mixed" + and branch % 8 == 0, + checkpoint=slots[family % len(slots)] if slots else Unset, + ) + ) + return requests + + +def _tokens(offset: int, length: int) -> torch.Tensor: + return (torch.arange(length, dtype=torch.long) + offset) % 32_000 + 100 + + +def _diff(actual: torch.Tensor, expected: torch.Tensor) -> Diff: + assert actual.shape == expected.shape, ( + f"shape mismatch: {actual.shape} != {expected.shape}" + ) + if not actual.numel(): + return Diff() + delta = (actual.float() - expected.float()).abs() + return Diff( + float(delta.mean() / expected.float().abs().mean().clamp_min(1e-18)), + float(delta.max()), + ) + + +def _cpu(tensor: object) -> torch.Tensor | None: + return tensor.detach().cpu() if isinstance(tensor, torch.Tensor) else None + + +def _ints(value: str) -> tuple[int, ...]: + return tuple(int(item) for item in value.split(",") if item.strip()) + + +def _topology() -> dict[str, int]: + from megatron.core import parallel_state as ps + + return { + "world": dist.get_world_size(), + "dp": int(ps.get_data_parallel_world_size()), + "tp": int(ps.get_tensor_model_parallel_world_size()), + "cp": int(ps.get_context_parallel_world_size()), + "ep": int(ps.get_expert_model_parallel_world_size()), + } + + +if __name__ == "__main__": + typer.run(main) diff --git a/dev/trainer_rank_fast_check.py b/dev/trainer_rank_fast_check.py new file mode 100644 index 000000000..60857d3b5 --- /dev/null +++ b/dev/trainer_rank_fast_check.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from importlib.util import find_spec +import subprocess +import sys + +FAST_TESTS = ( + "tests/unit/test_trainer_rank_validation.py", + "tests/unit/test_trainer_rank_weird_shapes.py", + "tests/unit/test_prefix_tree_packing.py", +) + +MEGATRON_FAST_TESTS = ( + "tests/unit/test_prefix_tree.py", + "tests/unit/test_prefix_tree_attention_builder.py", + "tests/unit/test_prefix_tree_grad_parity.py", +) + + +def main() -> None: + tests = (*FAST_TESTS, *(MEGATRON_FAST_TESTS if find_spec("megatron") else ())) + command = [sys.executable, "-m", "pytest", "--tb=short", *tests, *sys.argv[1:]] + raise SystemExit(subprocess.call(command)) + + +if __name__ == "__main__": + main() diff --git a/dev/trainer_rank_support.py b/dev/trainer_rank_support.py new file mode 100644 index 000000000..6bc30ce9f --- /dev/null +++ b/dev/trainer_rank_support.py @@ -0,0 +1,51 @@ +from typing import Any + +import torch +import torch.distributed as dist + +from art.trainer_rank import TrainerRank + + +def load_random_checkpoint_slots( + runtime: Any, + rank: TrainerRank, + count: int, + *, + lora_rank: int = 8, +) -> tuple[str, ...]: + assert count >= 0, "slots must be >= 0" + if count == 0: + return () + from art.megatron.lora import LoRAPublishPlanner + + gathered: list[list[Any] | None] = [None] * dist.get_world_size() + dist.all_gather_object( + gathered, LoRAPublishPlanner(runtime.model).global_metadata({}) + ) + metadata = {meta.key: meta for values in gathered if values for meta in values} + 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): + shape = list(meta.shape) + if meta.manifest["sharded"]: + axis = int(meta.manifest["export_shard_dim"]) + shape[axis] = sum( + map( + int, + meta.manifest.get("component_sizes") + or [shape[axis] * int(meta.manifest["shard_world_size"])], + ) + ) + is_a = ".lora_A." in meta.key + shape[0 if is_a else -1] = lora_rank + tensor = torch.randn( + shape, device=rank.device, dtype=dtype, generator=generator + ) + adapter[meta.key] = tensor if is_a else tensor.mul_(1e-3) + assert rank.load_checkpoint_slot(name, adapter) > 0, ( + "TrainerRank check requires installed LoRA adapter sites" + ) + return names diff --git a/src/art/trainer_rank/__init__.py b/src/art/trainer_rank/__init__.py new file mode 100644 index 000000000..b22483d43 --- /dev/null +++ b/src/art/trainer_rank/__init__.py @@ -0,0 +1,2815 @@ +from __future__ import annotations + +from collections.abc import ( + Iterable, + Iterator, + Mapping, + Sequence, +) +from dataclasses import dataclass +import os +from typing import ( + TYPE_CHECKING, + Any, + Generic, + Literal, + TypeVar, + cast, + overload, +) +import weakref + +import torch +import torch.distributed as dist + +from art.megatron.prefix_tree_packing import ( + PrefixTreePack, + _local_position_pairs, + estimate_prefix_tree_packed_tokens, + prefix_tree_pack, +) + +if TYPE_CHECKING: + from megatron.core.models.gpt.gpt_model import GPTModel + from megatron.core.packed_seq_params import PackedSeqParams + + from art.megatron.context_parallel.types import ( + ArtContextParallelState, + ParallelTopology, + ) + from art.megatron.lora import LoRASlotRef + from art.megatron.prefix_tree_state import PrefixTreeAttentionState + from art.megatron.train import TrainingRuntime + + +@dataclass(frozen=True) +class AdamParams: + learning_rate: float + beta1: float = 0.9 + beta2: float = 0.99 + weight_decay: float = 0.1 + grad_clip_norm: float = 0.1 + + +@dataclass(frozen=True) +class TopK: + logprobs: torch.Tensor + tokens: torch.Tensor + + +LogprobsT = TypeVar("LogprobsT", bound=torch.Tensor | None, covariant=True) +TopKT = TypeVar("TopKT", bound=TopK | None, covariant=True) +LogitsT = TypeVar("LogitsT", bound=torch.Tensor | None, covariant=True) +HiddenStatesT = TypeVar("HiddenStatesT", bound=torch.Tensor | None, covariant=True) +T = TypeVar("T") + +_MEMORY_PROFILE_TRUST_GROWTH = 8 + + +class _Unset: + pass + + +Unset = _Unset() +type AdapterSelection = str | None | _Unset + + +@dataclass(frozen=True) +class _LocalLoRASlotRef: + kind: Literal["checkpoint", "lora"] + name: str | None + + +@dataclass(frozen=True) +class ForwardOutput(Generic[LogprobsT, TopKT, LogitsT, HiddenStatesT]): + target_logprobs: LogprobsT + top_k: TopKT + logits: LogitsT + hidden_states: HiddenStatesT + + +@dataclass(slots=True) +class ForwardInput(Generic[LogprobsT, TopKT, LogitsT, HiddenStatesT]): + input_tokens: torch.Tensor + target_tokens: torch.Tensor | None = None + top_k: int | None = None + logits: bool = False + hidden_states: bool = False + checkpoint: AdapterSelection = Unset + lora: AdapterSelection = Unset + + @overload + def __new__( + cls, + *, + input_tokens: torch.Tensor, + target_tokens: None = None, + top_k: None = None, + logits: Literal[False] = False, + hidden_states: Literal[False] = False, + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, + ) -> "ForwardInput[None, None, None, None]": ... + + @overload + def __new__( + cls, + *, + input_tokens: torch.Tensor, + target_tokens: torch.Tensor, + top_k: None = None, + logits: Literal[False] = False, + hidden_states: Literal[False] = False, + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, + ) -> "ForwardInput[torch.Tensor, None, None, None]": ... + + @overload + def __new__( + cls, + *, + input_tokens: torch.Tensor, + target_tokens: None = None, + top_k: int, + logits: Literal[False] = False, + hidden_states: Literal[False] = False, + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, + ) -> "ForwardInput[None, TopK, None, None]": ... + + @overload + def __new__( + cls, + *, + input_tokens: torch.Tensor, + target_tokens: None = None, + top_k: None = None, + logits: Literal[True], + hidden_states: Literal[False] = False, + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, + ) -> "ForwardInput[None, None, torch.Tensor, None]": ... + + @overload + def __new__( + cls, + *, + input_tokens: torch.Tensor, + target_tokens: None = None, + top_k: None = None, + logits: Literal[False] = False, + hidden_states: Literal[True], + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, + ) -> "ForwardInput[None, None, None, torch.Tensor]": ... + + @overload + def __new__( + cls, + *, + input_tokens: torch.Tensor, + target_tokens: torch.Tensor, + top_k: int, + logits: Literal[False] = False, + hidden_states: Literal[False] = False, + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, + ) -> "ForwardInput[torch.Tensor, TopK, None, None]": ... + + @overload + def __new__( + cls, + *, + input_tokens: torch.Tensor, + target_tokens: torch.Tensor, + top_k: None = None, + logits: Literal[True], + hidden_states: Literal[False] = False, + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, + ) -> "ForwardInput[torch.Tensor, None, torch.Tensor, None]": ... + + @overload + def __new__( + cls, + *, + input_tokens: torch.Tensor, + target_tokens: torch.Tensor, + top_k: None = None, + logits: Literal[False] = False, + hidden_states: Literal[True], + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, + ) -> "ForwardInput[torch.Tensor, None, None, torch.Tensor]": ... + + @overload + def __new__( + cls, + *, + input_tokens: torch.Tensor, + target_tokens: None = None, + top_k: int, + logits: Literal[True], + hidden_states: Literal[False] = False, + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, + ) -> "ForwardInput[None, TopK, torch.Tensor, None]": ... + + @overload + def __new__( + cls, + *, + input_tokens: torch.Tensor, + target_tokens: None = None, + top_k: int, + logits: Literal[False] = False, + hidden_states: Literal[True], + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, + ) -> "ForwardInput[None, TopK, None, torch.Tensor]": ... + + @overload + def __new__( + cls, + *, + input_tokens: torch.Tensor, + target_tokens: None = None, + top_k: None = None, + logits: Literal[True], + hidden_states: Literal[True], + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, + ) -> "ForwardInput[None, None, torch.Tensor, torch.Tensor]": ... + + @overload + def __new__( + cls, + *, + input_tokens: torch.Tensor, + target_tokens: torch.Tensor, + top_k: int, + logits: Literal[True], + hidden_states: Literal[False] = False, + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, + ) -> "ForwardInput[torch.Tensor, TopK, torch.Tensor, None]": ... + + @overload + def __new__( + cls, + *, + input_tokens: torch.Tensor, + target_tokens: torch.Tensor, + top_k: int, + logits: Literal[False] = False, + hidden_states: Literal[True], + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, + ) -> "ForwardInput[torch.Tensor, TopK, None, torch.Tensor]": ... + + @overload + def __new__( + cls, + *, + input_tokens: torch.Tensor, + target_tokens: torch.Tensor, + top_k: None = None, + logits: Literal[True], + hidden_states: Literal[True], + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, + ) -> "ForwardInput[torch.Tensor, None, torch.Tensor, torch.Tensor]": ... + + @overload + def __new__( + cls, + *, + input_tokens: torch.Tensor, + target_tokens: None = None, + top_k: int, + logits: Literal[True], + hidden_states: Literal[True], + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, + ) -> "ForwardInput[None, TopK, torch.Tensor, torch.Tensor]": ... + + @overload + def __new__( + cls, + *, + input_tokens: torch.Tensor, + target_tokens: torch.Tensor, + top_k: int, + logits: Literal[True], + hidden_states: Literal[True], + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, + ) -> "ForwardInput[torch.Tensor, TopK, torch.Tensor, torch.Tensor]": ... + + @overload + def __new__( + cls, + *, + input_tokens: torch.Tensor, + target_tokens: torch.Tensor | None = None, + top_k: int | None = None, + logits: bool = False, + hidden_states: bool = False, + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, + ) -> "ForwardInput[torch.Tensor | None, TopK | None, torch.Tensor | None, torch.Tensor | None]": ... + + def __new__( + cls, + *, + input_tokens: torch.Tensor, + target_tokens: torch.Tensor | None = None, + top_k: int | None = None, + logits: bool = False, + hidden_states: bool = False, + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, + ) -> Any: + return object.__new__(cls) + + def __post_init__(self) -> None: + if self.top_k is not None and self.top_k < 1: + raise ValueError("top_k must be >= 1") + if self.checkpoint is not Unset and self.lora is not Unset: + raise ValueError("ForwardInput cannot set both checkpoint and lora") + + +type AnyForwardInput = ForwardInput[ + torch.Tensor | None, + TopK | None, + torch.Tensor | None, + torch.Tensor | None, +] +type AnyForwardOutput = ForwardOutput[ + torch.Tensor | None, + TopK | None, + torch.Tensor | None, + torch.Tensor | None, +] +type ForwardInputs = AnyForwardInput | Iterable["ForwardInputs"] +type ForwardOutputs = AnyForwardOutput | Sequence["ForwardOutputs"] +ForwardInputsT = TypeVar("ForwardInputsT", bound=ForwardInputs) +ForwardOutputsT = TypeVar("ForwardOutputsT", bound=ForwardOutputs) +MicroBatchInputsT = TypeVar("MicroBatchInputsT", bound=ForwardInputs, covariant=True) +MicroBatchOutputsT = TypeVar("MicroBatchOutputsT", bound=ForwardOutputs, covariant=True) + + +@dataclass(frozen=True) +class MicroBatch(Generic[MicroBatchInputsT, MicroBatchOutputsT]): + inputs: Sequence[MicroBatchInputsT] + outputs: Sequence[MicroBatchOutputsT] + indices: Sequence[int] + stats: "MicroBatchStats" + + def select(self, xs: Sequence[T]) -> Sequence[T]: + return [xs[i] for i in self.indices] + + +@dataclass(frozen=True) +class MicroBatchStats: + global_start: int + global_stop: int + global_count: int + local_count: int + packed_tokens: int + logical_tokens: int + estimated_required_bytes: int + available_bytes: int + rejected_candidates: int + cold_start: bool + + +class TrainerRankMemoryError(RuntimeError): + pass + + +class TrainerRankSlotStateError(RuntimeError): + pass + + +@dataclass(frozen=True) +class _MemoryCheck: + estimated_required_bytes: int + available_bytes: int + fits: bool + + +@dataclass(frozen=True) +class _MemoryProfile: + bytes_per_token: float + packed_tokens: int + + +@dataclass(frozen=True) +class _CandidateMicroBatch(Generic[ForwardInputsT]): + inputs: Sequence[ForwardInputsT] + indices: tuple[int, ...] + plan: "_FlatForwardPlan" + check: _MemoryCheck + stats_global_count: int + rejected_candidates: int + cold_start: bool + + +class _SlotGraphSentinel(torch.autograd.Function): + @staticmethod + def forward( + ctx: Any, + tensor: torch.Tensor, + marker: torch.Tensor, + ) -> torch.Tensor: + ctx.save_for_backward(marker) + return tensor + + @staticmethod + def backward(ctx: Any, *grad_outputs: Any) -> tuple[torch.Tensor, None]: + return cast(torch.Tensor, grad_outputs[0]), None + + +@dataclass(frozen=True) +class _DynamicOptimizer: + optimizer: torch.optim.Optimizer + master_params: tuple[torch.nn.Parameter, ...] + + +@dataclass(frozen=True) +class _PushedSlot: + trainer: "TrainerRank" + ref: "LoRASlotRef" + + def __enter__(self) -> "_PushedSlot": + return self + + def __exit__(self, *args: object) -> bool: + if not self.trainer._slot_stack or self.trainer._slot_stack[-1] != self.ref: + raise RuntimeError( + "Pushed LoRA/checkpoint stack changed before context exit" + ) + self.trainer.pop_pushed_lora_or_checkpoint() + return False + + +@dataclass(frozen=True) +class _ForwardItem: + request: AnyForwardInput + input_ids: torch.Tensor + labels: torch.Tensor | None + + +@dataclass(frozen=True) +class _PreparedPackedForward: + tokens: torch.Tensor + position_ids: torch.Tensor + attention_state: "PrefixTreeAttentionState | ArtContextParallelState" + packed_seq_params: "PackedSeqParams | None" + positions_by_item: tuple[torch.Tensor, ...] + source_positions_by_item: tuple[torch.Tensor, ...] + + +type _RowMatch = tuple[torch.Tensor, torch.Tensor, tuple[int, ...]] + + +@dataclass(frozen=True) +class _MemorySignature: + topology: tuple[int, int, int, int] + shared_prefix_max_depth: int + slot_group_count: int + request_mix: tuple[str, ...] + + +@dataclass(frozen=True) +class _ForwardGroupPlan: + slot_ref: "LoRASlotRef | None" + request_indices: tuple[int, ...] + items: tuple[_ForwardItem, ...] + packed: PrefixTreePack + + +@dataclass(frozen=True) +class _FlatForwardPlan: + request_count: int + groups: tuple[_ForwardGroupPlan, ...] + packed_tokens: int + logical_tokens: int + output_bytes: int + signature: _MemorySignature + + +class TrainerRank: + def __init__( + self, + runtime: TrainingRuntime, + *, + head_chunk_tokens: int = 512, + shared_prefix_max_depth: int = 1, + memory_safety_factor: float = 1.10, + memory_reserve_fraction: float = 0.03, + ) -> None: + if head_chunk_tokens < 1: + raise ValueError("head_chunk_tokens must be >= 1") + if shared_prefix_max_depth < 0: + raise ValueError("shared_prefix_max_depth must be >= 0") + if memory_safety_factor < 1.0: + raise ValueError("memory_safety_factor must be >= 1.0") + if not (0.0 <= memory_reserve_fraction < 1.0): + raise ValueError("memory_reserve_fraction must be in [0, 1)") + self.runtime: TrainingRuntime = runtime + self.head_chunk_tokens = head_chunk_tokens + self.shared_prefix_max_depth = shared_prefix_max_depth + self.memory_safety_factor = memory_safety_factor + self.memory_reserve_fraction = memory_reserve_fraction + self.device = next(runtime.model[0].parameters()).device + self._param_dtype_size = _dtype_size(next(runtime.model[0].parameters()).dtype) + try: + metadata_model = _language_model(runtime.model[0]) + except RuntimeError: + metadata_model = None + self._hidden_size = _hidden_size(metadata_model, runtime.provider) + self._padded_vocab_size = ( + None if metadata_model is None else _padded_vocab_size(metadata_model) + ) + self._num_layers = int( + getattr(getattr(metadata_model, "config", None), "num_layers", 0) + or getattr(runtime.provider, "num_layers", 1) + or 1 + ) + self._default_slot_ref: LoRASlotRef | None = None + self._slot_stack: list[LoRASlotRef] = [] + self._dynamic_optimizers: dict[str, _DynamicOptimizer] = {} + self._checkpoint_slot_params_by_name: dict[ + str, tuple[torch.nn.Parameter, ...] + ] = {} + self._pending_slot_graphs: dict[ + LoRASlotRef, list[weakref.ReferenceType[torch.Tensor]] + ] = {} + self._memory_profiles: dict[_MemorySignature, _MemoryProfile] = {} + self._last_global_micro_batch_size: int | None = None + self.zero_grad() + + def zero_grad(self) -> None: + for chunk in self.runtime.model: + zero_grad_buffer = getattr(chunk, "zero_grad_buffer", None) + if callable(zero_grad_buffer): + zero_grad_buffer() + optimizer = self.runtime.optimizer + if optimizer is not None: + optimizer.zero_grad() + for params in self._checkpoint_slot_params_by_name.values(): + for param in params: + param.grad = None + self._prune_slot_graphs() + + def set_checkpoint(self, name: str | None) -> None: + self._set_default_slot(self._slot_ref("checkpoint", name)) + + def set_lora(self, name: str | None) -> None: + self._set_default_slot(self._slot_ref("lora", name)) + + def push_checkpoint(self, name: str | None) -> _PushedSlot: + ref = self._slot_ref("checkpoint", name) + self._slot_stack.append(ref) + return _PushedSlot(self, ref) + + def push_lora(self, name: str | None) -> _PushedSlot: + ref = self._slot_ref("lora", name) + self._slot_stack.append(ref) + return _PushedSlot(self, ref) + + def pop_pushed_lora_or_checkpoint(self) -> None: + if not self._slot_stack: + raise RuntimeError("No pushed LoRA or checkpoint to pop") + self._slot_stack.pop() + + def load_checkpoint_slot( + self, + name: str, + adapter_model: dict[str, torch.Tensor], + *, + optimizer_state: Mapping[str, object] | None = None, + alpha: float | None = None, + ) -> int: + loaded = self._load_slot( + "checkpoint", name, adapter_model, trainable=True, alpha=alpha + ) + self._checkpoint_slot_params_by_name[name] = ( + self._validate_dynamic_slot_consistency("checkpoint", name, loaded) + ) + if optimizer_state is None: + self._dynamic_optimizers.pop(name, None) + else: + self._dynamic_optimizers[name] = self._restore_dynamic_optimizer( + name, optimizer_state + ) + return loaded + + def checkpoint_slot_optimizer_state(self, name: str) -> dict[str, object] | None: + if name not in self._checkpoint_slot_params_by_name: + raise ValueError(f"Unknown checkpoint slot: {name!r}") + dynamic = self._dynamic_optimizers.get(name) + if dynamic is None: + return None + return { + "format_version": 1, + "layout": self._dynamic_optimizer_layout(name), + "master_params": tuple( + param.detach().cpu().clone() for param in dynamic.master_params + ), + "optimizer": _state_to_cpu(dynamic.optimizer.state_dict()), + } + + def load_lora_slot( + self, + name: str, + adapter_model: dict[str, torch.Tensor], + *, + alpha: float | None = None, + ) -> int: + loaded = self._load_slot( + "lora", name, adapter_model, trainable=False, alpha=alpha + ) + self._validate_dynamic_slot_consistency("lora", name, loaded) + return loaded + + @overload + def forward_micro_batches( + self, + inputs: Iterable[ForwardInput[LogprobsT, TopKT, LogitsT, HiddenStatesT]], + ) -> Iterator[ + MicroBatch[ + ForwardInput[LogprobsT, TopKT, LogitsT, HiddenStatesT], + ForwardOutput[LogprobsT, TopKT, LogitsT, HiddenStatesT], + ] + ]: ... + + @overload + def forward_micro_batches( + self, + inputs: Iterable[ + Iterable[ForwardInput[LogprobsT, TopKT, LogitsT, HiddenStatesT]] + ], + ) -> Iterator[ + MicroBatch[ + Sequence[ForwardInput[LogprobsT, TopKT, LogitsT, HiddenStatesT]], + Sequence[ForwardOutput[LogprobsT, TopKT, LogitsT, HiddenStatesT]], + ] + ]: ... + + @overload + def forward_micro_batches( + self, + inputs: Iterable[ + Iterable[Iterable[ForwardInput[LogprobsT, TopKT, LogitsT, HiddenStatesT]]] + ], + ) -> Iterator[ + MicroBatch[ + Sequence[Sequence[ForwardInput[LogprobsT, TopKT, LogitsT, HiddenStatesT]]], + Sequence[Sequence[ForwardOutput[LogprobsT, TopKT, LogitsT, HiddenStatesT]]], + ] + ]: ... + + @overload + def forward_micro_batches( + self, + inputs: Iterable[ + Iterable[ + Iterable[ + Iterable[ForwardInput[LogprobsT, TopKT, LogitsT, HiddenStatesT]] + ] + ] + ], + ) -> Iterator[ + MicroBatch[ + Sequence[ + Sequence[ + Sequence[ForwardInput[LogprobsT, TopKT, LogitsT, HiddenStatesT]] + ] + ], + Sequence[ + Sequence[ + Sequence[ForwardOutput[LogprobsT, TopKT, LogitsT, HiddenStatesT]] + ] + ], + ] + ]: ... + + def forward_micro_batches( + self, + inputs: Iterable[ForwardInputs], + ) -> Iterator[MicroBatch[ForwardInputs, ForwardOutputs]]: + items = [_materialize(item) for item in inputs] + requests = list(_flatten(items)) + for _, indices in self._group_active_request_indices(requests): + for index in indices: + self._forward_item(requests[index]) + self._validate_replicated_top_level_count(len(items)) + start = 0 + while start < len(items): + candidate = self._select_next_micro_batch(items, start) + flat_outputs = iter( + self._run_flat_plan_with_memory_tracking( + candidate.plan, + context="forward_micro_batches", + ) + ) + outputs = [_unflatten(item, flat_outputs) for item in candidate.inputs] + stop = start + candidate.stats_global_count + if stop < len(items): + self._last_global_micro_batch_size = max( + self._last_global_micro_batch_size or 0, + candidate.stats_global_count, + ) + yield MicroBatch( + inputs=candidate.inputs, + outputs=outputs, + indices=candidate.indices, + stats=MicroBatchStats( + global_start=start, + global_stop=stop, + global_count=candidate.stats_global_count, + local_count=len(candidate.inputs), + packed_tokens=candidate.plan.packed_tokens, + logical_tokens=candidate.plan.logical_tokens, + estimated_required_bytes=candidate.check.estimated_required_bytes, + available_bytes=candidate.check.available_bytes, + rejected_candidates=candidate.rejected_candidates, + cold_start=candidate.cold_start, + ), + ) + start = stop + + @overload + def dp_rank_forward( + self, + inputs: Iterable[ForwardInput[LogprobsT, TopKT, LogitsT, HiddenStatesT]], + ) -> Sequence[ForwardOutput[LogprobsT, TopKT, LogitsT, HiddenStatesT]]: ... + + @overload + def dp_rank_forward( + self, + inputs: Iterable[ + Iterable[ForwardInput[LogprobsT, TopKT, LogitsT, HiddenStatesT]] + ], + ) -> Sequence[ + Sequence[ForwardOutput[LogprobsT, TopKT, LogitsT, HiddenStatesT]] + ]: ... + + @overload + def dp_rank_forward( + self, + inputs: Iterable[ + Iterable[Iterable[ForwardInput[LogprobsT, TopKT, LogitsT, HiddenStatesT]]] + ], + ) -> Sequence[ + Sequence[Sequence[ForwardOutput[LogprobsT, TopKT, LogitsT, HiddenStatesT]]] + ]: ... + + @overload + def dp_rank_forward( + self, + inputs: Iterable[ + Iterable[ + Iterable[ + Iterable[ForwardInput[LogprobsT, TopKT, LogitsT, HiddenStatesT]] + ] + ] + ], + ) -> Sequence[ + Sequence[ + Sequence[Sequence[ForwardOutput[LogprobsT, TopKT, LogitsT, HiddenStatesT]]] + ] + ]: ... + + def dp_rank_forward(self, inputs: ForwardInputs) -> ForwardOutputs: + materialized = _materialize(inputs) + plan = self._plan_flat_forward(list(_flatten(materialized))) + check = self._memory_check(plan) + if not check.fits: + self._raise_memory_error( + plan, + check, + context="dp_rank_forward", + message="forward is predicted to exceed available memory", + ) + outputs = iter( + self._run_flat_plan_with_memory_tracking( + plan, + context="dp_rank_forward", + ) + ) + return _unflatten(materialized, outputs) + + def dp_reduce( + self, + tensor: torch.Tensor, + *, + op: dist.ReduceOp.RedOpType = dist.ReduceOp.SUM, + ) -> None: + from megatron.core import parallel_state as ps + + dist.all_reduce( + tensor, + op=op, + group=ps.get_data_parallel_group(with_context_parallel=True), + ) + + def optim_step( + self, + *, + params: AdamParams, + scale_grads: float = 1.0, + checkpoints: Sequence[str] | None = None, + ) -> dict[str, float]: + selected_checkpoints = self._selected_dynamic_checkpoints(checkpoints) + return self._dynamic_optim_step( + selected_checkpoints, + params=params, + scale_grads=scale_grads, + ) + + def _load_slot( + self, + kind: Literal["checkpoint", "lora"], + name: str, + adapter_model: dict[str, torch.Tensor], + *, + trainable: bool, + alpha: float | None, + ) -> int: + if self._slot_stack: + raise RuntimeError("Cannot load a LoRA/checkpoint while a slot is pushed") + adapter_model = self._prepare_adapter_model(kind, name, adapter_model) + from art.megatron.lora import LORA_ALPHA, load_lora_slot_into_model + + ref = self._slot_ref(kind, name) + self._guard_slot_can_load(ref) + return load_lora_slot_into_model( + self.runtime.model, + ref, + adapter_model, + alpha=LORA_ALPHA if alpha is None else alpha, + requires_grad=trainable, + ) + + def _prepare_adapter_model( + self, + kind: Literal["checkpoint", "lora"], + name: str, + adapter_model: Mapping[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + templates = self._local_lora_adapter_templates() + keys = set(adapter_model) + expected = set(templates) + if dist.is_available() and dist.is_initialized(): + gathered: list[set[str] | None] = [None] * dist.get_world_size() + dist.all_gather_object(gathered, expected) + expected = set().union(*(value for value in gathered if value is not None)) + if unknown := sorted(keys - expected): + preview = ", ".join(repr(key) for key in unknown[:8]) + more = "" if len(unknown) <= 8 else f", ... +{len(unknown) - 8} more" + raise ValueError( + f"Adapter for {kind} slot {name!r} contains keys that do not match " + f"installed LoRA wrapper sites: {preview}{more}. Configure the " + "Megatron runtime with matching LoRA target modules before loading." + ) + return { + key: ( + tensor.to( + device=templates[key].device, + dtype=templates[key].dtype, + non_blocking=True, + ) + if key in templates + else tensor + ) + for key, tensor in adapter_model.items() + } + + def _local_lora_adapter_templates(self) -> dict[str, torch.Tensor]: + templates: dict[str, torch.Tensor] = {} + for chunk in self.runtime.model: + for module in chunk.modules(): + expected_weight_keys = getattr(module, "_expected_weight_keys", None) + if not callable(expected_weight_keys): + continue + for suffix, parameter_name in ( + ("lora_A", "A_T"), + ("lora_B", "B_T"), + ): + parameter = getattr(module, parameter_name, None) + if not isinstance(parameter, torch.Tensor): + continue + templates.update( + (str(key), parameter) for key in expected_weight_keys(suffix) + ) + return templates + + def _set_default_slot(self, ref: "LoRASlotRef") -> None: + if self._slot_stack: + raise RuntimeError("Cannot set a LoRA/checkpoint while a slot is pushed") + self._default_slot_ref = ref + + @staticmethod + def _slot_ref( + kind: Literal["checkpoint", "lora"], name: str | None + ) -> "LoRASlotRef": + try: + from art.megatron.lora import LoRASlotRef + except ModuleNotFoundError as exc: + if exc.name is None or not exc.name.startswith("megatron"): + raise + + return cast("LoRASlotRef", _LocalLoRASlotRef(kind=kind, name=name)) + + return LoRASlotRef(kind=kind, name=name) + + def _validate_dynamic_slot_consistency( + self, + kind: Literal["checkpoint", "lora"], + name: str, + loaded_sites: int, + ) -> tuple[torch.nn.Parameter, ...]: + from art.megatron.lora import iter_lora_slot_parameters + + ref = self._slot_ref(kind, name) + params = tuple(iter_lora_slot_parameters(self.runtime.model, ref)) + if not (dist.is_available() and dist.is_initialized()): + return params + + signature = tuple( + ( + tuple(param.shape), + str(param.dtype), + bool(getattr(param, "allreduce", True)), + str(getattr(param, "grad_sync_domain", "tp_default")), + str(getattr(param, "grad_sync_op", "none")), + ) + for param in params + ) + local = (int(loaded_sites), signature) + gathered: list[tuple[int, object] | None] = [None] * dist.get_world_size() + dist.all_gather_object(gathered, local) + ranks = [state for state in gathered if state is not None] + if all(state == ranks[0] for state in ranks[1:]): + return params + raise RuntimeError( + f"Dynamic LoRA slot {kind}:{name} is not loaded consistently across " + "distributed ranks. This usually means a sharded/exported LoRA state " + "dict was passed directly to TrainerRank; gather or materialize the " + "full adapter state before loading a dynamic slot. " + f"Loaded-site counts by rank: {[state[0] for state in ranks]}." + ) + + def _resolve_slot_ref(self, request: AnyForwardInput) -> "LoRASlotRef | None": + if request.checkpoint is not Unset: + return self._slot_ref("checkpoint", cast(str | None, request.checkpoint)) + if request.lora is not Unset: + return self._slot_ref("lora", cast(str | None, request.lora)) + if self._slot_stack: + return self._slot_stack[-1] + if self._default_slot_ref is not None: + return self._default_slot_ref + return self._slot_ref("checkpoint", None) + + def _selected_dynamic_checkpoints( + self, + checkpoints: Sequence[str] | None, + ) -> tuple[str, ...]: + loaded = set(self._checkpoint_slot_params_by_name) + if not loaded: + raise TrainerRankSlotStateError( + "TrainerRank.optim_step requires a loaded checkpoint slot. Call " + "load_checkpoint_slot(...) and run backward on outputs produced by " + "that slot before stepping." + ) + requested = ( + tuple(sorted(loaded)) + if checkpoints is None + else tuple(dict.fromkeys(checkpoints)) + ) + if not requested: + raise TrainerRankSlotStateError( + "TrainerRank.optim_step(checkpoints=...) received no checkpoint " + "names. Pass at least one loaded checkpoint slot." + ) + if unknown := set(requested) - loaded: + raise ValueError(f"Unknown checkpoint slots: {sorted(unknown)}") + flags = self._checkpoint_grad_flags(requested) + selected = tuple( + name for name, has_grad in zip(requested, flags, strict=True) if has_grad + ) + if checkpoints is None: + if selected: + return selected + raise TrainerRankSlotStateError( + "TrainerRank.optim_step found loaded checkpoint slots, but none " + "have gradients on any rank. Call loss.backward() first." + ) + if missing := [ + name + for name, has_grad in zip(requested, flags, strict=True) + if not has_grad + ]: + raise TrainerRankSlotStateError( + "TrainerRank.optim_step was asked to step checkpoint slots with no " + f"gradients on any rank: {missing}. Call loss.backward() for those " + "slots first, or omit them from checkpoints=[...]." + ) + return selected + + def _checkpoint_grad_flags(self, names: Sequence[str]) -> tuple[bool, ...]: + flags = torch.tensor( + [ + any( + param.grad is not None + for param in self._checkpoint_slot_params_by_name[name] + ) + for name in names + ], + device=self.device, + dtype=torch.int32, + ) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(flags, op=dist.ReduceOp.MAX) + return tuple(bool(flag) for flag in flags.tolist()) + + def _dynamic_optim_step( + self, + checkpoint_names: Sequence[str], + *, + params: AdamParams, + scale_grads: float, + ) -> dict[str, float]: + selected = [] + for name in checkpoint_names: + self._guard_checkpoint_can_step(name) + slot_params = self._checkpoint_slot_params_by_name[name] + slot_grads = self._reduce_dynamic_grads( + slot_params, scale_grads=scale_grads + ) + selected.append((name, slot_params, slot_grads)) + + all_params = tuple( + param for _, slot_params, _ in selected for param in slot_params + ) + all_grads = tuple(grad for _, _, slot_grads in selected for grad in slot_grads) + grad_norm = _distributed_grad_norm(all_params, all_grads) + if not torch.isfinite(torch.tensor(grad_norm)): + self.zero_grad() + return { + "learning_rate": float(params.learning_rate), + "grad_norm": float(grad_norm), + "update_successful": 0.0, + "num_zeros_in_grad": 0.0, + } + clip = ( + min(1.0, params.grad_clip_norm / (grad_norm + 1.0e-6)) + if params.grad_clip_norm > 0.0 + else 1.0 + ) + for name, model_params, grads in selected: + dynamic = self._dynamic_optimizer(name, params) + for master, grad in zip(dynamic.master_params, grads, strict=True): + master.grad = grad.mul(clip) + dynamic.optimizer.step() + dynamic.optimizer.zero_grad(set_to_none=True) + with torch.no_grad(): + for model, master in zip( + model_params, dynamic.master_params, strict=True + ): + model.copy_(master) + model.grad = None + self._prune_slot_graphs(self._slot_ref("checkpoint", name)) + return { + "learning_rate": float(params.learning_rate), + "grad_norm": float(grad_norm), + "update_successful": 1.0, + "num_zeros_in_grad": 0.0, + } + + def _dynamic_optimizer( + self, + name: str, + params: AdamParams, + ) -> _DynamicOptimizer: + dynamic = self._dynamic_optimizers.get(name) + if dynamic is None: + dynamic = self._new_dynamic_optimizer(name, params) + self._dynamic_optimizers[name] = dynamic + return dynamic + for group in dynamic.optimizer.param_groups: + group["lr"] = params.learning_rate + group["betas"] = (params.beta1, params.beta2) + group["weight_decay"] = params.weight_decay + return dynamic + + def _new_dynamic_optimizer( + self, + name: str, + params: AdamParams, + *, + master_params: Sequence[torch.Tensor] | None = None, + ) -> _DynamicOptimizer: + model_params = self._checkpoint_slot_params_by_name[name] + sources = model_params if master_params is None else tuple(master_params) + if len(sources) != len(model_params) or any( + not isinstance(source, torch.Tensor) for source in sources + ): + raise TrainerRankSlotStateError( + f"Optimizer state for checkpoint slot {name!r} has " + f"{len(sources)} master parameters; expected {len(model_params)}." + ) + masters = tuple( + torch.nn.Parameter( + source.detach().to(device=model.device, dtype=torch.float32).clone() + ) + for model, source in zip( + model_params, + sources, + strict=True, + ) + ) + optimizer = torch.optim.AdamW( + masters, + lr=params.learning_rate, + betas=(params.beta1, params.beta2), + weight_decay=params.weight_decay, + ) + return _DynamicOptimizer(optimizer, masters) + + def _restore_dynamic_optimizer( + self, + name: str, + state: Mapping[str, object], + ) -> _DynamicOptimizer: + if state.get("format_version") != 1: + raise TrainerRankSlotStateError( + f"Unsupported optimizer state format for checkpoint slot {name!r}." + ) + if state.get("layout") != self._dynamic_optimizer_layout(name): + raise TrainerRankSlotStateError( + f"Optimizer state for checkpoint slot {name!r} was saved for a " + "different topology or parameter layout. Save and restore one " + "optimizer shard per TrainerRank with matching TP/EP/ETP ranks." + ) + master_params = state.get("master_params") + optimizer_state = state.get("optimizer") + if not isinstance(master_params, Sequence) or not isinstance( + optimizer_state, Mapping + ): + raise TrainerRankSlotStateError( + f"Optimizer state for checkpoint slot {name!r} is incomplete." + ) + dynamic = self._new_dynamic_optimizer( + name, + AdamParams(learning_rate=0.0), + master_params=cast(Sequence[torch.Tensor], master_params), + ) + try: + dynamic.optimizer.load_state_dict( + {str(key): value for key, value in optimizer_state.items()} + ) + except ValueError as exc: + raise TrainerRankSlotStateError( + f"Optimizer state for checkpoint slot {name!r} does not match the " + "loaded slot parameter groups." + ) from exc + for param in dynamic.master_params: + for state_name, value in dynamic.optimizer.state.get(param, {}).items(): + if ( + isinstance(value, torch.Tensor) + and int(value.ndim) > 0 + and tuple(value.shape) != tuple(param.shape) + ): + raise TrainerRankSlotStateError( + f"Optimizer state {state_name!r} for checkpoint slot " + f"{name!r} has shape {tuple(value.shape)}, but the loaded " + f"slot parameter has shape {tuple(param.shape)}." + ) + return dynamic + + def _reduce_dynamic_grads( + self, + params: Sequence[torch.nn.Parameter], + *, + scale_grads: float, + ) -> tuple[torch.Tensor, ...]: + from megatron.core import parallel_state as ps + + from art.megatron.training.finalize_grads import ( + coalesced_all_reduce, + tensor_parallel_grad_sync, + ) + + buckets: dict[ + tuple[int, str, torch.dtype, torch.device], + tuple[object, dist.ReduceOp.RedOpType, list[torch.Tensor]], + ] = {} + + def add(group: object, op: dist.ReduceOp.RedOpType, grad: torch.Tensor) -> None: + key = (id(group), str(op), grad.dtype, grad.device) + buckets.setdefault(key, (group, op, []))[2].append(grad) + + grads = tuple( + ( + torch.zeros_like(param, dtype=torch.float32) + if param.grad is None + else param.grad.detach().float().mul(scale_grads) + ) + for param in params + ) + for param, grad in zip(params, grads, strict=True): + if bool(getattr(param, "allreduce", True)): + group = ps.get_data_parallel_group(with_context_parallel=True) + else: + group = ps.get_expert_data_parallel_group() + if group is not None and group.size() > 1: + add(group, dist.ReduceOp.SUM, grad) + + sync = tensor_parallel_grad_sync(param, name="dynamic LoRA") + if sync is not None: + group, reduce_op = sync + add(group, reduce_op, grad) + + for group, op, bucket_grads in buckets.values(): + coalesced_all_reduce(bucket_grads, group=group, op=op) + return grads + + def _dynamic_optimizer_layout(self, name: str) -> dict[str, object]: + return { + "parallel": _parallel_optimizer_coordinates(), + "parameters": tuple( + ( + tuple(param.shape), + str(param.dtype), + str(getattr(param, "lora_shard_domain", "tp")), + bool(getattr(param, "lora_tp_sharded", False)), + getattr(param, "lora_tp_shard_dim", None), + str(getattr(param, "lora_tp_shard_strategy", "uniform")), + tuple(getattr(param, "lora_tp_component_sizes", ())), + ) + for param in self._checkpoint_slot_params_by_name[name] + ), + } + + def _select_next_micro_batch( + self, + items: Sequence[ForwardInputsT], + start: int, + ) -> _CandidateMicroBatch[ForwardInputsT]: + dp_rank, dp_size = self._dp_rank_and_size() + remaining = len(items) - start + min_width = min(dp_size, remaining) + if min_width <= 0: + raise RuntimeError("cannot select an empty microbatch window") + base_granularity = 1 if remaining < 64 else 8 if remaining < 256 else 32 + granularity = max( + 1, + ((base_granularity + dp_size - 1) // dp_size) * dp_size, + ) + + def normalize(width: int) -> int: + width = max(min_width, min(width, remaining)) + if width in (min_width, remaining) or granularity <= 1: + return width + if width < granularity: + return width + return max(min_width, (width // granularity) * granularity) + + def local_slice(width: int) -> tuple[tuple[int, ...], list[ForwardInputsT]]: + stop = start + width + indices = tuple(range(start + dp_rank, stop, dp_size)) + return indices, [items[index] for index in indices] + + estimates: dict[int, tuple[_MemoryCheck, bool, bool] | None] = {} + plans: dict[int, _FlatForwardPlan] = {} + + def estimate(width: int) -> tuple[_MemoryCheck, bool, bool] | None: + width = normalize(width) + if width in estimates: + return estimates[width] + indices, local_inputs = local_slice(width) + values = self._estimate_flat_forward(list(_flatten(local_inputs))) + if not self._all_ranks_true(values is not None): + estimates[width] = None + return None + assert values is not None + packed_tokens, output_bytes, signature = values + result = ( + self._memory_check_required( + self._estimate_required_memory_bytes_from_values( + packed_tokens=packed_tokens, + output_bytes=output_bytes, + signature=signature, + ), + sync_across_dp=True, + ), + self._all_ranks_have_memory_profile( + packed_tokens=packed_tokens, + signature=signature, + ), + self._all_ranks_true(signature in self._memory_profiles), + ) + estimates[width] = result + return result + + rejected_widths: set[int] = set() + + def fits(width: int) -> tuple[bool, bool]: + width = normalize(width) + result = estimate(width) + if result is None: + plan = materialize(width) + check = self._memory_check(plan, sync_across_dp=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) + else: + check, trusted, profiled = result + if not check.fits: + rejected_widths.add(width) + return check.fits and (trusted or not profiled), trusted + + def materialize(width: int) -> _FlatForwardPlan: + width = normalize(width) + plan = plans.get(width) + if plan is None: + _, local_inputs = local_slice(width) + plan = self._plan_flat_forward(list(_flatten(local_inputs))) + plans[width] = plan + return plan + + def candidate(width: int) -> _CandidateMicroBatch[ForwardInputsT]: + width = normalize(width) + indices, local_inputs = local_slice(width) + plan = materialize(width) + estimated = estimates.get(width) + check = ( + estimated[0] + if estimated is not None + else self._memory_check(plan, sync_across_dp=True) + ) + cold_start = not self._all_ranks_have_memory_profile( + packed_tokens=plan.packed_tokens, + signature=plan.signature, + ) + return _CandidateMicroBatch( + inputs=local_inputs, + indices=indices, + plan=plan, + check=check, + stats_global_count=width, + rejected_candidates=len(rejected_widths), + cold_start=cold_start, + ) + + 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 not first.check.fits: + self._raise_memory_error( + first.plan, + first.check, + context="forward_micro_batches", + message="smallest DP microbatch is predicted to exceed available memory", + ) + if first.cold_start: + return first + + best = min_width + failed: int | None = None + width = normalize(self._last_global_micro_batch_size or min_width) + if width > best: + fit, trusted = fits(width) + if fit: + best = width + if not trusted: + return candidate(best) + else: + failed = width + + while failed is None and best < remaining: + width = normalize(max(best + 1, best * 2)) + if width == best: + break + fit, trusted = fits(width) + if fit: + best = width + if not trusted: + break + else: + failed = width + + if failed is not None: + while failed - best > 1: + width = normalize((best + failed) // 2) + if width in (best, failed): + break + if fits(width)[0]: + best = width + else: + failed = width + + return candidate(best) + + def _validate_replicated_top_level_count(self, count: int) -> None: + if not (dist.is_available() and dist.is_initialized()): + return + counts = [0 for _ in range(dist.get_world_size())] + dist.all_gather_object(counts, int(count)) + if len(set(counts)) == 1: + return + raise ValueError( + "forward_micro_batches requires the same top-level input count on every " + "distributed rank. Pass already-DP-local inputs to dp_rank_forward instead. " + f"Observed counts by rank: {counts}." + ) + + def _dp_rank_and_size(self) -> tuple[int, int]: + try: + from megatron.core import parallel_state as ps + + return int(ps.get_data_parallel_rank()), int( + ps.get_data_parallel_world_size() + ) + except (AssertionError, ImportError, RuntimeError, ValueError): + return 0, 1 + + def _plan_flat_forward( + self, requests: Sequence[AnyForwardInput] + ) -> _FlatForwardPlan: + plans: list[_ForwardGroupPlan] = [] + output_bytes = self._estimate_group_request_output_bytes(requests) + logical_tokens = sum(int(request.input_tokens.numel()) for request in requests) + groups = self._group_active_request_indices(requests) + for slot_ref, group_indices in groups: + items = tuple( + self._forward_item(requests[index]) for index in group_indices + ) + packed = prefix_tree_pack( + (item.input_ids for item in items), + max_depth=self.shared_prefix_max_depth, + ) + plans.append( + _ForwardGroupPlan( + slot_ref=slot_ref, + request_indices=tuple(group_indices), + items=items, + packed=packed, + ) + ) + + return _FlatForwardPlan( + request_count=len(requests), + groups=tuple(plans), + packed_tokens=sum(int(plan.packed.tokens.numel()) for plan in plans), + logical_tokens=logical_tokens, + output_bytes=output_bytes, + signature=self._memory_signature_from_requests( + requests, + slot_group_count=len(plans), + ), + ) + + def _estimate_flat_forward( + self, requests: Sequence[AnyForwardInput] + ) -> tuple[int, int, _MemorySignature] | None: + groups = self._group_active_request_indices(requests) + packed_tokens = 0 + for _, group_indices in groups: + group_packed_tokens = estimate_prefix_tree_packed_tokens( + (requests[index].input_tokens.reshape(-1) for index in group_indices), + max_depth=self.shared_prefix_max_depth, + ) + if group_packed_tokens is None: + return None + packed_tokens += group_packed_tokens + + return ( + packed_tokens, + self._estimate_group_request_output_bytes(requests), + self._memory_signature_from_requests( + requests, + slot_group_count=len(groups), + ), + ) + + def _group_active_request_indices( + self, + requests: Sequence[AnyForwardInput], + ) -> tuple[tuple["LoRASlotRef | None", tuple[int, ...]], ...]: + groups: dict[LoRASlotRef | None, list[int]] = {} + for index, request in enumerate(requests): + if ( + request.target_tokens is not None + or request.logits + or request.top_k is not None + or request.hidden_states + ): + groups.setdefault(self._resolve_slot_ref(request), []).append(index) + return tuple((slot_ref, tuple(indices)) for slot_ref, indices in groups.items()) + + def _run_flat_plan_with_memory_tracking( + self, + plan: _FlatForwardPlan, + *, + context: str, + ) -> list[AnyForwardOutput]: + if torch.cuda.is_available() and self.device.type == "cuda": + torch.cuda.synchronize(self.device) + baseline = int(torch.cuda.memory_allocated(self.device)) + torch.cuda.reset_peak_memory_stats(self.device) + else: + baseline = 0 + try: + outputs = self._execute_flat_plan(plan) + except torch.cuda.OutOfMemoryError as exc: + check = self._memory_check(plan) + self._raise_memory_error( + plan, + check, + context=context, + message="CUDA OOM occurred despite the planner estimate", + ) + raise AssertionError("unreachable") from exc + if torch.cuda.is_available() and self.device.type == "cuda": + torch.cuda.synchronize(self.device) + peak = int(torch.cuda.max_memory_allocated(self.device)) + self._update_memory_profile(plan, max(0, peak - baseline)) + return outputs + + def _execute_flat_plan(self, plan: _FlatForwardPlan) -> list[AnyForwardOutput]: + outputs = [ + ForwardOutput(None, None, None, None) for _ in range(plan.request_count) + ] + 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 + return outputs + + def _track_slot_graph_outputs( + self, + ref: "LoRASlotRef | None", + outputs: Sequence[AnyForwardOutput], + ) -> list[AnyForwardOutput]: + if ref is None or ref.name is None: + return list(outputs) + + marker: torch.Tensor | None = None + + def track(tensor: torch.Tensor | None) -> torch.Tensor | None: + nonlocal marker + if tensor is None or not tensor.requires_grad: + return tensor + if marker is None: + marker = tensor.new_empty(0) + return cast(torch.Tensor, _SlotGraphSentinel.apply(tensor, marker)) + + tracked_outputs = [ + ForwardOutput( + target_logprobs=track(output.target_logprobs), + top_k=( + None + if output.top_k is None + else TopK( + logprobs=cast(torch.Tensor, track(output.top_k.logprobs)), + tokens=output.top_k.tokens, + ) + ), + logits=track(output.logits), + hidden_states=track(output.hidden_states), + ) + for output in outputs + ] + if marker is not None: + self._slot_graphs().setdefault(ref, []).append(weakref.ref(marker)) + return tracked_outputs + + def _slot_graphs( + self, + ) -> dict["LoRASlotRef", list[weakref.ReferenceType[torch.Tensor]]]: + graphs = getattr(self, "_pending_slot_graphs", None) + if graphs is None: + graphs = {} + self._pending_slot_graphs = graphs + return graphs + + def _prune_slot_graphs(self, ref: "LoRASlotRef | None" = None) -> None: + graphs = self._slot_graphs() + refs = tuple(graphs) if ref is None else (ref,) + for current in refs: + live = [ + marker for marker in graphs.get(current, ()) if marker() is not None + ] + if live: + graphs[current] = live + else: + graphs.pop(current, None) + + def _has_live_slot_graph(self, ref: "LoRASlotRef") -> bool: + self._prune_slot_graphs(ref) + return bool(self._slot_graphs().get(ref)) + + def _guard_slot_can_load(self, ref: "LoRASlotRef") -> None: + if not self._has_live_slot_graph(ref): + return + raise TrainerRankSlotStateError( + f"Cannot load {ref.kind} slot {ref.name!r} while outputs from an " + "earlier forward using that slot still have a live backward graph. " + "Activation checkpoint recompute resolves slots by name, so replacing " + "the slot before backward can compute gradients with different LoRA " + "weights than the original forward. Finish backward first; if the " + "forward was abandoned, release all references to its outputs; or load " + "the new weights under a different slot name." + ) + + def _guard_checkpoint_can_step(self, name: str) -> None: + ref = self._slot_ref("checkpoint", name) + if not self._has_live_slot_graph(ref): + return + raise TrainerRankSlotStateError( + f"Cannot optim_step checkpoint slot {name!r} while outputs from an " + "earlier forward using that slot have not been backpropagated. Call " + "loss.backward() without retaining the graph before optim_step(); if " + "the forward was abandoned, release all references to its outputs." + ) + + def _estimate_group_request_output_bytes( + self, + requests: Sequence[AnyForwardInput], + ) -> int: + total = 0 + for request in requests: + seq_len = int(request.input_tokens.numel()) + if request.target_tokens is not None: + total += int(request.target_tokens.numel()) * _dtype_size(torch.float32) + if request.top_k is not None: + total += ( + seq_len + * int(request.top_k) + * (_dtype_size(torch.float32) + _dtype_size(torch.long)) + ) + if request.logits: + if self._padded_vocab_size is None: + raise RuntimeError("logits output memory requires a GPT model") + total += seq_len * self._padded_vocab_size * self._param_dtype_size + if request.hidden_states: + total += seq_len * self._hidden_size * self._param_dtype_size + return total + + def _memory_signature_from_requests( + self, + requests: Sequence[AnyForwardInput], + *, + slot_group_count: int, + ) -> _MemorySignature: + return _MemorySignature( + topology=self._topology_key(), + shared_prefix_max_depth=self.shared_prefix_max_depth, + slot_group_count=slot_group_count, + request_mix=tuple( + sorted({_request_mix_key(request) for request in requests}) + ), + ) + + def _topology_key(self) -> tuple[int, int, int, int]: + try: + topology = self._topology() + return cast( + tuple[int, int, int, int], + tuple( + int(getattr(topology, name)) for name in ("dp", "tp", "cp", "pp") + ), + ) + except (AssertionError, AttributeError, ImportError, RuntimeError, ValueError): + return (1, 1, 1, 1) + + def _memory_check( + self, + forward: _FlatForwardPlan, + *, + sync_across_dp: bool = False, + ) -> _MemoryCheck: + return self._memory_check_required( + self._estimate_required_memory_bytes_from_values( + packed_tokens=forward.packed_tokens, + output_bytes=forward.output_bytes, + signature=forward.signature, + ), + sync_across_dp=sync_across_dp, + ) + + def _memory_check_required( + self, + required: int, + *, + sync_across_dp: bool = False, + ) -> _MemoryCheck: + available = self._available_memory_bytes() + if dist.is_available() and dist.is_initialized(): + group = None if sync_across_dp else self._forward_memory_group() + values = torch.tensor( + [float(required), float(available)], + device=self.device if self.device.type == "cuda" else "cpu", + dtype=torch.float64, + ) + dist.all_reduce(values[0], op=dist.ReduceOp.MAX, group=group) + dist.all_reduce(values[1], op=dist.ReduceOp.MIN, group=group) + required = int(values[0].item()) + available = int(values[1].item()) + return _MemoryCheck( + estimated_required_bytes=required, + available_bytes=available, + fits=required <= available, + ) + + @staticmethod + def _forward_memory_group() -> object | None: + try: + from megatron.core import parallel_state as ps + + return ps.get_tensor_and_context_parallel_group(check_initialized=False) + except (AssertionError, ImportError, RuntimeError, ValueError): + return None + + def _raise_memory_error( + self, + plan: _FlatForwardPlan, + check: _MemoryCheck, + *, + context: str, + message: str, + ) -> None: + raise TrainerRankMemoryError( + f"{context}: {message}. " + f"packed_tokens={plan.packed_tokens} " + f"logical_tokens={plan.logical_tokens} " + f"output_gb={plan.output_bytes / 1024**3:.3f} " + f"estimated_required_gb={check.estimated_required_bytes / 1024**3:.3f} " + f"available_gb={check.available_bytes / 1024**3:.3f}. " + "Use smaller top-level items, reduce output requests, or call " + "dp_rank_forward with already-DP-local smaller inputs." + ) + + def _estimate_required_memory_bytes_from_values( + self, + *, + packed_tokens: int, + output_bytes: int, + signature: _MemorySignature, + ) -> int: + if packed_tokens <= 0: + return output_bytes + profiled = self._memory_profiles.get(signature) + activation_factor = max(4, min(16, self._num_layers // 4 + 4)) + static_compute = ( + packed_tokens + * self._hidden_size + * self._param_dtype_size + * activation_factor + ) + if ( + profiled is None + or profiled.packed_tokens * _MEMORY_PROFILE_TRUST_GROWTH < packed_tokens + ): + compute = static_compute + else: + compute = max(static_compute, int(profiled.bytes_per_token * packed_tokens)) + return int((output_bytes + compute) * self.memory_safety_factor) + + def _available_memory_bytes(self) -> int: + if not (torch.cuda.is_available() and self.device.type == "cuda"): + return 1 << 60 + free, total = torch.cuda.mem_get_info(self.device) + allocated = int(torch.cuda.memory_allocated(self.device)) + reserved = int(torch.cuda.memory_reserved(self.device)) + reusable_reserved = max(0, reserved - allocated) + reserve = int(total * self.memory_reserve_fraction) + return max(0, int(free) + reusable_reserved - reserve) + + def _all_ranks_have_memory_profile( + self, + *, + packed_tokens: int, + signature: _MemorySignature, + ) -> bool: + profile = self._memory_profiles.get(signature) + local = packed_tokens <= 0 or ( + profile is not None + and profile.packed_tokens * _MEMORY_PROFILE_TRUST_GROWTH >= packed_tokens + ) + return self._all_ranks_true(local) + + def _all_ranks_true(self, local: bool) -> bool: + if not (dist.is_available() and dist.is_initialized()): + return local + value = torch.tensor( + int(local), + device=self.device if self.device.type == "cuda" else "cpu", + dtype=torch.int32, + ) + dist.all_reduce(value, op=dist.ReduceOp.MIN) + return bool(value.item()) + + def _update_memory_profile( + self, plan: _FlatForwardPlan, peak_delta_bytes: int + ) -> None: + if plan.packed_tokens <= 0: + return + compute_delta = max(0, peak_delta_bytes - plan.output_bytes) + bytes_per_token = compute_delta / max(1, plan.packed_tokens) + previous = self._memory_profiles.get(plan.signature) + self._memory_profiles[plan.signature] = _MemoryProfile( + bytes_per_token=max( + bytes_per_token, + 0.0 if previous is None else previous.bytes_per_token, + ), + packed_tokens=max( + plan.packed_tokens, + 0 if previous is None else previous.packed_tokens, + ), + ) + + def _forward_item(self, request: AnyForwardInput) -> _ForwardItem: + if request.top_k is not None: + _validate_top_k(request.top_k, _language_model(self.runtime.model[0])) + input_ids = request.input_tokens.reshape(-1).to(dtype=torch.long) + if int(input_ids.numel()) == 0: + raise ValueError("input_tokens must not be empty") + labels = None + if request.target_tokens is not None: + labels = request.target_tokens.to(dtype=torch.long) + if int(labels.numel()) == 0: + raise ValueError("target_tokens must not be empty") + input_shape = tuple(request.input_tokens.shape) + if tuple(labels.shape) == input_shape: + labels = labels.reshape(-1) + elif ( + labels.ndim > request.input_tokens.ndim + and tuple(labels.shape[: request.input_tokens.ndim]) == input_shape + ): + labels = labels.reshape( + int(input_ids.numel()), *labels.shape[request.input_tokens.ndim :] + ) + elif labels.ndim < 1 or int(labels.shape[0]) != int(input_ids.numel()): + raise ValueError( + "target_tokens must match input_tokens or add trailing target " + f"dimensions: input_tokens={input_shape} " + f"target_tokens={tuple(labels.shape)}" + ) + return _ForwardItem(request=request, input_ids=input_ids, labels=labels) + + def _forward_packed( + self, + items: Sequence[_ForwardItem], + prepared: _PreparedPackedForward, + ) -> list[AnyForwardOutput]: + hidden_by_row = self._gather_sequence_parallel_hidden( + self._decoder_hidden(prepared) + ) + return self._project_head(items, prepared, hidden_by_row) + + def _decoder_hidden( + self, + prepared: _PreparedPackedForward, + ) -> torch.Tensor: + from art.megatron.train import _placeholder_attention_mask + + handler = self.runtime.model_support_handler + model = _language_model(self.runtime.model[0]) + attention_mask = _placeholder_attention_mask(self.device) + forward_kwargs = handler.get_forward_kwargs( + self.runtime.model[0], + attention_bias=prepared.attention_state, + ) + extra_block_kwargs = cast( + dict[str, object] | None, + forward_kwargs.pop("extra_block_kwargs", None), + ) + preprocessed = model._preprocess( + input_ids=prepared.tokens, + position_ids=prepared.position_ids, + packed_seq_params=cast("PackedSeqParams", prepared.packed_seq_params), + ) + ( + decoder_input, + rotary_pos_emb, + rotary_pos_cos, + rotary_pos_sin, + sequence_len_offset, + padding_mask, + ) = preprocessed[:6] + rotary_pos_cos_sin = preprocessed[6] if len(preprocessed) == 7 else None + return cast( + torch.Tensor, + model.decoder( + hidden_states=decoder_input, + attention_mask=attention_mask, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + rotary_pos_cos_sin=rotary_pos_cos_sin, + packed_seq_params=prepared.packed_seq_params, + sequence_len_offset=sequence_len_offset, + padding_mask=padding_mask, + **(extra_block_kwargs or {}), + ), + ) + + def _project_head( + self, + items: Sequence[_ForwardItem], + prepared: _PreparedPackedForward, + hidden_by_row: torch.Tensor, + ) -> list[AnyForwardOutput]: + model = _language_model(self.runtime.model[0]) + output_weight = ( + model.shared_embedding_or_output_weight() + if bool(model.share_embeddings_and_output_weights) + else None + ) + device = hidden_by_row.device + target_logprobs = [None for _ in items] + logits: list[torch.Tensor | None] = [None for _ in items] + top_k: list[TopK | None] = [None for _ in items] + label_rows: list[torch.Tensor | None] = [None for _ in items] + projected_rows: list[torch.Tensor] = [] + + for index, (item, positions_cpu) in enumerate( + zip(items, prepared.positions_by_item, strict=True) + ): + positions = positions_cpu.to(device=device) + if item.request.logits or item.request.top_k is not None: + projected_rows.append(positions) + if item.labels is not None: + source_positions = prepared.source_positions_by_item[index].to(device) + labels = item.labels.to(device=device).index_select(0, source_positions) + label_rows[index] = labels + target_logprobs[index] = torch.zeros( + tuple(labels.shape), + device=device, + dtype=torch.float32, + ) + if item.request.top_k is None and not item.request.logits: + if int(labels.shape[0]): + valid = labels != -100 + if labels.ndim > 1: + valid = valid.reshape(int(labels.shape[0]), -1).any(dim=1) + valid_offsets = torch.nonzero(valid, as_tuple=False).reshape(-1) + if int(valid_offsets.numel()): + projected_rows.append( + positions.index_select(0, valid_offsets) + ) + if item.request.logits: + logits[index] = torch.empty( + (int(positions.numel()), _padded_vocab_size(model)), + device=hidden_by_row.device, + dtype=hidden_by_row.dtype, + ) + if item.request.top_k is not None: + shape = (int(positions.numel()), item.request.top_k) + top_k[index] = TopK( + logprobs=torch.empty(shape, device=device, dtype=torch.float32), + tokens=torch.empty(shape, device=device, dtype=torch.long), + ) + + row_tensor = ( + torch.cat(projected_rows).unique(sorted=True) + if projected_rows + else torch.empty(0, dtype=torch.long, device=device) + ) + if int(row_tensor.numel()): + rows_cpu = row_tensor.detach().cpu() + cpu_matches = tuple( + _row_match( + positions.cpu(), + rows_cpu, + chunk_tokens=self.head_chunk_tokens, + ) + for positions in prepared.positions_by_item + ) + local_row_matches = tuple( + (source.to(device), row.to(device), bounds) + for source, row, bounds in cpu_matches + ) + logit_rows_cpu = torch.cat( + tuple( + match[1] + for item, match in zip(items, cpu_matches, strict=True) + if item.request.logits + ) + or (torch.empty(0, dtype=torch.long),) + ).unique(sorted=True) + self._project_vocab_parallel( + items, + hidden_by_row, + row_tensor, + row_matches=local_row_matches, + logit_rows=logit_rows_cpu.to(device), + logit_bounds=_chunk_boundaries( + logit_rows_cpu, + end=int(row_tensor.numel()), + chunk_tokens=self.head_chunk_tokens, + ), + output_weight=output_weight, + target_logprobs=target_logprobs, + top_k=top_k, + logits=logits, + label_rows=label_rows, + ) + + target_logprobs, top_k = _anchor_disconnected_outputs( + target_logprobs, + top_k, + hidden_by_row, + ) + return [ + ForwardOutput( + target_logprobs=target_logprobs[index], + top_k=top_k[index], + logits=logits[index], + hidden_states=( + _select_positions(hidden_by_row, positions) + if item.request.hidden_states + else None + ), + ) + for index, (item, positions) in enumerate( + zip(items, prepared.positions_by_item, strict=True) + ) + ] + + def _project_vocab_parallel( + self, + items: Sequence[_ForwardItem], + hidden_by_row: torch.Tensor, + rows: torch.Tensor, + *, + row_matches: Sequence[_RowMatch], + logit_rows: torch.Tensor, + logit_bounds: tuple[int, ...], + output_weight: torch.Tensor | None, + target_logprobs: list[torch.Tensor | None], + top_k: list[TopK | None], + logits: list[torch.Tensor | None], + label_rows: list[torch.Tensor | None], + ) -> None: + model = _language_model(self.runtime.model[0]) + max_top_k = max((int(item.request.top_k or 0) for item in items), default=0) + need_log_z = any( + item.labels is not None or item.request.top_k is not None for item in items + ) + for chunk_index, start in enumerate( + range(0, int(rows.numel()), self.head_chunk_tokens) + ): + chunk_rows = rows[start : start + self.head_chunk_tokens] + local_logits = self._local_logits_from_hidden_rows( + model, + _select_positions(hidden_by_row, chunk_rows), + output_weight=output_weight, + ) + log_z: torch.Tensor | None = None + local_topk: tuple[torch.Tensor, torch.Tensor] | None = None + if need_log_z: + topk_stats = _try_triton_local_topk_stats(local_logits, k=max_top_k) + logsumexp_stats = ( + cast( + tuple[torch.Tensor, torch.Tensor] | None, + _try_triton_stats("local_logsumexp_stats", local_logits), + ) + if topk_stats is None + else None + ) + stats = topk_stats if topk_stats is not None else logsumexp_stats + if stats is not None: + local_max, local_sum = stats[:2] + local_max = local_max.detach() + global_max = _all_reduce_tensor_parallel_max(local_max) + global_sum = _all_reduce_tensor_parallel_sum( + local_sum * torch.exp(local_max - global_max) + ) + log_z = global_max + torch.log(global_sum) + else: + log_z = _vocab_parallel_log_z(local_logits) + + if topk_stats is not None: + _, _, local_values, local_tokens = topk_stats + local_topk = (local_values, local_tokens) + elif logsumexp_stats is not None and max_top_k > 0: + local_k = min(max_top_k, int(local_logits.shape[1])) + local_values, local_tokens = torch.topk( + local_logits, k=local_k, dim=-1 + ) + local_topk = (local_values.float(), local_tokens) + + logit_start, logit_end = logit_bounds[chunk_index : chunk_index + 2] + logit_chunk_offsets = logit_rows[logit_start:logit_end] - start + chunk_logits: torch.Tensor | None = None + if int(logit_chunk_offsets.numel()): + chunk_logits = _batch_seq_logits( + self._gather_tensor_parallel_logits( + local_logits.index_select(0, logit_chunk_offsets).unsqueeze(1) + ), + seq_len=int(logit_chunk_offsets.numel()), + ).squeeze(0) + + for index, item in enumerate(items): + offsets, row_offsets, bounds = row_matches[index] + begin, finish = bounds[chunk_index : chunk_index + 2] + offsets = offsets[begin:finish] + chunk_offsets = row_offsets[begin:finish] - start + if int(offsets.numel()) == 0: + continue + item_logits = logits[index] + if item_logits is not None: + if chunk_logits is None: + raise RuntimeError("logits output requires gathered logits") + item_logits[offsets] = chunk_logits.index_select( + 0, + torch.searchsorted(logit_chunk_offsets, chunk_offsets), + ) + labels = label_rows[index] + item_logprobs = target_logprobs[index] + if item_logprobs is not None and labels is not None: + if log_z is None: + raise RuntimeError("target logprobs require logsumexp") + selected_log_z = log_z.index_select(0, chunk_offsets) + item_logprobs[offsets] = _vocab_parallel_target_logprobs( + local_logits, + labels.index_select(0, offsets), + selected_log_z, + row_offsets=chunk_offsets, + ) + k = item.request.top_k + if k is not None: + if log_z is None: + raise RuntimeError("top_k requires logsumexp") + selected_log_z = log_z.index_select(0, chunk_offsets) + if local_topk is not None: + local_values, local_tokens = local_topk + selected_values = local_values.index_select(0, chunk_offsets) + selected_tokens = local_tokens.index_select(0, chunk_offsets) + else: + selected_logits = local_logits.index_select(0, chunk_offsets) + selected_values, selected_tokens = torch.topk( + selected_logits.float(), + k=min(k, int(selected_logits.shape[1])), + dim=-1, + ) + values = _vocab_parallel_topk_from_local( + selected_values, + selected_tokens, + k=k, + log_z=selected_log_z, + vocab_start=_vocab_range(local_logits)[0], + ) + current = top_k[index] + if current is None: + raise RuntimeError("top_k output was not allocated") + current.logprobs[offsets] = values.logprobs + current.tokens[offsets] = values.tokens + + def _local_logits_from_hidden_rows( + self, + model: "GPTModel", + hidden: torch.Tensor, + *, + output_weight: torch.Tensor | None, + ) -> torch.Tensor: + output_layer = model.output_layer + sequence_parallel = bool(getattr(output_layer, "sequence_parallel", False)) + if sequence_parallel: + output_layer.sequence_parallel = False + try: + logits, _ = output_layer( + hidden.unsqueeze(1), + weight=output_weight, + runtime_gather_output=None, + ) + finally: + if sequence_parallel: + output_layer.sequence_parallel = True + return _batch_seq_logits( + model._scale_logits(logits), + seq_len=int(hidden.shape[0]), + ).squeeze(0) + + def _gather_sequence_parallel_hidden(self, hidden: torch.Tensor) -> torch.Tensor: + from megatron.core import parallel_state as ps + + if int(ps.get_tensor_model_parallel_world_size()) <= 1: + return hidden.squeeze(1) + from megatron.core import tensor_parallel + + gathered = tensor_parallel.gather_from_sequence_parallel_region( + hidden, + tensor_parallel_output_grad=False, + group=ps.get_tensor_model_parallel_group(check_initialized=False), + ) + return cast(torch.Tensor, gathered).squeeze(1) + + def _prepare_packed_forward( + self, + batch: PrefixTreePack, + ) -> _PreparedPackedForward: + topology = self._topology() + batch = _pad_packed_batch(batch, multiple=int(topology.tp)) + 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 + + handler = self.runtime.model_support_handler + provider = self.runtime.provider + return _PreparedPackedForward( + tokens=batch.tokens.to(self.device), + position_ids=batch.position_ids.to(self.device), + attention_state=create_prefix_tree_state( + group_ids=batch.group_ids, + parent_ids=batch.parent_ids, + target_device=self.device, + build_gdn_execution_spec=handler.build_gdn_execution_spec, + attention_head_dim=provider.kv_channels, + attention_value_head_dim=provider.kv_channels, + ), + packed_seq_params=None, + positions_by_item=batch.positions_by_sequence, + source_positions_by_item=tuple( + torch.arange( + int(positions.numel()), + dtype=torch.long, + device=positions.device, + ) + for positions in batch.positions_by_sequence + ), + ) + + def _prepare_context_parallel_forward( + self, + batch: PrefixTreePack, + *, + topology: "ParallelTopology", + ) -> _PreparedPackedForward: + from megatron.core import parallel_state as ps + + from art.megatron.context_parallel.runtime import ( + _dispatch_tensor, + prepare_cp_micro, + ) + from art.megatron.training.microbatches import ( + _art_flex_cp_block_mask_variants, + _context_parallel_config_for_provider, + _gdn_planner_config_for_provider, + ) + from art.preprocessing.pack import PackedTensors + + assistant_mask = torch.ones_like(batch.tokens, dtype=torch.bool) + sparse_micro: PackedTensors = { + "tokens": batch.tokens, + "group_ids": batch.group_ids, + "parent_ids": batch.parent_ids, + "input_pos": batch.position_ids, + "assistant_mask": assistant_mask, + "logprobs": torch.full_like( + batch.tokens, float("nan"), dtype=torch.float32 + ), + "advantages": torch.zeros_like(batch.tokens, dtype=torch.float32), + "weights": assistant_mask.to(dtype=torch.float32), + "pixel_values": [None], + "image_grid_thw": [None], + "moe_routing_replay": None, + } + handler = self.runtime.model_support_handler + provider = self.runtime.provider + prepared = prepare_cp_micro( + micro=sparse_micro, + topology=topology, + config=_context_parallel_config_for_provider(provider, self.device), + cp_group=ps.get_context_parallel_group(check_initialized=False), + cp_rank=ps.get_context_parallel_rank(), + build_gdn_execution_spec=handler.build_gdn_execution_spec, + gdn_planner_config=_gdn_planner_config_for_provider(provider, handler), + block_mask_variants=_art_flex_cp_block_mask_variants(provider, self.device), + target_device=self.device, + ) + if prepared.rank_plan is None: + raise RuntimeError("CP forward preparation did not return a rank plan") + local_positions = _dispatch_tensor( + torch.arange( + int(batch.tokens.shape[1]), + dtype=torch.long, + ).unsqueeze(0), + rank_plan=prepared.rank_plan, + pad_value=-1, + pad_multiple=prepared.pad_multiple, + ) + local_position_pairs = tuple( + _local_position_pairs(local_positions, positions) + for positions in batch.positions_by_sequence + ) + return _PreparedPackedForward( + tokens=prepared.tensors.tokens, + position_ids=prepared.tensors.input_pos, + attention_state=cast("ArtContextParallelState", prepared.attention_state), + packed_seq_params=prepared.packed_seq_params, + positions_by_item=tuple(pair[0] for pair in local_position_pairs), + source_positions_by_item=tuple(pair[1] for pair in local_position_pairs), + ) + + def _topology(self) -> "ParallelTopology": + from art.megatron.train import _infer_parallel_topology + + return _infer_parallel_topology(self.runtime.model) + + def _gather_tensor_parallel_logits(self, logits: torch.Tensor) -> torch.Tensor: + from megatron.core import parallel_state as ps + + if int(ps.get_tensor_model_parallel_world_size()) <= 1: + return logits + from megatron.core import tensor_parallel + + return cast( + torch.Tensor, + tensor_parallel.gather_from_tensor_model_parallel_region(logits), + ) + + +def _validate_top_k(top_k: int, model: "GPTModel") -> None: + vocab_size = _padded_vocab_size(model) + if top_k > vocab_size: + raise ValueError(f"top_k={top_k} exceeds vocabulary size {vocab_size}") + + +def _request_mix_key(request: AnyForwardInput) -> str: + parts = [] + if request.target_tokens is not None: + target = request.target_tokens + tail_shape = tuple(target.shape[request.input_tokens.ndim :]) + parts.append(f"target:{tail_shape or 'single'}") + if request.top_k is not None: + parts.append(f"topk:{int(request.top_k)}") + if request.logits: + parts.append("logits") + if request.hidden_states: + parts.append("hidden") + return "+".join(parts) if parts else "inactive" + + +def _pad_packed_batch( + batch: PrefixTreePack, + *, + multiple: int, +) -> PrefixTreePack: + if multiple <= 1: + return batch + seq_len = int(batch.tokens.shape[1]) + pad = -seq_len % multiple + if pad == 0: + return batch + + device = batch.tokens.device + next_group = ( + int(batch.group_ids.max().item()) + 1 if int(batch.group_ids.numel()) else 1 + ) + pad_group_ids = torch.arange( + next_group, + next_group + pad, + dtype=batch.group_ids.dtype, + device=device, + ).unsqueeze(0) + return PrefixTreePack( + tokens=torch.cat((batch.tokens, batch.tokens.new_zeros((1, pad))), dim=1), + group_ids=torch.cat((batch.group_ids, pad_group_ids), dim=1), + parent_ids=torch.cat((batch.parent_ids, pad_group_ids), dim=1), + position_ids=torch.cat( + (batch.position_ids, batch.position_ids.new_zeros((1, pad))), dim=1 + ), + positions_by_sequence=batch.positions_by_sequence, + ) + + +def _language_model(model: torch.nn.Module) -> "GPTModel": + module: object = model + while hasattr(module, "module"): + module = getattr(module, "module") + if hasattr(module, "_preprocess") and hasattr(module, "decoder"): + return cast("GPTModel", module) + language_model = getattr(module, "language_model", None) + if language_model is not None: + return cast("GPTModel", language_model) + raise RuntimeError("expected a Megatron GPT model") + + +def _padded_vocab_size(model: "GPTModel") -> int: + vocab_size = getattr(getattr(model, "config", None), "padded_vocab_size", None) + if vocab_size is None: + vocab_size = getattr(model, "vocab_size", None) + if vocab_size is None: + raise RuntimeError("could not determine full padded vocabulary size") + return int(vocab_size) + + +def _hidden_size(model: "GPTModel | None", provider: object) -> int: + for source in (getattr(model, "config", None), model, provider): + if source is None: + continue + hidden_size = getattr(source, "hidden_size", None) + if hidden_size is not None: + return int(hidden_size) + raise RuntimeError("could not determine hidden size") + + +def _dtype_size(dtype: torch.dtype) -> int: + return torch.empty((), dtype=dtype).element_size() + + +def _distributed_grad_norm( + params: Sequence[torch.nn.Parameter], + grads: Sequence[torch.Tensor], +) -> float: + if len(params) != len(grads): + raise ValueError("params and grads must have matching lengths") + included = [ + grad + for param, grad in zip(params, grads, strict=True) + if _include_in_distributed_grad_norm(param) + ] + device = grads[0].device if grads else torch.device("cpu") + squared = torch.zeros((), device=device, dtype=torch.float32) + for grad in included: + squared.add_(grad.float().square().sum()) + if dist.is_available() and dist.is_initialized(): + dist.all_reduce(squared, op=dist.ReduceOp.SUM) + return float(torch.sqrt(squared).item()) + + +def _include_in_distributed_grad_norm(param: torch.nn.Parameter) -> bool: + if not (dist.is_available() and dist.is_initialized()): + return True + from megatron.core import parallel_state as ps + + replica_group = ( + ps.get_data_parallel_group(with_context_parallel=True) + if bool(getattr(param, "allreduce", True)) + else ps.get_expert_data_parallel_group() + ) + if replica_group is not None and replica_group.size() > 1: + if replica_group.rank() != 0: + return False + if bool(getattr(param, "lora_tp_sharded", False)): + return True + shard_group = ( + ps.get_tensor_model_parallel_group(check_initialized=False) + if getattr(param, "lora_shard_domain", "tp") == "tp" + else ps.get_expert_tensor_parallel_group(check_initialized=False) + ) + return shard_group is None or shard_group.size() <= 1 or shard_group.rank() == 0 + + +def _parallel_optimizer_coordinates() -> tuple[int, ...]: + if not (dist.is_available() and dist.is_initialized()): + return (1, 0, 1, 0, 1, 0, 1, 0) + from megatron.core import parallel_state as ps + + expert_tp_group = ps.get_expert_tensor_parallel_group(check_initialized=False) + return ( + int(ps.get_tensor_model_parallel_world_size()), + int(ps.get_tensor_model_parallel_rank()), + int(ps.get_expert_model_parallel_world_size()), + int(ps.get_expert_model_parallel_rank()), + 1 if expert_tp_group is None else int(expert_tp_group.size()), + 0 if expert_tp_group is None else int(expert_tp_group.rank()), + int(ps.get_pipeline_model_parallel_world_size()), + int(ps.get_pipeline_model_parallel_rank()), + ) + + +def _state_to_cpu(value: object) -> object: + if isinstance(value, torch.Tensor): + return value.detach().cpu().clone() + if isinstance(value, Mapping): + return {key: _state_to_cpu(item) for key, item in value.items()} + if isinstance(value, tuple): + return tuple(_state_to_cpu(item) for item in value) + if isinstance(value, list): + return [_state_to_cpu(item) for item in value] + return value + + +def _vocab_parallel_target_logprobs( + local_logits: torch.Tensor, + labels: torch.Tensor, + log_z: torch.Tensor, + *, + row_offsets: torch.Tensor, +) -> torch.Tensor: + start, _ = _vocab_range(local_logits) + flat_labels = labels.reshape(int(labels.shape[0]), -1) + local_labels = flat_labels - start + owns_label = ( + (flat_labels != -100) + & (local_labels >= 0) + & (local_labels < int(local_logits.shape[1])) + ) + rows = row_offsets.reshape(-1, 1).expand_as(flat_labels) + target_logits = local_logits[ + rows, + local_labels.clamp(0, int(local_logits.shape[1]) - 1), + ].float() + target_logits = target_logits.masked_fill(~owns_label, 0.0).reshape(labels.shape) + target_logits = _all_reduce_tensor_parallel_sum(target_logits) + log_z = log_z.reshape(int(log_z.shape[0]), *((1,) * (int(labels.ndim) - 1))) + return (target_logits.float() - log_z).masked_fill(labels == -100, 0.0) + + +def _anchor_disconnected_outputs( + target_logprobs: list[torch.Tensor | None], + top_k: list[TopK | None], + hidden_by_row: torch.Tensor, +) -> tuple[list[torch.Tensor | None], list[TopK | None]]: + if not hidden_by_row.requires_grad: + return target_logprobs, top_k + anchor: torch.Tensor | None = None + + def anchor_tensor(tensor: torch.Tensor) -> torch.Tensor: + nonlocal anchor + if tensor.requires_grad: + return tensor + if anchor is None: + anchor = hidden_by_row.reshape(-1)[:1].float().sum() * 0.0 + return tensor + anchor + + for index, logprobs in enumerate(target_logprobs): + if logprobs is not None: + target_logprobs[index] = anchor_tensor(logprobs) + for index, item in enumerate(top_k): + if item is not None: + top_k[index] = TopK(anchor_tensor(item.logprobs), item.tokens) + return target_logprobs, top_k + + +def _try_triton_local_topk_stats( + local_logits: torch.Tensor, + *, + k: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None: + if k <= 0 or k > int( + os.environ.get("ART_TRAINER_RANK_TRITON_FUSED_TOPK_MAX", "10") + ): + return None + return cast( + tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None, + _try_triton_stats( + "local_topk_stats", + local_logits, + k=min(k, int(local_logits.shape[1])), + ), + ) + + +def _try_triton_stats( + name: str, + local_logits: torch.Tensor, + **kwargs: object, +) -> object | None: + if not local_logits.is_cuda: + return None + if os.environ.get("ART_TRAINER_RANK_TRITON_TOPK", "1").lower() in { + "0", + "false", + } or int(local_logits.shape[0]) < int( + os.environ.get("ART_TRAINER_RANK_TRITON_MIN_ROWS", "64") + ): + return None + try: + from art.trainer_rank import topk + + return getattr(topk, name)(local_logits, **kwargs) + except Exception: + if os.environ.get("ART_TRAINER_RANK_TRITON_TOPK", "1").lower() == "strict": + raise + return None + + +def _vocab_parallel_topk_from_local( + local_values: torch.Tensor, + local_tokens: torch.Tensor, + *, + k: int, + log_z: torch.Tensor, + vocab_start: int, +) -> TopK: + local_k = min(k, int(local_values.shape[1])) + local_values = local_values[:, :local_k] + local_tokens = local_tokens[:, :local_k] + vocab_start + + from megatron.core import parallel_state as ps + + tp_size = int(ps.get_tensor_model_parallel_world_size()) + if tp_size <= 1: + return TopK( + logprobs=local_values - log_z.unsqueeze(1), + tokens=local_tokens, + ) + + from megatron.core import tensor_parallel + + group = ps.get_tensor_model_parallel_group(check_initialized=False) + values = cast( + torch.Tensor, + tensor_parallel.gather_from_tensor_model_parallel_region( + local_values, + group=group, + ), + ) + gathered_tokens = [torch.empty_like(local_tokens) for _ in range(tp_size)] + dist.all_gather(gathered_tokens, local_tokens, group=group) + tokens = torch.cat(gathered_tokens, dim=1) + top_values, top_offsets = torch.topk(values, k=k, dim=-1) + return TopK( + logprobs=top_values - log_z.unsqueeze(1), + tokens=tokens.gather(1, top_offsets), + ) + + +def _vocab_parallel_log_z(local_logits: torch.Tensor) -> torch.Tensor: + local_logits = local_logits.float() + local_max = local_logits.max(dim=-1).values.detach() + global_max = _all_reduce_tensor_parallel_max(local_max) + local_sum = _local_vocab_exp_sum(local_logits, global_max) + global_sum = _all_reduce_tensor_parallel_sum(local_sum) + return global_max + torch.log(global_sum) + + +def _local_vocab_exp_sum( + local_logits: torch.Tensor, + global_max: torch.Tensor, +) -> torch.Tensor: + return torch.exp(local_logits.float() - global_max.unsqueeze(1)).sum(dim=-1) + + +def _vocab_range(local_logits: torch.Tensor) -> tuple[int, int]: + from megatron.core import parallel_state as ps + + local_size = int(local_logits.shape[1]) + rank = int(ps.get_tensor_model_parallel_rank()) + start = rank * local_size + return start, start + local_size + + +def _all_reduce_tensor_parallel_sum(tensor: torch.Tensor) -> torch.Tensor: + from megatron.core import parallel_state as ps + + if int(ps.get_tensor_model_parallel_world_size()) <= 1: + return tensor + from megatron.core import tensor_parallel + + return cast( + torch.Tensor, + tensor_parallel.reduce_from_tensor_model_parallel_region( + tensor, + group=ps.get_tensor_model_parallel_group(check_initialized=False), + ), + ) + + +def _all_reduce_tensor_parallel_max(tensor: torch.Tensor) -> torch.Tensor: + from megatron.core import parallel_state as ps + + if int(ps.get_tensor_model_parallel_world_size()) <= 1: + return tensor + output = tensor.clone() + dist.all_reduce( + output, + op=dist.ReduceOp.MAX, + group=ps.get_tensor_model_parallel_group(check_initialized=False), + ) + return output + + +def _row_match( + positions: torch.Tensor, + rows: torch.Tensor, + *, + chunk_tokens: int, +) -> _RowMatch: + row_offsets = torch.searchsorted(rows, positions) + in_bounds = row_offsets < int(rows.numel()) + source_offsets = torch.arange( + int(positions.numel()), device=positions.device, dtype=torch.long + )[in_bounds] + row_offsets = row_offsets[in_bounds] + keep = rows.index_select(0, row_offsets) == positions.index_select( + 0, source_offsets + ) + source_offsets, row_offsets = source_offsets[keep], row_offsets[keep] + if int(row_offsets.numel()) > 1: + order = row_offsets.argsort() + source_offsets = source_offsets.index_select(0, order) + row_offsets = row_offsets.index_select(0, order) + return ( + source_offsets, + row_offsets, + _chunk_boundaries( + row_offsets, + end=int(rows.numel()), + chunk_tokens=chunk_tokens, + ), + ) + + +def _chunk_boundaries( + offsets: torch.Tensor, + *, + end: int, + chunk_tokens: int, +) -> tuple[int, ...]: + edges = torch.arange(0, end, chunk_tokens, dtype=torch.long) + edges = torch.cat((edges, torch.tensor((end,), dtype=torch.long))) + return tuple(torch.searchsorted(offsets, edges).tolist()) + + +def _select_positions(values: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + if int(positions.numel()) == 0: + return values[:0] + return values.index_select(0, positions.to(device=values.device)) + + +def _batch_seq_logits(logits: torch.Tensor, *, seq_len: int) -> torch.Tensor: + if int(logits.ndim) != 3: + raise RuntimeError( + f"expected logits with shape [B, S, V] or [S, B, V], got {tuple(logits.shape)}" + ) + if int(logits.shape[0]) == 1 and int(logits.shape[1]) == seq_len: + return logits + if int(logits.shape[0]) == seq_len and int(logits.shape[1]) == 1: + return logits.transpose(0, 1).contiguous() + raise RuntimeError( + f"logits do not match sequence length {seq_len}: {tuple(logits.shape)}" + ) + + +def _materialize(inputs: ForwardInputs) -> ForwardInputs: + if isinstance(inputs, ForwardInput): + return inputs + return [_materialize(item) for item in _nested_forward_children(inputs)] + + +def _flatten(inputs: ForwardInputs) -> Iterator[AnyForwardInput]: + if isinstance(inputs, ForwardInput): + yield inputs + return + for item in _nested_forward_children(inputs): + yield from _flatten(item) + + +def _unflatten( + template: ForwardInputs, outputs: Iterator[AnyForwardOutput] +) -> ForwardOutputs: + if isinstance(template, ForwardInput): + return next(outputs) + return [_unflatten(item, outputs) for item in _nested_forward_children(template)] + + +def _nested_forward_children(inputs: ForwardInputs) -> Iterator[ForwardInputs]: + if isinstance(inputs, Mapping): + raise TypeError( + "dict was passed directly to TrainerRank; gather or materialize the " + "values into a list/tuple so nested forward output ordering is explicit" + ) + if isinstance(inputs, str | bytes): + raise TypeError( + "TrainerRank forward inputs must be ForwardInput objects or nested " + "iterables of ForwardInput objects, not strings" + ) + try: + return iter(cast(Iterable[ForwardInputs], inputs)) + except TypeError as exc: + raise TypeError( + "TrainerRank forward inputs must be ForwardInput objects or nested " + "iterables of ForwardInput objects" + ) from exc + + +__all__ = [ + "AdamParams", + "ForwardInput", + "ForwardOutput", + "MicroBatch", + "MicroBatchStats", + "TopK", + "TrainerRank", + "TrainerRankMemoryError", + "TrainerRankSlotStateError", +] diff --git a/src/art/trainer_rank/topk.py b/src/art/trainer_rank/topk.py new file mode 100644 index 000000000..e0a84722f --- /dev/null +++ b/src/art/trainer_rank/topk.py @@ -0,0 +1,283 @@ +from __future__ import annotations + +from typing import Any + +import torch +import triton +import triton.language as tl + +type LocalTopKStats = tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] +type LocalLogSumExpStats = tuple[torch.Tensor, torch.Tensor] + + +@triton.jit +def _stats_stage1_kernel( + logits_ptr, + partial_max_ptr, + partial_sum_ptr, + partial_values_ptr, + partial_tokens_ptr, + stride_row: tl.constexpr, + vocab_size: tl.constexpr, + n_blocks: tl.constexpr, + k: tl.constexpr, + block_v: tl.constexpr, +): + row = tl.program_id(0) + block = tl.program_id(1) + offsets = block * block_v + tl.arange(0, block_v) + mask = offsets < vocab_size + values = tl.load( + logits_ptr + row * stride_row + offsets, + mask=mask, + other=-float("inf"), + ).to(tl.float32) + + block_max = tl.max(values, axis=0) + block_sum = tl.sum(tl.exp(values - block_max), axis=0) + partial_offset = row * n_blocks + block + tl.store(partial_max_ptr + partial_offset, block_max) + tl.store(partial_sum_ptr + partial_offset, block_sum) + + work = values + arange = tl.arange(0, block_v) + for slot in tl.static_range(0, k): + top_value, top_index = tl.max( + work, + axis=0, + return_indices=True, + return_indices_tie_break_left=True, + ) + output_offset = (partial_offset * k) + slot + tl.store(partial_values_ptr + output_offset, top_value) + tl.store( + partial_tokens_ptr + output_offset, + (block * block_v + top_index).to(tl.int64), + ) + work = tl.where(arange == top_index, -float("inf"), work) + + +@triton.jit +def _stats_stage2_kernel( + partial_max_ptr, + partial_sum_ptr, + partial_values_ptr, + partial_tokens_ptr, + local_max_ptr, + local_sum_ptr, + values_ptr, + tokens_ptr, + n_blocks: tl.constexpr, + k: tl.constexpr, + block_b: tl.constexpr, + block_candidates: tl.constexpr, +): + row = tl.program_id(0) + + block_offsets = tl.arange(0, block_b) + block_mask = block_offsets < n_blocks + partial_base = row * n_blocks + block_max = tl.load( + partial_max_ptr + partial_base + block_offsets, + mask=block_mask, + other=-float("inf"), + ) + row_max = tl.max(block_max, axis=0) + block_sum = tl.load( + partial_sum_ptr + partial_base + block_offsets, + mask=block_mask, + other=0.0, + ) + row_sum = tl.sum(block_sum * tl.exp(block_max - row_max), axis=0) + tl.store(local_max_ptr + row, row_max) + tl.store(local_sum_ptr + row, row_sum) + + if k > 0: + candidate_offsets = tl.arange(0, block_candidates) + candidate_mask = candidate_offsets < n_blocks * k + candidate_base = row * n_blocks * k + candidates = tl.load( + partial_values_ptr + candidate_base + candidate_offsets, + mask=candidate_mask, + other=-float("inf"), + ) + work = candidates + for slot in tl.static_range(0, k): + top_value, top_index = tl.max( + work, + axis=0, + return_indices=True, + return_indices_tie_break_left=True, + ) + output_offset = row * k + slot + tl.store(values_ptr + output_offset, top_value) + tl.store( + tokens_ptr + output_offset, + tl.load(partial_tokens_ptr + candidate_base + top_index), + ) + work = tl.where(candidate_offsets == top_index, -float("inf"), work) + + +@triton.jit +def _stats_backward_kernel( + logits_ptr, + local_max_ptr, + tokens_ptr, + grad_sum_ptr, + grad_values_ptr, + grad_logits_ptr, + stride_row: tl.constexpr, + vocab_size: tl.constexpr, + k: tl.constexpr, + block_v: tl.constexpr, +): + row = tl.program_id(0) + block = tl.program_id(1) + offsets = block * block_v + tl.arange(0, block_v) + mask = offsets < vocab_size + + logits = tl.load( + logits_ptr + row * stride_row + offsets, + mask=mask, + other=-float("inf"), + ).to(tl.float32) + local_max = tl.load(local_max_ptr + row) + grad = tl.load(grad_sum_ptr + row).to(tl.float32) * tl.exp(logits - local_max) + + for slot in tl.static_range(0, k): + token = tl.load(tokens_ptr + row * k + slot) + value_grad = tl.load(grad_values_ptr + row * k + slot).to(tl.float32) + grad += tl.where(offsets == token, value_grad, 0.0) + + tl.store(grad_logits_ptr + row * stride_row + offsets, grad, mask=mask) + + +class _LocalStatsFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, local_logits: torch.Tensor, k: int): + local_max, local_sum, values, tokens = _local_stats_forward(local_logits, k=k) + ctx.save_for_backward(local_logits, local_max, tokens) + ctx.k = k + return local_max, local_sum, values, tokens + + @staticmethod + def backward(ctx: Any, *grad_outputs: Any) -> Any: + grad_local_max, grad_local_sum, grad_values, grad_tokens = grad_outputs + del grad_local_max, grad_tokens + logits, local_max, tokens = ctx.saved_tensors + k = int(ctx.k) + rows = int(logits.shape[0]) + vocab_size = int(logits.shape[1]) + block_v = 4096 + n_blocks = int(triton.cdiv(vocab_size, block_v)) + + if grad_local_sum is None: + grad_local_sum = torch.zeros_like(local_max) + if grad_values is None: + grad_values = torch.zeros( + (rows, k), + device=logits.device, + dtype=torch.float32, + ) + + grad_logits = torch.empty_like(logits) + _stats_backward_kernel[(rows, n_blocks)]( + logits, + local_max, + tokens, + grad_local_sum.contiguous(), + grad_values.contiguous(), + grad_logits, + logits.stride(0), + vocab_size=vocab_size, # ty: ignore[invalid-argument-type] + k=k, # ty: ignore[invalid-argument-type] + block_v=block_v, # ty: ignore[invalid-argument-type] + num_warps=8, # ty: ignore[unknown-argument] + ) + return grad_logits, None + + +def _check_local_logits(local_logits: torch.Tensor) -> torch.Tensor: + if local_logits.ndim != 2: + raise ValueError( + f"expected [rows, vocab] logits, got {tuple(local_logits.shape)}" + ) + if not local_logits.is_cuda: + raise ValueError("local top-k helpers require CUDA logits") + return local_logits.contiguous() + + +def _local_stats_forward(local_logits: torch.Tensor, *, k: int) -> LocalTopKStats: + logits = _check_local_logits(local_logits) + if k < 0 or k > int(local_logits.shape[1]): + raise ValueError( + f"k={k} is outside local vocab size {int(local_logits.shape[1])}" + ) + + rows = int(logits.shape[0]) + vocab_size = int(logits.shape[1]) + block_v = 4096 + n_blocks = int(triton.cdiv(vocab_size, block_v)) + block_b = int(triton.next_power_of_2(n_blocks)) + block_candidates = int(triton.next_power_of_2(n_blocks * k)) if k else 1 + + partial_shape = (rows, n_blocks) + partial_max = torch.empty(partial_shape, device=logits.device, dtype=torch.float32) + partial_sum = torch.empty_like(partial_max) + partial_topk_shape = (rows, n_blocks, k) if k else (1,) + partial_values = torch.empty( + partial_topk_shape, device=logits.device, dtype=torch.float32 + ) + partial_tokens = torch.empty( + partial_topk_shape, device=logits.device, dtype=torch.long + ) + local_max = torch.empty((rows,), device=logits.device, dtype=torch.float32) + local_sum = torch.empty_like(local_max) + values = torch.empty((rows, k), device=logits.device, dtype=torch.float32) + tokens = torch.empty((rows, k), device=logits.device, dtype=torch.long) + + _stats_stage1_kernel[(rows, n_blocks)]( + logits, + partial_max, + partial_sum, + partial_values, + partial_tokens, + stride_row=logits.stride(0), # ty: ignore[invalid-argument-type] + vocab_size=vocab_size, # ty: ignore[invalid-argument-type] + n_blocks=n_blocks, # ty: ignore[invalid-argument-type] + k=k, # ty: ignore[invalid-argument-type] + block_v=block_v, # ty: ignore[invalid-argument-type] + num_warps=8, # ty: ignore[unknown-argument] + ) + _stats_stage2_kernel[(rows,)]( + partial_max, + partial_sum, + partial_values, + partial_tokens, + local_max, + local_sum, + values, + tokens, + n_blocks=n_blocks, # ty: ignore[invalid-argument-type] + k=k, # ty: ignore[invalid-argument-type] + block_b=block_b, # ty: ignore[invalid-argument-type] + block_candidates=block_candidates, # ty: ignore[invalid-argument-type] + num_warps=8, # ty: ignore[unknown-argument] + ) + return local_max, local_sum, values, tokens + + +def local_topk_stats(local_logits: torch.Tensor, *, k: int) -> LocalTopKStats: + logits = local_logits.contiguous() + if not logits.requires_grad: + return _local_stats_forward(logits, k=k) + return _LocalStatsFunction.apply(logits, k) + + +def local_logsumexp_stats(local_logits: torch.Tensor) -> LocalLogSumExpStats: + logits = local_logits.contiguous() + if not logits.requires_grad: + local_max, local_sum, _, _ = _local_stats_forward(logits, k=0) + return local_max, local_sum + local_max, local_sum, _, _ = _LocalStatsFunction.apply(logits, 0) + return local_max, local_sum diff --git a/tests/integration/megatron/lora/test_dynamic_lora_slots.py b/tests/integration/megatron/lora/test_dynamic_lora_slots.py new file mode 100644 index 000000000..ca7cf921b --- /dev/null +++ b/tests/integration/megatron/lora/test_dynamic_lora_slots.py @@ -0,0 +1,421 @@ +from __future__ import annotations + +from contextlib import contextmanager +import os +from pathlib import Path +import socket +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("megatron.core") + +from megatron.core import parallel_state as ps # noqa: E402 +from torch.distributed import destroy_process_group, init_process_group # noqa: E402 +import torch.multiprocessing as mp # noqa: E402 + +from art.megatron.lora import LoRA, LoRASlotRef, use_lora_slot # noqa: E402 +from art.trainer_rank import ( # noqa: E402 + AdamParams, + TrainerRank, + _distributed_grad_norm, + _vocab_parallel_log_z, + _vocab_parallel_target_logprobs, + _vocab_parallel_topk_from_local, +) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required.") +def test_dynamic_lora_slots_capture_recompute_context_and_step_independently() -> None: + with _single_rank_model_parallel(): + device = torch.device("cuda") + lora = LoRA( + "dense", + in_features=4, + out_features=5, + rank=2, + alpha=32, + dtype=torch.float32, + device=device, + ) + ref_a = LoRASlotRef("checkpoint", "A") + ref_b = LoRASlotRef("checkpoint", "B") + lora.load_lora_slot( + ref_a, _adapter("dense", rank=1, seed=1), requires_grad=True + ) + lora.load_lora_slot( + ref_b, _adapter("dense", rank=4, seed=2), requires_grad=True + ) + + x = torch.randn(7, 4, device=device) + with use_lora_slot(LoRASlotRef("checkpoint", None)): + assert torch.equal(lora(x), torch.zeros(7, 5, device=device)) + with use_lora_slot(LoRASlotRef("lora", "missing")): + assert torch.equal(lora(x), torch.zeros(7, 5, device=device)) + + slot_a = lora._slot(ref_a) + assert slot_a is not None + with use_lora_slot(ref_a): + actual = lora(x) + expected = (x @ slot_a.A_T) @ slot_a.B_T * slot_a.scale + assert torch.allclose(actual, expected, atol=0, rtol=0) + assert slot_a.rank == 1 + assert slot_a.scale == 32.0 + assert lora._slot(ref_b).scale == 8.0 # type: ignore[union-attr] + + trainer = _trainer_for(lora, device) + cpu_adapter = { + key: value.cpu().double() + for key, value in _adapter("dense", rank=3, seed=7).items() + } + trainer.load_checkpoint_slot("CPU", cpu_adapter) + cpu_slot = lora._slot(LoRASlotRef("checkpoint", "CPU")) + assert cpu_slot is not None + assert cpu_slot.A_T.device == lora.A_T.device + assert cpu_slot.A_T.dtype == lora.A_T.dtype + with use_lora_slot(LoRASlotRef("checkpoint", "CPU")): + assert lora(x).is_cuda + + with trainer.push_checkpoint("A"): + assert trainer._slot_stack[-1] == ref_a + with trainer.push_lora(None): + assert trainer._slot_stack[-1].name is None + assert trainer._slot_stack[-1] == ref_a + assert trainer._slot_stack == [] + + from megatron.core.tensor_parallel.random import ( + checkpoint as megatron_checkpoint, + ) + from torch.utils.checkpoint import checkpoint as torch_checkpoint + + _assert_checkpoint_recomputes_with(ref_a, ref_b, lora, torch_checkpoint) + _assert_checkpoint_recomputes_with( + ref_a, ref_b, lora, megatron_checkpoint, False + ) + _assert_step_updates_only(ref_a, ref_b, lora, trainer) + _assert_reload_replaces_slot_optimizer(ref_a, lora, trainer) + + +@pytest.mark.parametrize("tp_size", (2, 4)) +def test_trainer_rank_tp_head_backward_matches_unsharded_oracle( + tp_size: int, + tmp_path: Path, +) -> None: + if not torch.cuda.is_available() or torch.cuda.device_count() < tp_size: + pytest.skip(f"requires {tp_size} CUDA devices") + init_file = tmp_path / f"tp_head_{tp_size}" + mp.spawn( + _tp_head_backward_worker, + args=(tp_size, f"file://{init_file}"), + nprocs=tp_size, + join=True, + ) + + +def _tp_head_backward_worker(rank: int, world: int, init_method: str) -> None: + torch.cuda.set_device(rank) + init_process_group( + "nccl", + rank=rank, + world_size=world, + init_method=init_method, + ) + try: + ps.initialize_model_parallel( + tensor_model_parallel_size=world, + pipeline_model_parallel_size=1, + context_parallel_size=1, + expert_model_parallel_size=1, + ) + device = torch.device("cuda", rank) + full = torch.tensor( + [ + [-1.2, 0.4, 2.1, -0.7, 1.3, 0.2, -2.0, 0.8], + [0.1, -0.5, 1.7, 0.3, -1.1, 2.4, 0.9, -0.2], + ], + device=device, + ) + local_size = int(full.shape[1]) // world + local = _local_shard(full, rank, local_size) + labels = torch.tensor([2, 5], device=device) + rows = torch.arange(int(full.shape[0]), device=device) + actual = _vocab_parallel_target_logprobs( + local, + labels, + _vocab_parallel_log_z(local), + row_offsets=rows, + ) + (-actual.sum()).backward() + + reference = full.detach().clone().requires_grad_() + (-torch.log_softmax(reference, dim=-1)[rows, labels].sum()).backward() + torch.testing.assert_close( + local.grad, + reference.grad[:, rank * local_size : (rank + 1) * local_size], + atol=1e-6, + rtol=1e-6, + ) + + local = _local_shard(full, rank, local_size) + local_values, local_tokens = torch.topk(local, k=min(2, local_size), dim=-1) + actual_topk = _vocab_parallel_topk_from_local( + local_values, + local_tokens, + k=2, + log_z=_vocab_parallel_log_z(local), + vocab_start=rank * local_size, + ) + (-actual_topk.logprobs.sum()).backward() + + reference = full.detach().clone().requires_grad_() + reference_values, reference_tokens = torch.topk( + torch.log_softmax(reference, dim=-1), k=2, dim=-1 + ) + (-reference_values.sum()).backward() + torch.testing.assert_close(actual_topk.tokens, reference_tokens) + torch.testing.assert_close( + local.grad, + reference.grad[:, rank * local_size : (rank + 1) * local_size], + atol=1e-6, + rtol=1e-6, + ) + + from megatron.core import tensor_parallel + + local_hidden = torch.randn(2, 1, 3, device=device, requires_grad=True) + gathered_hidden = tensor_parallel.gather_from_sequence_parallel_region( + local_hidden, + tensor_parallel_output_grad=False, + group=ps.get_tensor_model_parallel_group(check_initialized=False), + ).squeeze(1) + gathered_hidden.sum().backward() + torch.testing.assert_close(local_hidden.grad, torch.ones_like(local_hidden)) + + replicated = _grad_param(rank, device, sharded=False, sync_op="sum") + sharded = _grad_param(rank, device, sharded=True) + trainer = TrainerRank.__new__(TrainerRank) + reduced = trainer._reduce_dynamic_grads((replicated, sharded), scale_grads=0.5) + expected_replicated = 0.5 * sum(range(1, world + 1)) + torch.testing.assert_close( + reduced[0], torch.tensor([expected_replicated], device=device) + ) + torch.testing.assert_close( + reduced[1], torch.tensor([0.5 * (rank + 1)], device=device) + ) + norm = _distributed_grad_norm( + (replicated, sharded), + reduced, + ) + expected_norm = ( + expected_replicated**2 + + sum(float((0.5 * i) ** 2) for i in range(1, world + 1)) + ) ** 0.5 + assert norm == pytest.approx(expected_norm, rel=1e-6) + + _assert_replica_grad_reduction(rank, world, context_parallel=True) + _assert_replica_grad_reduction(rank, world, context_parallel=False) + _assert_distributed_optimizer_restore(device) + finally: + if getattr(ps, "model_parallel_is_initialized", lambda: False)(): + ps.destroy_model_parallel() + destroy_process_group() + + +def _assert_replica_grad_reduction( + rank: int, + world: int, + *, + context_parallel: bool, +) -> None: + ps.destroy_model_parallel() + torch.distributed.barrier() + ps.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=world if context_parallel else 1, + expert_model_parallel_size=1, + ) + device = torch.device("cuda", rank) + param = _grad_param(rank, device, sharded=False) + + trainer = TrainerRank.__new__(TrainerRank) + (reduced,) = trainer._reduce_dynamic_grads((param,), scale_grads=0.25) + expected = 0.25 * sum(range(1, world + 1)) + torch.testing.assert_close(reduced, torch.tensor([expected], device=device)) + assert _distributed_grad_norm((param,), (reduced,)) == pytest.approx(expected) + + +def _assert_distributed_optimizer_restore(device: torch.device) -> None: + ref = LoRASlotRef("checkpoint", "A") + adapter = _adapter("dense", rank=2, seed=11) + lora = LoRA("dense", 4, 5, 2, 32, torch.float32, device) + lora.load_lora_slot(ref, adapter, requires_grad=True) + trainer = _trainer_for(lora, device) + params = AdamParams(learning_rate=1e-3, weight_decay=0.0, grad_clip_norm=0.0) + x = torch.randn(3, 4, device=device) + + with use_lora_slot(ref): + lora(x).sum().backward() + trainer.optim_step(params=params, checkpoints=["A"]) + state = trainer.checkpoint_slot_optimizer_state("A") + assert state is not None + slot = lora._slot(ref) + assert slot is not None + adapter = { + "dense.lora_A.weight": slot.A_T.detach().T.contiguous(), + "dense.lora_B.weight": slot.B_T.detach().T.contiguous(), + } + with use_lora_slot(ref): + lora(x).sum().backward() + trainer.optim_step(params=params, checkpoints=["A"]) + + restored_lora = LoRA("dense", 4, 5, 2, 32, torch.float32, device) + restored = _trainer_for(restored_lora, device) + restored.load_checkpoint_slot("A", adapter, optimizer_state=state) + with use_lora_slot(ref): + restored_lora(x).sum().backward() + restored.optim_step(params=params, checkpoints=["A"]) + for expected, actual in zip( + lora.lora_slot_params(ref), restored_lora.lora_slot_params(ref), strict=True + ): + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +def _local_shard(full: torch.Tensor, rank: int, size: int) -> torch.Tensor: + return full[:, rank * size : (rank + 1) * size].clone().requires_grad_() + + +def _grad_param( + rank: int, device: torch.device, *, sharded: bool, sync_op: str = "none" +) -> torch.nn.Parameter: + param = torch.nn.Parameter(torch.ones(1, device=device)) + param.allreduce = True # type: ignore[attr-defined] + param.lora_shard_domain = "tp" # type: ignore[attr-defined] + param.lora_tp_sharded = sharded # type: ignore[attr-defined] + param.grad_sync_domain = "tp_default" # type: ignore[attr-defined] + param.grad_sync_op = sync_op # type: ignore[attr-defined] + param.grad = torch.tensor([float(rank + 1)], device=device) + return param + + +def _adapter(prefix: str, *, rank: int, seed: int) -> dict[str, torch.Tensor]: + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(seed) + return { + f"{prefix}.lora_A.weight": torch.randn( + rank, 4, generator=generator, device=device + ), + f"{prefix}.lora_B.weight": torch.randn( + 5, rank, generator=generator, device=device + ), + } + + +def _assert_checkpoint_recomputes_with( + expected_ref: LoRASlotRef, + ambient_ref: LoRASlotRef, + lora: LoRA, + checkpoint, + *checkpoint_args, +) -> None: + for param in lora.parameters(): + param.grad = None + x = torch.randn(3, 4, device="cuda", requires_grad=True) + with use_lora_slot(expected_ref): + y = checkpoint(lambda t: lora(t), *checkpoint_args, x) + with use_lora_slot(ambient_ref): + y.sum().backward() + assert lora._slot(expected_ref).A_T.grad is not None # type: ignore[union-attr] + assert lora._slot(ambient_ref).A_T.grad is None # type: ignore[union-attr] + + +def _assert_step_updates_only( + stepped_ref: LoRASlotRef, + frozen_ref: LoRASlotRef, + lora: LoRA, + trainer: TrainerRank, +) -> None: + for param in lora.parameters(): + param.grad = None + with use_lora_slot(stepped_ref): + lora(torch.randn(5, 4, device="cuda")).sum().backward() + before_stepped = [p.detach().clone() for p in lora.lora_slot_params(stepped_ref)] + before_frozen = [p.detach().clone() for p in lora.lora_slot_params(frozen_ref)] + trainer.optim_step( + params=AdamParams(learning_rate=1e-3, weight_decay=0.0, grad_clip_norm=1.0), + checkpoints=[stepped_ref.name or ""], + ) + assert any( + not torch.equal(before, after) + for before, after in zip( + before_stepped, lora.lora_slot_params(stepped_ref), strict=True + ) + ) + assert all( + torch.equal(before, after) + for before, after in zip( + before_frozen, lora.lora_slot_params(frozen_ref), strict=True + ) + ) + + +def _assert_reload_replaces_slot_optimizer( + ref: LoRASlotRef, + lora: LoRA, + trainer: TrainerRank, +) -> None: + assert ref.name is not None + old_params = trainer._checkpoint_slot_params_by_name[ref.name] + assert ref.name in trainer._dynamic_optimizers + + trainer.load_checkpoint_slot(ref.name, _adapter("dense", rank=3, seed=9)) + + new_params = trainer._checkpoint_slot_params_by_name[ref.name] + assert ref.name not in trainer._dynamic_optimizers + assert [tuple(param.shape) for param in new_params] == [(4, 3), (3, 5)] + assert all(old is not new for old, new in zip(old_params, new_params, strict=True)) + assert lora._slot(ref).rank == 3 # type: ignore[union-attr] + + +def _trainer_for(lora: LoRA, device: torch.device) -> TrainerRank: + trainer = TrainerRank.__new__(TrainerRank) + trainer.runtime = SimpleNamespace(model=[lora], optimizer=None) + trainer.device = device + trainer._slot_stack = [] + trainer._default_slot_ref = None + trainer._dynamic_optimizers = {} + trainer._checkpoint_slot_params_by_name = { + "A": tuple(lora.lora_slot_params(LoRASlotRef("checkpoint", "A"))), + "B": tuple(lora.lora_slot_params(LoRASlotRef("checkpoint", "B"))), + } + return trainer + + +@contextmanager +def _single_rank_model_parallel(): + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ["MASTER_PORT"] = str(_free_port()) + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("LOCAL_RANK", "0") + torch.cuda.set_device(0) + init_process_group("nccl", rank=0, world_size=1) + try: + ps.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=1, + expert_model_parallel_size=1, + ) + yield + finally: + if getattr(ps, "model_parallel_is_initialized", lambda: False)(): + ps.destroy_model_parallel() + destroy_process_group() + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py new file mode 100644 index 000000000..b0cff972f --- /dev/null +++ b/tests/unit/test_trainer_rank_validation.py @@ -0,0 +1,794 @@ +from __future__ import annotations + +from dataclasses import dataclass +import gc +import inspect +from types import SimpleNamespace +from typing import Any, cast + +import pytest +import torch + +from art.trainer_rank import ( + AdamParams, + ForwardInput, + ForwardOutput, + TopK, + TrainerRank, + TrainerRankMemoryError, + TrainerRankSlotStateError, + Unset, + _anchor_disconnected_outputs, + _MemoryCheck, + _MemoryProfile, + _validate_top_k, +) + + +class _Model: + vocab_size = 8 + + +class _FakeLoRASite(torch.nn.Module): + def __init__( + self, + prefix: str, + *, + device: torch.device | str = "cpu", + dtype: torch.dtype = torch.float32, + ) -> None: + super().__init__() + self.prefix = prefix + self.A_T = torch.nn.Parameter(torch.zeros(4, 2, device=device, dtype=dtype)) + self.B_T = torch.nn.Parameter(torch.zeros(2, 5, device=device, dtype=dtype)) + + def _expected_weight_keys(self, suffix: str) -> list[str]: + return [f"{self.prefix}.{suffix}.weight"] + + +class _NativeOptimizer: + config = None + param_groups: list[dict[str, object]] = [] + + def __init__(self) -> None: + self.step_calls = 0 + self.zero_grad_calls = 0 + + def step(self) -> tuple[bool, float, int | None]: + self.step_calls += 1 + raise AssertionError("TrainerRank must not step the native optimizer") + + def zero_grad(self) -> None: + self.zero_grad_calls += 1 + + +@dataclass(frozen=True) +class _SlotRef: + kind: str + name: str | None + + +def _runtime( + model: torch.nn.Module | None = None, + *, + optimizer: object | None = None, +) -> SimpleNamespace: + 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), + ) + + +def _target_request(token: int) -> ForwardInput[torch.Tensor, None, None, None]: + tokens = torch.tensor([token, token + 1], dtype=torch.long) + return ForwardInput(input_tokens=tokens, target_tokens=tokens) + + +def _indexed_outputs(plan: object, **_kwargs: object) -> list[ForwardOutput]: + return [ + ForwardOutput(torch.tensor([index], dtype=torch.float32), None, None, None) + for index in range(int(getattr(plan, "request_count"))) + ] + + +def _empty_outputs(plan: object, **_kwargs: object) -> list[ForwardOutput]: + return [ForwardOutput(None, None, None, None)] * int(getattr(plan, "request_count")) + + +def _stub_forward(mp, rank, out=_empty_outputs, dp=(0, 1), profiled=False) -> None: + mp.setattr(rank, "_dp_rank_and_size", lambda: dp) + mp.setattr(rank, "_run_flat_plan_with_memory_tracking", out) + if profiled: + mp.setattr(rank, "_all_ranks_have_memory_profile", lambda **_: True) + + +def _output_values(outputs: object) -> list[int]: + if isinstance(outputs, ForwardOutput): + target_logprobs = outputs.target_logprobs + assert isinstance(target_logprobs, torch.Tensor) + return [int(target_logprobs.item())] + values: list[int] = [] + for item in outputs: # type: ignore[union-attr] + values.extend(_output_values(item)) + return values + + +def _output_shape(outputs: object) -> object: + if isinstance(outputs, ForwardOutput): + return "output" + return [_output_shape(item) for item in outputs] # type: ignore[union-attr] + + +def _trainer_with_checkpoint( + monkeypatch: pytest.MonkeyPatch, + value: torch.Tensor, +) -> tuple[TrainerRank, torch.nn.Parameter]: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + param = torch.nn.Parameter(value.clone()) + 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), + ) + return trainer, param + + +def _tracked_targets( + trainer: TrainerRank, ref: _SlotRef, *scales: float +) -> list[torch.Tensor]: + tracked = trainer._track_slot_graph_outputs( + ref, # type: ignore[arg-type] + [ + ForwardOutput(torch.ones(1, requires_grad=True) * scale, None, None, None) + for scale in scales + ], + ) + return [cast(torch.Tensor, output.target_logprobs) for output in tracked] + + +def test_forward_input_validation() -> None: + with pytest.raises(ValueError, match="top_k must be >= 1"): + ForwardInput(input_tokens=torch.tensor([1]), top_k=0) + with pytest.raises(ValueError, match="cannot set both checkpoint and lora"): + ForwardInput(input_tokens=torch.tensor([1]), checkpoint="a", lora="b") + with pytest.raises(ValueError, match="top_k=9 exceeds vocabulary size 8"): + _validate_top_k(9, _Model()) # type: ignore[arg-type] + + +@pytest.mark.parametrize(("checkpoint", "expected"), ((Unset, Unset), (None, None))) +def test_forward_input_distinguishes_unset_and_base_checkpoint( + checkpoint: object, expected: object +) -> None: + request = ForwardInput(input_tokens=torch.tensor([1]), checkpoint=checkpoint) # type: ignore[arg-type] + + assert request.checkpoint is expected + assert request.lora is Unset + + +def test_forward_input_preserves_public_runtime_shape() -> None: + fields = tuple(ForwardInput.__dataclass_fields__) + assert tuple(inspect.signature(ForwardInput).parameters) == fields + assert ForwardInput.__match_args__ == fields + + +@pytest.mark.parametrize("depth", (0, 2)) +def test_trainer_rank_accepts_shared_prefix_depth(depth: int) -> None: + trainer = TrainerRank(_runtime(), shared_prefix_max_depth=depth) # type: ignore[arg-type] + + assert trainer.shared_prefix_max_depth == depth + + +def test_trainer_rank_adapter_stack_errors() -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + + with pytest.raises(RuntimeError, match="No pushed LoRA or checkpoint"): + trainer.pop_pushed_lora_or_checkpoint() + trainer._slot_stack.append(object()) # type: ignore[arg-type] + for load in (trainer.load_checkpoint_slot, trainer.load_lora_slot): + with pytest.raises(RuntimeError, match="Cannot load a LoRA/checkpoint"): + load("teacher", {}) + + +def test_trainer_rank_rejects_adapter_keys_without_installed_lora_site() -> None: + trainer = TrainerRank(_runtime(_FakeLoRASite("base.layer"))) # type: ignore[arg-type] + valid = { + "base.layer.lora_A.weight": torch.empty(1), + "base.layer.lora_B.weight": torch.empty(1), + } + trainer._prepare_adapter_model("checkpoint", "student", valid) + + with pytest.raises(ValueError, match="matching LoRA target modules"): + trainer._prepare_adapter_model( + "checkpoint", + "student", + {**valid, "base.other.lora_A.weight": torch.empty(1)}, + ) + + +def test_trainer_rank_normalizes_adapter_tensors_to_installed_site() -> None: + site = _FakeLoRASite("base.layer", dtype=torch.bfloat16) + trainer = TrainerRank(_runtime(site)) # type: ignore[arg-type] + adapter = { + "base.layer.lora_A.weight": torch.ones(3, 4, dtype=torch.float32), + "base.layer.lora_B.weight": torch.ones(5, 3, dtype=torch.float32), + } + + normalized = trainer._prepare_adapter_model("checkpoint", "student", adapter) + + assert all(tensor.device == site.A_T.device for tensor in normalized.values()) + assert all(tensor.dtype == torch.bfloat16 for tensor in normalized.values()) + + +def test_trainer_rank_default_forward_uses_explicit_base_slot() -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + + plan = trainer._plan_flat_forward([_target_request(1)]) + + assert len(plan.groups) == 1 + slot = plan.groups[0].slot_ref + assert slot is not None + assert getattr(slot, "kind") == "checkpoint" + assert getattr(slot, "name") is None + + +def test_optim_step_requires_loaded_checkpoint_slot() -> None: + optimizer = _NativeOptimizer() + trainer = TrainerRank(_runtime(optimizer=optimizer)) # type: ignore[arg-type] + + with pytest.raises(TrainerRankSlotStateError, match="loaded checkpoint slot"): + trainer.optim_step(params=AdamParams(learning_rate=1e-3)) + + assert optimizer.step_calls == 0 + + +def test_optim_step_rejects_loaded_slots_without_grads() -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + trainer._checkpoint_slot_params_by_name["student"] = ( + torch.nn.Parameter(torch.ones(2)), + ) + + with pytest.raises(TrainerRankSlotStateError, match="none have gradients"): + trainer.optim_step(params=AdamParams(learning_rate=1e-3)) + with pytest.raises(TrainerRankSlotStateError, match="no gradients"): + trainer.optim_step( + params=AdamParams(learning_rate=1e-3), + checkpoints=["student"], + ) + + +def test_optim_step_rejects_explicit_slot_subset_with_missing_grads( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + ready = torch.nn.Parameter(torch.ones(2)) + missing = torch.nn.Parameter(torch.ones(2)) + ready.grad = torch.ones_like(ready) + trainer._checkpoint_slot_params_by_name["ready"] = (ready,) + trainer._checkpoint_slot_params_by_name["missing"] = (missing,) + monkeypatch.setattr( + trainer, + "_reduce_dynamic_grads", + lambda params, **_kwargs: tuple(param.grad.float() for param in params), + ) + + with pytest.raises(TrainerRankSlotStateError, match="missing"): + trainer.optim_step( + params=AdamParams(learning_rate=1e-3), + checkpoints=["ready", "missing"], + ) + + +def test_optim_step_implicitly_steps_only_slots_with_grads( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + ready = torch.nn.Parameter(torch.ones(2)) + untouched = torch.nn.Parameter(torch.ones(2)) + ready.grad = torch.ones_like(ready) + trainer._checkpoint_slot_params_by_name["ready"] = (ready,) + trainer._checkpoint_slot_params_by_name["untouched"] = (untouched,) + monkeypatch.setattr( + trainer, + "_reduce_dynamic_grads", + lambda params, **_kwargs: tuple(param.grad.float() for param in params), + ) + + before_ready = ready.detach().clone() + before_untouched = untouched.detach().clone() + trainer.optim_step( + params=AdamParams(learning_rate=1e-2, weight_decay=0.0, grad_clip_norm=10.0) + ) + + assert "ready" in trainer._dynamic_optimizers + assert "untouched" not in trainer._dynamic_optimizers + assert not torch.equal(before_ready, ready) + torch.testing.assert_close(untouched, before_untouched) + + +def test_checkpoint_slot_optimizer_state_reproduces_exact_next_step( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adam = AdamParams( + learning_rate=3e-4, + beta1=0.8, + beta2=0.95, + weight_decay=0.1, + grad_clip_norm=10.0, + ) + + original, original_param = _trainer_with_checkpoint( + monkeypatch, torch.tensor([0.5, -0.25], dtype=torch.bfloat16) + ) + original_param.grad = torch.tensor([0.2, -0.4], dtype=torch.bfloat16) + original.optim_step(params=adam) + state = original.checkpoint_slot_optimizer_state("student") + assert state is not None + + restored, restored_param = _trainer_with_checkpoint( + monkeypatch, original_param.detach() + ) + restored._dynamic_optimizers["student"] = restored._restore_dynamic_optimizer( + "student", state + ) + for param in (original_param, restored_param): + param.grad = torch.tensor([-0.3, 0.1], dtype=torch.bfloat16) + original.optim_step(params=adam) + restored.optim_step(params=adam) + + torch.testing.assert_close(restored_param, original_param, atol=0, rtol=0) + original_state = original.checkpoint_slot_optimizer_state("student") + restored_state = restored.checkpoint_slot_optimizer_state("student") + assert original_state is not None and restored_state is not None + _assert_nested_tensors_equal(restored_state, original_state) + + +def test_dynamic_optimizer_keeps_fp32_master_weight_and_moments( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer, param = _trainer_with_checkpoint( + monkeypatch, torch.tensor([0.1], dtype=torch.bfloat16) + ) + + for _ in range(100): + param.grad = torch.ones_like(param) + trainer.optim_step( + params=AdamParams( + learning_rate=1e-5, + weight_decay=0.0, + grad_clip_norm=10.0, + ) + ) + + dynamic = trainer._dynamic_optimizers["student"] + assert dynamic.master_params[0].dtype == torch.float32 + assert param.item() < torch.tensor(0.1, dtype=torch.bfloat16).item() + state = dynamic.optimizer.state[dynamic.master_params[0]] + assert state["exp_avg"].dtype == torch.float32 + assert state["exp_avg_sq"].dtype == torch.float32 + + +@pytest.mark.parametrize( + ("corruption", "error"), + ( + ("layout", "topology or parameter layout"), + ("missing_master", "master parameters"), + ("shape", "topology or parameter layout"), + ), +) +def test_checkpoint_slot_optimizer_state_rejects_incompatible_state( + corruption: str, + error: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer, param = _trainer_with_checkpoint(monkeypatch, torch.ones(2)) + param.grad = torch.ones_like(param) + trainer.optim_step( + params=AdamParams(learning_rate=1e-2, weight_decay=0.0, grad_clip_norm=10.0) + ) + state = trainer.checkpoint_slot_optimizer_state("student") + assert state is not None + if corruption == "layout": + state["layout"] = {"different": True} + elif corruption == "missing_master": + state["master_params"] = () + restored, _ = _trainer_with_checkpoint( + monkeypatch, torch.ones(3 if corruption == "shape" else 2) + ) + with pytest.raises(TrainerRankSlotStateError, match=error): + restored._restore_dynamic_optimizer("student", state) + + +@pytest.mark.parametrize("operation", ("load", "step")) +def test_trainer_rank_rejects_mutating_slot_with_pending_graph( + operation: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + ref = _SlotRef("checkpoint", "teacher") + monkeypatch.setattr(trainer, "_slot_ref", lambda kind, name: _SlotRef(kind, name)) + target = _tracked_targets(trainer, ref, 2)[0] + guard = ( + (lambda: trainer._guard_slot_can_load(ref)) # type: ignore[arg-type] + if operation == "load" + else (lambda: trainer._guard_checkpoint_can_step("teacher")) + ) + + with pytest.raises(TrainerRankSlotStateError, match="Cannot"): + guard() + + target.sum().backward() + guard() + + +def test_trainer_rank_step_allows_missing_slot_graph_bookkeeping( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank.__new__(TrainerRank) + monkeypatch.setattr(trainer, "_slot_ref", lambda kind, name: _SlotRef(kind, name)) + + trainer._guard_checkpoint_can_step("student") + + +def test_trainer_rank_zero_grad_does_not_clear_live_slot_graphs() -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + ref = _SlotRef("lora", "teacher") + output = ForwardOutput( + None, + TopK( + torch.ones(1, requires_grad=True) * 2, + torch.ones(1, dtype=torch.long), + ), + None, + None, + ) + + tracked = trainer._track_slot_graph_outputs(ref, [output]) # type: ignore[arg-type] + trainer.zero_grad() + + assert tracked[0].top_k is not None + with pytest.raises(TrainerRankSlotStateError, match="live backward graph"): + trainer._guard_slot_can_load(ref) # type: ignore[arg-type] + + +def test_trainer_rank_retained_backward_keeps_slot_graph_guard() -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + ref = _SlotRef("checkpoint", "teacher") + target = _tracked_targets(trainer, ref, 2)[0] + + target.sum().backward(retain_graph=True) + with pytest.raises(TrainerRankSlotStateError, match="live backward graph"): + trainer._guard_slot_can_load(ref) # type: ignore[arg-type] + + target.sum().backward() + trainer._guard_slot_can_load(ref) # type: ignore[arg-type] + + +def test_trainer_rank_tracks_each_independent_output_graph() -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + ref = _SlotRef("checkpoint", "teacher") + first, second = _tracked_targets(trainer, ref, 2, 3) + + first.sum().backward() + with pytest.raises(TrainerRankSlotStateError, match="live backward graph"): + trainer._guard_slot_can_load(ref) # type: ignore[arg-type] + + second.sum().backward() + trainer._guard_slot_can_load(ref) # type: ignore[arg-type] + + +def test_trainer_rank_tracks_graph_after_output_is_replaced_by_loss() -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + ref = _SlotRef("checkpoint", "teacher") + target = _tracked_targets(trainer, ref, 2)[0] + loss = target.sum() + del target + gc.collect() + + with pytest.raises(TrainerRankSlotStateError, match="live backward graph"): + trainer._guard_slot_can_load(ref) # type: ignore[arg-type] + + loss.backward() + trainer._guard_slot_can_load(ref) # type: ignore[arg-type] + + +def test_trainer_rank_releases_abandoned_output_graph() -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + ref = _SlotRef("checkpoint", "teacher") + target = _tracked_targets(trainer, ref, 2)[0] + del target + gc.collect() + + trainer._guard_slot_can_load(ref) # type: ignore[arg-type] + + +def test_dp_rank_forward_preserves_nested_shape_for_inactive_requests() -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + request_a = ForwardInput(input_tokens=torch.tensor([1])) + request_b = ForwardInput(input_tokens=torch.tensor([2])) + + outputs = trainer.dp_rank_forward([[request_a], [request_b]]) + + assert len(outputs) == 2 + assert len(outputs[0]) == 1 + assert outputs[0][0].target_logprobs is None + assert outputs[1][0].target_logprobs is None + assert not hasattr(trainer, "forward") + assert not hasattr(trainer, "micro_batches") + + +def test_dp_rank_forward_supports_arbitrary_nested_depth( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + _stub_forward(monkeypatch, trainer, _indexed_outputs) + nested = [ + [[[[[_target_request(1)]]]]], + [[[[[_target_request(3), _target_request(5)]]]]], + ] + + outputs = cast(Any, trainer).dp_rank_forward(nested) + + assert _output_shape(outputs) == [ + [[[[["output"]]]]], + [[[[["output", "output"]]]]], + ] + assert _output_values(outputs) == [0, 1, 2] + + +def test_forward_micro_batches_uses_deterministic_dp_windows( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + _stub_forward(monkeypatch, trainer, dp=(1, 2)) + + batches = list( + trainer.forward_micro_batches([_target_request(i) for i in range(5)]) + ) + + assert [batch.indices for batch in batches] == [(1,), (3,), ()] + assert [len(batch.outputs) for batch in batches] == [1, 1, 0] + + +def test_forward_micro_batches_syncs_fit_decision_across_dp( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + _stub_forward(monkeypatch, trainer, dp=(1, 2), profiled=True) + sync_flags: list[bool] = [] + + def memory_check(required: int, *, sync_across_dp: bool = False) -> _MemoryCheck: + sync_flags.append(sync_across_dp) + return _MemoryCheck( + estimated_required_bytes=required, + available_bytes=1 << 30, + fits=True, + ) + + monkeypatch.setattr(trainer, "_memory_check_required", memory_check) + next(iter(trainer.forward_micro_batches([_target_request(i) for i in range(6)]))) + + assert sync_flags + assert all(sync_flags) + + +def test_forward_micro_batches_supports_arbitrary_nested_depth( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + _stub_forward(monkeypatch, trainer, _indexed_outputs, profiled=True) + expected = [ + [[[[[_target_request(1)]]]]], + [[[[[_target_request(3), _target_request(5)]]]]], + ] + nested = [(child for child in item) for item in expected] + + batches = list(cast(Any, trainer).forward_micro_batches(nested)) + + assert batches[0].inputs == expected + assert _output_shape(batches[0].outputs) == [ + [[[[["output"]]]]], + [[[[["output", "output"]]]]], + ] + assert _output_values(batches[0].outputs) == [0, 1, 2] + + +def test_forward_micro_batches_ramps_after_first_success( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + + def run(plan, **_kwargs): + trainer._memory_profiles[plan.signature] = _MemoryProfile( + bytes_per_token=0.0, + packed_tokens=plan.packed_tokens, + ) + return [ + ForwardOutput(None, None, None, None) for _ in range(plan.request_count) + ] + + _stub_forward(monkeypatch, trainer, run) + + batches = list( + trainer.forward_micro_batches([_target_request(i) for i in range(8)]) + ) + + assert batches[0].stats.global_count == 1 + assert batches[0].stats.cold_start + assert batches[1].stats.global_count > 1 + assert not batches[1].stats.cold_start + + +def test_forward_micro_batches_does_not_overtrust_tiny_memory_profile( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + monkeypatch.setattr(trainer, "_dp_rank_and_size", lambda: (0, 1)) + inputs = [_target_request(i) for i in range(64)] + tiny_plan = trainer._plan_flat_forward([inputs[0]]) + trainer._memory_profiles[tiny_plan.signature] = _MemoryProfile( + bytes_per_token=0.0, + packed_tokens=tiny_plan.packed_tokens, + ) + + candidate = trainer._select_next_micro_batch(inputs, 0) + + assert candidate.stats_global_count == 8 + assert candidate.plan.packed_tokens == 16 + + +def test_forward_micro_batches_tail_does_not_reset_stable_window( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + trainer._last_global_micro_batch_size = 64 + _stub_forward(monkeypatch, trainer, profiled=True) + monkeypatch.setattr( + trainer, + "_estimate_required_memory_bytes_from_values", + lambda **kwargs: kwargs["packed_tokens"], + ) + monkeypatch.setattr( + trainer, + "_memory_check_required", + lambda required, *, sync_across_dp=False: _MemoryCheck( + estimated_required_bytes=required, + available_bytes=128, + fits=required <= 128, + ), + ) + batches = list( + trainer.forward_micro_batches([_target_request(i) for i in range(130)]) + ) + + assert [batch.stats.global_count for batch in batches] == [64, 64, 2] + assert trainer._last_global_micro_batch_size == 64 + + +def test_forward_micro_batches_raises_when_smallest_batch_will_not_fit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + monkeypatch.setattr(trainer, "_dp_rank_and_size", lambda: (0, 1)) + monkeypatch.setattr( + trainer, + "_estimate_required_memory_bytes_from_values", + lambda **_kwargs: 4, + ) + monkeypatch.setattr( + trainer, + "_memory_check_required", + lambda required, *, sync_across_dp=False: _MemoryCheck( + estimated_required_bytes=required, + available_bytes=3, + fits=False, + ), + ) + with pytest.raises(TrainerRankMemoryError, match="smallest DP microbatch"): + next(iter(trainer.forward_micro_batches([_target_request(1)]))) + + +def test_forward_micro_batches_rejects_mismatched_replicated_counts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + trainer = TrainerRank(_runtime()) # type: ignore[arg-type] + import art.trainer_rank as trainer_rank + + monkeypatch.setattr(trainer_rank.dist, "is_available", lambda: True) + monkeypatch.setattr(trainer_rank.dist, "is_initialized", lambda: True) + monkeypatch.setattr(trainer_rank.dist, "get_world_size", lambda: 2) + + def gather(output, value): + output[:] = [value, value + 1] + + monkeypatch.setattr(trainer_rank.dist, "all_gather_object", gather) + + with pytest.raises(ValueError, match="same top-level input count"): + list(trainer.forward_micro_batches([_target_request(1)])) + + monkeypatch.setattr(trainer_rank.dist, "is_initialized", lambda: False) + _stub_forward(monkeypatch, trainer, dp=(1, 2)) + invalid = ForwardInput( + input_tokens=torch.tensor([1, 2]), target_tokens=torch.tensor([1, 2, 3]) + ) + with pytest.raises(ValueError, match="target_tokens"): + next(iter(trainer.forward_micro_batches([invalid, _target_request(1)]))) + + +def test_forward_plan_estimates_output_memory_for_request_combo() -> None: + class FakeGPT(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(())) + self.config = SimpleNamespace( + hidden_size=4, + num_layers=1, + padded_vocab_size=10, + ) + self.decoder = object() + + def _preprocess(self, *args: object, **kwargs: object) -> None: + return None + + trainer = TrainerRank(_runtime(FakeGPT())) # type: ignore[arg-type] + tokens = torch.tensor([[1, 2, 3]], dtype=torch.long) + labels = torch.stack((tokens, tokens + 1), dim=-1) + + request = ForwardInput( + input_tokens=tokens, + target_tokens=labels, + top_k=5, + logits=True, + hidden_states=True, + ) + plan = trainer._plan_flat_forward([request]) + estimate = trainer._estimate_flat_forward([request]) + + target_bytes = 3 * 2 * 4 + topk_bytes = 3 * 5 * (4 + 8) + logits_bytes = 3 * 10 * 4 + hidden_bytes = 3 * 4 * 4 + assert estimate is not None and estimate[0] == plan.packed_tokens + assert plan.output_bytes == target_bytes + topk_bytes + logits_bytes + hidden_bytes + + +def test_disconnected_outputs_keep_zero_graph_anchor() -> None: + hidden = torch.randn(2, 3, requires_grad=True) + disconnected = torch.zeros(4) + top_k = TopK(logprobs=torch.zeros(4, 2), tokens=torch.ones(4, 2, dtype=torch.long)) + + (anchored,), (anchored_top_k,) = _anchor_disconnected_outputs( + [disconnected], + [top_k], + hidden, + ) + + assert anchored is not None + assert anchored.requires_grad + assert anchored_top_k is not None + assert anchored_top_k.logprobs.requires_grad + torch.testing.assert_close(anchored, disconnected) + torch.testing.assert_close(anchored_top_k.logprobs, top_k.logprobs) + (anchored.sum() + anchored_top_k.logprobs.sum()).backward() + assert hidden.grad is not None + torch.testing.assert_close(hidden.grad, torch.zeros_like(hidden)) + + +def _assert_nested_tensors_equal(actual: object, expected: object) -> None: + if isinstance(expected, torch.Tensor): + assert isinstance(actual, torch.Tensor) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + elif isinstance(expected, dict): + assert isinstance(actual, dict) and actual.keys() == expected.keys() + actual_dict = cast(dict[Any, object], actual) + expected_dict = cast(dict[Any, object], expected) + for key in expected_dict: + _assert_nested_tensors_equal(actual_dict[key], expected_dict[key]) + elif isinstance(expected, tuple | list): + assert isinstance(actual, type(expected)) and len(actual) == len(expected) + for actual_item, expected_item in zip(actual, expected, strict=True): + _assert_nested_tensors_equal(actual_item, expected_item) + else: + assert actual == expected diff --git a/tests/unit/test_trainer_rank_weird_shapes.py b/tests/unit/test_trainer_rank_weird_shapes.py new file mode 100644 index 000000000..841169049 --- /dev/null +++ b/tests/unit/test_trainer_rank_weird_shapes.py @@ -0,0 +1,517 @@ +from __future__ import annotations + +from collections.abc import Callable, Iterable +from types import SimpleNamespace + +import pytest +import torch + +from art.megatron.prefix_tree_packing import ( + estimate_prefix_tree_packed_tokens, + prefix_tree_pack, +) +from art.trainer_rank import ( + AdapterSelection, + ForwardInput, + ForwardOutput, + TopK, + TrainerRank, + TrainerRankMemoryError, + Unset, + _flatten, + _MemoryCheck, + _MemoryProfile, +) + + +class _FakeGPT(torch.nn.Module): + def __init__(self, *, hidden_size: int = 8, vocab_size: int = 32) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros((), dtype=torch.float16)) + self.config = SimpleNamespace( + hidden_size=hidden_size, + num_layers=4, + padded_vocab_size=vocab_size, + ) + self.decoder = object() + + def _preprocess(self, *args: object, **kwargs: object) -> None: + return None + + +def _runtime() -> SimpleNamespace: + return SimpleNamespace( + model=[_FakeGPT()], + optimizer=None, + provider=SimpleNamespace(hidden_size=8, num_layers=4), + model_support_handler=SimpleNamespace(build_gdn_execution_spec=True), + ) + + +def _tokens(*values: int) -> torch.Tensor: + return torch.tensor(values, dtype=torch.long) + + +def _target_request( + tokens: torch.Tensor, + *, + target_count: int = 1, + top_k: int | None = None, + logits: bool = False, + hidden_states: bool = False, + checkpoint: AdapterSelection = Unset, + lora: AdapterSelection = Unset, +) -> ForwardInput: + labels = ( + tokens + if target_count == 1 + else torch.stack( + tuple(tokens + offset for offset in range(target_count)), + dim=-1, + ) + ) + return ForwardInput( + input_tokens=tokens, + target_tokens=labels, + top_k=top_k, + logits=logits, + hidden_states=hidden_states, + checkpoint=checkpoint, + lora=lora, + ) + + +def _set_packed_token_budget( + monkeypatch: pytest.MonkeyPatch, + rank: TrainerRank, + available: int | Callable[[], int], +) -> None: + monkeypatch.setattr( + rank, + "_estimate_required_memory_bytes_from_values", + lambda **kwargs: kwargs["packed_tokens"], + ) + + def check(required: int, *, sync_across_dp: bool = False) -> _MemoryCheck: + limit = available() if callable(available) else available + return _MemoryCheck(required, limit, required <= limit) + + monkeypatch.setattr(rank, "_memory_check_required", check) + + +def _ternary_tree_sequences() -> tuple[torch.Tensor, ...]: + # Shape: shared root, two continuation branches, and terminal nodes at + # several depths. This mirrors prompt -> continuation A/B -> terminal data. + root = [10, 11, 12] + left = root + [20, 21] + right = root + [30, 31, 32] + return ( + _tokens(*(root + [1])), + _tokens(*(left + [2])), + _tokens(*(left + [3, 4])), + _tokens(*(right + [5])), + _tokens(*(right + [6, 7])), + _tokens(80, 81), + ) + + +def _vineppo_like_inputs() -> list[list[ForwardInput]]: + groups: list[list[ForwardInput]] = [] + for prompt_index in range(4): + prompt = [100 + prompt_index, 200 + prompt_index, 201 + prompt_index] + trajectories = [] + for branch_index, completion_len in enumerate((1, 2, 4)): + completion = [300 + branch_index] * completion_len + tokens = _tokens(*(prompt + completion)) + trajectories.append( + _target_request( + tokens, + target_count=2 if branch_index == 2 else 1, + top_k=5 if branch_index == 1 else None, + hidden_states=branch_index == 0, + ) + ) + groups.append(trajectories) + return groups + + +def _random_tree_sequences(seed: int, *, max_depth: int) -> tuple[torch.Tensor, ...]: + generator = torch.Generator().manual_seed(seed) + out: list[torch.Tensor] = [] + + def randint(low: int, high: int) -> int: + return int(torch.randint(low, high + 1, (), generator=generator).item()) + + def segment(depth: int) -> list[int]: + return [depth * 100 + randint(1, 40) for _ in range(randint(1, 4))] + + def walk(prefix: list[int], depth: int) -> None: + if depth >= max_depth or randint(0, 2) == 0: + out.append(_tokens(*(prefix + segment(depth)))) + return + shared = prefix + segment(depth) + out.append(_tokens(*shared)) + walk(shared + [10 + depth], depth + 1) + walk(shared + [20 + depth], depth + 1) + + walk([], 0) + return tuple(out) + + +@pytest.mark.parametrize("max_depth", (0, 1, 2, 4)) +def test_pack_estimator_matches_ternary_and_random_trees(max_depth: int) -> None: + cases = [ + _ternary_tree_sequences(), + _random_tree_sequences(3, max_depth=4), + _random_tree_sequences(99, max_depth=5), + ] + + for sequences in cases: + pack = prefix_tree_pack(sequences, max_depth=max_depth) + + assert estimate_prefix_tree_packed_tokens( + sequences, max_depth=max_depth + ) == int(pack.tokens.numel()) + for sequence, positions in zip( + sequences, pack.positions_by_sequence, strict=True + ): + torch.testing.assert_close(pack.tokens.reshape(-1)[positions], sequence) + + +def test_shared_trainable_tokens_accumulate_independent_output_gradients() -> None: + sequences = ( + torch.tensor([1, 2, 3], dtype=torch.long), + torch.tensor([1, 2, 3], dtype=torch.long), + ) + pack = prefix_tree_pack(sequences, max_depth=4) + hidden = torch.randn(int(pack.tokens.numel()), 3, requires_grad=True) + weights = (2.0, 5.0) + + loss = sum( + weight * hidden.index_select(0, positions).sum() + for weight, positions in zip(weights, pack.positions_by_sequence, strict=True) + ) + loss.backward() + + expected = torch.zeros_like(hidden) + for weight, positions in zip(weights, pack.positions_by_sequence, strict=True): + expected.index_add_( + 0, + positions, + torch.full((int(positions.numel()), 3), weight, dtype=hidden.dtype), + ) + torch.testing.assert_close(hidden.grad, expected) + + +def test_planner_handles_vineppo_nested_shape_and_request_mix() -> None: + rank = TrainerRank(_runtime(), shared_prefix_max_depth=3) # type: ignore[arg-type] + inputs = _vineppo_like_inputs() + flat = list(_flatten(inputs)) + + plan = rank._plan_flat_forward(flat) + estimate = rank._estimate_flat_forward(flat) + + assert estimate is not None + packed_tokens, output_bytes, signature = estimate + assert packed_tokens == plan.packed_tokens + assert output_bytes == plan.output_bytes + assert signature == plan.signature + assert plan.request_count == 12 + assert plan.signature.request_mix == ( + "target:(2,)", + "target:single+hidden", + "target:single+topk:5", + ) + + +def test_forward_micro_batches_preserves_nested_vineppo_groups( + monkeypatch: pytest.MonkeyPatch, +) -> None: + rank = TrainerRank(_runtime(), shared_prefix_max_depth=2) # type: ignore[arg-type] + monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1)) + monkeypatch.setattr(rank, "_all_ranks_have_memory_profile", lambda **_kwargs: True) + monkeypatch.setattr( + rank, + "_memory_check", + lambda plan, *, sync_across_dp=False: _MemoryCheck( + plan.packed_tokens, 10_000, True + ), + ) + monkeypatch.setattr( + rank, + "_run_flat_plan_with_memory_tracking", + lambda plan, **_kwargs: [ + ForwardOutput(None, None, None, None) for _ in range(plan.request_count) + ], + ) + groups = _vineppo_like_inputs() + + micro_batches = list(rank.forward_micro_batches(groups)) + + assert [batch.indices for batch in micro_batches] == [(0, 1, 2, 3)] + assert micro_batches[0].select(groups) == groups + assert len(micro_batches[0].outputs) == 4 + assert all( + isinstance(group_outputs, list) and len(group_outputs) == 3 + for group_outputs in micro_batches[0].outputs + ) + + +def test_adaptive_planner_materializes_only_final_large_candidate( + monkeypatch: pytest.MonkeyPatch, +) -> None: + rank = TrainerRank(_runtime(), shared_prefix_max_depth=3) # type: ignore[arg-type] + rank._last_global_micro_batch_size = 32 + monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1)) + monkeypatch.setattr(rank, "_all_ranks_have_memory_profile", lambda **_kwargs: True) + plan_calls = 0 + estimate_calls = 0 + original_plan = rank._plan_flat_forward + original_estimate = rank._estimate_flat_forward + inputs = [ + _target_request( + _tokens(1, 2, 3, index % 7, index), + target_count=2 if index % 5 == 0 else 1, + top_k=3 if index % 4 == 0 else None, + hidden_states=index % 9 == 0, + ) + for index in range(96) + ] + limit = rank._estimate_flat_forward(inputs[:40]) + assert limit is not None + limit_packed_tokens = limit[0] + + def plan(requests): + nonlocal plan_calls + plan_calls += 1 + return original_plan(requests) + + def estimate(requests): + nonlocal estimate_calls + estimate_calls += 1 + return original_estimate(requests) + + monkeypatch.setattr(rank, "_plan_flat_forward", plan) + monkeypatch.setattr(rank, "_estimate_flat_forward", estimate) + _set_packed_token_budget(monkeypatch, rank, limit_packed_tokens) + + candidate = rank._select_next_micro_batch(inputs, 0) + + assert candidate.stats_global_count == 40 + assert plan_calls == 1 + assert estimate_calls <= 10 + assert candidate.rejected_candidates <= 8 + + +def test_adaptive_planner_globally_falls_back_when_one_rank_cannot_estimate( + monkeypatch: pytest.MonkeyPatch, +) -> None: + rank = TrainerRank(_runtime()) # type: ignore[arg-type] + monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 2)) + monkeypatch.setattr(rank, "_all_ranks_true", lambda _local: False) + plans = 0 + original = rank._plan_flat_forward + + def plan(requests): + nonlocal plans + plans += 1 + return original(requests) + + monkeypatch.setattr(rank, "_plan_flat_forward", plan) + candidate = rank._select_next_micro_batch( + [_target_request(_tokens(index)) for index in range(4)], 0 + ) + + assert candidate.stats_global_count == 2 + assert plans == 1 + + +def test_adaptive_planner_probes_new_heterogeneous_signatures( + monkeypatch: pytest.MonkeyPatch, +) -> None: + rank = TrainerRank(_runtime()) # type: ignore[arg-type] + monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1)) + monkeypatch.setattr(rank, "_resolve_slot_ref", lambda request: request.checkpoint) + inputs = [ + _target_request(_tokens(index), checkpoint=f"S{index % 4}") + for index in range(16) + ] + + first = rank._select_next_micro_batch(inputs, 0) + rank._memory_profiles[first.plan.signature] = _MemoryProfile(0.0, 1_000_000) + rank._last_global_micro_batch_size = 1 + second = rank._select_next_micro_batch(inputs, 1) + rank._memory_profiles[second.plan.signature] = _MemoryProfile(0.0, 1_000_000) + rank._last_global_micro_batch_size = 2 + third = rank._select_next_micro_batch(inputs, 3) + + assert [ + first.stats_global_count, + second.stats_global_count, + third.stats_global_count, + ] == [ + 1, + 2, + 4, + ] + + +def test_adaptive_planner_grows_stable_window_to_largest_aligned_fit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + rank = TrainerRank(_runtime(), shared_prefix_max_depth=1) # type: ignore[arg-type] + rank._last_global_micro_batch_size = 512 + monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1)) + monkeypatch.setattr(rank, "_all_ranks_have_memory_profile", lambda **_kwargs: True) + _set_packed_token_budget(monkeypatch, rank, 700) + + candidate = rank._select_next_micro_batch( + [_target_request(_tokens(index)) for index in range(900)], + 0, + ) + + assert candidate.stats_global_count == 672 + assert candidate.rejected_candidates <= 2 + + +def test_forward_micro_batches_shrinks_when_memory_budget_drops( + monkeypatch: pytest.MonkeyPatch, +) -> None: + rank = TrainerRank(_runtime(), shared_prefix_max_depth=2) # type: ignore[arg-type] + monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1)) + monkeypatch.setattr(rank, "_all_ranks_have_memory_profile", lambda **_kwargs: True) + inputs = [_target_request(_tokens(1, 2, 3, index)) for index in range(14)] + first_limit = rank._estimate_flat_forward(inputs[:8]) + tail_limit = rank._estimate_flat_forward(inputs[8:11]) + assert first_limit is not None + assert tail_limit is not None + first_limit_packed_tokens = first_limit[0] + tail_limit_packed_tokens = tail_limit[0] + available = {"packed_tokens": first_limit_packed_tokens} + plan_calls = 0 + original_plan = rank._plan_flat_forward + + def plan(requests): + nonlocal plan_calls + plan_calls += 1 + return original_plan(requests) + + def run(plan, **_kwargs): + if available["packed_tokens"] == first_limit_packed_tokens: + available["packed_tokens"] = tail_limit_packed_tokens + return [ + ForwardOutput(None, None, None, None) for _ in range(plan.request_count) + ] + + monkeypatch.setattr(rank, "_plan_flat_forward", plan) + _set_packed_token_budget(monkeypatch, rank, lambda: available["packed_tokens"]) + monkeypatch.setattr(rank, "_run_flat_plan_with_memory_tracking", run) + + batches = list(rank.forward_micro_batches(inputs)) + + assert [batch.stats.global_count for batch in batches] == [8, 3, 3] + assert [batch.stats.available_bytes for batch in batches] == [ + first_limit_packed_tokens, + tail_limit_packed_tokens, + tail_limit_packed_tokens, + ] + assert [batch.indices for batch in batches] == [ + tuple(range(8)), + (8, 9, 10), + (11, 12, 13), + ] + assert plan_calls == len(batches) + + +def test_heterogeneous_slots_split_packing_without_losing_output_estimates( + monkeypatch: pytest.MonkeyPatch, +) -> None: + rank = TrainerRank(_runtime(), shared_prefix_max_depth=4) # type: ignore[arg-type] + monkeypatch.setattr( + TrainerRank, + "_slot_ref", + staticmethod(lambda kind, name: (kind, name)), + ) + rank.set_checkpoint("student") + requests = [ + _target_request(_tokens(1, 2, 3), top_k=3), + _target_request(_tokens(1, 2, 4), checkpoint=None, logits=True), + _target_request(_tokens(1, 2, 5), lora="teacher", hidden_states=True), + _target_request(_tokens(1, 2, 6), checkpoint="critic", target_count=4), + ] + + plan = rank._plan_flat_forward(requests) + estimate = rank._estimate_flat_forward(requests) + + assert estimate is not None + packed_tokens, output_bytes, signature = estimate + assert packed_tokens == plan.packed_tokens + assert output_bytes == plan.output_bytes + assert signature == plan.signature + assert plan.signature.slot_group_count == 4 + assert {group.slot_ref for group in plan.groups} == { + ("checkpoint", "student"), + ("checkpoint", None), + ("lora", "teacher"), + ("checkpoint", "critic"), + } + + +@pytest.mark.parametrize("api", ("dp_rank_forward", "forward_micro_batches")) +def test_forward_raises_before_expected_oom_with_actionable_context( + api: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + rank = TrainerRank(_runtime()) # type: ignore[arg-type] + if api == "dp_rank_forward": + monkeypatch.setattr( + rank, + "_memory_check", + lambda plan, **_kwargs: _MemoryCheck(plan.output_bytes + 1, 0, False), + ) + else: + monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1)) + monkeypatch.setattr( + rank, + "_estimate_required_memory_bytes_from_values", + lambda **_kwargs: 99, + ) + monkeypatch.setattr( + rank, + "_memory_check_required", + lambda required, **_kwargs: _MemoryCheck(required, 1, False), + ) + request = [_target_request(_tokens(1, 2, 3), logits=True)] + + with pytest.raises(TrainerRankMemoryError) as exc_info: + ( + rank.dp_rank_forward(request) + if api == "dp_rank_forward" + else next(iter(rank.forward_micro_batches(request))) + ) + + message = str(exc_info.value) + assert api in message + assert "packed_tokens=" in message + assert "logical_tokens=" in message + assert "output_gb=" in message + assert "Use smaller top-level items" in message + + +def test_flatten_rejects_dicts_to_avoid_silent_top_level_shape_changes() -> None: + with pytest.raises(TypeError, match="dict was passed directly"): + list(_flatten({"bad": _target_request(_tokens(1, 2))})) # type: ignore[arg-type] + + +def test_no_output_requests_do_not_pack_or_consume_compute_memory() -> None: + rank = TrainerRank(_runtime()) # type: ignore[arg-type] + requests: Iterable[ForwardInput] = [ + ForwardInput(input_tokens=_tokens(1, 2, 3)), + ForwardInput(input_tokens=_tokens(1, 2, 4)), + ] + + plan = rank._plan_flat_forward(list(requests)) + + assert plan.groups == () + assert plan.packed_tokens == 0 + assert rank._memory_check(plan).estimated_required_bytes == 0