diff --git a/.github/workflows/build-gpu-image.yml b/.github/workflows/build-gpu-image.yml index d4d74523c..f2c4ce117 100644 --- a/.github/workflows/build-gpu-image.yml +++ b/.github/workflows/build-gpu-image.yml @@ -30,6 +30,11 @@ on: required: true default: true type: boolean + prewarm_modal: + description: "Prebuild the pushed image in Modal when auth is configured" + required: true + default: true + type: boolean prewarm_infras: description: "Space/comma-separated Kubernetes-backed infra targets to prewarm" required: true @@ -177,11 +182,16 @@ jobs: PULL_IMAGE_REPO: ${{ inputs.pull_image_repo || 'docker.io/bradhiltonnw/art-gpu' }} IMAGE_TAG: ${{ inputs.tag }} NO_CACHE: ${{ inputs.no_cache }} + MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} + MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} + PREWARM_MODAL_INPUT: ${{ inputs.prewarm_modal }} PREWARM_NODES: ${{ inputs.prewarm_nodes }} PREWARM_TIMEOUT: ${{ inputs.prewarm_timeout }} run: | IMAGE_TAG="${IMAGE_TAG:-latest}" NO_CACHE="${NO_CACHE:-false}" + export PREWARM_MODAL="${PREWARM_MODAL:-auto}" + PREWARM_MODAL_INPUT="${PREWARM_MODAL_INPUT:-true}" PREWARM_NODES="${PREWARM_NODES:-true}" PREWARM_TIMEOUT="${PREWARM_TIMEOUT:-30m}" @@ -197,6 +207,10 @@ jobs: args+=(--no-cache) fi + if [ "${PREWARM_MODAL_INPUT}" = "false" ]; then + args+=(--no-prewarm-modal) + fi + if [ "${PREWARM_NODES}" != "true" ]; then args+=(--no-prewarm-nodes) fi diff --git a/.gitignore b/.gitignore index d1f4ebd59..0dfae3afe 100644 --- a/.gitignore +++ b/.gitignore @@ -9,6 +9,8 @@ data/cache.db streaming-chat-completions/ unsloth_compiled_cache/ wandb/ +!/typings/wandb/ +!/typings/wandb/** docs/node_modules/ dist/ replays/ diff --git a/dev/megatron_review_perf.py b/dev/megatron_review_perf.py new file mode 100644 index 000000000..4da5a8ad0 --- /dev/null +++ b/dev/megatron_review_perf.py @@ -0,0 +1,868 @@ +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass +import json +from pathlib import Path +import time + +import numpy as np +import torch +from torch.nn.attention.flex_attention import AuxRequest, BlockMask +from torch.nn.attention.flex_attention import create_block_mask as torch_block_mask +import typer + +from art.megatron.context_parallel.block_mask import ( + build_block_mask_from_context, + prepare_block_mask_context, +) +from art.megatron.context_parallel.builder import build_shared_prefix_attention_spec +from art.megatron.context_parallel.executor import _build_stage_execution_spec +from art.megatron.context_parallel.runtime import ( + _RUNTIME_PLAN_CACHE, + get_or_build_runtime_plan, +) +from art.megatron.context_parallel.types import ( + ContextParallelConfig, + FlexMaskSpec, + ParallelTopology, + StageExecutionSpec, + StagePlan, +) +from art.megatron.flex_attn.compiled import ( + normalize_sparse_block_size, + sparse_compiled_flex_attention, +) +from art.megatron.shared_prefix_packing import SharedPrefixPack, pack_shared_prefixes + + +def main( + workload: str = "austin_198k", + max_depth: int = 1, + cp_size: int = 4, + block_size: int = 128, + prefix_families: int = 4, + prefix_len: int = 1024, + mid_prefixes_per_family: int = 1, + mid_prefix_len: int = 0, + branches_per_prefix: int = 8, + completion_len: int = 128, + warmup: int = 3, + repeat: int = 10, + shape_variants: int = 4, + validate_torch: bool = True, + validate_torch_token_cap: int = 32768, + run_flex: bool = True, + flex_token_cap: int = 8192, + flex_heads: int = 2, + flex_head_dim: int = 128, + flex_mask_variants: str = "current,causal_abs_only", + max_block_mask_build_ms: float | None = None, + max_cp_planning_cold_ms: float | None = None, + output_jsonl: Path = Path(".local/trainer_rank_review/block_mask_flex.jsonl"), +) -> None: + if warmup < 0 or repeat < 1: + raise ValueError("warmup must be >= 0 and repeat must be >= 1") + output_jsonl.parent.mkdir(parents=True, exist_ok=True) + + pack = _pack_workload( + workload=workload, + max_depth=max_depth, + prefix_families=prefix_families, + prefix_len=prefix_len, + mid_prefixes_per_family=mid_prefixes_per_family, + mid_prefix_len=mid_prefix_len, + branches_per_prefix=branches_per_prefix, + completion_len=completion_len, + ) + spec = build_shared_prefix_attention_spec( + group_ids=pack.group_ids, + parent_ids=pack.parent_ids, + ) + config = ContextParallelConfig(block_size=block_size) + topology = ParallelTopology(cp=cp_size) + base = { + "workload": workload, + "max_depth": max_depth, + "cp_size": cp_size, + "block_size": block_size, + "packed_tokens": int(pack.tokens.numel()), + "logical_tokens": _logical_tokens(pack), + "warmup": warmup, + "repeat": repeat, + "validate_torch": validate_torch, + "validate_torch_token_cap": validate_torch_token_cap, + } + + plan, plan_ms = _bench_cpu( + lambda: _build_cp_plan(pack, spec, topology, config), + warmup=warmup, + repeat=repeat, + before_each=_RUNTIME_PLAN_CACHE.clear, + ) + _write( + output_jsonl, + { + **base, + "case": "cp_planning_cold", + "ms": plan_ms, + **_plan_stats(plan), + }, + ) + + cached_plan, cached_plan_ms = _bench_cpu( + lambda: _build_cp_plan(pack, spec, topology, config), + warmup=warmup, + repeat=repeat, + ) + _write( + output_jsonl, + { + **base, + "case": "cp_planning_cached", + "ms": cached_plan_ms, + **_plan_stats(cached_plan), + }, + ) + + stage_masks, mask_ms = _bench_cpu( + lambda: _build_stage_masks(pack, plan, config), + warmup=warmup, + repeat=repeat, + ) + masks = tuple(mask for mask, _ in stage_masks) + torch_validation_skipped = _torch_validation_skip_reason( + validate_torch=validate_torch, + packed_tokens=int(pack.tokens.numel()), + token_cap=validate_torch_token_cap, + ) + if torch_validation_skipped is None: + for mask, slices in stage_masks: + _assert_matches_torch_block_mask(mask, slices=slices) + _write( + output_jsonl, + { + **base, + "case": "block_mask_build", + "ms": mask_ms, + "torch_validation_skipped": torch_validation_skipped, + **_mask_stats(masks), + }, + ) + _check_threshold("block_mask_build", mask_ms, max_block_mask_build_ms) + _check_threshold("cp_planning_cold", plan_ms, max_cp_planning_cold_ms) + + if run_flex: + for record in _flex_records( + pack, + plan, + config, + warmup=warmup, + repeat=repeat, + token_cap=flex_token_cap, + heads=flex_heads, + head_dim=flex_head_dim, + variants=_csv_values(flex_mask_variants), + ): + _write(output_jsonl, {**base, **record}) + + for variant in range(shape_variants): + variant_pack = _pack_workload( + workload="regular", + max_depth=max_depth, + prefix_families=prefix_families, + prefix_len=prefix_len + variant * 17, + mid_prefixes_per_family=mid_prefixes_per_family, + mid_prefix_len=mid_prefix_len + variant * 3, + branches_per_prefix=branches_per_prefix, + completion_len=completion_len + variant * 11, + ) + variant_spec = build_shared_prefix_attention_spec( + group_ids=variant_pack.group_ids, + parent_ids=variant_pack.parent_ids, + ) + variant_plan, variant_plan_ms = _bench_cpu( + lambda pack=variant_pack, spec=variant_spec: _build_cp_plan( + pack, + spec, + topology, + config, + ), + warmup=0, + repeat=1, + before_each=_RUNTIME_PLAN_CACHE.clear, + ) + variant_stage_masks, variant_mask_ms = _bench_cpu( + lambda pack=variant_pack, plan=variant_plan: _build_stage_masks( + pack, + plan, + config, + ), + warmup=0, + repeat=1, + ) + variant_masks = tuple(mask for mask, _ in variant_stage_masks) + variant_torch_validation_skipped = _torch_validation_skip_reason( + validate_torch=validate_torch, + packed_tokens=int(variant_pack.tokens.numel()), + token_cap=validate_torch_token_cap, + ) + if variant_torch_validation_skipped is None: + for mask, slices in variant_stage_masks: + _assert_matches_torch_block_mask(mask, slices=slices) + _write( + output_jsonl, + { + **base, + "case": "shape_variant", + "variant": variant, + "variant_packed_tokens": int(variant_pack.tokens.numel()), + "variant_logical_tokens": _logical_tokens(variant_pack), + "cp_planning_ms": variant_plan_ms, + "block_mask_build_ms": variant_mask_ms, + "torch_validation_skipped": variant_torch_validation_skipped, + **_plan_stats(variant_plan), + **_mask_stats(variant_masks), + }, + ) + + print(f"wrote review perf records to {output_jsonl}", flush=True) + + +def _pack_workload( + *, + workload: str, + max_depth: int, + prefix_families: int, + prefix_len: int, + mid_prefixes_per_family: int, + mid_prefix_len: int, + branches_per_prefix: int, + completion_len: int, +) -> SharedPrefixPack: + sequences = ( + _austin_sequences() + if workload == "austin_198k" + else _austin_varied_sequences() + if workload == "austin_varied" + else _regular_sequences( + prefix_families=prefix_families, + prefix_len=prefix_len, + mid_prefixes_per_family=mid_prefixes_per_family, + mid_prefix_len=mid_prefix_len, + branches_per_prefix=branches_per_prefix, + completion_len=completion_len, + ) + ) + return pack_shared_prefixes(sequences, max_depth=max_depth) + + +def _austin_sequences() -> tuple[torch.Tensor, ...]: + return tuple( + torch.cat( + ( + _tokens(family * 10_000_019, 5000), + _tokens(family * 10_000_019 + branch * 1009 + 17, 100), + ) + ) + for family in range(30) + for branch in range(16) + ) + + +def _austin_varied_sequences() -> tuple[torch.Tensor, ...]: + sequences: list[torch.Tensor] = [] + for family in range(30): + family_base = family * 10_000_019 + prefix_len = 4500 + ((family * 137) % 1001) + root = _tokens(family_base, prefix_len) + branch_count = 10 + ((family * 7) % 13) + for branch in range(branch_count): + completion_len = 32 + ((family * 19 + branch * 23) % 145) + sequences.append( + torch.cat( + ( + root, + _tokens( + family_base + branch * 1009 + 17, + completion_len, + ), + ) + ) + ) + return tuple(sequences) + + +def _regular_sequences( + *, + prefix_families: int, + prefix_len: int, + mid_prefixes_per_family: int, + mid_prefix_len: int, + branches_per_prefix: int, + completion_len: int, +) -> tuple[torch.Tensor, ...]: + sequences = [] + for family in range(max(1, prefix_families)): + family_base = family * 10_000_019 + root = _tokens(family_base, max(1, prefix_len)) + for mid in range(max(1, mid_prefixes_per_family)): + mid_prefix = _tokens( + family_base + 1_000_003 + mid * 100_003, + max(0, mid_prefix_len), + ) + prefix = torch.cat((root, mid_prefix)) + for branch in range(max(1, branches_per_prefix)): + sequences.append( + torch.cat( + ( + prefix, + _tokens( + family_base + mid * 100_003 + branch * 1009 + 17, + max(1, completion_len), + ), + ) + ) + ) + return tuple(sequences) + + +def _tokens(offset: int, length: int) -> torch.Tensor: + return (torch.arange(length, dtype=torch.long) + offset) % 32_000 + 100 + + +def _build_cp_plan( + pack: SharedPrefixPack, + spec: object, + topology: ParallelTopology, + config: ContextParallelConfig, +) -> object: + return get_or_build_runtime_plan( + spec, + topology=topology, + config=config, + original_seq_len=int(pack.tokens.numel()), + ) + + +def _build_stage_masks( + pack: SharedPrefixPack, + plan: object, + config: ContextParallelConfig, +) -> tuple[tuple[BlockMask, tuple[object, ...]], ...]: + masks = [] + context = prepare_block_mask_context( + group_ids=pack.group_ids[0], + parent_ids=pack.parent_ids[0], + ) + for rank_plan in plan: + for stage in rank_plan.stage_plans: + if stage.mask_metadata is None: + continue + execution_spec = _stage_execution_spec(stage, config) + mask_metadata = execution_spec.mask_metadata or stage.mask_metadata + if mask_metadata is None: + continue + mask = build_block_mask_from_context( + FlexMaskSpec( + q_len=execution_spec.q_len, + k_len=execution_spec.k_len, + block_size=_sparse_block_size(config), + slices=stage.slices, + exact_mask=mask_metadata, + ), + context=context, + device=torch.device("cpu"), + validate=False, + ) + if mask is not None: + masks.append((mask, tuple(stage.slices))) + return tuple(masks) + + +def _flex_records( + pack: SharedPrefixPack, + plan: object, + config: ContextParallelConfig, + *, + warmup: int, + repeat: int, + token_cap: int, + heads: int, + head_dim: int, + variants: Sequence[str], +) -> list[dict[str, object]]: + if not torch.cuda.is_available(): + return [{"case": "flex_attention_fwd_bwd", "skipped": "cuda_unavailable"}] + device = torch.device("cuda") + stage_cases = _build_stage_flex_cases( + pack, + plan, + config, + device=device, + ) + if not stage_cases: + return [{"case": "flex_attention_fwd_bwd", "skipped": "no_stage_masks"}] + largest_stage = max(max(case.q_len, case.k_len) for case in stage_cases) + if int(largest_stage) > int(token_cap): + return [ + { + "case": "flex_attention_fwd_bwd", + "skipped": "stage_tokens_exceed_flex_token_cap", + "flex_token_cap": int(token_cap), + "largest_stage_tokens": int(largest_stage), + } + ] + records: list[dict[str, object]] = [] + base_tensors = _stage_tensors( + stage_cases, + heads=heads, + head_dim=head_dim, + device=device, + ) + for variant in variants: + block_masks = [] + try: + block_masks = [ + _stage_variant_block_mask(case, variant, device=device) + for case in stage_cases + ] + except Exception as exc: + records.append( + { + "case": "flex_attention_fwd_bwd", + "flex_mask_variant": variant, + "compile_error": type(exc).__name__, + "compile_error_message": str(exc).splitlines()[0][:500], + "flex_heads": heads, + "flex_head_dim": head_dim, + } + ) + continue + qkv = [ + ( + q.detach().clone().requires_grad_(True), + k.detach().clone().requires_grad_(True), + v.detach().clone().requires_grad_(True), + ) + for q, k, v in base_tensors + ] + + def step() -> None: + loss = torch.zeros((), device=device, dtype=torch.float32) + for (q, k, v), block_mask in zip(qkv, block_masks, strict=True): + q.grad = None + k.grad = None + v.grad = None + out, _aux = sparse_compiled_flex_attention( + q, + k, + v, + block_mask=block_mask, + scale=float(head_dim) ** -0.5, + enable_gqa=False, + return_aux=AuxRequest(lse=True), + ) + loss = loss + out.float().sum() + loss.backward() + + try: + torch.cuda.reset_peak_memory_stats() + torch.cuda.synchronize() + first_started = time.perf_counter() + step() + torch.cuda.synchronize() + first_call_ms = round((time.perf_counter() - first_started) * 1000.0, 3) + ms = _bench_cuda(step, warmup=warmup, repeat=repeat) + except Exception as exc: + torch.cuda.empty_cache() + records.append( + { + "case": "flex_attention_fwd_bwd", + "flex_mask_variant": variant, + "compile_error": type(exc).__name__, + "compile_error_message": str(exc).splitlines()[0][:500], + "flex_heads": heads, + "flex_head_dim": head_dim, + **_stage_flex_stats(stage_cases), + } + ) + continue + records.append( + { + "case": "flex_attention_fwd_bwd", + "flex_mask_variant": variant, + "first_call_ms": first_call_ms, + "ms": ms, + "packed_tok_s": round(int(pack.tokens.numel()) * 1000.0 / ms, 3), + "flex_heads": heads, + "flex_head_dim": head_dim, + **_stage_flex_stats(stage_cases), + "peak_memory_gb": round(torch.cuda.max_memory_allocated() / 1024**3, 3), + } + ) + return records + + +@dataclass(frozen=True) +class _StageFlexCase: + rank: int + stage_index: int + q_len: int + k_len: int + logical_q_len: int + logical_k_len: int + block_mask: BlockMask + q_abs: np.ndarray + k_abs: np.ndarray + + +def _build_stage_flex_cases( + pack: SharedPrefixPack, + plan: object, + config: ContextParallelConfig, + *, + device: torch.device, +) -> tuple[_StageFlexCase, ...]: + cases: list[_StageFlexCase] = [] + context = prepare_block_mask_context( + group_ids=pack.group_ids[0], + parent_ids=pack.parent_ids[0], + ) + for rank_plan in plan: + for stage in rank_plan.stage_plans: + if stage.mask_metadata is None: + continue + execution_spec = _stage_execution_spec(stage, config) + mask_metadata = execution_spec.mask_metadata or stage.mask_metadata + if mask_metadata is None: + continue + mask = build_block_mask_from_context( + FlexMaskSpec( + q_len=execution_spec.q_len, + k_len=execution_spec.k_len, + block_size=_sparse_block_size(config), + slices=stage.slices, + exact_mask=mask_metadata, + ), + context=context, + device=device, + validate=False, + ) + if mask is None: + continue + q_abs = ( + mask_metadata.q_token_indices.detach() + .to(device="cpu", dtype=torch.int64) + .reshape(-1) + .numpy() + ) + k_abs = ( + mask_metadata.k_token_indices.detach() + .to(device="cpu", dtype=torch.int64) + .reshape(-1) + .numpy() + ) + cases.append( + _StageFlexCase( + rank=int(rank_plan.rank), + stage_index=int(stage.stage_index), + q_len=int(execution_spec.q_len), + k_len=int(execution_spec.k_len), + logical_q_len=int(stage.q_len), + logical_k_len=int(stage.k_len), + block_mask=mask, + q_abs=q_abs, + k_abs=k_abs, + ) + ) + return tuple(cases) + + +def _stage_tensors( + cases: Sequence[_StageFlexCase], + *, + heads: int, + head_dim: int, + device: torch.device, +) -> tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]: + generator = torch.Generator(device=device).manual_seed(17) + tensors = [] + for case in cases: + q_shape = (1, int(heads), int(case.q_len), int(head_dim)) + k_shape = (1, int(heads), int(case.k_len), int(head_dim)) + tensors.append( + ( + torch.randn( + q_shape, device=device, dtype=torch.bfloat16, generator=generator + ), + torch.randn( + k_shape, device=device, dtype=torch.bfloat16, generator=generator + ), + torch.randn( + k_shape, device=device, dtype=torch.bfloat16, generator=generator + ), + ) + ) + return tuple(tensors) + + +def _stage_variant_block_mask( + case: _StageFlexCase, + variant: str, + *, + device: torch.device, +) -> BlockMask: + if variant == "current": + return case.block_mask + q_abs = torch.as_tensor(case.q_abs, device=device, dtype=torch.int64) + k_abs = torch.as_tensor(case.k_abs, device=device, dtype=torch.int64) + if variant == "causal_abs_only": + + def mask_mod(batch_idx, head_idx, query_idx, kv_idx): + del batch_idx, head_idx + return q_abs[query_idx] >= k_abs[kv_idx] + + return _replace_block_mask_mod(case.block_mask, mask_mod) + raise ValueError(f"unknown flex_mask_variant {variant!r}") + + +def _stage_flex_stats(cases: Sequence[_StageFlexCase]) -> dict[str, object]: + return { + "flex_stage_count": len(cases), + "flex_stage_q_tokens": sum(case.q_len for case in cases), + "flex_stage_k_tokens": sum(case.k_len for case in cases), + "flex_stage_logical_q_tokens": sum(case.logical_q_len for case in cases), + "flex_stage_logical_k_tokens": sum(case.logical_k_len for case in cases), + "flex_stage_max_q_tokens": max(case.q_len for case in cases), + "flex_stage_max_k_tokens": max(case.k_len for case in cases), + "flex_stage_max_logical_q_tokens": max(case.logical_q_len for case in cases), + "flex_stage_max_logical_k_tokens": max(case.logical_k_len for case in cases), + } + + +def _sparse_block_size(config: ContextParallelConfig) -> tuple[int, int]: + return normalize_sparse_block_size( + config.attention_sparse_block_size or config.block_size + ) + + +def _stage_execution_spec( + stage: StagePlan, + config: ContextParallelConfig, +) -> StageExecutionSpec: + return _build_stage_execution_spec( + stage_plan=stage, + block_size=_sparse_block_size(config), + ) + + +def _replace_block_mask_mod(block_mask: BlockMask, mask_mod: object) -> BlockMask: + return BlockMask( + seq_lengths=block_mask.seq_lengths, + kv_num_blocks=block_mask.kv_num_blocks, + kv_indices=block_mask.kv_indices, + full_kv_num_blocks=block_mask.full_kv_num_blocks, + full_kv_indices=block_mask.full_kv_indices, + q_num_blocks=block_mask.q_num_blocks, + q_indices=block_mask.q_indices, + full_q_num_blocks=block_mask.full_q_num_blocks, + full_q_indices=block_mask.full_q_indices, + BLOCK_SIZE=block_mask.BLOCK_SIZE, + mask_mod=mask_mod, + ) + + +def _bench_cpu( + fn: Callable[[], object], + *, + warmup: int, + repeat: int, + before_each: Callable[[], object] | None = None, +) -> tuple[object, float]: + result = None + for _ in range(warmup): + if before_each is not None: + before_each() + result = fn() + elapsed = [] + for _ in range(repeat): + if before_each is not None: + before_each() + start = time.perf_counter() + result = fn() + elapsed.append((time.perf_counter() - start) * 1000.0) + assert result is not None + return result, round(sum(elapsed) / len(elapsed), 3) + + +def _bench_cuda(fn: Callable[[], object], *, warmup: int, repeat: int) -> float: + torch.cuda.reset_peak_memory_stats() + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + stop = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(repeat): + fn() + stop.record() + torch.cuda.synchronize() + return round(float(start.elapsed_time(stop)) / repeat, 3) + + +def _plan_stats(plan: object) -> dict[str, int]: + stage_count = 0 + remote_stage_count = 0 + mask_stage_count = 0 + for rank_plan in plan: + for stage in rank_plan.stage_plans: + stage_count += 1 + remote_stage_count += int(not stage.is_local_stage) + mask_stage_count += int(stage.mask_metadata is not None) + return { + "rank_count": len(plan), + "stage_count": stage_count, + "remote_stage_count": remote_stage_count, + "mask_stage_count": mask_stage_count, + } + + +def _mask_stats(masks: Sequence[BlockMask]) -> dict[str, int]: + return { + "mask_count": len(masks), + "partial_kv_blocks": sum(_block_count(mask, "kv_num_blocks") for mask in masks), + "full_kv_blocks": sum( + _block_count(mask, "full_kv_num_blocks") for mask in masks + ), + "partial_q_blocks": sum(_block_count(mask, "q_num_blocks") for mask in masks), + "full_q_blocks": sum(_block_count(mask, "full_q_num_blocks") for mask in masks), + } + + +def _block_count(block_mask: BlockMask, name: str) -> int: + counts = getattr(block_mask, name) + return 0 if counts is None else int(counts.sum().item()) + + +def _assert_matches_torch_block_mask( + block_mask: BlockMask, + *, + slices: Sequence[object] = (), +) -> None: + q_len, k_len = block_mask.seq_lengths + reference = torch_block_mask( + _slice_mask_mod(block_mask.mask_mod, slices), + B=int(block_mask.kv_num_blocks.shape[0]), + H=1, + Q_LEN=q_len, + KV_LEN=k_len, + device="cpu", + BLOCK_SIZE=block_mask.BLOCK_SIZE, + ) + for counts_name, indices_name in ( + ("kv_num_blocks", "kv_indices"), + ("full_kv_num_blocks", "full_kv_indices"), + ("q_num_blocks", "q_indices"), + ("full_q_num_blocks", "full_q_indices"), + ): + actual = _block_entries(block_mask, counts_name, indices_name) + expected = _block_entries(reference, counts_name, indices_name) + if actual != expected: + raise AssertionError(f"{counts_name}/{indices_name} mismatch") + + +def _slice_mask_mod(mask_mod: object, slices: Sequence[object]) -> object: + if not slices: + return mask_mod + + def sliced_mask_mod( + batch_idx: torch.Tensor, + head_idx: torch.Tensor, + query_idx: torch.Tensor, + kv_idx: torch.Tensor, + ) -> torch.Tensor: + in_slice = (query_idx < 0) & (kv_idx < 0) + for slice_ in slices: + in_slice |= ( + (query_idx >= int(slice_.q_range.start)) + & (query_idx < int(slice_.q_range.end)) + & (kv_idx >= int(slice_.k_range.start)) + & (kv_idx < int(slice_.k_range.end)) + ) + return in_slice & mask_mod(batch_idx, head_idx, query_idx, kv_idx) + + return sliced_mask_mod + + +def _block_entries( + block_mask: BlockMask, + counts_name: str, + indices_name: str, +) -> set[tuple[int, int, int, int]]: + counts = getattr(block_mask, counts_name) + indices = getattr(block_mask, indices_name) + if counts is None or indices is None: + return set() + entries = set() + for batch_index in range(int(counts.shape[0])): + for head_index in range(int(counts.shape[1])): + for block_index in range(int(counts.shape[2])): + block_count = int(counts[batch_index, head_index, block_index]) + for other_block in indices[ + batch_index, + head_index, + block_index, + :block_count, + ].tolist(): + entries.add( + ( + batch_index, + head_index, + block_index, + int(other_block), + ) + ) + return entries + + +def _logical_tokens(pack: SharedPrefixPack) -> int: + return sum(int(positions.numel()) for positions in pack.positions_by_sequence) + + +def _torch_validation_skip_reason( + *, + validate_torch: bool, + packed_tokens: int, + token_cap: int, +) -> str | None: + if not validate_torch: + return "disabled" + if token_cap > 0 and packed_tokens > token_cap: + return f"packed_tokens>{token_cap}" + return None + + +def _csv_values(value: str) -> tuple[str, ...]: + values = tuple(part.strip() for part in value.split(",") if part.strip()) + if not values: + raise ValueError("CSV option must contain at least one value") + return values + + +def _write(path: Path, payload: dict[str, object]) -> None: + line = json.dumps(payload, sort_keys=True) + with path.open("a", encoding="utf-8") as output: + output.write(line + "\n") + print(line, flush=True) + + +def _check_threshold(name: str, value_ms: float, limit_ms: float | None) -> None: + if limit_ms is not None and float(value_ms) > float(limit_ms): + raise RuntimeError( + f"{name} took {float(value_ms):.3f}ms, exceeding {float(limit_ms):.3f}ms" + ) + + +if __name__ == "__main__": + typer.run(main) diff --git a/pyproject.toml b/pyproject.toml index 57b67da75..7fcb65b40 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -277,6 +277,7 @@ allowed-unresolved-imports = [ "peft.**", "pyarrow.**", "torch.**", + "torchvision.**", "torchao.**", "transformers.**", "trl.**", diff --git a/scripts/build-gpu-image.sh b/scripts/build-gpu-image.sh index 448a166fa..5017ef3e9 100755 --- a/scripts/build-gpu-image.sh +++ b/scripts/build-gpu-image.sh @@ -10,11 +10,13 @@ Options: --image-repo REPO Image repository to publish --infra INFRA Kubernetes-backed SkyPilot infra (default: k8s/cks-wb3) --no-cache Disable registry-backed BuildKit cache + --no-prewarm-modal Skip prebuilding the pushed image in Modal --no-prewarm-nodes Skip pre-pulling the pushed image on GPU nodes --prewarm-infra INFRA Kubernetes-backed infra to prewarm; repeatable --pull-image-repo REPO Image repository for cluster pulls/prewarm + --prewarm-modal Require prebuilding the pushed image in Modal --prewarm-timeout DUR Timeout for the prewarm DaemonSet rollout (default: 30m) - --tag TAG Image tag to publish + --tag TAG Image tag to publish (default: latest) --help Show this help EOF } @@ -25,12 +27,13 @@ cluster_name="" infra="${SKY_INFRA:-k8s/cks-wb3}" image_repo="${ART_IMAGE_REPO:-}" pull_image_repo="${ART_PULL_IMAGE_REPO:-}" -image_tag="" +image_tag="${IMAGE_TAG:-latest}" docker_config_path="${DOCKER_CONFIG_PATH:-${HOME}/.docker/config.json}" buildkit_image="${BUILDKIT_IMAGE:-moby/buildkit:v0.29.0-rootless}" buildkit_namespace="${KUBECTL_NAMESPACE:-default}" buildkit_wait_timeout="${BUILDKIT_WAIT_TIMEOUT:-300s}" no_cache="${NO_CACHE:-false}" +prewarm_modal="${PREWARM_MODAL:-auto}" prewarm_nodes="${PREWARM_NODES:-true}" prewarm_infras=() if [[ -n "${PREWARM_INFRAS:-}" ]]; then @@ -73,6 +76,10 @@ while [[ $# -gt 0 ]]; do no_cache=true shift ;; + --no-prewarm-modal) + prewarm_modal=false + shift + ;; --no-prewarm-nodes) prewarm_nodes=false shift @@ -85,6 +92,10 @@ while [[ $# -gt 0 ]]; do pull_image_repo="$2" shift 2 ;; + --prewarm-modal) + prewarm_modal=true + shift + ;; --prewarm-timeout) prewarm_timeout="$2" shift 2 @@ -105,6 +116,14 @@ while [[ $# -gt 0 ]]; do esac done +case "${prewarm_modal}" in + auto|true|false) ;; + *) + echo "PREWARM_MODAL must be one of: auto, true, false" >&2 + exit 1 + ;; +esac + kube_context_from_infra() { local infra_value="$1" case "${infra_value}" in @@ -142,10 +161,6 @@ art_sha="$(git -C "${repo_root}" rev-parse HEAD)" art_short_sha="$(git -C "${repo_root}" rev-parse --short=12 HEAD)" timestamp="$(date +%m%d-%H%M%S)" -if [[ -z "${image_tag}" ]]; then - image_tag="skypilot-${art_short_sha}" -fi - if [[ -z "${cluster_name}" ]]; then cluster_name="art-gpu-build-${timestamp}" fi @@ -461,6 +476,38 @@ if [[ -n "${prewarm_refresh_tag_image}" ]]; then esac fi +modal_auth_available=false +if [[ "${prewarm_modal}" != "false" ]]; then + if uv run --with 'modal>=1.5.0' python - <<'PY' >/dev/null 2>&1; then +import modal + +modal.Workspace.from_context().hydrate() +PY + modal_auth_available=true + fi +fi + +if [[ "${prewarm_modal}" == "true" || "${modal_auth_available}" == "true" ]]; then + echo "Prewarming ${image_repo}:${image_tag} in Modal image cache" + MODAL_FORCE_BUILD=1 uv run --with 'modal>=1.5.0' python - "${image_repo}:${image_tag}" <<'PY' +import sys + +import modal + +image = ( + modal.Image.from_registry(sys.argv[1], add_python="3.12") + .apt_install("openssh-server", "sudo", "rsync", "curl", "procps", "patch", "lsof") +) +app = modal.App.lookup("skypilot-modal", create_if_missing=True) +with modal.enable_output(): + image.build(app) +PY +elif [[ "${prewarm_modal}" == "auto" ]]; then + echo "Skipping Modal image prewarm: Modal auth unavailable" +else + echo "Skipping Modal image prewarm" +fi + dump_prewarm_diagnostics() { echo "::group::Prewarm diagnostics" "${kubectl_cmd[@]}" get daemonset -n "${prewarm_namespace}" "${prewarm_name}" -o wide || true diff --git a/src/art/megatron/compile_workarounds.py b/src/art/megatron/compile_workarounds.py index a6d9d916f..acfbec203 100644 --- a/src/art/megatron/compile_workarounds.py +++ b/src/art/megatron/compile_workarounds.py @@ -211,7 +211,6 @@ def _sync_dealloc_fake( _install_self_attn_linear_proj_reduce_scatter_workaround() if "weighted_bias_swiglu_no_inner_forward_cast" in flags: _install_weighted_bias_swiglu_no_inner_forward_cast_workaround() - deepep_flags = {"deepep_permute_restore", "deepep_dispatch_combine"} & flags if deepep_flags: deepep_manager = _require_attr(token_dispatcher, "_DeepepManager") diff --git a/src/art/megatron/context_parallel/__init__.py b/src/art/megatron/context_parallel/__init__.py index 995b0c425..618ea47c9 100644 --- a/src/art/megatron/context_parallel/__init__.py +++ b/src/art/megatron/context_parallel/__init__.py @@ -1,22 +1,24 @@ -from .builder import build_dense_reference_mask, build_shared_prefix_attention_spec -from .layout_index import TokenLayoutIndex -from .runtime import build_context_parallel_token_layout_index -from .types import ( - ArtContextParallelState, - AttnMaskKind, - AttnSlice, - ContextParallelConfig, - ContextParallelRuntimeKey, - ContextParallelRuntimePlan, - DispatchedPackedTensors, - FlexMaskSpec, - PackedBatchAttentionSpec, - PackedRowAttentionSpec, - ParallelTopology, - PreparedMegatronBatch, - SharedPrefixBuilderConfig, - TokenRange, -) +from __future__ import annotations + +from importlib import import_module +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .builder import build_dense_reference_mask, build_prefix_tree_attention_spec + from .layout_index import TokenLayoutIndex + from .types import ( + ArtContextParallelState, + AttnMaskKind, + AttnSlice, + ContextParallelConfig, + DispatchedPackedTensors, + FlexMaskSpec, + PackedBatchAttentionSpec, + PackedRowAttentionSpec, + ParallelTopology, + PreparedMegatronBatch, + TokenRange, + ) __all__ = [ "ArtContextParallelState", @@ -28,13 +30,34 @@ "PackedRowAttentionSpec", "ParallelTopology", "PreparedMegatronBatch", - "SharedPrefixBuilderConfig", "ContextParallelConfig", - "ContextParallelRuntimeKey", - "ContextParallelRuntimePlan", "TokenRange", "TokenLayoutIndex", "build_dense_reference_mask", - "build_context_parallel_token_layout_index", - "build_shared_prefix_attention_spec", + "build_prefix_tree_attention_spec", ] + +_EXPORT_MODULES = { + "TokenLayoutIndex": ".layout_index", + "build_dense_reference_mask": ".builder", + "build_prefix_tree_attention_spec": ".builder", + "ArtContextParallelState": ".types", + "AttnMaskKind": ".types", + "AttnSlice": ".types", + "DispatchedPackedTensors": ".types", + "FlexMaskSpec": ".types", + "PackedBatchAttentionSpec": ".types", + "PackedRowAttentionSpec": ".types", + "ParallelTopology": ".types", + "PreparedMegatronBatch": ".types", + "ContextParallelConfig": ".types", + "TokenRange": ".types", +} + + +def __getattr__(name: str) -> Any: + if name not in _EXPORT_MODULES: + raise AttributeError(name) + value = getattr(import_module(_EXPORT_MODULES[name], __name__), name) + globals()[name] = value + return value diff --git a/src/art/megatron/context_parallel/block_mask.py b/src/art/megatron/context_parallel/block_mask.py index f86f63320..f5255ec35 100644 --- a/src/art/megatron/context_parallel/block_mask.py +++ b/src/art/megatron/context_parallel/block_mask.py @@ -1,61 +1,71 @@ from __future__ import annotations -from typing import cast +from dataclasses import dataclass import numpy as np import torch from torch.nn.attention.flex_attention import BlockMask from art.megatron.flex_attn.compiled import normalize_sparse_block_size +from art.megatron.prefix_tree import parse_prefix_tree_row from .types import AttnMaskKind, FlexMaskSpec -_INVALID_Q_GROUP = -(1 << 63) -_INVALID_Q_PARENT = _INVALID_Q_GROUP + 1 -_INVALID_K_GROUP = _INVALID_Q_GROUP + 2 -_INVALID_POSITION = _INVALID_Q_GROUP + 3 +_INVALID_ABS = -(1 << 63) +_INVALID_ENTER = -1 +_INVALID_EXIT = -1 +_INVALID_POS = -(1 << 63) -def _build_exact_mask_mod( +@dataclass(frozen=True, slots=True) +class PreparedBlockMaskContext: + source_len: int + group_enter_np: np.ndarray + group_exit_np: np.ndarray + input_pos_np: np.ndarray | None = None + + +def _build_interval_mask_mod( *, q_abs: np.ndarray, k_abs: np.ndarray, - q_group: np.ndarray, - q_parent: np.ndarray, - k_group: np.ndarray, + q_enter: np.ndarray, + k_enter: np.ndarray, + k_exit: np.ndarray, q_pos: np.ndarray | None, k_pos: np.ndarray | None, sliding_window: int | None, device: torch.device, ): - q_group_tensor = torch.as_tensor(q_group, device=device, dtype=torch.int64) - q_parent_tensor = torch.as_tensor(q_parent, device=device, dtype=torch.int64) - k_group_tensor = torch.as_tensor(k_group, device=device, dtype=torch.int64) - - if sliding_window is not None: - q_pos_tensor = torch.as_tensor(q_pos, device=device, dtype=torch.int64) - k_pos_tensor = torch.as_tensor(k_pos, device=device, dtype=torch.int64) - - def sliding_mask_mod( - batch_idx: torch.Tensor, - head_idx: torch.Tensor, - query_idx: torch.Tensor, - kv_idx: torch.Tensor, - ) -> torch.Tensor: - del batch_idx, head_idx - same_group = q_group_tensor[query_idx] == k_group_tensor[kv_idx] - parent_prefix = q_parent_tensor[query_idx] == k_group_tensor[kv_idx] - delta = q_pos_tensor[query_idx] - k_pos_tensor[kv_idx] - return ( - (same_group | parent_prefix) - & (delta >= 0) - & (delta < int(sliding_window)) - ) - - return sliding_mask_mod - - q_abs_tensor = torch.as_tensor(q_abs, device=device, dtype=torch.int64) - k_abs_tensor = torch.as_tensor(k_abs, device=device, dtype=torch.int64) + metadata_dtype = _smallest_mask_metadata_dtype( + q_abs, + k_abs, + q_enter, + k_enter, + k_exit, + *((q_pos, k_pos) if sliding_window is not None else ()), + ) + q_abs_tensor = torch.as_tensor(q_abs, device=device, dtype=metadata_dtype) + k_abs_tensor = torch.as_tensor(k_abs, device=device, dtype=metadata_dtype) + q_enter_tensor = torch.as_tensor(q_enter, device=device, dtype=metadata_dtype) + k_enter_tensor = torch.as_tensor(k_enter, device=device, dtype=metadata_dtype) + k_exit_tensor = torch.as_tensor(k_exit, device=device, dtype=metadata_dtype) + q_pos_tensor_or_none = ( + None + if q_pos is None + else torch.as_tensor(q_pos, device=device, dtype=metadata_dtype) + ) + k_pos_tensor_or_none = ( + None + if k_pos is None + else torch.as_tensor(k_pos, device=device, dtype=metadata_dtype) + ) + if sliding_window is not None and ( + q_pos_tensor_or_none is None or k_pos_tensor_or_none is None + ): + raise RuntimeError("Sliding-window prefix-tree masks require input_pos.") + q_pos_tensor = q_pos_tensor_or_none + k_pos_tensor = k_pos_tensor_or_none def mask_mod( batch_idx: torch.Tensor, @@ -66,22 +76,46 @@ def mask_mod( del batch_idx, head_idx q_abs_local = q_abs_tensor[query_idx] k_abs_local = k_abs_tensor[kv_idx] - same_group = q_group_tensor[query_idx] == k_group_tensor[kv_idx] - parent_prefix = q_parent_tensor[query_idx] == k_group_tensor[kv_idx] - return (q_abs_local >= k_abs_local) & (same_group | parent_prefix) + q_enter_local = q_enter_tensor[query_idx] + k_enter_local = k_enter_tensor[kv_idx] + k_exit_local = k_exit_tensor[kv_idx] + in_key_subtree = (k_enter_local <= q_enter_local) & ( + q_enter_local < k_exit_local + ) + allowed = (q_abs_local >= k_abs_local) & in_key_subtree + if sliding_window is None: + return allowed + assert q_pos_tensor is not None and k_pos_tensor is not None + delta = q_pos_tensor[query_idx] - k_pos_tensor[kv_idx] + return allowed & (delta >= 0) & (delta < int(sliding_window)) return mask_mod +def _smallest_mask_metadata_dtype(*arrays: np.ndarray | None) -> torch.dtype: + int32 = np.iinfo(np.int32) + for array in arrays: + if array is None or int(array.size) == 0: + continue + if int(array.min()) < int32.min or int(array.max()) > int32.max: + return torch.int64 + return torch.int32 + + def _dense_blocks_to_ordered( blocks: np.ndarray, *, device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor]: - counts = torch.from_numpy(blocks.sum(axis=-1).astype(np.int32)) - indices = torch.from_numpy( - np.argsort(-blocks.astype(np.int32), axis=-1, kind="stable").astype(np.int32) - ) + row_indices, column_indices = np.nonzero(blocks) + counts_np = np.bincount(row_indices, minlength=blocks.shape[0]).astype(np.int32) + indices_np = np.zeros(blocks.shape, dtype=np.int32) + if int(row_indices.size) > 0: + starts = np.concatenate(([0], np.cumsum(counts_np[:-1], dtype=np.int64))) + offsets = np.arange(int(row_indices.size), dtype=np.int64) - starts[row_indices] + indices_np[row_indices, offsets] = column_indices + counts = torch.from_numpy(counts_np) + indices = torch.from_numpy(indices_np) return ( counts.view(1, 1, -1).to(device=device), indices.view(1, 1, blocks.shape[0], blocks.shape[1]).to(device=device), @@ -101,142 +135,361 @@ def _select_with_invalid_np( return selected -def _build_q_block_group_state( +def _refine_interval_blocks( *, + partial_blocks: np.ndarray, + full_blocks: np.ndarray, q_abs: np.ndarray, - q_group: np.ndarray, - q_parent: np.ndarray, + k_abs: np.ndarray, + q_enter: np.ndarray, + k_enter: np.ndarray, + k_exit: np.ndarray, q_block: int, - q_blocks: int, -) -> tuple[np.ndarray, list[dict[int, int]], list[frozenset[int]]]: - q_min_by_block = np.empty((q_blocks,), dtype=np.int64) - q_allowed_max_by_group: list[dict[int, int]] = [] - q_all_allowed_groups: list[frozenset[int]] = [] - for block_idx in range(q_blocks): - start = block_idx * q_block - end = min((block_idx + 1) * q_block, int(q_abs.size)) - q = q_abs[start:end] - q_group_block = q_group[start:end] - q_parent_block = q_parent[start:end] - q_min_by_block[block_idx] = int(q.min()) if int(q.size) else 0 - max_by_group: dict[int, int] = {} - all_groups: list[int] = [] - for group_value in np.unique(np.concatenate((q_group_block, q_parent_block))): - allowed = (q_group_block == group_value) | (q_parent_block == group_value) - if bool(allowed.any()): - max_by_group[int(group_value)] = int(q[allowed].max()) - if bool(allowed.all()): - all_groups.append(int(group_value)) - q_allowed_max_by_group.append(max_by_group) - q_all_allowed_groups.append(frozenset(all_groups)) - return q_min_by_block, q_allowed_max_by_group, q_all_allowed_groups - - -def _build_k_block_group_state( + k_block: int, +) -> None: + if not bool((partial_blocks | full_blocks).any()): + return + + q_abs_blocks = _block_matrix( + q_abs, + block_size=q_block, + block_count=int(partial_blocks.shape[0]), + fill_value=_INVALID_ABS, + ) + q_enter_blocks = _block_matrix( + q_enter, + block_size=q_block, + block_count=int(partial_blocks.shape[0]), + fill_value=_INVALID_ENTER, + ) + k_abs_blocks = _block_matrix( + k_abs, + block_size=k_block, + block_count=int(partial_blocks.shape[1]), + fill_value=_INVALID_ABS, + ) + k_enter_blocks = _block_matrix( + k_enter, + block_size=k_block, + block_count=int(partial_blocks.shape[1]), + fill_value=_INVALID_ENTER, + ) + k_exit_blocks = _block_matrix( + k_exit, + block_size=k_block, + block_count=int(partial_blocks.shape[1]), + fill_value=_INVALID_EXIT, + ) + + q_valid = (q_abs_blocks >= 0) & (q_enter_blocks >= 0) + k_valid = ( + (k_abs_blocks >= 0) & (k_enter_blocks >= 0) & (k_exit_blocks > k_enter_blocks) + ) + q_all_valid = q_valid.all(axis=1) + k_all_valid = k_valid.all(axis=1) + q_min_abs = np.where(q_valid, q_abs_blocks, np.iinfo(np.int64).max).min(axis=1) + q_min_enter = np.where( + q_valid, + q_enter_blocks, + np.iinfo(np.int64).max, + ).min(axis=1) + q_max_enter = np.where(q_valid, q_enter_blocks, _INVALID_ENTER).max(axis=1) + k_max_abs = np.where(k_valid, k_abs_blocks, _INVALID_ABS).max(axis=1) + k_max_enter = np.where(k_valid, k_enter_blocks, _INVALID_ENTER).max(axis=1) + k_min_exit = np.where(k_valid, k_exit_blocks, np.iinfo(np.int64).max).min(axis=1) + safe_full = ( + q_all_valid[:, None] + & k_all_valid[None, :] + & (q_min_abs[:, None] >= k_max_abs[None, :]) + & (k_max_enter[None, :] <= q_min_enter[:, None]) + & (q_max_enter[:, None] < k_min_exit[None, :]) + ) + candidate_blocks = partial_blocks | (full_blocks & ~safe_full) + q_indices, k_indices = np.nonzero(candidate_blocks) + if int(q_indices.size) == 0: + return + + rows = np.arange(int(k_valid.shape[0])) + first_valid_offsets = k_valid.argmax(axis=1) + first_enter = k_enter_blocks[rows, first_valid_offsets] + first_exit = k_exit_blocks[rows, first_valid_offsets] + k_single_interval = k_valid.any(axis=1) & ( + (~k_valid) + | ( + (k_enter_blocks == first_enter[:, None]) + & (k_exit_blocks == first_exit[:, None]) + ) + ).all(axis=1) + + single_pair = k_single_interval[k_indices] + if bool(single_pair.any()): + single_q = q_indices[single_pair] + single_k = k_indices[single_pair] + q_abs_selected = q_abs_blocks[single_q] + q_enter_selected = q_enter_blocks[single_q] + in_subtree = ( + q_valid[single_q] + & (q_enter_selected >= first_enter[single_k, None]) + & (q_enter_selected < first_exit[single_k, None]) + ) + max_abs_in_subtree = np.where( + in_subtree, + q_abs_selected, + _INVALID_ABS, + ).max(axis=1) + k_min_abs = np.where(k_valid, k_abs_blocks, np.iinfo(np.int64).max).min(axis=1) + has_any = max_abs_in_subtree >= k_min_abs[single_k] + + is_full = ( + has_any + & q_all_valid[single_q] + & k_all_valid[single_k] + & (q_min_abs[single_q] >= k_max_abs[single_k]) + & (first_enter[single_k] <= q_min_enter[single_q]) + & (q_max_enter[single_q] < first_exit[single_k]) + ) + partial_blocks[single_q, single_k] = has_any & ~is_full + full_blocks[single_q, single_k] = is_full + + intervals_by_k: dict[int, tuple[tuple[int, int, int], ...]] = {} + + def k_intervals(k_idx: int) -> tuple[tuple[int, int, int], ...]: + cached = intervals_by_k.get(k_idx) + if cached is not None: + return cached + min_abs_by_interval: dict[tuple[int, int], int] = {} + for abs_value, enter_value, exit_value in zip( + k_abs_blocks[k_idx, k_valid[k_idx]], + k_enter_blocks[k_idx, k_valid[k_idx]], + k_exit_blocks[k_idx, k_valid[k_idx]], + strict=True, + ): + key = (int(enter_value), int(exit_value)) + prior = min_abs_by_interval.get(key) + min_abs_by_interval[key] = ( + int(abs_value) if prior is None else min(prior, int(abs_value)) + ) + cached = tuple( + (enter, exit, min_abs) + for (enter, exit), min_abs in min_abs_by_interval.items() + ) + intervals_by_k[k_idx] = cached + return cached + + for q_idx, k_idx in zip( + q_indices[~single_pair], + k_indices[~single_pair], + strict=True, + ): + q_valid_row = q_valid[q_idx] + intervals = k_intervals(int(k_idx)) + has_any = False + if bool(q_valid_row.any()) and intervals: + q_abs_row = q_abs_blocks[q_idx] + q_enter_row = q_enter_blocks[q_idx] + for enter, exit, min_abs in intervals: + in_subtree = q_valid_row & (q_enter_row >= enter) & (q_enter_row < exit) + if bool(in_subtree.any()) and int(q_abs_row[in_subtree].max()) >= int( + min_abs + ): + has_any = True + break + is_full = ( + has_any + and bool(q_all_valid[q_idx]) + and bool(k_all_valid[k_idx]) + and int(q_min_abs[q_idx]) >= int(k_max_abs[k_idx]) + and int(k_max_enter[k_idx]) <= int(q_min_enter[q_idx]) + and int(q_max_enter[q_idx]) < int(k_min_exit[k_idx]) + ) + partial_blocks[q_idx, k_idx] = has_any and not is_full + full_blocks[q_idx, k_idx] = bool(is_full) + + +def _refine_sliding_interval_blocks( *, + partial_blocks: np.ndarray, + full_blocks: np.ndarray, + q_abs: np.ndarray, k_abs: np.ndarray, - k_group: np.ndarray, + q_enter: np.ndarray, + k_enter: np.ndarray, + k_exit: np.ndarray, + q_pos: np.ndarray, + k_pos: np.ndarray, + q_block: int, k_block: int, - k_blocks: int, -) -> tuple[np.ndarray, list[dict[int, int]], list[tuple[int, ...]]]: - k_max_by_block = np.empty((k_blocks,), dtype=np.int64) - k_min_by_group: list[dict[int, int]] = [] - k_groups_by_block: list[tuple[int, ...]] = [] - for block_idx in range(k_blocks): - start = block_idx * k_block - end = min((block_idx + 1) * k_block, int(k_abs.size)) - k = k_abs[start:end] - k_group_block = k_group[start:end] - k_max_by_block[block_idx] = int(k.max()) if int(k.size) else 0 - min_by_group: dict[int, int] = {} - for group_value in np.unique(k_group_block): - min_by_group[int(group_value)] = int(k[k_group_block == group_value].min()) - k_min_by_group.append(min_by_group) - k_groups_by_block.append(tuple(min_by_group)) - return k_max_by_block, k_min_by_group, k_groups_by_block - - -def _exact_block_state( - *, - q_idx: int, - k_idx: int, - q_min_by_block: np.ndarray, - q_allowed_max_by_group: list[dict[int, int]], - q_all_allowed_groups: list[frozenset[int]], - k_max_by_block: np.ndarray, - k_min_by_group: list[dict[int, int]], - k_groups_by_block: list[tuple[int, ...]], -) -> tuple[bool, bool]: - q_allowed_max = q_allowed_max_by_group[q_idx] - k_min = k_min_by_group[k_idx] - if not any( - q_allowed_max.get(k_group_value, _INVALID_Q_GROUP) >= min_k - for k_group_value, min_k in k_min.items() - ): - return False, False - if int(q_min_by_block[q_idx]) < int(k_max_by_block[k_idx]): - return True, False - q_all_allowed = q_all_allowed_groups[q_idx] - return True, all( - k_group_value in q_all_allowed for k_group_value in k_groups_by_block[k_idx] + sliding_window: int, +) -> None: + candidates = partial_blocks | full_blocks + if not bool(candidates.any()): + return + + q_abs_blocks = _block_matrix( + q_abs, + block_size=q_block, + block_count=int(partial_blocks.shape[0]), + fill_value=_INVALID_ABS, + ) + q_enter_blocks = _block_matrix( + q_enter, + block_size=q_block, + block_count=int(partial_blocks.shape[0]), + fill_value=_INVALID_ENTER, + ) + q_pos_blocks = _block_matrix( + q_pos, + block_size=q_block, + block_count=int(partial_blocks.shape[0]), + fill_value=_INVALID_POS, + ) + k_abs_blocks = _block_matrix( + k_abs, + block_size=k_block, + block_count=int(partial_blocks.shape[1]), + fill_value=_INVALID_ABS, + ) + k_enter_blocks = _block_matrix( + k_enter, + block_size=k_block, + block_count=int(partial_blocks.shape[1]), + fill_value=_INVALID_ENTER, + ) + k_exit_blocks = _block_matrix( + k_exit, + block_size=k_block, + block_count=int(partial_blocks.shape[1]), + fill_value=_INVALID_EXIT, + ) + k_pos_blocks = _block_matrix( + k_pos, + block_size=k_block, + block_count=int(partial_blocks.shape[1]), + fill_value=_INVALID_POS, ) + q_valid = ( + (q_abs_blocks >= 0) & (q_enter_blocks >= 0) & (q_pos_blocks != _INVALID_POS) + ) + k_valid = ( + (k_abs_blocks >= 0) + & (k_enter_blocks >= 0) + & (k_exit_blocks > k_enter_blocks) + & (k_pos_blocks != _INVALID_POS) + ) -def _range_min_max( + q_indices, k_indices = np.nonzero(candidates) + partial_blocks[q_indices, k_indices] = False + full_blocks[q_indices, k_indices] = False + for q_idx, k_idx in zip(q_indices, k_indices, strict=True): + q_valid_row = q_valid[q_idx] + k_valid_row = k_valid[k_idx] + if not bool(q_valid_row.any()) or not bool(k_valid_row.any()): + continue + + q_abs_row = q_abs_blocks[q_idx][:, None] + q_enter_row = q_enter_blocks[q_idx][:, None] + q_pos_row = q_pos_blocks[q_idx][:, None] + k_abs_row = k_abs_blocks[k_idx][None, :] + k_enter_row = k_enter_blocks[k_idx][None, :] + k_exit_row = k_exit_blocks[k_idx][None, :] + k_pos_row = k_pos_blocks[k_idx][None, :] + delta = q_pos_row - k_pos_row + allowed = ( + q_valid_row[:, None] + & k_valid_row[None, :] + & (q_abs_row >= k_abs_row) + & (k_enter_row <= q_enter_row) + & (q_enter_row < k_exit_row) + & (delta >= 0) + & (delta < int(sliding_window)) + ) + if not bool(allowed.any()): + continue + if bool(q_valid_row.all()) and bool(k_valid_row.all()) and bool(allowed.all()): + full_blocks[q_idx, k_idx] = True + else: + partial_blocks[q_idx, k_idx] = True + + +def _is_strictly_increasing(values: np.ndarray) -> bool: + return int(values.size) <= 1 or bool(np.all(values[1:] > values[:-1])) + + +def _block_min_max( values: np.ndarray, starts: np.ndarray, ends: np.ndarray, ) -> tuple[np.ndarray, np.ndarray]: - mins = np.empty((int(starts.size),), dtype=np.int64) - maxes = np.empty((int(starts.size),), dtype=np.int64) + mins = np.empty(starts.shape, dtype=values.dtype) + maxes = np.empty(starts.shape, dtype=values.dtype) for index, (start, end) in enumerate(zip(starts, ends, strict=True)): block = values[int(start) : int(end)] - mins[index] = int(block.min()) if int(block.size) else 0 - maxes[index] = int(block.max()) if int(block.size) else 0 + mins[index] = block.min() + maxes[index] = block.max() return mins, maxes -def _exact_sliding_block_state( +def _block_matrix( + values: np.ndarray, *, - q_idx: int, - k_idx: int, - q_block: int, - k_block: int, - q_len: int, - k_len: int, - q_group: np.ndarray, - q_parent: np.ndarray, - k_group: np.ndarray, - q_pos: np.ndarray, - k_pos: np.ndarray, - sliding_window: int, -) -> tuple[bool, bool]: - q_start = int(q_idx) * int(q_block) - q_end = min(q_start + int(q_block), int(q_len)) - k_start = int(k_idx) * int(k_block) - k_end = min(k_start + int(k_block), int(k_len)) - q_group_block = q_group[q_start:q_end] - q_parent_block = q_parent[q_start:q_end] - k_group_block = k_group[k_start:k_end] - delta = q_pos[q_start:q_end, None] - k_pos[None, k_start:k_end] - allowed = ( - ( - (q_group_block[:, None] == k_group_block[None, :]) - | (q_parent_block[:, None] == k_group_block[None, :]) - ) - & (delta >= 0) - & (delta < int(sliding_window)) - ) - return bool(allowed.any()), bool(allowed.all()) + block_size: int, + block_count: int, + fill_value: int, +) -> np.ndarray: + padded = np.full(block_count * block_size, fill_value, dtype=np.int64) + padded[: int(values.size)] = values + return padded.reshape(block_count, block_size) + + +def _build_group_interval_arrays( + *, + row_tree, + length: int, +) -> tuple[np.ndarray, np.ndarray]: + enter_by_group: dict[int, int] = {} + exit_by_group: dict[int, int] = {} + segment_by_group = {segment.group_id: segment for segment in row_tree.segments} + children_by_group: dict[int, list[int]] = {} + roots: list[int] = [] + for segment in row_tree.segments: + if segment.ancestors: + children_by_group.setdefault(segment.parent_id, []).append(segment.group_id) + else: + roots.append(segment.group_id) + + next_enter = 0 + + def visit(group_id: int) -> None: + nonlocal next_enter + enter_by_group[group_id] = next_enter + next_enter += 1 + children = children_by_group.get(group_id, []) + children.sort(key=lambda child: segment_by_group[child].start) + for child_group_id in children: + visit(child_group_id) + exit_by_group[group_id] = next_enter + + roots.sort(key=lambda root: segment_by_group[root].start) + for root_group_id in roots: + visit(root_group_id) + + enter_by_token = np.full((length,), _INVALID_ENTER, dtype=np.int64) + exit_by_token = np.full((length,), _INVALID_EXIT, dtype=np.int64) + for segment in row_tree.segments: + enter_by_token[segment.start : segment.end] = enter_by_group[segment.group_id] + exit_by_token[segment.start : segment.end] = exit_by_group[segment.group_id] + if int(row_tree.valid_tokens) < int(length): + enter_by_token[int(row_tree.valid_tokens) :] = next_enter + exit_by_token[int(row_tree.valid_tokens) :] = next_enter + 1 + return enter_by_token, exit_by_token def _build_sparse_block_mask( spec: FlexMaskSpec, *, device: torch.device, - group_ids: torch.Tensor, - parent_ids: torch.Tensor, - input_pos: torch.Tensor | None, + context: PreparedBlockMaskContext, sliding_window: int | None, block_size: tuple[int, int], ) -> BlockMask: @@ -256,80 +509,52 @@ def _build_sparse_block_mask( ) q_abs = q_abs_tensor.numpy() k_abs = k_abs_tensor.numpy() - flat_group_ids = group_ids.detach().to(device="cpu", dtype=torch.int64).reshape(-1) - flat_parent_ids = ( - parent_ids.detach().to(device="cpu", dtype=torch.int64).reshape(-1) - ) - flat_group_ids_np = flat_group_ids.numpy() - flat_parent_ids_np = flat_parent_ids.numpy() - input_pos_for_window = cast(torch.Tensor, input_pos) - flat_input_pos_np = ( - input_pos_for_window.detach() - .to(device="cpu", dtype=torch.int64) - .reshape(-1) - .numpy() - if sliding_window is not None - else None - ) - q_group = _select_with_invalid_np( - flat_group_ids_np, + q_abs_sorted = _is_strictly_increasing(q_abs[q_abs >= 0]) + k_abs_sorted = _is_strictly_increasing(k_abs[k_abs >= 0]) + q_enter = _select_with_invalid_np( + context.group_enter_np, q_abs, - invalid_value=_INVALID_Q_GROUP, + invalid_value=_INVALID_ENTER, ) - q_parent = _select_with_invalid_np( - flat_parent_ids_np, - q_abs, - invalid_value=_INVALID_Q_PARENT, + k_enter = _select_with_invalid_np( + context.group_enter_np, + k_abs, + invalid_value=_INVALID_ENTER, ) - k_group = _select_with_invalid_np( - flat_group_ids_np, + k_exit = _select_with_invalid_np( + context.group_exit_np, k_abs, - invalid_value=_INVALID_K_GROUP, + invalid_value=_INVALID_EXIT, ) q_pos = ( _select_with_invalid_np( - cast(np.ndarray, flat_input_pos_np), + _require_input_pos(context), q_abs, - invalid_value=_INVALID_POSITION, + invalid_value=_INVALID_POS, ) if sliding_window is not None else None ) k_pos = ( _select_with_invalid_np( - cast(np.ndarray, flat_input_pos_np), + _require_input_pos(context), k_abs, - invalid_value=_INVALID_POSITION, + invalid_value=_INVALID_POS, ) if sliding_window is not None else None ) - mask_mod = _build_exact_mask_mod( + mask_mod = _build_interval_mask_mod( q_abs=q_abs, k_abs=k_abs, - q_group=q_group, - q_parent=q_parent, - k_group=k_group, + q_enter=q_enter, + k_enter=k_enter, + k_exit=k_exit, q_pos=q_pos, k_pos=k_pos, sliding_window=sliding_window, device=device, ) - q_min_by_block, q_allowed_max_by_group, q_all_allowed_groups = ( - _build_q_block_group_state( - q_abs=q_abs, - q_group=q_group, - q_parent=q_parent, - q_block=q_block, - q_blocks=q_blocks, - ) - ) - k_max_by_block, k_min_by_group, k_groups_by_block = _build_k_block_group_state( - k_abs=k_abs, - k_group=k_group, - k_block=k_block, - k_blocks=k_blocks, - ) if not spec.slices: raise RuntimeError( "Cannot build a CP attention block mask without stage slices" @@ -353,15 +578,11 @@ def _build_sparse_block_mask( if int(q_block_indices.size) == 0 or int(k_block_indices.size) == 0: continue q_block_start = q_block_indices * q_block - q_block_end = np.minimum( - (q_block_indices + 1) * q_block, - int(spec.q_len), - ) + q_block_end_raw = (q_block_indices + 1) * q_block + q_block_end = np.minimum(q_block_end_raw, int(spec.q_len)) k_block_start = k_block_indices * k_block - k_block_end = np.minimum( - (k_block_indices + 1) * k_block, - int(spec.k_len), - ) + k_block_end_raw = (k_block_indices + 1) * k_block + k_block_end = np.minimum(k_block_end_raw, int(spec.k_len)) q_overlap_start = np.maximum( q_block_start, q_start, @@ -378,20 +599,25 @@ def _build_sparse_block_mask( k_block_end, k_end, ) - q_min = q_abs[q_overlap_start] - q_max = q_abs[q_overlap_end - 1] - k_min = k_abs[k_overlap_start] - k_max = k_abs[k_overlap_end - 1] - if sliding_window is not None: - q_pos_for_window = cast(np.ndarray, q_pos) - k_pos_for_window = cast(np.ndarray, k_pos) - q_pos_min, q_pos_max = _range_min_max( - q_pos_for_window, + q_is_full = (q_overlap_start == q_block_start) & ( + q_overlap_end == q_block_end_raw + ) + k_is_full = (k_overlap_start == k_block_start) & ( + k_overlap_end == k_block_end_raw + ) + covers_block = q_is_full[:, None] & k_is_full[None, :] + if sliding_window is None: + window_has_any = None + window_is_full = None + else: + assert q_pos is not None and k_pos is not None + q_pos_min, q_pos_max = _block_min_max( + q_pos, q_overlap_start, q_overlap_end, ) - k_pos_min, k_pos_max = _range_min_max( - k_pos_for_window, + k_pos_min, k_pos_max = _block_min_max( + k_pos, k_overlap_start, k_overlap_end, ) @@ -401,20 +627,28 @@ def _build_sparse_block_mask( window_is_full = (q_pos_min[:, None] >= k_pos_max[None, :]) & ( q_pos_max[:, None] - k_pos_min[None, :] < int(sliding_window) ) - q_is_full = (q_overlap_start == q_block_start) & (q_overlap_end == q_block_end) - k_is_full = (k_overlap_start == k_block_start) & (k_overlap_end == k_block_end) - covers_block = q_is_full[:, None] & k_is_full[None, :] - if sliding_window is not None: - has_any = window_has_any - is_full = covers_block & window_is_full - elif slice_.mask_kind == AttnMaskKind.FULL: + if slice_.mask_kind == AttnMaskKind.FULL: has_any = np.ones( (int(q_block_indices.size), int(k_block_indices.size)), dtype=bool ) is_full = covers_block else: + q_min, q_max = ( + (q_abs[q_overlap_start], q_abs[q_overlap_end - 1]) + if q_abs_sorted + else _block_min_max(q_abs, q_overlap_start, q_overlap_end) + ) + k_min, k_max = ( + (k_abs[k_overlap_start], k_abs[k_overlap_end - 1]) + if k_abs_sorted + else _block_min_max(k_abs, k_overlap_start, k_overlap_end) + ) has_any = q_max[:, None] >= k_min[None, :] is_full = covers_block & (q_min[:, None] >= k_max[None, :]) + if sliding_window is not None: + assert window_has_any is not None and window_is_full is not None + has_any &= window_has_any + is_full &= window_is_full q_slice = slice(int(q_block_indices[0]), int(q_block_indices[-1]) + 1) k_slice = slice(int(k_block_indices[0]), int(k_block_indices[-1]) + 1) @@ -422,42 +656,40 @@ def _build_sparse_block_mask( partial_blocks[q_slice, k_slice] |= has_any full_blocks[q_slice, k_slice] |= is_full - ambiguous = (touch_counts > 1) & partial_blocks & ~full_blocks - for q_idx, k_idx in np.argwhere(ambiguous): - if sliding_window is None: - has_any, is_full = _exact_block_state( - q_idx=int(q_idx), - k_idx=int(k_idx), - q_min_by_block=q_min_by_block, - q_allowed_max_by_group=q_allowed_max_by_group, - q_all_allowed_groups=q_all_allowed_groups, - k_max_by_block=k_max_by_block, - k_min_by_group=k_min_by_group, - k_groups_by_block=k_groups_by_block, - ) - else: - has_any, is_full = _exact_sliding_block_state( - q_idx=int(q_idx), - k_idx=int(k_idx), - q_block=q_block, - k_block=k_block, - q_len=int(spec.q_len), - k_len=int(spec.k_len), - q_group=q_group, - q_parent=q_parent, - k_group=k_group, - q_pos=cast(np.ndarray, q_pos), - k_pos=cast(np.ndarray, k_pos), - sliding_window=int(sliding_window), - ) - partial_blocks[q_idx, k_idx] = False - full_blocks[q_idx, k_idx] = False - if is_full: - full_blocks[q_idx, k_idx] = True - elif has_any: - partial_blocks[q_idx, k_idx] = True - partial_blocks &= ~full_blocks + needs_refine = full_blocks | ((touch_counts > 1) & partial_blocks) + if bool(needs_refine.any()): + refined_partial = partial_blocks & needs_refine + refined_full = full_blocks & needs_refine + _refine_interval_blocks( + partial_blocks=refined_partial, + full_blocks=refined_full, + q_abs=q_abs, + k_abs=k_abs, + q_enter=q_enter, + k_enter=k_enter, + k_exit=k_exit, + q_block=q_block, + k_block=k_block, + ) + partial_blocks = (partial_blocks & ~needs_refine) | refined_partial + full_blocks = (full_blocks & ~needs_refine) | refined_full + if sliding_window is not None: + assert q_pos is not None and k_pos is not None + _refine_sliding_interval_blocks( + partial_blocks=partial_blocks, + full_blocks=full_blocks, + q_abs=q_abs, + k_abs=k_abs, + q_enter=q_enter, + k_enter=k_enter, + k_exit=k_exit, + q_pos=q_pos, + k_pos=k_pos, + q_block=q_block, + k_block=k_block, + sliding_window=int(sliding_window), + ) kv_num_blocks, kv_indices = _dense_blocks_to_ordered( partial_blocks, device=device, @@ -466,25 +698,154 @@ def _build_sparse_block_mask( full_blocks, device=device, ) - return BlockMask.from_kv_blocks( - kv_num_blocks, - kv_indices, - full_kv_num_blocks, - full_kv_indices, + q_num_blocks, q_indices = _dense_blocks_to_ordered( + partial_blocks.T, + device=device, + ) + full_q_num_blocks, full_q_indices = _dense_blocks_to_ordered( + full_blocks.T, + device=device, + ) + return BlockMask( + seq_lengths=(int(spec.q_len), int(spec.k_len)), + kv_num_blocks=kv_num_blocks, + kv_indices=kv_indices, + full_kv_num_blocks=full_kv_num_blocks, + full_kv_indices=full_kv_indices, + q_num_blocks=q_num_blocks, + q_indices=q_indices, + full_q_num_blocks=full_q_num_blocks, + full_q_indices=full_q_indices, BLOCK_SIZE=block_size, mask_mod=mask_mod, - seq_lengths=(int(spec.q_len), int(spec.k_len)), ) -def build_block_mask( - spec: FlexMaskSpec, +def prepare_block_mask_context( *, group_ids: torch.Tensor, parent_ids: torch.Tensor, input_pos: torch.Tensor | None = None, - sliding_window: int | None = None, +) -> PreparedBlockMaskContext: + if group_ids.ndim != 1 or parent_ids.ndim != 1: + raise RuntimeError( + "Shared-prefix sparse block masks require rank-1 group_ids and parent_ids." + ) + if int(group_ids.numel()) != int(parent_ids.numel()): + raise RuntimeError( + "Shared-prefix sparse block masks require equal group_ids and parent_ids lengths." + ) + if input_pos is not None and int(input_pos.numel()) != int(group_ids.numel()): + raise RuntimeError( + "Shared-prefix sparse block masks require input_pos to match group_ids length." + ) + flat_group_ids = group_ids.detach().to(device="cpu", dtype=torch.int64).reshape(-1) + flat_parent_ids = ( + parent_ids.detach().to(device="cpu", dtype=torch.int64).reshape(-1) + ) + flat_input_pos = ( + None + if input_pos is None + else input_pos.detach().to(device="cpu", dtype=torch.int64).reshape(-1) + ) + row_tree = parse_prefix_tree_row( + group_ids=flat_group_ids, + parent_ids=flat_parent_ids, + ) + group_enter_np, group_exit_np = _build_group_interval_arrays( + row_tree=row_tree, + length=int(flat_group_ids.numel()), + ) + return PreparedBlockMaskContext( + source_len=int(flat_group_ids.numel()), + group_enter_np=group_enter_np, + group_exit_np=group_exit_np, + input_pos_np=None if flat_input_pos is None else flat_input_pos.numpy(), + ) + + +def _require_input_pos(context: PreparedBlockMaskContext) -> np.ndarray: + if context.input_pos_np is None: + raise RuntimeError("Sliding-window prefix-tree block masks require input_pos.") + return context.input_pos_np + + +def _validate_exact_indices( + indices: torch.Tensor, + *, + name: str, + source_len: int, +) -> int: + if indices.ndim != 1: + raise RuntimeError(f"{name} exact token indices must be rank 1.") + if indices.dtype != torch.int64: + raise RuntimeError(f"{name} exact token indices must be int64.") + indices_cpu = indices.detach().to(device="cpu", dtype=torch.int64).contiguous() + invalid = indices_cpu < 0 + if bool(invalid.any().item()): + first_invalid = int(torch.nonzero(invalid, as_tuple=False)[0].item()) + if bool((indices_cpu[first_invalid:] >= 0).any().item()): + raise RuntimeError( + f"{name} exact token indices must use only contiguous tail padding." + ) + indices_cpu = indices_cpu[:first_invalid] + if int(indices_cpu.numel()) == 0: + return 0 + if int(indices_cpu.unique().numel()) != int(indices_cpu.numel()): + raise RuntimeError(f"{name} exact token indices must not contain duplicates.") + max_index = int(indices_cpu.max().item()) + if max_index >= int(source_len): + raise RuntimeError( + f"{name} exact token index {max_index} exceeds source metadata length {int(source_len)}." + ) + return int(indices_cpu.numel()) + + +def _validate_supported_mask_spec( + spec: FlexMaskSpec, + *, + source_len: int, +) -> None: + q_valid_len = _validate_exact_indices( + spec.exact_mask.q_token_indices, + name="q", + source_len=source_len, + ) + k_valid_len = _validate_exact_indices( + spec.exact_mask.k_token_indices, + name="k", + source_len=source_len, + ) + for slice_ in spec.slices: + if int(slice_.row_index) != 0: + raise RuntimeError( + "Shared-prefix sparse block masks support exactly one packed row." + ) + if slice_.mask_kind not in {AttnMaskKind.FULL, AttnMaskKind.CAUSAL}: + raise RuntimeError(f"Unsupported attention mask kind: {slice_.mask_kind}") + if ( + slice_.q_range.start < 0 + or slice_.q_range.end > int(spec.q_len) + or slice_.k_range.start < 0 + or slice_.k_range.end > int(spec.k_len) + or slice_.q_range.end < slice_.q_range.start + or slice_.k_range.end < slice_.k_range.start + ): + raise RuntimeError(f"Attention slice is outside mask bounds: {slice_}") + if slice_.q_range.end > q_valid_len or slice_.k_range.end > k_valid_len: + raise RuntimeError( + "Attention slices may not cover exact-index tail padding: " + f"slice={slice_}, q_valid_len={q_valid_len}, k_valid_len={k_valid_len}" + ) + + +def build_block_mask_from_context( + spec: FlexMaskSpec, + *, + context: PreparedBlockMaskContext, device: torch.device, + sliding_window: int | None = None, + validate: bool = True, ) -> BlockMask | None: if spec.q_len <= 0 or spec.k_len <= 0: return None @@ -498,13 +859,16 @@ def build_block_mask( "Exact stage k-token metadata length mismatch: " f"{int(spec.exact_mask.k_token_indices.numel())} != {int(spec.k_len)}" ) + if validate: + _validate_supported_mask_spec( + spec, + source_len=context.source_len, + ) block_size = normalize_sparse_block_size(spec.block_size) return _build_sparse_block_mask( spec, device=device, - group_ids=group_ids, - parent_ids=parent_ids, - input_pos=input_pos, + context=context, sliding_window=sliding_window, block_size=block_size, ) diff --git a/src/art/megatron/context_parallel/builder.py b/src/art/megatron/context_parallel/builder.py index 77ac1b623..27016aac5 100644 --- a/src/art/megatron/context_parallel/builder.py +++ b/src/art/megatron/context_parallel/builder.py @@ -2,110 +2,17 @@ import torch +from art.megatron.prefix_tree import parse_prefix_tree + from .types import ( AttnMaskKind, AttnSlice, PackedBatchAttentionSpec, PackedRowAttentionSpec, - SharedPrefixBuilderConfig, TokenRange, ) -def _valid_length( - group_ids: torch.Tensor, - parent_ids: torch.Tensor, - *, - ignore_padding_group_id: int, -) -> int: - valid_mask = group_ids != ignore_padding_group_id - valid_count = int(valid_mask.sum().item()) - if valid_count == 0: - return 0 - if not bool(valid_mask[:valid_count].all().item()): - raise RuntimeError("Padding tokens must be a contiguous tail") - return _infer_terminal_padding_length( - group_ids[:valid_count], - parent_ids[:valid_count], - ) - - -def _infer_terminal_padding_length( - group_row: torch.Tensor, - parent_row: torch.Tensor, -) -> int: - if group_row.numel() == 0: - return 0 - runs = _scan_runs(group_row, parent_row) - if len(runs) < 2: - return int(group_row.numel()) - last_start, _last_end, last_group_id, last_parent_id = runs[-1] - if last_parent_id >= 0: - return int(group_row.numel()) - terminal_pair = (last_group_id, last_parent_id) - if any( - (group_id, parent_id) == terminal_pair - for _start, _end, group_id, parent_id in runs[:-1] - ): - return last_start - return int(group_row.numel()) - - -def _scan_runs( - group_row: torch.Tensor, - parent_row: torch.Tensor, -) -> list[tuple[int, int, int, int]]: - length = int(group_row.numel()) - if length == 0: - return [] - - group_changes = group_row[1:] != group_row[:-1] - parent_changes = parent_row[1:] != parent_row[:-1] - inconsistent_parent = torch.nonzero( - torch.logical_not(group_changes) & parent_changes, - as_tuple=False, - ).flatten() - if int(inconsistent_parent.numel()) > 0: - mismatch_index = int(inconsistent_parent[0].item()) + 1 - prior_boundaries = torch.nonzero( - group_changes[: mismatch_index - 1], - as_tuple=False, - ).flatten() - start = ( - 0 - if int(prior_boundaries.numel()) == 0 - else int(prior_boundaries[-1].item()) + 1 - ) - group_id = int(group_row[start].item()) - raise RuntimeError( - "Found one group run with inconsistent parent ids: " - f"group_id={group_id}, start={start}, end={mismatch_index}" - ) - - run_starts = torch.cat( - ( - torch.zeros(1, dtype=torch.int64, device=group_row.device), - torch.nonzero(group_changes, as_tuple=False).flatten() + 1, - ) - ) - run_ends = torch.cat( - ( - run_starts[1:], - torch.tensor([length], dtype=torch.int64, device=group_row.device), - ) - ) - starts = run_starts.to(device="cpu").tolist() - ends = run_ends.to(device="cpu").tolist() - group_ids = group_row.index_select(0, run_starts).to(device="cpu").tolist() - parent_ids = parent_row.index_select(0, run_starts).to(device="cpu").tolist() - return [ - (int(start), int(end), int(group_id), int(parent_id)) - for start, end, group_id, parent_id in zip( - starts, ends, group_ids, parent_ids, strict=True - ) - ] - - def _sort_and_dedupe_slices(slices: list[AttnSlice]) -> tuple[AttnSlice, ...]: sorted_slices = sorted( slices, @@ -138,23 +45,11 @@ def _sort_and_dedupe_slices(slices: list[AttnSlice]) -> tuple[AttnSlice, ...]: return tuple(deduped) -def _is_prompt_run( - *, - start: int, - group_id: int, - parent_id: int, - ignore_padding_group_id: int, -) -> bool: - return group_id == parent_id or ( - start == 0 and parent_id == ignore_padding_group_id - ) - - -def build_shared_prefix_attention_spec( +def build_prefix_tree_attention_spec( *, group_ids: torch.Tensor, parent_ids: torch.Tensor, - config: SharedPrefixBuilderConfig = SharedPrefixBuilderConfig(), + ignore_padding_group_id: int = -1, ) -> PackedBatchAttentionSpec: if group_ids.shape != parent_ids.shape: raise RuntimeError( @@ -166,127 +61,49 @@ def build_shared_prefix_attention_spec( "group_ids and parent_ids must be rank-2 packed tensors, got " f"{group_ids.ndim}" ) - if int(group_ids.shape[0]) != 1: - raise RuntimeError( - "ART shared-prefix attention spec currently supports exactly one packed sequence, " - f"got batch={int(group_ids.shape[0])}." - ) - rows: list[PackedRowAttentionSpec] = [] - for row_index in range(group_ids.shape[0]): - group_row = group_ids[row_index] - parent_row = parent_ids[row_index] - valid_tokens = _valid_length( - group_row, - parent_row, - ignore_padding_group_id=config.ignore_padding_group_id, - ) - if valid_tokens == 0: + for row in parse_prefix_tree( + group_ids=group_ids, + parent_ids=parent_ids, + ignore_padding_group_id=ignore_padding_group_id, + ): + if row.valid_tokens == 0: rows.append( - PackedRowAttentionSpec(row_index=row_index, valid_tokens=0, slices=()) + PackedRowAttentionSpec( + row_index=row.row_index, valid_tokens=0, slices=() + ) ) continue - group_row = group_row[:valid_tokens] - parent_row = parent_row[:valid_tokens] - runs = _scan_runs(group_row, parent_row) - - group_run_count: dict[int, int] = {} - prompt_by_group_id: dict[int, tuple[tuple[int, int], int]] = {} - completion_ranges_by_prompt: dict[int, list[tuple[int, int]]] = {} - - for start, end, group_id, parent_id in runs: - group_run_count[group_id] = group_run_count.get(group_id, 0) + 1 - if _is_prompt_run( - start=start, - group_id=group_id, - parent_id=parent_id, - ignore_padding_group_id=config.ignore_padding_group_id, - ): - if group_id in prompt_by_group_id: - raise RuntimeError( - f"Prompt group_id {group_id} appears more than once in row {row_index}" - ) - family_index = len(prompt_by_group_id) - prompt_by_group_id[group_id] = ( - (start, end), - family_index, - ) - completion_ranges_by_prompt[group_id] = [] - - if config.require_contiguous_group_runs: - repeated_groups = { - group_id: count - for group_id, count in group_run_count.items() - if count > 1 and group_id != config.ignore_padding_group_id - } - if repeated_groups: - raise RuntimeError( - "Shared-prefix builder requires contiguous group runs per row, " - f"found repeats in row {row_index}: {repeated_groups}" - ) - - for start, end, group_id, parent_id in runs: - if _is_prompt_run( - start=start, - group_id=group_id, - parent_id=parent_id, - ignore_padding_group_id=config.ignore_padding_group_id, - ): - continue - prompt_entry = prompt_by_group_id.get(parent_id) - if prompt_entry is None: - raise RuntimeError( - "Completion run points to a missing prompt run: " - f"row={row_index}, group_id={group_id}, parent_id={parent_id}" - ) - completion_ranges_by_prompt[parent_id].append((start, end)) - + segment_by_group_id = {segment.group_id: segment for segment in row.segments} row_slices: list[AttnSlice] = [] - for prompt_group_id, ( - (prompt_start, prompt_end), - family_index, - ) in prompt_by_group_id.items(): - prompt_range = TokenRange(start=prompt_start, end=prompt_end) - row_slices.append( - AttnSlice( - q_range=prompt_range, - k_range=prompt_range, - mask_kind=AttnMaskKind.CAUSAL, - row_index=row_index, - family_index=family_index, - ) - ) - for completion_start, completion_end in completion_ranges_by_prompt[ - prompt_group_id - ]: - completion_range = TokenRange( - start=completion_start, - end=completion_end, - ) + for segment in row.segments: + q_range = TokenRange(start=segment.start, end=segment.end) + for ancestor_group_id in segment.ancestors: + ancestor = segment_by_group_id[ancestor_group_id] row_slices.append( AttnSlice( - q_range=completion_range, - k_range=prompt_range, + q_range=q_range, + k_range=TokenRange(start=ancestor.start, end=ancestor.end), mask_kind=AttnMaskKind.FULL, - row_index=row_index, - family_index=family_index, + row_index=row.row_index, + family_index=segment.family_index, ) ) - row_slices.append( - AttnSlice( - q_range=completion_range, - k_range=completion_range, - mask_kind=AttnMaskKind.CAUSAL, - row_index=row_index, - family_index=family_index, - ) + row_slices.append( + AttnSlice( + q_range=q_range, + k_range=q_range, + mask_kind=AttnMaskKind.CAUSAL, + row_index=row.row_index, + family_index=segment.family_index, ) + ) rows.append( PackedRowAttentionSpec( - row_index=row_index, - valid_tokens=valid_tokens, + row_index=row.row_index, + valid_tokens=row.valid_tokens, slices=_sort_and_dedupe_slices(row_slices), ) ) diff --git a/src/art/megatron/context_parallel/comm.py b/src/art/megatron/context_parallel/comm.py index c1767a4dc..8ea97067d 100644 --- a/src/art/megatron/context_parallel/comm.py +++ b/src/art/megatron/context_parallel/comm.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass from typing import Any, Protocol, cast @@ -141,37 +142,27 @@ def wait_post_process(self) -> tuple[torch.Tensor, torch.Tensor]: if range_.size() > 0 ) - def _apply_reduce() -> None: - dk_reduce = ( - dk_remote - if dk_remote.dtype == self.dk_local.dtype - else dk_remote.to(dtype=self.dk_local.dtype) - ) - dv_reduce = ( - dv_remote - if dv_remote.dtype == self.dv_local.dtype - else dv_remote.to(dtype=self.dv_local.dtype) - ) - reduce_fn = ( - range_reduce_sum_head_major_ - if self.input_layout == "head_major" - else range_reduce_sum_ - ) - reduce_fn( - dk_reduce, - output_tensor=self.dk_local, - ranges=flattened_ranges, - range_meta_cache=self.range_meta_cache, - ) - reduce_fn( - dv_reduce, - output_tensor=self.dv_local, - ranges=flattened_ranges, - range_meta_cache=self.range_meta_cache, - ) - return - - _apply_reduce() + reduce_fn = ( + range_reduce_sum_head_major_ + if self.input_layout == "head_major" + else range_reduce_sum_ + ) + reduce_fn( + dk_remote + if dk_remote.dtype == self.dk_local.dtype + else dk_remote.to(dtype=self.dk_local.dtype), + output_tensor=self.dk_local, + ranges=flattened_ranges, + range_meta_cache=self.range_meta_cache, + ) + reduce_fn( + dv_remote + if dv_remote.dtype == self.dv_local.dtype + else dv_remote.to(dtype=self.dv_local.dtype), + output_tensor=self.dv_local, + ranges=flattened_ranges, + range_meta_cache=self.range_meta_cache, + ) return self.dk_local, self.dv_local @@ -191,6 +182,60 @@ def _get_stream(self, tensor: torch.Tensor) -> torch.cuda.Stream | None: self._streams[device_index] = stream return stream + def _launch_exchange( + self, + *, + tensor: torch.Tensor, + recv_buffer: torch.Tensor, + total_send_rows: int, + make_send_buffer: Callable[[], torch.Tensor], + output_split_sizes: list[int], + input_split_sizes: list[int], + group: Any, + async_op: bool, + input_layout: str, + ) -> tuple[_Waitable | None, torch.Tensor, torch.cuda.Stream | None]: + stream = self._get_stream(tensor) if async_op else None + send_buffer = ( + tensor.new_empty( + _packed_peer_tensor_shape( + tensor=tensor, + total_rows=0, + input_layout=input_layout, + ) + ) + if total_send_rows <= 0 + else make_send_buffer() + ) + if stream is None: + return ( + _launch_peer_exchange( + recv_buffer=recv_buffer, + send_buffer=send_buffer, + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=group, + async_op=async_op, + ), + send_buffer, + None, + ) + + current_stream = torch.cuda.current_stream(tensor.device) + stream.wait_stream(current_stream) + send_buffer.record_stream(stream) + recv_buffer.record_stream(stream) + with torch.cuda.stream(stream): + handle = _launch_peer_exchange( + recv_buffer=recv_buffer, + send_buffer=send_buffer, + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=group, + async_op=True, + ) + return handle, send_buffer, stream + def launch_kv_fetch( self, *, @@ -230,70 +275,23 @@ def launch_kv_fetch( ) input_split_sizes = [split * 2 for split in plan.send_splits] output_split_sizes = [split * 2 for split in plan.recv_splits] - stream = self._get_stream(k_local) if async_op else None - if stream is not None: - current_stream = torch.cuda.current_stream(k_local.device) - if total_send_rows <= 0: - send_buffer = k_local.new_empty( - _packed_peer_tensor_shape( - tensor=k_local, - total_rows=0, - input_layout=input_layout, - ) - ) - else: - send_buffer = _pack_gathered_tensors_per_peer( - left_tensor=k_local, - right_tensor=v_local, - ranges_by_peer=plan.send_ranges_by_peer, - range_meta_cache=range_meta_cache, - input_layout=input_layout, - ) - stream.wait_stream(current_stream) - send_buffer.record_stream(stream) - recv_packed.record_stream(stream) - with torch.cuda.stream(stream): - handle = _launch_peer_exchange( - recv_buffer=recv_packed, - send_buffer=send_buffer, - output_split_sizes=output_split_sizes, - input_split_sizes=input_split_sizes, - group=group, - async_op=True, - ) - else: - if total_send_rows <= 0: - send_buffer = k_local.new_empty( - _packed_peer_tensor_shape( - tensor=k_local, - total_rows=0, - input_layout=input_layout, - ) - ) - handle = _launch_peer_exchange( - recv_buffer=recv_packed, - send_buffer=send_buffer, - output_split_sizes=output_split_sizes, - input_split_sizes=input_split_sizes, - group=group, - async_op=async_op, - ) - else: - send_buffer = _pack_gathered_tensors_per_peer( - left_tensor=k_local, - right_tensor=v_local, - ranges_by_peer=plan.send_ranges_by_peer, - range_meta_cache=range_meta_cache, - input_layout=input_layout, - ) - handle = _launch_peer_exchange( - recv_buffer=recv_packed, - send_buffer=send_buffer, - output_split_sizes=output_split_sizes, - input_split_sizes=input_split_sizes, - group=group, - async_op=async_op, - ) + handle, send_buffer, stream = self._launch_exchange( + tensor=k_local, + recv_buffer=recv_packed, + total_send_rows=total_send_rows, + make_send_buffer=lambda: _pack_gathered_tensors_per_peer( + left_tensor=k_local, + right_tensor=v_local, + ranges_by_peer=plan.send_ranges_by_peer, + range_meta_cache=range_meta_cache, + input_layout=input_layout, + ), + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=group, + async_op=async_op, + input_layout=input_layout, + ) return KvFetchWork( packed_buffer=recv_packed, recv_splits=plan.recv_splits, @@ -333,89 +331,33 @@ def launch_dkv_reduce( total_send_rows = int(sum(plan.send_splits)) recv_total = int(sum(plan.recv_splits)) - recv_packed = ( - dk_remote.new_empty( - _packed_peer_tensor_shape( - tensor=dk_remote, - total_rows=recv_total, - input_layout=input_layout, - ) - ) - if recv_total > 0 - else dk_remote.new_empty( - _packed_peer_tensor_shape( - tensor=dk_remote, - total_rows=0, - input_layout=input_layout, - ) + recv_packed = dk_remote.new_empty( + _packed_peer_tensor_shape( + tensor=dk_remote, + total_rows=recv_total, + input_layout=input_layout, ) ) input_split_sizes = [split * 2 for split in plan.send_splits] output_split_sizes = [split * 2 for split in plan.recv_splits] - stream = self._get_stream(dk_remote) if async_op else None - if stream is not None: - current_stream = torch.cuda.current_stream(dk_remote.device) - if total_send_rows <= 0: - send_buffer = dk_remote.new_empty( - _packed_peer_tensor_shape( - tensor=dk_remote, - total_rows=0, - input_layout=input_layout, - ) - ) - else: - send_buffer = _pack_split_tensors_by_peer( - left_tensor=dk_remote, - right_tensor=dv_remote, - splits=plan.send_splits, - input_layout=input_layout, - ) - stream.wait_stream(current_stream) - send_buffer.record_stream(stream) - recv_packed.record_stream(stream) - with torch.cuda.stream(stream): - handle = _launch_peer_exchange( - recv_buffer=recv_packed, - send_buffer=send_buffer, - output_split_sizes=output_split_sizes, - input_split_sizes=input_split_sizes, - group=group, - async_op=True, - ) - else: - if total_send_rows <= 0: - send_buffer = dk_remote.new_empty( - _packed_peer_tensor_shape( - tensor=dk_remote, - total_rows=0, - input_layout=input_layout, - ) - ) - handle = _launch_peer_exchange( - recv_buffer=recv_packed, - send_buffer=send_buffer, - output_split_sizes=output_split_sizes, - input_split_sizes=input_split_sizes, - group=group, - async_op=async_op, - ) - else: - send_buffer = _pack_split_tensors_by_peer( - left_tensor=dk_remote, - right_tensor=dv_remote, - splits=plan.send_splits, - input_layout=input_layout, - ) - handle = _launch_peer_exchange( - recv_buffer=recv_packed, - send_buffer=send_buffer, - output_split_sizes=output_split_sizes, - input_split_sizes=input_split_sizes, - group=group, - async_op=async_op, - ) + handle, send_buffer, stream = self._launch_exchange( + tensor=dk_remote, + recv_buffer=recv_packed, + total_send_rows=total_send_rows, + make_send_buffer=lambda: _pack_split_tensors_by_peer( + left_tensor=dk_remote, + right_tensor=dv_remote, + splits=plan.send_splits, + input_layout=input_layout, + ), + output_split_sizes=output_split_sizes, + input_split_sizes=input_split_sizes, + group=group, + async_op=async_op, + input_layout=input_layout, + ) return DkvReduceWork( - packed_buffer=recv_packed if recv_total > 0 else None, + packed_buffer=recv_packed, handle=handle, send_buffer=send_buffer, stream=stream, @@ -449,29 +391,6 @@ def range_gather_per_peer( return torch.cat(chunks, dim=0).contiguous() -def _split_tensor_to_peer( - input_tensor: torch.Tensor, - splits: tuple[int, ...], -) -> torch.Tensor: - if int(sum(splits)) == 0: - return input_tensor.new_empty((0, *input_tensor.shape[1:])) - if int(input_tensor.shape[0]) == int(sum(splits)): - return input_tensor.contiguous() - if len([split for split in splits if split > 0]) > 1: - raise RuntimeError( - f"Expected at most one non-zero send split for dKV reduce, got {splits}" - ) - pieces: list[torch.Tensor] = [] - cursor = 0 - for split in splits: - if split == 0: - pieces.append(input_tensor.new_empty((0, *input_tensor.shape[1:]))) - continue - pieces.append(input_tensor[cursor : cursor + split]) - cursor += split - return torch.cat(pieces, dim=0).contiguous() - - def _pack_gathered_tensors_per_peer( *, left_tensor: torch.Tensor, @@ -480,56 +399,16 @@ def _pack_gathered_tensors_per_peer( range_meta_cache: dict[Any, Any] | None = None, input_layout: str = "token_major", ) -> torch.Tensor: - if input_layout == "head_major": - return _pack_gathered_tensors_per_peer_head_major( - left_tensor=left_tensor, - right_tensor=right_tensor, - ranges_by_peer=ranges_by_peer, - range_meta_cache=range_meta_cache, - ) - if input_layout != "token_major": - raise ValueError(f"Unsupported gathered-pack input layout: {input_layout}") - total_rows = sum( - range_.size() for peer_ranges in ranges_by_peer for range_ in peer_ranges - ) - if total_rows == 0: - return left_tensor.new_empty((0, *left_tensor.shape[1:])) - packed = left_tensor.new_empty((total_rows * 2, *left_tensor.shape[1:])) - cursor = 0 - for peer_ranges in ranges_by_peer: - split = sum(range_.size() for range_ in peer_ranges) - if split <= 0: - continue - range_gather( - left_tensor, - peer_ranges, - output=packed[cursor : cursor + split], - range_meta_cache=range_meta_cache, - ) - range_gather( - right_tensor, - peer_ranges, - output=packed[cursor + split : cursor + split * 2], - range_meta_cache=range_meta_cache, - ) - cursor += split * 2 - return packed - - -def _pack_gathered_tensors_per_peer_head_major( - *, - left_tensor: torch.Tensor, - right_tensor: torch.Tensor, - ranges_by_peer: tuple[tuple[TokenRange, ...], ...], - range_meta_cache: dict[Any, Any] | None = None, -) -> torch.Tensor: + _validate_peer_layout(input_layout, context="gathered-pack input") total_rows = sum( range_.size() for peer_ranges in ranges_by_peer for range_ in peer_ranges ) - if total_rows == 0: - return left_tensor.new_empty((0, left_tensor.shape[0], left_tensor.shape[2])) packed = left_tensor.new_empty( - (total_rows * 2, left_tensor.shape[0], left_tensor.shape[2]) + _packed_peer_tensor_shape( + tensor=left_tensor, + total_rows=total_rows, + input_layout=input_layout, + ) ) cursor = 0 for peer_ranges in ranges_by_peer: @@ -537,18 +416,20 @@ def _pack_gathered_tensors_per_peer_head_major( if split <= 0: continue packed[cursor : cursor + split].copy_( - range_gather_head_major( + _gather_peer_rows( left_tensor, peer_ranges, + input_layout=input_layout, range_meta_cache=range_meta_cache, - ).permute(1, 0, 2) + ) ) packed[cursor + split : cursor + split * 2].copy_( - range_gather_head_major( + _gather_peer_rows( right_tensor, peer_ranges, + input_layout=input_layout, range_meta_cache=range_meta_cache, - ).permute(1, 0, 2) + ) ) cursor += split * 2 return packed @@ -561,79 +442,83 @@ def _pack_split_tensors_by_peer( splits: tuple[int, ...], input_layout: str = "token_major", ) -> torch.Tensor: - if input_layout == "head_major": - return _pack_split_tensors_by_peer_head_major( - left_tensor=left_tensor, - right_tensor=right_tensor, - splits=splits, - ) - if input_layout != "token_major": - raise ValueError(f"Unsupported split-pack input layout: {input_layout}") + _validate_peer_layout(input_layout, context="split-pack input") total_rows = int(sum(splits)) - if total_rows == 0: - return left_tensor.new_empty((0, *left_tensor.shape[1:])) - packed = left_tensor.new_empty((total_rows * 2, *left_tensor.shape[1:])) + packed = left_tensor.new_empty( + _packed_peer_tensor_shape( + tensor=left_tensor, + total_rows=total_rows, + input_layout=input_layout, + ) + ) cursor = 0 for split in splits: if split <= 0: continue packed[cursor * 2 : cursor * 2 + split].copy_( - left_tensor[cursor : cursor + split] + _slice_peer_rows(left_tensor, cursor, cursor + split, layout=input_layout) ) packed[cursor * 2 + split : cursor * 2 + split * 2].copy_( - right_tensor[cursor : cursor + split] + _slice_peer_rows(right_tensor, cursor, cursor + split, layout=input_layout) ) cursor += split - if cursor != int(left_tensor.shape[0]) or cursor != int(right_tensor.shape[0]): + left_rows = _peer_row_count(left_tensor, layout=input_layout) + right_rows = _peer_row_count(right_tensor, layout=input_layout) + if cursor != left_rows or cursor != right_rows: raise RuntimeError( "Packed split consumed the wrong number of rows: " - f"consumed={cursor}, left={int(left_tensor.shape[0])}, right={int(right_tensor.shape[0])}" + f"consumed={cursor}, left={left_rows}, right={right_rows}" ) return packed +def _validate_peer_layout(layout: str, *, context: str) -> None: + if layout not in {"token_major", "head_major"}: + raise ValueError(f"Unsupported {context} layout: {layout}") + + def _packed_peer_tensor_shape( *, tensor: torch.Tensor, total_rows: int, input_layout: str, ) -> tuple[int, ...]: + _validate_peer_layout(input_layout, context="peer tensor input") if input_layout == "head_major": return (total_rows * 2, int(tensor.shape[0]), int(tensor.shape[2])) - if input_layout != "token_major": - raise ValueError(f"Unsupported split-pack input layout: {input_layout}") return (total_rows * 2, *tuple(int(dim) for dim in tensor.shape[1:])) -def _pack_split_tensors_by_peer_head_major( +def _peer_row_count(tensor: torch.Tensor, *, layout: str) -> int: + return int(tensor.shape[1] if layout == "head_major" else tensor.shape[0]) + + +def _slice_peer_rows( + tensor: torch.Tensor, + start: int, + end: int, *, - left_tensor: torch.Tensor, - right_tensor: torch.Tensor, - splits: tuple[int, ...], + layout: str, ) -> torch.Tensor: - total_rows = int(sum(splits)) - if total_rows == 0: - return left_tensor.new_empty((0, left_tensor.shape[0], left_tensor.shape[2])) - packed = left_tensor.new_empty( - (total_rows * 2, left_tensor.shape[0], left_tensor.shape[2]) - ) - cursor = 0 - for split in splits: - if split <= 0: - continue - packed[cursor * 2 : cursor * 2 + split].copy_( - left_tensor[:, cursor : cursor + split].permute(1, 0, 2) - ) - packed[cursor * 2 + split : cursor * 2 + split * 2].copy_( - right_tensor[:, cursor : cursor + split].permute(1, 0, 2) - ) - cursor += split - if cursor != int(left_tensor.shape[1]) or cursor != int(right_tensor.shape[1]): - raise RuntimeError( - "Head-major split pack consumed the wrong number of rows: " - f"consumed={cursor}, left={int(left_tensor.shape[1])}, right={int(right_tensor.shape[1])}" - ) - return packed + if layout == "head_major": + return tensor[:, start:end].movedim(1, 0) + return tensor[start:end] + + +def _gather_peer_rows( + tensor: torch.Tensor, + ranges: tuple[TokenRange, ...], + *, + input_layout: str, + range_meta_cache: dict[Any, Any] | None, +) -> torch.Tensor: + if input_layout == "head_major": + return range_gather_head_major( + tensor, + ranges, + range_meta_cache=range_meta_cache, + ).movedim(1, 0) + return range_gather(tensor, ranges, range_meta_cache=range_meta_cache) def _unpack_packed_tensor_per_peer( @@ -642,15 +527,13 @@ def _unpack_packed_tensor_per_peer( *, output_layout: str = "token_major", ) -> tuple[torch.Tensor, torch.Tensor]: - if output_layout == "head_major": - return _unpack_packed_tensor_per_peer_head_major( + _validate_peer_layout(output_layout, context="packed-tensor output") + if int(packed_tensor.shape[0]) == 0: + empty = _new_unpacked_peer_tensor( packed_tensor, - splits, + total_rows=0, + output_layout=output_layout, ) - if output_layout != "token_major": - raise ValueError(f"Unsupported packed-tensor output layout: {output_layout}") - if int(packed_tensor.shape[0]) == 0: - empty = packed_tensor.new_empty((0, *packed_tensor.shape[1:])) return empty, empty total_rows = 0 cursor = 0 @@ -664,62 +547,59 @@ def _unpack_packed_tensor_per_peer( "Packed tensor unpack consumed the wrong number of rows: " f"consumed={cursor}, input={int(packed_tensor.shape[0])}" ) - left = packed_tensor.new_empty((total_rows, *packed_tensor.shape[1:])) - right = packed_tensor.new_empty((total_rows, *packed_tensor.shape[1:])) + left = _new_unpacked_peer_tensor( + packed_tensor, + total_rows=total_rows, + output_layout=output_layout, + ) + right = _new_unpacked_peer_tensor( + packed_tensor, + total_rows=total_rows, + output_layout=output_layout, + ) in_cursor = 0 out_cursor = 0 for split in splits: if split <= 0: continue - left[out_cursor : out_cursor + split].copy_( - packed_tensor[in_cursor : in_cursor + split] + _copy_from_peer_rows( + left, + out_cursor, + packed_tensor[in_cursor : in_cursor + split], + output_layout=output_layout, ) - right[out_cursor : out_cursor + split].copy_( - packed_tensor[in_cursor + split : in_cursor + split * 2] + _copy_from_peer_rows( + right, + out_cursor, + packed_tensor[in_cursor + split : in_cursor + split * 2], + output_layout=output_layout, ) in_cursor += split * 2 out_cursor += split return left, right -def _unpack_packed_tensor_per_peer_head_major( +def _new_unpacked_peer_tensor( packed_tensor: torch.Tensor, - splits: tuple[int, ...], -) -> tuple[torch.Tensor, torch.Tensor]: - if int(packed_tensor.shape[0]) == 0: - empty = packed_tensor.new_empty( - (packed_tensor.shape[1], 0, packed_tensor.shape[2]) - ) - return empty, empty - total_rows = 0 - cursor = 0 - for split in splits: - if split <= 0: - continue - cursor += split * 2 - total_rows += split - if cursor != int(packed_tensor.shape[0]): - raise RuntimeError( - "Packed tensor unpack consumed the wrong number of rows: " - f"consumed={cursor}, input={int(packed_tensor.shape[0])}" - ) - left = packed_tensor.new_empty( - (packed_tensor.shape[1], total_rows, packed_tensor.shape[2]) - ) - right = packed_tensor.new_empty( - (packed_tensor.shape[1], total_rows, packed_tensor.shape[2]) - ) - in_cursor = 0 - out_cursor = 0 - for split in splits: - if split <= 0: - continue - left[:, out_cursor : out_cursor + split].copy_( - packed_tensor[in_cursor : in_cursor + split].permute(1, 0, 2) - ) - right[:, out_cursor : out_cursor + split].copy_( - packed_tensor[in_cursor + split : in_cursor + split * 2].permute(1, 0, 2) + *, + total_rows: int, + output_layout: str, +) -> torch.Tensor: + if output_layout == "head_major": + return packed_tensor.new_empty( + (packed_tensor.shape[1], total_rows, *packed_tensor.shape[2:]) ) - in_cursor += split * 2 - out_cursor += split - return left, right + return packed_tensor.new_empty((total_rows, *packed_tensor.shape[1:])) + + +def _copy_from_peer_rows( + output: torch.Tensor, + start: int, + rows: torch.Tensor, + *, + output_layout: str, +) -> None: + if output_layout == "head_major": + output[:, start : start + int(rows.shape[0])].copy_(rows.movedim(0, 1)) + else: + output[start : start + int(rows.shape[0])].copy_(rows) diff --git a/src/art/megatron/context_parallel/core_attention.py b/src/art/megatron/context_parallel/core_attention.py index d35d708f0..e7efc7331 100644 --- a/src/art/megatron/context_parallel/core_attention.py +++ b/src/art/megatron/context_parallel/core_attention.py @@ -14,7 +14,7 @@ from art.megatron.flex_attn.attention import ( FlexAttentionWrapper, - SharedPrefixAttentionState, + PrefixTreeAttentionState, _configure_softmax_offset, _triton_num_stages_2_head_dims, ) @@ -109,13 +109,13 @@ def forward( softmax_offset=cast(Tensor | None, getattr(self, "softmax_offset")), ) else: - if isinstance(attention_bias, SharedPrefixAttentionState): + if isinstance(attention_bias, PrefixTreeAttentionState): block_mask = attention_bias.block_mask_for_window( getattr(self, "art_sliding_window", None) ) else: assert isinstance(attention_bias, BlockMask), ( - "Expected ArtContextParallelState, SharedPrefixAttentionState, or BlockMask in attention_bias." + "Expected ArtContextParallelState, PrefixTreeAttentionState, or BlockMask in attention_bias." ) block_mask = attention_bias q = query.permute(1, 2, 0, 3) diff --git a/src/art/megatron/context_parallel/executor.py b/src/art/megatron/context_parallel/executor.py index 13f7e332e..3636d7361 100644 --- a/src/art/megatron/context_parallel/executor.py +++ b/src/art/megatron/context_parallel/executor.py @@ -20,7 +20,7 @@ sparse_compiled_flex_attention, ) -from .block_mask import build_block_mask +from .block_mask import build_block_mask_from_context, prepare_block_mask_context from .comm import A2AVCommunicator from .range_ops import ( range_gather_head_major, @@ -748,19 +748,26 @@ def _build_stage_block_mask( raise RuntimeError( f"Stage {stage_plan.stage_index} is missing exact mask metadata" ) - mask = build_block_mask( + block_mask_context = state.execution_cache.block_mask_context + if block_mask_context is None: + block_mask_context = prepare_block_mask_context( + group_ids=state.group_ids, + parent_ids=state.parent_ids, + input_pos=state.input_pos, + ) + state.execution_cache.block_mask_context = block_mask_context + mask = build_block_mask_from_context( FlexMaskSpec( q_len=int(execution_spec.q_len), k_len=int(execution_spec.k_len), block_size=resolved_block_size, slices=stage_plan.slices, - exact_mask=mask_metadata.model_dump(mode="python"), + exact_mask=mask_metadata, ), - group_ids=state.group_ids, - parent_ids=state.parent_ids, - input_pos=state.input_pos, + context=block_mask_context, sliding_window=sliding_window, device=device, + validate=False, ) cache[cache_key] = mask return mask @@ -851,30 +858,6 @@ def prepare_context_parallel_execution_state( ) -def _causal_slice_pair_count(slice_: AttnSlice) -> int: - q_start = int(slice_.q_range.start) - q_end = int(slice_.q_range.end) - k_start = int(slice_.k_range.start) - k_end = int(slice_.k_range.end) - if q_end <= q_start or k_end <= k_start: - return 0 - - k_len = k_end - k_start - partial_q_start = max(q_start, k_start) - partial_q_end = min(q_end - 1, k_end - 2) - partial = 0 - if partial_q_start <= partial_q_end: - count = partial_q_end - partial_q_start + 1 - partial = count * (partial_q_start + partial_q_end + 2 - 2 * k_start) // 2 - - full_q_start = max(q_start, k_end - 1) - full_q_end = q_end - 1 - full = 0 - if full_q_start <= full_q_end: - full = (full_q_end - full_q_start + 1) * k_len - return int(partial + full) - - def _validate_stage_block_alignment( *, q_len: int, diff --git a/src/art/megatron/context_parallel/layout_index.py b/src/art/megatron/context_parallel/layout_index.py index 99fb2c35b..9f60550a0 100644 --- a/src/art/megatron/context_parallel/layout_index.py +++ b/src/art/megatron/context_parallel/layout_index.py @@ -1,10 +1,9 @@ from __future__ import annotations -from pydantic import BaseModel, ConfigDict +from dataclasses import dataclass -class TokenLayoutIndex(BaseModel): - model_config = ConfigDict(frozen=True) - +@dataclass(frozen=True) +class TokenLayoutIndex: ownership_ranges_by_rank: tuple[tuple[tuple[int, int, int], ...], ...] token_counts_by_rank: tuple[int, ...] diff --git a/src/art/megatron/context_parallel/runtime.py b/src/art/megatron/context_parallel/runtime.py index 98ee47c1d..52cbacd70 100644 --- a/src/art/megatron/context_parallel/runtime.py +++ b/src/art/megatron/context_parallel/runtime.py @@ -1,26 +1,23 @@ from __future__ import annotations from bisect import bisect_left, bisect_right +from dataclasses import dataclass, replace import hashlib import json from typing import Any, cast -import warnings -from pydantic import BaseModel, ConfigDict import torch from art.loss import shift_tensor from art.preprocessing.pack import PackedTensors -from .builder import build_shared_prefix_attention_spec +from .builder import build_prefix_tree_attention_spec from .layout_index import TokenLayoutIndex from .types import ( ArtContextParallelState, AttnMaskKind, AttnSlice, ContextParallelConfig, - ContextParallelRuntimeKey, - ContextParallelRuntimePlan, CpBlockMaskVariant, DispatchedPackedTensors, DkvReducePlan, @@ -29,17 +26,12 @@ PackedBatchAttentionSpec, PackedRowAttentionSpec, ParallelTopology, - PlannerProvenance, PreparedMegatronBatch, RankRuntimePlan, StagePlan, TokenRange, ) -_PLANNER_RUNTIME_BACKEND = "art_context_parallel" -_PLANNER_BEST_EFFORT_WARNING_KEYS: set[ - tuple[str, str, int, str, str, tuple[int, ...]] -] = set() _CHUNK_MASK_STATS_TORCH_THRESHOLD = 1024 _CP4_SEARCH_PROBE_CANDIDATE_LIMIT = 2 _CP4_SEARCH_PROBE_IMPROVEMENT_MS = 1.0 @@ -49,18 +41,16 @@ StageSliceKey = tuple[int, int, int, int, int, str, int] -class _PlanningBundle(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True) - +@dataclass(frozen=True) +class _PlanningBundle: spec: PackedBatchAttentionSpec - runtime_key: ContextParallelRuntimeKey - runtime_plan: ContextParallelRuntimePlan + rank_plans: tuple[RankRuntimePlan, ...] gdn_execution_spec: Any | None = None _PLANNING_BUNDLE_CACHE: dict[str, _PlanningBundle] = {} -_RUNTIME_PLAN_CACHE: dict[tuple[str, int], ContextParallelRuntimePlan] = {} -_GDN_RANK_PLAN_CACHE: dict[tuple[str, str, int | None, int], Any] = {} +_RUNTIME_PLAN_CACHE: dict[str, tuple[RankRuntimePlan, ...]] = {} +_GDN_RANK_PLAN_CACHE: dict[tuple[str, str, int | None, int, str], Any] = {} def _json_cache_key(payload: Any) -> str: @@ -112,174 +102,14 @@ def _planning_bundle_cache_key( { "group_ids": _metadata_tensor_digest(group_ids), "parent_ids": _metadata_tensor_digest(parent_ids), - "topology": topology.model_dump(mode="json"), - "config": config.model_dump(mode="json"), + "topology": _dataclass_payload(topology), + "config": _dataclass_payload(config), "original_seq_len": int(original_seq_len), "build_gdn_execution_spec": bool(build_gdn_execution_spec), } ) -def _rank_plan_cache_key( - *, - planning_key: str, - device: torch.device, - cp_rank: int, -) -> tuple[str, str, int | None, int]: - return (planning_key, device.type, device.index, int(cp_rank)) - - -def _config_for_runtime_cp( - *, - topology: ParallelTopology, - config: ContextParallelConfig, -) -> ContextParallelConfig: - cp_size = max(int(topology.cp), 1) - updates: dict[str, Any] = {} - applied_override = False - for override in config.planner_cp_overrides: - if int(override.cp_size) != cp_size: - continue - override_updates = override.model_dump(mode="python", exclude_none=True) - override_updates.pop("cp_size", None) - updates.update(override_updates) - applied_override = True - if not applied_override: - return config - updates.setdefault("planner_tuned_cp_sizes", (cp_size,)) - return config.model_copy(update=updates) - - -def _normalized_planner_metadata_value(value: str | None) -> str: - if value is None: - return "" - normalized = "".join( - character.lower() if character.isalnum() else " " - for character in str(value).strip() - ) - return " ".join(part for part in normalized.split() if part) - - -def _planner_metadata_matches( - expected: str | None, - actual: str | None, - *, - fuzzy: bool, -) -> bool: - normalized_expected = _normalized_planner_metadata_value(expected) - normalized_actual = _normalized_planner_metadata_value(actual) - if not normalized_expected or not normalized_actual: - return False - if normalized_expected == normalized_actual: - return True - return bool( - fuzzy - and ( - normalized_expected in normalized_actual - or normalized_actual in normalized_expected - ) - ) - - -def _planner_runtime_hardware() -> str | None: - if not torch.cuda.is_available(): - return None - try: - return str(torch.cuda.get_device_name(torch.cuda.current_device())) - except Exception: - return str(torch.cuda.get_device_name(0)) - - -def _planner_best_effort_warning_message(provenance: PlannerProvenance) -> str: - mismatch_reasons: list[str] = [] - if not provenance.backend_match: - mismatch_reasons.append( - f"backend runtime={provenance.runtime_backend!r} tuned={provenance.tuned_backend!r}" - ) - if not provenance.hardware_match: - mismatch_reasons.append( - f"hardware runtime={provenance.runtime_hardware!r} tuned={provenance.tuned_hardware!r}" - ) - if not provenance.cp_size_match: - mismatch_reasons.append( - f"cp_size runtime={int(provenance.runtime_cp_size)} tuned={list(provenance.tuned_cp_sizes)}" - ) - mismatch_text = ( - "; ".join(mismatch_reasons) if mismatch_reasons else "metadata missing" - ) - return ( - "ART context parallel planner coefficients are running in best-effort mode; " - f"{mismatch_text}. The runtime will continue with the configured coefficients." - ) - - -def _planner_provenance( - *, - topology: ParallelTopology, - config: ContextParallelConfig, - warn: bool = True, -) -> PlannerProvenance: - runtime_hardware = _planner_runtime_hardware() - tuned_cp_sizes = tuple( - sorted( - { - int(cp_size) - for cp_size in config.planner_tuned_cp_sizes - if int(cp_size) > 0 - } - ) - ) - provenance = PlannerProvenance( - runtime_backend=_PLANNER_RUNTIME_BACKEND, - runtime_hardware=runtime_hardware, - runtime_cp_size=max(int(topology.cp), 1), - tuned_backend=config.planner_tuned_backend, - tuned_hardware=config.planner_tuned_hardware, - tuned_cp_sizes=tuned_cp_sizes, - backend_match=_planner_metadata_matches( - config.planner_tuned_backend, - _PLANNER_RUNTIME_BACKEND, - fuzzy=False, - ), - hardware_match=_planner_metadata_matches( - config.planner_tuned_hardware, - runtime_hardware, - fuzzy=True, - ), - cp_size_match=bool(tuned_cp_sizes) - and max(int(topology.cp), 1) in tuned_cp_sizes, - using_best_effort=False, - ) - if ( - provenance.backend_match - and provenance.hardware_match - and provenance.cp_size_match - ): - return provenance - - warning_message = _planner_best_effort_warning_message(provenance) - warning_key = ( - _normalized_planner_metadata_value(provenance.runtime_backend), - _normalized_planner_metadata_value(provenance.runtime_hardware), - int(provenance.runtime_cp_size), - _normalized_planner_metadata_value(provenance.tuned_backend), - _normalized_planner_metadata_value(provenance.tuned_hardware), - provenance.tuned_cp_sizes, - ) - warning_emitted = False - if warn and warning_key not in _PLANNER_BEST_EFFORT_WARNING_KEYS: - _PLANNER_BEST_EFFORT_WARNING_KEYS.add(warning_key) - warnings.warn(warning_message, RuntimeWarning, stacklevel=3) - warning_emitted = True - return provenance.model_copy( - update={ - "using_best_effort": True, - "warning_message": warning_message, - "warning_emitted": warning_emitted, - } - ) - - def _normalized_chunk_size( *, valid_tokens: int, @@ -352,7 +182,7 @@ def _search_config_for_chunk_count( return config if all(int(getattr(config, key)) == int(value) for key, value in updates.items()): return config - return config.model_copy(update=updates) + return replace(config, **updates) def _best_improving_move( @@ -388,9 +218,9 @@ def _best_improving_move( candidate = list(current_owners) candidate[chunk_index] = dst_rank candidate_owners = tuple(candidate) - if not _assignment_uses_all_ranks( - candidate_owners, - cp_size=cp_size, + if ( + len(candidate_owners) >= cp_size + and len(set(candidate_owners)) != cp_size ): continue candidate_eval = evaluate_candidate( @@ -411,12 +241,10 @@ def _build_chunk_ranges( valid_tokens: int, chunk_size: int, ) -> tuple[TokenRange, ...]: - ranges: list[TokenRange] = [] - for start in range(0, valid_tokens, chunk_size): - ranges.append( - TokenRange(start=start, end=min(start + chunk_size, valid_tokens)) - ) - return tuple(ranges) + return tuple( + TokenRange(start=start, end=min(start + chunk_size, valid_tokens)) + for start in range(0, valid_tokens, chunk_size) + ) def _indexed_intersections( @@ -446,33 +274,6 @@ def _indexed_intersections( return intersections -def _slice_pair_count( - *, - mask_kind: AttnMaskKind, - q_range: TokenRange, - k_range: TokenRange, -) -> int: - if mask_kind is AttnMaskKind.FULL: - return int(q_range.size()) * int(k_range.size()) - return _causal_piece_pair_count( - q_range=q_range, - k_range=k_range, - ) - - -def _causal_piece_pair_count( - *, - q_range: TokenRange, - k_range: TokenRange, -) -> int: - return _causal_piece_pair_count_from_bounds( - q_start=int(q_range.start), - q_end=int(q_range.end), - k_start=int(k_range.start), - k_end=int(k_range.end), - ) - - def _causal_piece_pair_count_from_bounds( *, q_start: int, @@ -499,91 +300,15 @@ def _causal_piece_pair_count_from_bounds( return int(partial + full) -def _chunk_piece_decomposition( - *, - start: int, - end: int, - chunk_size: int, -) -> tuple[ - int, tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...], int -]: - first = start // chunk_size - last = (end - 1) // chunk_size - piece_starts: list[int] = [] - piece_ends: list[int] = [] - piece_lengths: list[int] = [] - piece_prefix_lengths: list[int] = [] - running_len = 0 - for chunk_index in range(first, last + 1): - piece_start = start if chunk_index == first else chunk_index * chunk_size - piece_end = end if chunk_index == last else (chunk_index + 1) * chunk_size - piece_len = piece_end - piece_start - if piece_len <= 0: - continue - running_len += piece_len - piece_starts.append(piece_start) - piece_ends.append(piece_end) - piece_lengths.append(piece_len) - piece_prefix_lengths.append(running_len) - return ( - first, - tuple(piece_starts), - tuple(piece_ends), - tuple(piece_lengths), - tuple(piece_prefix_lengths), - running_len, - ) - - -def _can_use_shared_prefix_chunk_pair_program( - row_spec: PackedRowAttentionSpec, -) -> bool: - slices = row_spec.slices - index = 0 - while index < len(slices): - prompt_slice = slices[index] - if ( - prompt_slice.family_index is None - or prompt_slice.mask_kind is not AttnMaskKind.CAUSAL - or prompt_slice.q_range != prompt_slice.k_range - ): - return False - prompt_family_index = prompt_slice.family_index - if prompt_family_index is None: - raise RuntimeError("shared-prefix prompt slices must carry family_index") - family_index = int(prompt_family_index) - prompt_start = int(prompt_slice.q_range.start) - prompt_end = int(prompt_slice.q_range.end) - index += 1 - while index < len(slices): - family_value = slices[index].family_index - if family_value is None or int(family_value) != family_index: - break - if index + 1 >= len(slices): - return False - full_slice = slices[index] - causal_slice = slices[index + 1] - if ( - full_slice.family_index != prompt_slice.family_index - or causal_slice.family_index != prompt_slice.family_index - or full_slice.mask_kind is not AttnMaskKind.FULL - or causal_slice.mask_kind is not AttnMaskKind.CAUSAL - or full_slice.q_range != causal_slice.q_range - or causal_slice.q_range != causal_slice.k_range - or int(full_slice.k_range.start) != prompt_start - or int(full_slice.k_range.end) != prompt_end - ): - return False - index += 2 - return True - - -def _build_chunk_pair_program_generic( +def _build_chunk_pair_program( row_spec: PackedRowAttentionSpec, *, - chunk_count: int, - chunk_size: int, + chunk_ranges: tuple[TokenRange, ...], ) -> tuple[torch.Tensor, list[float]]: + chunk_count = len(chunk_ranges) + if chunk_count == 0: + return torch.zeros((0, 0), dtype=torch.int64), [] + chunk_size = int(chunk_ranges[0].size()) pair_rows = [[0 for _ in range(chunk_count)] for _ in range(chunk_count)] q_weights = [0.0 for _ in range(chunk_count)] @@ -682,138 +407,6 @@ def _build_chunk_pair_program_generic( return torch.tensor(pair_rows, dtype=torch.int64), q_weights -def _build_chunk_pair_program( - row_spec: PackedRowAttentionSpec, - *, - chunk_ranges: tuple[TokenRange, ...], -) -> tuple[torch.Tensor, list[float]]: - chunk_count = len(chunk_ranges) - if chunk_count == 0: - return torch.zeros((0, 0), dtype=torch.int64), [] - chunk_size = int(chunk_ranges[0].size()) - if not _can_use_shared_prefix_chunk_pair_program(row_spec): - return _build_chunk_pair_program_generic( - row_spec, - chunk_count=chunk_count, - chunk_size=chunk_size, - ) - - pair_rows = [[0 for _ in range(chunk_count)] for _ in range(chunk_count)] - q_weights = [0.0 for _ in range(chunk_count)] - slices = row_spec.slices - index = 0 - while index < len(slices): - prompt_slice = slices[index] - ( - prompt_first, - prompt_starts, - prompt_ends, - prompt_lengths, - prompt_prefix, - prompt_total, - ) = _chunk_piece_decomposition( - start=int(prompt_slice.q_range.start), - end=int(prompt_slice.q_range.end), - chunk_size=chunk_size, - ) - for offset, q_chunk_index in enumerate( - range(prompt_first, prompt_first + len(prompt_lengths)) - ): - q_piece_len = prompt_lengths[offset] - row = pair_rows[q_chunk_index] - q_total = 0 - if offset > 0: - for k_offset in range(offset): - row[prompt_first + k_offset] += ( - q_piece_len * prompt_lengths[k_offset] - ) - q_total += q_piece_len * prompt_prefix[offset - 1] - pair_count = _causal_piece_pair_count_from_bounds( - q_start=prompt_starts[offset], - q_end=prompt_ends[offset], - k_start=prompt_starts[offset], - k_end=prompt_ends[offset], - ) - if pair_count > 0: - row[q_chunk_index] += pair_count - q_total += pair_count - if q_total > 0: - q_weights[q_chunk_index] += float(q_total) - - prompt_family_index = prompt_slice.family_index - if prompt_family_index is None: - raise RuntimeError("shared-prefix prompt slices must carry family_index") - family_index = int(prompt_family_index) - index += 1 - completion_chunk_indices: list[int] = [] - completion_chunk_totals: list[int] = [] - while index < len(slices): - family_value = slices[index].family_index - if family_value is None or int(family_value) != family_index: - break - full_slice = slices[index] - ( - completion_first, - completion_starts, - completion_ends, - completion_lengths, - completion_prefix, - _, - ) = _chunk_piece_decomposition( - start=int(full_slice.q_range.start), - end=int(full_slice.q_range.end), - chunk_size=chunk_size, - ) - for offset, q_chunk_index in enumerate( - range(completion_first, completion_first + len(completion_lengths)) - ): - q_piece_len = completion_lengths[offset] - if ( - completion_chunk_indices - and completion_chunk_indices[-1] == q_chunk_index - ): - completion_chunk_totals[-1] += q_piece_len - else: - completion_chunk_indices.append(q_chunk_index) - completion_chunk_totals.append(q_piece_len) - - for offset, q_chunk_index in enumerate( - range(completion_first, completion_first + len(completion_lengths)) - ): - q_piece_len = completion_lengths[offset] - row = pair_rows[q_chunk_index] - q_total = 0 - if offset > 0: - for k_offset in range(offset): - row[completion_first + k_offset] += ( - q_piece_len * completion_lengths[k_offset] - ) - q_total += q_piece_len * completion_prefix[offset - 1] - pair_count = _causal_piece_pair_count_from_bounds( - q_start=completion_starts[offset], - q_end=completion_ends[offset], - k_start=completion_starts[offset], - k_end=completion_ends[offset], - ) - if pair_count > 0: - row[q_chunk_index] += pair_count - q_total += pair_count - if q_total > 0: - q_weights[q_chunk_index] += float(q_total) - index += 2 - - for q_chunk_index, total_q_len in zip( - completion_chunk_indices, - completion_chunk_totals, - strict=True, - ): - row = pair_rows[q_chunk_index] - for k_offset, k_piece_len in enumerate(prompt_lengths): - row[prompt_first + k_offset] += total_q_len * k_piece_len - q_weights[q_chunk_index] += float(total_q_len * prompt_total) - return torch.tensor(pair_rows, dtype=torch.int64), q_weights - - def _collect_rank_stage_pieces( row_spec: PackedRowAttentionSpec, *, @@ -963,63 +556,6 @@ def _contiguous_chunk_assignment( return tuple(owners) -def _bucket_chunk_assignment( - *, - q_weights: list[float], - cp_size: int, -) -> tuple[int, ...]: - chunk_count = len(q_weights) - if chunk_count == 0: - return tuple() - if cp_size <= 1: - return tuple(0 for _ in range(chunk_count)) - rank_loads = [0.0 for _ in range(cp_size)] - rank_chunk_counts = [0 for _ in range(cp_size)] - owners = [-1 for _ in range(chunk_count)] - for chunk_index in sorted( - range(chunk_count), - key=lambda index: (-q_weights[index], index), - ): - rank = min( - range(cp_size), - key=lambda candidate: ( - rank_loads[candidate], - rank_chunk_counts[candidate], - candidate, - ), - ) - owners[chunk_index] = rank - rank_loads[rank] += q_weights[chunk_index] - rank_chunk_counts[rank] += 1 - return tuple(int(owner) for owner in owners) - - -def _striped_chunk_assignment( - *, - chunk_count: int, - cp_size: int, - group_size: int, -) -> tuple[int, ...]: - if chunk_count == 0: - return tuple() - if cp_size <= 1: - return tuple(0 for _ in range(chunk_count)) - group_size = max(1, int(group_size)) - return tuple( - ((chunk_index // group_size) % cp_size) for chunk_index in range(chunk_count) - ) - - -def _assignment_uses_all_ranks( - owners: tuple[int, ...], - *, - cp_size: int, -) -> bool: - if len(owners) < cp_size: - return True - return len({int(owner) for owner in owners}) == cp_size - - def _candidate_chunk_indices( *, owners: tuple[int, ...], @@ -1128,31 +664,6 @@ def _chunk_mask_stats( return token_count, range_count -def _merge_chunk_ranges_from_mask( - *, - chunk_ranges: tuple[TokenRange, ...], - chunk_mask: torch.Tensor, -) -> tuple[TokenRange, ...]: - chunk_indices = torch.nonzero(chunk_mask, as_tuple=False).flatten() - if int(chunk_indices.numel()) == 0: - return tuple() - ordered_chunk_indices = chunk_indices.tolist() - first_range = chunk_ranges[int(ordered_chunk_indices[0])] - current_start = int(first_range.start) - current_end = int(first_range.end) - merged: list[TokenRange] = [] - for chunk_index in ordered_chunk_indices[1:]: - range_ = chunk_ranges[int(chunk_index)] - if int(range_.start) <= current_end: - current_end = max(current_end, int(range_.end)) - continue - merged.append(TokenRange(start=current_start, end=current_end)) - current_start = int(range_.start) - current_end = int(range_.end) - merged.append(TokenRange(start=current_start, end=current_end)) - return tuple(merged) - - def _stage_cost_ms( *, pair_count: int, @@ -1530,31 +1041,6 @@ def _evaluate_plan( } -def _evaluate_plan_for_search( - *, - chunk_ranges: tuple[TokenRange, ...], - pair_matrix: list[list[int]] | torch.Tensor, - owners: tuple[int, ...], - wave_assignment: tuple[int, ...], - cp_size: int, - config: ContextParallelConfig, - pair_positive: torch.Tensor | None = None, - chunk_lengths: tuple[int, ...] | None = None, - chunk_lengths_tensor: torch.Tensor | None = None, -) -> dict[str, Any]: - return _evaluate_plan( - chunk_ranges=chunk_ranges, - pair_matrix=pair_matrix, - owners=owners, - wave_assignment=wave_assignment, - cp_size=cp_size, - config=config, - pair_positive=pair_positive, - chunk_lengths=chunk_lengths, - chunk_lengths_tensor=chunk_lengths_tensor, - ) - - def _search_chunk_assignment( *, chunk_ranges: tuple[TokenRange, ...], @@ -1563,7 +1049,6 @@ def _search_chunk_assignment( cp_size: int, config: ContextParallelConfig, ) -> tuple[tuple[int, ...], tuple[int, ...], dict[str, Any]]: - cp_size = int(cp_size) config = _search_config_for_chunk_count( config=config, chunk_count=len(chunk_ranges), @@ -1572,9 +1057,7 @@ def _search_chunk_assignment( 1, min(int(config.planner_max_remote_waves), len(chunk_ranges)) + 1, ) - best_owners: tuple[int, ...] = tuple() - best_waves: tuple[int, ...] = tuple() - best_eval: dict[str, Any] | None = None + best: tuple[tuple[int, ...], tuple[int, ...], dict[str, Any]] | None = None eval_cache: dict[tuple[tuple[int, ...], tuple[int, ...]], dict[str, Any]] = {} pair_counts = torch.as_tensor(pair_matrix, dtype=torch.int64) pair_positive = pair_counts > 0 @@ -1594,7 +1077,7 @@ def _evaluate_candidate( cached = eval_cache.get(cache_key) if cached is not None: return cached - cached = _evaluate_plan_for_search( + cached = _evaluate_plan( chunk_ranges=chunk_ranges, pair_matrix=pair_counts, owners=owners, @@ -1608,86 +1091,35 @@ def _evaluate_candidate( eval_cache[cache_key] = cached return cached - def _best_wave_assignment_for_owners( - owners: tuple[int, ...], - ) -> tuple[tuple[int, ...], dict[str, Any]]: - best_wave_assignment = tuple() - best_eval_local: dict[str, Any] | None = None - for wave_count in wave_count_candidates: - wave_assignment = _wave_assignment( - chunk_count=len(chunk_ranges), - wave_count=wave_count, - ) - candidate_eval = _evaluate_candidate( - owners=owners, - wave_assignment=wave_assignment, - ) - if best_eval_local is None or float(candidate_eval["score"]) + 1e-9 < float( - best_eval_local["score"] - ): - best_wave_assignment = wave_assignment - best_eval_local = candidate_eval - if best_eval_local is None: - raise RuntimeError("Failed to evaluate any wave assignment candidate.") - return best_wave_assignment, best_eval_local - - strategy = str(config.planner_assignment_strategy).strip().lower() - striped_owners = _striped_chunk_assignment( - chunk_count=len(chunk_ranges), - cp_size=cp_size, - group_size=int(config.planner_stripe_group_size), - ) - fixed_owners_by_strategy = { - "contiguous": _contiguous_chunk_assignment( - q_weights=q_weights, cp_size=cp_size - ), - "bucket": _bucket_chunk_assignment(q_weights=q_weights, cp_size=cp_size), - "striped": striped_owners, - } - if strategy in fixed_owners_by_strategy: - owners = fixed_owners_by_strategy[strategy] - best_waves, best_eval = _best_wave_assignment_for_owners(owners) - return owners, best_waves, best_eval - if strategy not in {"search", "search_with_striped_seed"}: - raise ValueError( - "Unsupported planner_assignment_strategy=" - f"{config.planner_assignment_strategy!r}." - ) - contiguous_owners = _contiguous_chunk_assignment( q_weights=q_weights, cp_size=cp_size, ) + if not contiguous_owners: + wave_assignment = _wave_assignment(chunk_count=len(chunk_ranges), wave_count=1) + return ( + contiguous_owners, + wave_assignment, + _evaluate_candidate( + owners=contiguous_owners, + wave_assignment=wave_assignment, + ), + ) + for wave_count in wave_count_candidates: wave_assignment = _wave_assignment( chunk_count=len(chunk_ranges), wave_count=wave_count, ) - initial_candidates = [ - initial_owners - for initial_owners in (contiguous_owners,) - if initial_owners - if _assignment_uses_all_ranks(initial_owners, cp_size=cp_size) - ] - if not initial_candidates: - continue - current_owners = min( - initial_candidates, - key=lambda owners: float( - _evaluate_candidate(owners=owners, wave_assignment=wave_assignment)[ - "score" - ] - ), - ) + current_owners = contiguous_owners current_eval = _evaluate_candidate( owners=current_owners, wave_assignment=wave_assignment, ) - if cp_size >= 8: - search_steps_remaining = 0 - else: - search_steps_remaining = int(config.planner_max_search_steps) + search_steps_remaining = ( + 0 if cp_size >= 8 else int(config.planner_max_search_steps) + ) if cp_size == 4 and search_steps_remaining > 0: probe_move = _best_improving_move( current_owners=current_owners, @@ -1725,27 +1157,13 @@ def _best_wave_assignment_for_owners( break current_owners, current_eval = best_move - if best_eval is None or float(current_eval["score"]) + 1e-9 < float( - best_eval["score"] + if best is None or float(current_eval["score"]) + 1e-9 < float( + best[2]["score"] ): - best_owners = current_owners - best_waves = wave_assignment - best_eval = current_eval - - if best_eval is None: - best_owners = _contiguous_chunk_assignment(q_weights=q_weights, cp_size=cp_size) - best_waves = _wave_assignment(chunk_count=len(chunk_ranges), wave_count=1) - best_eval = _evaluate_candidate( - owners=best_owners, - wave_assignment=best_waves, - ) - return best_owners, best_waves, best_eval - - -def _concatenate_peer_ranges( - ranges_by_peer: list[tuple[TokenRange, ...]] | tuple[tuple[TokenRange, ...], ...], -) -> tuple[tuple[TokenRange, ...], ...]: - return tuple(tuple(ranges) for ranges in ranges_by_peer) + best = (current_owners, wave_assignment, current_eval) + if best is None: + raise RuntimeError("Failed to evaluate any CP planner wave assignment.") + return best def _flatten_ranges_by_peer( @@ -1965,16 +1383,8 @@ def _build_rank_runtime_plan( _remap_subrange(range_, host_local_ranges) for range_ in local_global_k_ranges ), - kv_fetch_plan=KvFetchPlan( - send_splits=tuple(0 for _ in range(cp_size)), - recv_splits=tuple(0 for _ in range(cp_size)), - send_ranges_by_peer=tuple(tuple() for _ in range(cp_size)), - ), - dkv_reduce_plan=DkvReducePlan( - send_splits=tuple(0 for _ in range(cp_size)), - recv_splits=tuple(0 for _ in range(cp_size)), - recv_ranges_by_peer=tuple(tuple() for _ in range(cp_size)), - ), + kv_fetch_plan=None, + dkv_reduce_plan=None, remote_buffer_range=None, block_size=block_size, ) @@ -2082,14 +1492,8 @@ def _build_rank_runtime_plan( token_layout_index=token_layout_index, local_valid_lengths=(local_token_count,), local_row_ranges=local_row_ranges, - local_token_count=local_token_count, stage_plans=tuple(stage_plans), backward_stage_indices=tuple(backward_stage_indices + [0]), - remote_kv_fetch_plan=KvFetchPlan( - send_splits=aggregate_send_splits, - recv_splits=tuple(aggregate_recv_splits), - send_ranges_by_peer=aggregate_send_ranges, - ), remote_dkv_reduce_plan=DkvReducePlan( send_splits=tuple(aggregate_recv_splits), recv_splits=aggregate_send_splits, @@ -2098,58 +1502,6 @@ def _build_rank_runtime_plan( ) -def make_runtime_key( - spec: PackedBatchAttentionSpec, - *, - topology: ParallelTopology, - config: ContextParallelConfig, -) -> ContextParallelRuntimeKey: - if len(spec.rows) != 1: - raise RuntimeError( - "ART context parallel runtime keys expect exactly one packed sequence, " - f"got {len(spec.rows)} rows." - ) - row_signatures = tuple(_row_signature(row) for row in spec.rows) - return ContextParallelRuntimeKey( - topology=topology, - config=config, - row_signatures=row_signatures, - ) - - -def build_context_parallel_token_layout_index( - *, - group_ids: torch.Tensor, - parent_ids: torch.Tensor, - topology: ParallelTopology, - config: ContextParallelConfig, - original_seq_len: int, -) -> TokenLayoutIndex: - """Return the token ownership chosen by the real CP attention planner.""" - - spec = build_shared_prefix_attention_spec( - group_ids=group_ids, parent_ids=parent_ids - ) - if int(topology.cp) <= 1: - valid_tokens = int(spec.rows[0].valid_tokens) if spec.rows else 0 - return TokenLayoutIndex( - ownership_ranges_by_rank=(((0, valid_tokens, 0),) if valid_tokens else (),), - token_counts_by_rank=(valid_tokens,), - ) - runtime_config = _config_for_runtime_cp(topology=topology, config=config) - _row_spec, chunk_ranges, owners, _wave_assignment = _runtime_plan_assignment( - spec, - topology=topology, - config=runtime_config, - ) - del original_seq_len - return _build_runtime_token_layout_index( - chunk_ranges=chunk_ranges, - owners=owners, - cp_size=max(int(topology.cp), 1), - ) - - def prepare_cp_micro( *, micro: PackedTensors, @@ -2158,6 +1510,7 @@ def prepare_cp_micro( cp_group: Any, cp_rank: int, build_gdn_execution_spec: bool = False, + gdn_planner_config: Any | None = None, trace_token_uids: bool = False, prepare_execution_state: bool = True, block_mask_variants: tuple[CpBlockMaskVariant, ...] = (), @@ -2166,7 +1519,7 @@ def prepare_cp_micro( ) -> PreparedMegatronBatch: """Prepare one CP microbatch with a CPU-only planning phase. - The intended overlap contract is: build the shared-prefix/runtime plan from + The intended overlap contract is: build the prefix-tree/runtime plan from CPU metadata, then materialize local tensors and BlockMasks on `target_device`. Passing CUDA `group_ids` or `parent_ids` still works for older direct callers, but it reintroduces D2H syncs and invalidates the @@ -2179,6 +1532,7 @@ def prepare_cp_micro( cp_group=cp_group, cp_rank=cp_rank, build_gdn_execution_spec=build_gdn_execution_spec, + gdn_planner_config=gdn_planner_config, block_mask_variants=block_mask_variants, target_device=target_device, ) @@ -2193,7 +1547,7 @@ def prepare_cp_micro( ref_logprobs=ref_logprobs, ) if tensors.token_uids is not None: - state = state.model_copy(update={"trace_token_uids": tensors.token_uids}) + state = replace(state, trace_token_uids=tensors.token_uids) if prepare_execution_state: from .executor import prepare_context_parallel_execution_state @@ -2218,6 +1572,7 @@ def prepare_megatron_context_parallel_state( cp_group: Any, cp_rank: int, build_gdn_execution_spec: bool = False, + gdn_planner_config: Any | None = None, block_mask_variants: tuple[CpBlockMaskVariant, ...] = (), target_device: torch.device | None = None, ) -> tuple[ArtContextParallelState, RankRuntimePlan, PackedBatchAttentionSpec, int]: @@ -2226,7 +1581,7 @@ def prepare_megatron_context_parallel_state( This is the portion of CP prepare that must stay free of CUDA reads so the training loop can run it after enqueueing backward for the previous microbatch. If device metadata reaches this function, scalar reads, - cache-key hashing, and shared-prefix parsing can block the host on GPU work. + cache-key hashing, and prefix-tree parsing can block the host on GPU work. """ if int(topology.cp) <= 1: raise RuntimeError( @@ -2245,48 +1600,43 @@ def prepare_megatron_context_parallel_state( group_ids_cpu = _planning_metadata_cpu(micro["group_ids"]) parent_ids_cpu = _planning_metadata_cpu(micro["parent_ids"]) input_pos_cpu = _planning_metadata_cpu(micro["input_pos"]) - runtime_config = _config_for_runtime_cp(topology=topology, config=config) planning_key = _planning_bundle_cache_key( group_ids=group_ids_cpu, parent_ids=parent_ids_cpu, topology=topology, - config=runtime_config, + config=config, original_seq_len=int(micro["tokens"].shape[1]), build_gdn_execution_spec=build_gdn_execution_spec, ) bundle = _PLANNING_BUNDLE_CACHE.get(planning_key) if bundle is None: - spec = build_shared_prefix_attention_spec( + spec = build_prefix_tree_attention_spec( group_ids=group_ids_cpu, parent_ids=parent_ids_cpu, ) - runtime_key = make_runtime_key(spec, topology=topology, config=runtime_config) runtime_plan = get_or_build_runtime_plan( spec, topology=topology, - config=runtime_config, - runtime_key=runtime_key, + config=config, original_seq_len=int(micro["tokens"].shape[1]), ) gdn_execution_spec = None if build_gdn_execution_spec: - from art.megatron.gdn.gdn_shared_prefix import ( - parse_gdn_shared_prefix_segments, + from art.megatron.gdn.gdn_prefix_tree import ( + parse_gdn_prefix_tree_segments, ) - gdn_execution_spec = parse_gdn_shared_prefix_segments( + gdn_execution_spec = parse_gdn_prefix_tree_segments( group_ids_cpu, parent_ids_cpu, - min_completions_per_family=0, ) bundle = _PlanningBundle( spec=spec, - runtime_key=runtime_key, - runtime_plan=runtime_plan, + rank_plans=runtime_plan, gdn_execution_spec=gdn_execution_spec, ) _cache_put(_PLANNING_BUNDLE_CACHE, planning_key, bundle) - rank_plan = bundle.runtime_plan.rank_plans[int(cp_rank)] + rank_plan = bundle.rank_plans[int(cp_rank)] gdn_execution_plan = None if build_gdn_execution_spec: if bundle.gdn_execution_spec is None: @@ -2294,14 +1644,18 @@ def prepare_megatron_context_parallel_state( gdn_plan_device = ( target_device if target_device is not None else micro["tokens"].device ) - rank_gdn_key = _rank_plan_cache_key( - planning_key=planning_key, - device=gdn_plan_device, - cp_rank=int(cp_rank), + rank_gdn_key = ( + planning_key, + gdn_plan_device.type, + gdn_plan_device.index, + int(cp_rank), + _json_cache_key(_dataclass_payload(gdn_planner_config)) + if gdn_planner_config is not None + else "", ) gdn_execution_plan = _GDN_RANK_PLAN_CACHE.get(rank_gdn_key) if gdn_execution_plan is None: - from art.megatron.gdn.gdn_shared_prefix import ( + from art.megatron.gdn.gdn_prefix_tree import ( build_gdn_rank_execution_plan, ) @@ -2311,26 +1665,20 @@ def prepare_megatron_context_parallel_state( cp_rank=int(cp_rank), cp_size=int(topology.cp), attention_token_layout_index=rank_plan.token_layout_index, + planner_config=gdn_planner_config, ) _cache_put(_GDN_RANK_PLAN_CACHE, rank_gdn_key, gdn_execution_plan) - planner_provenance = _planner_provenance( - topology=topology, - config=runtime_config, - warn=int(cp_rank) == 0, - ) pad_multiple = int(topology.tp) if bool(topology.sp) and int(topology.tp) > 1 else 1 state = ArtContextParallelState( - runtime_key=bundle.runtime_key, rank_plan=rank_plan, cp_group=cp_group, - config=runtime_config, + config=config, group_ids=group_ids_cpu[0].contiguous(), parent_ids=parent_ids_cpu[0].contiguous(), input_pos=input_pos_cpu[0].contiguous(), block_mask_variants=block_mask_variants, gdn_execution_spec=bundle.gdn_execution_spec, gdn_execution_plan=gdn_execution_plan, - planner_provenance=planner_provenance, trace_token_uids=None, ) return state, rank_plan, bundle.spec, pad_multiple @@ -2378,115 +1726,43 @@ def dispatch_megatron_context_parallel_training_tensors( if trace_token_uids else None ) - local_tokens = _dispatch_tensor( - micro["tokens"], - rank_plan=rank_plan, - pad_value=0, - pad_multiple=pad_multiple, - dispatch_meta_cache=dispatch_meta_cache, - ) - local_labels = _dispatch_tensor( - labels, - rank_plan=rank_plan, - pad_value=-100, - pad_multiple=pad_multiple, - dispatch_meta_cache=dispatch_meta_cache, - ) - local_input_pos = _dispatch_tensor( - micro["input_pos"], - rank_plan=rank_plan, - pad_value=0, - pad_multiple=pad_multiple, - dispatch_meta_cache=dispatch_meta_cache, - ) - local_assistant_mask = _dispatch_tensor( - assistant_mask, - rank_plan=rank_plan, - pad_value=False, - pad_multiple=pad_multiple, - dispatch_meta_cache=dispatch_meta_cache, - ).to(dtype=torch.bool) - local_group_ids = _dispatch_tensor( - shifted_group_ids, - rank_plan=rank_plan, - pad_value=0, - pad_multiple=pad_multiple, - dispatch_meta_cache=dispatch_meta_cache, - ) - local_old_logprobs = _dispatch_tensor( - old_logprobs, - rank_plan=rank_plan, - pad_value=float("nan"), - pad_multiple=pad_multiple, - dispatch_meta_cache=dispatch_meta_cache, - ) - local_original_logprobs = ( - None - if original_logprobs is None - else _dispatch_tensor( - original_logprobs, - rank_plan=rank_plan, - pad_value=0.0, - pad_multiple=pad_multiple, - dispatch_meta_cache=dispatch_meta_cache, - ) - ) - local_ref_logprobs = ( - None - if ref_logprobs is None - else _dispatch_tensor( - ref_logprobs, + + def dispatch( + tensor: torch.Tensor, + pad_value: int | float | bool, + *, + move_to_target: bool = True, + ) -> torch.Tensor: + local = _dispatch_tensor( + tensor, rank_plan=rank_plan, - pad_value=float("nan"), + pad_value=pad_value, pad_multiple=pad_multiple, dispatch_meta_cache=dispatch_meta_cache, ) - ) - local_advantages = _dispatch_tensor( - advantages, - rank_plan=rank_plan, - pad_value=0.0, - pad_multiple=pad_multiple, - dispatch_meta_cache=dispatch_meta_cache, - ) - local_weights = _dispatch_tensor( - weights, - rank_plan=rank_plan, - pad_value=0.0, - pad_multiple=pad_multiple, - dispatch_meta_cache=dispatch_meta_cache, - ) + return _to_target_device(local, target_device) if move_to_target else local + + def maybe_dispatch( + tensor: torch.Tensor | None, + pad_value: int | float | bool, + ) -> torch.Tensor | None: + return None if tensor is None else dispatch(tensor, pad_value) + local_token_uids = ( - None - if token_uids is None - else _dispatch_tensor( - token_uids, - rank_plan=rank_plan, - pad_value=-1, - pad_multiple=pad_multiple, - dispatch_meta_cache=dispatch_meta_cache, - ) + None if token_uids is None else dispatch(token_uids, -1, move_to_target=False) ) return DispatchedPackedTensors( - tokens=_to_target_device(local_tokens, target_device), - labels=_to_target_device(local_labels, target_device), - input_pos=_to_target_device(local_input_pos, target_device), - assistant_mask=_to_target_device(local_assistant_mask, target_device), - group_ids=_to_target_device(local_group_ids, target_device), - old_logprobs=_to_target_device(local_old_logprobs, target_device), - advantages=_to_target_device(local_advantages, target_device), - weights=_to_target_device(local_weights, target_device), + tokens=dispatch(micro["tokens"], 0), + labels=dispatch(labels, -100), + input_pos=dispatch(micro["input_pos"], 0), + assistant_mask=dispatch(assistant_mask, False).to(dtype=torch.bool), + group_ids=dispatch(shifted_group_ids, 0), + old_logprobs=dispatch(old_logprobs, float("nan")), + advantages=dispatch(advantages, 0.0), + weights=dispatch(weights, 0.0), valid_lengths=rank_plan.local_valid_lengths, - original_logprobs=( - None - if local_original_logprobs is None - else _to_target_device(local_original_logprobs, target_device) - ), - ref_logprobs=( - None - if local_ref_logprobs is None - else _to_target_device(local_ref_logprobs, target_device) - ), + original_logprobs=maybe_dispatch(original_logprobs, 0.0), + ref_logprobs=maybe_dispatch(ref_logprobs, float("nan")), loss_all_reduce_group=cp_group, token_uids=None if local_token_uids is None else local_token_uids.contiguous(), ) @@ -2497,12 +1773,13 @@ def get_or_build_runtime_plan( *, topology: ParallelTopology, config: ContextParallelConfig, - runtime_key: ContextParallelRuntimeKey, original_seq_len: int, -) -> ContextParallelRuntimePlan: - key = ( - _json_cache_key(runtime_key.model_dump(mode="json")), - int(original_seq_len), +) -> tuple[RankRuntimePlan, ...]: + key = _runtime_plan_cache_key( + spec, + topology=topology, + config=config, + original_seq_len=original_seq_len, ) cached = _RUNTIME_PLAN_CACHE.get(key) if cached is not None: @@ -2517,25 +1794,6 @@ def get_or_build_runtime_plan( return plan -def get_or_build_rank_runtime_plan( - spec: PackedBatchAttentionSpec, - *, - topology: ParallelTopology, - config: ContextParallelConfig, - runtime_key: ContextParallelRuntimeKey, - original_seq_len: int, - target_rank: int, -) -> RankRuntimePlan: - del runtime_key - return _build_rank_runtime_plan_for_spec( - spec, - topology=topology, - config=config, - original_seq_len=original_seq_len, - target_rank=target_rank, - ) - - def _runtime_plan_assignment( spec: PackedBatchAttentionSpec, *, @@ -2581,45 +1839,13 @@ def _runtime_plan_assignment( return row_spec, chunk_ranges, owners, wave_assignment -def _build_rank_runtime_plan_for_spec( - spec: PackedBatchAttentionSpec, - *, - topology: ParallelTopology, - config: ContextParallelConfig, - original_seq_len: int, - target_rank: int, -) -> RankRuntimePlan: - row_spec, chunk_ranges, owners, wave_assignment = _runtime_plan_assignment( - spec, - topology=topology, - config=config, - ) - cp_size = max(int(topology.cp), 1) - token_layout_index = _build_runtime_token_layout_index( - chunk_ranges=chunk_ranges, - owners=owners, - cp_size=cp_size, - ) - return _build_rank_runtime_plan( - row_spec=row_spec, - chunk_ranges=chunk_ranges, - owners=owners, - wave_assignment=wave_assignment, - token_layout_index=token_layout_index, - cp_size=cp_size, - original_seq_len=original_seq_len, - target_rank=int(target_rank), - block_size=int(config.block_size), - ) - - def _build_runtime_plan( spec: PackedBatchAttentionSpec, *, topology: ParallelTopology, config: ContextParallelConfig, original_seq_len: int, -) -> ContextParallelRuntimePlan: +) -> tuple[RankRuntimePlan, ...]: row_spec, chunk_ranges, owners, wave_assignment = _runtime_plan_assignment( spec, topology=topology, @@ -2631,7 +1857,7 @@ def _build_runtime_plan( owners=owners, cp_size=cp_size, ) - rank_plans = [ + return tuple( _build_rank_runtime_plan( row_spec=row_spec, chunk_ranges=chunk_ranges, @@ -2644,12 +1870,6 @@ def _build_runtime_plan( block_size=int(config.block_size), ) for rank in range(cp_size) - ] - return ContextParallelRuntimePlan( - topology=topology, - config=config, - token_layout_index=token_layout_index, - rank_plans=tuple(rank_plans), ) @@ -2677,11 +1897,42 @@ def _build_runtime_token_layout_index( def _row_signature(row_spec: PackedRowAttentionSpec) -> str: payload = { "valid_tokens": row_spec.valid_tokens, - "slices": [slice_.model_dump(mode="json") for slice_ in row_spec.slices], + "slices": [_attn_slice_payload(slice_) for slice_ in row_spec.slices], } return json.dumps(payload, sort_keys=True) +def _runtime_plan_cache_key( + spec: PackedBatchAttentionSpec, + *, + topology: ParallelTopology, + config: ContextParallelConfig, + original_seq_len: int, +) -> str: + return _json_cache_key( + { + "topology": _dataclass_payload(topology), + "config": _dataclass_payload(config), + "row_signatures": tuple(_row_signature(row) for row in spec.rows), + "original_seq_len": int(original_seq_len), + } + ) + + +def _dataclass_payload(value: Any) -> dict[str, Any]: + return dict(value.__dict__) + + +def _attn_slice_payload(slice_: AttnSlice) -> dict[str, Any]: + return { + "q_range": _dataclass_payload(slice_.q_range), + "k_range": _dataclass_payload(slice_.k_range), + "mask_kind": slice_.mask_kind.value, + "row_index": slice_.row_index, + "family_index": slice_.family_index, + } + + def _range_key(range_: TokenRange) -> tuple[int, int]: return (int(range_.start), int(range_.end)) @@ -2723,101 +1974,6 @@ def _set_stage_token_indices( current_indices.copy_(source_indices) -def _token_costs(row_spec: PackedRowAttentionSpec) -> list[float]: - costs = [0.0] * row_spec.valid_tokens - for slice_ in row_spec.slices: - q_range = slice_.q_range - k_range = slice_.k_range - if slice_.mask_kind is AttnMaskKind.FULL: - cost = float(k_range.size()) - for q_idx in range(q_range.start, q_range.end): - costs[q_idx] += cost - continue - if q_range.size() != k_range.size(): - raise RuntimeError( - "The current planner only supports causal slices with matched q/k sizes, got " - f"{q_range} vs {k_range}" - ) - for q_idx in range(q_range.start, q_range.end): - costs[q_idx] += float(q_idx - q_range.start + 1) - return costs - - -def _split_row_by_cost( - row_spec: PackedRowAttentionSpec, - *, - cp_size: int, - block_size: int, -) -> tuple[TokenRange | None, ...]: - if cp_size == 1: - return (TokenRange(start=0, end=row_spec.valid_tokens),) - if row_spec.valid_tokens == 0: - return tuple(None for _ in range(cp_size)) - - costs = _token_costs(row_spec) - prefix = [0.0] - for cost in costs: - prefix.append(prefix[-1] + cost) - total_cost = prefix[-1] - boundaries = [0] - block_aligned_split = int(block_size) > 1 and row_spec.valid_tokens >= ( - cp_size * int(block_size) - ) - for split_index in range(1, cp_size): - remaining_ranks = cp_size - split_index - min_boundary = boundaries[-1] - max_boundary = row_spec.valid_tokens - remaining_ranks - if max_boundary <= min_boundary: - boundaries.append(min_boundary) - continue - target = ( - total_cost * split_index / cp_size - if total_cost > 0.0 - else row_spec.valid_tokens * split_index / cp_size - ) - best_boundary = min_boundary + 1 - best_error = float("inf") - candidate_boundaries = range(min_boundary + 1, max_boundary + 1) - if block_aligned_split: - aligned_start = ( - (min_boundary + 1 + block_size - 1) // block_size - ) * block_size - aligned_end = (max_boundary // block_size) * block_size - if aligned_start <= aligned_end: - candidate_boundaries = range(aligned_start, aligned_end + 1, block_size) - for boundary in candidate_boundaries: - current = prefix[boundary] if total_cost > 0.0 else float(boundary) - error = abs(current - target) - if error < best_error: - best_error = error - best_boundary = boundary - boundaries.append(best_boundary) - boundaries.append(row_spec.valid_tokens) - - ranges: list[TokenRange | None] = [] - for start, end in zip(boundaries[:-1], boundaries[1:]): - if end <= start: - ranges.append(None) - else: - ranges.append(TokenRange(start=start, end=end)) - return tuple(ranges) - - -def _intersections( - base_range: TokenRange, - owner_ranges: tuple[TokenRange | None, ...], -) -> list[tuple[int, TokenRange]]: - intersections: list[tuple[int, TokenRange]] = [] - for rank, owner_range in enumerate(owner_ranges): - if owner_range is None: - continue - start = max(base_range.start, owner_range.start) - end = min(base_range.end, owner_range.end) - if end > start: - intersections.append((rank, TokenRange(start=start, end=end))) - return intersections - - def _resolve_stage_mask_kind( *, mask_kind: AttnMaskKind, diff --git a/src/art/megatron/context_parallel/types.py b/src/art/megatron/context_parallel/types.py index e94cf5584..0673101be 100644 --- a/src/art/megatron/context_parallel/types.py +++ b/src/art/megatron/context_parallel/types.py @@ -1,10 +1,11 @@ from __future__ import annotations +from dataclasses import dataclass, field from enum import Enum from typing import Any from megatron.core.packed_seq_params import PackedSeqParams -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict import torch from .layout_index import TokenLayoutIndex @@ -16,9 +17,8 @@ class AttnMaskKind(str, Enum): CAUSAL = "causal" -class TokenRange(BaseModel): - model_config = ConfigDict(frozen=True) - +@dataclass(frozen=True) +class TokenRange: start: int end: int @@ -29,9 +29,8 @@ def is_empty(self) -> bool: return self.end <= self.start -class AttnSlice(BaseModel): - model_config = ConfigDict(frozen=True) - +@dataclass(frozen=True) +class AttnSlice: q_range: TokenRange k_range: TokenRange mask_kind: AttnMaskKind @@ -39,68 +38,25 @@ class AttnSlice(BaseModel): family_index: int | None = None -class PackedRowAttentionSpec(BaseModel): - model_config = ConfigDict(frozen=True) - +@dataclass(frozen=True) +class PackedRowAttentionSpec: row_index: int valid_tokens: int slices: tuple[AttnSlice, ...] -class PackedBatchAttentionSpec(BaseModel): - model_config = ConfigDict(frozen=True) - +@dataclass(frozen=True) +class PackedBatchAttentionSpec: rows: tuple[PackedRowAttentionSpec, ...] -class SharedPrefixBuilderConfig(BaseModel): - model_config = ConfigDict(frozen=True) - - ignore_padding_group_id: int = -1 - require_contiguous_group_runs: bool = True - - -class PlannerCpOverride(BaseModel): - model_config = ConfigDict(frozen=True) - - cp_size: int - block_size: int | None = None - planner_chunk_size: int | None = None - planner_chunk_budget_base: int | None = None - planner_chunk_budget_per_cp_rank: int | None = None - planner_assignment_strategy: str | None = None - planner_stripe_group_size: int | None = None - planner_max_search_steps: int | None = None - planner_candidate_chunk_limit: int | None = None - planner_max_remote_waves: int | None = None - planner_stage_overhead_ms: float | None = None - planner_comm_stage_overhead_ms: float | None = None - planner_interval_overhead_ms: float | None = None - planner_merge_q_token_ms: float | None = None - planner_fetch_token_ms: float | None = None - planner_reduce_token_ms: float | None = None - planner_local_pair_ms: float | None = None - planner_remote_pair_ms: float | None = None - planner_local_backward_pair_ms: float | None = None - planner_remote_backward_pair_ms: float | None = None - planner_remote_stage_token_floor: int | None = None - planner_remote_stage_pair_floor: int | None = None - planner_remote_stage_underfill_ms: float | None = None - planner_tuned_backend: str | None = None - planner_tuned_hardware: str | None = None - planner_tuned_cp_sizes: tuple[int, ...] | None = None - - -class ContextParallelConfig(BaseModel): - model_config = ConfigDict(frozen=True, extra="forbid") - +@dataclass(frozen=True) +class ContextParallelConfig: block_size: int = 128 attention_sparse_block_size: tuple[int, int] | None = None planner_chunk_size: int = 512 planner_chunk_budget_base: int = 128 planner_chunk_budget_per_cp_rank: int = 16 - planner_assignment_strategy: str = "search" - planner_stripe_group_size: int = 16 planner_max_search_steps: int = 8 planner_candidate_chunk_limit: int = 8 planner_max_remote_waves: int = 4 @@ -117,15 +73,10 @@ class ContextParallelConfig(BaseModel): planner_remote_stage_token_floor: int = 4096 planner_remote_stage_pair_floor: int = 4_000_000 planner_remote_stage_underfill_ms: float = 0.287151 - planner_tuned_backend: str | None = "art_context_parallel" - planner_tuned_hardware: str | None = "NVIDIA H200" - planner_tuned_cp_sizes: tuple[int, ...] = (2,) - planner_cp_overrides: tuple[PlannerCpOverride, ...] = () - -class ParallelTopology(BaseModel): - model_config = ConfigDict(frozen=True) +@dataclass(frozen=True) +class ParallelTopology: tp: int = 1 cp: int = 1 dp: int = 1 @@ -133,73 +84,50 @@ class ParallelTopology(BaseModel): sp: bool = False -class ContextParallelRuntimeKey(BaseModel): - model_config = ConfigDict(frozen=True) - - topology: ParallelTopology - config: ContextParallelConfig - row_signatures: tuple[str, ...] - - -class KvFetchPlan(BaseModel): - model_config = ConfigDict(frozen=True) - +@dataclass(frozen=True) +class KvFetchPlan: send_splits: tuple[int, ...] recv_splits: tuple[int, ...] send_ranges_by_peer: tuple[tuple[TokenRange, ...], ...] -class DkvReducePlan(BaseModel): - model_config = ConfigDict(frozen=True) - +@dataclass(frozen=True) +class DkvReducePlan: send_splits: tuple[int, ...] recv_splits: tuple[int, ...] recv_ranges_by_peer: tuple[tuple[TokenRange, ...], ...] -class StagePlan(BaseModel): - model_config = ConfigDict(frozen=True) - +@dataclass(frozen=True) +class StagePlan: stage_index: int source_rank: int - source_ranks: tuple[int, ...] = () is_local_stage: bool - wave_index: int | None = None slices: tuple[AttnSlice, ...] - global_q_ranges: tuple[TokenRange, ...] = () - global_k_ranges: tuple[TokenRange, ...] = () owner_local_q_ranges: tuple[TokenRange, ...] owner_local_k_ranges: tuple[TokenRange, ...] - mask_metadata: "ExactMaskMetadata | None" = None - remote_buffer_range: TokenRange | None = None q_len: int k_len: int + source_ranks: tuple[int, ...] = () + wave_index: int | None = None + global_q_ranges: tuple[TokenRange, ...] = () + global_k_ranges: tuple[TokenRange, ...] = () + mask_metadata: "ExactMaskMetadata | None" = None + remote_buffer_range: TokenRange | None = None kv_fetch_plan: KvFetchPlan | None = None dkv_reduce_plan: DkvReducePlan | None = None -class RankRuntimePlan(BaseModel): - model_config = ConfigDict(frozen=True) - +@dataclass(frozen=True) +class RankRuntimePlan: rank: int original_seq_len: int token_layout_index: TokenLayoutIndex local_valid_lengths: tuple[int, ...] local_row_ranges: tuple[TokenRange | None, ...] - local_token_count: int stage_plans: tuple[StagePlan, ...] - backward_stage_indices: tuple[int, ...] = () - remote_kv_fetch_plan: KvFetchPlan remote_dkv_reduce_plan: DkvReducePlan - - -class ContextParallelRuntimePlan(BaseModel): - model_config = ConfigDict(frozen=True) - - topology: ParallelTopology - config: ContextParallelConfig - token_layout_index: TokenLayoutIndex - rank_plans: tuple[RankRuntimePlan, ...] + backward_stage_indices: tuple[int, ...] = () class DispatchedPackedTensors(ContextParallelLossInputs): @@ -220,54 +148,33 @@ class DispatchedPackedTensors(ContextParallelLossInputs): token_uids: torch.Tensor | None = None -class ContextParallelExecutionCache(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True) - - block_masks: dict[Any, Any] = Field(default_factory=dict) - range_indices: dict[Any, torch.Tensor] = Field(default_factory=dict) - range_meta: dict[Any, tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]] = Field( +@dataclass +class ContextParallelExecutionCache: + block_mask_context: Any | None = None + block_masks: dict[Any, Any] = field(default_factory=dict) + range_indices: dict[Any, torch.Tensor] = field(default_factory=dict) + range_meta: dict[Any, tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]] = field( default_factory=dict ) - stage_execution_specs: dict[Any, "StageExecutionSpec"] = Field(default_factory=dict) + stage_execution_specs: dict[Any, "StageExecutionSpec"] = field(default_factory=dict) -class CpBlockMaskVariant(BaseModel): - model_config = ConfigDict(frozen=True) - - sliding_window: int | None = None +@dataclass(frozen=True) +class CpBlockMaskVariant: block_size: tuple[int, int] + sliding_window: int | None = None -class StageExecutionSpec(BaseModel): - model_config = ConfigDict(frozen=True) - +@dataclass(frozen=True) +class StageExecutionSpec: q_len: int k_len: int compile_key: str mask_metadata: "ExactMaskMetadata | None" = None -class PlannerProvenance(BaseModel): - model_config = ConfigDict(frozen=True) - - runtime_backend: str - runtime_hardware: str | None = None - runtime_cp_size: int - tuned_backend: str | None = None - tuned_hardware: str | None = None - tuned_cp_sizes: tuple[int, ...] = () - backend_match: bool - hardware_match: bool - cp_size_match: bool - using_best_effort: bool - warning_message: str | None = None - warning_emitted: bool = False - - -class ArtContextParallelState(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True) - - runtime_key: ContextParallelRuntimeKey +@dataclass +class ArtContextParallelState: rank_plan: RankRuntimePlan cp_group: Any config: ContextParallelConfig @@ -281,31 +188,28 @@ class ArtContextParallelState(BaseModel): gdn_input_layout: str | None = None gdn_output_layout: str | None = None gdn_attention_original_shape: tuple[int, int, int] | None = None - gdn_attention_original_shapes: dict[int, tuple[int, int, int]] = Field( + gdn_attention_original_shapes: dict[int, tuple[int, int, int]] = field( default_factory=dict ) gdn_attention_token_uids: torch.Tensor | None = None gdn_active_module: Any | None = None - planner_provenance: PlannerProvenance trace_token_uids: torch.Tensor | None = None - execution_cache: ContextParallelExecutionCache = Field( + execution_cache: ContextParallelExecutionCache = field( default_factory=ContextParallelExecutionCache ) -class PreparedMegatronBatch(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True) - +@dataclass +class PreparedMegatronBatch: tensors: DispatchedPackedTensors - packed_seq_params: PackedSeqParams | None = None attention_state: Any + packed_seq_params: PackedSeqParams | None = None rank_plan: RankRuntimePlan | None = None pad_multiple: int = 1 -class FlexMaskSpec(BaseModel): - model_config = ConfigDict(frozen=True) - +@dataclass(frozen=True) +class FlexMaskSpec: q_len: int k_len: int block_size: int | tuple[int, int] @@ -313,9 +217,8 @@ class FlexMaskSpec(BaseModel): exact_mask: "ExactMaskMetadata" -class ExactMaskMetadata(BaseModel): - model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True) - +@dataclass(frozen=True) +class ExactMaskMetadata: q_token_indices: torch.Tensor k_token_indices: torch.Tensor cache_key: str diff --git a/src/art/megatron/dsv4/compressor.py b/src/art/megatron/dsv4/compressor.py index b677f2521..fda1b2bf6 100644 --- a/src/art/megatron/dsv4/compressor.py +++ b/src/art/megatron/dsv4/compressor.py @@ -15,19 +15,33 @@ get_rope_cache_at_positions, ) from art.megatron.dsv4.utils import rotate_activation +from art.megatron.prefix_tree import ( + PrefixTreeRow, + PrefixTreeSegment, + parse_prefix_tree, +) class Dsv4CompressionLayout(NamedTuple): current_indices: torch.Tensor previous_indices: torch.Tensor - entry_group_ids: torch.Tensor - entry_parent_visible: torch.Tensor + entry_group_pre_order: torch.Tensor + entry_group_post_order: torch.Tensor entry_start_positions: torch.Tensor entry_end_positions: torch.Tensor - entry_valid: torch.Tensor + query_group_pre_order: torch.Tensor + + +class _Dsv4CompressionPlan(NamedTuple): + row: PrefixTreeRow + position_to_index_by_group: dict[int, dict[int, int]] + positions_by_group: dict[int, list[int]] + group_pre_order: dict[int, int] + group_post_order: dict[int, int] + query_group_pre_order: list[int] -class Dsv4SharedPrefixState(BaseModel): +class Dsv4PrefixTreeState(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) compression_layouts: dict[int, Dsv4CompressionLayout] @@ -43,232 +57,261 @@ def move_compression_layout_to_device( ) -def build_shared_prefix_compression_layouts( +def build_prefix_tree_compression_layouts( *, position_ids: torch.Tensor, group_ids: torch.Tensor, parent_ids: torch.Tensor, device: torch.device, ) -> dict[int, Dsv4CompressionLayout]: + plan = _build_prefix_tree_compression_plan( + position_ids=position_ids, + group_ids=group_ids, + parent_ids=parent_ids, + ) return { ratio: move_compression_layout_to_device( - build_shared_prefix_compression_layout( - position_ids=position_ids, - group_ids=group_ids, - parent_ids=parent_ids, - ratio=ratio, - duplicate_prompt_entries=ratio == 4, - ), + _emit_prefix_tree_compression_layout(plan=plan, ratio=ratio), device, ) for ratio in (4, 128) } +def _segment_positions( + position_row: torch.Tensor, + segment: PrefixTreeSegment, +) -> list[int]: + positions = [int(value) for value in position_row[segment.start : segment.end]] + if positions != list(range(positions[0], positions[-1] + 1)): + raise ValueError( + "DSV4 prefix-tree compression requires contiguous positions within " + f"group={segment.group_id}, got {positions[:4]}...{positions[-4:]}." + ) + return positions + + +def _segment_path( + row: PrefixTreeRow, + segment: PrefixTreeSegment, +) -> tuple[PrefixTreeSegment, ...]: + by_group = {candidate.group_id: candidate for candidate in row.segments} + return tuple( + by_group[group_id] for group_id in (*segment.ancestors, segment.group_id) + ) + + +def _branch_position_map( + *, + row: PrefixTreeRow, + segment: PrefixTreeSegment, + positions_by_group: dict[int, list[int]], +) -> dict[int, int]: + by_group = {candidate.group_id: candidate for candidate in row.segments} + mapping: dict[int, int] = {} + expected = 0 + for path_segment in _segment_path(row, segment): + positions = positions_by_group[path_segment.group_id] + if positions[0] != expected: + raise ValueError( + "DSV4 prefix-tree compression requires contiguous branch " + f"positions; expected {expected}, got {positions[0]} for " + f"group={path_segment.group_id}." + ) + for offset, logical_pos in enumerate(positions): + mapping[logical_pos] = by_group[path_segment.group_id].start + offset + expected = positions[-1] + 1 + return mapping + + +def _group_preorder_intervals( + row: PrefixTreeRow, +) -> tuple[dict[int, int], dict[int, int]]: + children_by_group: dict[int, list[int]] = { + segment.group_id: [] for segment in row.segments + } + roots: list[int] = [] + for segment in row.segments: + if segment.depth == 0: + roots.append(segment.group_id) + else: + children_by_group[segment.parent_id].append(segment.group_id) + + pre_order: dict[int, int] = {} + post_order: dict[int, int] = {} + next_order = 0 + + def visit(group_id: int) -> None: + nonlocal next_order + pre_order[group_id] = next_order + next_order += 1 + for child_group_id in children_by_group[group_id]: + visit(child_group_id) + post_order[group_id] = next_order + + for root_group_id in roots: + visit(root_group_id) + return pre_order, post_order + + +def _build_prefix_tree_compression_plan( + *, + position_ids: torch.Tensor, + group_ids: torch.Tensor, + parent_ids: torch.Tensor, +) -> _Dsv4CompressionPlan: + if position_ids.ndim != 2 or group_ids.ndim != 2 or parent_ids.ndim != 2: + raise ValueError("DSV4 prefix-tree metadata must be rank-2.") + if int(position_ids.shape[0]) != 1: + raise ValueError( + "DSV4 prefix-tree compression expects one packed sequence per " + f"Megatron microbatch, got batch={int(position_ids.shape[0])}." + ) + if tuple(position_ids.shape) != tuple(group_ids.shape) or tuple( + group_ids.shape + ) != tuple(parent_ids.shape): + raise ValueError( + "DSV4 prefix-tree position/group/parent metadata must share shape, " + f"got {tuple(position_ids.shape)}, {tuple(group_ids.shape)}, " + f"{tuple(parent_ids.shape)}." + ) + + position_cpu = position_ids.detach().cpu() + group_cpu = group_ids.detach().cpu() + parent_cpu = parent_ids.detach().cpu() + (row,) = parse_prefix_tree(group_ids=group_cpu, parent_ids=parent_cpu) + positions_by_group = { + segment.group_id: _segment_positions(position_cpu[0], segment) + for segment in row.segments + } + group_pre_order, group_post_order = _group_preorder_intervals(row) + query_group_pre_order = [-1] * int(position_ids.shape[1]) + for segment in row.segments: + pre_order = group_pre_order[segment.group_id] + for token_index in range(segment.start, segment.end): + query_group_pre_order[token_index] = pre_order + position_to_index_by_group = { + segment.group_id: _branch_position_map( + row=row, + segment=segment, + positions_by_group=positions_by_group, + ) + for segment in row.segments + } + return _Dsv4CompressionPlan( + row=row, + position_to_index_by_group=position_to_index_by_group, + positions_by_group=positions_by_group, + group_pre_order=group_pre_order, + group_post_order=group_post_order, + query_group_pre_order=query_group_pre_order, + ) + + def _logical_window_indices( + position_to_index: dict[int, int], *, - prompt_start: int, - prompt_len: int, - completion_start: int | None, logical_start: int, ratio: int, + allow_negative_padding: bool, ) -> list[int]: indices: list[int] = [] for logical_pos in range(logical_start, logical_start + ratio): - if logical_pos < 0: - indices.append(-1) - elif logical_pos < prompt_len: - indices.append(prompt_start + logical_pos) - elif completion_start is not None: - indices.append(completion_start + logical_pos - prompt_len) - else: - indices.append(-1) + index = position_to_index.get(logical_pos) + if index is None: + if logical_pos < 0 and allow_negative_padding: + indices.append(-1) + continue + raise ValueError( + "DSV4 prefix-tree compression window references missing " + f"logical position {logical_pos}." + ) + indices.append(index) return indices -def build_shared_prefix_compression_layout( +def build_prefix_tree_compression_layout( *, position_ids: torch.Tensor, group_ids: torch.Tensor, parent_ids: torch.Tensor, ratio: int, - duplicate_prompt_entries: bool = False, ) -> Dsv4CompressionLayout: - device = position_ids.device - bsz, seqlen = position_ids.shape - position_cpu = position_ids.detach().cpu() - group_cpu = group_ids.detach().cpu() - parent_cpu = parent_ids.detach().cpu() - current_rows: list[list[list[int]]] = [] - previous_rows: list[list[list[int]]] = [] - group_rows: list[list[int]] = [] - parent_visible_rows: list[list[bool]] = [] - start_rows: list[list[int]] = [] - end_rows: list[list[int]] = [] - - for b in range(bsz): - valid_mask = (group_cpu[b] != -1) & (parent_cpu[b] != -1) - padding = torch.nonzero(~valid_mask, as_tuple=False) - valid_len = int(padding[0].item()) if padding.numel() else seqlen - if bool(valid_mask[valid_len:].any().item()): - raise ValueError("DSV4 shared-prefix metadata must pad only at row end.") - segments: list[tuple[int, int, int, int, int, int]] = [] - cursor = 0 - while cursor < valid_len: - group = int(group_cpu[b, cursor].item()) - parent = int(parent_cpu[b, cursor].item()) - start = cursor - cursor += 1 - while cursor < valid_len and int(group_cpu[b, cursor].item()) == group: - cursor += 1 - end = cursor - start_pos = int(position_cpu[b, start].item()) - end_pos = int(position_cpu[b, end - 1].item()) - segments.append((group, parent, start, end, start_pos, end_pos)) - - row_current: list[list[int]] = [] - row_previous: list[list[int]] = [] - row_groups: list[int] = [] - row_parent_visible: list[bool] = [] - row_starts: list[int] = [] - row_ends: list[int] = [] - - def append_entry( - *, - entry_group: int, - parent_visible: bool, - current_indices: list[int], - previous_indices: list[int], - logical_start: int, - ) -> None: - row_current.append(current_indices) - row_previous.append(previous_indices) - row_groups.append(entry_group) - row_parent_visible.append(parent_visible) - row_starts.append(logical_start) - row_ends.append(logical_start + ratio - 1) - - seg_idx = 0 - while seg_idx < len(segments): - group, parent, prompt_start, prompt_end, _, _ = segments[seg_idx] - if group != parent: - seg_idx += 1 + return _emit_prefix_tree_compression_layout( + plan=_build_prefix_tree_compression_plan( + position_ids=position_ids, + group_ids=group_ids, + parent_ids=parent_ids, + ), + ratio=ratio, + ) + + +def _emit_prefix_tree_compression_layout( + *, + plan: _Dsv4CompressionPlan, + ratio: int, +) -> Dsv4CompressionLayout: + current_rows: list[list[int]] = [] + previous_rows: list[list[int]] = [] + entry_pre_orders: list[int] = [] + entry_post_orders: list[int] = [] + start_rows: list[int] = [] + end_rows: list[int] = [] + + for segment in plan.row.segments: + position_to_index = plan.position_to_index_by_group[segment.group_id] + segment_positions = plan.positions_by_group[segment.group_id] + segment_end_pos = segment_positions[-1] + parent_end_pos = ( + -1 + if not segment.ancestors + else plan.positions_by_group[segment.ancestors[-1]][-1] + ) + usable = ((segment_end_pos + 1) // ratio) * ratio + for logical_start in range(0, usable, ratio): + if logical_start + ratio - 1 <= parent_end_pos: continue - prompt_len = prompt_end - prompt_start - shared_usable = (prompt_len // ratio) * ratio - shared_windows: list[tuple[list[int], list[int], int]] = [] - for logical_start in range(0, shared_usable, ratio): - current_indices = _logical_window_indices( - prompt_start=prompt_start, - prompt_len=prompt_len, - completion_start=None, + current_rows.append( + _logical_window_indices( + position_to_index, logical_start=logical_start, ratio=ratio, + allow_negative_padding=False, ) - previous_indices = _logical_window_indices( - prompt_start=prompt_start, - prompt_len=prompt_len, - completion_start=None, + ) + previous_rows.append( + _logical_window_indices( + position_to_index, logical_start=logical_start - ratio, ratio=ratio, + allow_negative_padding=True, ) - shared_windows.append( - (current_indices, previous_indices, logical_start) - ) - append_entry( - entry_group=group, - parent_visible=not duplicate_prompt_entries, - current_indices=current_indices, - previous_indices=previous_indices, - logical_start=logical_start, - ) - - child_idx = seg_idx + 1 - while child_idx < len(segments): - child_group, child_parent, child_start, _, _, child_end_pos = segments[ - child_idx - ] - if child_group == child_parent or child_parent != group: - break - if duplicate_prompt_entries: - for ( - current_indices, - previous_indices, - logical_start, - ) in shared_windows: - append_entry( - entry_group=child_group, - parent_visible=False, - current_indices=current_indices, - previous_indices=previous_indices, - logical_start=logical_start, - ) - branch_usable = ((child_end_pos + 1) // ratio) * ratio - for logical_start in range(0, branch_usable, ratio): - if logical_start + ratio <= prompt_len: - continue - append_entry( - entry_group=child_group, - parent_visible=False, - current_indices=_logical_window_indices( - prompt_start=prompt_start, - prompt_len=prompt_len, - completion_start=child_start, - logical_start=logical_start, - ratio=ratio, - ), - previous_indices=_logical_window_indices( - prompt_start=prompt_start, - prompt_len=prompt_len, - completion_start=child_start, - logical_start=logical_start - ratio, - ratio=ratio, - ), - logical_start=logical_start, - ) - child_idx += 1 - seg_idx = child_idx - - current_rows.append(row_current) - previous_rows.append(row_previous) - group_rows.append(row_groups) - parent_visible_rows.append(row_parent_visible) - start_rows.append(row_starts) - end_rows.append(row_ends) - - max_entries = max((len(row) for row in current_rows), default=0) - current = torch.full((bsz, max_entries, ratio), -1, dtype=torch.long, device=device) - previous = torch.full_like(current, -1) - entry_groups = torch.full((bsz, max_entries), -1, dtype=torch.long, device=device) - parent_visible = torch.zeros((bsz, max_entries), dtype=torch.bool, device=device) - starts = torch.full_like(entry_groups, -1) - ends = torch.full_like(entry_groups, -1) - valid = torch.zeros((bsz, max_entries), dtype=torch.bool, device=device) - for b, row in enumerate(current_rows): - if not row: - continue - count = len(row) - current[b, :count] = torch.tensor(row, dtype=torch.long, device=device) - previous[b, :count] = torch.tensor( - previous_rows[b], dtype=torch.long, device=device - ) - entry_groups[b, :count] = torch.tensor( - group_rows[b], dtype=torch.long, device=device - ) - parent_visible[b, :count] = torch.tensor( - parent_visible_rows[b], dtype=torch.bool, device=device - ) - starts[b, :count] = torch.tensor(start_rows[b], dtype=torch.long, device=device) - ends[b, :count] = torch.tensor(end_rows[b], dtype=torch.long, device=device) - valid[b, :count] = True + ) + entry_pre_orders.append(plan.group_pre_order[segment.group_id]) + entry_post_orders.append(plan.group_post_order[segment.group_id]) + start_rows.append(logical_start) + end_rows.append(logical_start + ratio - 1) + + entry_count = len(current_rows) + current = torch.empty((entry_count, ratio), dtype=torch.long) + previous = torch.empty_like(current) + if entry_count: + current[:] = torch.tensor(current_rows, dtype=torch.long) + previous[:] = torch.tensor(previous_rows, dtype=torch.long) + entry_pre_order = torch.tensor(entry_pre_orders, dtype=torch.int32) + entry_post_order = torch.tensor(entry_post_orders, dtype=torch.int32) + starts = torch.tensor(start_rows, dtype=torch.long) + ends = torch.tensor(end_rows, dtype=torch.long) + query_pre_order = torch.tensor(plan.query_group_pre_order, dtype=torch.int32) return Dsv4CompressionLayout( current, previous, - entry_groups, - parent_visible, + entry_pre_order, + entry_post_order, starts, ends, - valid, + query_pre_order, ) @@ -276,42 +319,38 @@ def compressed_layout_visibility( layout: Dsv4CompressionLayout, *, position_ids: torch.Tensor, - group_ids: torch.Tensor, - parent_ids: torch.Tensor, q_start: int = 0, q_end: int | None = None, ) -> torch.Tensor: + if int(position_ids.shape[0]) != 1: + raise ValueError( + "DSV4 prefix-tree compressed visibility expects one packed sequence." + ) if q_end is None: q_end = position_ids.shape[1] - q_pos = position_ids[:, q_start:q_end].to(dtype=torch.long).unsqueeze(-1) - q_group = group_ids[:, q_start:q_end].to(dtype=torch.long).unsqueeze(-1) - q_parent = parent_ids[:, q_start:q_end].to(dtype=torch.long).unsqueeze(-1) - entry_group = layout.entry_group_ids.unsqueeze(1) - parent_visible = layout.entry_parent_visible.unsqueeze(1) - return ( - layout.entry_valid.unsqueeze(1) - & (layout.entry_end_positions.unsqueeze(1) <= q_pos) - & ((entry_group == q_group) | (parent_visible & (entry_group == q_parent))) + q_pos = position_ids[0, q_start:q_end].to(dtype=torch.long).unsqueeze(-1) + q_group_pre = layout.query_group_pre_order[q_start:q_end].unsqueeze(-1) + visible = ( + (layout.entry_end_positions.unsqueeze(0) <= q_pos) + & (layout.entry_group_pre_order.unsqueeze(0) <= q_group_pre) + & (q_group_pre < layout.entry_group_post_order.unsqueeze(0)) ) + return visible.unsqueeze(0) def compressed_layout_topk_idxs( layout: Dsv4CompressionLayout, *, position_ids: torch.Tensor, - group_ids: torch.Tensor, - parent_ids: torch.Tensor, offset: int, ) -> torch.Tensor: visibility = compressed_layout_visibility( layout, position_ids=position_ids, - group_ids=group_ids, - parent_ids=parent_ids, ) entry_ids = torch.arange( - layout.entry_group_ids.shape[1], - device=layout.entry_group_ids.device, + layout.entry_group_pre_order.shape[0], + device=layout.entry_group_pre_order.device, dtype=torch.long, ).view(1, 1, -1) return torch.where(visibility, entry_ids + offset, torch.full_like(entry_ids, -1)) @@ -372,7 +411,18 @@ def _gather_projected_tokens( ) -> torch.Tensor: bsz, seqlen, channels = tensor.shape if indices.numel() == 0: + if indices.ndim == 2: + return tensor.new_empty((bsz, *indices.shape, channels)) return tensor.new_empty((*indices.shape, channels)) + if indices.ndim == 2: + if bsz != 1: + raise ValueError( + "DSV4 prefix-tree compression expects one packed sequence per " + f"Megatron microbatch, got batch={bsz}." + ) + safe_indices = indices.clamp(0, max(seqlen - 1, 0)) + gathered = tensor[0].index_select(0, safe_indices.reshape(-1)) + return gathered.view(*indices.shape, channels).unsqueeze(0) batch_offsets = ( torch.arange(bsz, device=tensor.device, dtype=torch.long).view(bsz, 1, 1) * seqlen @@ -534,7 +584,7 @@ def _compress_projected( return kv - def _compress_shared_prefix_projected( + def _compress_prefix_tree_projected( self, kv: torch.Tensor, score: torch.Tensor, @@ -585,11 +635,6 @@ def _compress_shared_prefix_projected( slots_kv = current_kv slots_score = current_score + self.ape.view(1, 1, ratio, -1) - slots_score = torch.where( - layout.entry_valid[:, :, None, None], - slots_score, - torch.zeros_like(slots_score), - ) score_softmax = slots_score.softmax(dim=2, dtype=torch.float32).to( slots_kv.dtype ) @@ -599,11 +644,6 @@ def _compress_shared_prefix_projected( self, position_ids=layout.entry_start_positions, device=kv.device ) apply_rotary_emb(compressed[..., -self.rope_head_dim :], freqs_cis) - compressed = torch.where( - layout.entry_valid.unsqueeze(-1), - compressed, - torch.zeros_like(compressed), - ) if self.rotate: compressed = rotate_activation(compressed) @@ -618,9 +658,7 @@ def forward_raw( ) -> torch.Tensor: kv, score = self._project_raw(x) if shared_layout is not None: - return self._compress_shared_prefix_projected( - kv, score, layout=shared_layout - ) + return self._compress_prefix_tree_projected(kv, score, layout=shared_layout) usable = (x.shape[1] // self.compress_ratio) * self.compress_ratio freqs_cis = get_rope_cache(self, seqlen=usable, device=x.device)[ : usable : self.compress_ratio diff --git a/src/art/megatron/dsv4/deepseek_v4.py b/src/art/megatron/dsv4/deepseek_v4.py index 85d9d313a..1ade6a202 100644 --- a/src/art/megatron/dsv4/deepseek_v4.py +++ b/src/art/megatron/dsv4/deepseek_v4.py @@ -33,7 +33,7 @@ from art.megatron.dsv4.compressor import ( DeepSeekV4Compressor, Dsv4CompressionLayout, - Dsv4SharedPrefixState, + Dsv4PrefixTreeState, compressed_layout_topk_idxs, ) from art.megatron.dsv4.kernel.tilelang_sparse_mla import sparse_attn_tilelang @@ -44,6 +44,7 @@ get_rope_cache_at_positions, ) from art.megatron.dsv4.v4_indexer import V4Indexer +from art.megatron.prefix_tree import PrefixTreeRow, PrefixTreeSegment, parse_prefix_tree def _window_topk_idxs( @@ -73,7 +74,7 @@ def _compress_topk_idxs( return compress.unsqueeze(0).expand(bsz, -1, -1).to(torch.int32) -def _shared_prefix_tensors( +def _prefix_tree_tensors( attention_bias: Any, *, bsz: int, @@ -88,14 +89,14 @@ def _shared_prefix_tensors( parent_ids = parent_ids.to(device=device, dtype=torch.long) if group_ids.shape != (bsz, seqlen) or parent_ids.shape != (bsz, seqlen): raise ValueError( - "DSV4 shared-prefix metadata must match attention input shape: " + "DSV4 prefix-tree metadata must match attention input shape: " f"group_ids={tuple(group_ids.shape)} parent_ids={tuple(parent_ids.shape)} " f"expected={(bsz, seqlen)}." ) return group_ids, parent_ids -def _shared_prefix_window_topk_idxs( +def _prefix_tree_window_topk_idxs( position_ids: torch.Tensor, group_ids: torch.Tensor, parent_ids: torch.Tensor, @@ -104,50 +105,82 @@ def _shared_prefix_window_topk_idxs( ) -> torch.Tensor: bsz, seqlen = group_ids.shape width = min(seqlen, window_size) - arange = torch.arange(seqlen, device=group_ids.device).expand(bsz, -1) - valid_token = (group_ids != -1) & (parent_ids != -1) - group_start = valid_token.clone() - group_start[:, 1:] &= group_ids[:, 1:] != group_ids[:, :-1] - start_values = torch.where(group_start, arange, torch.zeros_like(arange)) - group_start_idx = torch.cummax(start_values, dim=1).values - prompt_start = group_start & (group_ids == parent_ids) - prompt_values = torch.where(prompt_start, arange, torch.zeros_like(arange)) - prompt_start_idx = torch.cummax(prompt_values, dim=1).values - - q_pos = position_ids.to(device=group_ids.device, dtype=torch.long) - group_start_pos = q_pos.gather(1, group_start_idx) - offsets = torch.arange(width, device=group_ids.device) - target_pos = q_pos.unsqueeze(-1) - width + 1 + offsets - from_group = target_pos >= group_start_pos.unsqueeze(-1) - group_k = group_start_idx.unsqueeze(-1) + target_pos - group_start_pos.unsqueeze(-1) - prompt_k = prompt_start_idx.unsqueeze(-1) + target_pos - k_idx = torch.where(from_group, group_k, prompt_k) - valid = ( - valid_token.unsqueeze(-1) - & (target_pos >= 0) - & (target_pos <= q_pos.unsqueeze(-1)) - & (k_idx >= 0) - & (k_idx < seqlen) + position_cpu = position_ids.detach().cpu() + rows = parse_prefix_tree( + group_ids=group_ids.detach().cpu(), + parent_ids=parent_ids.detach().cpu(), ) - return torch.where(valid, k_idx, torch.full_like(k_idx, -1)).to(torch.int32) + topk = torch.full((bsz, seqlen, width), -1, dtype=torch.int32) + if width == 0: + return topk.to(device=group_ids.device) + + def branch_indices( + row: PrefixTreeRow, + segment: PrefixTreeSegment, + by_group: dict[int, PrefixTreeSegment], + positions_by_group: dict[int, torch.Tensor], + ) -> torch.Tensor: + branch_len = int(positions_by_group[segment.group_id][-1].item()) + 1 + indices = torch.full((branch_len,), -1, dtype=torch.int32) + for group_id in (*segment.ancestors, segment.group_id): + path_segment = by_group[group_id] + logical_positions = positions_by_group[group_id] + physical_indices = torch.arange( + path_segment.start, + path_segment.end, + dtype=torch.int32, + ) + indices[logical_positions] = physical_indices + if bool((indices < 0).any().item()): + raise ValueError( + "DSV4 prefix-tree SWA path must cover every logical position " + f"from 0 through {branch_len - 1}." + ) + return indices + + for row in rows: + by_group = {candidate.group_id: candidate for candidate in row.segments} + positions_by_group = { + candidate.group_id: position_cpu[ + row.row_index, candidate.start : candidate.end + ].to(torch.long) + for candidate in row.segments + } + for segment in row.segments: + indices = branch_indices(row, segment, by_group, positions_by_group) + padded = torch.cat( + ( + torch.full((width - 1,), -1, dtype=torch.int32), + indices, + ) + ) + windows = padded.unfold(0, width, 1) + q_positions = position_cpu[row.row_index, segment.start : segment.end].to( + torch.long + ) + topk[row.row_index, segment.start : segment.end] = windows.index_select( + 0, + q_positions, + ) + return topk.to(device=group_ids.device) -def _dsv4_shared_prefix_state(attention_bias: Any) -> Dsv4SharedPrefixState: +def _dsv4_prefix_tree_state(attention_bias: Any) -> Dsv4PrefixTreeState: model_state = getattr(attention_bias, "model_state", None) state = model_state.get("dsv4") if isinstance(model_state, dict) else None - if not isinstance(state, Dsv4SharedPrefixState): + if not isinstance(state, Dsv4PrefixTreeState): raise RuntimeError( - "DSV4 shared-prefix state is missing. Build it once per packed " - "sequence through the model-support shared-prefix state hook." + "DSV4 prefix-tree state is missing. Build it once per packed " + "sequence through the model-support prefix-tree state hook." ) return state def _dsv4_topk_cache(attention_bias: Any) -> dict[Any, Any]: - return _dsv4_shared_prefix_state(attention_bias).topk_idx_cache + return _dsv4_prefix_tree_state(attention_bias).topk_idx_cache -def _shared_prefix_window_topk_idxs_cached( +def _prefix_tree_window_topk_idxs_cached( attention_bias: Any, position_ids: torch.Tensor, group_ids: torch.Tensor, @@ -159,7 +192,7 @@ def _shared_prefix_window_topk_idxs_cached( key = ("raw_swa", int(window_size)) cached = cache.get(key) if cached is None: - cached = _shared_prefix_window_topk_idxs( + cached = _prefix_tree_window_topk_idxs( position_ids, group_ids, parent_ids, @@ -169,13 +202,11 @@ def _shared_prefix_window_topk_idxs_cached( return cached -def _shared_prefix_compressed_topk_idxs_cached( +def _prefix_tree_compressed_topk_idxs_cached( attention_bias: Any, layout: Dsv4CompressionLayout, *, position_ids: torch.Tensor, - group_ids: torch.Tensor, - parent_ids: torch.Tensor, ratio: int, offset: int, ) -> torch.Tensor: @@ -187,8 +218,6 @@ def _shared_prefix_compressed_topk_idxs_cached( compressed_layout_topk_idxs( layout, position_ids=position_ids, - group_ids=group_ids, - parent_ids=parent_ids, offset=offset, ) .to(torch.int32) @@ -198,26 +227,20 @@ def _shared_prefix_compressed_topk_idxs_cached( return cached -def _shared_prefix_i32_metadata_cached( +def _prefix_tree_i32_metadata_cached( attention_bias: Any, position_ids: torch.Tensor, - group_ids: torch.Tensor, - parent_ids: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> torch.Tensor: cache = _dsv4_topk_cache(attention_bias) key = "indexer_q_metadata_i32" cached = cache.get(key) if cached is None: - cached = ( - position_ids.to(torch.int32).contiguous(), - group_ids.to(torch.int32).contiguous(), - parent_ids.to(torch.int32).contiguous(), - ) + cached = position_ids.to(torch.int32).contiguous() cache[key] = cached return cached -def _shared_layout_indexer_metadata_cached( +def _prefix_tree_layout_indexer_metadata_cached( attention_bias: Any, layout: Dsv4CompressionLayout, *, @@ -228,26 +251,26 @@ def _shared_layout_indexer_metadata_cached( cached = cache.get(key) if cached is None: cached = ( - layout.entry_group_ids.to(torch.int32).contiguous(), - layout.entry_parent_visible.to(torch.int32).contiguous(), + layout.query_group_pre_order.to(torch.int32).contiguous(), + layout.entry_group_pre_order.to(torch.int32).contiguous(), + layout.entry_group_post_order.to(torch.int32).contiguous(), layout.entry_end_positions.to(torch.int32).contiguous(), - layout.entry_valid.to(torch.int32).contiguous(), ) cache[key] = cached return cached -def _shared_prefix_compression_layout( +def _prefix_tree_compression_layout( attention_bias: Any, *, ratio: int, ) -> Dsv4CompressionLayout: - layouts = _dsv4_shared_prefix_state(attention_bias).compression_layouts + layouts = _dsv4_prefix_tree_state(attention_bias).compression_layouts if ratio not in layouts: raise RuntimeError( - "DSV4 shared-prefix compression layout was not prepared on the " + "DSV4 prefix-tree compression layout was not prepared on the " f"attention state for ratio={ratio}. Build it once per packed " - "sequence through the model-support shared-prefix state hook." + "sequence through the model-support prefix-tree state hook." ) layout = layouts[ratio] if not isinstance(layout, Dsv4CompressionLayout): @@ -486,16 +509,12 @@ def forward( win = self.window_size ratio = self.compress_ratio rd = self.rope_head_dim - shared_prefix = _shared_prefix_tensors( + prefix_tree = _prefix_tree_tensors( attention_bias, bsz=bsz, seqlen=seqlen_local, device=x.device ) shared_layout: Dsv4CompressionLayout | None = None - if ( - self.compress_ratio - and shared_prefix is not None - and position_ids is not None - ): - shared_layout = _shared_prefix_compression_layout( + if self.compress_ratio and prefix_tree is not None and position_ids is not None: + shared_layout = _prefix_tree_compression_layout( attention_bias, ratio=ratio, ) @@ -519,9 +538,9 @@ def forward( seqlen_global = seqlen_local q_positions = torch.arange(seqlen_local, device=x.device) - if shared_prefix is not None and position_ids is not None: - group_ids, parent_ids = shared_prefix - topk_idxs = _shared_prefix_window_topk_idxs_cached( + if prefix_tree is not None and position_ids is not None: + group_ids, parent_ids = prefix_tree + topk_idxs = _prefix_tree_window_topk_idxs_cached( attention_bias, position_ids, group_ids, @@ -544,32 +563,27 @@ def forward( qr_sbd, group=self.tp_group ) if isinstance(self.indexer, V4Indexer): - group_ids = parent_ids = None shared_layout_i32 = None index_position_ids = position_ids - if shared_prefix is not None: - group_ids, parent_ids = shared_prefix - if position_ids is not None and shared_layout is not None: - index_position_ids, group_ids, parent_ids = ( - _shared_prefix_i32_metadata_cached( - attention_bias, - position_ids, - group_ids, - parent_ids, - ) - ) - shared_layout_i32 = _shared_layout_indexer_metadata_cached( - attention_bias, - shared_layout, - ratio=ratio, - ) + if ( + prefix_tree is not None + and position_ids is not None + and shared_layout is not None + ): + index_position_ids = _prefix_tree_i32_metadata_cached( + attention_bias, + position_ids, + ) + shared_layout_i32 = _prefix_tree_layout_indexer_metadata_cached( + attention_bias, + shared_layout, + ratio=ratio, + ) compress_topk_idxs = self.indexer( x_sbd, qr_sbd, position_ids=index_position_ids, shared_layout=shared_layout, - group_ids=group_ids, - parent_ids=parent_ids, shared_layout_i32=shared_layout_i32, ) else: @@ -594,22 +608,15 @@ def forward( compress_topk_idxs + kv_compress_offset, ) else: - if ( - shared_layout is None - or shared_prefix is None - or position_ids is None - ): + if shared_layout is None or prefix_tree is None or position_ids is None: compress_topk_idxs = _compress_topk_idxs( q_positions, ratio=ratio, bsz=bsz ) else: - group_ids, parent_ids = shared_prefix - compress_topk_idxs = _shared_prefix_compressed_topk_idxs_cached( + compress_topk_idxs = _prefix_tree_compressed_topk_idxs_cached( attention_bias, shared_layout, position_ids=position_ids, - group_ids=group_ids, - parent_ids=parent_ids, ratio=ratio, offset=kv_compress_offset, ) diff --git a/src/art/megatron/dsv4/kernel/tilelang_indexer_fwd.py b/src/art/megatron/dsv4/kernel/tilelang_indexer_fwd.py index 6d2e4a51c..1f8ff482a 100644 --- a/src/art/megatron/dsv4/kernel/tilelang_indexer_fwd.py +++ b/src/art/megatron/dsv4/kernel/tilelang_indexer_fwd.py @@ -234,7 +234,7 @@ def indexer_topk_interface(q, kv, weights, cu_seqlen_ks, cu_seqlen_ke, topk): tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True, }, ) -def tl_shared_prefix_indexer_fwd_impl( +def tl_prefix_tree_indexer_fwd_impl( heads, index_dim, block_N=256, @@ -256,18 +256,16 @@ def tl_shared_prefix_indexer_fwd_impl( logits_shape = [seq_len, seq_len_kv] @T.prim_func - def tl_shared_prefix_indexer_fwd_kernel( + def tl_prefix_tree_indexer_fwd_kernel( IndexQ: T.Tensor(index_q_shape, dtype), # type: ignore IndexK: T.Tensor(index_k_shape, dtype), # type: ignore Logits: T.Tensor(logits_shape, accum_dtype), # type: ignore Weights: T.Tensor([seq_len, heads], accum_dtype), # type: ignore QPosition: T.Tensor([seq_len], index_dtype), # type: ignore - QGroup: T.Tensor([seq_len], index_dtype), # type: ignore - QParent: T.Tensor([seq_len], index_dtype), # type: ignore - EntryGroup: T.Tensor([seq_len_kv], index_dtype), # type: ignore - EntryParentVisible: T.Tensor([seq_len_kv], index_dtype), # type: ignore + QGroupPreOrder: T.Tensor([seq_len], index_dtype), # type: ignore + EntryGroupPreOrder: T.Tensor([seq_len_kv], index_dtype), # type: ignore + EntryGroupPostOrder: T.Tensor([seq_len_kv], index_dtype), # type: ignore EntryEndPosition: T.Tensor([seq_len_kv], index_dtype), # type: ignore - EntryValid: T.Tensor([seq_len_kv], index_dtype), # type: ignore ): with T.Kernel(T.ceildiv(seq_len, block_Q), threads=threads) as bx: index_q_shared = T.alloc_shared([block_Q * heads, index_dim], dtype) @@ -313,17 +311,11 @@ def tl_shared_prefix_indexer_fwd_kernel( q_i = seq_len_i + bq_i k_i = nbn_i * block_N + bn_i if k_i < seq_len_kv: - entry_group = EntryGroup[k_i] + q_group_pre_order = QGroupPreOrder[q_i] visible = ( - (EntryValid[k_i] != 0) - and (EntryEndPosition[k_i] <= QPosition[q_i]) - and ( - (entry_group == QGroup[q_i]) - or ( - (EntryParentVisible[k_i] != 0) - and (entry_group == QParent[q_i]) - ) - ) + (EntryEndPosition[k_i] <= QPosition[q_i]) + and (EntryGroupPreOrder[k_i] <= q_group_pre_order) + and (q_group_pre_order < EntryGroupPostOrder[k_i]) ) Logits[q_i, k_i] = T.if_then_else( visible, @@ -331,7 +323,7 @@ def tl_shared_prefix_indexer_fwd_kernel( -T.infinity(accum_dtype), ) - return tl_shared_prefix_indexer_fwd_kernel + return tl_prefix_tree_indexer_fwd_kernel def _i32_contiguous(tensor): @@ -340,23 +332,21 @@ def _i32_contiguous(tensor): return tensor.contiguous() -def shared_prefix_indexer_fwd_interface( +def prefix_tree_indexer_fwd_interface( q, kv, weights, position_ids, - group_ids, - parent_ids, - entry_group_ids, - entry_parent_visible, + query_group_pre_order, + entry_group_pre_order, + entry_group_post_order, entry_end_positions, - entry_valid, ): - """Compute shared-prefix-aware indexer logits for one batch element. + """Compute prefix-tree-aware indexer logits for one batch element. The output is intentionally block-local [query_block, compressed_kv] fp32. Callers run topk on this block to avoid materializing full [S, compressed_kv] - logits for long shared-prefix packed sequences. + logits for long prefix-tree packed sequences. """ if q.dtype is torch.float32: q = q.to(torch.bfloat16) @@ -371,13 +361,13 @@ def shared_prefix_indexer_fwd_interface( block_q = max(1, 128 // heads) if seq_len % block_q != 0: raise ValueError( - "DSV4 shared-prefix indexer TileLang query length must be divisible " + "DSV4 prefix-tree indexer TileLang query length must be divisible " f"by {block_q}, got {seq_len}." ) logits = torch.empty([seq_len, seq_len_kv], device=q.device, dtype=torch.float32) with preserve_tilelang_env(): - tl_indexer_fwd_kernel = tl_shared_prefix_indexer_fwd_impl( + tl_indexer_fwd_kernel = tl_prefix_tree_indexer_fwd_impl( heads=heads, index_dim=index_dim, ) @@ -387,42 +377,36 @@ def shared_prefix_indexer_fwd_interface( logits, weights.float(), _i32_contiguous(position_ids), - _i32_contiguous(group_ids), - _i32_contiguous(parent_ids), - _i32_contiguous(entry_group_ids), - _i32_contiguous(entry_parent_visible), + _i32_contiguous(query_group_pre_order), + _i32_contiguous(entry_group_pre_order), + _i32_contiguous(entry_group_post_order), _i32_contiguous(entry_end_positions), - _i32_contiguous(entry_valid), ) return logits -def shared_prefix_indexer_topk_interface( +def prefix_tree_indexer_topk_interface( q, kv, weights, position_ids, - group_ids, - parent_ids, - entry_group_ids, - entry_parent_visible, + query_group_pre_order, + entry_group_pre_order, + entry_group_post_order, entry_end_positions, - entry_valid, topk, ): - """Compute shared-prefix-aware DSV4 indexer top-k compressed ids.""" + """Compute prefix-tree-aware DSV4 indexer top-k compressed ids.""" return _topk_indices_from_logits( - shared_prefix_indexer_fwd_interface( + prefix_tree_indexer_fwd_interface( q, kv, weights, position_ids, - group_ids, - parent_ids, - entry_group_ids, - entry_parent_visible, + query_group_pre_order, + entry_group_pre_order, + entry_group_post_order, entry_end_positions, - entry_valid, ), topk, ) diff --git a/src/art/megatron/dsv4/kernel/tilelang_sparse_mla_bwd.py b/src/art/megatron/dsv4/kernel/tilelang_sparse_mla_bwd.py index 0d14e5e4a..a9c2c743c 100644 --- a/src/art/megatron/dsv4/kernel/tilelang_sparse_mla_bwd.py +++ b/src/art/megatron/dsv4/kernel/tilelang_sparse_mla_bwd.py @@ -23,8 +23,6 @@ @tilelang.jit(out_idx=[-1]) def preprocess( - B, - S, H, D, block_ND=32, @@ -34,6 +32,8 @@ def preprocess( ): assert dtype == T.bfloat16 assert accum_dtype == T.float32 + B = T.dynamic("batch") + S = T.dynamic("seq_len") shape = [B, S, H, D] @T.prim_func @@ -77,8 +77,6 @@ def preprocess_kernel( @tilelang.jit(out_idx=[-1]) def postprocess( - B, - S_kv, D, block_N=64, threads=128, @@ -87,6 +85,8 @@ def postprocess( ): assert dtype == T.bfloat16 assert accum_dtype == T.float32 + B = T.dynamic("batch") + S_kv = T.dynamic("seq_len_kv") dkv_shape = [B, S_kv, D] @T.prim_func @@ -112,9 +112,6 @@ def postprocess_kernel( }, ) def bwd( - B, - S, - S_kv, H, D, topk, @@ -131,6 +128,9 @@ def bwd( ) assert dtype == T.bfloat16 assert accum_dtype == T.float32 + B = T.dynamic("batch") + S = T.dynamic("seq_len") + S_kv = T.dynamic("seq_len_kv") if sm_scale is None: sm_scale = D ** (-0.5) @@ -356,17 +356,16 @@ def sparse_mqa_bwd_interface(q, kv, attn_sink, o, do, topk_idxs, lse, sm_scale=N topk = padded_topk with preserve_tilelang_env(): - preprocess_kernel = preprocess(B, S, H, D, dtype=dtype) - postprocess_kernel = postprocess(B, S_kv, D, dtype=dtype) + # Keep sequence lengths dynamic so changing packed workloads reuse the + # same generated kernels. Model/tile dimensions remain static. + preprocess_kernel = preprocess(H, D, dtype=dtype) + postprocess_kernel = postprocess(D, dtype=dtype) delta = preprocess_kernel(o, do) dkv = torch.zeros_like(kv, dtype=torch.float32) d_attn_sink = torch.zeros_like(attn_sink) if topk <= block_size: with preserve_tilelang_env(): bwd_kernel = bwd( - B, - S, - S_kv, H, D, topk, @@ -382,9 +381,6 @@ def sparse_mqa_bwd_interface(q, kv, attn_sink, o, do, topk_idxs, lse, sm_scale=N chunk_count = topk // block_size with preserve_tilelang_env(): bwd_kernel = bwd( - B, - S, - S_kv, H, D, block_size, diff --git a/src/art/megatron/dsv4/v4_indexer.py b/src/art/megatron/dsv4/v4_indexer.py index 6a90a055e..cacfacd66 100644 --- a/src/art/megatron/dsv4/v4_indexer.py +++ b/src/art/megatron/dsv4/v4_indexer.py @@ -15,7 +15,7 @@ ) from art.megatron.dsv4.kernel.tilelang_indexer_fwd import ( indexer_topk_interface, - shared_prefix_indexer_topk_interface, + prefix_tree_indexer_topk_interface, ) from art.megatron.dsv4.rope import ( apply_rotary_emb, @@ -48,8 +48,7 @@ def _exact_indexer_topk( *, shared_layout: Dsv4CompressionLayout | None = None, position_ids: torch.Tensor | None = None, - group_ids: torch.Tensor | None = None, - parent_ids: torch.Tensor | None = None, + q_offset: int = 0, ) -> torch.Tensor: """Return frozen DSV4 indexer topk using the reference score equation. @@ -94,17 +93,15 @@ def _exact_indexer_topk( kv_ids.unsqueeze(0) >= cu_seqlen_ks[q_start:q_end, None] ) & (kv_ids.unsqueeze(0) < cu_seqlen_ke[q_start:q_end, None]) else: - if position_ids is None or group_ids is None or parent_ids is None: + if position_ids is None: raise ValueError( - "DSV4 shared-prefix indexer requires position/group metadata." + "DSV4 prefix-tree indexer requires position ids." ) visible = compressed_layout_visibility( shared_layout, - position_ids=position_ids[b : b + 1], - group_ids=group_ids[b : b + 1], - parent_ids=parent_ids[b : b + 1], - q_start=q_start, - q_end=q_end, + position_ids=position_ids, + q_start=q_offset + q_start, + q_end=q_offset + q_end, )[0] scores.masked_fill_(~visible, float("-inf")) top_scores, top_indices = scores.topk(actual_topk, dim=-1) @@ -127,8 +124,6 @@ def _tilelang_indexer_topk( *, shared_layout: Dsv4CompressionLayout | None = None, position_ids: torch.Tensor | None = None, - group_ids: torch.Tensor | None = None, - parent_ids: torch.Tensor | None = None, shared_layout_i32: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: @@ -147,33 +142,38 @@ def _tilelang_indexer_topk( with torch.no_grad(): for b in range(batch): if shared_layout is not None: - if position_ids is None or group_ids is None or parent_ids is None: + if batch != 1: raise ValueError( - "DSV4 shared-prefix indexer requires position/group metadata." + "DSV4 prefix-tree indexer expects one packed sequence per " + f"Megatron microbatch, got batch={batch}." ) + if position_ids is None: + raise ValueError("DSV4 prefix-tree indexer requires position ids.") if shared_layout_i32 is None: - entry_group_ids = shared_layout.entry_group_ids.to(torch.int32) - entry_parent_visible = shared_layout.entry_parent_visible.to( + query_group_pre_order = shared_layout.query_group_pre_order.to( + torch.int32 + ) + entry_group_pre_order = shared_layout.entry_group_pre_order.to( + torch.int32 + ) + entry_group_post_order = shared_layout.entry_group_post_order.to( torch.int32 ) entry_end_positions = shared_layout.entry_end_positions.to( torch.int32 ) - entry_valid = shared_layout.entry_valid.to(torch.int32) else: ( - entry_group_ids, - entry_parent_visible, + query_group_pre_order, + entry_group_pre_order, + entry_group_post_order, entry_end_positions, - entry_valid, ) = shared_layout_i32 position_b = position_ids[b].to(torch.int32).contiguous() - group_b = group_ids[b].to(torch.int32).contiguous() - parent_b = parent_ids[b].to(torch.int32).contiguous() - entry_group_b = entry_group_ids[b].contiguous() - entry_parent_visible_b = entry_parent_visible[b].contiguous() - entry_end_b = entry_end_positions[b].contiguous() - entry_valid_b = entry_valid[b].contiguous() + query_group_pre_b = query_group_pre_order.contiguous() + entry_group_pre_b = entry_group_pre_order.contiguous() + entry_group_post_b = entry_group_post_order.contiguous() + entry_end_b = entry_end_positions.contiguous() for q_start in range(0, fast_seqlen, q_block): q_end = min(q_start + q_block, fast_seqlen) q_slice = q[q_start:q_end, b].contiguous() @@ -189,17 +189,15 @@ def _tilelang_indexer_topk( topk, ) else: - out[b, q_start:q_end] = shared_prefix_indexer_topk_interface( + out[b, q_start:q_end] = prefix_tree_indexer_topk_interface( q_slice, k_slice, weights_slice, position_b[q_start:q_end], - group_b[q_start:q_end], - parent_b[q_start:q_end], - entry_group_b, - entry_parent_visible_b, + query_group_pre_b[q_start:q_end], + entry_group_pre_b, + entry_group_post_b, entry_end_b, - entry_valid_b, topk, ) if fast_seqlen != seqlen: @@ -211,11 +209,8 @@ def _tilelang_indexer_topk( cu_seqlen_ke[fast_seqlen:], topk, shared_layout=shared_layout, - position_ids=None - if position_ids is None - else position_ids[:, fast_seqlen:], - group_ids=None if group_ids is None else group_ids[:, fast_seqlen:], - parent_ids=None if parent_ids is None else parent_ids[:, fast_seqlen:], + position_ids=position_ids, + q_offset=fast_seqlen, ) return out @@ -292,8 +287,6 @@ def forward( packed_seq_params=None, position_ids: torch.Tensor | None = None, shared_layout: Dsv4CompressionLayout | None = None, - group_ids: torch.Tensor | None = None, - parent_ids: torch.Tensor | None = None, shared_layout_i32: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ): @@ -372,7 +365,5 @@ def forward( self.index_topk, shared_layout=shared_layout, position_ids=position_ids, - group_ids=group_ids, - parent_ids=parent_ids, shared_layout_i32=shared_layout_i32, ) diff --git a/src/art/megatron/flex_attn/attention.py b/src/art/megatron/flex_attn/attention.py index cf49df713..b284813f6 100644 --- a/src/art/megatron/flex_attn/attention.py +++ b/src/art/megatron/flex_attn/attention.py @@ -20,8 +20,8 @@ ) -class SharedPrefixAttentionState(BaseModel): - """Shared-prefix sparsity metadata for one packed ART training sample.""" +class PrefixTreeAttentionState(BaseModel): + """Prefix-tree sparsity metadata for one packed ART training sample.""" model_config = ConfigDict(arbitrary_types_allowed=True) block_mask: BlockMask @@ -124,14 +124,14 @@ def _configure_softmax_offset( raise ValueError(f"Unsupported softmax_type: {config.softmax_type}") -def create_shared_prefix_attention_state( +def create_prefix_tree_attention_state( group_ids: Tensor, parent_ids: Tensor, *, input_pos: Tensor | None = None, sliding_windows: tuple[int, ...] = (), -) -> SharedPrefixAttentionState: - """Build a compiled block mask for ART shared-prefix packing. +) -> PrefixTreeAttentionState: + """Build a compiled block mask for ART prefix-tree packing. Initialized on the device of the group_ids tensor. @@ -140,9 +140,9 @@ def create_shared_prefix_attention_state( parent_ids: `[B, S]` parent group id for each token in a packed sequence. """ - from art.megatron.shared_prefix_state import create_shared_prefix_state + from art.megatron.prefix_tree_state import create_prefix_tree_state - return create_shared_prefix_state( + return create_prefix_tree_state( group_ids, parent_ids, input_pos=input_pos, @@ -220,7 +220,7 @@ def forward( key: `[S, B, Hkv, D]` value: `[S, B, Hkv, D]` attention_mask: unused placeholder tensor kept for Megatron checkpoint API. - attention_bias: `SharedPrefixAttentionState` or `BlockMask`. + attention_bias: `PrefixTreeAttentionState` or `BlockMask`. """ del attention_mask, attn_mask_type @@ -228,7 +228,7 @@ def forward( "PackedSeqParams is not used in ART Megatron flex path." ) - if isinstance(attention_bias, SharedPrefixAttentionState): + if isinstance(attention_bias, PrefixTreeAttentionState): block_mask = attention_bias.block_mask_for_window( getattr(self, "art_sliding_window", None) ) diff --git a/src/art/megatron/gdn/__init__.py b/src/art/megatron/gdn/__init__.py index cd3a0873a..a62769edb 100644 --- a/src/art/megatron/gdn/__init__.py +++ b/src/art/megatron/gdn/__init__.py @@ -1,17 +1,15 @@ """ART helpers for Megatron GatedDeltaNet integration.""" from .fla_cp import chunk_gated_delta_rule_native_cp -from .gdn_shared_prefix import ( +from .gdn_prefix_tree import ( GdnPackedExecutionSpec, - GdnPackedFamilySpec, GdnPlannerConfig, GdnRankExecutionPlan, GdnSegmentBucketPlan, GdnSegmentSpec, - build_gdn_cp_segment_schedule, build_gdn_rank_execution_plan, move_gdn_rank_execution_plan_to_device, - parse_gdn_shared_prefix_segments, + parse_gdn_prefix_tree_segments, ) from .layout import exchange_rank_tensor_all_to_all from .operator import run_gdn_layer @@ -19,15 +17,13 @@ __all__ = [ "chunk_gated_delta_rule_native_cp", "GdnPackedExecutionSpec", - "GdnPackedFamilySpec", "GdnPlannerConfig", "GdnRankExecutionPlan", "GdnSegmentSpec", "GdnSegmentBucketPlan", - "build_gdn_cp_segment_schedule", "build_gdn_rank_execution_plan", "exchange_rank_tensor_all_to_all", "move_gdn_rank_execution_plan_to_device", - "parse_gdn_shared_prefix_segments", + "parse_gdn_prefix_tree_segments", "run_gdn_layer", ] diff --git a/src/art/megatron/gdn/conv_gelu.py b/src/art/megatron/gdn/conv_gelu.py index 2795665e8..fea7d9d22 100644 --- a/src/art/megatron/gdn/conv_gelu.py +++ b/src/art/megatron/gdn/conv_gelu.py @@ -53,16 +53,18 @@ def _segment_for_token( cu_seqlens, token, SEGMENTS, - SEARCH_STEPS: tl.constexpr, + SEARCH_STEPS, ): lo = tl.zeros(token.shape, dtype=tl.int64) hi = lo + SEGMENTS.to(tl.int64) - 1 - for _ in tl.static_range(0, SEARCH_STEPS): + step = 0 + while step < SEARCH_STEPS: mid = (lo + hi + 1) // 2 mid_start = tl.load(cu_seqlens + mid) take_upper = mid_start <= token lo = tl.where(take_upper, mid, lo) hi = tl.where(take_upper, hi, mid - 1) + step += 1 return lo @@ -73,7 +75,7 @@ def _packed_conv_token_metadata_kernel( token_local_t, TOTAL_TOKENS, SEGMENTS, - SEARCH_STEPS: tl.constexpr, + SEARCH_STEPS, BLOCK_N: tl.constexpr, ): pid_n = tl.program_id(0) diff --git a/src/art/megatron/gdn/fla_cp.py b/src/art/megatron/gdn/fla_cp.py index 3629185c6..bb053018e 100644 --- a/src/art/megatron/gdn/fla_cp.py +++ b/src/art/megatron/gdn/fla_cp.py @@ -187,11 +187,10 @@ def forward( chunk_indices=chunk_indices, ) summary = _fwd_summary(k=k, w=w, u=u, g=g_cumsum, cu_seqlens=cu_seqlens) - gathered_summary = _all_gather_summary(summary, group) + prefix_summary = _prefix_summary_exclusive(summary, group) local_initial = _scan_fwd_initial_state( - gathered_summary, + prefix_summary, initial_state, - rank=dist.get_rank(group), # ty: ignore[possibly-missing-attribute] ) h, v_new, local_final_state = chunk_gated_delta_rule_fwd_h( k=k, @@ -275,11 +274,12 @@ def backward(ctx: Any, *grad_outputs: Tensor | None) -> tuple[Any, ...]: scale=ctx.scale, cu_seqlens=ctx.cu_seqlens, ) - gathered_bwd_summary = _all_gather_summary(bwd_summary, ctx.group) + suffix_summary, full_bwd_summary = _suffix_summary_exclusive_and_full( + bwd_summary, ctx.group + ) local_dht = _scan_bwd_local_final_grad( - gathered_bwd_summary, + suffix_summary, external_dht, - rank=dist.get_rank(ctx.group), # ty: ignore[possibly-missing-attribute] ) dq, dk, dv, db, dg, _dh0, _dA_log, _ddt_bias = chunk_gated_delta_rule_bwd( q=q, @@ -296,7 +296,7 @@ def backward(ctx: Any, *grad_outputs: Tensor | None) -> tuple[Any, ...]: chunk_indices=ctx.chunk_indices, use_exp2=False, ) - dh0 = _scan_bwd_initial_state_grad(gathered_bwd_summary, external_dht) + dh0 = _scan_bwd_initial_state_grad(full_bwd_summary, external_dht) return ( dq.to(q), dk.to(k), @@ -425,25 +425,122 @@ def _bwd_summary( return summary -def _all_gather_summary(summary: Tensor, group: Any) -> Tensor: +def _prefix_summary_exclusive(summary: Tensor, group: Any) -> Tensor | None: + inclusive = _prefix_summary_inclusive(summary, group) + rank = dist.get_rank(group) # ty: ignore[possibly-missing-attribute] world_size = dist.get_world_size(group) # ty: ignore[possibly-missing-attribute] - gathered = torch.empty( - world_size, - *summary.shape, - device=summary.device, - dtype=summary.dtype, + return _exchange_summary( + inclusive, + group=group, + send_to=rank + 1 if rank + 1 < world_size else None, + recv_from=rank - 1 if rank > 0 else None, ) - dist.all_gather_into_tensor( # ty: ignore[possibly-missing-attribute] - gathered, summary.contiguous(), group=group + + +def _prefix_summary_inclusive(summary: Tensor, group: Any) -> Tensor: + rank = dist.get_rank(group) # ty: ignore[possibly-missing-attribute] + world_size = dist.get_world_size(group) # ty: ignore[possibly-missing-attribute] + aggregate = summary.contiguous() + offset = 1 + while offset < world_size: + received = _exchange_summary( + aggregate, + group=group, + send_to=rank + offset if rank + offset < world_size else None, + recv_from=rank - offset if rank >= offset else None, + ) + if received is not None: + aggregate = _compose_summary(aggregate, received) + offset *= 2 + return aggregate + + +def _suffix_summary_exclusive_and_full( + summary: Tensor, group: Any +) -> tuple[Tensor | None, Tensor]: + inclusive = _suffix_summary_inclusive(summary, group) + rank = dist.get_rank(group) # ty: ignore[possibly-missing-attribute] + world_size = dist.get_world_size(group) # ty: ignore[possibly-missing-attribute] + exclusive = _exchange_summary( + inclusive, + group=group, + send_to=rank - 1 if rank > 0 else None, + recv_from=rank + 1 if rank + 1 < world_size else None, ) - return gathered + full = inclusive if rank == 0 else torch.empty_like(summary) + dist.broadcast(full, src=0, group=group) # ty: ignore[possibly-missing-attribute] + return exclusive, full -def _scan_fwd_initial_state(summaries: Tensor, h0: Tensor, *, rank: int) -> Tensor: - multi = summaries.ndim == 5 +def _suffix_summary_inclusive(summary: Tensor, group: Any) -> Tensor: + rank = dist.get_rank(group) # ty: ignore[possibly-missing-attribute] + world_size = dist.get_world_size(group) # ty: ignore[possibly-missing-attribute] + aggregate = summary.contiguous() + offset = 1 + while offset < world_size: + received = _exchange_summary( + aggregate, + group=group, + send_to=rank - offset if rank >= offset else None, + recv_from=rank + offset if rank + offset < world_size else None, + ) + if received is not None: + aggregate = _compose_summary(aggregate, received) + offset *= 2 + return aggregate + + +def _exchange_summary( + summary: Tensor, + *, + group: Any, + send_to: int | None, + recv_from: int | None, +) -> Tensor | None: + world_size = dist.get_world_size(group) # ty: ignore[possibly-missing-attribute] + count = int(summary.numel()) + input_splits = [0] * world_size + output_splits = [0] * world_size + if send_to is None: + send_buffer = summary.new_empty((0,)) + else: + input_splits[int(send_to)] = count + send_buffer = summary.reshape(-1).contiguous() + if recv_from is not None: + output_splits[int(recv_from)] = count + recv_buffer = summary.new_empty((sum(output_splits),)) + dist.all_to_all_single( # ty: ignore[possibly-missing-attribute] + recv_buffer, + send_buffer, + output_split_sizes=output_splits, + input_split_sizes=input_splits, + group=group, + ) + if recv_from is None: + return None + return recv_buffer.view_as(summary) + + +def _compose_summary(after: Tensor, before: Tensor) -> Tensor: + value_dim = int(before.shape[-1]) - int(before.shape[-2]) + before_he = before[..., :value_dim] + before_transition = before[..., value_dim:] + after_he = after[..., :value_dim] + after_transition = after[..., value_dim:] + he = torch.matmul(after_transition.float(), before_he.float()) + after_he.float() + transition = torch.matmul( + after_transition.float(), + before_transition.float(), + ) + return torch.cat((he, transition), dim=-1).contiguous() + + +def _scan_fwd_initial_state(summary: Tensor | None, h0: Tensor) -> Tensor: + if summary is None: + return h0.float() + multi = summary.ndim == 4 state = h0.float() if multi else h0[0].float() - for peer in range(rank): - state = _apply_summary(summaries[peer], state) + state = _apply_summary(summary, state) return state if multi else state.unsqueeze(0) @@ -457,23 +554,21 @@ def _broadcast_chain_final_state(final_state: Tensor | None, group: Any) -> Tens def _scan_bwd_local_final_grad( - summaries: Tensor, + summary: Tensor | None, dht: Tensor, - *, - rank: int, ) -> Tensor: - multi = summaries.ndim == 5 + if summary is None: + return dht.float() + multi = summary.ndim == 4 state = dht.float() if multi else dht[0].float() - for peer in range(int(summaries.shape[0]) - 1, rank, -1): - state = _apply_summary(summaries[peer], state) + state = _apply_summary(summary, state) return state if multi else state.unsqueeze(0) -def _scan_bwd_initial_state_grad(summaries: Tensor, dht: Tensor) -> Tensor: - multi = summaries.ndim == 5 +def _scan_bwd_initial_state_grad(summary: Tensor, dht: Tensor) -> Tensor: + multi = summary.ndim == 4 state = dht.float() if multi else dht[0].float() - for peer in range(int(summaries.shape[0]) - 1, -1, -1): - state = _apply_summary(summaries[peer], state) + state = _apply_summary(summary, state) return state if multi else state.unsqueeze(0) diff --git a/src/art/megatron/gdn/gdn_prefix_tree.py b/src/art/megatron/gdn/gdn_prefix_tree.py new file mode 100644 index 000000000..eeb7a1a1c --- /dev/null +++ b/src/art/megatron/gdn/gdn_prefix_tree.py @@ -0,0 +1,2484 @@ +from __future__ import annotations + +from bisect import bisect_left +from dataclasses import dataclass, replace +from typing import Any, Literal, NamedTuple, cast + +import torch + +from art.megatron.context_parallel.layout_index import TokenLayoutIndex +from art.megatron.prefix_tree import parse_prefix_tree + +GdnSegmentKind = Literal["prefix", "completion"] +# FLA's public chunk_gated_delta_rule hard-codes 64-token WY chunks. +FLA_CHUNK_SIZE = 64 +# Fitted from Qwen3.5 35B/397B GDN lab points. Doubling state elements +# increases exposed recurrent runtime by about 1.7x, not 2x, because the +# kernels recover some parallelism at larger state shapes. Communication terms +# below still use exact bytes moved. +_RUNTIME_STATE_THROUGHPUT_EXPONENT = 0.75 + + +def _dtype_bytes(dtype: Any) -> int: + if dtype is None: + return 2 + if isinstance(dtype, str): + dtype_by_name = { + "float16": torch.float16, + "fp16": torch.float16, + "bfloat16": torch.bfloat16, + "bf16": torch.bfloat16, + "float32": torch.float32, + "fp32": torch.float32, + } + if dtype not in dtype_by_name: + raise ValueError(f"unsupported GDN planner dtype {dtype!r}") + dtype = dtype_by_name[dtype] + if dtype in (torch.float16, torch.bfloat16): + return 2 + if dtype == torch.float32: + return 4 + return torch.tensor([], dtype=dtype).element_size() + + +@dataclass(frozen=True) +class GdnSegmentSpec: + """Contiguous logical GDN segment in one packed row.""" + + row_index: int + family_index: int + group_id: int + parent_id: int + start: int + end: int + kind: GdnSegmentKind + child_index: int | None = None + + @property + def length(self) -> int: + return self.end - self.start + + +@dataclass(frozen=True) +class GdnPackedExecutionSpec: + """Parsed prefix-tree GDN execution metadata for a packed batch.""" + + batch_size: int + sequence_length: int + valid_lengths: tuple[int, ...] + tree_segments: tuple[GdnSegmentSpec, ...] + tree_parent_indices: tuple[int, ...] + tree_depths: tuple[int, ...] + + @property + def family_count(self) -> int: + return len(self.tree_segments) + + @property + def real_token_count(self) -> int: + return sum(self.valid_lengths) + + +@dataclass(frozen=True) +class GdnSegmentBucketPlan: + """Device-local index tensors for a variable-length GDN segment batch.""" + + length: int + lengths: torch.Tensor + lengths_cpu: torch.Tensor + real_mask: torch.Tensor + cu_seqlens: torch.Tensor + cu_seqlens_cpu: torch.Tensor + row_indices: torch.Tensor + position_indices: torch.Tensor + family_indices: torch.Tensor + real_token_count_static: int + lengths_by_rank_cpu: torch.Tensor | None = None + family_indices_cpu: torch.Tensor | None = None + parent_indices: torch.Tensor | None = None + parent_indices_cpu: torch.Tensor | None = None + needs_final_state: bool = True + output_mask: torch.Tensor | None = None + + @property + def segment_count(self) -> int: + return int(self.family_indices.numel()) + + @property + def real_token_count(self) -> int: + return self.real_token_count_static + + +@dataclass(frozen=True) +class GdnStateExchangePlan: + """Sparse CP exchange for tree parent states needed by remote children.""" + + source_family_indices: tuple[int, ...] + dest_family_indices: tuple[int, ...] + exchange: Any + reverse_exchange: Any + + +@dataclass(frozen=True) +class _ExplicitBucketColumn: + row_index: int + family_index: int + parent_index: int + positions: tuple[int, ...] + output_mask: tuple[bool, ...] + needs_final_state: bool + + @property + def length(self) -> int: + return len(self.positions) + + +@dataclass(frozen=True) +class GdnPlannerConfig: + """Runtime model for one packed-row GDN execution plan. + + The defaults reproduce the Qwen3.5-35B reference shape used to fit the + GDN CP lab: hidden=2048, BF16 hidden exchange, 32 local value heads, and + key/value dim 128. Production callers should construct this with + ``from_model_shape`` so the same fitted hardware model scales by the actual + GDN state size instead of acting as a per-model profile. + """ + + cp_chain_beam_width: int = 2 + cp_chain_beam_branch_factor: int = 4 + cp_chain_beam_candidate_limit: int = 16 + cp_chain_beam_max_steps: int = 4 + # Chain buckets add extra collectives and kernel shapes; require a + # measurable runtime win before selecting them over local execution. + cp_chain_min_runtime_delta_ms: float = 4.0 + runtime_hidden_bytes_per_token: int = 4096 + runtime_layout_exchange_count: int = 4 + # Global all-to-all token counts are priced against aggregate CP bandwidth. + runtime_layout_bandwidth_bytes_per_ms: float = 448_000_000.0 + runtime_layout_collective_latency_ms: float = 0.0 + # Recurrent rates are fitted fwd+bwd exposed runtime, including the CP + # recurrent communication work that is not captured by token counts alone. + runtime_local_recurrent_tokens_per_ms: float = 1_500.0 + runtime_chain_recurrent_tokens_per_ms: float = 1_400.0 + runtime_local_bucket_launch_ms: float = 0.20 + runtime_chain_bucket_launch_ms: float = 0.20 + runtime_local_segment_launch_ms: float = 0.005 + runtime_cp_summary_bytes_per_segment: int = 4_194_304 + runtime_cp_summary_exchange_count_per_bucket: int = 8 + # Summary collectives move small state tensors and do not sustain the large + # hidden-state all-to-all bandwidth used by the layout exchange term. + runtime_cp_summary_bandwidth_bytes_per_ms: float = 80_000_000.0 + runtime_cp_summary_collective_latency_ms: float = 0.0 + runtime_cp_summary_compute_segments_per_ms: float = 320.0 + runtime_cp_suffix_scan_latency_ms: float = 2.0 + runtime_cp_suffix_scan_segments_per_ms: float = 15.0 + runtime_parent_state_bytes_per_exchange: int = 262_144 + runtime_parent_state_bandwidth_bytes_per_ms: float = 56_000_000.0 + runtime_parent_state_latency_ms: float = 0.0 + + @classmethod + def from_provider( + cls, + provider: Any, + *, + dtype_bytes: int | None = None, + **overrides: Any, + ) -> GdnPlannerConfig: + return cls.from_model_shape( + hidden_size=int(getattr(provider, "hidden_size")), + tensor_model_parallel_size=int( + getattr(provider, "tensor_model_parallel_size", 1) or 1 + ), + linear_num_key_heads=int(getattr(provider, "linear_num_key_heads")), + linear_num_value_heads=int(getattr(provider, "linear_num_value_heads")), + linear_key_head_dim=int(getattr(provider, "linear_key_head_dim")), + linear_value_head_dim=int(getattr(provider, "linear_value_head_dim")), + linear_conv_kernel_dim=int( + getattr(provider, "linear_conv_kernel_dim", 4) or 4 + ), + dtype_bytes=( + dtype_bytes + if dtype_bytes is not None + else _dtype_bytes(getattr(provider, "params_dtype", None)) + ), + **overrides, + ) + + @classmethod + def from_model_shape( + cls, + *, + hidden_size: int, + tensor_model_parallel_size: int, + linear_num_key_heads: int, + linear_num_value_heads: int, + linear_key_head_dim: int, + linear_value_head_dim: int, + linear_conv_kernel_dim: int = 4, + dtype_bytes: int = 2, + **overrides: Any, + ) -> GdnPlannerConfig: + """Build the fitted runtime model for an arbitrary GDN tensor shape.""" + + tp_size = max(1, int(tensor_model_parallel_size)) + if linear_num_key_heads % tp_size != 0: + raise ValueError( + "linear_num_key_heads must be divisible by tensor parallel size, " + f"got {linear_num_key_heads} and {tp_size}" + ) + if linear_num_value_heads % tp_size != 0: + raise ValueError( + "linear_num_value_heads must be divisible by tensor parallel size, " + f"got {linear_num_value_heads} and {tp_size}" + ) + local_key_heads = int(linear_num_key_heads) // tp_size + local_value_heads = int(linear_num_value_heads) // tp_size + key_dim = int(linear_key_head_dim) + value_dim = int(linear_value_head_dim) + recurrent_state_elements = local_value_heads * key_dim * value_dim + summary_state_elements = local_value_heads * key_dim * (value_dim + key_dim) + ref_recurrent_elements = 32 * 128 * 128 + ref_summary_elements = 32 * 128 * (128 + 128) + recurrent_scale = ( + ref_recurrent_elements / max(1, recurrent_state_elements) + ) ** _RUNTIME_STATE_THROUGHPUT_EXPONENT + summary_scale = ( + ref_summary_elements / max(1, summary_state_elements) + ) ** _RUNTIME_STATE_THROUGHPUT_EXPONENT + qkv_channels = 2 * local_key_heads * key_dim + local_value_heads * value_dim + conv_state_bytes = ( + qkv_channels * max(0, int(linear_conv_kernel_dim) - 1) * int(dtype_bytes) + ) + recurrent_state_bytes = recurrent_state_elements * 4 + ref_parent_state_bytes = ( + (2 * 16 * 128 + 32 * 128) * (4 - 1) * 2 + ) + ref_recurrent_elements * 4 + parent_state_bandwidth = 56_000_000.0 * ref_parent_state_bytes / 262_144.0 + config = cls( + runtime_hidden_bytes_per_token=int(hidden_size) * int(dtype_bytes), + runtime_local_recurrent_tokens_per_ms=1_500.0 * recurrent_scale, + runtime_chain_recurrent_tokens_per_ms=1_400.0 * recurrent_scale, + runtime_cp_summary_bytes_per_segment=summary_state_elements * 4, + runtime_cp_summary_compute_segments_per_ms=320.0 * summary_scale, + runtime_cp_suffix_scan_segments_per_ms=15.0 * summary_scale, + runtime_parent_state_bytes_per_exchange=( + conv_state_bytes + recurrent_state_bytes + ), + runtime_parent_state_bandwidth_bytes_per_ms=parent_state_bandwidth, + ) + return ( + cast(GdnPlannerConfig, replace(config, **overrides)) + if overrides + else config + ) + + +@dataclass(frozen=True) +class GdnRankExecutionPlan: + """Rank-local planned execution metadata for prefix-tree GDN.""" + + cp_rank: int + cp_size: int + batch_size: int + sequence_length: int + real_token_mask: torch.Tensor + packed_batch_size: int | None = None + packed_sequence_length: int | None = None + attention_to_gdn: Any | None = None + gdn_to_attention: Any | None = None + attention_token_ranges: tuple[tuple[int, int, int], ...] = () + gdn_token_ranges: tuple[tuple[int, int, int], ...] = () + attention_token_count: int = 0 + gdn_token_count: int = 0 + tree_segment_buckets_by_depth: tuple[tuple[GdnSegmentBucketPlan, ...], ...] = () + tree_chain_buckets_by_depth: tuple[tuple[GdnSegmentBucketPlan, ...], ...] = () + tree_state_exchanges_by_depth: tuple[GdnStateExchangePlan | None, ...] = () + + @property + def attention_token_indices(self) -> tuple[int, ...]: + return _tokens_from_rank_ranges(self.attention_token_ranges) + + @property + def gdn_token_indices(self) -> tuple[int, ...]: + return _tokens_from_rank_ranges(self.gdn_token_ranges) + + +@dataclass(frozen=True) +class _AttentionLayoutIndex: + """Counting index for CP attention token ownership.""" + + token_ranges_by_rank: tuple[tuple[tuple[int, int], ...], ...] + token_range_ends_by_rank: tuple[tuple[int, ...], ...] + + +GdnSegmentDecisionKey = tuple[int, int, int] + + +class _GdnChainSearchDecision(NamedTuple): + chain_segment_keys: frozenset[GdnSegmentDecisionKey] + score: float + + +class _GdnChainSearchMetadata(NamedTuple): + depth_count: int + children_by_node: tuple[tuple[int, ...], ...] + root_indices: tuple[int, ...] + subtree_token_counts: tuple[int, ...] + subtree_indices_by_node: tuple[tuple[int, ...], ...] + subtree_attention_counts: tuple[tuple[int, ...], ...] + segment_keys: tuple[GdnSegmentDecisionKey, ...] + segment_lengths: tuple[int, ...] + segment_depths: tuple[int, ...] + segment_attention_counts: tuple[tuple[int, ...], ...] + + +class _GdnLocalRuntimeEstimate(NamedTuple): + rank_work: tuple[int, ...] + rank_bucket_counts: tuple[int, ...] + rank_segment_counts: tuple[int, ...] + + +class _GdnChainRuntimeEstimate(NamedTuple): + rank_work: tuple[int, ...] + bucket_count: int + segment_count: int + + +def _tokens_from_rank_ranges( + ranges: tuple[tuple[int, int, int], ...], +) -> tuple[int, ...]: + return tuple(token for start, end, _ in ranges for token in range(start, end)) + + +def build_gdn_rank_execution_plan( + spec: GdnPackedExecutionSpec, + *, + device: torch.device | str, + cp_rank: int = 0, + cp_size: int = 1, + attention_token_layout_index: TokenLayoutIndex | None = None, + planner_config: GdnPlannerConfig | None = None, +) -> GdnRankExecutionPlan: + """Build rank-local tensor metadata from a parsed prefix-tree DAG. + + Planning is CPU-bound and must run once per packed training sequence. CP>1 + emits mixed work: native FLA CP chain buckets for long segments and local + fork buckets for short work where CP collectives would be inefficient. + """ + + planner_config = planner_config or GdnPlannerConfig() + target_device = torch.device(device) + if target_device.type != "cpu": + cpu_plan = build_gdn_rank_execution_plan( + spec, + device="cpu", + cp_rank=cp_rank, + cp_size=cp_size, + attention_token_layout_index=attention_token_layout_index, + planner_config=planner_config, + ) + return move_gdn_rank_execution_plan_to_device(cpu_plan, target_device) + return _build_tree_rank_execution_plan( + spec, + device=device, + cp_rank=cp_rank, + cp_size=cp_size, + attention_token_layout_index=attention_token_layout_index, + planner_config=planner_config, + ) + + +def _build_tree_rank_execution_plan( + spec: GdnPackedExecutionSpec, + *, + device: torch.device | str, + cp_rank: int, + cp_size: int, + attention_token_layout_index: TokenLayoutIndex | None, + planner_config: GdnPlannerConfig, +) -> GdnRankExecutionPlan: + if cp_size < 1: + raise ValueError(f"cp_size must be >= 1, got {cp_size}") + if cp_rank < 0 or cp_rank >= cp_size: + raise ValueError(f"cp_rank must be in [0, {cp_size}), got {cp_rank}") + if not spec.tree_segments: + raise ValueError("tree GDN planning requires tree segments") + if len(spec.tree_parent_indices) != len(spec.tree_segments): + raise ValueError("tree parent metadata length must match tree segments") + if len(spec.tree_depths) != len(spec.tree_segments): + raise ValueError("tree depth metadata length must match tree segments") + + from art.megatron.gdn.layout import ( + _reverse_exchange_plan, + build_local_rank_cp_exchange_plan_from_dest_ranges, + ) + + source_layout = _attention_source_layout( + spec, + cp_size=cp_size, + attention_token_layout_index=attention_token_layout_index, + ) + attention_layout_index = _build_attention_layout_index_from_token_layout( + source_layout + ) + segment_attention_counts = _segment_attention_rank_counts( + spec, + cp_size=cp_size, + attention_layout_index=attention_layout_index, + ) + chain_segment_keys = _select_chain_segment_keys( + spec, + cp_size=cp_size, + attention_layout_index=attention_layout_index, + segment_attention_counts=segment_attention_counts, + planner_config=planner_config, + ) + + depth_count = max(spec.tree_depths, default=0) + 1 + rank_loads = [0] * cp_size + owner_by_node = [-1] * len(spec.tree_segments) + chained_nodes = [False] * len(spec.tree_segments) + tree_has_children = [False] * len(spec.tree_segments) + for parent_index in spec.tree_parent_indices: + if parent_index >= 0: + tree_has_children[parent_index] = True + gdn_ranges_by_rank: list[list[tuple[int, int, int]]] = [[] for _ in range(cp_size)] + segments_by_rank_depth: list[list[list[GdnSegmentSpec]]] = [ + [[] for _ in range(depth_count)] for _ in range(cp_size) + ] + chain_segments_by_depth: list[list[GdnSegmentSpec]] = [ + [] for _ in range(depth_count) + ] + cross_rank_token_count = 0 + + children_by_node: list[list[int]] = [[] for _ in spec.tree_segments] + root_indices: list[int] = [] + for node_index, parent_index in enumerate(spec.tree_parent_indices): + if parent_index < 0: + root_indices.append(node_index) + else: + children_by_node[parent_index].append(node_index) + + def subtree_indices(root_index: int) -> tuple[int, ...]: + ordered: list[int] = [] + stack = [root_index] + while stack: + node_index = stack.pop() + ordered.append(node_index) + stack.extend(reversed(children_by_node[node_index])) + return tuple(ordered) + + def assign_local_group(node_indices: tuple[int, ...]) -> None: + nonlocal cross_rank_token_count + segments = tuple(spec.tree_segments[index] for index in node_indices) + owner = _best_segment_owner( + segments, + rank_loads, + segment_attention_counts=segment_attention_counts, + planner_config=planner_config, + ) + for segment in segments: + owner_by_node[segment.family_index] = owner + segments_by_rank_depth[owner][ + spec.tree_depths[segment.family_index] + ].append(segment) + cross_rank_token_count += _append_local_segment( + gdn_ranges_by_rank, + rank_loads, + owner, + segment, + spec, + segment_attention_counts=segment_attention_counts, + ) + + subtree_token_counts = [segment.length for segment in spec.tree_segments] + for node_index in reversed(range(len(spec.tree_segments))): + for child_index in children_by_node[node_index]: + subtree_token_counts[node_index] += subtree_token_counts[child_index] + chain_nodes = [ + _segment_key(segment) in chain_segment_keys for segment in spec.tree_segments + ] + subtree_has_chained_node = list(chain_nodes) + for node_index in reversed(range(len(spec.tree_segments))): + subtree_has_chained_node[node_index] = subtree_has_chained_node[ + node_index + ] or any( + subtree_has_chained_node[child_index] + for child_index in children_by_node[node_index] + ) + target_rank_load = spec.real_token_count / max(1, cp_size) + max_local_group_tokens = max(1, int(target_rank_load)) + + def assign_tree(node_index: int) -> None: + nonlocal cross_rank_token_count + segment = spec.tree_segments[node_index] + if chain_nodes[node_index]: + chained_nodes[segment.family_index] = True + chain_segments_by_depth[spec.tree_depths[segment.family_index]].append( + segment + ) + cross_rank_token_count += _append_chain_segment( + gdn_ranges_by_rank, + rank_loads, + segment, + spec, + attention_layout_index=attention_layout_index, + ) + for child_index in children_by_node[node_index]: + assign_tree(child_index) + return + + if ( + not subtree_has_chained_node[node_index] + and subtree_token_counts[node_index] <= max_local_group_tokens + ): + assign_local_group(subtree_indices(node_index)) + return + + assign_local_group((node_index,)) + for child_index in children_by_node[node_index]: + assign_tree(child_index) + + for root_index in root_indices: + assign_tree(root_index) + + gdn_ranges_by_rank_by_position = tuple( + tuple(ranges) for ranges in gdn_ranges_by_rank + ) + gdn_ranges_by_rank_by_source = tuple( + tuple(sorted(ranges)) for ranges in gdn_ranges_by_rank + ) + + attention_to_gdn = build_local_rank_cp_exchange_plan_from_dest_ranges( + source_layout=source_layout, + device=device, + local_rank=cp_rank, + dest_ranges_by_rank=gdn_ranges_by_rank_by_position, + cross_rank_token_count=cross_rank_token_count, + ) + local_token_ranges = gdn_ranges_by_rank_by_source[cp_rank] + if cp_size == 1: + tree_segment_buckets_by_depth = _build_chunk_aligned_cp1_tree_buckets( + spec, + tuple(tree_has_children), + device=device, + planner_config=planner_config, + ) + else: + tree_segment_buckets_by_depth = tuple( + _build_tree_bucket_plans( + tuple(segments_by_rank_depth[cp_rank][depth]), + spec.tree_parent_indices, + tuple(tree_has_children), + local_token_ranges=local_token_ranges, + sequence_length=spec.sequence_length, + device=device, + ) + for depth in range(depth_count) + ) + tree_chain_buckets_by_depth = ( + tuple( + _build_tree_bucket_plans( + tuple(chain_segments_by_depth[depth]), + spec.tree_parent_indices, + tuple(tree_has_children), + local_token_ranges=local_token_ranges, + sequence_length=spec.sequence_length, + device=device, + token_ranges_by_rank=tuple( + tuple(ranges) for ranges in gdn_ranges_by_rank_by_source + ), + split_by_final_state=False, + ) + for depth in range(depth_count) + ) + if cp_size > 1 + else tuple(() for _ in range(depth_count)) + ) + tree_state_exchanges_by_depth = _build_tree_state_exchanges_by_depth( + spec, + owner_by_node=tuple(owner_by_node), + chained_nodes=tuple(chained_nodes), + cp_rank=cp_rank, + cp_size=cp_size, + depth_count=depth_count, + device=device, + ) + if cp_size == 1: + valid_lengths = torch.tensor( + spec.valid_lengths, device=device, dtype=torch.long + ) + positions = torch.arange(spec.sequence_length, device=device, dtype=torch.long) + real_token_mask = positions.unsqueeze(0) < valid_lengths.unsqueeze(1) + else: + real_token_mask = torch.ones( + 1, + rank_loads[cp_rank], + device=device, + dtype=torch.bool, + ) + + return GdnRankExecutionPlan( + cp_rank=cp_rank, + cp_size=cp_size, + batch_size=1 if cp_size > 1 else spec.batch_size, + sequence_length=rank_loads[cp_rank] if cp_size > 1 else spec.sequence_length, + packed_batch_size=spec.batch_size, + packed_sequence_length=spec.sequence_length, + real_token_mask=real_token_mask, + attention_to_gdn=attention_to_gdn, + gdn_to_attention=_reverse_exchange_plan(attention_to_gdn), + attention_token_ranges=source_layout.ownership_ranges_by_rank[cp_rank], + gdn_token_ranges=gdn_ranges_by_rank_by_position[cp_rank], + attention_token_count=source_layout.token_counts_by_rank[cp_rank], + gdn_token_count=rank_loads[cp_rank], + tree_segment_buckets_by_depth=tree_segment_buckets_by_depth, + tree_chain_buckets_by_depth=tree_chain_buckets_by_depth, + tree_state_exchanges_by_depth=tree_state_exchanges_by_depth, + ) + + +def move_gdn_rank_execution_plan_to_device( + plan: GdnRankExecutionPlan, + device: torch.device | str, +) -> GdnRankExecutionPlan: + """Move planner tensors to the execution device after CPU planning.""" + + from art.megatron.gdn.layout import move_cp_exchange_plan_to_device + + return replace( + plan, + real_token_mask=_move_planner_tensor(plan.real_token_mask, device), + attention_to_gdn=move_cp_exchange_plan_to_device(plan.attention_to_gdn, device), + gdn_to_attention=move_cp_exchange_plan_to_device(plan.gdn_to_attention, device), + tree_segment_buckets_by_depth=tuple( + _move_bucket_plans(buckets, device) + for buckets in plan.tree_segment_buckets_by_depth + ), + tree_chain_buckets_by_depth=tuple( + _move_bucket_plans(buckets, device) + for buckets in plan.tree_chain_buckets_by_depth + ), + tree_state_exchanges_by_depth=tuple( + _move_state_exchange_plan(exchange, device) + for exchange in plan.tree_state_exchanges_by_depth + ), + ) + + +def _move_state_exchange_plan( + exchange: GdnStateExchangePlan | None, + device: torch.device | str, +) -> GdnStateExchangePlan | None: + if exchange is None: + return None + from art.megatron.gdn.layout import move_cp_exchange_plan_to_device + + return replace( + exchange, + exchange=move_cp_exchange_plan_to_device(exchange.exchange, device), + reverse_exchange=move_cp_exchange_plan_to_device( + exchange.reverse_exchange, device + ), + ) + + +def _move_bucket_plans( + buckets: tuple[GdnSegmentBucketPlan, ...], + device: torch.device | str, +) -> tuple[GdnSegmentBucketPlan, ...]: + return tuple( + replace( + bucket, + lengths=_move_planner_tensor(bucket.lengths, device), + real_mask=_move_planner_tensor(bucket.real_mask, device), + cu_seqlens=_move_planner_tensor(bucket.cu_seqlens, device), + row_indices=_move_planner_tensor(bucket.row_indices, device), + position_indices=_move_planner_tensor(bucket.position_indices, device), + family_indices=_move_planner_tensor(bucket.family_indices, device), + parent_indices=( + _move_planner_tensor(bucket.parent_indices, device) + if bucket.parent_indices is not None + else None + ), + output_mask=( + _move_planner_tensor(bucket.output_mask, device) + if bucket.output_mask is not None + else None + ), + ) + for bucket in buckets + ) + + +def parse_gdn_prefix_tree_segments( + group_ids: torch.Tensor, + parent_ids: torch.Tensor, +) -> GdnPackedExecutionSpec: + """Parse ART packed prefix-tree metadata into generic GDN tree nodes.""" + + groups = _rank2_long_cpu("group_ids", group_ids) + parents = _rank2_long_cpu("parent_ids", parent_ids) + if tuple(groups.shape) != tuple(parents.shape): + raise ValueError( + "group_ids and parent_ids must have the same shape, got " + f"{tuple(groups.shape)} and {tuple(parents.shape)}" + ) + + batch_size, sequence_length = (int(groups.shape[0]), int(groups.shape[1])) + rows = parse_prefix_tree(group_ids=groups, parent_ids=parents) + tree_segments: list[GdnSegmentSpec] = [] + tree_parent_indices: list[int] = [] + tree_depths: list[int] = [] + valid_lengths: list[int] = [] + node_by_row_group: dict[tuple[int, int], int] = {} + child_counts_by_parent: dict[int, int] = {} + + for row in rows: + valid_lengths.append(row.valid_tokens) + for segment in row.segments: + node_index = len(tree_segments) + is_root = segment.depth == 0 + parent_node_index = ( + -1 if is_root else node_by_row_group[(row.row_index, segment.parent_id)] + ) + child_index = None + if not is_root: + child_index = child_counts_by_parent.get(parent_node_index, 0) + child_counts_by_parent[parent_node_index] = child_index + 1 + tree_segments.append( + GdnSegmentSpec( + row_index=row.row_index, + family_index=node_index, + group_id=segment.group_id, + parent_id=segment.parent_id, + start=segment.start, + end=segment.end, + kind="prefix" if is_root else "completion", + child_index=child_index, + ) + ) + tree_parent_indices.append(parent_node_index) + tree_depths.append(segment.depth) + node_by_row_group[(row.row_index, segment.group_id)] = node_index + + return GdnPackedExecutionSpec( + batch_size=batch_size, + sequence_length=sequence_length, + valid_lengths=tuple(valid_lengths), + tree_segments=tuple(tree_segments), + tree_parent_indices=tuple(tree_parent_indices), + tree_depths=tuple(tree_depths), + ) + + +def _attention_source_layout( + spec: GdnPackedExecutionSpec, + *, + cp_size: int, + attention_token_layout_index: TokenLayoutIndex | None, +) -> TokenLayoutIndex: + if attention_token_layout_index is not None: + layout_cp_size = len(attention_token_layout_index.token_counts_by_rank) + layout_token_count = sum( + int(count) for count in attention_token_layout_index.token_counts_by_rank + ) + if layout_cp_size != cp_size: + raise ValueError( + "attention token layout index cp_size must match GDN cp_size, got " + f"{layout_cp_size} and {cp_size}" + ) + if layout_token_count != spec.real_token_count: + raise ValueError( + "attention token layout index token count must match GDN real token " + f"count, got {layout_token_count} and {spec.real_token_count}" + ) + return attention_token_layout_index + if cp_size != 1: + raise ValueError("GDN CP planning requires attention_token_layout_index") + ranges_by_rank = (((0, int(spec.real_token_count), 0),),) + return TokenLayoutIndex( + ownership_ranges_by_rank=ranges_by_rank, + token_counts_by_rank=tuple( + sum(int(end) - int(start) for start, end, _ in ranges) + for ranges in ranges_by_rank + ), + ) + + +def _can_chain_tree_segment( + segment: GdnSegmentSpec, + *, + cp_size: int, +) -> bool: + return segment.length >= cp_size and segment.length // FLA_CHUNK_SIZE >= cp_size + + +def _best_segment_owner( + segments: tuple[GdnSegmentSpec, ...], + rank_loads: list[int], + *, + segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], + planner_config: GdnPlannerConfig, +) -> int: + segment_length = sum(segment.length for segment in segments) + if len(segments) == 1: + on_rank_tokens = segment_attention_counts[_segment_key(segments[0])] + else: + rank_count = len(rank_loads) + counts_by_rank = [0] * rank_count + for segment in segments: + segment_counts = segment_attention_counts[_segment_key(segment)] + for rank in range(rank_count): + counts_by_rank[rank] += segment_counts[rank] + on_rank_tokens = tuple(counts_by_rank) + return _best_search_owner( + segment_length, + on_rank_tokens, + rank_loads, + planner_config=planner_config, + ) + + +def _best_search_owner( + segment_length: int, + on_rank_tokens: tuple[int, ...], + rank_loads: list[int], + *, + planner_config: GdnPlannerConfig, +) -> int: + best: tuple[float, float, int, int, int] | None = None + for rank, tokens in enumerate(on_rank_tokens): + projected_loads = list(rank_loads) + projected_loads[rank] += segment_length + rank_runtime_ms = max( + load / planner_config.runtime_local_recurrent_tokens_per_ms + for load in projected_loads + ) + cross_rank_tokens = segment_length - int(tokens) + exchange_runtime_ms = _predict_layout_exchange_runtime_ms( + cross_rank_tokens, + planner_config, + ) + candidate = ( + rank_runtime_ms + exchange_runtime_ms, + exchange_runtime_ms, + cross_rank_tokens, + -int(tokens), + rank, + ) + if best is None or candidate < best: + best = candidate + if best is None: + return _least_loaded_rank(rank_loads) + return best[-1] + + +def _select_chain_segment_keys( + spec: GdnPackedExecutionSpec, + *, + cp_size: int, + attention_layout_index: _AttentionLayoutIndex, + segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], + planner_config: GdnPlannerConfig, +) -> frozenset[GdnSegmentDecisionKey]: + if cp_size <= 1: + return frozenset() + legal_chain_segments = tuple( + segment + for segment in spec.tree_segments + if _can_chain_tree_segment(segment, cp_size=cp_size) + ) + if not legal_chain_segments: + return frozenset() + chain_rank_counts_by_key: dict[GdnSegmentDecisionKey, tuple[int, ...]] = {} + chain_cross_rank_tokens_by_key: dict[GdnSegmentDecisionKey, int] = {} + for segment in legal_chain_segments: + key = _segment_key(segment) + ( + chain_rank_counts_by_key[key], + chain_cross_rank_tokens_by_key[key], + ) = _chain_segment_rank_counts_and_cross_rank_tokens( + segment, + spec, + cp_size=cp_size, + attention_layout_index=attention_layout_index, + ) + + search_metadata = _build_chain_search_metadata( + spec, + segment_attention_counts=segment_attention_counts, + ) + score_cache: dict[frozenset[GdnSegmentDecisionKey], _GdnChainSearchDecision] = {} + + def decision_for( + chain_segment_keys: frozenset[GdnSegmentDecisionKey], + ) -> _GdnChainSearchDecision: + cached = score_cache.get(chain_segment_keys) + if cached is not None: + return cached + decision = _GdnChainSearchDecision( + chain_segment_keys=chain_segment_keys, + score=_score_chain_segment_keys( + spec, + cp_size=cp_size, + search_metadata=search_metadata, + chain_rank_counts_by_key=chain_rank_counts_by_key, + chain_cross_rank_tokens_by_key=chain_cross_rank_tokens_by_key, + chain_segment_keys=chain_segment_keys, + planner_config=planner_config, + ), + ) + score_cache[chain_segment_keys] = decision + return decision + + legal_chain_keys = frozenset( + _segment_key(segment) for segment in legal_chain_segments + ) + best = decision_for(frozenset()) + all_chain = decision_for(legal_chain_keys) + if best.score - all_chain.score > planner_config.cp_chain_min_runtime_delta_ms: + best = all_chain + beam = _best_chain_search_decisions( + (decision_for(frozenset()), all_chain), + limit=planner_config.cp_chain_beam_width, + ) + candidate_groups = _bounded_chain_candidate_groups( + spec, + legal_chain_segments, + segment_attention_counts=segment_attention_counts, + chain_rank_counts_by_key=chain_rank_counts_by_key, + planner_config=planner_config, + ) + stale_steps = 0 + for _ in range(max(0, planner_config.cp_chain_beam_max_steps)): + if not candidate_groups: + break + expanded: dict[frozenset[GdnSegmentDecisionKey], _GdnChainSearchDecision] = {} + for decision in beam: + neighbors: list[_GdnChainSearchDecision] = [] + for segment_keys in _chain_beam_neighbor_groups( + decision.chain_segment_keys, + candidate_groups=candidate_groups, + branch_factor=planner_config.cp_chain_beam_branch_factor, + ): + next_keys = ( + decision.chain_segment_keys - segment_keys + if segment_keys.issubset(decision.chain_segment_keys) + else decision.chain_segment_keys | segment_keys + ) + neighbors.append(decision_for(frozenset(next_keys))) + for neighbor in _best_chain_search_decisions( + neighbors, + limit=planner_config.cp_chain_beam_branch_factor, + ): + expanded[neighbor.chain_segment_keys] = neighbor + if not expanded: + break + beam = _best_chain_search_decisions( + (*beam, *expanded.values()), + limit=planner_config.cp_chain_beam_width, + ) + step_best = beam[0] + if best.score - step_best.score > planner_config.cp_chain_min_runtime_delta_ms: + best = step_best + stale_steps = 0 + else: + stale_steps += 1 + if stale_steps >= 2: + break + return best.chain_segment_keys + + +def _build_chain_search_metadata( + spec: GdnPackedExecutionSpec, + *, + segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], +) -> _GdnChainSearchMetadata: + children_by_node: list[list[int]] = [[] for _ in spec.tree_segments] + root_indices: list[int] = [] + for node_index, parent_index in enumerate(spec.tree_parent_indices): + if parent_index < 0: + root_indices.append(node_index) + else: + children_by_node[parent_index].append(node_index) + segment_keys = tuple(_segment_key(segment) for segment in spec.tree_segments) + segment_lengths = tuple(segment.length for segment in spec.tree_segments) + segment_depths = tuple( + spec.tree_depths[segment.family_index] for segment in spec.tree_segments + ) + segment_attention_counts_by_node = tuple( + segment_attention_counts[key] for key in segment_keys + ) + subtree_token_counts = list(segment_lengths) + subtree_indices_by_node: list[tuple[int, ...]] = [ + (node_index,) for node_index in range(len(spec.tree_segments)) + ] + subtree_attention_counts = [ + tuple(counts) for counts in segment_attention_counts_by_node + ] + for node_index in reversed(range(len(spec.tree_segments))): + for child_index in children_by_node[node_index]: + subtree_token_counts[node_index] += subtree_token_counts[child_index] + subtree_indices_by_node[node_index] = ( + *subtree_indices_by_node[node_index], + *subtree_indices_by_node[child_index], + ) + subtree_attention_counts[node_index] = tuple( + parent + child + for parent, child in zip( + subtree_attention_counts[node_index], + subtree_attention_counts[child_index], + strict=True, + ) + ) + return _GdnChainSearchMetadata( + depth_count=max(spec.tree_depths, default=0) + 1, + children_by_node=tuple(tuple(children) for children in children_by_node), + root_indices=tuple(root_indices), + subtree_token_counts=tuple(subtree_token_counts), + subtree_indices_by_node=tuple(subtree_indices_by_node), + subtree_attention_counts=tuple(subtree_attention_counts), + segment_keys=segment_keys, + segment_lengths=segment_lengths, + segment_depths=segment_depths, + segment_attention_counts=segment_attention_counts_by_node, + ) + + +def _score_chain_segment_keys( + spec: GdnPackedExecutionSpec, + *, + cp_size: int, + search_metadata: _GdnChainSearchMetadata, + chain_rank_counts_by_key: dict[GdnSegmentDecisionKey, tuple[int, ...]], + chain_cross_rank_tokens_by_key: dict[GdnSegmentDecisionKey, int], + chain_segment_keys: frozenset[GdnSegmentDecisionKey], + planner_config: GdnPlannerConfig, +) -> float: + chain_nodes = [key in chain_segment_keys for key in search_metadata.segment_keys] + subtree_has_chained_node = list(chain_nodes) + for node_index in reversed(range(len(spec.tree_segments))): + subtree_has_chained_node[node_index] = subtree_has_chained_node[ + node_index + ] or any( + subtree_has_chained_node[child_index] + for child_index in search_metadata.children_by_node[node_index] + ) + + rank_loads = [0] * cp_size + owner_by_node = [-1] * len(spec.tree_segments) + local_lengths_by_rank_depth: list[list[list[int]]] = [ + [[] for _ in range(search_metadata.depth_count)] for _ in range(cp_size) + ] + chain_rank_counts_by_depth: list[list[tuple[int, ...]]] = [ + [] for _ in range(search_metadata.depth_count) + ] + cross_rank_token_count = 0 + target_rank_load = spec.real_token_count / max(1, cp_size) + max_local_group_tokens = max(1, int(target_rank_load)) + + def add_local_group( + node_indices: tuple[int, ...], + segment_length: int, + on_rank_tokens: tuple[int, ...], + ) -> None: + nonlocal cross_rank_token_count + owner = _best_search_owner( + segment_length, + on_rank_tokens, + rank_loads, + planner_config=planner_config, + ) + rank_loads[owner] += segment_length + cross_rank_token_count += segment_length - on_rank_tokens[owner] + for node_index in node_indices: + owner_by_node[node_index] = owner + local_lengths_by_rank_depth[owner][ + search_metadata.segment_depths[node_index] + ].append(search_metadata.segment_lengths[node_index]) + + def add_chain_segment(segment: GdnSegmentSpec) -> None: + nonlocal cross_rank_token_count + key = _segment_key(segment) + chain_rank_counts_by_depth[spec.tree_depths[segment.family_index]].append( + chain_rank_counts_by_key[key] + ) + cross_rank_token_count += chain_cross_rank_tokens_by_key[key] + for rank, token_count in enumerate(chain_rank_counts_by_key[key]): + rank_loads[rank] += token_count + + def assign_tree(node_index: int) -> None: + segment = spec.tree_segments[node_index] + if chain_nodes[node_index]: + add_chain_segment(segment) + for child_index in search_metadata.children_by_node[node_index]: + assign_tree(child_index) + return + if ( + not subtree_has_chained_node[node_index] + and search_metadata.subtree_token_counts[node_index] + <= max_local_group_tokens + ): + add_local_group( + search_metadata.subtree_indices_by_node[node_index], + search_metadata.subtree_token_counts[node_index], + search_metadata.subtree_attention_counts[node_index], + ) + return + add_local_group( + (node_index,), + search_metadata.segment_lengths[node_index], + search_metadata.segment_attention_counts[node_index], + ) + for child_index in search_metadata.children_by_node[node_index]: + assign_tree(child_index) + + for root_index in search_metadata.root_indices: + assign_tree(root_index) + + local_runtime = _estimate_local_runtime_from_lengths( + tuple( + tuple(tuple(depth_segments) for depth_segments in rank_segments) + for rank_segments in local_lengths_by_rank_depth + ) + ) + chain_runtime = _estimate_chain_runtime_from_counts( + tuple(tuple(depth_counts) for depth_counts in chain_rank_counts_by_depth), + cp_size=cp_size, + ) + parent_state_exchange_count = _count_parent_state_exchanges( + spec, + owner_by_node=tuple(owner_by_node), + chain_nodes=tuple(chain_nodes), + ) + return _predict_gdn_plan_runtime_ms( + local_runtime, + chain_runtime, + cross_rank_token_count=cross_rank_token_count, + parent_state_exchange_count=parent_state_exchange_count, + planner_config=planner_config, + ) + + +def _add_local_search_load( + rank_loads: list[int], + owner: int, + segment: GdnSegmentSpec, + *, + segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], +) -> int: + rank_loads[owner] += segment.length + return segment.length - segment_attention_counts[_segment_key(segment)][owner] + + +def _estimate_local_rank_kernel_work( + local_segments_by_rank_depth: tuple[tuple[tuple[GdnSegmentSpec, ...], ...], ...], +) -> tuple[tuple[int, ...], int, int]: + estimate = _estimate_local_runtime_from_lengths( + tuple( + tuple( + tuple(segment.length for segment in segments) + for segments in rank_segments + ) + for rank_segments in local_segments_by_rank_depth + ) + ) + return ( + estimate.rank_work, + max(estimate.rank_bucket_counts, default=0), + max(estimate.rank_segment_counts, default=0), + ) + + +def _estimate_local_rank_kernel_work_from_lengths( + local_lengths_by_rank_depth: tuple[tuple[tuple[int, ...], ...], ...], +) -> tuple[tuple[int, ...], int, int]: + estimate = _estimate_local_runtime_from_lengths(local_lengths_by_rank_depth) + return ( + estimate.rank_work, + max(estimate.rank_bucket_counts, default=0), + max(estimate.rank_segment_counts, default=0), + ) + + +def _estimate_local_runtime_from_lengths( + local_lengths_by_rank_depth: tuple[tuple[tuple[int, ...], ...], ...], +) -> _GdnLocalRuntimeEstimate: + rank_work: list[int] = [] + rank_bucket_counts: list[int] = [] + rank_segment_counts: list[int] = [] + for lengths_by_depth in local_lengths_by_rank_depth: + work = 0 + bucket_count = 0 + segment_count = 0 + for lengths in lengths_by_depth: + bucket_work, depth_bucket_count = _padded_work_from_lengths(lengths) + work += bucket_work + bucket_count += depth_bucket_count + segment_count += len(lengths) + rank_work.append(work) + rank_bucket_counts.append(bucket_count) + rank_segment_counts.append(segment_count) + return _GdnLocalRuntimeEstimate( + rank_work=tuple(rank_work), + rank_bucket_counts=tuple(rank_bucket_counts), + rank_segment_counts=tuple(rank_segment_counts), + ) + + +def _estimate_chain_rank_kernel_work( + chain_segments_by_depth: tuple[tuple[GdnSegmentSpec, ...], ...], + *, + chain_rank_counts_by_key: dict[GdnSegmentDecisionKey, tuple[int, ...]], + cp_size: int, +) -> tuple[tuple[int, ...], int]: + estimate = _estimate_chain_runtime_from_counts( + tuple( + tuple( + chain_rank_counts_by_key[_segment_key(segment)] for segment in segments + ) + for segments in chain_segments_by_depth + ), + cp_size=cp_size, + ) + return estimate.rank_work, estimate.bucket_count + + +def _estimate_chain_rank_kernel_work_from_counts( + chain_rank_counts_by_depth: tuple[tuple[tuple[int, ...], ...], ...], + *, + cp_size: int, +) -> tuple[tuple[int, ...], int]: + estimate = _estimate_chain_runtime_from_counts( + chain_rank_counts_by_depth, + cp_size=cp_size, + ) + return estimate.rank_work, estimate.bucket_count + + +def _estimate_chain_runtime_from_counts( + chain_rank_counts_by_depth: tuple[tuple[tuple[int, ...], ...], ...], + *, + cp_size: int, +) -> _GdnChainRuntimeEstimate: + rank_work = [0] * cp_size + bucket_count = 0 + segment_count = 0 + for rank_counts_by_segment in chain_rank_counts_by_depth: + if not rank_counts_by_segment: + continue + bucket_count += 1 + segment_count += len(rank_counts_by_segment) + for rank in range(cp_size): + lengths = tuple(counts[rank] for counts in rank_counts_by_segment) + rank_work[rank] += max(lengths, default=0) * len(lengths) + return _GdnChainRuntimeEstimate( + rank_work=tuple(rank_work), + bucket_count=bucket_count, + segment_count=segment_count, + ) + + +def _predict_gdn_plan_runtime_ms( + local_runtime: _GdnLocalRuntimeEstimate, + chain_runtime: _GdnChainRuntimeEstimate, + *, + cross_rank_token_count: int, + parent_state_exchange_count: int, + planner_config: GdnPlannerConfig, +) -> float: + """Predict exposed fwd+bwd runtime for a GDN CP plan in milliseconds. + + This is a runtime cost model fitted from GDN CP lab measurements. It models + the critical per-rank recurrent path separately from exposed communication + and chain-summary work, so chain/local decisions compare predicted wall time + rather than abstract balance or token-count penalties. + """ + + rank_runtime_ms = 0.0 + rank_count = max( + len(local_runtime.rank_work), + len(chain_runtime.rank_work), + ) + for rank in range(rank_count): + local_work = ( + local_runtime.rank_work[rank] if rank < len(local_runtime.rank_work) else 0 + ) + local_bucket_count = ( + local_runtime.rank_bucket_counts[rank] + if rank < len(local_runtime.rank_bucket_counts) + else 0 + ) + local_segment_count = ( + local_runtime.rank_segment_counts[rank] + if rank < len(local_runtime.rank_segment_counts) + else 0 + ) + chain_work = ( + chain_runtime.rank_work[rank] if rank < len(chain_runtime.rank_work) else 0 + ) + rank_runtime_ms = max( + rank_runtime_ms, + _predict_local_rank_runtime_ms( + local_work, + bucket_count=local_bucket_count, + segment_count=local_segment_count, + planner_config=planner_config, + ) + + chain_work / planner_config.runtime_chain_recurrent_tokens_per_ms, + ) + return ( + rank_runtime_ms + + _predict_chain_exposed_runtime_ms( + chain_runtime, + planner_config=planner_config, + ) + + _predict_layout_exchange_runtime_ms( + cross_rank_token_count, + planner_config, + ) + + _predict_parent_state_exchange_runtime_ms( + parent_state_exchange_count, + planner_config, + ) + ) + + +def _predict_local_rank_runtime_ms( + recurrent_work: int, + *, + bucket_count: int, + segment_count: int, + planner_config: GdnPlannerConfig, +) -> float: + return ( + recurrent_work / planner_config.runtime_local_recurrent_tokens_per_ms + + bucket_count * planner_config.runtime_local_bucket_launch_ms + + segment_count * planner_config.runtime_local_segment_launch_ms + ) + + +def _predict_chain_exposed_runtime_ms( + chain_runtime: _GdnChainRuntimeEstimate, + *, + planner_config: GdnPlannerConfig, +) -> float: + if chain_runtime.bucket_count == 0: + return 0.0 + summary_bytes = ( + chain_runtime.segment_count + * planner_config.runtime_cp_summary_bytes_per_segment + * planner_config.runtime_cp_summary_exchange_count_per_bucket + ) + summary_exchange_ms = ( + summary_bytes / planner_config.runtime_cp_summary_bandwidth_bytes_per_ms + + chain_runtime.bucket_count + * planner_config.runtime_cp_summary_exchange_count_per_bucket + * planner_config.runtime_cp_summary_collective_latency_ms + ) + return ( + chain_runtime.bucket_count * planner_config.runtime_chain_bucket_launch_ms + + summary_exchange_ms + + chain_runtime.segment_count + / planner_config.runtime_cp_summary_compute_segments_per_ms + + chain_runtime.bucket_count * planner_config.runtime_cp_suffix_scan_latency_ms + + chain_runtime.segment_count + / planner_config.runtime_cp_suffix_scan_segments_per_ms + ) + + +def _predict_layout_exchange_runtime_ms( + cross_rank_token_count: int, + planner_config: GdnPlannerConfig, +) -> float: + if cross_rank_token_count <= 0: + return 0.0 + bytes_moved = ( + cross_rank_token_count + * planner_config.runtime_hidden_bytes_per_token + * planner_config.runtime_layout_exchange_count + ) + return ( + bytes_moved / planner_config.runtime_layout_bandwidth_bytes_per_ms + + planner_config.runtime_layout_exchange_count + * planner_config.runtime_layout_collective_latency_ms + ) + + +def _predict_parent_state_exchange_runtime_ms( + exchange_count: int, + planner_config: GdnPlannerConfig, +) -> float: + if exchange_count <= 0: + return 0.0 + return ( + exchange_count + * planner_config.runtime_parent_state_bytes_per_exchange + / planner_config.runtime_parent_state_bandwidth_bytes_per_ms + + exchange_count * planner_config.runtime_parent_state_latency_ms + ) + + +def _padded_work_from_lengths(lengths: tuple[int, ...]) -> tuple[int, int]: + real_lengths = tuple(length for length in lengths if length > 0) + if not real_lengths: + return 0, 0 + return max(real_lengths) * len(real_lengths), 1 + + +def _count_parent_state_exchanges( + spec: GdnPackedExecutionSpec, + *, + owner_by_node: tuple[int, ...], + chain_nodes: tuple[bool, ...], +) -> int: + exchange_count = 0 + for child_index, parent_index in enumerate(spec.tree_parent_indices): + if parent_index < 0 or chain_nodes[parent_index]: + continue + parent_owner = owner_by_node[parent_index] + if parent_owner < 0: + continue + if chain_nodes[child_index] or owner_by_node[child_index] != parent_owner: + exchange_count += 1 + return exchange_count + + +def _chain_beam_neighbor_groups( + chain_segment_keys: frozenset[GdnSegmentDecisionKey], + *, + candidate_groups: tuple[frozenset[GdnSegmentDecisionKey], ...], + branch_factor: int, +) -> tuple[frozenset[GdnSegmentDecisionKey], ...]: + selected: list[frozenset[GdnSegmentDecisionKey]] = [] + for group in candidate_groups: + if group and not group.issubset(chain_segment_keys): + selected.append(group) + if len(selected) >= branch_factor: + return tuple(selected) + for group in reversed(candidate_groups): + if group and group.intersection(chain_segment_keys) and group not in selected: + selected.append(group) + if len(selected) >= branch_factor: + break + return tuple(selected) + + +def _best_chain_search_decisions( + decisions: Any, + *, + limit: int, +) -> tuple[_GdnChainSearchDecision, ...]: + return tuple( + sorted( + decisions, + key=lambda decision: ( + decision.score, + len(decision.chain_segment_keys), + tuple(sorted(decision.chain_segment_keys)), + ), + )[: max(1, limit)] + ) + + +def _bounded_chain_candidate_groups( + spec: GdnPackedExecutionSpec, + legal_chain_segments: tuple[GdnSegmentSpec, ...], + *, + segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], + chain_rank_counts_by_key: dict[GdnSegmentDecisionKey, tuple[int, ...]], + planner_config: GdnPlannerConfig, +) -> tuple[frozenset[GdnSegmentDecisionKey], ...]: + legal_key_set = frozenset(_segment_key(segment) for segment in legal_chain_segments) + if not legal_key_set: + return () + root_keys = frozenset( + _segment_key(segment) + for segment in legal_chain_segments + if spec.tree_parent_indices[segment.family_index] < 0 + ) + nonroot_keys = legal_key_set - root_keys + groups: list[frozenset[GdnSegmentDecisionKey]] = [] + for group in (legal_key_set, root_keys, nonroot_keys): + if group and group not in groups: + groups.append(group) + for group in _ranked_chain_beam_groups( + spec, + legal_chain_segments, + segment_attention_counts=segment_attention_counts, + chain_rank_counts_by_key=chain_rank_counts_by_key, + ): + if group and group not in groups: + groups.append(group) + if len(groups) >= planner_config.cp_chain_beam_candidate_limit: + break + return tuple(groups[: max(1, planner_config.cp_chain_beam_candidate_limit)]) + + +def _ranked_chain_beam_groups( + spec: GdnPackedExecutionSpec, + legal_chain_segments: tuple[GdnSegmentSpec, ...], + *, + segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], + chain_rank_counts_by_key: dict[GdnSegmentDecisionKey, tuple[int, ...]], +) -> tuple[frozenset[GdnSegmentDecisionKey], ...]: + priority_by_key = { + _segment_key(segment): _chain_beam_segment_priority( + segment, + segment_attention_counts=segment_attention_counts, + chain_rank_counts_by_key=chain_rank_counts_by_key, + ) + for segment in legal_chain_segments + } + groups: set[frozenset[GdnSegmentDecisionKey]] = { + frozenset((key,)) for key in priority_by_key + } + children_by_parent: dict[int, list[GdnSegmentDecisionKey]] = {} + for segment in legal_chain_segments: + parent_index = spec.tree_parent_indices[segment.family_index] + if parent_index >= 0: + children_by_parent.setdefault(parent_index, []).append( + _segment_key(segment) + ) + for child_keys in children_by_parent.values(): + if len(child_keys) > 1: + groups.add(frozenset(child_keys)) + return tuple( + sorted( + groups, + key=lambda group: _chain_beam_group_priority( + group, + priority_by_key=priority_by_key, + ), + reverse=True, + ) + ) + + +def _chain_beam_group_priority( + group: frozenset[GdnSegmentDecisionKey], + *, + priority_by_key: dict[GdnSegmentDecisionKey, tuple[int, int, int, int]], +) -> tuple[int, int, int, int, int]: + priorities = tuple(priority_by_key[key] for key in group) + return ( + sum(priority[0] for priority in priorities), + sum(priority[1] for priority in priorities), + max((priority[2] for priority in priorities), default=0), + sum(priority[3] for priority in priorities), + len(group), + ) + + +def _chain_beam_segment_priority( + segment: GdnSegmentSpec, + *, + segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], + chain_rank_counts_by_key: dict[GdnSegmentDecisionKey, tuple[int, ...]], +) -> tuple[int, int, int, int]: + key = _segment_key(segment) + chain_max_load = max(chain_rank_counts_by_key[key], default=0) + best_attention_locality = max(segment_attention_counts[key], default=0) + chain_load_relief = segment.length - chain_max_load + minimum_local_exchange = segment.length - best_attention_locality + return ( + chain_load_relief, + segment.length, + best_attention_locality, + -minimum_local_exchange, + ) + + +def _build_tree_state_exchanges_by_depth( + spec: GdnPackedExecutionSpec, + *, + owner_by_node: tuple[int, ...], + chained_nodes: tuple[bool, ...], + cp_rank: int, + cp_size: int, + depth_count: int, + device: torch.device | str, +) -> tuple[GdnStateExchangePlan | None, ...]: + if cp_size <= 1: + return tuple(None for _ in range(depth_count)) + + from art.megatron.gdn.layout import ( + GdnCpExchangePlan, + _make_peer_transfer, + _reverse_exchange_plan, + ) + + families_by_depth_pair: list[dict[tuple[int, int], set[int]]] = [ + {} for _ in range(depth_count) + ] + for child_index, parent_index in enumerate(spec.tree_parent_indices): + if parent_index < 0 or chained_nodes[parent_index]: + continue + source_rank = owner_by_node[parent_index] + depth = spec.tree_depths[child_index] + if source_rank < 0: + raise ValueError( + "tree state exchange requires every local parent to have an owner" + ) + if chained_nodes[child_index]: + for dest_rank in range(cp_size): + if source_rank == dest_rank: + continue + families_by_depth_pair[depth].setdefault( + (source_rank, dest_rank), set() + ).add(parent_index) + continue + dest_rank = owner_by_node[child_index] + if dest_rank < 0: + raise ValueError( + "tree state exchange requires every local child to have an owner" + ) + if source_rank != dest_rank: + families_by_depth_pair[depth].setdefault( + (source_rank, dest_rank), set() + ).add(parent_index) + + state_exchanges: list[GdnStateExchangePlan | None] = [] + for pair_families in families_by_depth_pair: + if not pair_families: + state_exchanges.append(None) + continue + source_families_by_rank = [set[int]() for _ in range(cp_size)] + dest_families_by_rank = [set[int]() for _ in range(cp_size)] + for (source_rank, dest_rank), parent_indices in pair_families.items(): + source_families_by_rank[source_rank].update(parent_indices) + dest_families_by_rank[dest_rank].update(parent_indices) + source_families = tuple( + tuple(sorted(families)) for families in source_families_by_rank + ) + dest_families = tuple( + tuple(sorted(families)) for families in dest_families_by_rank + ) + source_positions = ( + {family: index for index, family in enumerate(families)} + for families in source_families + ) + dest_positions = ( + {family: index for index, family in enumerate(families)} + for families in dest_families + ) + source_position_by_rank = tuple(source_positions) + dest_position_by_rank = tuple(dest_positions) + transfers = [] + transfer_count = 0 + for (source_rank, dest_rank), parent_indices in sorted(pair_families.items()): + ordered = tuple(sorted(parent_indices)) + transfer_count += len(ordered) + transfers.append( + _make_peer_transfer( + source_rank=source_rank, + dest_rank=dest_rank, + source_positions=torch.tensor( + [ + source_position_by_rank[source_rank][family] + for family in ordered + ], + dtype=torch.long, + ), + dest_positions=torch.tensor( + [ + dest_position_by_rank[dest_rank][family] + for family in ordered + ], + dtype=torch.long, + ), + source_count=len(source_families[source_rank]), + dest_count=len(dest_families[dest_rank]), + device=device, + ) + ) + exchange = GdnCpExchangePlan( + cp_size=cp_size, + source_token_counts_by_rank=tuple( + len(families) for families in source_families + ), + dest_token_counts_by_rank=tuple( + len(families) for families in dest_families + ), + transfers=tuple(transfers), + cross_rank_token_count_override=transfer_count, + ) + state_exchanges.append( + GdnStateExchangePlan( + source_family_indices=source_families[cp_rank], + dest_family_indices=dest_families[cp_rank], + exchange=exchange, + reverse_exchange=_reverse_exchange_plan(exchange), + ) + ) + return tuple(state_exchanges) + + +def _build_attention_layout_index_from_token_layout( + layout: TokenLayoutIndex, +) -> _AttentionLayoutIndex: + ranges_by_rank = tuple( + tuple(sorted((int(start), int(end)) for start, end, _ in rank_ranges)) + for rank_ranges in layout.ownership_ranges_by_rank + ) + return _AttentionLayoutIndex( + token_ranges_by_rank=ranges_by_rank, + token_range_ends_by_rank=tuple( + tuple(end for _, end in ranges) for ranges in ranges_by_rank + ), + ) + + +def _segment_attention_rank_counts( + spec: GdnPackedExecutionSpec, + *, + cp_size: int, + attention_layout_index: _AttentionLayoutIndex, +) -> dict[tuple[int, int, int], tuple[int, ...]]: + del cp_size + segments = spec.tree_segments + if not segments: + return {} + starts = torch.tensor( + [_segment_token_start(segment, spec.sequence_length) for segment in segments], + dtype=torch.long, + ) + lengths = torch.tensor([segment.length for segment in segments], dtype=torch.long) + ends = starts + lengths + counts_by_rank = [] + for ranges in attention_layout_index.token_ranges_by_rank: + counts_by_rank.append(_rank_range_overlap_counts(starts, ends, ranges)) + counts_tensor = torch.stack(counts_by_rank, dim=1) + totals = counts_tensor.sum(dim=1) + if not torch.equal(totals, lengths): + bad_index = int(torch.nonzero(totals != lengths, as_tuple=False)[0].item()) + raise ValueError( + "attention layout is missing a real token required by GDN; " + f"segment={_segment_key(segments[bad_index])}" + ) + counts = counts_tensor.tolist() + return { + _segment_key(segment): tuple(int(value) for value in counts[index]) + for index, segment in enumerate(segments) + } + + +def _rank_range_overlap_counts( + starts: torch.Tensor, + ends: torch.Tensor, + ranges: tuple[tuple[int, int], ...], +) -> torch.Tensor: + if not ranges: + return torch.zeros_like(starts) + range_starts = torch.tensor([start for start, _ in ranges], dtype=torch.long) + range_ends = torch.tensor([end for _, end in ranges], dtype=torch.long) + range_lengths = range_ends - range_starts + prefix = torch.cat((range_lengths.new_zeros(1), torch.cumsum(range_lengths, dim=0))) + + def owned_before(points: torch.Tensor) -> torch.Tensor: + indices = torch.searchsorted(range_ends, points, right=False) + counts = prefix.index_select(0, indices) + active = indices < int(range_starts.numel()) + if bool(active.any().item()): + active_indices = indices[active] + active_starts = range_starts.index_select(0, active_indices) + active_ends = range_ends.index_select(0, active_indices) + counts[active] += torch.minimum( + torch.clamp(points[active] - active_starts, min=0), + active_ends - active_starts, + ) + return counts + + return owned_before(ends) - owned_before(starts) + + +def _segment_key(segment: GdnSegmentSpec) -> tuple[int, int, int]: + return (segment.row_index, segment.start, segment.end) + + +def _append_chain_segment( + gdn_ranges_by_rank: list[list[tuple[int, int, int]]], + rank_loads: list[int], + segment: GdnSegmentSpec, + spec: GdnPackedExecutionSpec, + *, + attention_layout_index: _AttentionLayoutIndex | None = None, +) -> int: + rank_ranges, cross_rank_tokens = _chain_segment_rank_ranges_and_cross_rank_tokens( + segment, + spec, + cp_size=len(gdn_ranges_by_rank), + attention_layout_index=attention_layout_index, + ) + for rank, (shard_start, shard_end) in enumerate(rank_ranges): + shard_length = shard_end - shard_start + if shard_length <= 0: + raise ValueError( + "CP chain planning requires non-empty shards; " + f"segment={segment.kind}:{segment.family_index} " + f"length={segment.length} cp_size={len(gdn_ranges_by_rank)}" + ) + position_start = rank_loads[rank] + gdn_ranges_by_rank[rank].append((shard_start, shard_end, position_start)) + rank_loads[rank] += shard_length + return cross_rank_tokens + + +def _chain_segment_rank_counts_and_cross_rank_tokens( + segment: GdnSegmentSpec, + spec: GdnPackedExecutionSpec, + *, + cp_size: int, + attention_layout_index: _AttentionLayoutIndex | None, +) -> tuple[tuple[int, ...], int]: + rank_ranges, cross_rank_tokens = _chain_segment_rank_ranges_and_cross_rank_tokens( + segment, + spec, + cp_size=cp_size, + attention_layout_index=attention_layout_index, + ) + return tuple(end - start for start, end in rank_ranges), cross_rank_tokens + + +def _chain_segment_rank_ranges_and_cross_rank_tokens( + segment: GdnSegmentSpec, + spec: GdnPackedExecutionSpec, + *, + cp_size: int, + attention_layout_index: _AttentionLayoutIndex | None, +) -> tuple[tuple[tuple[int, int], ...], int]: + token_start = _segment_token_start(segment, spec.sequence_length) + attention_shards = _attention_contiguous_chain_shards( + token_start, + segment.length, + cp_size=cp_size, + attention_layout_index=attention_layout_index, + ) + if attention_shards is not None: + return tuple((shard.start, shard.stop) for shard in attention_shards), 0 + shard_lengths = _fla_aligned_chain_shard_lengths(segment.length, cp_size=cp_size) + rank_ranges: list[tuple[int, int]] = [] + cross_rank_tokens = 0 + start = 0 + for rank, shard_length in enumerate(shard_lengths): + end = start + shard_length + if start >= end: + raise ValueError( + "CP chain planning requires non-empty shards; " + f"segment={segment.kind}:{segment.family_index} " + f"length={segment.length} cp_size={cp_size}" + ) + shard_start = token_start + start + shard_end = shard_start + shard_length + rank_ranges.append((shard_start, shard_end)) + if attention_layout_index is not None: + cross_rank_tokens += shard_length - _range_overlap_count( + shard_start, + shard_end, + attention_layout_index.token_ranges_by_rank[rank], + attention_layout_index.token_range_ends_by_rank[rank], + ) + start = end + return tuple(rank_ranges), cross_rank_tokens + + +def _fla_aligned_chain_shard_lengths(length: int, *, cp_size: int) -> tuple[int, ...]: + full_chunks = int(length) // FLA_CHUNK_SIZE + if full_chunks < int(cp_size): + raise ValueError( + "CP chain planning requires at least one full FLA chunk per rank; " + f"length={length} cp_size={cp_size}" + ) + base_chunks = full_chunks // int(cp_size) + extra_chunks = full_chunks % int(cp_size) + chunk_counts = tuple( + base_chunks + (1 if rank < extra_chunks else 0) for rank in range(int(cp_size)) + ) + lengths = [count * FLA_CHUNK_SIZE for count in chunk_counts] + lengths[-1] += int(length) - full_chunks * FLA_CHUNK_SIZE + return tuple(lengths) + + +def _attention_contiguous_chain_shards( + token_start: int, + token_count: int, + *, + cp_size: int, + attention_layout_index: _AttentionLayoutIndex | None, +) -> tuple[range, ...] | None: + if attention_layout_index is None: + return None + segment_end = token_start + token_count + shards: list[range] = [] + cursor = token_start + for rank in range(cp_size): + overlaps = _range_overlaps( + token_start, + segment_end, + attention_layout_index.token_ranges_by_rank[rank], + ) + if len(overlaps) != 1: + return None + start, end = overlaps[0] + if start != cursor or end <= start: + return None + shards.append(range(start, end)) + cursor = end + if cursor != segment_end: + return None + if any(len(shard) % FLA_CHUNK_SIZE != 0 for shard in shards[:-1]): + return None + return tuple(shards) + + +def _append_local_segment( + gdn_ranges_by_rank: list[list[tuple[int, int, int]]], + rank_loads: list[int], + rank: int, + segment: GdnSegmentSpec, + spec: GdnPackedExecutionSpec, + *, + segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], +) -> int: + token_start = _segment_token_start(segment, spec.sequence_length) + position_start = rank_loads[rank] + gdn_ranges_by_rank[rank].append( + (token_start, token_start + segment.length, position_start) + ) + rank_loads[rank] += segment.length + return segment.length - segment_attention_counts[_segment_key(segment)][rank] + + +def _least_loaded_rank(rank_loads: list[int]) -> int: + return min(range(len(rank_loads)), key=lambda rank: (rank_loads[rank], rank)) + + +def _build_tree_bucket_plans( + segments: tuple[GdnSegmentSpec, ...], + tree_parent_indices: tuple[int, ...], + tree_has_children: tuple[bool, ...], + *, + local_token_ranges: tuple[tuple[int, int, int], ...] | None, + sequence_length: int, + device: torch.device | str, + token_ranges_by_rank: tuple[tuple[tuple[int, int, int], ...], ...] | None = None, + split_by_final_state: bool = True, +) -> tuple[GdnSegmentBucketPlan, ...]: + segment_buckets = ( + _batch_tree_segments_by_padded_work( + segments, + tree_has_children, + ) + if split_by_final_state + else _batch_segments_by_padded_work(segments) + ) + return tuple( + _bucket_with_tree_parent_indices( + ( + _build_segment_bucket_plan(bucket, device=device) + if local_token_ranges is None + else _build_position_bucket_plan( + bucket, + local_token_ranges, + sequence_length=sequence_length, + device=device, + token_ranges_by_rank=token_ranges_by_rank, + ) + ), + bucket, + tree_parent_indices, + tree_has_children, + device=device, + ) + for bucket in segment_buckets + ) + + +def _build_chunk_aligned_cp1_tree_buckets( + spec: GdnPackedExecutionSpec, + tree_has_children: tuple[bool, ...], + *, + device: torch.device | str, + planner_config: GdnPlannerConfig, +) -> tuple[tuple[GdnSegmentBucketPlan, ...], ...]: + depth_count = max(spec.tree_depths, default=0) + 1 + children_by_node: list[list[int]] = [[] for _ in spec.tree_segments] + for node_index, parent_index in enumerate(spec.tree_parent_indices): + if parent_index >= 0: + children_by_node[parent_index].append(node_index) + + regular_by_depth: list[list[GdnSegmentSpec]] = [[] for _ in range(depth_count)] + boundary_by_depth: list[list[GdnSegmentSpec]] = [[] for _ in range(depth_count)] + child_columns_by_depth: list[list[_ExplicitBucketColumn]] = [ + [] for _ in range(depth_count) + ] + + for node_index, segment in enumerate(spec.tree_segments): + depth = spec.tree_depths[node_index] + parent_index = spec.tree_parent_indices[node_index] + if parent_index >= 0: + continue + if not tree_has_children[node_index]: + regular_by_depth[depth].append(segment) + continue + boundary_end = _prefix_chunk_boundary_end(segment) + if boundary_end > segment.start: + boundary_by_depth[depth].append( + replace(segment, start=segment.start, end=boundary_end) + ) + + for parent_index, child_indices in enumerate(children_by_node): + if not child_indices: + continue + parent = spec.tree_segments[parent_index] + parent_depth = spec.tree_depths[parent_index] + child_depth = min(parent_depth + 1, depth_count - 1) + parent_tree_parent = spec.tree_parent_indices[parent_index] + if parent_tree_parent < 0: + boundary_end = _prefix_chunk_boundary_end(parent) + tail_positions = tuple(range(boundary_end, parent.end)) + explicit_parent = parent.family_index if boundary_end > parent.start else -1 + else: + tail_positions = () + explicit_parent = parent.family_index + for child_offset, child_index in enumerate(child_indices): + child = spec.tree_segments[child_index] + child_positions = tail_positions + tuple(range(child.start, child.end)) + child_columns_by_depth[child_depth].append( + _ExplicitBucketColumn( + row_index=child.row_index, + family_index=child.family_index, + parent_index=explicit_parent, + positions=child_positions, + output_mask=( + ((child_offset == 0),) * len(tail_positions) + + (True,) * child.length + ), + needs_final_state=tree_has_children[child.family_index], + ) + ) + + return tuple( + ( + *_build_tree_bucket_plans( + tuple(boundary_by_depth[depth]), + spec.tree_parent_indices, + tree_has_children, + local_token_ranges=None, + sequence_length=spec.sequence_length, + device=device, + ), + *_build_tree_bucket_plans( + tuple(regular_by_depth[depth]), + spec.tree_parent_indices, + tree_has_children, + local_token_ranges=None, + sequence_length=spec.sequence_length, + device=device, + ), + *_build_explicit_bucket_plans( + tuple(child_columns_by_depth[depth]), + device=device, + ), + ) + for depth in range(depth_count) + ) + + +def _build_explicit_bucket_plans( + columns: tuple[_ExplicitBucketColumn, ...], + *, + device: torch.device | str, +) -> tuple[GdnSegmentBucketPlan, ...]: + return tuple( + _build_explicit_bucket_plan( + batch, + needs_final_state=any(column.needs_final_state for column in batch), + device=device, + ) + for batch in _batch_explicit_columns( + columns, + ) + ) + + +def _batch_explicit_columns( + columns: tuple[_ExplicitBucketColumn, ...], +) -> tuple[tuple[_ExplicitBucketColumn, ...], ...]: + if not columns: + return () + return (columns,) + + +def _build_explicit_bucket_plan( + columns: tuple[_ExplicitBucketColumn, ...], + *, + needs_final_state: bool, + device: torch.device | str, +) -> GdnSegmentBucketPlan: + lengths_cpu = torch.tensor([column.length for column in columns], dtype=torch.long) + max_length = int(lengths_cpu.max().item()) + row_indices_cpu = torch.zeros(max_length, len(columns), dtype=torch.long) + position_indices_cpu = torch.zeros(max_length, len(columns), dtype=torch.long) + output_mask_cpu = torch.zeros(max_length, len(columns), dtype=torch.bool) + for column_index, column in enumerate(columns): + length = column.length + row_indices_cpu[:length, column_index] = column.row_index + position_indices_cpu[:length, column_index] = torch.tensor( + column.positions, dtype=torch.long + ) + output_mask_cpu[:length, column_index] = torch.tensor( + column.output_mask, dtype=torch.bool + ) + plan = _build_bucket_plan( + tuple( + GdnSegmentSpec( + row_index=column.row_index, + family_index=column.family_index, + group_id=column.family_index, + parent_id=column.parent_index, + start=0, + end=column.length, + kind="completion", + ) + for column in columns + ), + lengths_cpu=lengths_cpu, + row_indices_cpu=row_indices_cpu, + position_indices_cpu=position_indices_cpu, + device=device, + ) + parent_indices_cpu = torch.tensor( + [column.parent_index for column in columns], dtype=torch.long + ) + return replace( + plan, + parent_indices=_move_planner_tensor(parent_indices_cpu, device), + parent_indices_cpu=parent_indices_cpu, + output_mask=_move_planner_tensor( + _pack_bucket_column_major(output_mask_cpu, lengths_cpu), + device, + ), + needs_final_state=needs_final_state, + ) + + +def _prefix_chunk_boundary_end(segment: GdnSegmentSpec) -> int: + aligned_length = (segment.length // FLA_CHUNK_SIZE) * FLA_CHUNK_SIZE + return segment.start + aligned_length + + +def _bucket_with_tree_parent_indices( + plan: GdnSegmentBucketPlan, + segments: tuple[GdnSegmentSpec, ...], + tree_parent_indices: tuple[int, ...], + tree_has_children: tuple[bool, ...], + *, + device: torch.device | str, +) -> GdnSegmentBucketPlan: + parent_indices = torch.tensor( + [tree_parent_indices[segment.family_index] for segment in segments], + dtype=torch.long, + ) + return replace( + plan, + parent_indices=_move_planner_tensor(parent_indices, device), + parent_indices_cpu=parent_indices, + needs_final_state=any( + tree_has_children[segment.family_index] for segment in segments + ), + ) + + +def _build_position_bucket_plan( + segments: tuple[GdnSegmentSpec, ...], + local_token_ranges: tuple[tuple[int, int, int], ...], + *, + sequence_length: int, + device: torch.device | str, + token_ranges_by_rank: tuple[tuple[tuple[int, int, int], ...], ...] | None = None, +) -> GdnSegmentBucketPlan: + range_positions = { + (start, end): position for start, end, position in local_token_ranges + } + starts: list[int] = [] + lengths: list[int] = [] + for segment in segments: + token_start = _segment_token_start(segment, sequence_length) + token_end = token_start + segment.length + position_start = range_positions.get((token_start, token_end)) + if position_start is None: + break + starts.append(position_start) + lengths.append(segment.length) + else: + starts_cpu = torch.tensor(starts, dtype=torch.long) + lengths_cpu = torch.tensor(lengths, dtype=torch.long) + offsets_cpu = torch.arange(max(lengths), dtype=torch.long).unsqueeze(1) + position_indices_cpu = torch.where( + offsets_cpu < lengths_cpu.unsqueeze(0), + starts_cpu.unsqueeze(0) + offsets_cpu, + torch.zeros_like(offsets_cpu), + ) + return _build_bucket_plan( + segments, + lengths_cpu=lengths_cpu, + row_indices_cpu=torch.zeros_like(position_indices_cpu), + position_indices_cpu=position_indices_cpu, + lengths_by_rank_cpu=_bucket_lengths_by_rank_cpu( + segments, + token_ranges_by_rank, + sequence_length=sequence_length, + ), + device=device, + ) + + local_positions_by_segment: list[torch.Tensor] = [] + local_range_ends = tuple(token_end for _, token_end, _ in local_token_ranges) + for segment in segments: + positions = _local_positions_for_segment( + segment, + sequence_length=sequence_length, + local_token_ranges=local_token_ranges, + local_range_ends=local_range_ends, + ) + if not int(positions.numel()): + raise ValueError( + "planned GDN bucket contains a segment with no local tokens; " + f"family={segment.family_index} kind={segment.kind}" + ) + local_positions_by_segment.append(positions) + + lengths_cpu = torch.tensor( + [int(positions.numel()) for positions in local_positions_by_segment], + dtype=torch.long, + ) + max_length = int(lengths_cpu.max().item()) + position_indices_cpu = torch.zeros(max_length, len(segments), dtype=torch.long) + for column, positions in enumerate(local_positions_by_segment): + position_indices_cpu[: int(positions.numel()), column] = positions + return _build_bucket_plan( + segments, + lengths_cpu=lengths_cpu, + row_indices_cpu=torch.zeros_like(position_indices_cpu), + position_indices_cpu=position_indices_cpu, + lengths_by_rank_cpu=_bucket_lengths_by_rank_cpu( + segments, + token_ranges_by_rank, + sequence_length=sequence_length, + ), + device=device, + ) + + +def _bucket_lengths_by_rank_cpu( + segments: tuple[GdnSegmentSpec, ...], + token_ranges_by_rank: tuple[tuple[tuple[int, int, int], ...], ...] | None, + *, + sequence_length: int, +) -> torch.Tensor | None: + if token_ranges_by_rank is None: + return None + lengths_by_rank = [] + for rank_ranges_with_positions in token_ranges_by_rank: + rank_ranges = tuple( + (start, end) for start, end, _position in rank_ranges_with_positions + ) + rank_lengths = [] + for segment in segments: + start = _segment_token_start(segment, sequence_length) + end = start + segment.length + rank_lengths.append( + sum( + max(0, min(end, range_end) - max(start, range_start)) + for range_start, range_end in rank_ranges + ) + ) + lengths_by_rank.append(rank_lengths) + return torch.tensor(lengths_by_rank, dtype=torch.long) + + +def _move_planner_tensor( + tensor: torch.Tensor, device: torch.device | str +) -> torch.Tensor: + target = torch.device(device) + if target.type == "cpu": + return tensor + return tensor.to(device=target) + + +def _batch_segments_by_padded_work( + segments: tuple[GdnSegmentSpec, ...], +) -> tuple[tuple[GdnSegmentSpec, ...], ...]: + if not segments: + return () + ordered = sorted( + segments, key=lambda segment: (segment.length, segment.family_index) + ) + return (tuple(ordered),) + + +def _batch_tree_segments_by_padded_work( + segments: tuple[GdnSegmentSpec, ...], + tree_has_children: tuple[bool, ...], +) -> tuple[tuple[GdnSegmentSpec, ...], ...]: + del tree_has_children + return _batch_segments_by_padded_work(segments) + + +def _build_segment_bucket_plan( + segments: tuple[GdnSegmentSpec, ...], *, device: torch.device | str +) -> GdnSegmentBucketPlan: + lengths_cpu = torch.tensor( + [segment.length for segment in segments], dtype=torch.long + ) + max_length = int(lengths_cpu.max().item()) + starts_cpu = torch.tensor([segment.start for segment in segments], dtype=torch.long) + rows_cpu = torch.tensor( + [segment.row_index for segment in segments], dtype=torch.long + ) + offsets_cpu = torch.arange(max_length, dtype=torch.long).unsqueeze(1) + return _build_bucket_plan( + segments, + lengths_cpu=lengths_cpu, + row_indices_cpu=rows_cpu.unsqueeze(0).expand(max_length, -1).contiguous(), + position_indices_cpu=starts_cpu.unsqueeze(0) + offsets_cpu, + device=device, + ) + + +def _build_bucket_plan( + segments: tuple[GdnSegmentSpec, ...], + *, + lengths_cpu: torch.Tensor, + row_indices_cpu: torch.Tensor, + position_indices_cpu: torch.Tensor, + device: torch.device | str, + lengths_by_rank_cpu: torch.Tensor | None = None, +) -> GdnSegmentBucketPlan: + max_length = int(lengths_cpu.max().item()) + if ( + int(row_indices_cpu.shape[0]) < max_length + or int(position_indices_cpu.shape[0]) < max_length + ): + raise ValueError("bucket index tensors are shorter than max segment length") + row_indices_cpu = _pack_bucket_column_major(row_indices_cpu, lengths_cpu) + position_indices_cpu = _pack_bucket_column_major(position_indices_cpu, lengths_cpu) + real_mask_cpu = torch.ones(int(lengths_cpu.sum().item()), dtype=torch.bool) + cu_seqlens_cpu = torch.cat( + [lengths_cpu.new_zeros(1), torch.cumsum(lengths_cpu, dim=0)] + ) + family_indices_cpu = torch.tensor( + [segment.family_index for segment in segments], + dtype=torch.long, + ) + return GdnSegmentBucketPlan( + length=max_length, + lengths=_move_planner_tensor(lengths_cpu, device), + lengths_cpu=lengths_cpu, + lengths_by_rank_cpu=lengths_by_rank_cpu, + real_mask=_move_planner_tensor(real_mask_cpu, device), + cu_seqlens=_move_planner_tensor(cu_seqlens_cpu, device), + cu_seqlens_cpu=cu_seqlens_cpu, + row_indices=_move_planner_tensor(row_indices_cpu, device), + position_indices=_move_planner_tensor(position_indices_cpu, device), + family_indices=_move_planner_tensor(family_indices_cpu, device), + family_indices_cpu=family_indices_cpu, + real_token_count_static=int(lengths_cpu.sum().item()), + ) + + +def _pack_bucket_column_major( + values_cpu: torch.Tensor, + lengths_cpu: torch.Tensor, +) -> torch.Tensor: + pieces = [ + values_cpu[: int(length), column] + for column, length in enumerate(lengths_cpu.tolist()) + if int(length) > 0 + ] + if not pieces: + return values_cpu.new_empty((0,)) + return torch.cat(pieces, dim=0).contiguous() + + +def _segment_token_start(segment: GdnSegmentSpec, sequence_length: int) -> int: + return segment.row_index * sequence_length + segment.start + + +def _range_overlap_count( + start: int, + end: int, + ranges: tuple[tuple[int, int], ...], + range_ends: tuple[int, ...], +) -> int: + count = 0 + range_index = bisect_left(range_ends, start + 1) + for range_start, range_end in ranges[range_index:]: + if range_start >= end: + break + count += min(end, range_end) - max(start, range_start) + return count + + +def _range_overlaps( + start: int, + end: int, + ranges: tuple[tuple[int, int], ...], +) -> list[tuple[int, int]]: + overlaps = [ + (max(start, range_start), min(end, range_end)) + for range_start, range_end in ranges + if max(start, range_start) < min(end, range_end) + ] + overlaps.sort() + return overlaps + + +def _local_positions_for_segment( + segment: GdnSegmentSpec, + *, + sequence_length: int, + local_token_ranges: tuple[tuple[int, int, int], ...], + local_range_ends: tuple[int, ...], +) -> torch.Tensor: + segment_start = _segment_token_start(segment, sequence_length) + segment_end = segment_start + segment.length + pieces = [] + range_index = bisect_left(local_range_ends, segment_start + 1) + for token_start, token_end, position_start in local_token_ranges[range_index:]: + if token_start >= segment_end: + break + overlap_start = max(segment_start, token_start) + overlap_end = min(segment_end, token_end) + if overlap_start >= overlap_end: + continue + pieces.append( + torch.arange( + position_start + overlap_start - token_start, + position_start + overlap_end - token_start, + dtype=torch.long, + ) + ) + if not pieces: + return torch.empty((0,), dtype=torch.long) + if len(pieces) == 1: + return pieces[0] + return torch.cat(pieces) + + +def _rank2_long_cpu(name: str, tensor: torch.Tensor) -> torch.Tensor: + if not torch.is_tensor(tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.ndim != 2: + raise ValueError(f"{name} must be rank 2 [batch, sequence], got {tensor.ndim}") + if tensor.dtype not in ( + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.long, + ): + raise TypeError(f"{name} must contain integer ids, got dtype={tensor.dtype}") + return tensor.detach().to(device="cpu", dtype=torch.long) diff --git a/src/art/megatron/gdn/gdn_shared_prefix.py b/src/art/megatron/gdn/gdn_shared_prefix.py deleted file mode 100644 index 3fb693891..000000000 --- a/src/art/megatron/gdn/gdn_shared_prefix.py +++ /dev/null @@ -1,4363 +0,0 @@ -from __future__ import annotations - -from bisect import bisect_left -from typing import Any, Literal, TypeVar - -from pydantic import BaseModel, ConfigDict, Field -import torch - -from art.megatron.context_parallel.layout_index import TokenLayoutIndex - -GdnSegmentKind = Literal["prefix", "completion"] -GdnSegmentDecisionKey = tuple[int, int, int] -# FLA's public chunk_gated_delta_rule hard-codes 64-token WY chunks. -FLA_CHUNK_SIZE = 64 -_PydanticModelT = TypeVar("_PydanticModelT", bound=BaseModel) - - -class GdnSegmentSpec(BaseModel): - """Contiguous logical GDN segment in one packed row.""" - - model_config = ConfigDict(frozen=True) - - row_index: int = Field(ge=0) - family_index: int = Field(ge=0) - group_id: int - parent_id: int - start: int = Field(ge=0) - end: int = Field(ge=1) - kind: GdnSegmentKind - child_index: int | None = Field(default=None, ge=0) - - @property - def length(self) -> int: - return self.end - self.start - - def linear_indices(self, sequence_length: int) -> tuple[int, ...]: - base = self.row_index * sequence_length - return tuple(range(base + self.start, base + self.end)) - - -class GdnPackedFamilySpec(BaseModel): - """One shared-prefix family plus child completion segments.""" - - model_config = ConfigDict(frozen=True) - - row_index: int = Field(ge=0) - family_index: int = Field(ge=0) - prefix: GdnSegmentSpec - completions: tuple[GdnSegmentSpec, ...] - - @property - def completion_count(self) -> int: - return len(self.completions) - - @property - def token_count(self) -> int: - return self.prefix.length + sum(segment.length for segment in self.completions) - - -class GdnPackedExecutionSpec(BaseModel): - """Parsed shared-prefix GDN execution metadata for a packed batch.""" - - model_config = ConfigDict(frozen=True) - - batch_size: int = Field(ge=1) - sequence_length: int = Field(ge=1) - valid_lengths: tuple[int, ...] - families: tuple[GdnPackedFamilySpec, ...] - - @property - def family_count(self) -> int: - return len(self.families) - - @property - def completion_count(self) -> int: - return sum(family.completion_count for family in self.families) - - @property - def real_token_count(self) -> int: - return sum(self.valid_lengths) - - @property - def max_segment_length(self) -> int: - lengths = [ - segment.length - for family in self.families - for segment in (family.prefix, *family.completions) - ] - return max(lengths, default=0) - - def segments(self) -> tuple[GdnSegmentSpec, ...]: - return tuple( - segment - for family in self.families - for segment in (family.prefix, *family.completions) - ) - - -_GDN_SEGMENT_SPEC_FIELDS = frozenset( - { - "row_index", - "family_index", - "group_id", - "parent_id", - "start", - "end", - "kind", - "child_index", - } -) -_GDN_PACKED_FAMILY_SPEC_FIELDS = frozenset( - { - "row_index", - "family_index", - "prefix", - "completions", - } -) - - -def _trusted_pydantic_construct( - model_type: type[_PydanticModelT], - fields_set: frozenset[str], - **values: Any, -) -> _PydanticModelT: - model = model_type.__new__(model_type) - object.__setattr__(model, "__dict__", values) - object.__setattr__(model, "__pydantic_fields_set__", fields_set) - object.__setattr__(model, "__pydantic_extra__", None) - object.__setattr__(model, "__pydantic_private__", None) - return model - - -class GdnSegmentBucketPlan(BaseModel): - """Device-local index tensors for a variable-length GDN segment batch.""" - - model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True) - - length: int = Field(ge=1) - lengths: torch.Tensor - lengths_cpu: torch.Tensor - lengths_by_rank_cpu: torch.Tensor | None = None - real_mask: torch.Tensor - cu_seqlens: torch.Tensor - cu_seqlens_cpu: torch.Tensor - row_indices: torch.Tensor - position_indices: torch.Tensor - family_indices: torch.Tensor - real_token_count_static: int = Field(ge=0) - output_mask: torch.Tensor | None = None - - @property - def segment_count(self) -> int: - return int(self.family_indices.numel()) - - @property - def real_token_count(self) -> int: - return self.real_token_count_static - - -class GdnParentStateTransferPlan(BaseModel): - """Prefix-state rows transferred from one CP rank to another.""" - - model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True) - - source_rank: int = Field(ge=0) - dest_rank: int = Field(ge=0) - family_indices: tuple[int, ...] - family_indices_tensor: torch.Tensor | None = None - - -class GdnPlannerConfig(BaseModel): - """Tunable cost coefficients for one packed-row GDN execution plan.""" - - model_config = ConfigDict(frozen=True) - - max_padding_ratio: float = Field(default=2.0, gt=1.0) - max_segments_per_batch: int = Field(default=4096, ge=1) - cp_chain_min_tokens_per_rank: int = Field(default=32, ge=1) - cp_chain_min_total_tokens: int = Field(default=32768, ge=1) - cp_chain_min_prefix_only_tokens: int = Field(default=32768, ge=1) - local_fork_launch_penalty_tokens: int = Field(default=256, ge=0) - cp_collective_latency_tokens: int = Field(default=512, ge=0) - parent_state_exchange_penalty_tokens: int = Field(default=16384, ge=0) - layout_cross_rank_token_cost: float = Field(default=6.0, ge=0.0) - rank_idle_token_cost: float = Field(default=1.0, ge=0.0) - empty_rank_penalty_tokens: int = Field(default=65536, ge=0) - max_zero_exchange_load_imbalance: float = Field(default=1.5, ge=1.0) - local_completion_rebalance_min_imbalance: float = Field(default=1.08, ge=1.0) - cp_chain_beam_width: int = Field(default=2, ge=1) - cp_chain_beam_branch_factor: int = Field(default=4, ge=1) - cp_chain_beam_candidate_limit: int = Field(default=16, ge=1) - cp_chain_beam_max_steps: int = Field(default=4, ge=0) - cp_chain_beam_min_score_delta_tokens: float = Field(default=512.0, ge=0.0) - cp_chain_min_score_delta_ms: float = Field(default=0.25, ge=0.0) - planner_local_token_ms: float = Field(default=0.00065, ge=0.0) - planner_chain_token_ms: float = Field(default=0.00055, ge=0.0) - planner_local_bucket_ms: float = Field(default=0.25, ge=0.0) - planner_chain_bucket_ms: float = Field(default=22.0, ge=0.0) - planner_local_segment_ms: float = Field(default=0.010, ge=0.0) - planner_layout_cross_rank_token_ms: float = Field(default=0.00008, ge=0.0) - planner_parent_state_exchange_base_ms: float = Field(default=40.0, ge=0.0) - planner_parent_state_exchange_ms: float = Field(default=0.5, ge=0.0) - planner_empty_rank_ms: float = Field(default=32.0, ge=0.0) - - -class GdnRankExecutionPlan(BaseModel): - """Rank-local planned execution metadata for shared-prefix GDN.""" - - model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True) - - cp_rank: int = Field(ge=0) - cp_size: int = Field(ge=1) - batch_size: int = Field(ge=1) - sequence_length: int = Field(ge=0) - packed_batch_size: int | None = Field(default=None, ge=1) - packed_sequence_length: int | None = Field(default=None, ge=1) - real_token_mask: torch.Tensor - family_count: int = Field(ge=0) - completion_count: int = Field(ge=0) - local_prefix_buckets: tuple[GdnSegmentBucketPlan, ...] = () - local_completion_buckets: tuple[GdnSegmentBucketPlan, ...] = () - ready_local_completion_buckets: tuple[GdnSegmentBucketPlan, ...] = () - remote_local_completion_buckets: tuple[GdnSegmentBucketPlan, ...] = () - chain_prefix_buckets: tuple[GdnSegmentBucketPlan, ...] = () - chain_completion_buckets: tuple[GdnSegmentBucketPlan, ...] = () - prefix_table_is_dense_ordered: bool - attention_to_gdn: Any | None = None - gdn_to_attention: Any | None = None - attention_token_ranges: tuple[tuple[int, int, int], ...] = () - gdn_token_ranges: tuple[tuple[int, int, int], ...] = () - attention_token_count: int = Field(default=0, ge=0) - gdn_token_count: int = Field(default=0, ge=0) - parent_state_exchange_family_indices: tuple[int, ...] = () - parent_state_transfers: tuple[GdnParentStateTransferPlan, ...] = () - prefix_boundary_buckets: tuple[GdnSegmentBucketPlan, ...] = () - prefix_tail_buckets: tuple[GdnSegmentBucketPlan, ...] = () - completion_with_prefix_tail_buckets: tuple[GdnSegmentBucketPlan, ...] = () - remote_prefix_tail_buckets: tuple[GdnSegmentBucketPlan, ...] = () - remote_completion_with_prefix_tail_buckets: tuple[GdnSegmentBucketPlan, ...] = () - remote_prefix_tail_exchange: Any | None = None - remote_prefix_tail_backward_exchange: Any | None = None - remote_prefix_tail_state_transfers: tuple[GdnParentStateTransferPlan, ...] = () - - @property - def attention_token_indices(self) -> tuple[int, ...]: - return _tokens_from_rank_ranges(self.attention_token_ranges) - - @property - def gdn_token_indices(self) -> tuple[int, ...]: - return _tokens_from_rank_ranges(self.gdn_token_ranges) - - -class GdnCpSegmentSchedule(BaseModel): - """CPU-side ownership and bucket schedule for one CP GDN plan.""" - - model_config = ConfigDict(frozen=True) - - gdn_token_counts_by_rank: tuple[int, ...] - gdn_token_ranges_by_rank: tuple[tuple[tuple[int, int, int], ...], ...] = () - cross_rank_token_count: int = Field(ge=0) - chain_prefix_buckets: tuple[tuple[GdnSegmentSpec, ...], ...] - chain_completion_buckets: tuple[tuple[GdnSegmentSpec, ...], ...] - local_prefix_segments_by_rank: tuple[tuple[GdnSegmentSpec, ...], ...] - local_completion_segments_by_rank: tuple[tuple[GdnSegmentSpec, ...], ...] - parent_state_exchange_family_indices: tuple[int, ...] = () - parent_state_transfers: tuple[GdnParentStateTransferPlan, ...] = () - - -class _GdnCpSegmentSearchDecision(BaseModel): - model_config = ConfigDict(frozen=True) - - chain_segment_keys: frozenset[GdnSegmentDecisionKey] - co_locate_local_families: bool - score: float - - -class _ExplicitBucketColumn(BaseModel): - model_config = ConfigDict(frozen=True) - - row_index: int - family_index: int - positions: tuple[int, ...] - output_mask: tuple[bool, ...] - - @property - def length(self) -> int: - return len(self.positions) - - -def _explicit_bucket_column( - *, - row_index: int, - family_index: int, - positions: tuple[int, ...], - output_mask: tuple[bool, ...], -) -> _ExplicitBucketColumn: - return _ExplicitBucketColumn.model_construct( - row_index=row_index, - family_index=family_index, - positions=positions, - output_mask=output_mask, - ) - - -class _AttentionLayoutIndex(BaseModel): - """Counting index for CP attention token ownership.""" - - model_config = ConfigDict(frozen=True) - - token_ranges_by_rank: tuple[tuple[tuple[int, int], ...], ...] - token_range_ends_by_rank: tuple[tuple[int, ...], ...] - range_count: int = Field(ge=0) - - -def _layout_cp_size(layout: TokenLayoutIndex) -> int: - return len(layout.token_counts_by_rank) - - -def _layout_token_count(layout: TokenLayoutIndex) -> int: - return sum(int(count) for count in layout.token_counts_by_rank) - - -def _tokens_from_rank_ranges( - ranges: tuple[tuple[int, int, int], ...], -) -> tuple[int, ...]: - return tuple(token for start, end, _ in ranges for token in range(start, end)) - - -def _token_layout_from_rank_ranges( - ranges_by_rank: tuple[tuple[tuple[int, int, int], ...], ...], -) -> TokenLayoutIndex: - return TokenLayoutIndex( - ownership_ranges_by_rank=ranges_by_rank, - token_counts_by_rank=tuple( - _ranges_token_count(ranges) for ranges in ranges_by_rank - ), - ) - - -def _ranges_token_count(ranges: tuple[tuple[int, int, int], ...]) -> int: - return sum(int(end) - int(start) for start, end, _ in ranges) - - -def build_gdn_rank_execution_plan( - spec: GdnPackedExecutionSpec, - *, - device: torch.device | str, - cp_rank: int = 0, - cp_size: int = 1, - attention_token_layout_index: TokenLayoutIndex | None = None, - cp_segment_schedule: GdnCpSegmentSchedule | None = None, - planner_config: GdnPlannerConfig | None = None, -) -> GdnRankExecutionPlan: - """Build rank-local tensor metadata from a parsed shared-prefix DAG. - - Planning is CPU-bound and must run once per packed training sequence. CP>1 - emits mixed work: native FLA CP chain buckets for long segments and local - fork buckets for short work where CP collectives would be inefficient. - """ - - planner_config = planner_config or GdnPlannerConfig() - target_device = torch.device(device) - if target_device.type != "cpu": - cpu_plan = build_gdn_rank_execution_plan( - spec, - device="cpu", - cp_rank=cp_rank, - cp_size=cp_size, - attention_token_layout_index=attention_token_layout_index, - cp_segment_schedule=cp_segment_schedule, - planner_config=planner_config, - ) - return move_gdn_rank_execution_plan_to_device(cpu_plan, target_device) - if cp_size != 1 or cp_rank != 0: - return _build_cp_rank_execution_plan( - spec, - device=device, - cp_rank=cp_rank, - cp_size=cp_size, - attention_token_layout_index=attention_token_layout_index, - cp_segment_schedule=cp_segment_schedule, - planner_config=planner_config, - ) - ( - prefix_boundary_buckets, - prefix_tail_buckets, - completion_with_prefix_tail_buckets, - ) = _build_chunk_aligned_cp1_bucket_plans( - spec, - device=device, - planner_config=planner_config, - ) - valid_lengths = torch.tensor( - spec.valid_lengths, - device=device, - dtype=torch.long, - ) - positions = torch.arange(spec.sequence_length, device=device, dtype=torch.long) - local_range_list: list[tuple[int, int, int]] = [] - local_position = 0 - for row_index, length in enumerate(spec.valid_lengths): - if length: - start = row_index * spec.sequence_length - local_range_list.append((start, start + length, local_position)) - local_position += length - local_ranges = tuple(local_range_list) - return GdnRankExecutionPlan.model_construct( - cp_rank=cp_rank, - cp_size=cp_size, - batch_size=spec.batch_size, - sequence_length=spec.sequence_length, - packed_batch_size=spec.batch_size, - packed_sequence_length=spec.sequence_length, - real_token_mask=positions.unsqueeze(0) < valid_lengths.unsqueeze(1), - family_count=spec.family_count, - completion_count=spec.completion_count, - local_prefix_buckets=(), - local_completion_buckets=(), - ready_local_completion_buckets=(), - remote_local_completion_buckets=(), - chain_prefix_buckets=(), - chain_completion_buckets=(), - prefix_table_is_dense_ordered=False, - attention_token_ranges=local_ranges, - gdn_token_ranges=local_ranges, - attention_token_count=spec.real_token_count, - gdn_token_count=spec.real_token_count, - prefix_boundary_buckets=prefix_boundary_buckets, - prefix_tail_buckets=prefix_tail_buckets, - completion_with_prefix_tail_buckets=completion_with_prefix_tail_buckets, - ) - - -def move_gdn_rank_execution_plan_to_device( - plan: GdnRankExecutionPlan, - device: torch.device | str, -) -> GdnRankExecutionPlan: - """Move planner tensors to the execution device after CPU planning.""" - - from art.megatron.gdn.layout import move_cp_exchange_plan_to_device - - return GdnRankExecutionPlan.model_construct( - cp_rank=plan.cp_rank, - cp_size=plan.cp_size, - batch_size=plan.batch_size, - sequence_length=plan.sequence_length, - packed_batch_size=plan.packed_batch_size, - packed_sequence_length=plan.packed_sequence_length, - real_token_mask=_move_planner_tensor(plan.real_token_mask, device), - family_count=plan.family_count, - completion_count=plan.completion_count, - local_prefix_buckets=_move_bucket_plans(plan.local_prefix_buckets, device), - local_completion_buckets=_move_bucket_plans( - plan.local_completion_buckets, device - ), - ready_local_completion_buckets=_move_bucket_plans( - plan.ready_local_completion_buckets, device - ), - remote_local_completion_buckets=_move_bucket_plans( - plan.remote_local_completion_buckets, device - ), - chain_prefix_buckets=_move_bucket_plans(plan.chain_prefix_buckets, device), - chain_completion_buckets=_move_bucket_plans( - plan.chain_completion_buckets, device - ), - prefix_table_is_dense_ordered=plan.prefix_table_is_dense_ordered, - attention_to_gdn=move_cp_exchange_plan_to_device(plan.attention_to_gdn, device), - gdn_to_attention=move_cp_exchange_plan_to_device(plan.gdn_to_attention, device), - attention_token_ranges=plan.attention_token_ranges, - gdn_token_ranges=plan.gdn_token_ranges, - attention_token_count=plan.attention_token_count, - gdn_token_count=plan.gdn_token_count, - parent_state_exchange_family_indices=plan.parent_state_exchange_family_indices, - parent_state_transfers=_move_parent_state_transfers( - plan.parent_state_transfers, device - ), - prefix_boundary_buckets=_move_bucket_plans( - plan.prefix_boundary_buckets, device - ), - prefix_tail_buckets=_move_bucket_plans(plan.prefix_tail_buckets, device), - completion_with_prefix_tail_buckets=_move_bucket_plans( - plan.completion_with_prefix_tail_buckets, device - ), - remote_prefix_tail_buckets=_move_bucket_plans( - plan.remote_prefix_tail_buckets, device - ), - remote_completion_with_prefix_tail_buckets=_move_bucket_plans( - plan.remote_completion_with_prefix_tail_buckets, device - ), - remote_prefix_tail_exchange=move_cp_exchange_plan_to_device( - plan.remote_prefix_tail_exchange, device - ), - remote_prefix_tail_backward_exchange=move_cp_exchange_plan_to_device( - plan.remote_prefix_tail_backward_exchange, device - ), - remote_prefix_tail_state_transfers=_move_parent_state_transfers( - plan.remote_prefix_tail_state_transfers, device - ), - ) - - -def _move_bucket_plans( - buckets: tuple[GdnSegmentBucketPlan, ...], - device: torch.device | str, -) -> tuple[GdnSegmentBucketPlan, ...]: - return tuple( - GdnSegmentBucketPlan.model_construct( - length=bucket.length, - lengths=_move_planner_tensor(bucket.lengths, device), - lengths_cpu=bucket.lengths_cpu, - lengths_by_rank_cpu=bucket.lengths_by_rank_cpu, - real_mask=_move_planner_tensor(bucket.real_mask, device), - cu_seqlens=_move_planner_tensor(bucket.cu_seqlens, device), - cu_seqlens_cpu=bucket.cu_seqlens_cpu, - row_indices=_move_planner_tensor(bucket.row_indices, device), - position_indices=_move_planner_tensor(bucket.position_indices, device), - family_indices=_move_planner_tensor(bucket.family_indices, device), - real_token_count_static=bucket.real_token_count, - output_mask=( - _move_planner_tensor(bucket.output_mask, device) - if bucket.output_mask is not None - else None - ), - ) - for bucket in buckets - ) - - -def _move_parent_state_transfers( - transfers: tuple[GdnParentStateTransferPlan, ...], - device: torch.device | str, -) -> tuple[GdnParentStateTransferPlan, ...]: - return tuple( - GdnParentStateTransferPlan.model_construct( - source_rank=transfer.source_rank, - dest_rank=transfer.dest_rank, - family_indices=transfer.family_indices, - family_indices_tensor=( - _move_planner_tensor(transfer.family_indices_tensor, device) - if transfer.family_indices_tensor is not None - else None - ), - ) - for transfer in transfers - ) - - -def _build_local_attention_layout_rank_execution_plan( - spec: GdnPackedExecutionSpec, - *, - device: torch.device | str, - cp_rank: int, - cp_size: int, - attention_token_layout_index: TokenLayoutIndex | None, - planner_config: GdnPlannerConfig, -) -> GdnRankExecutionPlan | None: - if cp_size <= 1 or not spec.families: - return None - if any( - _has_chainable_segment(family, cp_size=cp_size, planner_config=planner_config) - for family in spec.families - ): - return None - - from art.megatron.gdn.layout import ( - _reverse_exchange_plan, - build_local_rank_cp_exchange_plan_from_dest_ranges, - ) - - source_layout = _attention_source_layout( - spec, - cp_size=cp_size, - attention_token_layout_index=attention_token_layout_index, - planner_config=planner_config, - ) - attention_layout_index = _build_attention_layout_index_from_token_layout( - source_layout, - max_ranges=max(1, 2 * spec.real_token_count // len(tuple(spec.segments()))), - ) - segment_attention_counts = _segment_attention_rank_counts( - spec, - cp_size=cp_size, - attention_layout_index=attention_layout_index, - ) - best = _assign_local_attention_segments( - spec, - cp_size=cp_size, - segment_attention_counts=segment_attention_counts, - co_locate_local_families=False, - planner_config=planner_config, - ) - co_located = _assign_local_attention_segments( - spec, - cp_size=cp_size, - segment_attention_counts=segment_attention_counts, - co_locate_local_families=True, - planner_config=planner_config, - ) - if co_located[4] < best[4]: - best = co_located - ( - prefix_owner_by_family, - completion_owners_by_family, - _, - cross_rank_token_count, - _, - ) = best - - local_prefix_segments: list[GdnSegmentSpec] = [] - local_completion_segments: list[GdnSegmentSpec] = [] - prefix_segments_by_rank: list[list[GdnSegmentSpec]] = [[] for _ in range(cp_size)] - completion_segments_by_rank: list[list[GdnSegmentSpec]] = [ - [] for _ in range(cp_size) - ] - gdn_ranges_by_rank: list[list[tuple[int, int, int]]] = [[] for _ in range(cp_size)] - rank_loads = [0] * cp_size - parent_state_exchange_families: set[int] = set() - parent_state_transfer_families: dict[tuple[int, int], set[int]] = {} - - def append_segment(rank: int, segment: GdnSegmentSpec) -> None: - token_start = _segment_token_start(segment, spec.sequence_length) - position_start = rank_loads[rank] - gdn_ranges_by_rank[rank].append( - (token_start, token_start + segment.length, position_start) - ) - rank_loads[rank] += segment.length - - for family in spec.families: - prefix_owner = prefix_owner_by_family[family.family_index] - if prefix_owner == cp_rank: - local_prefix_segments.append(family.prefix) - prefix_segments_by_rank[prefix_owner].append(family.prefix) - append_segment(prefix_owner, family.prefix) - completion_owners = completion_owners_by_family[family.family_index] - for completion, completion_owner in zip( - family.completions, completion_owners, strict=True - ): - if completion_owner == cp_rank: - local_completion_segments.append(completion) - completion_segments_by_rank[completion_owner].append(completion) - append_segment(completion_owner, completion) - if completion_owner != prefix_owner: - parent_state_exchange_families.add(family.family_index) - parent_state_transfer_families.setdefault( - (prefix_owner, completion_owner), set() - ).add(family.family_index) - - local_token_ranges = tuple(gdn_ranges_by_rank[cp_rank]) - local_token_count = rank_loads[cp_rank] - schedule = GdnCpSegmentSchedule.model_construct( - gdn_token_counts_by_rank=tuple(rank_loads), - gdn_token_ranges_by_rank=tuple(tuple(ranges) for ranges in gdn_ranges_by_rank), - cross_rank_token_count=cross_rank_token_count, - chain_prefix_buckets=(), - chain_completion_buckets=(), - local_prefix_segments_by_rank=tuple( - tuple(segments) for segments in prefix_segments_by_rank - ), - local_completion_segments_by_rank=tuple( - tuple(segments) for segments in completion_segments_by_rank - ), - parent_state_exchange_family_indices=tuple( - sorted(parent_state_exchange_families) - ), - parent_state_transfers=_build_parent_state_transfer_plans( - parent_state_transfer_families - ), - ) - if parent_state_transfer_families: - ( - remote_prefix_tail_buckets, - remote_completion_with_prefix_tail_buckets, - remote_prefix_tail_exchange, - remote_prefix_tail_backward_exchange, - remote_prefix_tail_state_transfers, - remote_prefix_tail_families, - ) = _build_remote_prefix_tail_plans( - spec, - schedule, - cp_rank=cp_rank, - device=device, - planner_config=planner_config, - ) - else: - ( - remote_prefix_tail_buckets, - remote_completion_with_prefix_tail_buckets, - remote_prefix_tail_exchange, - remote_prefix_tail_backward_exchange, - remote_prefix_tail_state_transfers, - remote_prefix_tail_families, - ) = _empty_remote_prefix_tail_plans() - attention_to_gdn = build_local_rank_cp_exchange_plan_from_dest_ranges( - source_layout=source_layout, - device=device, - dest_ranges_by_rank=tuple(tuple(ranges) for ranges in gdn_ranges_by_rank), - local_rank=cp_rank, - cross_rank_token_count=cross_rank_token_count, - ) - gdn_to_attention = _reverse_exchange_plan(attention_to_gdn) - local_prefix_family_indices = { - segment.family_index for segment in local_prefix_segments - } - local_prefix_buckets = _batch_segments_by_padded_work( - (), - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - chunk_local_completion_segments = tuple( - segment - for segment in local_completion_segments - if segment.family_index in local_prefix_family_indices - ) - plain_local_completion_segments = tuple( - segment - for segment in local_completion_segments - if segment.family_index not in local_prefix_family_indices - and segment.family_index not in remote_prefix_tail_families - ) - ready_completion_segments, remote_completion_segments = ( - _split_ready_and_remote_completion_segments( - plain_local_completion_segments, - local_prefix_segments=(), - chain_prefix_buckets=(), - ) - ) - ready_completion_buckets = _batch_segments_by_padded_work( - ready_completion_segments, - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - remote_completion_buckets = _batch_segments_by_padded_work( - remote_completion_segments, - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - prefix_family_order = tuple( - segment.family_index for bucket in local_prefix_buckets for segment in bucket - ) - ready_completion_bucket_plans = _build_position_bucket_plans( - ready_completion_buckets, - local_token_ranges, - sequence_length=spec.sequence_length, - device=device, - ) - remote_completion_bucket_plans = _build_position_bucket_plans( - remote_completion_buckets, - local_token_ranges, - sequence_length=spec.sequence_length, - device=device, - ) - ( - prefix_boundary_buckets, - prefix_tail_buckets, - completion_with_prefix_tail_buckets, - ) = _build_chunk_aligned_position_bucket_plans( - tuple(local_prefix_segments), - chunk_local_completion_segments, - local_token_ranges, - sequence_length=spec.sequence_length, - device=device, - planner_config=planner_config, - ) - return GdnRankExecutionPlan.model_construct( - cp_rank=cp_rank, - cp_size=cp_size, - batch_size=1, - sequence_length=local_token_count, - packed_batch_size=spec.batch_size, - packed_sequence_length=spec.sequence_length, - real_token_mask=torch.ones( - 1, local_token_count, device=device, dtype=torch.bool - ), - family_count=spec.family_count, - completion_count=spec.completion_count, - local_prefix_buckets=_build_position_bucket_plans( - local_prefix_buckets, - local_token_ranges, - sequence_length=spec.sequence_length, - device=device, - ), - local_completion_buckets=( - ready_completion_bucket_plans + remote_completion_bucket_plans - ), - ready_local_completion_buckets=ready_completion_bucket_plans, - remote_local_completion_buckets=remote_completion_bucket_plans, - chain_prefix_buckets=(), - chain_completion_buckets=(), - prefix_table_is_dense_ordered=( - not local_prefix_segments - and prefix_family_order == tuple(range(spec.family_count)) - ), - attention_to_gdn=attention_to_gdn, - gdn_to_attention=gdn_to_attention, - attention_token_ranges=source_layout.ownership_ranges_by_rank[cp_rank], - gdn_token_ranges=local_token_ranges, - attention_token_count=source_layout.token_counts_by_rank[cp_rank], - gdn_token_count=local_token_count, - parent_state_exchange_family_indices=tuple( - sorted(parent_state_exchange_families - remote_prefix_tail_families) - ), - parent_state_transfers=_filter_parent_state_transfers( - _build_parent_state_transfer_plans(parent_state_transfer_families), - excluded_families=remote_prefix_tail_families, - device=device, - ), - prefix_boundary_buckets=prefix_boundary_buckets, - prefix_tail_buckets=prefix_tail_buckets, - completion_with_prefix_tail_buckets=completion_with_prefix_tail_buckets, - remote_prefix_tail_buckets=remote_prefix_tail_buckets, - remote_completion_with_prefix_tail_buckets=remote_completion_with_prefix_tail_buckets, - remote_prefix_tail_exchange=remote_prefix_tail_exchange, - remote_prefix_tail_backward_exchange=remote_prefix_tail_backward_exchange, - remote_prefix_tail_state_transfers=remote_prefix_tail_state_transfers, - ) - - -def _assign_local_attention_segments( - spec: GdnPackedExecutionSpec, - *, - cp_size: int, - segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], - co_locate_local_families: bool, - planner_config: GdnPlannerConfig, -) -> tuple[ - tuple[int, ...], - tuple[tuple[int, ...], ...], - tuple[int, ...], - int, - float, -]: - rank_loads = [0] * cp_size - has_prefix = [False] * cp_size - has_completion = [False] * cp_size - prefix_owner_by_family: list[int] = [] - completion_owners_by_family: list[tuple[int, ...]] = [] - parent_state_exchange_families: set[int] = set() - cross_rank_token_count = 0 - - def append_owner(rank: int, segment: GdnSegmentSpec) -> None: - nonlocal cross_rank_token_count - rank_loads[rank] += segment.length - cross_rank_token_count += ( - segment.length - segment_attention_counts[_segment_key(segment)][rank] - ) - - for family in spec.families: - if co_locate_local_families: - owner = _best_segment_owner( - (family.prefix, *family.completions), - rank_loads, - segment_attention_counts=segment_attention_counts, - planner_config=planner_config, - ) - prefix_owner_by_family.append(owner) - completion_owners = tuple(owner for _ in family.completions) - completion_owners_by_family.append(completion_owners) - has_prefix[owner] = True - for segment in (family.prefix, *family.completions): - append_owner(owner, segment) - if family.completions: - has_completion[owner] = True - continue - - prefix_owner = _best_segment_owner( - (family.prefix,), - rank_loads, - segment_attention_counts=segment_attention_counts, - planner_config=planner_config, - ) - prefix_owner_by_family.append(prefix_owner) - has_prefix[prefix_owner] = True - append_owner(prefix_owner, family.prefix) - completion_owners = [] - for completion in family.completions: - owner = _best_segment_owner( - (completion,), - rank_loads, - segment_attention_counts=segment_attention_counts, - planner_config=planner_config, - ) - completion_owners.append(owner) - has_completion[owner] = True - append_owner(owner, completion) - if owner != prefix_owner: - parent_state_exchange_families.add(family.family_index) - completion_owners_by_family.append(tuple(completion_owners)) - - del has_prefix, has_completion - score = _score_local_segment_assignment( - spec, - cp_size=cp_size, - prefix_owner_by_family=tuple(prefix_owner_by_family), - completion_owners_by_family=tuple(completion_owners_by_family), - rank_loads=tuple(rank_loads), - cross_rank_token_count=cross_rank_token_count, - parent_state_exchange_family_count=len(parent_state_exchange_families), - planner_config=planner_config, - ) - return ( - tuple(prefix_owner_by_family), - tuple(completion_owners_by_family), - tuple(sorted(parent_state_exchange_families)), - cross_rank_token_count, - score, - ) - - -def _score_local_segment_assignment( - spec: GdnPackedExecutionSpec, - *, - cp_size: int, - prefix_owner_by_family: tuple[int, ...], - completion_owners_by_family: tuple[tuple[int, ...], ...], - rank_loads: tuple[int, ...], - cross_rank_token_count: int, - parent_state_exchange_family_count: int, - planner_config: GdnPlannerConfig, -) -> float: - local_prefix_segments_by_rank: list[list[GdnSegmentSpec]] = [ - [] for _ in range(cp_size) - ] - local_completion_segments_by_rank: list[list[GdnSegmentSpec]] = [ - [] for _ in range(cp_size) - ] - for family in spec.families: - prefix_owner = prefix_owner_by_family[family.family_index] - local_prefix_segments_by_rank[prefix_owner].append(family.prefix) - completion_owners = completion_owners_by_family[family.family_index] - for completion, completion_owner in zip( - family.completions, completion_owners, strict=True - ): - local_completion_segments_by_rank[completion_owner].append(completion) - ( - local_work_by_rank, - local_bucket_count, - local_segment_count, - ) = _estimate_local_rank_kernel_work( - tuple(tuple(segments) for segments in local_prefix_segments_by_rank), - tuple(tuple(segments) for segments in local_completion_segments_by_rank), - planner_config=planner_config, - ) - return _score_cp_segment_stats( - rank_local_work=local_work_by_rank, - rank_chain_work=tuple(0 for _ in range(cp_size)), - rank_real_tokens=rank_loads, - cross_rank_token_count=cross_rank_token_count, - parent_state_exchange_family_count=parent_state_exchange_family_count, - local_bucket_count=local_bucket_count, - local_segment_count=local_segment_count, - chain_bucket_count=0, - planner_config=planner_config, - ) - - -def _can_zero_exchange_colocate_families( - spec: GdnPackedExecutionSpec, - *, - cp_size: int, - segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], -) -> bool: - for family in spec.families: - family_rank_counts = [0] * cp_size - for segment in (family.prefix, *family.completions): - segment_counts = segment_attention_counts[_segment_key(segment)] - for rank in range(cp_size): - family_rank_counts[rank] += segment_counts[rank] - if max(family_rank_counts, default=0) != family.token_count: - return False - return True - - -def parse_gdn_shared_prefix_segments( - group_ids: torch.Tensor, - parent_ids: torch.Tensor, - *, - min_completions_per_family: int = 0, -) -> GdnPackedExecutionSpec: - """Parse ART packed shared-prefix metadata into a GDN segment DAG. - - The parser is intentionally strict: GDN state routing depends on prompt-family - boundaries, so malformed metadata should fail before execution can silently - leak recurrent or conv state across siblings or independent families. - """ - - groups = _rank2_long_cpu("group_ids", group_ids) - parents = _rank2_long_cpu("parent_ids", parent_ids) - if tuple(groups.shape) != tuple(parents.shape): - raise ValueError( - "group_ids and parent_ids must have the same shape, got " - f"{tuple(groups.shape)} and {tuple(parents.shape)}" - ) - - batch_size, sequence_length = (int(groups.shape[0]), int(groups.shape[1])) - valid_lengths: list[int] = [] - families: list[GdnPackedFamilySpec] = [] - for row_index in range(batch_size): - row_group_ids = groups[row_index] - row_parent_ids = parents[row_index] - valid_length = _validate_padding_tensor( - row_index, row_group_ids, row_parent_ids - ) - valid_lengths.append(valid_length) - if valid_length == 0: - continue - families.extend( - _parse_row_tensor( - row_index=row_index, - group_ids=row_group_ids, - parent_ids=row_parent_ids, - valid_length=valid_length, - first_family_index=len(families), - min_completions_per_family=min_completions_per_family, - ) - ) - - return GdnPackedExecutionSpec( - batch_size=batch_size, - sequence_length=sequence_length, - valid_lengths=tuple(valid_lengths), - families=tuple(families), - ) - - -def _build_segment_bucket_plans( - segment_buckets: tuple[tuple[GdnSegmentSpec, ...], ...], - *, - device: torch.device | str, -) -> tuple[GdnSegmentBucketPlan, ...]: - return tuple( - _build_segment_bucket_plan(bucket[0].length, bucket, device=device) - for bucket in segment_buckets - ) - - -def _build_chunk_aligned_cp1_bucket_plans( - spec: GdnPackedExecutionSpec, - *, - device: torch.device | str, - planner_config: GdnPlannerConfig, -) -> tuple[ - tuple[GdnSegmentBucketPlan, ...], - tuple[GdnSegmentBucketPlan, ...], - tuple[GdnSegmentBucketPlan, ...], -]: - boundary_segments: list[GdnSegmentSpec] = [] - tail_segments: list[GdnSegmentSpec] = [] - completion_columns: list[_ExplicitBucketColumn] = [] - for family in spec.families: - prefix = family.prefix - boundary_end = _prefix_chunk_boundary_end(prefix) - if boundary_end > prefix.start: - boundary_segments.append( - _segment_with_bounds(prefix, prefix.start, boundary_end) - ) - prefix_tail_positions = tuple(range(boundary_end, prefix.end)) - if prefix_tail_positions and not family.completions: - tail_segments.append(_segment_with_bounds(prefix, boundary_end, prefix.end)) - for child_offset, completion in enumerate(family.completions): - completion_positions = prefix_tail_positions + tuple( - range(completion.start, completion.end) - ) - completion_columns.append( - _explicit_bucket_column( - row_index=completion.row_index, - family_index=completion.family_index, - positions=completion_positions, - output_mask=( - ((child_offset == 0),) * len(prefix_tail_positions) - + (True,) * completion.length - ), - ) - ) - boundary_buckets = _batch_segments_by_padded_work( - tuple(boundary_segments), - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - tail_buckets = _batch_segments_by_padded_work( - tuple(tail_segments), - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - completion_column_batches = _batch_explicit_bucket_columns( - tuple(completion_columns), - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - return ( - _build_segment_bucket_plans(boundary_buckets, device=device), - _build_segment_bucket_plans(tail_buckets, device=device), - _build_explicit_bucket_plans(completion_column_batches, device=device), - ) - - -def _build_chunk_aligned_position_bucket_plans( - prefix_segments: tuple[GdnSegmentSpec, ...], - completion_segments: tuple[GdnSegmentSpec, ...], - local_token_ranges: tuple[tuple[int, int, int], ...], - *, - sequence_length: int, - device: torch.device | str, - planner_config: GdnPlannerConfig, -) -> tuple[ - tuple[GdnSegmentBucketPlan, ...], - tuple[GdnSegmentBucketPlan, ...], - tuple[GdnSegmentBucketPlan, ...], -]: - local_range_ends = tuple(token_end for _, token_end, _ in local_token_ranges) - local_range_positions = { - (token_start, token_end): position_start - for token_start, token_end, position_start in local_token_ranges - } - completions_by_family: dict[int, list[GdnSegmentSpec]] = {} - for completion in completion_segments: - completions_by_family.setdefault(completion.family_index, []).append(completion) - boundary_segments: list[GdnSegmentSpec] = [] - tail_segments: list[GdnSegmentSpec] = [] - completion_columns: list[_ExplicitBucketColumn] = [] - for prefix in prefix_segments: - boundary_end = _prefix_chunk_boundary_end(prefix) - if boundary_end > prefix.start: - boundary_segments.append( - _segment_with_bounds(prefix, prefix.start, boundary_end) - ) - family_completions = tuple(completions_by_family.get(prefix.family_index, ())) - prefix_tail_positions = _local_positions_for_span( - prefix.row_index, - boundary_end, - prefix.end, - sequence_length=sequence_length, - local_token_ranges=local_token_ranges, - local_range_ends=local_range_ends, - local_range_positions=local_range_positions, - ) - if prefix_tail_positions and not family_completions: - tail_segments.append(_segment_with_bounds(prefix, boundary_end, prefix.end)) - for child_offset, completion in enumerate(family_completions): - completion_positions = _local_positions_for_span( - completion.row_index, - completion.start, - completion.end, - sequence_length=sequence_length, - local_token_ranges=local_token_ranges, - local_range_ends=local_range_ends, - local_range_positions=local_range_positions, - ) - positions = prefix_tail_positions + completion_positions - completion_columns.append( - _explicit_bucket_column( - row_index=0, - family_index=completion.family_index, - positions=positions, - output_mask=( - ((child_offset == 0),) * len(prefix_tail_positions) - + (True,) * len(completion_positions) - ), - ) - ) - boundary_buckets = _batch_segments_by_padded_work( - tuple(boundary_segments), - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - tail_buckets = _batch_segments_by_padded_work( - tuple(tail_segments), - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - completion_column_batches = _batch_explicit_bucket_columns( - tuple(completion_columns), - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - return ( - _build_position_bucket_plans( - boundary_buckets, - local_token_ranges, - sequence_length=sequence_length, - device=device, - ), - _build_position_bucket_plans( - tail_buckets, - local_token_ranges, - sequence_length=sequence_length, - device=device, - ), - _build_explicit_bucket_plans(completion_column_batches, device=device), - ) - - -def _build_remote_prefix_tail_plans( - spec: GdnPackedExecutionSpec, - schedule: GdnCpSegmentSchedule, - *, - cp_rank: int, - device: torch.device | str, - planner_config: GdnPlannerConfig, -) -> tuple[ - tuple[GdnSegmentBucketPlan, ...], - tuple[GdnSegmentBucketPlan, ...], - Any | None, - Any | None, - tuple[GdnParentStateTransferPlan, ...], - frozenset[int], -]: - from art.megatron.gdn.layout import ( - GdnCpExchangePlan, - GdnCpPeerTransfer, - _reverse_exchange_plan, - ) - - family_by_index = {family.family_index: family for family in spec.families} - prefix_owner_by_family = _prefix_owner_by_family(schedule) - source_positions_by_pair: dict[tuple[int, int], list[int]] = {} - dest_positions_by_pair: dict[tuple[int, int], list[int]] = {} - dest_counts = [0 for _ in schedule.gdn_token_counts_by_rank] - state_transfer_families: dict[tuple[int, int], set[int]] = {} - remote_tail_family_indices: set[int] = set() - local_tail_columns: list[_ExplicitBucketColumn] = [] - local_completion_columns: list[_ExplicitBucketColumn] = [] - tail_positions_by_dest_family: dict[tuple[int, int], tuple[int, ...]] = {} - local_tail_column_families: set[int] = set() - rank_ranges = schedule.gdn_token_ranges_by_rank - rank_range_ends = tuple( - tuple(end for _, end, _ in ranges) for ranges in rank_ranges - ) - rank_range_positions = tuple( - { - (token_start, token_end): position_start - for token_start, token_end, position_start in ranges - } - for ranges in rank_ranges - ) - - for dest_rank, completions in enumerate(schedule.local_completion_segments_by_rank): - for completion in completions: - source_rank = prefix_owner_by_family.get(completion.family_index) - if source_rank is None or source_rank == dest_rank: - continue - family = family_by_index[completion.family_index] - boundary_end = _prefix_chunk_boundary_end(family.prefix) - if boundary_end == family.prefix.end: - continue - dest_family = (dest_rank, family.family_index) - dest_positions = tail_positions_by_dest_family.get(dest_family) - if dest_positions is None: - source_positions = _local_positions_for_span( - family.prefix.row_index, - boundary_end, - family.prefix.end, - sequence_length=spec.sequence_length, - local_token_ranges=rank_ranges[source_rank], - local_range_ends=rank_range_ends[source_rank], - local_range_positions=rank_range_positions[source_rank], - ) - if len(source_positions) != family.prefix.end - boundary_end: - raise ValueError( - "remote prefix-tail exchange could not locate all source tokens " - f"for family {family.family_index}" - ) - dest_start = dest_counts[dest_rank] - dest_positions = tuple( - range(dest_start, dest_start + len(source_positions)) - ) - tail_positions_by_dest_family[dest_family] = dest_positions - dest_counts[dest_rank] += len(source_positions) - pair = (source_rank, dest_rank) - source_positions_by_pair.setdefault(pair, []).extend(source_positions) - dest_positions_by_pair.setdefault(pair, []).extend(dest_positions) - state_transfer_families.setdefault(pair, set()).add(family.family_index) - remote_tail_family_indices.add(family.family_index) - - if dest_rank != cp_rank: - continue - completion_positions = _local_positions_for_span( - completion.row_index, - completion.start, - completion.end, - sequence_length=spec.sequence_length, - local_token_ranges=rank_ranges[dest_rank], - local_range_ends=rank_range_ends[dest_rank], - local_range_positions=rank_range_positions[dest_rank], - ) - if len(completion_positions) != completion.length: - raise ValueError( - "remote prefix-tail bucket could not locate all completion tokens " - f"for family {family.family_index}" - ) - remote_base = int(schedule.gdn_token_counts_by_rank[dest_rank]) - if ( - len(dest_positions) > 0 - and family.family_index not in local_tail_column_families - ): - local_tail_column_families.add(family.family_index) - local_tail_columns.append( - _explicit_bucket_column( - row_index=0, - family_index=family.family_index, - positions=tuple(remote_base + pos for pos in dest_positions), - output_mask=(False,) * len(dest_positions), - ) - ) - local_completion_columns.append( - _explicit_bucket_column( - row_index=0, - family_index=family.family_index, - positions=completion_positions, - output_mask=(True,) * len(completion_positions), - ) - ) - - if not source_positions_by_pair: - return (), (), None, None, (), frozenset() - - transfers = tuple( - GdnCpPeerTransfer.model_construct( - source_rank=source_rank, - dest_rank=dest_rank, - token_count=len(source_positions), - source_positions_tensor=_move_planner_tensor( - torch.tensor(source_positions, dtype=torch.long), device - ), - dest_positions_tensor=_move_planner_tensor( - torch.tensor( - dest_positions_by_pair[(source_rank, dest_rank)], - dtype=torch.long, - ), - device, - ), - ) - for (source_rank, dest_rank), source_positions in sorted( - source_positions_by_pair.items() - ) - ) - exchange = GdnCpExchangePlan.model_construct( - cp_size=len(schedule.gdn_token_counts_by_rank), - source_token_counts_by_rank=schedule.gdn_token_counts_by_rank, - dest_token_counts_by_rank=tuple(dest_counts), - transfers=transfers, - cross_rank_token_count_override=sum(dest_counts), - ) - tail_column_batches = _batch_explicit_bucket_columns( - tuple(local_tail_columns), - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - completion_column_batches = _batch_explicit_bucket_columns( - tuple(local_completion_columns), - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - return ( - _build_explicit_bucket_plans(tail_column_batches, device=device), - _build_explicit_bucket_plans(completion_column_batches, device=device), - exchange, - _reverse_exchange_plan(exchange), - _transfer_plans_to_device( - _build_parent_state_transfer_plans(state_transfer_families), - device=device, - ), - frozenset(remote_tail_family_indices), - ) - - -def _empty_remote_prefix_tail_plans() -> tuple[ - tuple[GdnSegmentBucketPlan, ...], - tuple[GdnSegmentBucketPlan, ...], - Any | None, - Any | None, - tuple[GdnParentStateTransferPlan, ...], - frozenset[int], -]: - return (), (), None, None, (), frozenset() - - -def _prefix_owner_by_family(schedule: GdnCpSegmentSchedule) -> dict[int, int]: - owners: dict[int, int] = {} - for rank, segments in enumerate(schedule.local_prefix_segments_by_rank): - for segment in segments: - owners[segment.family_index] = rank - return owners - - -def _filter_parent_state_transfers( - transfers: tuple[GdnParentStateTransferPlan, ...], - *, - excluded_families: frozenset[int], - device: torch.device | str, -) -> tuple[GdnParentStateTransferPlan, ...]: - if not excluded_families: - return _transfer_plans_to_device(transfers, device=device) - kept: dict[tuple[int, int], set[int]] = {} - for transfer in transfers: - families = set(transfer.family_indices) - excluded_families - if families: - kept.setdefault((transfer.source_rank, transfer.dest_rank), set()).update( - families - ) - return _transfer_plans_to_device( - _build_parent_state_transfer_plans(kept), device=device - ) - - -def _local_positions_for_span( - row_index: int, - start: int, - end: int, - *, - sequence_length: int, - local_token_ranges: tuple[tuple[int, int, int], ...], - local_range_ends: tuple[int, ...], - local_range_positions: dict[tuple[int, int], int] | None = None, -) -> tuple[int, ...]: - if start == end: - return () - token_start = row_index * sequence_length + start - token_end = row_index * sequence_length + end - if local_range_positions is not None: - position_start = local_range_positions.get((token_start, token_end)) - if position_start is not None: - return tuple(range(position_start, position_start + end - start)) - range_index = bisect_left(local_range_ends, token_start + 1) - if range_index < len(local_token_ranges): - range_start, range_end, position_start = local_token_ranges[range_index] - if range_start <= token_start and token_end <= range_end: - local_start = position_start + token_start - range_start - return tuple(range(local_start, local_start + end - start)) - segment = _trusted_pydantic_construct( - GdnSegmentSpec, - _GDN_SEGMENT_SPEC_FIELDS, - row_index=row_index, - family_index=0, - group_id=0, - parent_id=0, - start=start, - end=end, - kind="prefix", - child_index=None, - ) - return tuple( - int(position) - for position in _local_positions_for_segment( - segment, - sequence_length=sequence_length, - local_token_ranges=local_token_ranges, - local_range_ends=local_range_ends, - ).tolist() - ) - - -def _prefix_chunk_boundary_end(prefix: GdnSegmentSpec) -> int: - aligned_length = (prefix.length // FLA_CHUNK_SIZE) * FLA_CHUNK_SIZE - return prefix.start + aligned_length - - -def _segment_with_bounds( - segment: GdnSegmentSpec, start: int, end: int -) -> GdnSegmentSpec: - return _trusted_pydantic_construct( - GdnSegmentSpec, - _GDN_SEGMENT_SPEC_FIELDS, - row_index=segment.row_index, - family_index=segment.family_index, - group_id=segment.group_id, - parent_id=segment.parent_id, - start=start, - end=end, - kind=segment.kind, - child_index=segment.child_index, - ) - - -def _batch_explicit_bucket_columns( - columns: tuple[_ExplicitBucketColumn, ...], - *, - max_padding_ratio: float = 1.25, - max_segments_per_batch: int = 128, -) -> tuple[tuple[_ExplicitBucketColumn, ...], ...]: - if not columns: - return () - ordered = sorted( - columns, - key=lambda column: (column.length, column.family_index, column.row_index), - ) - batches: list[list[_ExplicitBucketColumn]] = [] - current: list[_ExplicitBucketColumn] = [] - current_tokens = 0 - current_max = 0 - for column in ordered: - next_count = len(current) + 1 - next_tokens = current_tokens + column.length - next_max = max(current_max, column.length) - padded = next_max * next_count - can_extend = not current or ( - next_count <= max_segments_per_batch - and padded <= max_padding_ratio * next_tokens - ) - if not can_extend: - batches.append(current) - current = [] - current_tokens = 0 - current_max = 0 - current.append(column) - current_tokens += column.length - current_max = max(current_max, column.length) - if current: - batches.append(current) - return tuple(tuple(batch) for batch in batches) - - -def _build_explicit_bucket_plans( - bucket_columns: tuple[tuple[_ExplicitBucketColumn, ...], ...], - *, - device: torch.device | str, -) -> tuple[GdnSegmentBucketPlan, ...]: - return tuple( - _build_explicit_bucket_plan(columns, device=device) - for columns in bucket_columns - ) - - -def _build_explicit_bucket_plan( - columns: tuple[_ExplicitBucketColumn, ...], - *, - device: torch.device | str, -) -> GdnSegmentBucketPlan: - max_length = max(column.length for column in columns) - column_count = len(columns) - lengths = [column.length for column in columns] - lengths_cpu = torch.tensor(lengths, dtype=torch.long) - offsets_cpu = torch.arange(max_length, dtype=torch.long).unsqueeze(1) - real_mask_cpu = offsets_cpu < lengths_cpu.unsqueeze(0) - padded_element_count = max_length * column_count - row_indices = [0] * padded_element_count - position_indices = [0] * padded_element_count - output_mask = [False] * padded_element_count - for column_index, column in enumerate(columns): - length = column.length - column_slice = slice(column_index, length * column_count, column_count) - row_indices[column_slice] = [column.row_index] * length - position_indices[column_slice] = column.positions - output_mask[column_slice] = column.output_mask - row_indices_cpu = torch.tensor(row_indices, dtype=torch.long).reshape( - max_length, column_count - ) - position_indices_cpu = torch.tensor(position_indices, dtype=torch.long).reshape( - max_length, column_count - ) - output_mask_cpu = torch.tensor(output_mask, dtype=torch.bool).reshape( - max_length, column_count - ) - family_indices_cpu = torch.tensor( - [column.family_index for column in columns], dtype=torch.long - ) - cu_seqlens_cpu = torch.cat( - [lengths_cpu.new_zeros(1), torch.cumsum(lengths_cpu, dim=0)] - ) - return GdnSegmentBucketPlan.model_construct( - length=max_length, - lengths=_move_planner_tensor(lengths_cpu, device), - lengths_cpu=lengths_cpu, - lengths_by_rank_cpu=None, - real_mask=_move_planner_tensor(real_mask_cpu, device), - cu_seqlens=_move_planner_tensor(cu_seqlens_cpu, device), - cu_seqlens_cpu=cu_seqlens_cpu, - row_indices=_move_planner_tensor(row_indices_cpu, device), - position_indices=_move_planner_tensor(position_indices_cpu, device), - family_indices=_move_planner_tensor(family_indices_cpu, device), - real_token_count_static=int(lengths_cpu.sum().item()), - output_mask=_move_planner_tensor(output_mask_cpu, device), - ) - - -def _attention_source_layout( - spec: GdnPackedExecutionSpec, - *, - cp_size: int, - attention_token_layout_index: TokenLayoutIndex | None, - planner_config: GdnPlannerConfig, -) -> TokenLayoutIndex: - if attention_token_layout_index is not None: - if _layout_cp_size(attention_token_layout_index) != cp_size: - raise ValueError( - "attention token layout index cp_size must match GDN cp_size, got " - f"{_layout_cp_size(attention_token_layout_index)} and {cp_size}" - ) - if _layout_token_count(attention_token_layout_index) != spec.real_token_count: - raise ValueError( - "attention token layout index token count must match GDN real token " - f"count, got {_layout_token_count(attention_token_layout_index)} and " - f"{spec.real_token_count}" - ) - return attention_token_layout_index - return _token_layout_from_rank_ranges( - _default_attention_layout_ranges( - spec, - cp_size=cp_size, - planner_config=planner_config, - ) - ) - - -def _build_cp_rank_execution_plan( - spec: GdnPackedExecutionSpec, - *, - device: torch.device | str, - cp_rank: int, - cp_size: int, - attention_token_layout_index: TokenLayoutIndex | None, - cp_segment_schedule: GdnCpSegmentSchedule | None, - planner_config: GdnPlannerConfig, -) -> GdnRankExecutionPlan: - if cp_size < 1: - raise ValueError(f"cp_size must be >= 1, got {cp_size}") - if cp_rank < 0 or cp_rank >= cp_size: - raise ValueError(f"cp_rank must be in [0, {cp_size}), got {cp_rank}") - if ( - attention_token_layout_index is not None - and _layout_cp_size(attention_token_layout_index) != cp_size - ): - raise ValueError( - "attention token layout index cp_size must match GDN cp_size, got " - f"{_layout_cp_size(attention_token_layout_index)} and {cp_size}" - ) - - from art.megatron.gdn.layout import ( - _reverse_exchange_plan, - build_local_rank_cp_exchange_plan_from_dest_ranges, - ) - - has_explicit_attention_layout = attention_token_layout_index is not None - if cp_segment_schedule is None and not has_explicit_attention_layout: - local_family_plan = _build_local_family_rank_execution_plan( - spec, - device=device, - cp_rank=cp_rank, - cp_size=cp_size, - planner_config=planner_config, - ) - if local_family_plan is not None: - return local_family_plan - if cp_segment_schedule is None and has_explicit_attention_layout: - local_layout_plan = _build_local_attention_layout_rank_execution_plan( - spec, - device=device, - cp_rank=cp_rank, - cp_size=cp_size, - attention_token_layout_index=attention_token_layout_index, - planner_config=planner_config, - ) - if local_layout_plan is not None: - return local_layout_plan - - source_layout = _attention_source_layout( - spec, - cp_size=cp_size, - attention_token_layout_index=attention_token_layout_index, - planner_config=planner_config, - ) - if cp_segment_schedule is None: - schedule = _build_cp_segment_schedule( - spec, - cp_size=cp_size, - attention_layout_index=_build_attention_layout_index_from_token_layout( - source_layout, - max_ranges=max( - 1, - (2 * spec.real_token_count) // max(1, len(spec.segments())), - ), - ), - planner_config=planner_config, - ) - else: - schedule = cp_segment_schedule - if len(schedule.gdn_token_counts_by_rank) != cp_size: - raise ValueError(f"CP GDN schedule must contain {cp_size} ranks") - attention_to_gdn = build_local_rank_cp_exchange_plan_from_dest_ranges( - source_layout=source_layout, - device=device, - local_rank=cp_rank, - dest_ranges_by_rank=schedule.gdn_token_ranges_by_rank, - cross_rank_token_count=schedule.cross_rank_token_count, - ) - gdn_to_attention = _reverse_exchange_plan(attention_to_gdn) - local_token_ranges = schedule.gdn_token_ranges_by_rank[cp_rank] - local_gdn_token_count = schedule.gdn_token_counts_by_rank[cp_rank] - if schedule.parent_state_exchange_family_indices: - ( - remote_prefix_tail_buckets, - remote_completion_with_prefix_tail_buckets, - remote_prefix_tail_exchange, - remote_prefix_tail_backward_exchange, - remote_prefix_tail_state_transfers, - remote_prefix_tail_families, - ) = _build_remote_prefix_tail_plans( - spec, - schedule, - cp_rank=cp_rank, - device=device, - planner_config=planner_config, - ) - else: - ( - remote_prefix_tail_buckets, - remote_completion_with_prefix_tail_buckets, - remote_prefix_tail_exchange, - remote_prefix_tail_backward_exchange, - remote_prefix_tail_state_transfers, - remote_prefix_tail_families, - ) = _empty_remote_prefix_tail_plans() - - chain_prefix_buckets = tuple( - bucket for bucket in schedule.chain_prefix_buckets if bucket - ) - chain_completion_buckets = tuple( - bucket for bucket in schedule.chain_completion_buckets if bucket - ) - local_prefix_segments = tuple(schedule.local_prefix_segments_by_rank[cp_rank]) - local_prefix_family_indices = { - segment.family_index for segment in local_prefix_segments - } - local_prefix_buckets = _batch_segments_by_padded_work( - () if local_prefix_segments else (), - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - local_completion_segments = tuple( - schedule.local_completion_segments_by_rank[cp_rank] - ) - chunk_local_completion_segments = tuple( - segment - for segment in local_completion_segments - if segment.family_index in local_prefix_family_indices - ) - plain_local_completion_segments = tuple( - segment - for segment in local_completion_segments - if segment.family_index not in local_prefix_family_indices - and segment.family_index not in remote_prefix_tail_families - ) - ready_completion_segments, remote_completion_segments = ( - _split_ready_and_remote_completion_segments( - plain_local_completion_segments, - local_prefix_segments=(), - chain_prefix_buckets=chain_prefix_buckets, - ) - ) - ready_local_completion_buckets = _batch_segments_by_padded_work( - ready_completion_segments, - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - remote_local_completion_buckets = _batch_segments_by_padded_work( - remote_completion_segments, - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - local_completion_buckets = ( - ready_local_completion_buckets + remote_local_completion_buckets - ) - prefix_family_order = tuple( - segment.family_index - for bucket in ( - *chain_prefix_buckets, - *local_prefix_buckets, - ) - for segment in bucket - ) - ( - prefix_boundary_buckets, - prefix_tail_buckets, - completion_with_prefix_tail_buckets, - ) = _build_chunk_aligned_position_bucket_plans( - local_prefix_segments, - chunk_local_completion_segments, - local_token_ranges, - sequence_length=spec.sequence_length, - device=device, - planner_config=planner_config, - ) - return GdnRankExecutionPlan.model_construct( - cp_rank=cp_rank, - cp_size=cp_size, - batch_size=1, - sequence_length=local_gdn_token_count, - packed_batch_size=spec.batch_size, - packed_sequence_length=spec.sequence_length, - real_token_mask=torch.ones( - 1, local_gdn_token_count, device=device, dtype=torch.bool - ), - family_count=spec.family_count, - completion_count=spec.completion_count, - local_prefix_buckets=_build_position_bucket_plans( - local_prefix_buckets, - local_token_ranges, - sequence_length=spec.sequence_length, - device=device, - ), - local_completion_buckets=_build_position_bucket_plans( - local_completion_buckets, - local_token_ranges, - sequence_length=spec.sequence_length, - device=device, - ), - ready_local_completion_buckets=_build_position_bucket_plans( - ready_local_completion_buckets, - local_token_ranges, - sequence_length=spec.sequence_length, - device=device, - ), - remote_local_completion_buckets=_build_position_bucket_plans( - remote_local_completion_buckets, - local_token_ranges, - sequence_length=spec.sequence_length, - device=device, - ), - chain_prefix_buckets=_build_position_bucket_plans( - chain_prefix_buckets, - local_token_ranges, - sequence_length=spec.sequence_length, - device=device, - token_ranges_by_rank=schedule.gdn_token_ranges_by_rank, - ), - chain_completion_buckets=_build_position_bucket_plans( - chain_completion_buckets, - local_token_ranges, - sequence_length=spec.sequence_length, - device=device, - token_ranges_by_rank=schedule.gdn_token_ranges_by_rank, - ), - prefix_table_is_dense_ordered=( - not local_prefix_segments - and prefix_family_order == tuple(range(spec.family_count)) - ), - attention_to_gdn=attention_to_gdn, - gdn_to_attention=gdn_to_attention, - attention_token_ranges=source_layout.ownership_ranges_by_rank[cp_rank], - gdn_token_ranges=local_token_ranges, - attention_token_count=source_layout.token_counts_by_rank[cp_rank], - gdn_token_count=local_gdn_token_count, - parent_state_exchange_family_indices=( - tuple( - family_index - for family_index in schedule.parent_state_exchange_family_indices - if family_index not in remote_prefix_tail_families - ) - ), - parent_state_transfers=_filter_parent_state_transfers( - schedule.parent_state_transfers, - excluded_families=remote_prefix_tail_families, - device=device, - ), - prefix_boundary_buckets=prefix_boundary_buckets, - prefix_tail_buckets=prefix_tail_buckets, - completion_with_prefix_tail_buckets=completion_with_prefix_tail_buckets, - remote_prefix_tail_buckets=remote_prefix_tail_buckets, - remote_completion_with_prefix_tail_buckets=remote_completion_with_prefix_tail_buckets, - remote_prefix_tail_exchange=remote_prefix_tail_exchange, - remote_prefix_tail_backward_exchange=remote_prefix_tail_backward_exchange, - remote_prefix_tail_state_transfers=remote_prefix_tail_state_transfers, - ) - - -def build_gdn_cp_segment_schedule( - spec: GdnPackedExecutionSpec, - *, - cp_size: int, - attention_token_layout_index: TokenLayoutIndex | None = None, - planner_config: GdnPlannerConfig | None = None, -) -> GdnCpSegmentSchedule: - planner_config = planner_config or GdnPlannerConfig() - source_layout = _attention_source_layout( - spec, - cp_size=cp_size, - attention_token_layout_index=attention_token_layout_index, - planner_config=planner_config, - ) - return _build_cp_segment_schedule( - spec, - cp_size=cp_size, - attention_layout_index=_build_attention_layout_index_from_token_layout( - source_layout, - max_ranges=max( - 1, (2 * spec.real_token_count) // max(1, len(spec.segments())) - ), - ), - planner_config=planner_config, - ) - - -def _build_cp_segment_schedule( - spec: GdnPackedExecutionSpec, - *, - cp_size: int, - attention_layout_index: _AttentionLayoutIndex, - planner_config: GdnPlannerConfig, -) -> GdnCpSegmentSchedule: - segment_attention_counts = _segment_attention_rank_counts( - spec, - cp_size=cp_size, - attention_layout_index=attention_layout_index, - ) - legal_chain_segments = tuple( - segment - for family in spec.families - for segment in (family.prefix, *family.completions) - if ( - _can_chain_prefix_segment( - segment, cp_size=cp_size, planner_config=planner_config - ) - if segment.kind == "prefix" - else _can_chain_segment( - segment, cp_size=cp_size, planner_config=planner_config - ) - ) - ) - decision = _beam_search_cp_segment_schedule_decision( - spec, - cp_size=cp_size, - attention_layout_index=attention_layout_index, - segment_attention_counts=segment_attention_counts, - legal_chain_segments=legal_chain_segments, - planner_config=planner_config, - ) - return _materialize_cp_segment_schedule( - spec, - cp_size=cp_size, - attention_layout_index=attention_layout_index, - segment_attention_counts=segment_attention_counts, - chain_segment_keys=decision.chain_segment_keys, - co_locate_local_families=decision.co_locate_local_families, - planner_config=planner_config, - ) - - -def _beam_search_cp_segment_schedule_decision( - spec: GdnPackedExecutionSpec, - *, - cp_size: int, - attention_layout_index: _AttentionLayoutIndex, - segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], - legal_chain_segments: tuple[GdnSegmentSpec, ...], - planner_config: GdnPlannerConfig, -) -> _GdnCpSegmentSearchDecision: - legal_chain_keys = frozenset( - _segment_key(segment) for segment in legal_chain_segments - ) - chain_rank_counts_by_key: dict[GdnSegmentDecisionKey, tuple[int, ...]] = {} - chain_cross_rank_tokens_by_key: dict[GdnSegmentDecisionKey, int] = {} - for segment in legal_chain_segments: - key = _segment_key(segment) - ( - chain_rank_counts_by_key[key], - chain_cross_rank_tokens_by_key[key], - ) = _chain_segment_rank_counts_and_cross_rank_tokens( - segment, - spec, - cp_size=cp_size, - attention_layout_index=attention_layout_index, - ) - - score_cache: dict[ - frozenset[GdnSegmentDecisionKey], _GdnCpSegmentSearchDecision - ] = {} - - def decision_for( - chain_segment_keys: frozenset[GdnSegmentDecisionKey], - ) -> _GdnCpSegmentSearchDecision: - cached = score_cache.get(chain_segment_keys) - if cached is not None: - return cached - non_colocated_score = _score_cp_segment_decisions( - spec, - cp_size=cp_size, - segment_attention_counts=segment_attention_counts, - chain_rank_counts_by_key=chain_rank_counts_by_key, - chain_cross_rank_tokens_by_key=chain_cross_rank_tokens_by_key, - chain_segment_keys=chain_segment_keys, - co_locate_local_families=False, - planner_config=planner_config, - ) - colocated_score = _score_cp_segment_decisions( - spec, - cp_size=cp_size, - segment_attention_counts=segment_attention_counts, - chain_rank_counts_by_key=chain_rank_counts_by_key, - chain_cross_rank_tokens_by_key=chain_cross_rank_tokens_by_key, - chain_segment_keys=chain_segment_keys, - co_locate_local_families=True, - planner_config=planner_config, - ) - co_locate = colocated_score < non_colocated_score - decision = _GdnCpSegmentSearchDecision.model_construct( - chain_segment_keys=chain_segment_keys, - co_locate_local_families=co_locate, - score=colocated_score if co_locate else non_colocated_score, - ) - score_cache[chain_segment_keys] = decision - return decision - - best = decision_for(frozenset()) - beam_by_keys = {best.chain_segment_keys: best} - if legal_chain_keys: - all_chain = decision_for(legal_chain_keys) - beam_by_keys[all_chain.chain_segment_keys] = all_chain - if best.score - all_chain.score > planner_config.cp_chain_min_score_delta_ms: - best = all_chain - candidate_groups = _bounded_chain_candidate_groups( - spec, - legal_chain_segments, - segment_attention_counts=segment_attention_counts, - chain_rank_counts_by_key=chain_rank_counts_by_key, - planner_config=planner_config, - ) - beam = _best_cp_segment_search_decisions( - beam_by_keys.values(), - limit=planner_config.cp_chain_beam_width, - ) - stale_steps = 0 - for _ in range(planner_config.cp_chain_beam_max_steps): - if not candidate_groups: - break - expanded: dict[ - frozenset[GdnSegmentDecisionKey], _GdnCpSegmentSearchDecision - ] = {} - for decision in beam: - neighbors = [] - for segment_keys in _chain_beam_neighbor_groups( - decision.chain_segment_keys, - candidate_groups=candidate_groups, - branch_factor=planner_config.cp_chain_beam_branch_factor, - ): - if segment_keys.issubset(decision.chain_segment_keys): - next_keys = decision.chain_segment_keys - segment_keys - else: - next_keys = decision.chain_segment_keys | segment_keys - neighbors.append(decision_for(frozenset(next_keys))) - for neighbor in _best_cp_segment_search_decisions( - neighbors, - limit=planner_config.cp_chain_beam_branch_factor, - ): - expanded[neighbor.chain_segment_keys] = neighbor - if not expanded: - break - beam = _best_cp_segment_search_decisions( - (*beam, *expanded.values()), - limit=planner_config.cp_chain_beam_width, - ) - step_best = beam[0] - if best.score - step_best.score > planner_config.cp_chain_min_score_delta_ms: - best = step_best - stale_steps = 0 - else: - stale_steps += 1 - if stale_steps >= 2: - break - return best - - -def _chain_beam_neighbor_groups( - chain_segment_keys: frozenset[GdnSegmentDecisionKey], - *, - candidate_groups: tuple[frozenset[GdnSegmentDecisionKey], ...], - branch_factor: int, -) -> tuple[frozenset[GdnSegmentDecisionKey], ...]: - selected: list[frozenset[GdnSegmentDecisionKey]] = [] - for group in candidate_groups: - if group and not group.issubset(chain_segment_keys): - selected.append(group) - if len(selected) >= branch_factor: - return tuple(selected) - for group in reversed(candidate_groups): - if group and group.intersection(chain_segment_keys) and group not in selected: - selected.append(group) - if len(selected) >= branch_factor: - break - return tuple(selected) - - -def _best_cp_segment_search_decisions( - decisions: Any, - *, - limit: int, -) -> tuple[_GdnCpSegmentSearchDecision, ...]: - return tuple( - sorted( - decisions, - key=lambda decision: ( - decision.score, - len(decision.chain_segment_keys), - tuple(sorted(decision.chain_segment_keys)), - ), - )[:limit] - ) - - -def _bounded_chain_candidate_groups( - spec: GdnPackedExecutionSpec, - legal_chain_segments: tuple[GdnSegmentSpec, ...], - *, - segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], - chain_rank_counts_by_key: dict[GdnSegmentDecisionKey, tuple[int, ...]], - planner_config: GdnPlannerConfig, -) -> tuple[frozenset[GdnSegmentDecisionKey], ...]: - legal_key_set = frozenset(_segment_key(segment) for segment in legal_chain_segments) - if not legal_key_set: - return () - prefix_keys = frozenset( - _segment_key(family.prefix) - for family in spec.families - if _segment_key(family.prefix) in legal_key_set - ) - completion_keys = legal_key_set - prefix_keys - groups: list[frozenset[GdnSegmentDecisionKey]] = [] - for group in (legal_key_set, prefix_keys, completion_keys): - if group and group not in groups: - groups.append(group) - for group in _ranked_chain_beam_groups( - spec, - legal_chain_segments, - segment_attention_counts=segment_attention_counts, - chain_rank_counts_by_key=chain_rank_counts_by_key, - planner_config=planner_config, - ): - if group and group not in groups: - groups.append(group) - return tuple(groups[: planner_config.cp_chain_beam_candidate_limit]) - - -def _ranked_chain_beam_groups( - spec: GdnPackedExecutionSpec, - legal_chain_segments: tuple[GdnSegmentSpec, ...], - *, - segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], - chain_rank_counts_by_key: dict[GdnSegmentDecisionKey, tuple[int, ...]], - planner_config: GdnPlannerConfig, -) -> tuple[frozenset[GdnSegmentDecisionKey], ...]: - if not legal_chain_segments: - return () - priority_by_key = { - _segment_key(segment): _chain_beam_segment_priority( - segment, - segment_attention_counts=segment_attention_counts, - chain_rank_counts_by_key=chain_rank_counts_by_key, - ) - for segment in legal_chain_segments - } - legal_key_set = frozenset(priority_by_key) - groups: set[frozenset[GdnSegmentDecisionKey]] = { - frozenset((key,)) for key in legal_key_set - } - for family in spec.families: - completion_keys = frozenset( - _segment_key(completion) - for completion in family.completions - if _segment_key(completion) in legal_key_set - ) - if len(completion_keys) > 1: - groups.add(completion_keys) - family_keys = completion_keys - prefix_key = _segment_key(family.prefix) - if prefix_key in legal_key_set: - family_keys = family_keys | frozenset((prefix_key,)) - if len(family_keys) > 1: - groups.add(family_keys) - ranked = tuple( - sorted( - groups, - key=lambda group: _chain_beam_group_priority( - group, priority_by_key=priority_by_key - ), - reverse=True, - ) - ) - limit = planner_config.cp_chain_beam_candidate_limit - if len(ranked) <= limit: - return ranked - high_count = (limit + 1) // 2 - low_count = limit - high_count - selected = [*ranked[:high_count]] - for group in ranked[-low_count:]: - if group not in selected: - selected.append(group) - return tuple(selected) - - -def _chain_beam_group_priority( - group: frozenset[GdnSegmentDecisionKey], - *, - priority_by_key: dict[GdnSegmentDecisionKey, tuple[int, int, int, int]], -) -> tuple[int, int, int, int, int]: - priorities = tuple(priority_by_key[key] for key in group) - return ( - sum(priority[0] for priority in priorities), - sum(priority[1] for priority in priorities), - max((priority[2] for priority in priorities), default=0), - sum(priority[3] for priority in priorities), - len(group), - ) - - -def _chain_beam_segment_priority( - segment: GdnSegmentSpec, - *, - segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], - chain_rank_counts_by_key: dict[GdnSegmentDecisionKey, tuple[int, ...]], -) -> tuple[int, int, int, int]: - key = _segment_key(segment) - chain_max_load = max(chain_rank_counts_by_key[key], default=0) - best_attention_locality = max(segment_attention_counts[key], default=0) - chain_load_relief = segment.length - chain_max_load - minimum_local_exchange = segment.length - best_attention_locality - return ( - chain_load_relief, - segment.length, - best_attention_locality, - -minimum_local_exchange, - ) - - -def _score_cp_segment_decisions( - spec: GdnPackedExecutionSpec, - *, - cp_size: int, - segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], - chain_rank_counts_by_key: dict[GdnSegmentDecisionKey, tuple[int, ...]], - chain_cross_rank_tokens_by_key: dict[GdnSegmentDecisionKey, int], - chain_segment_keys: frozenset[GdnSegmentDecisionKey], - co_locate_local_families: bool, - planner_config: GdnPlannerConfig, -) -> float: - rank_loads = [0] * cp_size - local_prefix_segments_by_rank: list[list[GdnSegmentSpec]] = [ - [] for _ in range(cp_size) - ] - local_completion_segments_by_rank: list[list[GdnSegmentSpec]] = [ - [] for _ in range(cp_size) - ] - chain_prefix_segments: list[GdnSegmentSpec] = [] - chain_completion_segments: list[GdnSegmentSpec] = [] - parent_state_exchange_families: set[int] = set() - cross_rank_token_count = 0 - - for family in spec.families: - prefix_key = _segment_key(family.prefix) - chain_prefix = prefix_key in chain_segment_keys - local_completions = tuple( - completion - for completion in family.completions - if _segment_key(completion) not in chain_segment_keys - ) - prefix_owner: int | None = None - if chain_prefix: - chain_prefix_segments.append(family.prefix) - cross_rank_token_count += _add_chain_search_load( - rank_loads, - family.prefix, - chain_rank_counts_by_key=chain_rank_counts_by_key, - chain_cross_rank_tokens_by_key=chain_cross_rank_tokens_by_key, - ) - else: - owner_segments = ( - (family.prefix, *local_completions) - if co_locate_local_families - else (family.prefix,) - ) - prefix_owner = _best_segment_owner( - owner_segments, - rank_loads, - segment_attention_counts=segment_attention_counts, - planner_config=planner_config, - ) - local_prefix_segments_by_rank[prefix_owner].append(family.prefix) - cross_rank_token_count += _add_local_search_load( - rank_loads, - prefix_owner, - family.prefix, - segment_attention_counts=segment_attention_counts, - ) - for completion in family.completions: - completion_key = _segment_key(completion) - if completion_key in chain_segment_keys: - chain_completion_segments.append(completion) - cross_rank_token_count += _add_chain_search_load( - rank_loads, - completion, - chain_rank_counts_by_key=chain_rank_counts_by_key, - chain_cross_rank_tokens_by_key=chain_cross_rank_tokens_by_key, - ) - if not chain_prefix: - parent_state_exchange_families.add(family.family_index) - continue - if co_locate_local_families and not chain_prefix: - if prefix_owner is None: - raise RuntimeError( - "co-located local completion planning lost the prefix owner" - ) - owner = prefix_owner - else: - owner = _best_segment_owner( - (completion,), - rank_loads, - segment_attention_counts=segment_attention_counts, - planner_config=planner_config, - ) - if not chain_prefix: - if prefix_owner is None: - raise RuntimeError( - "local completion planning lost the prefix owner" - ) - if owner != prefix_owner: - parent_state_exchange_families.add(family.family_index) - local_completion_segments_by_rank[owner].append(completion) - cross_rank_token_count += _add_local_search_load( - rank_loads, - owner, - completion, - segment_attention_counts=segment_attention_counts, - ) - ( - local_work_by_rank, - local_bucket_count, - local_segment_count, - ) = _estimate_local_rank_kernel_work( - tuple(tuple(segments) for segments in local_prefix_segments_by_rank), - tuple(tuple(segments) for segments in local_completion_segments_by_rank), - planner_config=planner_config, - ) - chain_work_by_rank, chain_bucket_count = _estimate_chain_rank_kernel_work( - cp_size=cp_size, - chain_prefix_segments=tuple(chain_prefix_segments), - chain_completion_segments=tuple(chain_completion_segments), - chain_rank_counts_by_key=chain_rank_counts_by_key, - planner_config=planner_config, - ) - return _score_cp_segment_stats( - rank_local_work=local_work_by_rank, - rank_chain_work=chain_work_by_rank, - rank_real_tokens=tuple(rank_loads), - cross_rank_token_count=cross_rank_token_count, - parent_state_exchange_family_count=len(parent_state_exchange_families), - local_bucket_count=local_bucket_count, - local_segment_count=local_segment_count, - chain_bucket_count=chain_bucket_count, - planner_config=planner_config, - ) - - -def _estimate_local_rank_kernel_work( - local_prefix_segments_by_rank: tuple[tuple[GdnSegmentSpec, ...], ...], - local_completion_segments_by_rank: tuple[tuple[GdnSegmentSpec, ...], ...], - *, - planner_config: GdnPlannerConfig, -) -> tuple[tuple[int, ...], int, int]: - rank_work: list[int] = [] - rank_bucket_counts: list[int] = [] - rank_segment_counts: list[int] = [] - for prefix_segments, completion_segments in zip( - local_prefix_segments_by_rank, - local_completion_segments_by_rank, - strict=True, - ): - prefix_family_indices = {segment.family_index for segment in prefix_segments} - chunk_local_completion_segments = tuple( - segment - for segment in completion_segments - if segment.family_index in prefix_family_indices - ) - plain_local_completion_segments = tuple( - segment - for segment in completion_segments - if segment.family_index not in prefix_family_indices - ) - chunk_work, chunk_bucket_count = _estimate_chunk_aligned_local_work( - prefix_segments, - chunk_local_completion_segments, - planner_config=planner_config, - ) - completion_work, completion_bucket_count = _padded_work_from_lengths( - tuple(segment.length for segment in plain_local_completion_segments), - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - rank_work.append(chunk_work + completion_work) - rank_bucket_counts.append(chunk_bucket_count + completion_bucket_count) - rank_segment_counts.append(len(prefix_segments) + len(completion_segments)) - return ( - tuple(rank_work), - max(rank_bucket_counts, default=0), - max(rank_segment_counts, default=0), - ) - - -def _estimate_chunk_aligned_local_work( - prefix_segments: tuple[GdnSegmentSpec, ...], - completion_segments: tuple[GdnSegmentSpec, ...], - *, - planner_config: GdnPlannerConfig, -) -> tuple[int, int]: - completions_by_family: dict[int, list[GdnSegmentSpec]] = {} - for completion in completion_segments: - completions_by_family.setdefault(completion.family_index, []).append(completion) - boundary_lengths: list[int] = [] - tail_lengths: list[int] = [] - completion_column_lengths: list[int] = [] - for prefix in prefix_segments: - boundary_end = _prefix_chunk_boundary_end(prefix) - boundary_length = boundary_end - prefix.start - if boundary_length > 0: - boundary_lengths.append(boundary_length) - tail_length = prefix.end - boundary_end - family_completions = tuple(completions_by_family.get(prefix.family_index, ())) - if tail_length > 0 and not family_completions: - tail_lengths.append(tail_length) - for completion in family_completions: - completion_column_lengths.append(tail_length + completion.length) - boundary_work, boundary_bucket_count = _padded_work_from_lengths( - tuple(boundary_lengths), - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - tail_work, tail_bucket_count = _padded_work_from_lengths( - tuple(tail_lengths), - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - completion_work, completion_bucket_count = _padded_work_from_lengths( - tuple(completion_column_lengths), - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - return ( - boundary_work + tail_work + completion_work, - boundary_bucket_count + tail_bucket_count + completion_bucket_count, - ) - - -def _estimate_chain_rank_kernel_work( - *, - cp_size: int, - chain_prefix_segments: tuple[GdnSegmentSpec, ...], - chain_completion_segments: tuple[GdnSegmentSpec, ...], - chain_rank_counts_by_key: dict[GdnSegmentDecisionKey, tuple[int, ...]], - planner_config: GdnPlannerConfig, -) -> tuple[tuple[int, ...], int]: - rank_work = [0] * cp_size - bucket_count = 0 - for segments in (chain_prefix_segments, chain_completion_segments): - buckets = _batch_segments_by_padded_work( - segments, - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - bucket_count += len(buckets) - for bucket in buckets: - for rank in range(cp_size): - lengths = tuple( - chain_rank_counts_by_key[_segment_key(segment)][rank] - for segment in bucket - ) - rank_work[rank] += max(lengths, default=0) * len(lengths) - return tuple(rank_work), bucket_count - - -def _padded_work_from_lengths( - lengths: tuple[int, ...], - *, - max_padding_ratio: float, - max_segments_per_batch: int, -) -> tuple[int, int]: - if not lengths: - return 0, 0 - ordered = sorted(length for length in lengths if length > 0) - if not ordered: - return 0, 0 - bucket_count = 0 - padded_work = 0 - current_count = 0 - current_tokens = 0 - current_max = 0 - for length in ordered: - next_count = current_count + 1 - next_tokens = current_tokens + length - next_max = max(current_max, length) - next_padded = next_max * next_count - can_extend = current_count == 0 or ( - next_count <= max_segments_per_batch - and next_padded <= max_padding_ratio * next_tokens - ) - if not can_extend: - bucket_count += 1 - padded_work += current_max * current_count - current_count = 0 - current_tokens = 0 - current_max = 0 - current_count += 1 - current_tokens += length - current_max = max(current_max, length) - if current_count: - bucket_count += 1 - padded_work += current_max * current_count - return padded_work, bucket_count - - -def _add_chain_search_load( - rank_loads: list[int], - segment: GdnSegmentSpec, - *, - chain_rank_counts_by_key: dict[GdnSegmentDecisionKey, tuple[int, ...]], - chain_cross_rank_tokens_by_key: dict[GdnSegmentDecisionKey, int], -) -> int: - key = _segment_key(segment) - for rank, token_count in enumerate(chain_rank_counts_by_key[key]): - rank_loads[rank] += token_count - return chain_cross_rank_tokens_by_key[key] - - -def _add_local_search_load( - rank_loads: list[int], - rank: int, - segment: GdnSegmentSpec, - *, - segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], -) -> int: - rank_loads[rank] += segment.length - return segment.length - segment_attention_counts[_segment_key(segment)][rank] - - -def _chain_segment_rank_counts_and_cross_rank_tokens( - segment: GdnSegmentSpec, - spec: GdnPackedExecutionSpec, - *, - cp_size: int, - attention_layout_index: _AttentionLayoutIndex, -) -> tuple[tuple[int, ...], int]: - token_start = _segment_token_start(segment, spec.sequence_length) - attention_shards = _attention_contiguous_chain_shards( - token_start, - segment.length, - cp_size=cp_size, - attention_layout_index=attention_layout_index, - ) - if attention_shards is not None: - return tuple(len(shard) for shard in attention_shards), 0 - shard_lengths = _fla_aligned_chain_shard_lengths(segment.length, cp_size=cp_size) - cross_rank_tokens = 0 - start = 0 - for rank, shard_length in enumerate(shard_lengths): - end = start + shard_length - shard_start = token_start + start - cross_rank_tokens += shard_length - _attention_overlap_count( - attention_layout_index, - rank, - shard_start, - shard_start + shard_length, - ) - start = end - return shard_lengths, cross_rank_tokens - - -def _materialize_cp_segment_schedule( - spec: GdnPackedExecutionSpec, - *, - cp_size: int, - attention_layout_index: _AttentionLayoutIndex, - segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], - chain_segment_keys: frozenset[GdnSegmentDecisionKey], - co_locate_local_families: bool, - planner_config: GdnPlannerConfig, -) -> GdnCpSegmentSchedule: - gdn_ranges_by_rank: list[list[tuple[int, int, int]]] = [[] for _ in range(cp_size)] - rank_loads = [0] * cp_size - local_prefix_segments_by_rank: list[list[GdnSegmentSpec]] = [ - [] for _ in range(cp_size) - ] - local_completion_segments_by_rank: list[list[GdnSegmentSpec]] = [ - [] for _ in range(cp_size) - ] - chain_prefix_segments: list[GdnSegmentSpec] = [] - chain_completion_segments: list[GdnSegmentSpec] = [] - parent_state_exchange_families: set[int] = set() - parent_state_transfer_families: dict[tuple[int, int], set[int]] = {} - cross_rank_token_count = 0 - - for family in spec.families: - prefix_key = _segment_key(family.prefix) - chain_prefix = prefix_key in chain_segment_keys - local_completions = tuple( - completion - for completion in family.completions - if _segment_key(completion) not in chain_segment_keys - ) - prefix_owner: int | None = None - if chain_prefix: - chain_prefix_segments.append(family.prefix) - cross_rank_token_count += _append_chain_segment( - gdn_ranges_by_rank, - rank_loads, - family.prefix, - spec, - attention_layout_index=attention_layout_index, - ) - else: - owner_segments = ( - (family.prefix, *local_completions) - if co_locate_local_families - else (family.prefix,) - ) - prefix_owner = _best_segment_owner( - owner_segments, - rank_loads, - segment_attention_counts=segment_attention_counts, - planner_config=planner_config, - ) - local_prefix_segments_by_rank[prefix_owner].append(family.prefix) - cross_rank_token_count += _append_local_segment( - gdn_ranges_by_rank, - rank_loads, - prefix_owner, - family.prefix, - spec, - segment_attention_counts=segment_attention_counts, - ) - for completion in family.completions: - if _segment_key(completion) in chain_segment_keys: - chain_completion_segments.append(completion) - cross_rank_token_count += _append_chain_segment( - gdn_ranges_by_rank, - rank_loads, - completion, - spec, - attention_layout_index=attention_layout_index, - ) - if not chain_prefix: - if prefix_owner is None: - raise RuntimeError( - "local-prefix/chained-completion planning lost the prefix owner" - ) - parent_state_exchange_families.add(family.family_index) - for dest_rank in range(cp_size): - if dest_rank == prefix_owner: - continue - parent_state_transfer_families.setdefault( - (prefix_owner, dest_rank), set() - ).add(family.family_index) - continue - if co_locate_local_families and not chain_prefix: - if prefix_owner is None: - raise RuntimeError( - "co-located local completion planning lost the prefix owner" - ) - owner = prefix_owner - else: - owner = _best_segment_owner( - (completion,), - rank_loads, - segment_attention_counts=segment_attention_counts, - planner_config=planner_config, - ) - if not chain_prefix: - if prefix_owner is None: - raise RuntimeError( - "local completion planning lost the prefix owner" - ) - if owner != prefix_owner: - parent_state_exchange_families.add(family.family_index) - parent_state_transfer_families.setdefault( - (prefix_owner, owner), set() - ).add(family.family_index) - local_completion_segments_by_rank[owner].append(completion) - cross_rank_token_count += _append_local_segment( - gdn_ranges_by_rank, - rank_loads, - owner, - completion, - spec, - segment_attention_counts=segment_attention_counts, - ) - - return GdnCpSegmentSchedule.model_construct( - gdn_token_counts_by_rank=tuple(rank_loads), - gdn_token_ranges_by_rank=tuple(tuple(ranges) for ranges in gdn_ranges_by_rank), - cross_rank_token_count=cross_rank_token_count, - chain_prefix_buckets=_batch_segments_by_padded_work( - tuple(chain_prefix_segments), - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ), - chain_completion_buckets=_batch_segments_by_padded_work( - tuple(chain_completion_segments), - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ), - local_prefix_segments_by_rank=tuple( - tuple(segments) for segments in local_prefix_segments_by_rank - ), - local_completion_segments_by_rank=tuple( - tuple(segments) for segments in local_completion_segments_by_rank - ), - parent_state_exchange_family_indices=tuple( - sorted(parent_state_exchange_families) - ), - parent_state_transfers=_build_parent_state_transfer_plans( - parent_state_transfer_families - ), - ) - - -def _build_local_family_rank_execution_plan( - spec: GdnPackedExecutionSpec, - *, - device: torch.device | str, - cp_rank: int, - cp_size: int, - planner_config: GdnPlannerConfig, -) -> GdnRankExecutionPlan | None: - if cp_size <= 1 or not spec.families: - return None - target_rank_load = spec.real_token_count / cp_size - loads = [0] * cp_size - prefix_owner_by_family: list[int] = [] - completion_owners_by_family: list[tuple[int, ...]] = [] - for family in spec.families: - if _has_chainable_segment( - family, cp_size=cp_size, planner_config=planner_config - ): - return None - prefix_locality_limit = max( - planner_config.max_zero_exchange_load_imbalance * target_rank_load, - min(64.0, float(spec.real_token_count)), - ) - if family.prefix.length > prefix_locality_limit: - return None - owner = _least_loaded_rank(loads) - prefix_owner_by_family.append(owner) - completion_owners_by_family.append(tuple(owner for _ in family.completions)) - loads[owner] += family.token_count - - if max(loads, default=0) > ( - planner_config.local_completion_rebalance_min_imbalance * target_rank_load - ): - completion_owners_by_family = list( - _rebalance_local_completion_segments( - spec, - prefix_owner_by_family=tuple(prefix_owner_by_family), - completion_owners_by_family=tuple(completion_owners_by_family), - initial_loads=tuple(loads), - planner_config=planner_config, - ) - ) - rank_assignments = _materialize_local_family_rank_assignments( - spec, - cp_size=cp_size, - prefix_owner_by_family=tuple(prefix_owner_by_family), - completion_owners_by_family=tuple(completion_owners_by_family), - ) - local_token_count, local_token_ranges, prefix_segments, completion_segments = ( - rank_assignments[cp_rank] - ) - parent_state_transfer_families: dict[tuple[int, int], set[int]] = {} - for family in spec.families: - prefix_owner = prefix_owner_by_family[family.family_index] - completion_owners = completion_owners_by_family[family.family_index] - for completion_owner in sorted(set(completion_owners)): - if completion_owner == prefix_owner: - continue - parent_state_transfer_families.setdefault( - (prefix_owner, completion_owner), set() - ).add(family.family_index) - - from art.megatron.gdn.layout import GdnCpExchangePlan, GdnCpPeerTransfer - - token_counts_by_rank = tuple(assignment[0] for assignment in rank_assignments) - identity_exchange = GdnCpExchangePlan.model_construct( - cp_size=cp_size, - source_token_counts_by_rank=token_counts_by_rank, - dest_token_counts_by_rank=token_counts_by_rank, - transfers=tuple( - GdnCpPeerTransfer.model_construct( - source_rank=rank, - dest_rank=rank, - token_count=token_count, - source_positions_tensor=None, - dest_positions_tensor=None, - ) - for rank, token_count in enumerate(token_counts_by_rank) - if token_count - ), - ) - parent_state_exchange_family_indices = tuple( - sorted( - family_index - for family_indices in parent_state_transfer_families.values() - for family_index in family_indices - ) - ) - schedule = GdnCpSegmentSchedule.model_construct( - gdn_token_counts_by_rank=token_counts_by_rank, - gdn_token_ranges_by_rank=tuple( - assignment[1] for assignment in rank_assignments - ), - cross_rank_token_count=0, - chain_prefix_buckets=(), - chain_completion_buckets=(), - local_prefix_segments_by_rank=tuple( - assignment[2] for assignment in rank_assignments - ), - local_completion_segments_by_rank=tuple( - assignment[3] for assignment in rank_assignments - ), - parent_state_exchange_family_indices=parent_state_exchange_family_indices, - parent_state_transfers=_build_parent_state_transfer_plans( - parent_state_transfer_families - ), - ) - if parent_state_exchange_family_indices: - ( - remote_prefix_tail_buckets, - remote_completion_with_prefix_tail_buckets, - remote_prefix_tail_exchange, - remote_prefix_tail_backward_exchange, - remote_prefix_tail_state_transfers, - remote_prefix_tail_families, - ) = _build_remote_prefix_tail_plans( - spec, - schedule, - cp_rank=cp_rank, - device=device, - planner_config=planner_config, - ) - else: - ( - remote_prefix_tail_buckets, - remote_completion_with_prefix_tail_buckets, - remote_prefix_tail_exchange, - remote_prefix_tail_backward_exchange, - remote_prefix_tail_state_transfers, - remote_prefix_tail_families, - ) = _empty_remote_prefix_tail_plans() - local_prefix_family_indices = {segment.family_index for segment in prefix_segments} - chunk_local_completion_segments = tuple( - segment - for segment in completion_segments - if segment.family_index in local_prefix_family_indices - ) - suffix_only_completion_segments = tuple( - segment - for segment in completion_segments - if segment.family_index not in local_prefix_family_indices - and segment.family_index not in remote_prefix_tail_families - ) - ready_completion_segments, remote_completion_segments = ( - _split_ready_and_remote_completion_segments( - suffix_only_completion_segments, - local_prefix_segments=(), - chain_prefix_buckets=(), - ) - ) - ready_completion_buckets = _batch_segments_by_padded_work( - ready_completion_segments, - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - remote_completion_buckets = _batch_segments_by_padded_work( - remote_completion_segments, - max_padding_ratio=planner_config.max_padding_ratio, - max_segments_per_batch=planner_config.max_segments_per_batch, - ) - ready_completion_bucket_plans = _build_position_bucket_plans( - ready_completion_buckets, - local_token_ranges, - sequence_length=spec.sequence_length, - device=device, - ) - remote_completion_bucket_plans = _build_position_bucket_plans( - remote_completion_buckets, - local_token_ranges, - sequence_length=spec.sequence_length, - device=device, - ) - local_completion_bucket_plans = ( - ready_completion_bucket_plans + remote_completion_bucket_plans - ) - ( - prefix_boundary_buckets, - prefix_tail_buckets, - completion_with_prefix_tail_buckets, - ) = _build_chunk_aligned_position_bucket_plans( - prefix_segments, - chunk_local_completion_segments, - local_token_ranges, - sequence_length=spec.sequence_length, - device=device, - planner_config=planner_config, - ) - return GdnRankExecutionPlan.model_construct( - cp_rank=cp_rank, - cp_size=cp_size, - batch_size=1, - sequence_length=local_token_count, - packed_batch_size=spec.batch_size, - packed_sequence_length=spec.sequence_length, - real_token_mask=torch.ones( - 1, local_token_count, device=device, dtype=torch.bool - ), - family_count=spec.family_count, - completion_count=spec.completion_count, - local_prefix_buckets=(), - local_completion_buckets=local_completion_bucket_plans, - ready_local_completion_buckets=ready_completion_bucket_plans, - remote_local_completion_buckets=remote_completion_bucket_plans, - chain_prefix_buckets=(), - chain_completion_buckets=(), - prefix_table_is_dense_ordered=( - tuple(segment.family_index for segment in prefix_segments) - == tuple(range(spec.family_count)) - ), - attention_to_gdn=identity_exchange, - gdn_to_attention=identity_exchange, - attention_token_ranges=local_token_ranges, - gdn_token_ranges=local_token_ranges, - attention_token_count=local_token_count, - gdn_token_count=local_token_count, - parent_state_exchange_family_indices=tuple( - family_index - for family_index in parent_state_exchange_family_indices - if family_index not in remote_prefix_tail_families - ), - parent_state_transfers=_filter_parent_state_transfers( - _build_parent_state_transfer_plans(parent_state_transfer_families), - excluded_families=remote_prefix_tail_families, - device=device, - ), - prefix_boundary_buckets=prefix_boundary_buckets, - prefix_tail_buckets=prefix_tail_buckets, - completion_with_prefix_tail_buckets=completion_with_prefix_tail_buckets, - remote_prefix_tail_buckets=remote_prefix_tail_buckets, - remote_completion_with_prefix_tail_buckets=remote_completion_with_prefix_tail_buckets, - remote_prefix_tail_exchange=remote_prefix_tail_exchange, - remote_prefix_tail_backward_exchange=remote_prefix_tail_backward_exchange, - remote_prefix_tail_state_transfers=remote_prefix_tail_state_transfers, - ) - - -def _rebalance_local_completion_segments( - spec: GdnPackedExecutionSpec, - *, - prefix_owner_by_family: tuple[int, ...], - completion_owners_by_family: tuple[tuple[int, ...], ...], - initial_loads: tuple[int, ...], - planner_config: GdnPlannerConfig, -) -> tuple[tuple[int, ...], ...]: - owners = [list(family_owners) for family_owners in completion_owners_by_family] - loads = list(initial_loads) - remote_owners_by_family = [ - { - owner - for owner in family_owners - if owner != prefix_owner_by_family[family_index] - } - for family_index, family_owners in enumerate(owners) - ] - transfer_count = sum( - len(remote_owners) for remote_owners in remote_owners_by_family - ) - - def score(candidate_loads: list[int], candidate_transfer_count: int) -> float: - max_load = max(candidate_loads, default=0) - idle_tokens = sum(max_load - load for load in candidate_loads) - return ( - max_load - + planner_config.rank_idle_token_cost * idle_tokens - + planner_config.parent_state_exchange_penalty_tokens - * candidate_transfer_count - ) - - best_score = score(loads, transfer_count) - while True: - best_move: ( - tuple[int, int, int, tuple[int, ...], list[int], int, float] | None - ) = None - for family in spec.families: - family_owners = owners[family.family_index] - prefix_owner = prefix_owner_by_family[family.family_index] - original_remote_owners = remote_owners_by_family[family.family_index] - for source in sorted(set(family_owners)): - source_children = [ - child_index - for child_index, owner in enumerate(family_owners) - if owner == source - ] - ordered_children = sorted( - source_children, - key=lambda child_index: family.completions[child_index].length, - reverse=True, - ) - for dest in range(len(loads)): - if dest == source: - continue - moved_tokens = 0 - moved_children = [] - for child_index in ordered_children: - moved_tokens += family.completions[child_index].length - moved_children.append(child_index) - candidate_loads = list(loads) - candidate_loads[source] -= moved_tokens - candidate_loads[dest] += moved_tokens - candidate_remote_owners = set(original_remote_owners) - if source != prefix_owner and len(moved_children) == len( - source_children - ): - candidate_remote_owners.discard(source) - if dest != prefix_owner: - candidate_remote_owners.add(dest) - candidate_transfer_count = ( - transfer_count - - len(original_remote_owners) - + len(candidate_remote_owners) - ) - candidate_score = score( - candidate_loads, candidate_transfer_count - ) - if candidate_score >= best_score: - continue - if best_move is None or candidate_score < best_move[-1]: - best_move = ( - family.family_index, - source, - dest, - tuple(moved_children), - candidate_loads, - candidate_transfer_count, - candidate_score, - ) - if best_move is None: - return tuple(tuple(item) for item in owners) - ( - family_index, - _source, - dest, - moved_children, - loads, - transfer_count, - best_score, - ) = best_move - for child_index in moved_children: - owners[family_index][child_index] = dest - prefix_owner = prefix_owner_by_family[family_index] - remote_owners_by_family[family_index] = { - owner for owner in set(owners[family_index]) if owner != prefix_owner - } - - -def _materialize_local_family_rank_assignments( - spec: GdnPackedExecutionSpec, - *, - cp_size: int, - prefix_owner_by_family: tuple[int, ...], - completion_owners_by_family: tuple[tuple[int, ...], ...], -) -> tuple[ - tuple[ - int, - tuple[tuple[int, int, int], ...], - tuple[GdnSegmentSpec, ...], - tuple[GdnSegmentSpec, ...], - ], - ..., -]: - token_ranges_by_rank: list[list[tuple[int, int, int]]] = [ - [] for _ in range(cp_size) - ] - token_counts_by_rank = [0] * cp_size - prefix_segments_by_rank: list[list[GdnSegmentSpec]] = [[] for _ in range(cp_size)] - completion_segments_by_rank: list[list[GdnSegmentSpec]] = [ - [] for _ in range(cp_size) - ] - sequence_length = spec.sequence_length - for family in spec.families: - prefix_owner = prefix_owner_by_family[family.family_index] - prefix_segments_by_rank[prefix_owner].append(family.prefix) - prefix_token_start = ( - family.prefix.row_index * sequence_length + family.prefix.start - ) - prefix_position_start = token_counts_by_rank[prefix_owner] - token_ranges_by_rank[prefix_owner].append( - ( - prefix_token_start, - prefix_token_start + family.prefix.length, - prefix_position_start, - ) - ) - token_counts_by_rank[prefix_owner] = ( - prefix_position_start + family.prefix.length - ) - for completion, completion_owner in zip( - family.completions, - completion_owners_by_family[family.family_index], - strict=True, - ): - completion_segments_by_rank[completion_owner].append(completion) - completion_token_start = ( - completion.row_index * sequence_length + completion.start - ) - completion_position_start = token_counts_by_rank[completion_owner] - token_ranges_by_rank[completion_owner].append( - ( - completion_token_start, - completion_token_start + completion.length, - completion_position_start, - ) - ) - token_counts_by_rank[completion_owner] = ( - completion_position_start + completion.length - ) - return tuple( - ( - token_counts_by_rank[rank], - tuple(token_ranges_by_rank[rank]), - tuple(prefix_segments_by_rank[rank]), - tuple(completion_segments_by_rank[rank]), - ) - for rank in range(cp_size) - ) - - -def _empty_local_family_rank_execution_plan( - spec: GdnPackedExecutionSpec, - *, - device: torch.device | str, - cp_rank: int, - cp_size: int, -) -> GdnRankExecutionPlan: - from art.megatron.gdn.layout import GdnCpExchangePlan - - identity_exchange = GdnCpExchangePlan.model_construct( - cp_size=cp_size, - source_token_counts_by_rank=tuple(0 for _ in range(cp_size)), - dest_token_counts_by_rank=tuple(0 for _ in range(cp_size)), - transfers=(), - ) - return GdnRankExecutionPlan.model_construct( - cp_rank=cp_rank, - cp_size=cp_size, - batch_size=1, - sequence_length=0, - packed_batch_size=spec.batch_size, - packed_sequence_length=spec.sequence_length, - real_token_mask=torch.ones(1, 0, device=device, dtype=torch.bool), - family_count=spec.family_count, - completion_count=spec.completion_count, - local_prefix_buckets=(), - local_completion_buckets=(), - ready_local_completion_buckets=(), - remote_local_completion_buckets=(), - chain_prefix_buckets=(), - chain_completion_buckets=(), - prefix_table_is_dense_ordered=False, - attention_to_gdn=identity_exchange, - gdn_to_attention=identity_exchange, - attention_token_ranges=(), - gdn_token_ranges=(), - attention_token_count=0, - gdn_token_count=0, - parent_state_exchange_family_indices=(), - parent_state_transfers=(), - ) - - -def _can_chain_segment( - segment: GdnSegmentSpec, - *, - cp_size: int, - planner_config: GdnPlannerConfig, -) -> bool: - min_tokens = ( - planner_config.cp_chain_min_prefix_only_tokens - if segment.kind == "prefix" - else planner_config.cp_chain_min_total_tokens - ) - if segment.length < min_tokens: - return False - if segment.length < cp_size: - return False - if segment.length // FLA_CHUNK_SIZE < cp_size: - return False - per_rank = segment.length / cp_size - if per_rank < planner_config.cp_chain_min_tokens_per_rank: - return False - return True - - -def _build_parent_state_transfer_plans( - families_by_peer: dict[tuple[int, int], set[int]], -) -> tuple[GdnParentStateTransferPlan, ...]: - return tuple( - GdnParentStateTransferPlan( - source_rank=source_rank, - dest_rank=dest_rank, - family_indices=tuple(sorted(family_indices)), - ) - for (source_rank, dest_rank), family_indices in sorted(families_by_peer.items()) - if source_rank != dest_rank and family_indices - ) - - -def _split_ready_and_remote_completion_segments( - completion_segments: tuple[GdnSegmentSpec, ...], - *, - local_prefix_segments: tuple[GdnSegmentSpec, ...], - chain_prefix_buckets: tuple[tuple[GdnSegmentSpec, ...], ...], -) -> tuple[tuple[GdnSegmentSpec, ...], tuple[GdnSegmentSpec, ...]]: - ready_family_indices = { - segment.family_index for segment in local_prefix_segments - } | {segment.family_index for bucket in chain_prefix_buckets for segment in bucket} - ready = [] - remote = [] - for segment in completion_segments: - if segment.family_index in ready_family_indices: - ready.append(segment) - else: - remote.append(segment) - return tuple(ready), tuple(remote) - - -def _transfer_plans_to_device( - transfers: tuple[GdnParentStateTransferPlan, ...], - *, - device: torch.device | str, -) -> tuple[GdnParentStateTransferPlan, ...]: - return tuple( - transfer.model_copy( - update={ - "family_indices_tensor": _move_planner_tensor( - torch.tensor(transfer.family_indices, dtype=torch.long), - device, - ) - } - ) - for transfer in transfers - ) - - -def _has_chainable_segment( - family: GdnPackedFamilySpec, - *, - cp_size: int, - planner_config: GdnPlannerConfig, -) -> bool: - return _can_chain_prefix_segment( - family.prefix, cp_size=cp_size, planner_config=planner_config - ) or any( - _can_chain_segment(completion, cp_size=cp_size, planner_config=planner_config) - for completion in family.completions - ) - - -def _can_chain_prefix_segment( - segment: GdnSegmentSpec, - *, - cp_size: int, - planner_config: GdnPlannerConfig, -) -> bool: - return _can_chain_segment(segment, cp_size=cp_size, planner_config=planner_config) - - -def _score_cp_segment_stats( - *, - rank_local_work: tuple[int, ...], - rank_chain_work: tuple[int, ...], - rank_real_tokens: tuple[int, ...], - cross_rank_token_count: int, - parent_state_exchange_family_count: int, - local_bucket_count: int, - local_segment_count: int, - chain_bucket_count: int, - planner_config: GdnPlannerConfig, -) -> float: - empty_rank_count = sum(1 for token_count in rank_real_tokens if token_count == 0) - return ( - _rank_kernel_ms( - rank_local_work, - rank_chain_work, - local_token_ms=planner_config.planner_local_token_ms, - chain_token_ms=planner_config.planner_chain_token_ms, - ) - + planner_config.planner_local_bucket_ms * local_bucket_count - + planner_config.planner_chain_bucket_ms * chain_bucket_count - + planner_config.planner_local_segment_ms * local_segment_count - + planner_config.planner_layout_cross_rank_token_ms * cross_rank_token_count - + ( - planner_config.planner_parent_state_exchange_base_ms - + planner_config.planner_parent_state_exchange_ms - * parent_state_exchange_family_count - if parent_state_exchange_family_count - else 0.0 - ) - + planner_config.planner_empty_rank_ms * empty_rank_count - ) - - -def _rank_kernel_ms( - rank_local_work: tuple[int, ...], - rank_chain_work: tuple[int, ...], - *, - local_token_ms: float, - chain_token_ms: float, -) -> float: - return max( - ( - local_work * local_token_ms + chain_work * chain_token_ms - for local_work, chain_work in zip( - rank_local_work, rank_chain_work, strict=True - ) - ), - default=0.0, - ) - - -def _best_segment_owner( - segments: tuple[GdnSegmentSpec, ...], - rank_loads: list[int], - *, - segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], - planner_config: GdnPlannerConfig, -) -> int: - segment_length = sum(segment.length for segment in segments) - if len(segments) == 1: - on_rank_tokens = segment_attention_counts[_segment_key(segments[0])] - else: - rank_count = len(rank_loads) - counts_by_rank = [0] * rank_count - for segment in segments: - segment_counts = segment_attention_counts[_segment_key(segment)] - for rank in range(rank_count): - counts_by_rank[rank] += segment_counts[rank] - on_rank_tokens = tuple(counts_by_rank) - best: tuple[float, int, int, int, int] | None = None - for rank, tokens in enumerate(on_rank_tokens): - projected_loads = list(rank_loads) - projected_loads[rank] += segment_length - max_load = max(projected_loads, default=0) - idle_tokens = sum(max_load - load for load in projected_loads) - cross_rank_tokens = segment_length - int(tokens) - empty_rank_count = sum(1 for load in projected_loads if load == 0) - score = ( - max_load * planner_config.planner_local_token_ms - + idle_tokens - * planner_config.rank_idle_token_cost - * planner_config.planner_local_token_ms - + cross_rank_tokens * planner_config.planner_layout_cross_rank_token_ms - + empty_rank_count * planner_config.planner_empty_rank_ms - ) - candidate = ( - score, - max_load, - cross_rank_tokens, - -int(tokens), - rank, - ) - if best is None or candidate < best: - best = candidate - if best is None: - return _least_loaded_rank(rank_loads) - return best[-1] - - -def _build_attention_layout_index_from_token_layout( - layout: TokenLayoutIndex, - *, - max_ranges: int, -) -> _AttentionLayoutIndex: - del max_ranges - ranges_by_rank = tuple( - tuple(sorted((int(start), int(end)) for start, end, _ in rank_ranges)) - for rank_ranges in layout.ownership_ranges_by_rank - ) - range_count = sum(len(ranges) for ranges in ranges_by_rank) - return _AttentionLayoutIndex.model_construct( - token_ranges_by_rank=ranges_by_rank, - token_range_ends_by_rank=tuple( - tuple(end for _, end in ranges) for ranges in ranges_by_rank - ), - range_count=range_count, - ) - - -def _segment_attention_rank_counts( - spec: GdnPackedExecutionSpec, - *, - cp_size: int, - attention_layout_index: _AttentionLayoutIndex, -) -> dict[tuple[int, int, int], tuple[int, ...]]: - del cp_size - segments = tuple(spec.segments()) - if not segments: - return {} - starts = torch.tensor( - [_segment_token_start(segment, spec.sequence_length) for segment in segments], - dtype=torch.long, - ) - lengths = torch.tensor([segment.length for segment in segments], dtype=torch.long) - ends = starts + lengths - counts_by_rank = [] - for ranges in attention_layout_index.token_ranges_by_rank: - counts_by_rank.append(_rank_range_overlap_counts(starts, ends, ranges)) - counts_tensor = torch.stack(counts_by_rank, dim=1) - totals = counts_tensor.sum(dim=1) - if not torch.equal(totals, lengths): - bad_index = int(torch.nonzero(totals != lengths, as_tuple=False)[0].item()) - raise ValueError( - "attention layout is missing a real token required by GDN; " - f"segment={_segment_key(segments[bad_index])}" - ) - counts = counts_tensor.tolist() - return { - _segment_key(segment): tuple(int(value) for value in counts[index]) - for index, segment in enumerate(segments) - } - - -def _rank_range_overlap_counts( - starts: torch.Tensor, - ends: torch.Tensor, - ranges: tuple[tuple[int, int], ...], -) -> torch.Tensor: - if not ranges: - return torch.zeros_like(starts) - range_starts = torch.tensor([start for start, _ in ranges], dtype=torch.long) - range_ends = torch.tensor([end for _, end in ranges], dtype=torch.long) - range_lengths = range_ends - range_starts - prefix = torch.cat((range_lengths.new_zeros(1), torch.cumsum(range_lengths, dim=0))) - - def owned_before(points: torch.Tensor) -> torch.Tensor: - indices = torch.searchsorted(range_ends, points, right=False) - counts = prefix.index_select(0, indices) - active = indices < int(range_starts.numel()) - if bool(active.any().item()): - active_indices = indices[active] - active_starts = range_starts.index_select(0, active_indices) - active_ends = range_ends.index_select(0, active_indices) - counts[active] += torch.minimum( - torch.clamp(points[active] - active_starts, min=0), - active_ends - active_starts, - ) - return counts - - return owned_before(ends) - owned_before(starts) - - -def _segment_key(segment: GdnSegmentSpec) -> tuple[int, int, int]: - return (segment.row_index, segment.start, segment.end) - - -def _default_attention_layout_ranges( - spec: GdnPackedExecutionSpec, - *, - cp_size: int, - planner_config: GdnPlannerConfig, -) -> tuple[tuple[tuple[int, int, int], ...], ...]: - ranks: list[list[tuple[int, int, int]]] = [[] for _ in range(cp_size)] - loads = [0] * cp_size - target_rank_load = spec.real_token_count / cp_size - - def append_segment(rank: int, token_start: int, token_count: int) -> None: - ranks[rank].append((token_start, token_start + token_count, loads[rank])) - loads[rank] += token_count - - def should_split_segment(segment: GdnSegmentSpec) -> bool: - if segment.length <= planner_config.max_zero_exchange_load_imbalance * ( - target_rank_load - ): - return False - if segment.kind == "prefix": - return _can_chain_prefix_segment( - segment, cp_size=cp_size, planner_config=planner_config - ) - return _can_chain_segment( - segment, cp_size=cp_size, planner_config=planner_config - ) - - for family in spec.families: - has_split_segment = any( - should_split_segment(segment) - for segment in (family.prefix, *family.completions) - ) - if not has_split_segment: - if _should_co_locate_non_chain_family( - family, - total_real_tokens=spec.real_token_count, - cp_size=cp_size, - planner_config=planner_config, - ): - owner = _least_loaded_rank(loads) - for segment in (family.prefix, *family.completions): - token_start = _segment_token_start(segment, spec.sequence_length) - append_segment(owner, token_start, segment.length) - continue - for segment in (family.prefix, *family.completions): - token_start = _segment_token_start(segment, spec.sequence_length) - owner = _least_loaded_rank(loads) - append_segment(owner, token_start, segment.length) - continue - for segment in (family.prefix, *family.completions): - token_start = _segment_token_start(segment, spec.sequence_length) - if should_split_segment(segment): - _append_split_default_attention_segment( - ranks, loads, token_start, segment.length - ) - continue - owner = _least_loaded_rank(loads) - append_segment(owner, token_start, segment.length) - return tuple(tuple(ranges) for ranges in ranks) - - -def _should_co_locate_non_chain_family( - family: GdnPackedFamilySpec, - *, - total_real_tokens: int, - cp_size: int, - planner_config: GdnPlannerConfig, -) -> bool: - target_rank_load = total_real_tokens / cp_size - return family.token_count <= ( - planner_config.max_zero_exchange_load_imbalance * target_rank_load - ) - - -def _append_split_default_attention_segment( - ranks: list[list[tuple[int, int, int]]], - loads: list[int], - token_start: int, - token_count: int, -) -> None: - cp_size = len(ranks) - for rank in range(cp_size): - start = (token_count * rank) // cp_size - end = (token_count * (rank + 1)) // cp_size - ranks[rank].append((token_start + start, token_start + end, loads[rank])) - loads[rank] += end - start - - -def _append_chain_segment( - gdn_ranges_by_rank: list[list[tuple[int, int, int]]], - rank_loads: list[int], - segment: GdnSegmentSpec, - spec: GdnPackedExecutionSpec, - *, - attention_layout_index: _AttentionLayoutIndex | None = None, -) -> int: - token_start = _segment_token_start(segment, spec.sequence_length) - cp_size = len(gdn_ranges_by_rank) - attention_shards = _attention_contiguous_chain_shards( - token_start, - segment.length, - cp_size=cp_size, - attention_layout_index=attention_layout_index, - ) - if attention_shards is not None: - for rank, shard in enumerate(attention_shards): - position_start = rank_loads[rank] - gdn_ranges_by_rank[rank].append((shard.start, shard.stop, position_start)) - rank_loads[rank] += len(shard) - return 0 - cross_rank_tokens = 0 - shard_lengths = _fla_aligned_chain_shard_lengths(segment.length, cp_size=cp_size) - start = 0 - for rank, shard_length in enumerate(shard_lengths): - end = start + shard_length - if start >= end: - raise ValueError( - "CP chain planning requires non-empty shards; " - f"segment={segment.kind}:{segment.family_index} " - f"length={segment.length} cp_size={cp_size}" - ) - shard_start = token_start + start - position_start = rank_loads[rank] - gdn_ranges_by_rank[rank].append( - (shard_start, shard_start + shard_length, position_start) - ) - rank_loads[rank] += shard_length - if attention_layout_index is not None: - cross_rank_tokens += shard_length - _attention_overlap_count( - attention_layout_index, - rank, - shard_start, - shard_start + shard_length, - ) - start = end - return cross_rank_tokens - - -def _chain_rank_token_indices( - segment: GdnSegmentSpec, - spec: GdnPackedExecutionSpec, - *, - cp_rank: int, - cp_size: int, -) -> range: - token_start = _segment_token_start(segment, spec.sequence_length) - lengths = _fla_aligned_chain_shard_lengths(segment.length, cp_size=cp_size) - start = sum(lengths[:cp_rank]) - end = start + lengths[cp_rank] - if start >= end: - raise ValueError( - "CP chain planning requires non-empty shards; " - f"segment={segment.kind}:{segment.family_index} " - f"length={segment.length} cp_size={cp_size}" - ) - return range(token_start + start, token_start + end) - - -def _fla_aligned_chain_shard_lengths(length: int, *, cp_size: int) -> tuple[int, ...]: - full_chunks = int(length) // FLA_CHUNK_SIZE - if full_chunks < int(cp_size): - raise ValueError( - "CP chain planning requires at least one full FLA chunk per rank; " - f"length={length} cp_size={cp_size}" - ) - base_chunks = full_chunks // int(cp_size) - extra_chunks = full_chunks % int(cp_size) - chunk_counts = tuple( - base_chunks + (1 if rank < extra_chunks else 0) for rank in range(int(cp_size)) - ) - lengths = [count * FLA_CHUNK_SIZE for count in chunk_counts] - lengths[-1] += int(length) - full_chunks * FLA_CHUNK_SIZE - return tuple(lengths) - - -def _attention_contiguous_chain_shards( - token_start: int, - token_count: int, - *, - cp_size: int, - attention_layout_index: _AttentionLayoutIndex | None, -) -> tuple[range, ...] | None: - if attention_layout_index is None: - return None - segment_end = token_start + token_count - shards: list[range] = [] - cursor = token_start - for rank in range(cp_size): - overlap = _attention_single_contiguous_overlap( - attention_layout_index, - rank, - token_start, - segment_end, - ) - if overlap is None: - return None - start, end = overlap - if start != cursor or end <= start: - return None - shards.append(range(start, end)) - cursor = end - if cursor != segment_end: - return None - if any(len(shard) % FLA_CHUNK_SIZE != 0 for shard in shards[:-1]): - return None - return tuple(shards) - - -def _attention_single_contiguous_overlap( - index: _AttentionLayoutIndex, - rank: int, - start: int, - end: int, -) -> tuple[int, int] | None: - overlaps = _range_overlaps(start, end, index.token_ranges_by_rank[rank]) - if len(overlaps) != 1: - return None - return overlaps[0] - - -def _append_local_segment( - gdn_ranges_by_rank: list[list[tuple[int, int, int]]], - rank_loads: list[int], - rank: int, - segment: GdnSegmentSpec, - spec: GdnPackedExecutionSpec, - *, - segment_attention_counts: dict[tuple[int, int, int], tuple[int, ...]], -) -> int: - token_start = _segment_token_start(segment, spec.sequence_length) - position_start = rank_loads[rank] - gdn_ranges_by_rank[rank].append( - (token_start, token_start + segment.length, position_start) - ) - rank_loads[rank] += segment.length - return segment.length - segment_attention_counts[_segment_key(segment)][rank] - - -def _least_loaded_rank(rank_loads: list[int]) -> int: - return min(range(len(rank_loads)), key=lambda rank: (rank_loads[rank], rank)) - - -def _owner_rank( - local_prefix_segments_by_rank: list[list[GdnSegmentSpec]], - prefix: GdnSegmentSpec, -) -> int: - for rank, segments in enumerate(local_prefix_segments_by_rank): - if prefix in segments: - return rank - raise RuntimeError("local prefix owner was not recorded") - - -def _build_position_bucket_plans( - segment_buckets: tuple[tuple[GdnSegmentSpec, ...], ...], - local_token_ranges: tuple[tuple[int, int, int], ...], - *, - sequence_length: int, - device: torch.device | str, - token_ranges_by_rank: tuple[tuple[tuple[int, int, int], ...], ...] | None = None, -) -> tuple[GdnSegmentBucketPlan, ...]: - return tuple( - _build_position_bucket_plan( - bucket, - local_token_ranges, - sequence_length=sequence_length, - device=device, - token_ranges_by_rank=token_ranges_by_rank, - ) - for bucket in segment_buckets - ) - - -def _build_position_bucket_plan( - segments: tuple[GdnSegmentSpec, ...], - local_token_ranges: tuple[tuple[int, int, int], ...], - *, - sequence_length: int, - device: torch.device | str, - token_ranges_by_rank: tuple[tuple[tuple[int, int, int], ...], ...] | None = None, -) -> GdnSegmentBucketPlan: - exact_plan = _build_exact_range_position_bucket_plan( - segments, - local_token_ranges, - sequence_length=sequence_length, - device=device, - token_ranges_by_rank=token_ranges_by_rank, - ) - if exact_plan is not None: - return exact_plan - local_positions_by_segment = [] - lengths = [] - local_range_ends = tuple(token_end for _, token_end, _ in local_token_ranges) - for segment in segments: - positions = _local_positions_for_segment( - segment, - sequence_length=sequence_length, - local_token_ranges=local_token_ranges, - local_range_ends=local_range_ends, - ) - length = int(positions.numel()) - if not length: - raise ValueError( - "planned GDN bucket contains a segment with no local tokens; " - f"family={segment.family_index} kind={segment.kind}" - ) - local_positions_by_segment.append(positions) - lengths.append(length) - max_length = max(lengths) - lengths_cpu = torch.tensor(lengths, dtype=torch.long) - offsets_cpu = torch.arange(max_length, dtype=torch.long).unsqueeze(1) - real_mask_cpu = offsets_cpu < lengths_cpu.unsqueeze(0) - position_indices_cpu = torch.zeros(max_length, len(segments), dtype=torch.long) - for column, positions in enumerate(local_positions_by_segment): - position_indices_cpu[: int(positions.numel()), column] = positions - cu_seqlens_cpu = torch.cat( - [lengths_cpu.new_zeros(1), torch.cumsum(lengths_cpu, dim=0)] - ) - lengths_by_rank_cpu = _bucket_lengths_by_rank_cpu( - segments, - token_ranges_by_rank, - sequence_length=sequence_length, - ) - row_indices_cpu = torch.zeros(max_length, len(segments), dtype=torch.long) - family_indices_cpu = torch.tensor( - [segment.family_index for segment in segments], - dtype=torch.long, - ) - return GdnSegmentBucketPlan.model_construct( - length=max_length, - lengths=_move_planner_tensor(lengths_cpu, device), - lengths_cpu=lengths_cpu, - lengths_by_rank_cpu=lengths_by_rank_cpu, - real_mask=_move_planner_tensor(real_mask_cpu, device), - cu_seqlens=_move_planner_tensor(cu_seqlens_cpu, device), - cu_seqlens_cpu=cu_seqlens_cpu, - row_indices=_move_planner_tensor(row_indices_cpu, device), - position_indices=_move_planner_tensor(position_indices_cpu, device), - family_indices=_move_planner_tensor(family_indices_cpu, device), - real_token_count_static=sum(lengths), - ) - - -def _build_exact_range_position_bucket_plan( - segments: tuple[GdnSegmentSpec, ...], - local_token_ranges: tuple[tuple[int, int, int], ...], - *, - sequence_length: int, - device: torch.device | str, - token_ranges_by_rank: tuple[tuple[tuple[int, int, int], ...], ...] | None = None, -) -> GdnSegmentBucketPlan | None: - range_positions = { - (start, end): position for start, end, position in local_token_ranges - } - starts = [] - lengths = [] - for segment in segments: - token_start = _segment_token_start(segment, sequence_length) - token_end = token_start + segment.length - position_start = range_positions.get((token_start, token_end)) - if position_start is None: - return None - starts.append(position_start) - lengths.append(segment.length) - max_length = max(lengths) - starts_cpu = torch.tensor(starts, dtype=torch.long) - lengths_cpu = torch.tensor(lengths, dtype=torch.long) - offsets_cpu = torch.arange(max_length, dtype=torch.long).unsqueeze(1) - real_mask_cpu = offsets_cpu < lengths_cpu.unsqueeze(0) - position_indices_cpu = torch.where( - real_mask_cpu, - starts_cpu.unsqueeze(0) + offsets_cpu, - torch.zeros_like(offsets_cpu), - ) - cu_seqlens_cpu = torch.cat( - [lengths_cpu.new_zeros(1), torch.cumsum(lengths_cpu, dim=0)] - ) - lengths_by_rank_cpu = _bucket_lengths_by_rank_cpu( - segments, - token_ranges_by_rank, - sequence_length=sequence_length, - ) - row_indices_cpu = torch.zeros(max_length, len(segments), dtype=torch.long) - family_indices_cpu = torch.tensor( - [segment.family_index for segment in segments], - dtype=torch.long, - ) - return GdnSegmentBucketPlan.model_construct( - length=max_length, - lengths=_move_planner_tensor(lengths_cpu, device), - lengths_cpu=lengths_cpu, - lengths_by_rank_cpu=lengths_by_rank_cpu, - real_mask=_move_planner_tensor(real_mask_cpu, device), - cu_seqlens=_move_planner_tensor(cu_seqlens_cpu, device), - cu_seqlens_cpu=cu_seqlens_cpu, - row_indices=_move_planner_tensor(row_indices_cpu, device), - position_indices=_move_planner_tensor(position_indices_cpu, device), - family_indices=_move_planner_tensor(family_indices_cpu, device), - real_token_count_static=sum(lengths), - ) - - -def _bucket_lengths_by_rank_cpu( - segments: tuple[GdnSegmentSpec, ...], - token_ranges_by_rank: tuple[tuple[tuple[int, int, int], ...], ...] | None, - *, - sequence_length: int, -) -> torch.Tensor | None: - if token_ranges_by_rank is None: - return None - lengths_by_rank = [] - for rank_ranges_with_positions in token_ranges_by_rank: - rank_ranges = tuple( - (start, end) for start, end, _position in rank_ranges_with_positions - ) - rank_lengths = [] - for segment in segments: - start = _segment_token_start(segment, sequence_length) - end = start + segment.length - rank_lengths.append( - sum( - max(0, min(end, range_end) - max(start, range_start)) - for range_start, range_end in rank_ranges - ) - ) - lengths_by_rank.append(rank_lengths) - return torch.tensor(lengths_by_rank, dtype=torch.long) - - -def _move_planner_tensor( - tensor: torch.Tensor, device: torch.device | str -) -> torch.Tensor: - target = torch.device(device) - if target.type == "cpu": - return tensor - return tensor.to(device=target) - - -def _batch_segments_by_padded_work( - segments: tuple[GdnSegmentSpec, ...], - *, - max_padding_ratio: float = 1.25, - max_segments_per_batch: int = 128, -) -> tuple[tuple[GdnSegmentSpec, ...], ...]: - if not segments: - return () - ordered = sorted( - segments, key=lambda segment: (segment.length, segment.family_index) - ) - batches: list[list[GdnSegmentSpec]] = [] - current: list[GdnSegmentSpec] = [] - current_tokens = 0 - current_max = 0 - for segment in ordered: - next_count = len(current) + 1 - next_tokens = current_tokens + segment.length - next_max = max(current_max, segment.length) - padded = next_max * next_count - can_extend = not current or ( - next_count <= max_segments_per_batch - and padded <= max_padding_ratio * next_tokens - ) - if not can_extend: - batches.append(current) - current = [] - current_tokens = 0 - current_max = 0 - current.append(segment) - current_tokens += segment.length - current_max = max(current_max, segment.length) - if current: - batches.append(current) - return tuple(tuple(batch) for batch in batches) - - -def _build_segment_bucket_plan( - length: int, segments: tuple[GdnSegmentSpec, ...], *, device: torch.device | str -) -> GdnSegmentBucketPlan: - max_length = max(segment.length for segment in segments) - lengths_cpu = torch.tensor( - [segment.length for segment in segments], dtype=torch.long - ) - starts_cpu = torch.tensor([segment.start for segment in segments], dtype=torch.long) - rows_cpu = torch.tensor( - [segment.row_index for segment in segments], dtype=torch.long - ) - offsets_cpu = torch.arange(max_length, dtype=torch.long).unsqueeze(1) - real_mask_cpu = offsets_cpu < lengths_cpu.unsqueeze(0) - positions_cpu = starts_cpu.unsqueeze(0) + offsets_cpu - family_indices_cpu = torch.tensor( - [segment.family_index for segment in segments], - dtype=torch.long, - ) - cu_seqlens_cpu = torch.cat( - [lengths_cpu.new_zeros(1), torch.cumsum(lengths_cpu, dim=0)] - ) - return GdnSegmentBucketPlan.model_construct( - length=max_length, - lengths=_move_planner_tensor(lengths_cpu, device), - lengths_cpu=lengths_cpu, - lengths_by_rank_cpu=None, - real_mask=_move_planner_tensor(real_mask_cpu, device), - cu_seqlens=_move_planner_tensor(cu_seqlens_cpu, device), - cu_seqlens_cpu=cu_seqlens_cpu, - row_indices=_move_planner_tensor( - rows_cpu.unsqueeze(0).expand(max_length, -1).contiguous(), device - ), - position_indices=_move_planner_tensor(positions_cpu, device), - family_indices=_move_planner_tensor(family_indices_cpu, device), - real_token_count_static=sum(segment.length for segment in segments), - ) - - -def _segment_token_start(segment: GdnSegmentSpec, sequence_length: int) -> int: - return segment.row_index * sequence_length + segment.start - - -def _attention_overlap_count( - index: _AttentionLayoutIndex, - rank: int, - start: int, - end: int, -) -> int: - return _range_overlap_count( - start, - end, - index.token_ranges_by_rank[rank], - index.token_range_ends_by_rank[rank], - ) - - -def _range_overlap_count( - start: int, - end: int, - ranges: tuple[tuple[int, int], ...], - range_ends: tuple[int, ...], -) -> int: - count = 0 - range_index = bisect_left(range_ends, start + 1) - for range_start, range_end in ranges[range_index:]: - if range_start >= end: - break - count += min(end, range_end) - max(start, range_start) - return count - - -def _range_overlaps( - start: int, - end: int, - ranges: tuple[tuple[int, int], ...], -) -> list[tuple[int, int]]: - overlaps = [ - (max(start, range_start), min(end, range_end)) - for range_start, range_end in ranges - if max(start, range_start) < min(end, range_end) - ] - overlaps.sort() - return overlaps - - -def _local_token_ranges( - local_gdn_tokens: tuple[int, ...], -) -> tuple[tuple[int, int, int], ...]: - if not local_gdn_tokens: - return () - ranges = [] - token_start = local_gdn_tokens[0] - token_end = token_start + 1 - position_start = 0 - for position, token in enumerate(local_gdn_tokens[1:], start=1): - if token == token_end: - token_end += 1 - continue - ranges.append((token_start, token_end, position_start)) - token_start = token - token_end = token + 1 - position_start = position - ranges.append((token_start, token_end, position_start)) - return tuple(ranges) - - -def _local_positions_for_segment( - segment: GdnSegmentSpec, - *, - sequence_length: int, - local_token_ranges: tuple[tuple[int, int, int], ...], - local_range_ends: tuple[int, ...], -) -> torch.Tensor: - segment_start = _segment_token_start(segment, sequence_length) - segment_end = segment_start + segment.length - pieces = [] - range_index = bisect_left(local_range_ends, segment_start + 1) - for token_start, token_end, position_start in local_token_ranges[range_index:]: - if token_start >= segment_end: - break - overlap_start = max(segment_start, token_start) - overlap_end = min(segment_end, token_end) - if overlap_start >= overlap_end: - continue - pieces.append( - torch.arange( - position_start + overlap_start - token_start, - position_start + overlap_end - token_start, - dtype=torch.long, - ) - ) - if not pieces: - return torch.empty((0,), dtype=torch.long) - if len(pieces) == 1: - return pieces[0] - return torch.cat(pieces) - - -def _rank2_long_cpu(name: str, tensor: torch.Tensor) -> torch.Tensor: - if not torch.is_tensor(tensor): - raise TypeError(f"{name} must be a torch.Tensor") - if tensor.ndim != 2: - raise ValueError(f"{name} must be rank 2 [batch, sequence], got {tensor.ndim}") - if tensor.dtype not in ( - torch.int8, - torch.int16, - torch.int32, - torch.int64, - torch.long, - ): - raise TypeError(f"{name} must contain integer ids, got dtype={tensor.dtype}") - return tensor.detach().to(device="cpu", dtype=torch.long) - - -def _validate_padding_tensor( - row_index: int, - group_ids: torch.Tensor, - parent_ids: torch.Tensor, -) -> int: - padding_positions = torch.nonzero(group_ids == -1, as_tuple=False) - valid_length = ( - int(padding_positions[0].item()) - if int(padding_positions.numel()) > 0 - else int(group_ids.numel()) - ) - if valid_length == 0: - if bool(torch.any(parent_ids != -1).item()): - raise ValueError(f"row {row_index}: padding parent_ids must be -1") - return 0 - if bool(torch.any(group_ids[valid_length:] != -1).item()): - raise ValueError( - f"row {row_index}: valid tokens must be contiguous before padding" - ) - if bool(torch.any(parent_ids[:valid_length] == -1).item()): - raise ValueError( - f"row {row_index}: valid tokens must have non-padding parent_ids" - ) - if bool(torch.any(parent_ids[valid_length:] != -1).item()): - raise ValueError(f"row {row_index}: padding parent_ids must be -1") - return valid_length - - -def _validate_padding( - row_index: int, - group_ids: list[int], - parent_ids: list[int], -) -> int: - valid_length = 0 - for group_id in group_ids: - if group_id == -1: - break - valid_length += 1 - if valid_length == 0: - if any(parent_id != -1 for parent_id in parent_ids): - raise ValueError(f"row {row_index}: padding parent_ids must be -1") - return 0 - if any(group_id != -1 for group_id in group_ids[valid_length:]): - raise ValueError( - f"row {row_index}: valid tokens must be contiguous before padding" - ) - if any(parent_id == -1 for parent_id in parent_ids[:valid_length]): - raise ValueError( - f"row {row_index}: valid tokens must have non-padding parent_ids" - ) - if any(parent_id != -1 for parent_id in parent_ids[valid_length:]): - raise ValueError(f"row {row_index}: padding parent_ids must be -1") - return valid_length - - -def _parse_row_tensor( - *, - row_index: int, - group_ids: torch.Tensor, - parent_ids: torch.Tensor, - valid_length: int, - first_family_index: int, - min_completions_per_family: int, -) -> list[GdnPackedFamilySpec]: - valid_groups = group_ids[:valid_length] - valid_parents = parent_ids[:valid_length] - if valid_length > 1: - same_group = valid_groups[1:] == valid_groups[:-1] - parent_changed = same_group & (valid_parents[1:] != valid_parents[:-1]) - if bool(torch.any(parent_changed).item()): - position = int(torch.nonzero(parent_changed, as_tuple=False)[0].item()) + 1 - group_id = int(valid_groups[position].item()) - previous_parent = int(valid_parents[position - 1].item()) - current_parent = int(valid_parents[position].item()) - raise ValueError( - f"row {row_index}: group {group_id} changes parent from " - f"{previous_parent} to {current_parent}" - ) - boundaries = torch.nonzero(~same_group, as_tuple=False).flatten() + 1 - starts_tensor = torch.cat( - (valid_groups.new_zeros(1), boundaries.to(valid_groups.dtype)) - ) - ends_tensor = torch.cat( - ( - boundaries.to(valid_groups.dtype), - valid_groups.new_tensor([valid_length]), - ) - ) - else: - starts_tensor = valid_groups.new_zeros(1) - ends_tensor = valid_groups.new_tensor([valid_length]) - - starts = tuple(int(value) for value in starts_tensor.tolist()) - ends = tuple(int(value) for value in ends_tensor.tolist()) - segment_group_ids = tuple(int(valid_groups[start].item()) for start in starts) - segment_parent_ids = tuple(int(valid_parents[start].item()) for start in starts) - families: list[GdnPackedFamilySpec] = [] - seen_groups: set[int] = set() - segment_cursor = 0 - while segment_cursor < len(starts): - group_id = segment_group_ids[segment_cursor] - parent_id = segment_parent_ids[segment_cursor] - start = starts[segment_cursor] - end = ends[segment_cursor] - if group_id in seen_groups: - raise ValueError(f"row {row_index}: group_id {group_id} is non-contiguous") - if group_id != parent_id: - raise ValueError( - f"row {row_index}: completion group {group_id} appears before " - f"its prefix parent {parent_id}" - ) - seen_groups.add(group_id) - family_index = first_family_index + len(families) - prefix = _trusted_pydantic_construct( - GdnSegmentSpec, - _GDN_SEGMENT_SPEC_FIELDS, - row_index=row_index, - family_index=family_index, - group_id=group_id, - parent_id=parent_id, - start=start, - end=end, - kind="prefix", - child_index=None, - ) - segment_cursor += 1 - completions: list[GdnSegmentSpec] = [] - while segment_cursor < len(starts): - child_group_id = segment_group_ids[segment_cursor] - child_parent_id = segment_parent_ids[segment_cursor] - child_start = starts[segment_cursor] - child_end = ends[segment_cursor] - if child_group_id == child_parent_id: - break - if child_parent_id != group_id: - raise ValueError( - f"row {row_index}: completion group {child_group_id} has " - f"parent {child_parent_id}, expected active prefix {group_id}" - ) - if child_group_id in seen_groups: - raise ValueError( - f"row {row_index}: group_id {child_group_id} is non-contiguous" - ) - seen_groups.add(child_group_id) - completions.append( - _trusted_pydantic_construct( - GdnSegmentSpec, - _GDN_SEGMENT_SPEC_FIELDS, - row_index=row_index, - family_index=family_index, - group_id=child_group_id, - parent_id=child_parent_id, - start=child_start, - end=child_end, - kind="completion", - child_index=len(completions), - ) - ) - segment_cursor += 1 - if len(completions) < min_completions_per_family: - raise ValueError( - f"row {row_index}: prefix group {group_id} has {len(completions)} " - f"completion(s), expected at least {min_completions_per_family}" - ) - families.append( - _trusted_pydantic_construct( - GdnPackedFamilySpec, - _GDN_PACKED_FAMILY_SPEC_FIELDS, - row_index=row_index, - family_index=family_index, - prefix=prefix, - completions=tuple(completions), - ) - ) - return families - - -def _parse_row( - *, - row_index: int, - group_ids: list[int], - parent_ids: list[int], - valid_length: int, - first_family_index: int, - min_completions_per_family: int, -) -> list[GdnPackedFamilySpec]: - families: list[GdnPackedFamilySpec] = [] - seen_groups: set[int] = set() - cursor = 0 - while cursor < valid_length: - group_id, parent_id, start, end = _read_segment( - row_index, group_ids, parent_ids, valid_length, cursor - ) - if group_id in seen_groups: - raise ValueError(f"row {row_index}: group_id {group_id} is non-contiguous") - if group_id != parent_id: - raise ValueError( - f"row {row_index}: completion group {group_id} appears before " - f"its prefix parent {parent_id}" - ) - seen_groups.add(group_id) - family_index = first_family_index + len(families) - prefix = GdnSegmentSpec( - row_index=row_index, - family_index=family_index, - group_id=group_id, - parent_id=parent_id, - start=start, - end=end, - kind="prefix", - ) - cursor = end - completions: list[GdnSegmentSpec] = [] - while cursor < valid_length: - child_group_id, child_parent_id, child_start, child_end = _read_segment( - row_index, group_ids, parent_ids, valid_length, cursor - ) - if child_group_id == child_parent_id: - break - if child_parent_id != group_id: - raise ValueError( - f"row {row_index}: completion group {child_group_id} has " - f"parent {child_parent_id}, expected active prefix {group_id}" - ) - if child_group_id in seen_groups: - raise ValueError( - f"row {row_index}: group_id {child_group_id} is non-contiguous" - ) - seen_groups.add(child_group_id) - completions.append( - GdnSegmentSpec( - row_index=row_index, - family_index=family_index, - group_id=child_group_id, - parent_id=child_parent_id, - start=child_start, - end=child_end, - kind="completion", - child_index=len(completions), - ) - ) - cursor = child_end - if len(completions) < min_completions_per_family: - raise ValueError( - f"row {row_index}: prefix group {group_id} has {len(completions)} " - f"completion(s), expected at least {min_completions_per_family}" - ) - families.append( - GdnPackedFamilySpec( - row_index=row_index, - family_index=family_index, - prefix=prefix, - completions=tuple(completions), - ) - ) - return families - - -def _read_segment( - row_index: int, - group_ids: list[int], - parent_ids: list[int], - valid_length: int, - cursor: int, -) -> tuple[int, int, int, int]: - group_id = int(group_ids[cursor]) - parent_id = int(parent_ids[cursor]) - if group_id < 0 or parent_id < 0: - raise ValueError(f"row {row_index}: segment ids must be non-negative") - start = cursor - cursor += 1 - while cursor < valid_length and int(group_ids[cursor]) == group_id: - current_parent = int(parent_ids[cursor]) - if current_parent != parent_id: - raise ValueError( - f"row {row_index}: group {group_id} changes parent from " - f"{parent_id} to {current_parent}" - ) - cursor += 1 - return group_id, parent_id, start, cursor diff --git a/src/art/megatron/gdn/layout.py b/src/art/megatron/gdn/layout.py index c3469a451..7119218f6 100644 --- a/src/art/megatron/gdn/layout.py +++ b/src/art/megatron/gdn/layout.py @@ -1,9 +1,9 @@ from __future__ import annotations from collections.abc import Sequence +from dataclasses import dataclass, replace from typing import Any -from pydantic import BaseModel, ConfigDict, Field, model_validator import torch from torch import Tensor from torch.distributed import ( @@ -20,47 +20,47 @@ from art.megatron.context_parallel.layout_index import TokenLayoutIndex -class GdnCpPeerTransfer(BaseModel): +@dataclass(frozen=True) +class GdnCpPeerTransfer: """Token rows sent from one source rank to one destination rank.""" - model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True) - - source_rank: int = Field(ge=0) - dest_rank: int = Field(ge=0) - token_count: int = Field(ge=0) + source_rank: int + dest_rank: int + token_count: int + source_positions_cpu: tuple[int, ...] | None = None + dest_positions_cpu: tuple[int, ...] | None = None source_positions_tensor: Tensor | None = None dest_positions_tensor: Tensor | None = None - @model_validator(mode="after") - def _same_lengths(self) -> "GdnCpPeerTransfer": + def __post_init__(self) -> None: lengths = {int(self.token_count)} + if self.source_positions_cpu is not None: + lengths.add(len(self.source_positions_cpu)) + if self.dest_positions_cpu is not None: + lengths.add(len(self.dest_positions_cpu)) if self.source_positions_tensor is not None: lengths.add(int(self.source_positions_tensor.numel())) if self.dest_positions_tensor is not None: lengths.add(int(self.dest_positions_tensor.numel())) if len(lengths) != 1: raise ValueError("token, source, and destination position counts differ") - return self -class GdnCpExchangePlan(BaseModel): +@dataclass(frozen=True) +class GdnCpExchangePlan: """Permutation/all-to-all metadata between two distributed token layouts.""" - model_config = ConfigDict(frozen=True) - - cp_size: int = Field(ge=1) + cp_size: int source_token_counts_by_rank: tuple[int, ...] dest_token_counts_by_rank: tuple[int, ...] transfers: tuple[GdnCpPeerTransfer, ...] - cross_rank_token_count_override: int | None = Field(default=None, ge=0) + cross_rank_token_count_override: int | None = None - @model_validator(mode="after") - def _rank_counts(self) -> "GdnCpExchangePlan": + def __post_init__(self) -> None: if len(self.source_token_counts_by_rank) != self.cp_size: raise ValueError("source token count length must equal cp_size") if len(self.dest_token_counts_by_rank) != self.cp_size: raise ValueError("destination token count length must equal cp_size") - return self @property def cross_rank_token_count(self) -> int: @@ -73,11 +73,10 @@ def cross_rank_token_count(self) -> int: ) -class GdnSpExchangePlan(BaseModel): +@dataclass(frozen=True) +class GdnSpExchangePlan: """Sequence-parallel view of an existing CP exchange plan.""" - model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True) - plan: GdnCpExchangePlan rank: int @@ -206,7 +205,7 @@ def build_local_rank_cp_exchange_plan_from_dest_ranges( device=device, ) ) - return GdnCpExchangePlan.model_construct( + return GdnCpExchangePlan( cp_size=cp_size, source_token_counts_by_rank=source_layout.token_counts_by_rank, dest_token_counts_by_rank=dest_counts, @@ -238,23 +237,33 @@ def _make_peer_transfer( source_count=source_count, dest_count=dest_count, ): + source_cpu = None + dest_cpu = None source_tensor = None dest_tensor = None else: + source_cpu = _tensor_positions_tuple(source_positions) + dest_cpu = _tensor_positions_tuple(dest_positions) target = torch.device(device) if device is not None else torch.device("cpu") source_tensor = source_positions.to( device=target, dtype=torch.long ).contiguous() dest_tensor = dest_positions.to(device=target, dtype=torch.long).contiguous() - return GdnCpPeerTransfer.model_construct( + return GdnCpPeerTransfer( source_rank=source_rank, dest_rank=dest_rank, token_count=token_count, + source_positions_cpu=source_cpu, + dest_positions_cpu=dest_cpu, source_positions_tensor=source_tensor, dest_positions_tensor=dest_tensor, ) +def _tensor_positions_tuple(tensor: Tensor) -> tuple[int, ...]: + return tuple(int(value) for value in tensor.detach().cpu().tolist()) + + def _is_full_identity_transfer( *, source_rank: int, @@ -277,16 +286,18 @@ def _is_full_identity_transfer( def _reverse_exchange_plan(plan: GdnCpExchangePlan) -> GdnCpExchangePlan: - return GdnCpExchangePlan.model_construct( + return GdnCpExchangePlan( cp_size=plan.cp_size, source_token_counts_by_rank=_dest_counts_by_rank(plan), dest_token_counts_by_rank=_source_counts_by_rank(plan), cross_rank_token_count_override=plan.cross_rank_token_count_override, transfers=tuple( - GdnCpPeerTransfer.model_construct( + GdnCpPeerTransfer( source_rank=transfer.dest_rank, dest_rank=transfer.source_rank, token_count=_transfer_token_count(transfer), + source_positions_cpu=transfer.dest_positions_cpu, + dest_positions_cpu=transfer.source_positions_cpu, source_positions_tensor=transfer.dest_positions_tensor, dest_positions_tensor=transfer.source_positions_tensor, ) @@ -485,15 +496,11 @@ def move_cp_exchange_plan_to_device( if plan is None: return None target = torch.device(device) - return GdnCpExchangePlan.model_construct( - cp_size=plan.cp_size, - source_token_counts_by_rank=_source_counts_by_rank(plan), - dest_token_counts_by_rank=_dest_counts_by_rank(plan), + return replace( + plan, transfers=tuple( - GdnCpPeerTransfer.model_construct( - source_rank=transfer.source_rank, - dest_rank=transfer.dest_rank, - token_count=transfer.token_count, + replace( + transfer, source_positions_tensor=_move_optional_index_tensor( transfer.source_positions_tensor, target ), @@ -503,7 +510,6 @@ def move_cp_exchange_plan_to_device( ) for transfer in plan.transfers ), - cross_rank_token_count_override=plan.cross_rank_token_count_override, ) @@ -532,7 +538,7 @@ def shard_cp_exchange_plan_for_sequence_parallel( """ if tp_size <= 1: - return GdnSpExchangePlan.model_construct(plan=plan, rank=cp_rank) + return GdnSpExchangePlan(plan=plan, rank=cp_rank) _check_rank(plan, cp_rank) if tp_rank < 0 or tp_rank >= tp_size: raise ValueError(f"tp_rank must be in [0, {tp_size}), got {tp_rank}") @@ -603,7 +609,7 @@ def shard_cp_exchange_plan_for_sequence_parallel( # A CP-local reorder can still move rows between TP ranks, and local CP plans do # not contain enough global TP information for every rank to independently # prove that no peer exchange is needed. - sp_plan = GdnCpExchangePlan.model_construct( + sp_plan = GdnCpExchangePlan( cp_size=world_size, source_token_counts_by_rank=source_counts, dest_token_counts_by_rank=dest_counts, @@ -612,7 +618,7 @@ def shard_cp_exchange_plan_for_sequence_parallel( ), cross_rank_token_count_override=1, ) - return GdnSpExchangePlan.model_construct(plan=sp_plan, rank=composite_rank) + return GdnSpExchangePlan(plan=sp_plan, rank=composite_rank) def recv_split_sizes_for_rank(plan: GdnCpExchangePlan, rank: int) -> tuple[int, ...]: @@ -750,10 +756,15 @@ def _is_implicit_full_identity_transfer( ) -def _transfer_positions_tuple(tensor: Tensor | None) -> tuple[int, ...]: +def _transfer_positions_tuple( + positions: tuple[int, ...] | None, + tensor: Tensor | None, +) -> tuple[int, ...]: + if positions is not None: + return positions if tensor is None: return () - return tuple(int(value) for value in tensor.detach().cpu().tolist()) + return _tensor_positions_tuple(tensor) def _transfer_index_tensor( @@ -895,17 +906,6 @@ def _exchange_rank_tensor_local( ) -def _copy_rank_self_transfers( - local_tensor: Tensor, - plan: GdnCpExchangePlan, - *, - rank: int, -) -> Tensor: - return _init_rank_exchange_output( - local_tensor, plan, rank=rank, accumulate=False, zero_init=False - ) - - def _init_rank_exchange_output( local_tensor: Tensor, plan: GdnCpExchangePlan, @@ -1028,7 +1028,10 @@ def _transfer_dest_positions_for_duplicate_check( dest_count=_dest_count_for_rank(plan, transfer.dest_rank), ): return tuple(range(token_count)) - positions = _transfer_positions_tuple(transfer.dest_positions_tensor) + positions = _transfer_positions_tuple( + transfer.dest_positions_cpu, + transfer.dest_positions_tensor, + ) if len(positions) != token_count: raise ValueError("GDN CP transfer destination positions must match token_count") return positions diff --git a/src/art/megatron/gdn/operator.py b/src/art/megatron/gdn/operator.py index e8a122f5c..302f43380 100644 --- a/src/art/megatron/gdn/operator.py +++ b/src/art/megatron/gdn/operator.py @@ -1,22 +1,24 @@ from __future__ import annotations from types import MethodType -from typing import Any, Callable, Literal, NamedTuple, Sequence, cast +from typing import Any, Callable, Iterable, Literal, NamedTuple, Sequence, cast import torch from torch import Tensor import torch.distributed as dist import torch.nn.functional as F +import triton +import triton.language as tl from .conv_gelu import packed_varlen_causal_conv from .fla_cp import chunk_gated_delta_rule_native_cp -from .gdn_shared_prefix import ( +from .gdn_prefix_tree import ( GdnPackedExecutionSpec, - GdnParentStateTransferPlan, GdnRankExecutionPlan, GdnSegmentBucketPlan, + GdnStateExchangePlan, build_gdn_rank_execution_plan, - parse_gdn_shared_prefix_segments, + parse_gdn_prefix_tree_segments, ) from .segment_layout import ( gather_bucket_streams_compact as _gather_bucket_streams_compact_fused, @@ -46,8 +48,8 @@ def set_gdn_trace_token_uid_hooks(hooks: Any | None) -> Any | None: return previous -def install_shared_prefix_gdn_hooks(model_chunks: Sequence[Any]) -> None: - """Patch Megatron GatedDeltaNet modules to honor ART shared-prefix packing.""" +def install_prefix_tree_gdn_hooks(model_chunks: Sequence[Any]) -> None: + """Patch Megatron GatedDeltaNet modules to honor ART prefix-tree packing.""" gated_delta_net_type = _optional_gated_delta_net_type() if gated_delta_net_type is None: @@ -56,12 +58,12 @@ def install_shared_prefix_gdn_hooks(model_chunks: Sequence[Any]) -> None: for module in chunk.modules(): if not isinstance(module, gated_delta_net_type): continue - if getattr(module, "_art_shared_prefix_gdn_hooked", False): + if getattr(module, "_art_prefix_tree_gdn_hooked", False): continue original_forward = module.forward module._art_physical_forward = original_forward - module.forward = MethodType(_shared_prefix_forward, module) - module._art_shared_prefix_gdn_hooked = True + module.forward = MethodType(_prefix_tree_forward, module) + module._art_prefix_tree_gdn_hooked = True def install_gdn_island_hooks(model_chunks: Sequence[Any]) -> None: @@ -301,7 +303,8 @@ def _empty_safe_norm_forward( return original_forward(input_, *args, **kwargs) -def _shared_prefix_forward( +@torch.compiler.disable +def _prefix_tree_forward( self: Any, hidden_states: Tensor, attention_mask: Tensor, @@ -336,11 +339,9 @@ def _shared_prefix_forward( del attention_mask, key_value_states, sequence_len_offset, kwargs if inference_context is not None or inference_params is not None: - raise NotImplementedError("ART shared-prefix GDN does not support inference.") + raise NotImplementedError("ART prefix-tree GDN does not support inference.") if packed_seq_params is not None: - raise NotImplementedError( - "PackedSeqParams is not used in ART shared-prefix GDN." - ) + raise NotImplementedError("PackedSeqParams is not used in ART prefix-tree GDN.") current_layout = _normalize_cp_layout( getattr(attention_bias, "gdn_hidden_layout", "attention") ) @@ -355,7 +356,7 @@ def _shared_prefix_forward( _mark_cp_layout_active( attention_bias, hidden_states, gdn=self, layout=input_layout ) - output = gdn_shared_prefix_forward( + output = gdn_prefix_tree_forward( self, hidden_states, group_ids=cast(Tensor, group_ids), @@ -377,7 +378,7 @@ def _shared_prefix_forward( @torch.compiler.disable -def gdn_shared_prefix_forward( +def gdn_prefix_tree_forward( gdn: Any, hidden_states: Tensor, *, @@ -390,7 +391,7 @@ def gdn_shared_prefix_forward( input_layout: Literal["attention", "gdn"] = "attention", output_layout: Literal["attention", "gdn"] = "attention", ) -> tuple[Tensor, Tensor | None]: - """Run one GDN layer over ART shared-prefix packed rows.""" + """Run one GDN layer over ART prefix-tree packed rows.""" return run_gdn_layer( gdn, @@ -420,7 +421,7 @@ def run_gdn_layer( input_layout: Literal["attention", "gdn"] = "attention", output_layout: Literal["attention", "gdn"] = "attention", ) -> tuple[Tensor, Tensor | None]: - """Run one production shared-prefix GDN layer.""" + """Run one production prefix-tree GDN layer.""" _disable_reentrant_te_linear_transpose_cache(gdn) if hidden_states.ndim != 3: @@ -447,7 +448,7 @@ def run_gdn_layer( or int(group_ids.shape[1]) != expected_group_seq_len ): raise ValueError( - "shared-prefix GDN group_ids shape must match the logical sequence " + "prefix-tree GDN group_ids shape must match the logical sequence " "processed by Megatron GDN after sequence-parallel input gather, got " f"hidden={tuple(hidden_states.shape)} " f"group_ids={tuple(group_ids.shape)} " @@ -456,16 +457,14 @@ def run_gdn_layer( if require_prebuilt_plan and execution_plan is None: raise ValueError( - "ART shared-prefix GDN production path requires a prebuilt " - "GDN execution plan on SharedPrefixAttentionState. Build it once " - "per packed sequence via create_shared_prefix_state(..., " + "ART prefix-tree GDN production path requires a prebuilt " + "GDN execution plan on PrefixTreeAttentionState. Build it once " + "per packed sequence via create_prefix_tree_state(..., " "build_gdn_execution_spec=True)." ) if execution_spec is None and execution_plan is None: - execution_spec = parse_gdn_shared_prefix_segments( - group_ids, parent_ids, min_completions_per_family=0 - ) + execution_spec = parse_gdn_prefix_tree_segments(group_ids, parent_ids) if ( execution_spec is not None and requested_cp_size == 1 @@ -510,136 +509,396 @@ def run_gdn_layer( ) if input_layout != "attention" or output_layout != "attention": raise ValueError("GDN layout controls require a CP execution plan") - return _run_planned_prefixes_and_completions(gdn, hidden_states, execution_plan) + return _run_tree_prefixes(gdn, hidden_states, execution_plan) -def _run_planned_prefixes_and_completions( +def _run_tree_prefixes( gdn: Any, hidden_states: Tensor, plan: GdnRankExecutionPlan, ) -> tuple[Tensor, Tensor | None]: - if _has_chunk_aligned_local_plan(plan): - return _run_chunk_aligned_prefixes_and_completions(gdn, hidden_states, plan) - raise ValueError( - "shared-prefix GDN requires a chunk-aligned execution plan; " - "prefix/completion bucket execution has been removed" + qkv, gate, beta, recurrent_g = _project_gdn_inputs(gdn, hidden_states) + gate = gate.clone() + recurrent_output = torch.zeros_like(gate) + recurrent_output, _cp_dependency = _run_tree_depth_buckets( + gdn, + qkv, + beta, + recurrent_g, + recurrent_output, + plan, + state_reference=hidden_states, ) + return _project_gdn_output(gdn, recurrent_output, gate, plan) -def _has_chunk_aligned_local_plan(plan: GdnRankExecutionPlan) -> bool: - return bool( - plan.prefix_boundary_buckets - or plan.prefix_tail_buckets - or plan.completion_with_prefix_tail_buckets +def _run_tree_depth_buckets( + gdn: Any, + qkv: Tensor, + beta: Tensor, + recurrent_g: Tensor, + recurrent_output: Tensor, + plan: GdnRankExecutionPlan, + *, + state_reference: Tensor, + group: Any | None = None, + cp_dependency: Tensor | None = None, +) -> tuple[Tensor, Tensor | None]: + state_cache = _TreeStateChunkCache( + device=state_reference.device, ) + for depth, buckets in enumerate(plan.tree_segment_buckets_by_depth): + if depth < len(plan.tree_state_exchanges_by_depth): + cp_dependency = state_cache.exchange_remote_parent_states( + gdn, + plan.tree_state_exchanges_by_depth[depth], + state_reference=state_reference, + rank=plan.cp_rank, + group=group, + cp_dependency=cp_dependency, + ) + if depth < len(plan.tree_chain_buckets_by_depth): + for bucket in plan.tree_chain_buckets_by_depth[depth]: + recurrent_output, cp_dependency = _run_tree_bucket( + gdn, + qkv, + beta, + recurrent_g, + recurrent_output, + state_cache, + bucket, + state_reference=state_reference, + group=group, + cp_dependency=cp_dependency, + recurrent_cp=True, + scale_parent_state_gradient=1.0 / plan.cp_size, + ) + + for bucket in buckets: + recurrent_output, cp_dependency = _run_tree_bucket( + gdn, + qkv, + beta, + recurrent_g, + recurrent_output, + state_cache, + bucket, + state_reference=state_reference, + cp_dependency=cp_dependency, + ) -def _run_chunk_aligned_prefixes_and_completions( + return recurrent_output, cp_dependency + + +def _run_tree_bucket( gdn: Any, - hidden_states: Tensor, - plan: GdnRankExecutionPlan, + qkv: Tensor, + beta: Tensor, + recurrent_g: Tensor, + recurrent_output: Tensor, + state_cache: "_TreeStateChunkCache", + bucket: GdnSegmentBucketPlan, + *, + state_reference: Tensor, + group: Any | None = None, + cp_dependency: Tensor | None = None, + recurrent_cp: bool = False, + scale_parent_state_gradient: float | None = None, ) -> tuple[Tensor, Tensor | None]: - qkv, gate, beta, recurrent_g = _project_gdn_inputs(gdn, hidden_states) - gate = gate.clone() - recurrent_output = torch.zeros_like(gate) - boundary_family_chunks: list[Tensor] = [] - boundary_conv_chunks: list[Tensor] = [] - boundary_rec_chunks: list[Tensor] = [] - - for bucket in plan.prefix_boundary_buckets: - prefix_qkv, prefix_beta, prefix_g = _gather_bucket_streams( - qkv, beta, recurrent_g, bucket + parent_conv, parent_rec = state_cache.parent_states( + gdn, + bucket, + state_reference=state_reference, + ) + if _bucket_has_parent_state(bucket): + parent_conv, parent_rec = _couple_parent_states(parent_conv, parent_rec) + if scale_parent_state_gradient is not None: + parent_conv = _scale_state_gradient( + parent_conv, + scale_parent_state_gradient, + ) + parent_rec = _scale_state_gradient(parent_rec, scale_parent_state_gradient) + segment_qkv, segment_beta, segment_g = _gather_bucket_streams( + qkv, + beta, + recurrent_g, + bucket, + ) + if cp_dependency is not None: + segment_qkv = _add_autograd_dependency(segment_qkv, cp_dependency) + segment_beta = _add_autograd_dependency(segment_beta, cp_dependency) + segment_g = _add_autograd_dependency(segment_g, cp_dependency) + parent_conv = _add_autograd_dependency(parent_conv, cp_dependency) + parent_rec = _add_autograd_dependency(parent_rec, cp_dependency) + segment_out, segment_conv, segment_rec = run_gdn_bucket( + bucket, + (segment_qkv, segment_beta, segment_g), + (parent_conv, parent_rec), + gdn=gdn, + group=group, + recurrent_cp=recurrent_cp, + output_final_state=bucket.needs_final_state or recurrent_cp, + ) + if bucket.needs_final_state and (segment_conv is None or segment_rec is None): + raise RuntimeError("tree GDN execution must return final states") + if ( + bucket.needs_final_state + and segment_conv is not None + and segment_rec is not None + ): + cp_dependency = _make_autograd_dependency( + segment_out, segment_conv, segment_rec ) - zero_conv = _zero_conv_state( - gdn, hidden_states, batch_size=bucket.segment_count + else: + cp_dependency = _make_autograd_dependency(segment_out) + recurrent_output = _scatter_bucket_recurrent_output( + recurrent_output, + bucket, + segment_out, + ) + if bucket.needs_final_state: + state_cache.append( + bucket, + cast(Tensor, segment_conv), + cast(Tensor, segment_rec), + ) + return recurrent_output, cp_dependency + + +class _TreeStateChunkCache: + def __init__(self, *, device: torch.device) -> None: + self._device = device + self._conv_chunks: list[Tensor] = [] + self._rec_chunks: list[Tensor] = [] + self._source_by_family: list[tuple[int, int] | None] = [] + + def append(self, bucket: GdnSegmentBucketPlan, conv: Tensor, rec: Tensor) -> None: + self.append_families(_bucket_family_indices_cpu(bucket), conv, rec) + + def append_families( + self, family_indices: Sequence[int], conv: Tensor, rec: Tensor + ) -> None: + if len(family_indices) == 0: + return + if int(conv.shape[0]) != len(family_indices): + raise ValueError( + "tree GDN state cache conv batch must match family count, got " + f"{tuple(conv.shape)} and {len(family_indices)} families" + ) + if int(rec.shape[0]) != len(family_indices): + raise ValueError( + "tree GDN state cache recurrent batch must match family count, got " + f"{tuple(rec.shape)} and {len(family_indices)} families" + ) + chunk_index = len(self._conv_chunks) + self._conv_chunks.append(conv) + self._rec_chunks.append(rec) + max_family = max(int(index) for index in family_indices) + if max_family >= len(self._source_by_family): + self._source_by_family.extend( + None for _ in range(max_family + 1 - len(self._source_by_family)) + ) + for source_row, family_index in enumerate(family_indices): + self._source_by_family[int(family_index)] = (chunk_index, source_row) + + def exchange_remote_parent_states( + self, + gdn: Any, + exchange: GdnStateExchangePlan | None, + *, + state_reference: Tensor, + rank: int, + group: Any | None, + cp_dependency: Tensor | None, + ) -> Tensor | None: + if exchange is None: + return cp_dependency + from .layout import exchange_rank_tensor_all_to_all + + source_conv, source_rec = self.states_for_families( + gdn, + exchange.source_family_indices, + state_reference=state_reference, + ) + if cp_dependency is not None: + source_conv = _add_autograd_dependency(source_conv, cp_dependency) + source_rec = _add_autograd_dependency(source_rec, cp_dependency) + remote_conv = exchange_rank_tensor_all_to_all( + source_conv, + exchange.exchange, + rank=rank, + group=group, + backward_plan=exchange.reverse_exchange, ) - zero_rec = _zero_recurrent_state( - gdn, hidden_states, batch_size=bucket.segment_count + remote_rec = exchange_rank_tensor_all_to_all( + source_rec, + exchange.exchange, + rank=rank, + group=group, + backward_plan=exchange.reverse_exchange, ) - prefix_out, prefix_conv, prefix_rec = run_gdn_bucket( - bucket, - (prefix_qkv, prefix_beta, prefix_g), - (zero_conv, zero_rec), - gdn=gdn, - output_final_state=True, - ) - if prefix_conv is None or prefix_rec is None: - raise RuntimeError("prefix boundary GDN execution must return final states") - recurrent_output = _scatter_bucket_recurrent_output( - recurrent_output, bucket, prefix_out - ) - boundary_family_chunks.append(bucket.family_indices) - boundary_conv_chunks.append(prefix_conv) - boundary_rec_chunks.append(prefix_rec) - - boundary_conv_table = _materialize_indexed_family_state_table( - plan=plan, - family_chunks=boundary_family_chunks, - state_chunks=boundary_conv_chunks, - zero_state=_zero_conv_state(gdn, hidden_states, batch_size=plan.family_count), - ) - boundary_rec_table = _materialize_indexed_family_state_table( - plan=plan, - family_chunks=boundary_family_chunks, - state_chunks=boundary_rec_chunks, - zero_state=_zero_recurrent_state( - gdn, hidden_states, batch_size=plan.family_count - ), - ) - - tail_family_chunks: list[Tensor] = [] - tail_conv_chunks: list[Tensor] = [] - tail_rec_chunks: list[Tensor] = [] - for bucket in plan.prefix_tail_buckets: - tail_qkv, tail_beta, tail_g = _gather_bucket_streams( - qkv, beta, recurrent_g, bucket - ) - tail_conv = boundary_conv_table.index_select(0, bucket.family_indices) - tail_rec = boundary_rec_table.index_select(0, bucket.family_indices) - tail_out, tail_conv, tail_rec = run_gdn_bucket( - bucket, - (tail_qkv, tail_beta, tail_g), - (tail_conv, tail_rec), - gdn=gdn, - output_final_state=True, + self.append_families(exchange.dest_family_indices, remote_conv, remote_rec) + dependency = _make_zero_autograd_dependency( + source_conv, source_rec, remote_conv, remote_rec ) - if tail_conv is None or tail_rec is None: - raise RuntimeError("prefix tail GDN execution must return final states") - recurrent_output = _scatter_bucket_recurrent_output( - recurrent_output, bucket, tail_out + return dependency if cp_dependency is None else dependency + cp_dependency + + def states_for_families( + self, + gdn: Any, + family_indices: Sequence[int], + *, + state_reference: Tensor, + ) -> tuple[Tensor, Tensor]: + if len(family_indices) == 0: + conv = _zero_conv_state(gdn, state_reference, batch_size=0) + rec = _zero_recurrent_state(gdn, state_reference, batch_size=0) + return conv.requires_grad_(True), rec.requires_grad_(True) + return self._mixed_parent_states( + gdn, + tuple(int(index) for index in family_indices), + state_reference=state_reference, + batch_size=len(family_indices), + roots_allowed=False, ) - tail_family_chunks.append(bucket.family_indices) - tail_conv_chunks.append(tail_conv) - tail_rec_chunks.append(tail_rec) - prefix_conv_table = _replace_indexed_family_states( - boundary_conv_table, - family_chunks=tail_family_chunks, - state_chunks=tail_conv_chunks, - ) - prefix_rec_table = _replace_indexed_family_states( - boundary_rec_table, - family_chunks=tail_family_chunks, - state_chunks=tail_rec_chunks, - ) + def parent_states( + self, + gdn: Any, + bucket: GdnSegmentBucketPlan, + *, + state_reference: Tensor, + ) -> tuple[Tensor, Tensor]: + parent_indices = bucket.parent_indices + if parent_indices is None: + raise RuntimeError("tree GDN bucket is missing parent indices") + parent_indices_cpu = _bucket_parent_indices_cpu(bucket) + batch_size = bucket.segment_count + if all(parent_index < 0 for parent_index in parent_indices_cpu): + return ( + _zero_conv_state(gdn, state_reference, batch_size=batch_size), + _zero_recurrent_state(gdn, state_reference, batch_size=batch_size), + ) - for bucket in plan.completion_with_prefix_tail_buckets: - completion_conv = prefix_conv_table.index_select(0, bucket.family_indices) - completion_rec = prefix_rec_table.index_select(0, bucket.family_indices) - completion_qkv, completion_beta, completion_g = _gather_bucket_streams( - qkv, beta, recurrent_g, bucket - ) - completion_out, _, _ = run_gdn_bucket( - bucket, - (completion_qkv, completion_beta, completion_g), - (completion_conv, completion_rec), - gdn=gdn, - output_final_state=False, + return self._mixed_parent_states( + gdn, + parent_indices_cpu, + state_reference=state_reference, + batch_size=batch_size, ) - recurrent_output = _scatter_bucket_recurrent_output( - recurrent_output, bucket, completion_out + + def _mixed_parent_states( + self, + gdn: Any, + parent_indices_cpu: tuple[int, ...], + *, + state_reference: Tensor, + batch_size: int, + roots_allowed: bool = True, + ) -> tuple[Tensor, Tensor]: + sources_by_chunk: dict[int, list[tuple[int, int]]] = {} + missing_parents: list[int] = [] + for dest_row, parent_index in enumerate(parent_indices_cpu): + if parent_index < 0: + if roots_allowed: + continue + missing_parents.append(parent_index) + continue + source = ( + self._source_by_family[parent_index] + if parent_index < len(self._source_by_family) + else None + ) + if source is None: + missing_parents.append(parent_index) + continue + chunk_index, source_row = source + sources_by_chunk.setdefault(chunk_index, []).append((dest_row, source_row)) + if missing_parents: + raise RuntimeError( + "tree GDN append-only execution is missing parent state for " + f"families {tuple(missing_parents)}" + ) + + single_source_chunk = next(iter(sources_by_chunk.values())) + if len(sources_by_chunk) == 1 and len(single_source_chunk) == batch_size: + chunk_index, pairs = next(iter(sources_by_chunk.items())) + return ( + _select_state_rows(self._conv_chunks[chunk_index], pairs), + _select_state_rows(self._rec_chunks[chunk_index], pairs), + ) + + conv = _zero_conv_state(gdn, state_reference, batch_size=batch_size) + rec = _zero_recurrent_state(gdn, state_reference, batch_size=batch_size) + for chunk_index, pairs in sources_by_chunk.items(): + dest_rows = _long_tensor( + (dest_row for dest_row, _ in pairs), + device=self._device, + ) + source_rows = _long_tensor( + (source_row for _, source_row in pairs), + device=self._device, + ) + conv = conv.index_copy( + 0, + dest_rows, + self._conv_chunks[chunk_index].index_select(0, source_rows), + ) + rec = rec.index_copy( + 0, + dest_rows, + self._rec_chunks[chunk_index].index_select(0, source_rows), + ) + return conv, rec + + +def _select_state_rows(chunk: Tensor, pairs: Sequence[tuple[int, int]]) -> Tensor: + source_rows = tuple(source_row for _, source_row in pairs) + if len(set(source_rows)) == 1: + return chunk.narrow(0, source_rows[0], 1).expand( + len(source_rows), + *tuple(chunk.shape[1:]), ) - return _project_gdn_output(gdn, recurrent_output, gate, plan) + first_row = source_rows[0] + if source_rows == tuple(range(first_row, first_row + len(source_rows))): + return chunk.narrow(0, first_row, len(source_rows)) + return chunk.index_select( + 0, + _long_tensor(source_rows, device=chunk.device), + ) + + +def _bucket_family_indices_cpu(bucket: GdnSegmentBucketPlan) -> tuple[int, ...]: + family_indices = bucket.family_indices_cpu + if family_indices is None: + family_indices = bucket.family_indices.detach().cpu() + return tuple(int(index) for index in family_indices.tolist()) + + +def _bucket_parent_indices_cpu(bucket: GdnSegmentBucketPlan) -> tuple[int, ...]: + parent_indices = bucket.parent_indices + if parent_indices is None: + raise RuntimeError("tree GDN bucket is missing parent indices") + parent_indices_cpu = bucket.parent_indices_cpu + if parent_indices_cpu is None: + parent_indices_cpu = parent_indices.detach().cpu() + return tuple(int(index) for index in parent_indices_cpu.tolist()) + + +def _long_tensor(values: Iterable[int], *, device: torch.device) -> Tensor: + return torch.tensor(tuple(values), dtype=torch.long, device=device) + + +def _bucket_has_parent_state(bucket: GdnSegmentBucketPlan) -> bool: + return any(parent_index >= 0 for parent_index in _bucket_parent_indices_cpu(bucket)) + + +def _bucket_has_uniform_lengths(bucket: GdnSegmentBucketPlan) -> bool: + lengths_cpu = bucket.lengths_cpu + if lengths_cpu is None: + lengths_cpu = bucket.lengths.detach().cpu() + return all(int(length) == int(bucket.length) for length in lengths_cpu.tolist()) def _run_cp_planned_prefixes_and_completions( @@ -679,385 +938,21 @@ def _run_cp_planned_prefixes_and_completions( if empty_gdn_rank else _empty_autograd_dependency(qkv) ) - qkv_with_remote_tail = qkv - beta_with_remote_tail = beta - recurrent_g_with_remote_tail = recurrent_g - if plan.remote_prefix_tail_exchange is not None: - remote_qkv, remote_beta, remote_g = _exchange_remote_prefix_tail_streams( - qkv, - beta, - recurrent_g, - plan=plan, - group=group, - ) - qkv_with_remote_tail = torch.cat([qkv, remote_qkv.unsqueeze(0)], dim=1) - beta_with_remote_tail = torch.cat([beta, remote_beta.unsqueeze(0)], dim=1) - recurrent_g_with_remote_tail = torch.cat( - [recurrent_g, remote_g.unsqueeze(0)], dim=1 - ) - cp_dependency = cp_dependency + _make_zero_autograd_dependency( - remote_qkv, remote_beta, remote_g - ) + if not plan.tree_segment_buckets_by_depth: + raise ValueError("CP prefix-tree GDN requires a tree execution plan") gate = gate.clone() recurrent_output = torch.zeros_like(gate) - prefix_family_chunks: list[Tensor] = [] - prefix_conv_chunks: list[Tensor] = [] - prefix_rec_chunks: list[Tensor] = [] - - for bucket in plan.chain_prefix_buckets: - prefix_qkv, prefix_beta, prefix_g = _gather_bucket_streams( - qkv, beta, recurrent_g, bucket - ) - zero_conv = _zero_conv_state(gdn, qkv, batch_size=bucket.segment_count) - zero_rec = _zero_recurrent_state(gdn, qkv, batch_size=bucket.segment_count) - prefix_out, prefix_conv, prefix_rec = run_gdn_bucket( - bucket, - (prefix_qkv, prefix_beta, prefix_g), - (zero_conv, zero_rec), - gdn=gdn, - group=group, - recurrent_cp=True, - output_final_state=True, - ) - if prefix_conv is None or prefix_rec is None: - raise RuntimeError("CP prefix GDN execution must return final states") - prefix_out = _add_autograd_dependency(prefix_out, cp_dependency) - prefix_conv = _add_autograd_dependency(prefix_conv, cp_dependency) - prefix_rec = _add_autograd_dependency(prefix_rec, cp_dependency) - cp_dependency = _make_autograd_dependency(prefix_out, prefix_conv, prefix_rec) - recurrent_output = _scatter_bucket_recurrent_output( - recurrent_output, bucket, prefix_out - ) - prefix_family_chunks.append(bucket.family_indices) - prefix_conv_chunks.append(prefix_conv) - prefix_rec_chunks.append(prefix_rec) - - boundary_family_chunks: list[Tensor] = [] - boundary_conv_chunks: list[Tensor] = [] - boundary_rec_chunks: list[Tensor] = [] - for bucket in plan.prefix_boundary_buckets: - prefix_qkv, prefix_beta, prefix_g = _gather_bucket_streams( - qkv, beta, recurrent_g, bucket - ) - zero_conv = _zero_conv_state(gdn, qkv, batch_size=bucket.segment_count) - zero_rec = _zero_recurrent_state(gdn, qkv, batch_size=bucket.segment_count) - prefix_out, prefix_conv, prefix_rec = run_gdn_bucket( - bucket, - (prefix_qkv, prefix_beta, prefix_g), - (zero_conv, zero_rec), - gdn=gdn, - output_final_state=True, - ) - if prefix_conv is None or prefix_rec is None: - raise RuntimeError("local prefix GDN execution must return final states") - prefix_out = _add_autograd_dependency(prefix_out, cp_dependency) - prefix_conv = _add_autograd_dependency(prefix_conv, cp_dependency) - prefix_rec = _add_autograd_dependency(prefix_rec, cp_dependency) - recurrent_output = _scatter_bucket_recurrent_output( - recurrent_output, bucket, prefix_out - ) - boundary_family_chunks.append(bucket.family_indices) - boundary_conv_chunks.append(prefix_conv) - boundary_rec_chunks.append(prefix_rec) - prefix_family_chunks.append(bucket.family_indices) - prefix_conv_chunks.append(prefix_conv) - prefix_rec_chunks.append(prefix_rec) - - if ( - plan.prefix_tail_buckets - or plan.remote_prefix_tail_buckets - or plan.completion_with_prefix_tail_buckets - or plan.remote_completion_with_prefix_tail_buckets - or plan.remote_prefix_tail_state_transfers - ): - boundary_conv_table = _materialize_indexed_family_state_table( - plan=plan, - family_chunks=boundary_family_chunks, - state_chunks=boundary_conv_chunks, - zero_state=_zero_conv_state(gdn, qkv, batch_size=plan.family_count), - ) - boundary_rec_table = _materialize_indexed_family_state_table( - plan=plan, - family_chunks=boundary_family_chunks, - state_chunks=boundary_rec_chunks, - zero_state=_zero_recurrent_state(gdn, qkv, batch_size=plan.family_count), - ) - remote_boundary_conv_table = boundary_conv_table - remote_boundary_rec_table = boundary_rec_table - if plan.remote_prefix_tail_state_transfers: - ( - remote_boundary_conv_table, - remote_boundary_rec_table, - remote_boundary_dependency, - ) = _exchange_parent_state_rows( - boundary_conv_table, - boundary_rec_table, - transfers=plan.remote_prefix_tail_state_transfers, - group=group, - ) - cp_dependency = cp_dependency + remote_boundary_dependency - tail_family_chunks: list[Tensor] = [] - tail_conv_chunks: list[Tensor] = [] - tail_rec_chunks: list[Tensor] = [] - for bucket in plan.prefix_tail_buckets: - tail_qkv, tail_beta, tail_g = _gather_bucket_streams( - qkv, beta, recurrent_g, bucket - ) - tail_conv = boundary_conv_table.index_select(0, bucket.family_indices) - tail_rec = boundary_rec_table.index_select(0, bucket.family_indices) - tail_out, tail_conv, tail_rec = run_gdn_bucket( - bucket, - (tail_qkv, tail_beta, tail_g), - (tail_conv, tail_rec), - gdn=gdn, - output_final_state=True, - ) - if tail_conv is None or tail_rec is None: - raise RuntimeError("local prefix tail GDN execution must return states") - tail_out = _add_autograd_dependency(tail_out, cp_dependency) - tail_conv = _add_autograd_dependency(tail_conv, cp_dependency) - tail_rec = _add_autograd_dependency(tail_rec, cp_dependency) - recurrent_output = _scatter_bucket_recurrent_output( - recurrent_output, bucket, tail_out - ) - tail_family_chunks.append(bucket.family_indices) - tail_conv_chunks.append(tail_conv) - tail_rec_chunks.append(tail_rec) - prefix_family_chunks.append(bucket.family_indices) - prefix_conv_chunks.append(tail_conv) - prefix_rec_chunks.append(tail_rec) - for bucket in plan.remote_prefix_tail_buckets: - tail_qkv, tail_beta, tail_g = _gather_bucket_streams( - qkv_with_remote_tail, - beta_with_remote_tail, - recurrent_g_with_remote_tail, - bucket, - ) - tail_conv = remote_boundary_conv_table.index_select( - 0, bucket.family_indices - ) - tail_rec = remote_boundary_rec_table.index_select(0, bucket.family_indices) - tail_out, tail_conv, tail_rec = run_gdn_bucket( - bucket, - (tail_qkv, tail_beta, tail_g), - (tail_conv, tail_rec), - gdn=gdn, - output_final_state=True, - ) - if tail_conv is None or tail_rec is None: - raise RuntimeError( - "remote prefix tail GDN execution must return states" - ) - tail_out = _add_autograd_dependency(tail_out, cp_dependency) - tail_conv = _add_autograd_dependency(tail_conv, cp_dependency) - tail_rec = _add_autograd_dependency(tail_rec, cp_dependency) - tail_family_chunks.append(bucket.family_indices) - tail_conv_chunks.append(tail_conv) - tail_rec_chunks.append(tail_rec) - prefix_family_chunks.append(bucket.family_indices) - prefix_conv_chunks.append(tail_conv) - prefix_rec_chunks.append(tail_rec) - prefix_conv_table = _replace_indexed_family_states( - boundary_conv_table, - family_chunks=tail_family_chunks, - state_chunks=tail_conv_chunks, - ) - prefix_rec_table = _replace_indexed_family_states( - boundary_rec_table, - family_chunks=tail_family_chunks, - state_chunks=tail_rec_chunks, - ) - for bucket in plan.completion_with_prefix_tail_buckets: - completion_conv = prefix_conv_table.index_select(0, bucket.family_indices) - completion_rec = prefix_rec_table.index_select(0, bucket.family_indices) - completion_conv, completion_rec = _couple_parent_states( - completion_conv, completion_rec - ) - completion_qkv, completion_beta, completion_g = _gather_bucket_streams( - qkv, beta, recurrent_g, bucket - ) - completion_out, _, _ = run_gdn_bucket( - bucket, - (completion_qkv, completion_beta, completion_g), - (completion_conv, completion_rec), - gdn=gdn, - output_final_state=False, - ) - completion_out = _add_autograd_dependency(completion_out, cp_dependency) - recurrent_output = _scatter_bucket_recurrent_output( - recurrent_output, bucket, completion_out - ) - for bucket in plan.remote_completion_with_prefix_tail_buckets: - completion_conv = prefix_conv_table.index_select(0, bucket.family_indices) - completion_rec = prefix_rec_table.index_select(0, bucket.family_indices) - completion_conv, completion_rec = _couple_parent_states( - completion_conv, completion_rec - ) - completion_qkv, completion_beta, completion_g = _gather_bucket_streams( - qkv, - beta, - recurrent_g, - bucket, - ) - completion_out, _, _ = run_gdn_bucket( - bucket, - (completion_qkv, completion_beta, completion_g), - (completion_conv, completion_rec), - gdn=gdn, - output_final_state=False, - ) - completion_out = _add_autograd_dependency(completion_out, cp_dependency) - recurrent_output = _scatter_bucket_recurrent_output( - recurrent_output, bucket, completion_out - ) - - for bucket in plan.local_prefix_buckets: - prefix_qkv, prefix_beta, prefix_g = _gather_bucket_streams( - qkv, beta, recurrent_g, bucket - ) - zero_conv = _zero_conv_state(gdn, qkv, batch_size=bucket.segment_count) - zero_rec = _zero_recurrent_state(gdn, qkv, batch_size=bucket.segment_count) - prefix_out, prefix_conv, prefix_rec = run_gdn_bucket( - bucket, - (prefix_qkv, prefix_beta, prefix_g), - (zero_conv, zero_rec), - gdn=gdn, - output_final_state=True, - ) - if prefix_conv is None or prefix_rec is None: - raise RuntimeError("local prefix GDN execution must return final states") - prefix_out = _add_autograd_dependency(prefix_out, cp_dependency) - prefix_conv = _add_autograd_dependency(prefix_conv, cp_dependency) - prefix_rec = _add_autograd_dependency(prefix_rec, cp_dependency) - recurrent_output = _scatter_bucket_recurrent_output( - recurrent_output, bucket, prefix_out - ) - prefix_family_chunks.append(bucket.family_indices) - prefix_conv_chunks.append(prefix_conv) - prefix_rec_chunks.append(prefix_rec) - - if not prefix_conv_chunks and not plan.parent_state_exchange_family_indices: - projected, out_bias = _project_cp_gdn_output( - gdn, - recurrent_output, - gate, - plan, - group=group, - output_layout=output_layout, - ) - projected = _add_autograd_dependency(projected, cp_dependency) - return projected, out_bias - - prefix_conv_table = _materialize_ordered_family_state_table( - family_chunks=prefix_family_chunks, - state_chunks=prefix_conv_chunks, - zero_state=_zero_conv_state(gdn, qkv, batch_size=plan.family_count), - ) - prefix_rec_table = _materialize_ordered_family_state_table( - family_chunks=prefix_family_chunks, - state_chunks=prefix_rec_chunks, - zero_state=_zero_recurrent_state(gdn, qkv, batch_size=plan.family_count), - ) - parent_state_exchanged = False - if plan.chain_completion_buckets and plan.parent_state_exchange_family_indices: - if not plan.parent_state_transfers: - raise ValueError("CP parent-state exchange requires planned transfers") - prefix_conv_table, prefix_rec_table, exchange_dependency = ( - _exchange_parent_state_rows( - prefix_conv_table, - prefix_rec_table, - transfers=plan.parent_state_transfers, - group=group, - ) - ) - cp_dependency = cp_dependency + exchange_dependency - parent_state_exchanged = True - for bucket in plan.chain_completion_buckets: - completion_qkv, completion_beta, completion_g = _gather_bucket_streams( - qkv, beta, recurrent_g, bucket - ) - completion_conv = prefix_conv_table.index_select(0, bucket.family_indices) - completion_rec = prefix_rec_table.index_select(0, bucket.family_indices) - completion_conv, completion_rec = _couple_parent_states( - completion_conv, completion_rec - ) - completion_conv = _scale_state_gradient(completion_conv, 1.0 / plan.cp_size) - completion_rec = _scale_state_gradient(completion_rec, 1.0 / plan.cp_size) - completion_out, _, _ = run_gdn_bucket( - bucket, - (completion_qkv, completion_beta, completion_g), - (completion_conv, completion_rec), - gdn=gdn, - group=group, - recurrent_cp=True, - output_final_state=False, - ) - completion_out = _add_autograd_dependency(completion_out, cp_dependency) - cp_dependency = _make_autograd_dependency(completion_out) - recurrent_output = _scatter_bucket_recurrent_output( - recurrent_output, bucket, completion_out - ) - - ready_completion_buckets = ( - plan.ready_local_completion_buckets - if plan.ready_local_completion_buckets or plan.remote_local_completion_buckets - else plan.local_completion_buckets + recurrent_output, cp_dependency = _run_tree_depth_buckets( + gdn, + qkv, + beta, + recurrent_g, + recurrent_output, + plan, + state_reference=qkv, + group=group, + cp_dependency=cp_dependency, ) - for bucket in ready_completion_buckets: - completion_qkv, completion_beta, completion_g = _gather_bucket_streams( - qkv, beta, recurrent_g, bucket - ) - completion_conv = prefix_conv_table.index_select(0, bucket.family_indices) - completion_rec = prefix_rec_table.index_select(0, bucket.family_indices) - completion_conv, completion_rec = _couple_parent_states( - completion_conv, completion_rec - ) - completion_out, _, _ = run_gdn_bucket( - bucket, - (completion_qkv, completion_beta, completion_g), - (completion_conv, completion_rec), - gdn=gdn, - output_final_state=False, - ) - completion_out = _add_autograd_dependency(completion_out, cp_dependency) - recurrent_output = _scatter_bucket_recurrent_output( - recurrent_output, bucket, completion_out - ) - - if plan.parent_state_exchange_family_indices and not parent_state_exchanged: - if not plan.parent_state_transfers: - raise ValueError("CP parent-state exchange requires planned transfers") - prefix_conv_table, prefix_rec_table, exchange_dependency = ( - _exchange_parent_state_rows( - prefix_conv_table, - prefix_rec_table, - transfers=plan.parent_state_transfers, - group=group, - ) - ) - cp_dependency = cp_dependency + exchange_dependency - - for bucket in plan.remote_local_completion_buckets: - completion_qkv, completion_beta, completion_g = _gather_bucket_streams( - qkv, beta, recurrent_g, bucket - ) - completion_conv = prefix_conv_table.index_select(0, bucket.family_indices) - completion_rec = prefix_rec_table.index_select(0, bucket.family_indices) - completion_conv, completion_rec = _couple_parent_states( - completion_conv, completion_rec - ) - completion_out, _, _ = run_gdn_bucket( - bucket, - (completion_qkv, completion_beta, completion_g), - (completion_conv, completion_rec), - gdn=gdn, - output_final_state=False, - ) - completion_out = _add_autograd_dependency(completion_out, cp_dependency) - recurrent_output = _scatter_bucket_recurrent_output( - recurrent_output, bucket, completion_out - ) - projected, out_bias = _project_cp_gdn_output( gdn, recurrent_output, @@ -1065,8 +960,8 @@ def _run_cp_planned_prefixes_and_completions( plan, group=group, output_layout=output_layout, + dependency=cp_dependency, ) - projected = _add_autograd_dependency(projected, cp_dependency) return projected, out_bias @@ -1659,12 +1554,6 @@ def _local_layout_token_count_for_hidden( return (real_count + _tp_world_size(projection) - 1) // _tp_world_size(projection) -def _attention_original_shape_from_plan( - hidden_states: Tensor, plan: GdnRankExecutionPlan -) -> tuple[int, int, int]: - return (int(plan.attention_token_count), 1, int(hidden_states.shape[-1])) - - def _restore_hidden_from_cp_flat( flat: Tensor, original_shape: tuple[int, int, int] ) -> Tensor: @@ -1847,9 +1736,7 @@ def _gather_bucket_streams( recurrent_g.reshape(-1, int(recurrent_g.shape[-1])), bucket.row_indices, bucket.position_indices, - bucket.cu_seqlens, token_count=int(bucket.real_token_count), - segment_count=int(bucket.segment_count), sequence_length=int(qkv.shape[1]), ) @@ -1922,6 +1809,7 @@ def _project_cp_gdn_output( *, group: Any, output_layout: Literal["attention", "gdn"], + dependency: Tensor | None = None, ) -> tuple[Tensor, Tensor | None]: batch_size, seq_len, _, _ = recurrent_output.shape token_uids = ( @@ -1933,6 +1821,8 @@ def _project_cp_gdn_output( norm_out = _apply_gated_rms_norm(gdn, recurrent_output, gate) norm_out = norm_out.reshape(batch_size, seq_len, _local_value_dim(gdn)) norm_out = norm_out.transpose(0, 1).contiguous() + if dependency is not None: + norm_out = _add_autograd_dependency(norm_out, dependency) if token_uids is not None: token_uids = _replicated_layout_token_uids(plan, "gdn", hidden_states=norm_out) _attach_trace_token_uids(norm_out, token_uids) @@ -2271,6 +2161,36 @@ def _local_value_dim(gdn: Any) -> int: return _local_value_heads(gdn) * int(gdn.value_head_dim) +def _prepare_dense_recurrent_inputs( + qkv: Tensor, + beta: Tensor, + recurrent_g: Tensor, + *, + key_heads: int, + value_heads: int, + key_dim: int, + value_dim: int, +) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + key_channels = int(key_heads) * int(key_dim) + value_channels = int(value_heads) * int(value_dim) + query = qkv[..., :key_channels].reshape(*qkv.shape[:2], key_heads, key_dim) + key = qkv[..., key_channels : 2 * key_channels].reshape( + *qkv.shape[:2], + key_heads, + key_dim, + ) + value = qkv[..., 2 * key_channels : 2 * key_channels + value_channels].reshape( + *qkv.shape[:2], + value_heads, + value_dim, + ) + repeat = int(value_heads) // int(key_heads) + if repeat != 1: + query = query.repeat_interleave(repeat, dim=2) + key = key.repeat_interleave(repeat, dim=2) + return query, key, value, beta, recurrent_g + + def _scatter_bucket_recurrent_output( output: Tensor, bucket: GdnSegmentBucketPlan, bucket_output: Tensor ) -> Tensor: @@ -2280,7 +2200,6 @@ def _scatter_bucket_recurrent_output( bucket.row_indices, bucket.position_indices, _bucket_output_mask(bucket), - bucket.cu_seqlens, ) @@ -2289,269 +2208,6 @@ def _bucket_output_mask(bucket: GdnSegmentBucketPlan) -> Tensor: return bucket.real_mask if output_mask is None else output_mask -def _materialize_indexed_family_state_table( - *, - plan: GdnRankExecutionPlan, - family_chunks: list[Tensor], - state_chunks: list[Tensor], - zero_state: Tensor, -) -> Tensor: - table = zero_state.detach() - if not state_chunks: - return table.requires_grad_(True) - values = torch.cat(state_chunks, dim=0) - family_indices = torch.cat(family_chunks, dim=0) - return table.index_copy(0, family_indices, values) - - -def _materialize_ordered_family_state_table( - *, - family_chunks: list[Tensor], - state_chunks: list[Tensor], - zero_state: Tensor, -) -> Tensor: - if len(family_chunks) != len(state_chunks): - raise RuntimeError("family and state chunk counts must match") - table = zero_state.detach().requires_grad_(True) - for family_indices, states in zip(family_chunks, state_chunks, strict=True): - table = table.index_copy(0, family_indices, states) - return table - - -def _replace_indexed_family_states( - table: Tensor, - *, - family_chunks: list[Tensor], - state_chunks: list[Tensor], -) -> Tensor: - if not state_chunks: - return table - return table.index_copy( - 0, - torch.cat(family_chunks, dim=0), - torch.cat(state_chunks, dim=0), - ) - - -def _exchange_parent_state_rows( - conv_table: Tensor, - rec_table: Tensor, - *, - transfers: tuple[GdnParentStateTransferPlan, ...], - group: Any, -) -> tuple[Tensor, Tensor, Tensor]: - if not transfers: - return conv_table, rec_table, _empty_autograd_dependency(conv_table) - conv_table, rec_table = _ParentStateExchange.apply( - conv_table, rec_table, transfers, group - ) - return conv_table, rec_table, _make_autograd_dependency(conv_table, rec_table) - - -def _exchange_remote_prefix_tail_streams( - qkv: Tensor, - beta: Tensor, - recurrent_g: Tensor, - *, - plan: GdnRankExecutionPlan, - group: Any, -) -> tuple[Tensor, Tensor, Tensor]: - from .layout import exchange_rank_tensor_all_to_all - - if plan.remote_prefix_tail_exchange is None: - return ( - qkv.new_empty((0, int(qkv.shape[-1]))), - beta.new_empty((0, int(beta.shape[-1]))), - recurrent_g.new_empty((0, int(recurrent_g.shape[-1]))), - ) - if plan.remote_prefix_tail_backward_exchange is None: - raise ValueError("remote prefix-tail exchange requires a backward plan") - qkv_flat = qkv.reshape(-1, int(qkv.shape[-1])) - beta_flat = beta.reshape(-1, int(beta.shape[-1])) - g_flat = recurrent_g.reshape(-1, int(recurrent_g.shape[-1])) - kwargs = { - "plan": plan.remote_prefix_tail_exchange, - "rank": plan.cp_rank, - "group": group, - "backward_plan": plan.remote_prefix_tail_backward_exchange, - } - return ( - exchange_rank_tensor_all_to_all(qkv_flat, **kwargs), - exchange_rank_tensor_all_to_all(beta_flat, **kwargs), - exchange_rank_tensor_all_to_all(g_flat, **kwargs), - ) - - -class _ParentStateExchange(torch.autograd.Function): - @staticmethod - def forward( - ctx: Any, - conv_table: Tensor, - rec_table: Tensor, - transfers: tuple[GdnParentStateTransferPlan, ...], - group: Any, - ) -> tuple[Tensor, Tensor]: - ctx.group = group - ctx.transfers = transfers - ctx.save_for_backward(conv_table, rec_table) - return ( - _exchange_parent_state_tensor_forward( - conv_table, - transfers, - group=group, - ), - _exchange_parent_state_tensor_forward( - rec_table, - transfers, - group=group, - ), - ) - - @staticmethod - def backward( - ctx: Any, *grad_outputs: Tensor | None - ) -> tuple[Tensor | None, Tensor | None, None, None]: - grad_conv, grad_rec = grad_outputs - conv_ref, rec_ref = ctx.saved_tensors - return ( - _exchange_parent_state_tensor_backward( - _zero_if_none(grad_conv, conv_ref), - ctx.transfers, - group=ctx.group, - ), - _exchange_parent_state_tensor_backward( - _zero_if_none(grad_rec, rec_ref), - ctx.transfers, - group=ctx.group, - ), - None, - None, - ) - - -def _exchange_parent_state_tensor_forward( - table: Tensor, - transfers: tuple[GdnParentStateTransferPlan, ...], - *, - group: Any, -) -> Tensor: - rank = torch.distributed.get_rank(group) # ty: ignore[possibly-missing-attribute] - output = table.clone() - recvs = _exchange_parent_state_rows_all_to_all( - table, transfers, rank=rank, reverse=False, group=group - ) - for transfer, rows in recvs: - index = _parent_state_index_tensor(transfer, device=table.device) - output.index_copy_(0, index, rows) - return output - - -def _exchange_parent_state_tensor_backward( - grad_output: Tensor, - transfers: tuple[GdnParentStateTransferPlan, ...], - *, - group: Any, -) -> Tensor: - rank = torch.distributed.get_rank(group) # ty: ignore[possibly-missing-attribute] - grad_input = grad_output.clone() - for transfer in transfers: - if transfer.dest_rank != rank: - continue - index = _parent_state_index_tensor(transfer, device=grad_output.device) - grad_input.index_fill_(0, index, 0) - recvs = _exchange_parent_state_rows_all_to_all( - grad_output, transfers, rank=rank, reverse=True, group=group - ) - for transfer, rows in recvs: - index = _parent_state_index_tensor(transfer, device=grad_output.device) - grad_input.index_add_(0, index, rows) - return grad_input - - -def _zero_if_none(grad: Tensor | None, reference: Tensor) -> Tensor: - if grad is None: - return reference.new_zeros(reference.shape) - return grad.contiguous() - - -def _exchange_parent_state_rows_all_to_all( - table: Tensor, - transfers: tuple[GdnParentStateTransferPlan, ...], - *, - rank: int, - reverse: bool, - group: Any, -) -> list[tuple[GdnParentStateTransferPlan, Tensor]]: - world_size = torch.distributed.get_world_size(group) # ty: ignore[possibly-missing-attribute] - send_counts = [0 for _ in range(world_size)] - recv_counts = [0 for _ in range(world_size)] - send_pieces: list[Tensor] = [] - for peer_rank in range(world_size): - for transfer in transfers: - send_rank = transfer.dest_rank if reverse else transfer.source_rank - recv_rank = transfer.source_rank if reverse else transfer.dest_rank - if send_rank == recv_rank: - continue - row_count = len(transfer.family_indices) - if rank == send_rank and peer_rank == recv_rank: - index = _parent_state_index_tensor(transfer, device=table.device) - send_pieces.append(table.index_select(0, index).contiguous()) - send_counts[peer_rank] += row_count - if rank == recv_rank and peer_rank == send_rank: - recv_counts[peer_rank] += row_count - - trailing_shape = tuple(table.shape[1:]) - send_buffer = ( - torch.cat(send_pieces, dim=0) - if send_pieces - else table.new_empty((0, *trailing_shape)) - ) - recv_buffer = table.new_empty((sum(recv_counts), *trailing_shape)) - work = torch.distributed.all_to_all_single( # ty: ignore[possibly-missing-attribute] - recv_buffer, - send_buffer, - output_split_sizes=recv_counts, - input_split_sizes=send_counts, - group=group, - async_op=True, - ) - work.wait() - - recvs: list[tuple[GdnParentStateTransferPlan, Tensor]] = [] - offset = 0 - for peer_rank, count in enumerate(recv_counts): - peer_end = offset + count - for transfer in transfers: - send_rank = transfer.dest_rank if reverse else transfer.source_rank - recv_rank = transfer.source_rank if reverse else transfer.dest_rank - if send_rank == recv_rank: - continue - if rank != recv_rank or peer_rank != send_rank: - continue - rows = len(transfer.family_indices) - recvs.append((transfer, recv_buffer[offset : offset + rows])) - offset += rows - if offset != peer_end: - raise RuntimeError( - "parent-state exchange unpack mismatch: " - f"rank={rank} peer={peer_rank} consumed={offset} expected={peer_end}" - ) - return recvs - - -def _parent_state_index_tensor( - transfer: GdnParentStateTransferPlan, - *, - device: torch.device, -) -> Tensor: - if ( - transfer.family_indices_tensor is not None - and transfer.family_indices_tensor.device == device - ): - return transfer.family_indices_tensor - return torch.tensor(transfer.family_indices, device=device, dtype=torch.long) - - def run_gdn_bucket( bucket: GdnSegmentBucketPlan, projected_streams: tuple[Tensor, Tensor, Tensor], @@ -2597,14 +2253,17 @@ def run_gdn_bucket( conv_output_final_state = output_final_state chain_conv_final: Tensor | None = None + chain_gradient_dependency: Tensor | None = None if recurrent_cp: - conv_initial, chain_conv_final = _chain_conv_initial_and_final( - qkv, - bucket.cu_seqlens_cpu, - bucket.lengths_by_rank_cpu, - conv_initial, - group=group, - output_final_state=output_final_state, + conv_initial, chain_conv_final, chain_gradient_dependency = ( + _chain_conv_initial_and_final( + qkv, + bucket.cu_seqlens_cpu, + bucket.lengths_by_rank_cpu, + conv_initial, + group=group, + output_final_state=output_final_state, + ) ) conv_output_final_state = False @@ -2618,15 +2277,31 @@ def run_gdn_bucket( if recurrent_cp: conv_final = chain_conv_final - query, key, value, beta, recurrent_g = _prepare_packed_recurrent_inputs_fused( - qkv, - beta, - recurrent_g, - key_heads=_local_key_heads(gdn), - value_heads=_local_value_heads(gdn), - key_dim=int(gdn.key_head_dim), - value_dim=int(gdn.value_head_dim), - ) + dense_local_bucket = not recurrent_cp and _bucket_has_uniform_lengths(bucket) + if dense_local_bucket: + query, key, value, beta, recurrent_g = _prepare_dense_recurrent_inputs( + qkv.reshape(batch_size, int(bucket.length), int(qkv.shape[-1])), + beta.reshape(batch_size, int(bucket.length), int(beta.shape[-1])), + recurrent_g.reshape( + batch_size, + int(bucket.length), + int(recurrent_g.shape[-1]), + ), + key_heads=_local_key_heads(gdn), + value_heads=_local_value_heads(gdn), + key_dim=int(gdn.key_head_dim), + value_dim=int(gdn.value_head_dim), + ) + else: + query, key, value, beta, recurrent_g = _prepare_packed_recurrent_inputs_fused( + qkv, + beta, + recurrent_g, + key_heads=_local_key_heads(gdn), + value_heads=_local_value_heads(gdn), + key_dim=int(gdn.key_head_dim), + value_dim=int(gdn.value_head_dim), + ) if gdn.use_qk_l2norm: query = _l2norm(query.contiguous()) key = _l2norm(key.contiguous()) @@ -2657,8 +2332,27 @@ def run_gdn_bucket( initial_state=recurrent_initial, output_final_state=output_final_state, use_qk_l2norm_in_kernel=False, - cu_seqlens=bucket.cu_seqlens, - ) + cu_seqlens=None if dense_local_bucket else bucket.cu_seqlens, + ) + if dense_local_bucket: + recurrent_out = recurrent_out.reshape( + 1, + token_count, + int(recurrent_out.shape[-2]), + int(recurrent_out.shape[-1]), + ) + if chain_gradient_dependency is not None: + recurrent_out = _add_autograd_dependency( + recurrent_out, + chain_gradient_dependency, + ) + if conv_final is not None: + conv_final = _add_autograd_dependency(conv_final, chain_gradient_dependency) + if recurrent_final is not None: + recurrent_final = _add_autograd_dependency( + recurrent_final, + chain_gradient_dependency, + ) return recurrent_out, conv_final, recurrent_final @@ -2670,15 +2364,22 @@ def _chain_conv_initial_and_final( *, group: Any, output_final_state: bool, -) -> tuple[Tensor, Tensor | None]: +) -> tuple[Tensor, Tensor | None, Tensor]: if group is None: raise ValueError("CP chain conv state requires a process group") if not dist.is_available() or not dist.is_initialized(): # ty: ignore[possibly-missing-attribute] raise RuntimeError("torch.distributed must be initialized for CP chain conv") - parent_initial = _AllReduceGradient.apply(parent_initial, group) + parent_initial, gradient_dependency = _AllReduceGradient.apply( + parent_initial, + group, + ) tail_width = int(parent_initial.shape[-1]) if tail_width <= 0: - return parent_initial, parent_initial if output_final_state else None + return ( + parent_initial, + parent_initial if output_final_state else None, + gradient_dependency, + ) if lengths_by_rank_cpu is None: raise ValueError("CP chain conv requires static all-rank bucket lengths") if cu_seqlens_cpu.device.type != "cpu" or lengths_by_rank_cpu.device.type != "cpu": @@ -2705,7 +2406,7 @@ def _chain_conv_initial_and_final( if output_final_state else None ) - return conv_initial, conv_final + return conv_initial, conv_final, gradient_dependency def _local_packed_conv_tail( @@ -2733,20 +2434,219 @@ def _scan_conv_tail_batch( *, stop_rank: int, ) -> Tensor: - states = [] tail_width = int(parent_initial.shape[-1]) + if tail_width <= 0: + return parent_initial + source_ids_cpu, source_pos_cpu = _conv_tail_source_map( + lengths_by_rank_cpu, + stop_rank=int(stop_rank), + tail_width=tail_width, + ) + return _ConvTailScan.apply( + parent_initial, + tails_by_rank, + source_ids_cpu.to(device=parent_initial.device, non_blocking=True), + source_pos_cpu.to(device=parent_initial.device, non_blocking=True), + ) + + +def _conv_tail_source_map( + lengths_by_rank_cpu: Tensor, + *, + stop_rank: int, + tail_width: int, +) -> tuple[Tensor, Tensor]: + if lengths_by_rank_cpu.device.type != "cpu": + raise ValueError("conv tail source-map lengths must stay on CPU") + segments = int(lengths_by_rank_cpu.shape[1]) + source_ids = torch.empty((segments, tail_width), dtype=torch.int32) + source_pos = torch.empty((segments, tail_width), dtype=torch.int32) host_lengths = lengths_by_rank_cpu.tolist() - for segment in range(int(parent_initial.shape[0])): - state = parent_initial[segment] - for peer in range(int(stop_rank)): - valid = int(host_lengths[peer][segment]) - if valid <= 0: + for segment in range(segments): + total = tail_width + sum( + int(host_lengths[peer][segment]) for peer in range(int(stop_rank)) + ) + first = total - tail_width + for out_pos in range(tail_width): + absolute = first + out_pos + if absolute < tail_width: + source_ids[segment, out_pos] = 0 + source_pos[segment, out_pos] = absolute continue - state = torch.cat([state, tails_by_rank[peer, segment, :, :valid]], dim=-1)[ - :, -tail_width: - ] - states.append(state) - return torch.stack(states, dim=0) + remaining = absolute - tail_width + for peer in range(int(stop_rank)): + valid = int(host_lengths[peer][segment]) + if remaining < valid: + source_ids[segment, out_pos] = peer + 1 + source_pos[segment, out_pos] = remaining + break + remaining -= valid + else: + raise RuntimeError("conv tail source-map construction fell off history") + return source_ids, source_pos + + +@triton.jit(do_not_specialize=["segments"]) +def _conv_tail_scan_kernel( + parent_initial, + tails_by_rank, + source_ids, + source_pos, + output, + segments, + channels: tl.constexpr, + tail_width: tl.constexpr, + BLOCK_C: tl.constexpr, +) -> None: + segment = tl.program_id(0) + channel_block = tl.program_id(1) + offsets = channel_block * BLOCK_C + tl.arange(0, BLOCK_C) + channel_mask = offsets < channels + for tail_pos in tl.static_range(0, tail_width): + source_id = tl.load(source_ids + segment * tail_width + tail_pos) + pos = tl.load(source_pos + segment * tail_width + tail_pos) + parent_values = tl.load( + parent_initial + (segment * channels + offsets) * tail_width + pos, + mask=channel_mask & (source_id == 0), + other=0.0, + ) + tail_values = tl.load( + tails_by_rank + + (((source_id - 1) * segments + segment) * channels + offsets) * tail_width + + pos, + mask=channel_mask & (source_id != 0), + other=0.0, + ) + tl.store( + output + (segment * channels + offsets) * tail_width + tail_pos, + parent_values + tail_values, + mask=channel_mask, + ) + + +@triton.jit(do_not_specialize=["segments"]) +def _conv_tail_scan_backward_kernel( + grad_output, + source_ids, + source_pos, + grad_parent_initial, + grad_tails_by_rank, + segments, + channels: tl.constexpr, + tail_width: tl.constexpr, + BLOCK_C: tl.constexpr, +) -> None: + segment = tl.program_id(0) + channel_block = tl.program_id(1) + offsets = channel_block * BLOCK_C + tl.arange(0, BLOCK_C) + channel_mask = offsets < channels + for tail_pos in tl.static_range(0, tail_width): + source_id = tl.load(source_ids + segment * tail_width + tail_pos) + pos = tl.load(source_pos + segment * tail_width + tail_pos) + values = tl.load( + grad_output + (segment * channels + offsets) * tail_width + tail_pos, + mask=channel_mask, + other=0.0, + ) + tl.store( + grad_parent_initial + (segment * channels + offsets) * tail_width + pos, + values, + mask=channel_mask & (source_id == 0), + ) + tl.store( + grad_tails_by_rank + + (((source_id - 1) * segments + segment) * channels + offsets) * tail_width + + pos, + values, + mask=channel_mask & (source_id != 0), + ) + + +class _ConvTailScan(torch.autograd.Function): + @staticmethod + def forward( + ctx: Any, + parent_initial: Tensor, + tails_by_rank: Tensor, + source_ids: Tensor, + source_pos: Tensor, + ) -> Tensor: + if parent_initial.ndim != 3: + raise ValueError( + f"parent_initial must be [segments, channels, tail], got {tuple(parent_initial.shape)}" + ) + if tails_by_rank.ndim != 4: + raise ValueError( + f"tails_by_rank must be [cp, segments, channels, tail], got {tuple(tails_by_rank.shape)}" + ) + segments, channels, tail_width = ( + int(parent_initial.shape[0]), + int(parent_initial.shape[1]), + int(parent_initial.shape[2]), + ) + if tuple(tails_by_rank.shape[1:]) != tuple(parent_initial.shape): + raise ValueError( + "tails_by_rank trailing dimensions must match parent_initial, got " + f"{tuple(tails_by_rank.shape)} and {tuple(parent_initial.shape)}" + ) + if tuple(source_ids.shape) != (segments, tail_width) or tuple( + source_pos.shape + ) != (segments, tail_width): + raise ValueError( + "conv tail source maps must be [segments, tail_width], got " + f"{tuple(source_ids.shape)} and {tuple(source_pos.shape)}" + ) + output = torch.empty_like(parent_initial) + block_c = 256 + _conv_tail_scan_kernel[(segments, triton.cdiv(channels, block_c))]( + parent_initial, + tails_by_rank, + source_ids.contiguous(), + source_pos.contiguous(), + output, + segments, + channels, + tail_width, + BLOCK_C=block_c, + num_warps=8, + ) + ctx.save_for_backward(source_ids, source_pos) + ctx.parent_shape = tuple(parent_initial.shape) + ctx.tails_shape = tuple(tails_by_rank.shape) + return output + + @staticmethod + def backward(ctx: Any, *grad_outputs: Tensor | None) -> tuple[Any, ...]: + grad_output = grad_outputs[0] + if grad_output is None: + raise RuntimeError("conv tail scan backward expected output gradient") + source_ids, source_pos = ctx.saved_tensors + grad_output = grad_output.contiguous() + grad_parent = torch.zeros( + ctx.parent_shape, device=grad_output.device, dtype=grad_output.dtype + ) + grad_tails = torch.zeros( + ctx.tails_shape, device=grad_output.device, dtype=grad_output.dtype + ) + segments, channels, tail_width = ( + int(ctx.parent_shape[0]), + int(ctx.parent_shape[1]), + int(ctx.parent_shape[2]), + ) + block_c = 256 + _conv_tail_scan_backward_kernel[(segments, triton.cdiv(channels, block_c))]( + grad_output, + source_ids, + source_pos, + grad_parent, + grad_tails, + segments, + channels, + tail_width, + BLOCK_C=block_c, + num_warps=8, + ) + return grad_parent, grad_tails, None, None class _AllGatherReplicated(torch.autograd.Function): @@ -2782,14 +2682,20 @@ def backward(ctx: Any, *grad_outputs: Tensor) -> tuple[Tensor, None]: class _AllReduceGradient(torch.autograd.Function): @staticmethod - def forward(ctx: Any, tensor: Tensor, group: Any) -> Tensor: + def forward(ctx: Any, tensor: Tensor, group: Any) -> tuple[Tensor, Tensor]: ctx.group = group - return tensor + ctx.save_for_backward(tensor) + return tensor, tensor.new_zeros(()) @staticmethod - def backward(ctx: Any, *grad_outputs: Tensor) -> tuple[Tensor, None]: - (grad_output,) = grad_outputs - grad_input = grad_output.contiguous() + def backward(ctx: Any, *grad_outputs: Tensor | None) -> tuple[Tensor, None]: + grad_output, _grad_dependency = grad_outputs + (reference,) = ctx.saved_tensors + grad_input = ( + reference.new_zeros(reference.shape) + if grad_output is None + else grad_output.contiguous() + ) dist.all_reduce( # ty: ignore[possibly-missing-attribute] grad_input, op=dist.ReduceOp.SUM, # ty: ignore[possibly-missing-attribute] @@ -2932,7 +2838,5 @@ def _chunk_gated_delta_rule(*args: Any, **kwargs: Any) -> tuple[Tensor, Tensor | try: from fla.ops.gated_delta_rule import chunk_gated_delta_rule except ImportError as exc: - raise ImportError( - "FLA is required for ART shared-prefix GDN execution." - ) from exc + raise ImportError("FLA is required for ART prefix-tree GDN execution.") from exc return chunk_gated_delta_rule(*args, **kwargs) diff --git a/src/art/megatron/gdn/segment_layout.py b/src/art/megatron/gdn/segment_layout.py index 9bbb3517a..0b5feb760 100644 --- a/src/art/megatron/gdn/segment_layout.py +++ b/src/art/megatron/gdn/segment_layout.py @@ -8,28 +8,13 @@ import triton.language as tl -@triton.jit(do_not_specialize=["segment_count"]) -def _segment_from_cu(cu_seqlens, n, segment_count): - lo = n * 0 - hi = lo + segment_count - for _ in tl.static_range(0, 16): - mid = (lo + hi) // 2 - start = tl.load(cu_seqlens + mid) - take_upper = start <= n - lo = tl.where(take_upper, mid, lo) - hi = tl.where(take_upper, hi, mid) - return lo, n - tl.load(cu_seqlens + lo) - - -@triton.jit(do_not_specialize=["token_count", "segment_count", "sequence_length"]) +@triton.jit(do_not_specialize=["token_count", "sequence_length"]) def _gather_compact_qkv_kernel( qkv_flat, row_indices, position_indices, - cu_seqlens, out, token_count, - segment_count, sequence_length, channels: tl.constexpr, BLOCK_N: tl.constexpr, @@ -38,10 +23,8 @@ def _gather_compact_qkv_kernel( n = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) d = tl.program_id(1) * BLOCK_D + tl.arange(0, BLOCK_D) token_mask = n < token_count - segment, offset = _segment_from_cu(cu_seqlens, n, segment_count) - p = offset * segment_count + segment - row = tl.load(row_indices + p, mask=token_mask, other=0) - pos = tl.load(position_indices + p, mask=token_mask, other=0) + row = tl.load(row_indices + n, mask=token_mask, other=0) + pos = tl.load(position_indices + n, mask=token_mask, other=0) src = row * sequence_length + pos n64 = n.to(tl.int64) d64 = d.to(tl.int64) @@ -55,15 +38,13 @@ def _gather_compact_qkv_kernel( tl.store(out + n64[:, None] * channels + d64[None, :], values, mask=mask) -@triton.jit(do_not_specialize=["token_count", "segment_count", "sequence_length"]) +@triton.jit(do_not_specialize=["token_count", "sequence_length"]) def _gather_compact_aux_kernel( x_flat, row_indices, position_indices, - cu_seqlens, out, token_count, - segment_count, sequence_length, width: tl.constexpr, BLOCK_N: tl.constexpr, @@ -72,10 +53,8 @@ def _gather_compact_aux_kernel( n = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) d = tl.program_id(1) * BLOCK_D + tl.arange(0, BLOCK_D) token_mask = n < token_count - segment, offset = _segment_from_cu(cu_seqlens, n, segment_count) - p = offset * segment_count + segment - row = tl.load(row_indices + p, mask=token_mask, other=0) - pos = tl.load(position_indices + p, mask=token_mask, other=0) + row = tl.load(row_indices + n, mask=token_mask, other=0) + pos = tl.load(position_indices + n, mask=token_mask, other=0) src = row * sequence_length + pos n64 = n.to(tl.int64) d64 = d.to(tl.int64) @@ -89,15 +68,13 @@ def _gather_compact_aux_kernel( tl.store(out + n64[:, None] * width + d64[None, :], values, mask=mask) -@triton.jit(do_not_specialize=["token_count", "segment_count", "sequence_length"]) +@triton.jit(do_not_specialize=["token_count", "sequence_length"]) def _scatter_compact_qkv_grad_kernel( grad_out, row_indices, position_indices, - cu_seqlens, grad_flat, token_count, - segment_count, sequence_length, channels: tl.constexpr, BLOCK_N: tl.constexpr, @@ -106,10 +83,8 @@ def _scatter_compact_qkv_grad_kernel( n = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) d = tl.program_id(1) * BLOCK_D + tl.arange(0, BLOCK_D) token_mask = n < token_count - segment, offset = _segment_from_cu(cu_seqlens, n, segment_count) - p = offset * segment_count + segment - row = tl.load(row_indices + p, mask=token_mask, other=0) - pos = tl.load(position_indices + p, mask=token_mask, other=0) + row = tl.load(row_indices + n, mask=token_mask, other=0) + pos = tl.load(position_indices + n, mask=token_mask, other=0) dst = row * sequence_length + pos n64 = n.to(tl.int64) d64 = d.to(tl.int64) @@ -128,15 +103,13 @@ def _scatter_compact_qkv_grad_kernel( ) -@triton.jit(do_not_specialize=["token_count", "segment_count", "sequence_length"]) +@triton.jit(do_not_specialize=["token_count", "sequence_length"]) def _scatter_compact_aux_grad_kernel( grad_out, row_indices, position_indices, - cu_seqlens, grad_flat, token_count, - segment_count, sequence_length, width: tl.constexpr, BLOCK_N: tl.constexpr, @@ -145,10 +118,8 @@ def _scatter_compact_aux_grad_kernel( n = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) d = tl.program_id(1) * BLOCK_D + tl.arange(0, BLOCK_D) token_mask = n < token_count - segment, offset = _segment_from_cu(cu_seqlens, n, segment_count) - p = offset * segment_count + segment - row = tl.load(row_indices + p, mask=token_mask, other=0) - pos = tl.load(position_indices + p, mask=token_mask, other=0) + row = tl.load(row_indices + n, mask=token_mask, other=0) + pos = tl.load(position_indices + n, mask=token_mask, other=0) dst = row * sequence_length + pos n64 = n.to(tl.int64) d64 = d.to(tl.int64) @@ -319,18 +290,14 @@ def _prepare_packed_qkv_backward_kernel( tl.store(grad_qkv + n64[:, None] * channels + c64[None, :], values, mask=mask) -@triton.jit( - do_not_specialize=["token_count", "segment_count", "output_sequence_length"] -) +@triton.jit(do_not_specialize=["token_count", "output_sequence_length"]) def _scatter_bucket_output_compact_forward_kernel( output, bucket_output, row_indices, position_indices, output_mask, - cu_seqlens, token_count, - segment_count, output_sequence_length, heads: tl.constexpr, dim: tl.constexpr, @@ -339,12 +306,10 @@ def _scatter_bucket_output_compact_forward_kernel( ): n = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) hd = tl.program_id(1) * BLOCK_D + tl.arange(0, BLOCK_D) - segment, offset = _segment_from_cu(cu_seqlens, n, segment_count) - p = offset * segment_count + segment token_mask = n < token_count - write = tl.load(output_mask + p, mask=token_mask, other=0).to(tl.int1) - row = tl.load(row_indices + p, mask=token_mask, other=0) - pos = tl.load(position_indices + p, mask=token_mask, other=0) + write = tl.load(output_mask + n, mask=token_mask, other=0).to(tl.int1) + row = tl.load(row_indices + n, mask=token_mask, other=0) + pos = tl.load(position_indices + n, mask=token_mask, other=0) h = hd // dim d = hd - h * dim n64 = n.to(tl.int64) @@ -368,9 +333,7 @@ def _scatter_bucket_output_compact_forward_kernel( ) -@triton.jit( - do_not_specialize=["token_count", "segment_count", "output_sequence_length"] -) +@triton.jit(do_not_specialize=["token_count", "output_sequence_length"]) def _scatter_bucket_output_compact_backward_kernel( grad_output, grad_base, @@ -378,9 +341,7 @@ def _scatter_bucket_output_compact_backward_kernel( row_indices, position_indices, output_mask, - cu_seqlens, token_count, - segment_count, output_sequence_length, heads: tl.constexpr, dim: tl.constexpr, @@ -389,12 +350,10 @@ def _scatter_bucket_output_compact_backward_kernel( ): n = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N) hd = tl.program_id(1) * BLOCK_D + tl.arange(0, BLOCK_D) - segment, offset = _segment_from_cu(cu_seqlens, n, segment_count) - p = offset * segment_count + segment token_mask = n < token_count - write = tl.load(output_mask + p, mask=token_mask, other=0).to(tl.int1) - row = tl.load(row_indices + p, mask=token_mask, other=0) - pos = tl.load(position_indices + p, mask=token_mask, other=0) + write = tl.load(output_mask + n, mask=token_mask, other=0).to(tl.int1) + row = tl.load(row_indices + n, mask=token_mask, other=0) + pos = tl.load(position_indices + n, mask=token_mask, other=0) h = hd // dim d = hd - h * dim n64 = n.to(tl.int64) @@ -428,9 +387,7 @@ def forward( recurrent_g_flat: Tensor, row_indices: Tensor, position_indices: Tensor, - cu_seqlens: Tensor, token_count: int, - segment_count: int, sequence_length: int, ) -> tuple[Tensor, Tensor, Tensor]: _validate_cuda("qkv_flat", qkv_flat) @@ -439,10 +396,18 @@ def forward( recurrent_g_flat = recurrent_g_flat.contiguous() row_indices = row_indices.contiguous() position_indices = position_indices.contiguous() - cu_seqlens = cu_seqlens.contiguous() token_count = int(token_count) - segment_count = int(segment_count) sequence_length = int(sequence_length) + if tuple(row_indices.shape) != (token_count,): + raise ValueError( + "packed row_indices must be [tokens], got " + f"{tuple(row_indices.shape)} for {token_count}" + ) + if tuple(position_indices.shape) != (token_count,): + raise ValueError( + "packed position_indices must be [tokens], got " + f"{tuple(position_indices.shape)} for {token_count}" + ) qkv_channels = int(qkv_flat.shape[-1]) aux_width = int(beta_flat.shape[-1]) qkv = torch.empty( @@ -465,10 +430,8 @@ def forward( qkv_flat, row_indices, position_indices, - cu_seqlens, qkv, token_count, - segment_count, sequence_length, qkv_channels, BLOCK_N=block_n, @@ -483,10 +446,8 @@ def forward( beta_flat, row_indices, position_indices, - cu_seqlens, beta, token_count, - segment_count, sequence_length, aux_width, BLOCK_N=block_n, @@ -497,19 +458,16 @@ def forward( recurrent_g_flat, row_indices, position_indices, - cu_seqlens, recurrent_g, token_count, - segment_count, sequence_length, aux_width, BLOCK_N=block_n, BLOCK_D=block_aux, num_warps=4, ) - ctx.save_for_backward(row_indices, position_indices, cu_seqlens) + ctx.save_for_backward(row_indices, position_indices) ctx.token_count = token_count - ctx.segment_count = segment_count ctx.sequence_length = sequence_length ctx.qkv_flat_count = int(qkv_flat.shape[0]) ctx.beta_flat_count = int(beta_flat.shape[0]) @@ -533,7 +491,7 @@ def backward( None, ]: grad_qkv_bucket, grad_beta_bucket, grad_g_bucket = grad_outputs - row_indices, position_indices, cu_seqlens = ctx.saved_tensors + row_indices, position_indices = ctx.saved_tensors block_n, block_qkv, block_aux = 32, 64, 32 grad_qkv = None if ctx.needs_input_grad[0] and grad_qkv_bucket is not None: @@ -548,10 +506,8 @@ def backward( grad_qkv_bucket, row_indices, position_indices, - cu_seqlens, grad_qkv, ctx.token_count, - ctx.segment_count, ctx.sequence_length, ctx.qkv_channels, BLOCK_N=block_n, @@ -571,10 +527,8 @@ def backward( grad_beta_bucket, row_indices, position_indices, - cu_seqlens, grad_beta, ctx.token_count, - ctx.segment_count, ctx.sequence_length, ctx.aux_width, BLOCK_N=block_n, @@ -594,10 +548,8 @@ def backward( grad_g_bucket, row_indices, position_indices, - cu_seqlens, grad_g, ctx.token_count, - ctx.segment_count, ctx.sequence_length, ctx.aux_width, BLOCK_N=block_n, @@ -782,7 +734,6 @@ def forward( row_indices: Tensor, position_indices: Tensor, output_mask: Tensor, - cu_seqlens: Tensor, ) -> Tensor: _validate_cuda("output", output) output = output.contiguous() @@ -790,7 +741,6 @@ def forward( row_indices = row_indices.contiguous() position_indices = position_indices.contiguous() output_mask = output_mask.contiguous() - cu_seqlens = cu_seqlens.contiguous() if bucket_output.ndim != 4 or int(bucket_output.shape[0]) != 1: raise ValueError( "bucket_output must have shape [1, tokens, heads, dim], got " @@ -799,16 +749,20 @@ def forward( output_batch, output_sequence_length, heads, dim = output.shape del output_batch token_count = int(bucket_output.shape[1]) - segment_count = int(cu_seqlens.numel()) - 1 - if tuple(row_indices.shape) != tuple(position_indices.shape): + if tuple(row_indices.shape) != (token_count,): + raise ValueError( + "packed row_indices must be [tokens], got " + f"{tuple(row_indices.shape)} for {token_count}" + ) + if tuple(position_indices.shape) != (token_count,): raise ValueError( - "row_indices and position_indices must have the same shape, got " - f"{tuple(row_indices.shape)} and {tuple(position_indices.shape)}" + "packed position_indices must be [tokens], got " + f"{tuple(position_indices.shape)} for {token_count}" ) - if tuple(output_mask.shape) != tuple(row_indices.shape): + if tuple(output_mask.shape) != (token_count,): raise ValueError( - "output_mask must match row_indices shape, got " - f"{tuple(output_mask.shape)} and {tuple(row_indices.shape)}" + "packed output_mask must be [tokens], got " + f"{tuple(output_mask.shape)} for {token_count}" ) out = output.clone() block_n, block_d = 16, 64 @@ -820,9 +774,7 @@ def forward( row_indices, position_indices, output_mask, - cu_seqlens, token_count, - segment_count, output_sequence_length, heads, dim, @@ -830,11 +782,10 @@ def forward( BLOCK_D=block_d, num_warps=4, ) - ctx.save_for_backward(row_indices, position_indices, output_mask, cu_seqlens) + ctx.save_for_backward(row_indices, position_indices, output_mask) ctx.output_shape = tuple(output.shape) ctx.bucket_output_shape = tuple(bucket_output.shape) ctx.token_count = token_count - ctx.segment_count = segment_count return out @staticmethod @@ -844,7 +795,7 @@ def backward( if len(grad_outputs) != 1 or grad_outputs[0] is None: raise RuntimeError("expected compact scatter output gradient") grad_out = grad_outputs[0] - row_indices, position_indices, output_mask, cu_seqlens = ctx.saved_tensors + row_indices, position_indices, output_mask = ctx.saved_tensors _, output_sequence_length, heads, dim = ctx.output_shape grad_out = grad_out.contiguous() grad_base = grad_out.clone() @@ -862,9 +813,7 @@ def backward( row_indices, position_indices, output_mask, - cu_seqlens, ctx.token_count, - ctx.segment_count, output_sequence_length, heads, dim, @@ -881,10 +830,8 @@ def gather_bucket_streams_compact( recurrent_g_flat: Tensor, row_indices: Tensor, position_indices: Tensor, - cu_seqlens: Tensor, *, token_count: int, - segment_count: int, sequence_length: int, ) -> tuple[Tensor, Tensor, Tensor]: return _CompactBucketStreamGather.apply( @@ -893,9 +840,7 @@ def gather_bucket_streams_compact( recurrent_g_flat, row_indices, position_indices, - cu_seqlens, token_count, - segment_count, sequence_length, ) @@ -906,7 +851,6 @@ def scatter_bucket_output_compact( row_indices: Tensor, position_indices: Tensor, output_mask: Tensor, - cu_seqlens: Tensor, ) -> Tensor: return _CompactScatterBucketOutput.apply( output, @@ -914,7 +858,6 @@ def scatter_bucket_output_compact( row_indices, position_indices, output_mask, - cu_seqlens, ) diff --git a/src/art/megatron/lora.py b/src/art/megatron/lora.py index 73aa66f23..d0951fda9 100644 --- a/src/art/megatron/lora.py +++ b/src/art/megatron/lora.py @@ -1,9 +1,14 @@ -from collections.abc import Sequence +from collections.abc import Iterator, Sequence +from contextlib import contextmanager +import contextvars +from dataclasses import dataclass, replace +import functools +import importlib import json import math import os import re -from typing import Any, Literal, NamedTuple, cast +from typing import Any, Callable, Literal, NamedTuple, TypeVar, cast from megatron.bridge.models.gpt_provider import GPTModelProvider from megatron.core import parallel_state as ps @@ -22,9 +27,7 @@ ) from megatron.core.transformer.attention import SelfAttention from megatron.core.transformer.moe.experts import TEGroupedMLP -from megatron.core.transformer.moe.shared_experts import SharedExpertMLP from megatron.core.transformer.transformer_layer import TransformerLayer -from pydantic import BaseModel, ConfigDict import torch from .kernels.cute_grouped_lora_quack import ( @@ -42,6 +45,8 @@ ShardDomain = Literal["tp", "expert_tp"] GradSyncDomain = Literal["tp_default", "expert_tp"] GradSyncOp = Literal["none", "sum", "avg"] +LoraSlotKind = Literal["checkpoint", "lora"] +_F = TypeVar("_F", bound=Callable[..., Any]) TP_DEFAULT_GRAD_SYNC_DOMAIN: GradSyncDomain = "tp_default" EXPERT_TP_GRAD_SYNC_DOMAIN: GradSyncDomain = "expert_tp" @@ -50,11 +55,114 @@ GRAD_SYNC_OP_AVG: GradSyncOp = "avg" -class LoRAParallelSpec(BaseModel): - # This spec only describes TP / expert-TP behavior. - # DP/CP vs expert-DP behavior is selected separately via `allreduce`. - model_config = ConfigDict(frozen=True) +@dataclass(frozen=True) +class LoRASlotRef: + kind: LoraSlotKind + name: str | None + +_CURRENT_LORA_SLOT: contextvars.ContextVar[LoRASlotRef | None] = contextvars.ContextVar( + "art_megatron_current_lora_slot", default=None +) + + +@contextmanager +def use_lora_slot(ref: LoRASlotRef | None) -> Iterator[None]: + token = _CURRENT_LORA_SLOT.set(ref) + try: + yield + finally: + _CURRENT_LORA_SLOT.reset(token) + + +def _with_captured_lora_slot(function: _F) -> _F: + context = _CURRENT_LORA_SLOT.get() + + @functools.wraps(function) + def wrapped(*args: Any, **kwargs: Any) -> Any: + token = _CURRENT_LORA_SLOT.set(context) + try: + return function(*args, **kwargs) + finally: + _CURRENT_LORA_SLOT.reset(token) + + return cast(_F, wrapped) + + +def _patch_function_once(module: Any, name: str, wrapper: Callable[[_F], _F]) -> None: + original = getattr(module, name, None) + if original is None or getattr(original, "_art_lora_slot_context_patch", False): + return + patched = wrapper(original) + setattr(patched, "_art_lora_slot_context_patch", True) + setattr(module, name, patched) + + +def install_lora_checkpoint_context_hooks() -> None: + """Preserve the selected dynamic LoRA slot across activation recompute.""" + + def wrap_checkpoint(original: _F, function_index: int) -> _F: + @functools.wraps(original) + def checkpoint(*args: Any, **kwargs: Any) -> Any: + if len(args) > function_index: + args = ( + *args[:function_index], + _with_captured_lora_slot(args[function_index]), + *args[function_index + 1 :], + ) + elif "function" in kwargs: + kwargs = { + **kwargs, + "function": _with_captured_lora_slot(kwargs["function"]), + } + elif "forward_func" in kwargs: + kwargs = { + **kwargs, + "forward_func": _with_captured_lora_slot(kwargs["forward_func"]), + } + else: + raise TypeError("checkpoint wrapper could not find callable argument") + return original(*args, **kwargs) + + return cast(_F, checkpoint) + + def patch(target: str, name: str, function_index: int) -> None: + try: + module_name, _, attr_path = target.partition(":") + target_obj = importlib.import_module(module_name) + for attr in attr_path.split(".") if attr_path else (): + target_obj = getattr(target_obj, attr, None) + if target_obj is None: + return + _patch_function_once( + target_obj, + name, + lambda original: wrap_checkpoint(original, function_index), + ) + except Exception: + pass + + for target, name, function_index in ( + ("torch.utils.checkpoint", "checkpoint", 0), + ("megatron.core.tensor_parallel", "checkpoint", 0), + ("megatron.core.tensor_parallel.random", "checkpoint", 0), + ( + "megatron.core.tensor_parallel.random:CheckpointWithoutOutput", + "checkpoint", + 1, + ), + ("megatron.core.transformer.transformer_block", "te_checkpoint", 0), + ("transformer_engine.pytorch.distributed", "checkpoint", 0), + ): + patch(target, name, function_index) + + +install_lora_checkpoint_context_hooks() + + +@dataclass(frozen=True) +class LoRAParallelSpec: + # This only describes TP / expert-TP; DP/CP vs expert-DP is selected by `allreduce`. shard_domain: ShardDomain = "tp" sharded: bool = False shard_dim: int | None = None @@ -72,10 +180,7 @@ class LoraShardMeta(NamedTuple): @property def numel(self) -> int: - total = 1 - for dim in self.shape: - total *= dim - return total + return math.prod(self.shape) class _LoraPublishTemplate(NamedTuple): @@ -123,14 +228,6 @@ def _get_shard_rank(domain: ShardDomain) -> int: return group.rank() -def _get_shard_group(domain: ShardDomain) -> Any | None: - if not _distributed_initialized(): - return None - if domain == "tp": - return ps.get_tensor_model_parallel_group() - return ps.get_expert_tensor_parallel_group(check_initialized=False) - - def _dtype_name(dtype: torch.dtype) -> str: return str(dtype).removeprefix("torch.") @@ -242,13 +339,8 @@ def _set_lora_parallel_metadata( setattr(param, "lora_tp_shard_dim", parallel_spec.shard_dim) setattr(param, "grad_sync_domain", parallel_spec.grad_sync_domain) setattr(param, "grad_sync_op", parallel_spec.grad_sync_op) - # Megatron DDP routing flag: - # - allreduce=True: sync with regular DP/CP replicas. - # - allreduce=False: sync with expert-DP replicas. - # TP / expert-TP replica handling is controlled by grad_sync_* metadata. setattr(param, "allreduce", allreduce) - # Megatron's native TP finalize path consumes this attr. setattr( param, "average_gradients_across_tp_domain", @@ -259,16 +351,12 @@ def _set_lora_parallel_metadata( ), ) - # Megatron optimizer and checkpoint logic rely on tensor model-parallel metadata - # to distinguish true shards from TP-duplicate params. if parallel_spec.sharded: shard_dim = parallel_spec.shard_dim if shard_dim is None: raise ValueError("LoRAParallelSpec.shard_dim must be set when sharded=True") setattr(param, "tensor_model_parallel", True) setattr(param, "partition_dim", _normalize_axis(shard_dim, param.ndim)) - # stride > 1 means the dim is split into blocks and each tp rank holds a shard of the block - # this might happen for fused e.g. gate_(up|proj), but loras are individual per module setattr(param, "partition_stride", 1) else: setattr(param, "tensor_model_parallel", False) @@ -307,6 +395,59 @@ def _exported_shard_dim(param: torch.nn.Parameter) -> int: return 1 - axis +def _copy_lora_param_metadata( + source: torch.nn.Parameter, + target: torch.nn.Parameter, +) -> None: + for name in ( + "lora_shard_domain", + "lora_tp_sharded", + "lora_tp_replicated", + "lora_tp_shard_dim", + "grad_sync_domain", + "grad_sync_op", + "allreduce", + "average_gradients_across_tp_domain", + "tensor_model_parallel", + "partition_dim", + "partition_stride", + "lora_tp_shard_strategy", + "lora_tp_component_sizes", + ): + if hasattr(source, name): + setattr(target, name, getattr(source, name)) + setattr(target, "_art_dynamic_lora_slot", True) + + +class LoRASlot(torch.nn.Module): + def __init__( + self, + *, + ref: LoRASlotRef, + a_t: torch.Tensor, + b_t: torch.Tensor, + alpha: float, + a_template: torch.nn.Parameter, + b_template: torch.nn.Parameter, + requires_grad: bool, + ) -> None: + super().__init__() + self.ref = ref + self.alpha = float(alpha) + self.A_T = torch.nn.Parameter(a_t.detach().clone(), requires_grad=requires_grad) + self.B_T = torch.nn.Parameter(b_t.detach().clone(), requires_grad=requires_grad) + _copy_lora_param_metadata(a_template, self.A_T) + _copy_lora_param_metadata(b_template, self.B_T) + + @property + def rank(self) -> int: + return int(self.A_T.shape[-1]) + + @property + def scale(self) -> float: + return self.alpha / self.rank + + class LoRA(torch.nn.Module): def __init__( self, @@ -327,7 +468,12 @@ def __init__( "adapter_model_prefix must contain the '{expert}' format placeholder if num_local_experts > 1" ) self.adapter_model_prefix = adapter_model_prefix + self.alpha = float(alpha) + self.in_features = int(in_features) + self.out_features = int(out_features) self.scale = alpha / rank + self._slot_modules = torch.nn.ModuleDict() + self._slot_keys: dict[LoRASlotRef, str] = {} self.A_T = torch.nn.Parameter( torch.zeros( num_local_experts, in_features, rank, dtype=dtype, device=device @@ -362,7 +508,11 @@ def _broadcast_if_replicated(self, param: torch.nn.Parameter) -> None: world_size = _get_shard_world_size(domain) if world_size <= 1: return - group = _get_shard_group(domain) + group = ( + ps.get_tensor_model_parallel_group() + if domain == "tp" + else ps.get_expert_tensor_parallel_group(check_initialized=False) + ) if group is None: raise RuntimeError( f"{self.adapter_model_prefix}: missing process group for replicated parameter domain={domain}" @@ -395,43 +545,106 @@ def _expected_weight_keys(self, suffix: str) -> list[str]: ] return [f"{self.adapter_model_prefix}.{suffix}.weight"] + def load_lora_slot( + self, + ref: LoRASlotRef, + adapter_model: dict[str, torch.Tensor], + *, + alpha: float = LORA_ALPHA, + requires_grad: bool, + ) -> bool: + if ref.name is None: + raise ValueError("base-model slot refs do not own LoRA tensors") + weights = self._adapter_weights(adapter_model, require=False) + if weights is None: + return False + a_t = self._localized_weight(weights[0], into=self.A_T) + b_t = self._localized_weight(weights[1], into=self.B_T) + slot_key = self._slot_keys.get(ref) + if slot_key is None: + slot_key = f"slot_{len(self._slot_keys)}" + self._slot_keys[ref] = slot_key + elif self._has_live_slot_grads(ref): + raise RuntimeError( + f"Cannot overwrite live LoRA slot {ref.kind}:{ref.name} for " + f"{self.adapter_model_prefix}; clear grads/backward graph first." + ) + self._slot_modules[slot_key] = LoRASlot( + ref=ref, + a_t=a_t, + b_t=b_t, + alpha=alpha, + a_template=self.A_T, + b_template=self.B_T, + requires_grad=requires_grad, + ) + return True + + def lora_slot_params(self, ref: LoRASlotRef) -> list[torch.nn.Parameter]: + slot = self._slot(ref) + if slot is None: + return [] + return [slot.A_T, slot.B_T] + + def _slot(self, ref: LoRASlotRef) -> LoRASlot | None: + key = self._slot_keys.get(ref) + if key is None: + return None + return cast(LoRASlot, self._slot_modules[key]) + + def _has_live_slot_grads(self, ref: LoRASlotRef) -> bool: + slot = self._slot(ref) + return slot is not None and any( + param.grad is not None for param in (slot.A_T, slot.B_T) + ) + def load_lora(self, adapter_model: dict[str, torch.Tensor]) -> None: - missing_keys = [ + weights = self._adapter_weights(adapter_model, require=False) + if weights is None: + self.reset_lora_parameters() + return + self._load_weight(weights[0], into=self.A_T) + self._load_weight(weights[1], into=self.B_T) + + def _adapter_weights( + self, + adapter_model: dict[str, torch.Tensor], + *, + require: bool, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + all_keys = [ key for suffix in ("lora_A", "lora_B") for key in self._expected_weight_keys(suffix) - if key not in adapter_model ] - if missing_keys: + missing = [key for key in all_keys if key not in adapter_model] + if len(missing) == len(all_keys) and not require: + return None + if missing: + state = "Missing" if require else "Incomplete" raise KeyError( - f"Missing LoRA adapter keys for {self.adapter_model_prefix}: {sorted(missing_keys)}" + f"{state} LoRA adapter keys for {self.adapter_model_prefix}: " + f"{sorted(missing)}" ) - self.load_weights( - adapter_model, - suffix="lora_A", - into=self.A_T, - ) - self.load_weights( - adapter_model, - suffix="lora_B", - into=self.B_T, + return ( + self._adapter_weight(adapter_model, suffix="lora_A"), + self._adapter_weight(adapter_model, suffix="lora_B"), ) - def load_weights( + def _adapter_weight( self, adapter_model: dict[str, torch.Tensor], *, suffix: str, - into: torch.nn.Parameter, - ) -> None: + ) -> torch.Tensor: keys = self._expected_weight_keys(suffix) if self.num_local_experts > 1: - weight = torch.stack([adapter_model[key].T for key in keys]) - else: - weight = adapter_model[keys[0]].T - self.load_weight(weight, into=into) + return torch.stack([adapter_model[key].T for key in keys]) + return adapter_model[keys[0]].T - def load_weight(self, weight: torch.Tensor, *, into: torch.nn.Parameter) -> None: + def _localized_weight( + self, weight: torch.Tensor, *, into: torch.nn.Parameter + ) -> torch.Tensor: domain = into.lora_shard_domain # ty: ignore[unresolved-attribute] if into.lora_tp_sharded: # ty: ignore[unresolved-attribute] axis = into.lora_tp_shard_dim # ty: ignore[unresolved-attribute] @@ -470,11 +683,10 @@ def load_weight(self, weight: torch.Tensor, *, into: torch.nn.Parameter) -> None raise ValueError( f"{self.adapter_model_prefix}: unsupported shard strategy={strategy}" ) - elif tuple(weight.shape) != tuple(into.shape): - raise ValueError( - f"{self.adapter_model_prefix}: unsharded load shape mismatch, got {tuple(weight.shape)} " - f"expected {tuple(into.shape)}" - ) + return weight.contiguous() + + def _load_weight(self, weight: torch.Tensor, *, into: torch.nn.Parameter) -> None: + weight = self._localized_weight(weight, into=into) if tuple(weight.shape) != tuple(into.shape): raise ValueError( f"{self.adapter_model_prefix}: sharded load shape mismatch, got {tuple(weight.shape)} " @@ -575,9 +787,26 @@ def sharded_lora_grad_dict(self) -> dict[str, torch.Tensor]: grads[key] = local_grad.T return grads + def active_lora_tensors( + self, + ) -> tuple[torch.Tensor, torch.Tensor, float] | None: + ref = _CURRENT_LORA_SLOT.get() + if ref is None: + return self.A_T, self.B_T, self.scale + if ref.name is None: + return None + slot = self._slot(ref) + if slot is None: + return None + return slot.A_T, slot.B_T, slot.scale + def forward( self, x: torch.Tensor, tokens_per_expert: list[int] | torch.Tensor | None = None ) -> torch.Tensor: + active = self.active_lora_tensors() + if active is None: + return x.new_zeros((*x.shape[:-1], self.out_features)) + a_t, b_t, scale = active if tokens_per_expert is not None: assert self.num_local_experts > 1, ( "tokens_per_expert is only supported if num_local_experts > 1" @@ -586,12 +815,10 @@ def forward( if isinstance(bsz, list): bsz = torch.tensor(bsz, dtype=torch.int64, device="cpu") if x.shape[0] == 0: - return x.new_zeros((x.shape[0], self.B_T.shape[-1])) - return quack_grouped_lora(x, self.A_T, self.B_T, bsz, scale=self.scale) - out = (x @ self.A_T) @ self.B_T - if self.scale == 1.0: - return out - return out * self.scale + return x.new_zeros((*x.shape[:-1], self.out_features)) + return quack_grouped_lora(x, a_t, b_t, bsz, scale=scale) + out = (x @ a_t) @ b_t + return out if scale == 1.0 else out * scale class LoRAPublishPlanner: @@ -667,52 +894,47 @@ def _metadata_for_template( template: _LoraPublishTemplate, adapter_model: dict[str, torch.Tensor], ) -> list[LoraShardMeta]: - if template.num_local_experts > 1: - return self._expert_metadata_for_template(template, adapter_model) - return self._dense_metadata_for_template(template, adapter_model) - - def _dense_metadata_for_template( - self, - template: _LoraPublishTemplate, - adapter_model: dict[str, torch.Tensor], - ) -> list[LoraShardMeta]: - tp_ranks = self._dense_tp_ranks() shard_ranks = range(template.shard_world_size) if template.sharded else (0,) + if template.num_local_experts <= 1: + tp_ranks = ( + _process_group_ranks(ps.get_tensor_model_parallel_group()) + if _distributed_initialized() + else (0,) + ) + owners = [ + ( + f"{template.adapter_model_prefix}.{template.suffix}", + tp_ranks[shard_rank], + shard_rank, + ) + for shard_rank in shard_ranks + ] + else: + ep_world_size = 1 + if _distributed_initialized(): + ep_world_size = ps.get_expert_model_parallel_world_size() + owners = [ + ( + f"{template.adapter_model_prefix.format(expert=expert)}.{template.suffix}", + self._expert_owner_rank(ep_rank, shard_rank), + shard_rank, + ) + for ep_rank in range(ep_world_size) + for local_expert in range(template.num_local_experts) + for expert in [ep_rank * template.num_local_experts + local_expert] + for shard_rank in shard_ranks + ] return [ self._make_metadata( template, - key=f"{template.adapter_model_prefix}.{template.suffix}", - owner_rank=tp_ranks[shard_rank], + key=key, + owner_rank=owner_rank, shard_rank=shard_rank, adapter_model=adapter_model, ) - for shard_rank in shard_ranks + for key, owner_rank, shard_rank in owners ] - def _expert_metadata_for_template( - self, - template: _LoraPublishTemplate, - adapter_model: dict[str, torch.Tensor], - ) -> list[LoraShardMeta]: - ep_world_size = self._expert_model_world_size() - shard_ranks = range(template.shard_world_size) if template.sharded else (0,) - metadata: list[LoraShardMeta] = [] - for ep_rank in range(ep_world_size): - for local_expert in range(template.num_local_experts): - expert = ep_rank * template.num_local_experts + local_expert - key = f"{template.adapter_model_prefix.format(expert=expert)}.{template.suffix}" - for shard_rank in shard_ranks: - metadata.append( - self._make_metadata( - template, - key=key, - owner_rank=self._expert_owner_rank(ep_rank, shard_rank), - shard_rank=shard_rank, - adapter_model=adapter_model, - ) - ) - return metadata - @staticmethod def _make_metadata( template: _LoraPublishTemplate, @@ -722,6 +944,18 @@ def _make_metadata( shard_rank: int, adapter_model: dict[str, torch.Tensor], ) -> LoraShardMeta: + manifest: dict[str, Any] = { + "sharded": template.sharded, + "shard_world_size": template.shard_world_size if template.sharded else 1, + "shard_rank": shard_rank if template.sharded else 0, + } + if template.sharded: + manifest["export_shard_dim"] = template.export_shard_dim + manifest["export_shard_strategy"] = ( + template.export_shard_strategy or "uniform" + ) + if template.component_sizes: + manifest["component_sizes"] = list(template.component_sizes) return LoraShardMeta( key=key, owner_rank=owner_rank, @@ -731,22 +965,10 @@ def _make_metadata( if key in adapter_model else template.dtype_name ), - manifest=_publish_manifest(template, shard_rank=shard_rank), + manifest=manifest, block=_block_for_key(key), ) - @staticmethod - def _dense_tp_ranks() -> tuple[int, ...]: - if not _distributed_initialized(): - return (0,) - return _process_group_ranks(ps.get_tensor_model_parallel_group()) - - @staticmethod - def _expert_model_world_size() -> int: - if not _distributed_initialized(): - return 1 - return ps.get_expert_model_parallel_world_size() - @staticmethod def _expert_owner_rank(ep_rank: int, shard_rank: int) -> int: if not _distributed_initialized(): @@ -793,24 +1015,6 @@ def _exported_param_shape(module: LoRA, param: torch.nn.Parameter) -> tuple[int, return tuple(int(dim) for dim in param.T.shape) -def _publish_manifest( - template: _LoraPublishTemplate, - *, - shard_rank: int, -) -> dict[str, Any]: - manifest: dict[str, Any] = { - "sharded": template.sharded, - "shard_world_size": template.shard_world_size if template.sharded else 1, - "shard_rank": shard_rank if template.sharded else 0, - } - if template.sharded: - manifest["export_shard_dim"] = template.export_shard_dim - manifest["export_shard_strategy"] = template.export_shard_strategy or "uniform" - if template.component_sizes: - manifest["component_sizes"] = list(template.component_sizes) - return manifest - - @torch.compiler.disable def _expert_grouped_lora_forward( lora: LoRA, @@ -834,15 +1038,110 @@ def _expert_grouped_lora_dual_forward( counts = torch.tensor(counts, dtype=torch.int64, device="cpu") if x.shape[0] == 0: return x.new_zeros((x.shape[0], module.linear_fc1.out_features)) + gate = module.gate_lora.active_lora_tensors() + up = module.up_lora.active_lora_tensors() + if gate is None or up is None: + return torch.cat( + [ + module.gate_lora(x, tokens_per_expert=counts), + module.up_lora(x, tokens_per_expert=counts), + ], + dim=-1, + ) + gate_a_t, gate_b_t, gate_scale = gate + up_a_t, up_b_t, up_scale = up return quack_grouped_lora_dual( x, - module.gate_lora.A_T, - module.gate_lora.B_T, - module.up_lora.A_T, - module.up_lora.B_T, + gate_a_t, + gate_b_t, + up_a_t, + up_b_t, counts, - scale_gate=module.gate_lora.scale, - scale_up=module.up_lora.scale, + scale_gate=gate_scale, + scale_up=up_scale, + ) + + +def _parallel_lora( + *, + adapter_model_prefix: str, + linear: Any, + out_features: int, + rank: int, + alpha: float, + layout: Literal["column", "row"], + shard_domain: ShardDomain = "tp", + grad_sync_domain: GradSyncDomain = TP_DEFAULT_GRAD_SYNC_DOMAIN, + allreduce: bool = True, + num_local_experts: int = 1, +) -> LoRA: + weight = getattr(linear, "weight0", None) + if weight is None: + weight = getattr(linear, "weight", None) + assert isinstance(weight, torch.Tensor) + row_layout = layout == "row" + a_parallel_spec = LoRAParallelSpec( + shard_domain=shard_domain, + sharded=row_layout, + shard_dim=-2 if row_layout else None, + grad_sync_domain=grad_sync_domain, + grad_sync_op=GRAD_SYNC_OP_NONE if row_layout else GRAD_SYNC_OP_SUM, + ) + b_parallel_spec = replace( + a_parallel_spec, + sharded=not row_layout, + shard_dim=None if row_layout else -1, + grad_sync_domain=grad_sync_domain, + grad_sync_op=GRAD_SYNC_OP_SUM if row_layout else GRAD_SYNC_OP_NONE, + ) + return LoRA( + adapter_model_prefix=adapter_model_prefix, + in_features=linear.in_features, + out_features=out_features, + rank=rank, + alpha=alpha, + dtype=weight.dtype, + device=weight.device, + num_local_experts=num_local_experts, + a_parallel_spec=a_parallel_spec, + b_parallel_spec=b_parallel_spec, + allreduce=allreduce, + ) + + +def _parallel_lora_pair( + *, + adapter_model_prefix: str, + linear: Any, + out_features: int, + rank: int, + alpha: float, + layout: Literal["column", "row"], + suffixes: tuple[str, str], + num_local_experts: int = 1, +) -> tuple[LoRA, LoRA]: + expert_parallel = num_local_experts > 1 + return cast( + tuple[LoRA, LoRA], + tuple( + _parallel_lora( + adapter_model_prefix=f"{adapter_model_prefix}.{suffix}", + linear=linear, + out_features=out_features, + rank=rank, + alpha=alpha, + layout=layout, + shard_domain="expert_tp" if expert_parallel else "tp", + grad_sync_domain=( + EXPERT_TP_GRAD_SYNC_DOMAIN + if expert_parallel + else TP_DEFAULT_GRAD_SYNC_DOMAIN + ), + allreduce=not expert_parallel, + num_local_experts=num_local_experts, + ) + for suffix in suffixes + ), ) @@ -860,33 +1159,13 @@ def __init__( self.provider = provider self.linear_proj = linear_proj self.reduce_output = reduce_output - assert isinstance(linear_proj.weight, torch.Tensor) - a_parallel_spec = LoRAParallelSpec( - shard_domain="tp", - sharded=True, - shard_dim=-2, - grad_sync_domain=TP_DEFAULT_GRAD_SYNC_DOMAIN, - grad_sync_op=GRAD_SYNC_OP_NONE, # only need DP-type reductions - ) - b_parallel_spec = a_parallel_spec.model_copy( - update={ - "sharded": False, - "shard_dim": None, - "grad_sync_op": GRAD_SYNC_OP_SUM, # sum replicated TP contributions - } - ) - self.lora = LoRA( + self.lora = _parallel_lora( adapter_model_prefix=adapter_model_prefix, - in_features=linear_proj.in_features, + linear=linear_proj, out_features=linear_proj.out_features, rank=rank, alpha=alpha, - dtype=linear_proj.weight.dtype, - device=linear_proj.weight.device, - a_parallel_spec=a_parallel_spec, - b_parallel_spec=b_parallel_spec, - # Non-expert LoRA params use Megatron's dense DP/CP gradient buckets. - allreduce=True, + layout="row", ) def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: @@ -1011,12 +1290,11 @@ def _build_qkv_lora( grad_sync_domain=TP_DEFAULT_GRAD_SYNC_DOMAIN, grad_sync_op=GRAD_SYNC_OP_SUM, # sum replicated TP contributions ) - b_parallel_spec = a_parallel_spec.model_copy( - update={ - "sharded": True, - "shard_dim": -1, - "grad_sync_op": GRAD_SYNC_OP_NONE, # only need DP-type reductions - } + b_parallel_spec = replace( + a_parallel_spec, + sharded=True, + shard_dim=-1, + grad_sync_op=GRAD_SYNC_OP_NONE, # only need DP-type reductions ) return LoRA( adapter_model_prefix=adapter_model_prefix, @@ -1119,13 +1397,13 @@ def __init__( z_out_features_per_partition = ( gated_delta_net.v_dim // ps.get_tensor_model_parallel_world_size() ) - assert isinstance(in_proj.weight, torch.Tensor) - self.qkv_lora = self._build_in_proj_lora( + self.qkv_lora = _parallel_lora( adapter_model_prefix=f"{adapter_model_prefix}.in_proj_qkv", - in_proj=in_proj, + linear=in_proj, + out_features=qkv_out_features_per_partition, rank=rank, alpha=alpha, - out_features=qkv_out_features_per_partition, + layout="column", ) _set_lora_shard_strategy_metadata( self.qkv_lora.B_T, @@ -1136,49 +1414,13 @@ def __init__( gated_delta_net.v_dim, ), ) - self.z_lora = self._build_in_proj_lora( + self.z_lora = _parallel_lora( adapter_model_prefix=f"{adapter_model_prefix}.in_proj_z", - in_proj=in_proj, - rank=rank, - alpha=alpha, + linear=in_proj, out_features=z_out_features_per_partition, - ) - - @staticmethod - def _build_in_proj_lora( - *, - adapter_model_prefix: str, - in_proj: TELayerNormColumnParallelLinear, - rank: int, - alpha: float, - out_features: int, - ) -> LoRA: - assert isinstance(in_proj.weight, torch.Tensor) - a_parallel_spec = LoRAParallelSpec( - shard_domain="tp", - sharded=False, - shard_dim=None, - grad_sync_domain=TP_DEFAULT_GRAD_SYNC_DOMAIN, - grad_sync_op=GRAD_SYNC_OP_SUM, - ) - b_parallel_spec = a_parallel_spec.model_copy( - update={ - "sharded": True, - "shard_dim": -1, - "grad_sync_op": GRAD_SYNC_OP_NONE, - } - ) - return LoRA( - adapter_model_prefix=adapter_model_prefix, - in_features=in_proj.in_features, - out_features=out_features, rank=rank, alpha=alpha, - dtype=in_proj.weight.dtype, - device=in_proj.weight.device, - a_parallel_spec=a_parallel_spec, - b_parallel_spec=b_parallel_spec, - allreduce=True, + layout="column", ) def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: @@ -1208,133 +1450,62 @@ def __init__( rank: int, alpha: float, num_local_experts: int, + fused_gate_up: bool = False, ) -> None: super().__init__() - assert linear_fc1 is not None - self.linear_fc1 = linear_fc1 - self.gate_lora = self._build_fc1_lora( - adapter_model_prefix=f"{adapter_model_prefix}.{{expert}}.gate_proj", - linear_fc1=linear_fc1, - rank=rank, - alpha=alpha, - num_local_experts=num_local_experts, - ) - self.up_lora = self._build_fc1_lora( - adapter_model_prefix=f"{adapter_model_prefix}.{{expert}}.up_proj", - linear_fc1=linear_fc1, - rank=rank, - alpha=alpha, - num_local_experts=num_local_experts, - ) - self.uses_direct_quack_grouped_lora_dual = True - - @staticmethod - def _build_fc1_lora( - *, - adapter_model_prefix: str, - linear_fc1: TEColumnParallelGroupedLinear, - rank: int, - alpha: float, - num_local_experts: int, - ) -> LoRA: - assert linear_fc1 is not None - assert isinstance(linear_fc1.weight0, torch.Tensor) - a_parallel_spec = LoRAParallelSpec( - shard_domain="expert_tp", - sharded=False, - shard_dim=None, - grad_sync_domain=EXPERT_TP_GRAD_SYNC_DOMAIN, - grad_sync_op=GRAD_SYNC_OP_SUM, # we handle this with extended finalize_grads - ) - b_parallel_spec = a_parallel_spec.model_copy( - update={ - "sharded": True, - "shard_dim": -1, - "grad_sync_domain": EXPERT_TP_GRAD_SYNC_DOMAIN, - "grad_sync_op": GRAD_SYNC_OP_NONE, # only need DP-type reductions - } - ) - return LoRA( - adapter_model_prefix=adapter_model_prefix, - in_features=linear_fc1.in_features, - out_features=linear_fc1.out_features // 2, - rank=rank, - alpha=alpha, - dtype=linear_fc1.weight0.dtype, - device=linear_fc1.weight0.device, - num_local_experts=num_local_experts, - a_parallel_spec=a_parallel_spec, - b_parallel_spec=b_parallel_spec, - # Expert LoRA params use Megatron's expert-DP gradient buckets. - allreduce=False, - ) - - def forward( - self, x: torch.Tensor, tokens_per_expert: list[int] | torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor | None]: - base_out, bias_out = self.linear_fc1(x, tokens_per_expert) - adapter_out = _expert_grouped_lora_dual_forward(self, x, tokens_per_expert) - return base_out + adapter_out, bias_out - - -class MLPExpertsLinearFC1FusedLoRA(torch.nn.Module): - def __init__( - self, - adapter_model_prefix: str, - linear_fc1: TEColumnParallelGroupedLinear, - rank: int, - alpha: float, - num_local_experts: int, - ) -> None: - super().__init__() - assert linear_fc1 is not None - assert isinstance(linear_fc1.weight0, torch.Tensor) self.linear_fc1 = linear_fc1 - a_parallel_spec = LoRAParallelSpec( - shard_domain="expert_tp", - sharded=False, - shard_dim=None, - grad_sync_domain=EXPERT_TP_GRAD_SYNC_DOMAIN, - grad_sync_op=GRAD_SYNC_OP_SUM, - ) - b_parallel_spec = a_parallel_spec.model_copy( - update={ - "sharded": True, - "shard_dim": -1, - "grad_sync_domain": EXPERT_TP_GRAD_SYNC_DOMAIN, - "grad_sync_op": GRAD_SYNC_OP_NONE, - } - ) - self.lora = LoRA( - adapter_model_prefix=f"{adapter_model_prefix}.{{expert}}.gate_up_proj", - in_features=linear_fc1.in_features, - out_features=linear_fc1.out_features, - rank=rank, - alpha=alpha, - dtype=linear_fc1.weight0.dtype, - device=linear_fc1.weight0.device, - num_local_experts=num_local_experts, - a_parallel_spec=a_parallel_spec, - b_parallel_spec=b_parallel_spec, - allreduce=False, - ) - gate_out_features = linear_fc1.out_features // 2 - expert_tp_world_size = _get_shard_world_size("expert_tp") - _set_lora_shard_strategy_metadata( - self.lora.B_T, - strategy="componentwise", - component_sizes=( - gate_out_features * expert_tp_world_size, - gate_out_features * expert_tp_world_size, - ), - ) + self.fused_gate_up = bool(fused_gate_up) + if self.fused_gate_up: + self.lora = _parallel_lora( + adapter_model_prefix=f"{adapter_model_prefix}.{{expert}}.gate_up_proj", + linear=linear_fc1, + out_features=linear_fc1.out_features, + rank=rank, + alpha=alpha, + layout="column", + shard_domain="expert_tp", + grad_sync_domain=EXPERT_TP_GRAD_SYNC_DOMAIN, + allreduce=False, + num_local_experts=num_local_experts, + ) + gate_out_features = linear_fc1.out_features // 2 + expert_tp_world_size = _get_shard_world_size("expert_tp") + _set_lora_shard_strategy_metadata( + self.lora.B_T, + strategy="componentwise", + component_sizes=( + gate_out_features * expert_tp_world_size, + gate_out_features * expert_tp_world_size, + ), + ) + else: + self.gate_lora, self.up_lora = _parallel_lora_pair( + adapter_model_prefix=f"{adapter_model_prefix}.{{expert}}", + linear=linear_fc1, + out_features=linear_fc1.out_features // 2, + rank=rank, + alpha=alpha, + layout="column", + suffixes=("gate_proj", "up_proj"), + num_local_experts=num_local_experts, + ) def forward( self, x: torch.Tensor, tokens_per_expert: list[int] | torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor | None]: - base_out, bias_out = self.linear_fc1(x, tokens_per_expert) - adapter_out = _expert_grouped_lora_forward( - self.lora, x, tokens_per_expert, self.linear_fc1.out_features + base_out, bias_out = cast( + Callable[ + [torch.Tensor, list[int] | torch.Tensor], + tuple[torch.Tensor, torch.Tensor | None], + ], + self.linear_fc1, + )(x, tokens_per_expert) + adapter_out = ( + _expert_grouped_lora_forward( + self.lora, x, tokens_per_expert, self.linear_fc1.out_features + ) + if self.fused_gate_up + else _expert_grouped_lora_dual_forward(self, x, tokens_per_expert) ) return base_out + adapter_out, bias_out @@ -1349,43 +1520,30 @@ def __init__( num_local_experts: int, ) -> None: super().__init__() - assert linear_fc2 is not None - assert isinstance(linear_fc2.weight0, torch.Tensor) self.linear_fc2 = linear_fc2 - a_parallel_spec = LoRAParallelSpec( - shard_domain="expert_tp", - sharded=True, - shard_dim=-2, - grad_sync_domain=EXPERT_TP_GRAD_SYNC_DOMAIN, - grad_sync_op=GRAD_SYNC_OP_NONE, # only need DP-type reductions - ) - b_parallel_spec = a_parallel_spec.model_copy( - update={ - "sharded": False, - "shard_dim": None, - "grad_sync_domain": EXPERT_TP_GRAD_SYNC_DOMAIN, - "grad_sync_op": GRAD_SYNC_OP_SUM, # we handle this with extended finalize_grads - } - ) - self.lora = LoRA( + self.lora = _parallel_lora( adapter_model_prefix=f"{adapter_model_prefix}.{{expert}}.down_proj", - in_features=linear_fc2.in_features, + linear=linear_fc2, out_features=linear_fc2.out_features, rank=rank, alpha=alpha, - dtype=linear_fc2.weight0.dtype, - device=linear_fc2.weight0.device, - num_local_experts=num_local_experts, - a_parallel_spec=a_parallel_spec, - b_parallel_spec=b_parallel_spec, - # Expert LoRA params use Megatron's expert-DP gradient buckets. + layout="row", + shard_domain="expert_tp", + grad_sync_domain=EXPERT_TP_GRAD_SYNC_DOMAIN, allreduce=False, + num_local_experts=num_local_experts, ) def forward( self, x: torch.Tensor, tokens_per_expert: list[int] | torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor | None]: - base_out, bias_out = self.linear_fc2(x, tokens_per_expert) + base_out, bias_out = cast( + Callable[ + [torch.Tensor, list[int] | torch.Tensor], + tuple[torch.Tensor, torch.Tensor | None], + ], + self.linear_fc2, + )(x, tokens_per_expert) adapter_out = _expert_grouped_lora_forward( self.lora, x, tokens_per_expert, self.linear_fc2.out_features ) @@ -1407,56 +1565,28 @@ def __init__( linear_fc1.return_layernorm_output = True linear_fc1.return_layernorm_output_gathered = True self.linear_fc1 = linear_fc1 - self.gate_lora = self._build_fc1_lora( - adapter_model_prefix=f"{adapter_model_prefix}.gate_proj", - linear_fc1=linear_fc1, - rank=rank, - alpha=alpha, - ) - self.up_lora = self._build_fc1_lora( - adapter_model_prefix=f"{adapter_model_prefix}.up_proj", - linear_fc1=linear_fc1, - rank=rank, - alpha=alpha, - ) - - @staticmethod - def _build_fc1_lora( - *, - adapter_model_prefix: str, - linear_fc1: TEColumnParallelLinear | TELayerNormColumnParallelLinear, - rank: int, - alpha: float, - ) -> LoRA: - assert isinstance(linear_fc1.weight, torch.Tensor) - a_parallel_spec = LoRAParallelSpec( - shard_domain="tp", - sharded=False, - shard_dim=None, - grad_sync_domain=TP_DEFAULT_GRAD_SYNC_DOMAIN, - grad_sync_op=GRAD_SYNC_OP_SUM, - ) - b_parallel_spec = a_parallel_spec.model_copy( - update={ - "sharded": True, - "shard_dim": -1, - "grad_sync_op": GRAD_SYNC_OP_NONE, - } - ) - return LoRA( + self.gate_lora, self.up_lora = _parallel_lora_pair( adapter_model_prefix=adapter_model_prefix, - in_features=linear_fc1.in_features, + linear=linear_fc1, out_features=linear_fc1.out_features // 2, rank=rank, alpha=alpha, - dtype=linear_fc1.weight.dtype, - device=linear_fc1.weight.device, - a_parallel_spec=a_parallel_spec, - b_parallel_spec=b_parallel_spec, - allreduce=True, + layout="column", + suffixes=("gate_proj", "up_proj"), ) def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: + if int(x.numel()) == 0: + zero = x.sum() * 0.0 + weight = getattr(self.linear_fc1, "weight", None) + if isinstance(weight, torch.Tensor): + zero = zero + weight.to(dtype=x.dtype).sum() * 0.0 + for lora in (self.gate_lora, self.up_lora): + zero = zero + lora.A_T.to(dtype=x.dtype).sum() * 0.0 + zero = zero + lora.B_T.to(dtype=x.dtype).sum() * 0.0 + return zero.expand( + *x.shape[:-1], self.linear_fc1.out_features + ).clone(), None base_output, bias_out = self.linear_fc1(x) if isinstance(base_output, tuple): base_out, lora_input = base_output @@ -1608,8 +1738,14 @@ def wrap_grouped_moe_experts( target_modules: set[str], rank: int, alpha: int, + fused_gate_up: bool = False, ) -> None: - if _targets_include(target_modules, "gate_proj", "up_proj"): + wrap_fc1 = ( + _targets_include(target_modules, "experts") + if fused_gate_up + else _targets_include(target_modules, "gate_proj", "up_proj") + ) + if wrap_fc1: mlp_experts_linear_fc1 = _unwrap_attr( experts.linear_fc1, "linear_fc1", @@ -1621,58 +1757,27 @@ def wrap_grouped_moe_experts( rank=rank, alpha=alpha, num_local_experts=experts.num_local_experts, + fused_gate_up=fused_gate_up, ) - if _targets_include(target_modules, "down_proj"): - mlp_experts_linear_fc2 = _unwrap_attr( - experts.linear_fc2, - "linear_fc2", - TERowParallelGroupedLinear, # type: ignore[arg-type] - ) - experts.linear_fc2 = MLPExpertsLinearFC2LoRA( - adapter_model_prefix=f"{adapter_model_prefix}.mlp.experts", - linear_fc2=mlp_experts_linear_fc2, - rank=rank, - alpha=alpha, - num_local_experts=experts.num_local_experts, - ) - - -def wrap_grouped_moe_experts_3d( - experts: TEGroupedMLP, - *, - adapter_model_prefix: str, - target_modules: set[str], - rank: int, - alpha: int, -) -> None: - if _targets_include(target_modules, "experts"): - mlp_experts_linear_fc1 = _unwrap_attr( - experts.linear_fc1, - "linear_fc1", - TEColumnParallelGroupedLinear, # type: ignore[arg-type] - ) - experts.linear_fc1 = MLPExpertsLinearFC1FusedLoRA( - adapter_model_prefix=f"{adapter_model_prefix}.mlp.experts", - linear_fc1=mlp_experts_linear_fc1, - rank=rank, - alpha=alpha, - num_local_experts=experts.num_local_experts, - ) - mlp_experts_linear_fc2 = _unwrap_attr( + wrap_fc2 = ( + wrap_fc1 if fused_gate_up else _targets_include(target_modules, "down_proj") + ) + if wrap_fc2: + linear_fc2 = _unwrap_attr( experts.linear_fc2, "linear_fc2", TERowParallelGroupedLinear, # type: ignore[arg-type] ) experts.linear_fc2 = MLPExpertsLinearFC2LoRA( adapter_model_prefix=f"{adapter_model_prefix}.mlp.experts", - linear_fc2=mlp_experts_linear_fc2, + linear_fc2=linear_fc2, rank=rank, alpha=alpha, num_local_experts=experts.num_local_experts, ) -def wrap_dense_mlp( +def wrap_split_mlp_lora( mlp: Any, *, adapter_model_prefix: str, @@ -1682,34 +1787,71 @@ def wrap_dense_mlp( alpha: int, ) -> None: if _targets_include(target_modules, "gate_proj", "up_proj"): - mlp_linear_fc1 = _unwrap_attr( + linear_fc1 = _unwrap_attr( mlp.linear_fc1, "linear_fc1", (TEColumnParallelLinear, TELayerNormColumnParallelLinear), ) mlp.linear_fc1 = SharedExpertsLinearFC1LoRA( - adapter_model_prefix=f"{adapter_model_prefix}.mlp", - linear_fc1=mlp_linear_fc1, + adapter_model_prefix=adapter_model_prefix, + linear_fc1=linear_fc1, rank=rank, alpha=alpha, ) if _targets_include(target_modules, "down_proj"): - mlp_linear_fc2 = _unwrap_attr( + linear_fc2 = _unwrap_attr( mlp.linear_fc2, "linear_fc2", TERowParallelLinear, ) mlp.linear_fc2 = SharedExpertsLinearFC2LoRA( - adapter_model_prefix=f"{adapter_model_prefix}.mlp", - linear_fc2=mlp_linear_fc2, + adapter_model_prefix=adapter_model_prefix, + linear_fc2=linear_fc2, rank=rank, alpha=alpha, provider=provider, ) +def wrap_grouped_moe_experts_3d( + experts: TEGroupedMLP, + *, + adapter_model_prefix: str, + target_modules: set[str], + rank: int, + alpha: int, +) -> None: + wrap_grouped_moe_experts( + experts, + adapter_model_prefix=adapter_model_prefix, + target_modules=target_modules, + rank=rank, + alpha=alpha, + fused_gate_up=True, + ) + + +def wrap_dense_mlp( + mlp: Any, + *, + adapter_model_prefix: str, + provider: GPTModelProvider, + target_modules: set[str], + rank: int, + alpha: int, +) -> None: + wrap_split_mlp_lora( + mlp, + adapter_model_prefix=f"{adapter_model_prefix}.mlp", + provider=provider, + target_modules=target_modules, + rank=rank, + alpha=alpha, + ) + + def wrap_shared_experts_mlp( - shared_experts: SharedExpertMLP, + shared_experts: Any, *, adapter_model_prefix: str, provider: GPTModelProvider, @@ -1717,31 +1859,14 @@ def wrap_shared_experts_mlp( rank: int, alpha: int, ) -> None: - if _targets_include(target_modules, "gate_proj", "up_proj"): - shared_experts_linear_fc1 = _unwrap_attr( - shared_experts.linear_fc1, - "linear_fc1", - (TEColumnParallelLinear, TELayerNormColumnParallelLinear), - ) - shared_experts.linear_fc1 = SharedExpertsLinearFC1LoRA( - adapter_model_prefix=f"{adapter_model_prefix}.mlp.shared_expert", - linear_fc1=shared_experts_linear_fc1, - rank=rank, - alpha=alpha, - ) - if _targets_include(target_modules, "down_proj"): - shared_experts_linear_fc2 = _unwrap_attr( - shared_experts.linear_fc2, - "linear_fc2", - TERowParallelLinear, - ) - shared_experts.linear_fc2 = SharedExpertsLinearFC2LoRA( - adapter_model_prefix=f"{adapter_model_prefix}.mlp.shared_expert", - linear_fc2=shared_experts_linear_fc2, - rank=rank, - alpha=alpha, - provider=provider, - ) + wrap_split_mlp_lora( + shared_experts, + adapter_model_prefix=f"{adapter_model_prefix}.mlp.shared_experts", + provider=provider, + target_modules=target_modules, + rank=rank, + alpha=alpha, + ) def apply_lora_adapters( @@ -1761,3 +1886,43 @@ def apply_lora_adapters( alpha=LORA_ALPHA, ) return list(model) + + +def load_lora_slot_into_model( + model: Sequence[torch.nn.Module], + ref: LoRASlotRef, + adapter_model: dict[str, torch.Tensor], + *, + alpha: float = LORA_ALPHA, + requires_grad: bool, +) -> int: + loaded = 0 + for chunk in model: + for module in chunk.modules(): + if isinstance(module, LoRA) and module.load_lora_slot( + ref, + adapter_model, + alpha=alpha, + requires_grad=requires_grad, + ): + loaded += 1 + if loaded == 0 and ref.name is not None: + raise RuntimeError(f"LoRA slot {ref.kind}:{ref.name} loaded no adapter sites") + return loaded + + +def iter_lora_slot_parameters( + model: Sequence[torch.nn.Module], + ref: LoRASlotRef, +) -> Iterator[torch.nn.Parameter]: + seen: set[int] = set() + for chunk in model: + for module in chunk.modules(): + if not isinstance(module, LoRA): + continue + for param in module.lora_slot_params(ref): + param_id = id(param) + if param_id in seen: + continue + seen.add(param_id) + yield param diff --git a/src/art/megatron/model_support/handlers/default_dense.py b/src/art/megatron/model_support/handlers/default_dense.py index cc7128215..ba168ac16 100644 --- a/src/art/megatron/model_support/handlers/default_dense.py +++ b/src/art/megatron/model_support/handlers/default_dense.py @@ -8,9 +8,9 @@ FlexAttentionCompileCrashConfig, HfWeightSource, LayerFamilyInstance, + PrefixTreeModelStateContext, RolloutWeightsMode, SharedExpertCompileState, - SharedPrefixModelStateContext, ) _CONTEXT_PARALLEL_ATTENTION_WORKAROUND_FLAG = "context_parallel_attention" @@ -125,9 +125,9 @@ def install_preprocess_patch(self, model_chunks: Sequence[Any]) -> None: del model_chunks return None - def build_shared_prefix_model_state( + def build_prefix_tree_model_state( self, - context: SharedPrefixModelStateContext, + context: PrefixTreeModelStateContext, ) -> dict[str, Any]: del context return {} @@ -195,7 +195,7 @@ def apply_lora_adapters( from art.megatron.lora import ( _adapter_model_prefix, - wrap_dense_mlp, + wrap_split_mlp_lora, wrap_standard_self_attention, ) @@ -204,18 +204,19 @@ def apply_lora_adapters( for module in chunk.modules(): if not isinstance(module, TransformerLayer): continue + adapter_model_prefix = _adapter_model_prefix(module) wrap_standard_self_attention( module.self_attention, - adapter_model_prefix=_adapter_model_prefix(module), + adapter_model_prefix=adapter_model_prefix, provider=provider, target_modules=target_set, rank=rank, alpha=alpha, ) _require_dense_mlp(module) - wrap_dense_mlp( + wrap_split_mlp_lora( module.mlp, - adapter_model_prefix=_adapter_model_prefix(module), + adapter_model_prefix=f"{adapter_model_prefix}.mlp", provider=provider, target_modules=target_set, rank=rank, @@ -226,32 +227,9 @@ def build_adapter_weights_by_base( self, model_chunks: Sequence[Any], ) -> dict[str, list[Any]]: - from megatron.core.transformer.transformer_layer import TransformerLayer + from art.megatron.weights import adapter_export - from art.megatron.weights.adapter_export import ( - add_dense_mlp_adapter_weights, - add_standard_self_attention_adapter_weights, - layer_base_prefix, - ) - - adapter_weights_by_base: dict[str, list[Any]] = {} - for chunk in model_chunks: - for module_name, module in chunk.named_modules(): - if not isinstance(module, TransformerLayer): - continue - layer_prefix = layer_base_prefix(module, module_name=module_name) - _require_dense_mlp(module) - add_standard_self_attention_adapter_weights( - adapter_weights_by_base, - layer_prefix=layer_prefix, - self_attention=module.self_attention, - ) - add_dense_mlp_adapter_weights( - adapter_weights_by_base, - layer_prefix=layer_prefix, - mlp=module.mlp, - ) - return adapter_weights_by_base + return adapter_export.build_transformer_layer_adapter_weights(model_chunks) def compile_workaround_config( self, @@ -301,7 +279,7 @@ def apply_lora_adapters( from art.megatron.lora import ( _adapter_model_prefix, wrap_grouped_moe_experts, - wrap_shared_experts_mlp, + wrap_split_mlp_lora, wrap_standard_self_attention, ) @@ -328,9 +306,9 @@ def apply_lora_adapters( ) shared_experts = getattr(module.mlp, "shared_experts", None) if shared_experts is not None: - wrap_shared_experts_mlp( + wrap_split_mlp_lora( shared_experts, - adapter_model_prefix=adapter_model_prefix, + adapter_model_prefix=f"{adapter_model_prefix}.mlp.shared_expert", provider=provider, target_modules=target_set, rank=rank, @@ -341,40 +319,13 @@ def build_adapter_weights_by_base( self, model_chunks: Sequence[Any], ) -> dict[str, list[Any]]: - from megatron.core.transformer.transformer_layer import TransformerLayer + from art.megatron.weights import adapter_export - from art.megatron.weights.adapter_export import ( - add_grouped_moe_adapter_weights, - add_shared_experts_adapter_weights, - add_standard_self_attention_adapter_weights, - layer_base_prefix, + return adapter_export.build_transformer_layer_adapter_weights( + model_chunks, + grouped_moe=True, ) - adapter_weights_by_base: dict[str, list[Any]] = {} - for chunk in model_chunks: - for module_name, module in chunk.named_modules(): - if not isinstance(module, TransformerLayer): - continue - layer_prefix = layer_base_prefix(module, module_name=module_name) - add_standard_self_attention_adapter_weights( - adapter_weights_by_base, - layer_prefix=layer_prefix, - self_attention=module.self_attention, - ) - add_grouped_moe_adapter_weights( - adapter_weights_by_base, - layer_prefix=layer_prefix, - experts=_require_moe_experts(module), - ) - shared_experts = getattr(module.mlp, "shared_experts", None) - if shared_experts is not None: - add_shared_experts_adapter_weights( - adapter_weights_by_base, - layer_prefix=layer_prefix, - shared_experts=shared_experts, - ) - return adapter_weights_by_base - def _require_dense_mlp(module: Any) -> None: if getattr(module.mlp, "experts", None) is not None: diff --git a/src/art/megatron/model_support/handlers/dsv4.py b/src/art/megatron/model_support/handlers/dsv4.py index 500df992e..c6a604049 100644 --- a/src/art/megatron/model_support/handlers/dsv4.py +++ b/src/art/megatron/model_support/handlers/dsv4.py @@ -15,7 +15,7 @@ from art.megatron.model_support.spec import ( CompileWorkaroundConfig, LayerFamilyInstance, - SharedPrefixModelStateContext, + PrefixTreeModelStateContext, ) _ORACLE_HIDDEN_SIZE = 512 @@ -107,22 +107,22 @@ def configure_tokenizer( return tokenizer return get_dsv4_tokenizer(tokenizer) - def build_shared_prefix_model_state( + def build_prefix_tree_model_state( self, - context: SharedPrefixModelStateContext, + context: PrefixTreeModelStateContext, ) -> dict[str, Any]: if context.input_pos is None: raise RuntimeError( - "DSV4 shared-prefix compression layouts require input_pos." + "DSV4 prefix-tree compression layouts require input_pos." ) from art.megatron.dsv4.compressor import ( - Dsv4SharedPrefixState, - build_shared_prefix_compression_layouts, + Dsv4PrefixTreeState, + build_prefix_tree_compression_layouts, ) return { - "dsv4": Dsv4SharedPrefixState( - compression_layouts=build_shared_prefix_compression_layouts( + "dsv4": Dsv4PrefixTreeState( + compression_layouts=build_prefix_tree_compression_layouts( position_ids=context.input_pos, group_ids=context.group_ids, parent_ids=context.parent_ids, diff --git a/src/art/megatron/model_support/handlers/gemma4.py b/src/art/megatron/model_support/handlers/gemma4.py index 33f57ece3..6534538f3 100644 --- a/src/art/megatron/model_support/handlers/gemma4.py +++ b/src/art/megatron/model_support/handlers/gemma4.py @@ -71,7 +71,7 @@ r"(?Plora_[AB])\.weight$" ) _SHARED_EXPERT_FC1_LORA_A_KEY_RE = re.compile( - r"^.*\.layers\.(?P\d+)\.mlp\.(?:shared_expert\.)?" + r"^.*\.layers\.(?P\d+)\.mlp\.(?:shared_experts\.)?" r"(?:gate_proj|up_proj)\.lora_A\.weight$" ) _SELF_ATTN_V_LORA_KEY_RE = re.compile( @@ -356,14 +356,7 @@ def flex_attention_compile_crash_config( self, provider: Any, ) -> FlexAttentionCompileCrashConfig: - if ( - _gemma4_compile_crash_signature(provider) - in _GEMMA4_TRITON_NUM_STAGES_2_SIGNATURES - ): - return FlexAttentionCompileCrashConfig( - triton_num_stages_2_head_dims=(int(provider.global_head_dim),) - ) - return FlexAttentionCompileCrashConfig() + return _gemma4_flex_attention_compile_crash_config(provider) GEMMA4_MOE_HANDLER = Gemma4MoeHandler() @@ -553,14 +546,7 @@ def flex_attention_compile_crash_config( self, provider: Any, ) -> FlexAttentionCompileCrashConfig: - if ( - _gemma4_compile_crash_signature(provider) - in _GEMMA4_TRITON_NUM_STAGES_2_SIGNATURES - ): - return FlexAttentionCompileCrashConfig( - triton_num_stages_2_head_dims=(int(provider.global_head_dim),) - ) - return FlexAttentionCompileCrashConfig() + return _gemma4_flex_attention_compile_crash_config(provider) def get_forward_kwargs(self, model: Any, **kwargs: Any) -> dict[str, Any]: return _gemma4_forward_kwargs(model, **kwargs) @@ -991,6 +977,19 @@ def _gemma4_attention_pattern(provider: Any) -> tuple[int, int]: return (int(pattern[0]), int(pattern[1])) +def _gemma4_flex_attention_compile_crash_config( + provider: Any, +) -> FlexAttentionCompileCrashConfig: + if ( + _gemma4_compile_crash_signature(provider) + in _GEMMA4_TRITON_NUM_STAGES_2_SIGNATURES + ): + return FlexAttentionCompileCrashConfig( + triton_num_stages_2_head_dims=(int(provider.global_head_dim),) + ) + return FlexAttentionCompileCrashConfig() + + def _gemma4_compile_crash_signature(provider: Any) -> tuple[Any, ...]: return ( "moe" if int(getattr(provider, "num_moe_experts", 0) or 0) > 0 else "dense", @@ -1136,9 +1135,10 @@ def _wrap_gemma4_attention_output_lora( def _to_vllm_key(key: str) -> str: - key = key.replace(".mlp.shared_expert.", ".mlp.").replace( - ".mlp.experts", - ".moe.experts", + key = ( + key.replace(".mlp.shared_experts.", ".mlp.") + .replace(".mlp.shared_expert.", ".mlp.") + .replace(".mlp.experts", ".moe.experts") ) return _HF_TEXT_EXPERT_KEY_RE.sub(r"\g.moe.experts", key) @@ -1146,7 +1146,7 @@ def _to_vllm_key(key: str) -> str: def _from_vllm_key(key: str) -> str: key = key.replace(".moe.experts", ".mlp.experts") return _DENSE_MLP_LORA_KEY_RE.sub( - r"\g.shared_expert.\g.\g.weight", + r"\g.shared_experts.\g.\g.weight", key, ) @@ -1227,10 +1227,8 @@ def _gemma4_shared_expert_prenorm_corrections( ) weight_map = dict(index["weight_map"]) text_config = _gemma4_text_config_dict(base_model_name_or_path) - num_layers = int(text_config["num_hidden_layers"]) norm_keys_by_file: dict[str, list[tuple[int, str, str]]] = {} - - for layer in range(num_layers): + for layer in range(int(text_config["num_hidden_layers"])): for suffix in ( "pre_feedforward_layernorm", "pre_feedforward_layernorm_2", @@ -1243,6 +1241,7 @@ def _gemma4_shared_expert_prenorm_corrections( norm_keys_by_file.setdefault(weight_map[key], []).append( (layer, suffix, key) ) + norm_weights: dict[tuple[int, str], torch.Tensor] = {} for filename, entries in norm_keys_by_file.items(): with safe_open( @@ -1256,24 +1255,10 @@ def _gemma4_shared_expert_prenorm_corrections( return tuple( norm_weights[(layer, "pre_feedforward_layernorm")] / norm_weights[(layer, "pre_feedforward_layernorm_2")] - for layer in range(num_layers) + for layer in range(int(text_config["num_hidden_layers"])) ) -def _shared_expert_fc1_prenorm_correction( - *, - adapter_config: dict[str, Any], - layer: int, - device: torch.device, -) -> torch.Tensor: - # Megatron Bridge folds pffl/pffl2 into shared-expert FC1 base weights because - # MCore feeds pffl2-normalized activations while HF/vLLM feeds pffl-normalized - # activations. LoRA-A needs the same basis change at the HF/vLLM boundary. - return _gemma4_shared_expert_prenorm_corrections( - str(adapter_config["base_model_name_or_path"]) - )[layer].to(device=device) - - def _rescale_shared_expert_fc1_lora_a( key: str, tensor: torch.Tensor, @@ -1284,11 +1269,10 @@ def _rescale_shared_expert_fc1_lora_a( match = _SHARED_EXPERT_FC1_LORA_A_KEY_RE.match(key) if match is None: return tensor - correction = _shared_expert_fc1_prenorm_correction( - adapter_config=adapter_config, - layer=int(match.group("layer")), - device=tensor.device, + corrections = _gemma4_shared_expert_prenorm_corrections( + str(adapter_config["base_model_name_or_path"]) ) + correction = corrections[int(match.group("layer"))].to(device=tensor.device) factor = correction.reciprocal() if to_vllm else correction return (tensor.float() * factor.unsqueeze(0)).to(tensor.dtype).contiguous() diff --git a/src/art/megatron/model_support/handlers/qwen3_5.py b/src/art/megatron/model_support/handlers/qwen3_5.py index ad200499a..9cf2c7064 100644 --- a/src/art/megatron/model_support/handlers/qwen3_5.py +++ b/src/art/megatron/model_support/handlers/qwen3_5.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Callable from copy import copy from functools import lru_cache import re @@ -50,6 +51,9 @@ r"^(?P.*\.mlp\.experts)\.(?P\d+)\." r"(?Pgate_proj|up_proj|down_proj)\.(?Plora_[AB])\.weight$" ) +_ART_MOE_MODULES = ("gate_up_proj", "down_proj") +_VLLM_EXPERT_MODULES = ("gate_proj", "up_proj", "down_proj") +_LORA_NAMES = ("lora_A", "lora_B") class Qwen35BaseHandler(DefaultDenseHandler): @@ -80,7 +84,7 @@ def to_vllm_lora_tensors( *, adapter_config: dict[str, Any], ) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: - if _group_art_moe_tensors(tensors): + if _group_expert_lora_tensors(tensors, _ART_MOE_EXPERT_KEY_RE): raise TypeError("Dense Qwen3.5 handler received MoE LoRA tensors") transformed: dict[str, torch.Tensor] = {} for key, tensor in tensors.items(): @@ -118,10 +122,10 @@ def install_preprocess_patch(self, model_chunks: Sequence[Any]) -> None: from art.megatron.gdn.operator import ( install_gdn_island_hooks, - install_shared_prefix_gdn_hooks, + install_prefix_tree_gdn_hooks, ) - install_shared_prefix_gdn_hooks(model_chunks) + install_prefix_tree_gdn_hooks(model_chunks) install_gdn_island_hooks(model_chunks) for chunk in list(model_chunks): module: Any = chunk @@ -313,44 +317,14 @@ def build_adapter_weights_by_base( self, model_chunks: Sequence[Any], ) -> dict[str, list[Any]]: - from megatron.core.ssm.gated_delta_net import GatedDeltaNet - from megatron.core.transformer.attention import SelfAttention - from megatron.core.transformer.transformer_layer import TransformerLayer - - from art.megatron.lora import _is_language_transformer_layer_name - from art.megatron.weights.adapter_export import ( - add_gated_delta_net_adapter_weights, - add_standard_self_attention_adapter_weights, - layer_base_prefix, - ) + from art.megatron.weights import adapter_export _ensure_bridge_qwen35_adapter_name_map() - adapter_weights_by_base: dict[str, list[Any]] = {} - for chunk in model_chunks: - for module_name, module in chunk.named_modules(): - if not isinstance(module, TransformerLayer): - continue - if not _is_language_transformer_layer_name(module_name): - continue - layer_prefix = layer_base_prefix(module, module_name=module_name) - if isinstance(module.self_attention, SelfAttention): - add_standard_self_attention_adapter_weights( - adapter_weights_by_base, - layer_prefix=layer_prefix, - self_attention=module.self_attention, - ) - elif isinstance(module.self_attention, GatedDeltaNet): - add_gated_delta_net_adapter_weights( - adapter_weights_by_base, - layer_prefix=layer_prefix, - self_attention=module.self_attention, - ) - self._add_mlp_adapter_weights( - adapter_weights_by_base, - layer_prefix=layer_prefix, - module=module, - ) - return adapter_weights_by_base + return adapter_export.build_transformer_layer_adapter_weights( + model_chunks, + grouped_moe=self.is_moe, + language_layers_only=True, + ) def _wrap_mlp_lora( self, @@ -362,34 +336,18 @@ def _wrap_mlp_lora( rank: int, alpha: int, ) -> None: - from art.megatron.lora import wrap_dense_mlp + from art.megatron.lora import wrap_split_mlp_lora _require_dense_mlp(module) - wrap_dense_mlp( + wrap_split_mlp_lora( module.mlp, - adapter_model_prefix=adapter_model_prefix, + adapter_model_prefix=f"{adapter_model_prefix}.mlp", provider=provider, target_modules=target_modules, rank=rank, alpha=alpha, ) - def _add_mlp_adapter_weights( - self, - adapter_weights_by_base: dict[str, list[Any]], - *, - layer_prefix: str, - module: Any, - ) -> None: - from art.megatron.weights.adapter_export import add_dense_mlp_adapter_weights - - _require_dense_mlp(module) - add_dense_mlp_adapter_weights( - adapter_weights_by_base, - layer_prefix=layer_prefix, - mlp=module.mlp, - ) - def get_forward_kwargs(self, model: Any, **kwargs: Any) -> dict[str, Any]: unwrapped = model while hasattr(unwrapped, "module"): @@ -483,54 +441,27 @@ def _wrap_mlp_lora( rank: int, alpha: int, ) -> None: - from art.megatron.lora import ( - wrap_grouped_moe_experts_3d, - wrap_shared_experts_mlp, - ) + from art.megatron.lora import wrap_grouped_moe_experts, wrap_split_mlp_lora - wrap_grouped_moe_experts_3d( + wrap_grouped_moe_experts( _require_moe_experts(module), adapter_model_prefix=adapter_model_prefix, target_modules=target_modules, rank=rank, alpha=alpha, + fused_gate_up=True, ) shared_experts = getattr(module.mlp, "shared_experts", None) if shared_experts is not None: - wrap_shared_experts_mlp( + wrap_split_mlp_lora( shared_experts, - adapter_model_prefix=adapter_model_prefix, + adapter_model_prefix=f"{adapter_model_prefix}.mlp.shared_expert", provider=provider, target_modules=target_modules, rank=rank, alpha=alpha, ) - def _add_mlp_adapter_weights( - self, - adapter_weights_by_base: dict[str, list[Any]], - *, - layer_prefix: str, - module: Any, - ) -> None: - from art.megatron.weights.adapter_export import ( - add_grouped_moe_adapter_weights, - add_shared_experts_adapter_weights, - ) - - add_grouped_moe_adapter_weights( - adapter_weights_by_base, - layer_prefix=layer_prefix, - experts=_require_moe_experts(module), - ) - shared_experts = getattr(module.mlp, "shared_experts", None) - if shared_experts is not None: - add_shared_experts_adapter_weights( - adapter_weights_by_base, - layer_prefix=layer_prefix, - shared_experts=shared_experts, - ) - def compile_workaround_config( self, provider: Any, @@ -573,10 +504,6 @@ def _from_vllm_key(key: str) -> str: ) -def _is_lora_weight_key(key: str) -> bool: - return key.endswith((".lora_A.weight", ".lora_B.weight")) - - def _is_self_attn_q_proj_lora_b(key: str) -> bool: return key.endswith(".self_attn.q_proj.lora_B.weight") @@ -725,12 +652,16 @@ def _vllm_moe_config( return config -def _group_art_moe_tensors( +type _ExpertLoraGroups = dict[str, dict[int, dict[str, dict[str, torch.Tensor]]]] + + +def _group_expert_lora_tensors( tensors: dict[str, torch.Tensor], -) -> dict[str, dict[int, dict[str, dict[str, torch.Tensor]]]]: - grouped: dict[str, dict[int, dict[str, dict[str, torch.Tensor]]]] = {} + pattern: re.Pattern[str], +) -> _ExpertLoraGroups: + grouped: _ExpertLoraGroups = {} for key, tensor in tensors.items(): - match = _ART_MOE_EXPERT_KEY_RE.match(key) + match = pattern.match(key) if match is None: continue grouped.setdefault(match.group("prefix"), {}).setdefault( @@ -740,27 +671,57 @@ def _group_art_moe_tensors( return grouped +def _expert_lora_key(prefix: str, expert: int, module: str, lora_name: str) -> str: + return f"{prefix}.{expert}.{module}.{lora_name}.weight" + + +def _convert_remaining_lora_tensors( + transformed: dict[str, torch.Tensor], + tensors: dict[str, torch.Tensor], + *, + used_keys: set[str], + convert: Callable[ + [str, torch.Tensor], + tuple[str, torch.Tensor], + ], + reject_fused_moe: bool = False, +) -> None: + for key, tensor in tensors.items(): + if key in used_keys: + continue + if reject_fused_moe and _VLLM_MOE_KEY_RE.match(key) is not None: + raise RuntimeError( + "Mixed fused and per-expert Qwen3.5 vLLM MoE LoRA tensors" + ) + converted_key, converted = convert(key, tensor) + if converted_key in transformed: + raise RuntimeError( + f"Duplicate Qwen3.5 LoRA tensor after conversion: {converted_key}" + ) + transformed[converted_key] = converted + + def _to_vllm_lora_tensors( tensors: dict[str, torch.Tensor], *, adapter_config: dict[str, Any], ) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: - grouped = _group_art_moe_tensors(tensors) + grouped = _group_expert_lora_tensors(tensors, _ART_MOE_EXPERT_KEY_RE) has_shared_experts = _has_shared_expert_lora_tensors(tensors) transformed: dict[str, torch.Tensor] = {} + convert = lambda key, tensor: _to_vllm_lora_tensor( + key, + tensor, + adapter_config=adapter_config, + ) if not grouped: has_fused_experts = any(_VLLM_MOE_KEY_RE.match(key) for key in tensors) - for key, tensor in tensors.items(): - vllm_key, tensor = _to_vllm_lora_tensor( - key, - tensor, - adapter_config=adapter_config, - ) - if vllm_key in transformed: - raise RuntimeError( - f"Duplicate Qwen3.5 LoRA tensor after conversion: {vllm_key}" - ) - transformed[vllm_key] = tensor + _convert_remaining_lora_tensors( + transformed, + tensors, + used_keys=set(), + convert=convert, + ) return transformed, ( _vllm_moe_config( adapter_config, @@ -772,17 +733,21 @@ def _to_vllm_lora_tensors( used_keys: set[str] = set() for prefix, experts in grouped.items(): vllm_prefix = _to_vllm_key(prefix) - gate_up_a: list[torch.Tensor] = [] - gate_up_b: list[torch.Tensor] = [] - down_a: list[torch.Tensor] = [] - down_b: list[torch.Tensor] = [] + blocks = { + ("gate_up_proj", "lora_A"): [], + ("gate_up_proj", "lora_B"): [], + ("down_proj", "lora_A"): [], + ("down_proj", "lora_B"): [], + } for expert in sorted(experts): modules = experts[expert] try: - gate_up_a_tensor = modules["gate_up_proj"]["lora_A"] gate_up_b_tensor = modules["gate_up_proj"]["lora_B"] - d_a = modules["down_proj"]["lora_A"] - d_b = modules["down_proj"]["lora_B"] + expert_tensors = { + (module_name, lora_name): modules[module_name][lora_name] + for module_name in _ART_MOE_MODULES + for lora_name in _LORA_NAMES + } except KeyError as exc: raise RuntimeError( f"Incomplete Qwen3.5 MoE LoRA block for {prefix}.{expert}" @@ -792,34 +757,29 @@ def _to_vllm_lora_tensors( f"{prefix}.{expert}: gate/up lora_B rows " f"{gate_up_b_tensor.shape[0]} are not even" ) - gate_up_a.append(gate_up_a_tensor.contiguous()) - gate_up_b.append(gate_up_b_tensor.contiguous()) - down_a.append(d_a.contiguous()) - down_b.append(d_b.contiguous()) - for module_name in ("gate_up_proj", "down_proj"): - for lora_name in ("lora_A", "lora_B"): - used_keys.add(f"{prefix}.{expert}.{module_name}.{lora_name}.weight") + for slot, tensor in expert_tensors.items(): + blocks[slot].append(tensor.contiguous()) + used_keys.add(_expert_lora_key(prefix, expert, *slot)) transformed[f"{vllm_prefix}.base_layer.lora_A.weight"] = torch.cat( - gate_up_a, + blocks[("gate_up_proj", "lora_A")], dim=0, ).contiguous() transformed[f"{vllm_prefix}.base_layer.lora_B.weight"] = _pack_vllm_3d_lora_b( - gate_up_b + blocks[("gate_up_proj", "lora_B")] ) transformed[f"{vllm_prefix}.lora_A.weight"] = torch.cat( - down_a, + blocks[("down_proj", "lora_A")], dim=0, ).contiguous() - transformed[f"{vllm_prefix}.lora_B.weight"] = _pack_vllm_3d_lora_b(down_b) - for key, tensor in tensors.items(): - if key in used_keys: - continue - vllm_key, tensor = _to_vllm_lora_tensor( - key, - tensor, - adapter_config=adapter_config, + transformed[f"{vllm_prefix}.lora_B.weight"] = _pack_vllm_3d_lora_b( + blocks[("down_proj", "lora_B")] ) - transformed[vllm_key] = tensor + _convert_remaining_lora_tensors( + transformed, + tensors, + used_keys=used_keys, + convert=convert, + ) return transformed, _vllm_moe_config( adapter_config, has_shared_experts=has_shared_experts, @@ -831,15 +791,12 @@ def _from_vllm_lora_tensors( *, adapter_config: dict[str, Any], ) -> dict[str, torch.Tensor]: - expert_grouped: dict[str, dict[int, dict[str, dict[str, torch.Tensor]]]] = {} - for key, tensor in tensors.items(): - match = _VLLM_MOE_EXPERT_KEY_RE.match(key) - if match is None: - continue - expert_grouped.setdefault(match.group("prefix"), {}).setdefault( - int(match.group("expert")), - {}, - ).setdefault(match.group("module"), {})[match.group("lora")] = tensor + convert = lambda key, tensor: _from_vllm_lora_tensor( + key, + tensor, + adapter_config=adapter_config, + ) + expert_grouped = _group_expert_lora_tensors(tensors, _VLLM_MOE_EXPERT_KEY_RE) if expert_grouped: transformed: dict[str, torch.Tensor] = {} used_keys: set[str] = set() @@ -847,16 +804,17 @@ def _from_vllm_lora_tensors( art_prefix = _from_vllm_key(prefix) for expert, modules in experts.items(): try: - gate_a = modules["gate_proj"]["lora_A"] - gate_b = modules["gate_proj"]["lora_B"] - up_a = modules["up_proj"]["lora_A"] - up_b = modules["up_proj"]["lora_B"] - down_a = modules["down_proj"]["lora_A"] - down_b = modules["down_proj"]["lora_B"] + expert_tensors = { + (module_name, lora_name): modules[module_name][lora_name] + for module_name in _VLLM_EXPERT_MODULES + for lora_name in _LORA_NAMES + } except KeyError as exc: raise RuntimeError( f"Incomplete Qwen3.5 vLLM MoE LoRA block for {prefix}.{expert}" ) from exc + gate_a = expert_tensors[("gate_proj", "lora_A")] + up_a = expert_tensors[("up_proj", "lora_A")] if not torch.equal(gate_a, up_a): raise RuntimeError( "Qwen3.5 Megatron gate_up_proj requires gate/up " @@ -866,32 +824,29 @@ def _from_vllm_lora_tensors( _clone(gate_a) ) transformed[f"{art_prefix}.{expert}.gate_up_proj.lora_B.weight"] = ( - torch.cat([gate_b, up_b], dim=0).contiguous() + torch.cat( + [ + expert_tensors[("gate_proj", "lora_B")], + expert_tensors[("up_proj", "lora_B")], + ], + dim=0, + ).contiguous() ) transformed[f"{art_prefix}.{expert}.down_proj.lora_A.weight"] = _clone( - down_a + expert_tensors[("down_proj", "lora_A")] ) transformed[f"{art_prefix}.{expert}.down_proj.lora_B.weight"] = _clone( - down_b - ) - for module_name in ("gate_proj", "up_proj", "down_proj"): - for lora_name in ("lora_A", "lora_B"): - used_keys.add( - f"{prefix}.{expert}.{module_name}.{lora_name}.weight" - ) - for key, tensor in tensors.items(): - if key in used_keys: - continue - if _VLLM_MOE_KEY_RE.match(key) is not None: - raise RuntimeError( - "Mixed fused and per-expert Qwen3.5 vLLM MoE LoRA tensors" + expert_tensors[("down_proj", "lora_B")] ) - art_key, tensor = _from_vllm_lora_tensor( - key, - tensor, - adapter_config=adapter_config, - ) - transformed[art_key] = tensor + for slot in expert_tensors: + used_keys.add(_expert_lora_key(prefix, expert, *slot)) + _convert_remaining_lora_tensors( + transformed, + tensors, + used_keys=used_keys, + convert=convert, + reject_fused_moe=True, + ) return transformed grouped: dict[str, dict[str, torch.Tensor]] = {} @@ -905,13 +860,12 @@ def _from_vllm_lora_tensors( grouped.setdefault(match.group("prefix"), {})[slot] = tensor if not grouped: transformed: dict[str, torch.Tensor] = {} - for key, tensor in tensors.items(): - art_key, tensor = _from_vllm_lora_tensor( - key, - tensor, - adapter_config=adapter_config, - ) - transformed[art_key] = tensor + _convert_remaining_lora_tensors( + transformed, + tensors, + used_keys=set(), + convert=convert, + ) return transformed rank = int(adapter_config["r"]) @@ -970,15 +924,12 @@ def _from_vllm_lora_tensors( f"{prefix}.lora_B.weight", } ) - for key, tensor in tensors.items(): - if key in used_keys: - continue - art_key, tensor = _from_vllm_lora_tensor( - key, - tensor, - adapter_config=adapter_config, - ) - transformed[art_key] = tensor + _convert_remaining_lora_tensors( + transformed, + tensors, + used_keys=used_keys, + convert=convert, + ) return transformed diff --git a/src/art/megatron/model_support/spec.py b/src/art/megatron/model_support/spec.py index e9d0e8b4b..4e9844bbc 100644 --- a/src/art/megatron/model_support/spec.py +++ b/src/art/megatron/model_support/spec.py @@ -46,7 +46,7 @@ class ArchitectureReport(BaseModel): unresolved_risks: list[str] = Field(default_factory=list) -class SharedPrefixModelStateContext(BaseModel): +class PrefixTreeModelStateContext(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) group_ids: Any @@ -104,6 +104,7 @@ class ModelSupportSpec(BaseModel): class ModelSupportHandler(Protocol): key: str is_moe: bool + build_gdn_execution_spec: bool cp_supported: bool native_vllm_lora_status: NativeVllmLoraStatus @@ -153,9 +154,9 @@ def vllm_server_args(self) -> dict[str, object]: ... def install_preprocess_patch(self, model_chunks: Sequence[Any]) -> None: ... - def build_shared_prefix_model_state( + def build_prefix_tree_model_state( self, - context: SharedPrefixModelStateContext, + context: PrefixTreeModelStateContext, ) -> dict[str, Any]: ... def correctness_precision(self) -> Literal["bf16", "fp32"]: ... diff --git a/src/art/megatron/prefix_tree.py b/src/art/megatron/prefix_tree.py new file mode 100644 index 000000000..4683ca79f --- /dev/null +++ b/src/art/megatron/prefix_tree.py @@ -0,0 +1,234 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True, slots=True) +class PrefixTreeSegment: + group_id: int + parent_id: int + start: int + end: int + family_index: int + ancestors: tuple[int, ...] + + @property + def depth(self) -> int: + return len(self.ancestors) + + +@dataclass(frozen=True, slots=True) +class PrefixTreeRow: + row_index: int + valid_tokens: int + segments: tuple[PrefixTreeSegment, ...] + + +def parse_prefix_tree( + *, + group_ids: torch.Tensor, + parent_ids: torch.Tensor, + ignore_padding_group_id: int = -1, +) -> tuple[PrefixTreeRow, ...]: + if group_ids.shape != parent_ids.shape: + raise RuntimeError( + "group_ids and parent_ids must share shape, got " + f"{tuple(group_ids.shape)} vs {tuple(parent_ids.shape)}" + ) + if group_ids.ndim != 2: + raise RuntimeError( + "group_ids and parent_ids must be rank-2 packed tensors, got " + f"{group_ids.ndim}" + ) + return tuple( + parse_prefix_tree_row( + group_ids=group_ids[row_index], + parent_ids=parent_ids[row_index], + row_index=row_index, + ignore_padding_group_id=ignore_padding_group_id, + ) + for row_index in range(int(group_ids.shape[0])) + ) + + +def parse_prefix_tree_row( + *, + group_ids: torch.Tensor, + parent_ids: torch.Tensor, + row_index: int = 0, + ignore_padding_group_id: int = -1, +) -> PrefixTreeRow: + if group_ids.shape != parent_ids.shape: + raise RuntimeError( + "group_ids and parent_ids must share shape, got " + f"{tuple(group_ids.shape)} vs {tuple(parent_ids.shape)}" + ) + if group_ids.ndim != 1: + raise RuntimeError( + f"group_ids and parent_ids must be rank-1 row tensors, got {group_ids.ndim}" + ) + + valid_tokens = _valid_length( + group_ids, + parent_ids, + ignore_padding_group_id=ignore_padding_group_id, + ) + if valid_tokens == 0: + return PrefixTreeRow(row_index=row_index, valid_tokens=0, segments=()) + + runs = _scan_runs(group_ids[:valid_tokens], parent_ids[:valid_tokens]) + first_segment_by_group: dict[int, PrefixTreeSegment] = {} + family_by_group: dict[int, int] = {} + ancestors_by_group: dict[int, tuple[int, ...]] = {} + segments: list[PrefixTreeSegment] = [] + next_family_index = 0 + + seen_groups: set[int] = set() + repeated_groups: dict[int, int] = {} + for _start, _end, group_id, _parent_id in runs: + if group_id in seen_groups and group_id != ignore_padding_group_id: + repeated_groups[group_id] = repeated_groups.get(group_id, 1) + 1 + seen_groups.add(group_id) + if repeated_groups: + raise RuntimeError( + "Prefix-tree metadata requires contiguous group runs per row, " + f"found repeats in row {row_index}: {repeated_groups}" + ) + + for start, end, group_id, parent_id in runs: + is_root = group_id == parent_id or ( + start == 0 and parent_id == ignore_padding_group_id + ) + if is_root: + family_index = next_family_index + next_family_index += 1 + ancestors: tuple[int, ...] = () + else: + parent_segment = first_segment_by_group.get(parent_id) + if parent_segment is None: + raise RuntimeError( + "Prefix-tree run points to a missing parent run: " + f"row={row_index}, group_id={group_id}, parent_id={parent_id}" + ) + if int(parent_segment.end) > int(start): + raise RuntimeError( + "Prefix-tree parent run must end before its child starts: " + f"row={row_index}, group_id={group_id}, parent_id={parent_id}" + ) + family_index = family_by_group[parent_id] + ancestors = (*ancestors_by_group[parent_id], parent_id) + + segment = PrefixTreeSegment( + group_id=group_id, + parent_id=parent_id, + start=start, + end=end, + family_index=family_index, + ancestors=ancestors, + ) + first_segment_by_group[group_id] = segment + family_by_group[group_id] = family_index + ancestors_by_group[group_id] = ancestors + segments.append(segment) + + return PrefixTreeRow( + row_index=row_index, + valid_tokens=valid_tokens, + segments=tuple(segments), + ) + + +def _valid_length( + group_ids: torch.Tensor, + parent_ids: torch.Tensor, + *, + ignore_padding_group_id: int, +) -> int: + valid_mask = group_ids != ignore_padding_group_id + valid_count = int(valid_mask.sum().item()) + if valid_count == 0: + return 0 + if not bool(valid_mask[:valid_count].all().item()): + raise RuntimeError("Padding tokens must be a contiguous tail") + return _infer_terminal_padding_length( + group_ids[:valid_count], + parent_ids[:valid_count], + ) + + +def _infer_terminal_padding_length( + group_row: torch.Tensor, + parent_row: torch.Tensor, +) -> int: + if group_row.numel() == 0: + return 0 + runs = _scan_runs(group_row, parent_row) + if len(runs) < 2: + return int(group_row.numel()) + last_start, _last_end, last_group_id, last_parent_id = runs[-1] + if last_parent_id >= 0: + return int(group_row.numel()) + terminal_pair = (last_group_id, last_parent_id) + if any( + (group_id, parent_id) == terminal_pair + for _start, _end, group_id, parent_id in runs[:-1] + ): + return last_start + return int(group_row.numel()) + + +def _scan_runs( + group_row: torch.Tensor, + parent_row: torch.Tensor, +) -> list[tuple[int, int, int, int]]: + length = int(group_row.numel()) + if length == 0: + return [] + + group_changes = group_row[1:] != group_row[:-1] + parent_changes = parent_row[1:] != parent_row[:-1] + inconsistent_parent = torch.nonzero( + torch.logical_not(group_changes) & parent_changes, + as_tuple=False, + ).flatten() + if int(inconsistent_parent.numel()) > 0: + mismatch_index = int(inconsistent_parent[0].item()) + 1 + prior_boundaries = torch.nonzero( + group_changes[: mismatch_index - 1], + as_tuple=False, + ).flatten() + start = ( + 0 + if int(prior_boundaries.numel()) == 0 + else int(prior_boundaries[-1].item()) + 1 + ) + group_id = int(group_row[start].item()) + raise RuntimeError( + "Found one group run with inconsistent parent ids: " + f"group_id={group_id}, start={start}, end={mismatch_index}" + ) + + run_starts = torch.cat( + ( + torch.zeros(1, dtype=torch.int64, device=group_row.device), + torch.nonzero(group_changes, as_tuple=False).flatten() + 1, + ) + ) + run_ends = torch.cat( + ( + run_starts[1:], + torch.tensor([length], dtype=torch.int64, device=group_row.device), + ) + ) + starts = run_starts.to(device="cpu").tolist() + ends = run_ends.to(device="cpu").tolist() + group_ids = group_row.index_select(0, run_starts).to(device="cpu").tolist() + parent_ids = parent_row.index_select(0, run_starts).to(device="cpu").tolist() + return [ + (int(start), int(end), int(group_id), int(parent_id)) + for start, end, group_id, parent_id in zip( + starts, ends, group_ids, parent_ids, strict=True + ) + ] diff --git a/src/art/megatron/prefix_tree_packing.py b/src/art/megatron/prefix_tree_packing.py new file mode 100644 index 000000000..20b4a4e3a --- /dev/null +++ b/src/art/megatron/prefix_tree_packing.py @@ -0,0 +1,321 @@ +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class PrefixTreePack: + tokens: torch.Tensor + group_ids: torch.Tensor + parent_ids: torch.Tensor + position_ids: torch.Tensor + positions_by_sequence: tuple[torch.Tensor, ...] + + +@dataclass(frozen=True) +class _PrefixSegment: + sequence_indices: tuple[int, ...] + start: int + end: int + group_id: int + parent_id: int + + +def prefix_tree_pack( + sequences: Iterable[torch.Tensor], + *, + max_depth: int, + shareable_lengths: Iterable[int] | None = None, +) -> PrefixTreePack: + """Pack token sequences by storing prefix trees once. + + This is the small packing step that lets `TrainerRank.dp_rank_forward()` run one + model pass over a compact prefix tree instead of replaying the same prompt + tokens for every request. Think of each input sequence as a path through a + tree: when several paths start with the same tokens, this function writes + that shared segment once, then writes each branch after it. + + Args: + sequences: 1-D token tensors to pack. + max_depth: How many nested prefix-tree levels to emit. `0` disables + tree sharing and writes each sequence as its own root segment. `1` + shares the first common segment in each branch; larger values allow + branches to contain shared sub-branches. + + Returns: + `tokens` is the compact model input, shaped `[1, packed_length]`. + `group_ids` and `parent_ids` describe the prefix tree to prefix-tree + attention. Positions in the same emitted segment share a group, and each + group points at the parent segment it continues from. Root groups point + to themselves. + `position_ids` keeps each token's original sequence position for + positional embeddings/rotary attention. + `positions_by_sequence` is the reverse index used after the model call + to unpack logits, logprobs, or hidden states back into one tensor per + original request. + + The implementation is a tiny radix-tree walk. It finds the longest prefix + shared by the active sequences, emits that segment once, then partitions the + remaining sequences by their next token while preserving first-seen order. + Single sequences, empty branches, and branches past `max_depth` are emitted + as ordinary unshared tails. + """ + if max_depth < 0: + raise ValueError("max_depth must be >= 0") + + tensors = tuple(_sequence_tensor(sequence) for sequence in sequences) + if not tensors: + return _empty_pack() + share_limits = _shareable_lengths(tensors, shareable_lengths) + + device = tensors[0].device + rows = tuple(tensor.detach().cpu().tolist() for tensor in tensors) + segments = _prefix_segments( + rows, max_depth=max_depth, shareable_lengths=share_limits + ) + if not segments: + return _empty_pack(len(tensors), device=device) + + token_chunks: list[torch.Tensor] = [] + group_chunks: list[torch.Tensor] = [] + parent_chunks: list[torch.Tensor] = [] + position_chunks: list[torch.Tensor] = [] + positions_by_sequence: list[list[torch.Tensor]] = [[] for _ in tensors] + cursor = 0 + + for planned in segments: + segment = tensors[planned.sequence_indices[0]][planned.start : planned.end] + packed_positions = torch.arange(cursor, cursor + len(segment), device=device) + token_chunks.append(segment) + group_chunks.append(torch.full_like(segment, planned.group_id)) + parent_chunks.append(torch.full_like(segment, planned.parent_id)) + position_chunks.append(torch.arange(planned.start, planned.end, device=device)) + for sequence_index in planned.sequence_indices: + positions_by_sequence[sequence_index].append(packed_positions) + cursor += len(segment) + + return PrefixTreePack( + tokens=torch.cat(token_chunks).unsqueeze(0), + group_ids=torch.cat(group_chunks).unsqueeze(0), + parent_ids=torch.cat(parent_chunks).unsqueeze(0), + position_ids=torch.cat(position_chunks).unsqueeze(0), + positions_by_sequence=tuple( + torch.cat(chunks) + if chunks + else torch.empty(0, dtype=torch.long, device=device) + for chunks in positions_by_sequence + ), + ) + + +def estimate_prefix_tree_packed_tokens( + sequences: Iterable[torch.Tensor], + *, + max_depth: int, + shareable_lengths: Iterable[int] | None = None, +) -> int | None: + """Return the exact packed token count without building a packed batch. + + The estimator intentionally only handles CPU tensors. For CUDA tensors, many + tiny prefix probes would launch many tiny kernels, so callers should fall + back to full packing instead. + """ + if max_depth < 0: + raise ValueError("max_depth must be >= 0") + + tensors: list[torch.Tensor] = [] + rows: list[list[int]] = [] + for sequence in sequences: + tensor = _sequence_tensor(sequence) + if tensor.device.type != "cpu": + return None + tensors.append(tensor) + rows.append(tensor.tolist()) + share_limits = _shareable_lengths(tuple(tensors), shareable_lengths) + + return sum( + segment.end - segment.start + for segment in _prefix_segments( + tuple(rows), + max_depth=max_depth, + shareable_lengths=share_limits, + ) + ) + + +def _local_position_pairs( + local_global_positions: torch.Tensor, + item_positions: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + flat = local_global_positions.reshape(-1).to(device=item_positions.device) + local_positions = torch.nonzero(flat >= 0, as_tuple=False).reshape(-1) + global_positions = flat.index_select(0, local_positions) + source_offsets, local_offsets = _matching_offsets(item_positions, global_positions) + return ( + local_positions.index_select(0, local_offsets).to("cpu"), + source_offsets.to("cpu"), + ) + + +def _prefix_segments( + rows: tuple[list[int], ...], + *, + max_depth: int, + shareable_lengths: tuple[int, ...], +) -> tuple[_PrefixSegment, ...]: + lengths = tuple(len(row) for row in rows) + segments: list[_PrefixSegment] = [] + next_group_id = 1 + + def emit( + indices: tuple[int, ...], + start: int, + end: int, + parent_group_id: int | None, + ) -> int: + nonlocal next_group_id + group_id = next_group_id + next_group_id += 1 + segments.append( + _PrefixSegment( + sequence_indices=indices, + start=start, + end=end, + group_id=group_id, + parent_id=group_id if parent_group_id is None else parent_group_id, + ) + ) + return group_id + + def shared_end(indices: tuple[int, ...], start: int) -> int: + end = min(min(lengths[index], shareable_lengths[index]) for index in indices) + low = high = rows[indices[0]] + for index in indices[1:]: + row = rows[index] + if row < low: + low = row + elif row > high: + high = row + while start < end: + if low[start] != high[start]: + break + start += 1 + return start + + def branch_groups(indices: tuple[int, ...], start: int) -> list[tuple[int, ...]]: + groups: dict[int, list[int]] = {} + order: list[int] = [] + for index in indices: + token = rows[index][start] + if token not in groups: + groups[token] = [] + order.append(token) + groups[token].append(index) + return [tuple(groups[token]) for token in order] + + def walk( + indices: tuple[int, ...], + start: int, + parent_group_id: int | None, + depth: int, + ) -> None: + active = tuple(index for index in indices if lengths[index] > start) + if not active: + return + shareable = tuple(index for index in active if shareable_lengths[index] > start) + for index in active: + if index not in shareable: + emit((index,), start, lengths[index], parent_group_id) + if not shareable: + return + if ( + max_depth == 0 + or len(shareable) == 1 + or (parent_group_id is not None and depth >= max_depth) + ): + for index in shareable: + emit((index,), start, lengths[index], parent_group_id) + return + + end = shared_end(shareable, start) + if end > start: + walk( + shareable, + end, + emit(shareable, start, end, parent_group_id), + depth + 1, + ) + return + for group in branch_groups(shareable, start): + walk(group, start, parent_group_id, depth) + + walk(tuple(range(len(rows))), 0, None, 0) + return tuple(segments) + + +def _empty_pack( + sequence_count: int = 0, + *, + device: torch.device | None = None, +) -> PrefixTreePack: + flat = torch.empty(0, dtype=torch.long, device=device) + row = flat.unsqueeze(0) + return PrefixTreePack( + tokens=row, + group_ids=row, + parent_ids=row, + position_ids=row, + positions_by_sequence=tuple(flat for _ in range(sequence_count)), + ) + + +def _sequence_tensor(tensor: torch.Tensor) -> torch.Tensor: + if tensor.ndim != 1: + raise ValueError( + f"prefix_tree_pack expects 1-D tensors, got {tuple(tensor.shape)}" + ) + return tensor.detach().to(dtype=torch.long).contiguous() + + +def _shareable_lengths( + tensors: tuple[torch.Tensor, ...], + lengths: Iterable[int] | None, +) -> tuple[int, ...]: + if lengths is None: + return tuple(int(tensor.numel()) for tensor in tensors) + values = tuple(int(length) for length in lengths) + if len(values) != len(tensors): + raise ValueError( + "shareable_lengths must have one entry per sequence, " + f"got {len(values)} for {len(tensors)} sequences" + ) + return tuple( + max(0, min(value, int(tensor.numel()))) + for value, tensor in zip(values, tensors, strict=True) + ) + + +def _matching_offsets( + positions: torch.Tensor, + chunk_rows: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + if int(positions.numel()) == 0 or int(chunk_rows.numel()) == 0: + empty = torch.empty(0, dtype=torch.long, device=positions.device) + return empty, empty + sorted_rows, order = chunk_rows.sort() + indices = torch.searchsorted(sorted_rows, positions) + in_bounds = indices < int(sorted_rows.numel()) + source_offsets = torch.arange( + int(positions.numel()), + device=positions.device, + dtype=torch.long, + )[in_bounds] + found = indices[in_bounds] + keep = sorted_rows.index_select(0, found) == positions.index_select( + 0, + source_offsets, + ) + return source_offsets[keep], order.index_select(0, found[keep]) diff --git a/src/art/megatron/shared_prefix_state.py b/src/art/megatron/prefix_tree_state.py similarity index 57% rename from src/art/megatron/shared_prefix_state.py rename to src/art/megatron/prefix_tree_state.py index 97dbb6d93..cb8c2af3b 100644 --- a/src/art/megatron/shared_prefix_state.py +++ b/src/art/megatron/prefix_tree_state.py @@ -1,7 +1,8 @@ -"""Shared-prefix packed-sequence state for ART attention and GDN integration.""" +"""Prefix-tree packed-sequence state for ART attention and GDN integration.""" from __future__ import annotations +from dataclasses import replace import gc from typing import Any @@ -10,8 +11,11 @@ from torch import Tensor from torch.nn.attention.flex_attention import BlockMask -from art.megatron.context_parallel.block_mask import build_block_mask -from art.megatron.context_parallel.builder import build_shared_prefix_attention_spec +from art.megatron.context_parallel.block_mask import ( + build_block_mask_from_context, + prepare_block_mask_context, +) +from art.megatron.context_parallel.builder import build_prefix_tree_attention_spec from art.megatron.context_parallel.layout_index import TokenLayoutIndex from art.megatron.context_parallel.types import ( AttnMaskKind, @@ -21,21 +25,22 @@ TokenRange, ) from art.megatron.flex_attn.attention import ( - SharedPrefixAttentionState as FlexSharedPrefixAttentionState, + PrefixTreeAttentionState as FlexPrefixTreeAttentionState, ) from art.megatron.flex_attn.compiled import flash_sparse_block_size_for_head_dim -from art.megatron.gdn.gdn_shared_prefix import ( +from art.megatron.gdn.gdn_prefix_tree import ( GdnPackedExecutionSpec, + GdnPlannerConfig, GdnRankExecutionPlan, build_gdn_rank_execution_plan, move_gdn_rank_execution_plan_to_device, - parse_gdn_shared_prefix_segments, + parse_gdn_prefix_tree_segments, ) -from art.megatron.model_support.spec import SharedPrefixModelStateContext +from art.megatron.model_support.spec import PrefixTreeModelStateContext -class SharedPrefixAttentionState(FlexSharedPrefixAttentionState): - """Shared-prefix sparsity and optional GDN execution metadata.""" +class PrefixTreeAttentionState(FlexPrefixTreeAttentionState): + """Prefix-tree sparsity and optional GDN execution metadata.""" group_ids: Tensor parent_ids: Tensor @@ -53,7 +58,7 @@ class SharedPrefixAttentionState(FlexSharedPrefixAttentionState): gdn_active_module: Any | None = None -def create_shared_prefix_state( +def create_prefix_tree_state( group_ids: Tensor, parent_ids: Tensor, *, @@ -65,18 +70,19 @@ def create_shared_prefix_state( attention_token_layout_index: TokenLayoutIndex | None = None, attention_head_dim: int | None = None, attention_value_head_dim: int | None = None, -) -> SharedPrefixAttentionState: - """Build shared-prefix attention mask state plus optional reusable GDN plan.""" + gdn_planner_config: GdnPlannerConfig | None = None, +) -> PrefixTreeAttentionState: + """Build prefix-tree attention mask state plus optional reusable GDN plan.""" device = group_ids.device if target_device is None else torch.device(target_device) group_ids_cpu = _metadata_cpu(group_ids) parent_ids_cpu = _metadata_cpu(parent_ids) - input_pos_cpu = _metadata_cpu(input_pos) if input_pos is not None else None - block_size = _shared_prefix_block_size( + input_pos_cpu = None if input_pos is None else _metadata_cpu(input_pos) + block_size = _prefix_tree_block_size( device, attention_head_dim=attention_head_dim, attention_value_head_dim=attention_value_head_dim, ) - block_mask = _build_sparse_shared_prefix_block_mask( + block_mask = _build_sparse_prefix_tree_block_mask( group_ids_cpu=group_ids_cpu, parent_ids_cpu=parent_ids_cpu, input_pos_cpu=input_pos_cpu, @@ -85,7 +91,7 @@ def create_shared_prefix_state( block_size=block_size, ) sliding_block_masks = { - window: _build_sparse_shared_prefix_block_mask( + window: _build_sparse_prefix_tree_block_mask( group_ids_cpu=group_ids_cpu, parent_ids_cpu=parent_ids_cpu, input_pos_cpu=input_pos_cpu, @@ -95,16 +101,16 @@ def create_shared_prefix_state( ) for window in tuple(dict.fromkeys(int(window) for window in sliding_windows)) } - cp_rank, cp_size, cp_group = _gdn_cp_rank_size_group() - gdn_execution_spec = _build_gdn_execution_spec_once( - group_ids_cpu, - parent_ids_cpu, - build=build_gdn_execution_spec, - cp_rank=cp_rank, - cp_size=cp_size, - cp_group=cp_group, + cp_rank, cp_size = _gdn_cp_rank_size() + gdn_execution_spec = ( + parse_gdn_prefix_tree_segments( + group_ids_cpu, + parent_ids_cpu, + ) + if build_gdn_execution_spec + else None ) - return SharedPrefixAttentionState( + return PrefixTreeAttentionState( block_mask=block_mask, sliding_block_masks=sliding_block_masks, group_ids=group_ids_cpu, @@ -125,8 +131,8 @@ def create_shared_prefix_state( device=device, cp_rank=cp_rank, cp_size=cp_size, - cp_group=cp_group, attention_token_layout_index=attention_token_layout_index, + planner_config=gdn_planner_config, ), ) @@ -145,8 +151,8 @@ def _build_model_state_once( if model_support_handler is None: return {} return dict( - model_support_handler.build_shared_prefix_model_state( - SharedPrefixModelStateContext( + model_support_handler.build_prefix_tree_model_state( + PrefixTreeModelStateContext( input_pos=input_pos, group_ids=group_ids, parent_ids=parent_ids, @@ -166,7 +172,7 @@ def _metadata_cpu(tensor: Tensor) -> Tensor: return tensor.contiguous() -def _build_sparse_shared_prefix_block_mask( +def _build_sparse_prefix_tree_block_mask( *, group_ids_cpu: Tensor, parent_ids_cpu: Tensor, @@ -175,63 +181,117 @@ def _build_sparse_shared_prefix_block_mask( device: torch.device, block_size: tuple[int, int], ): - batch_spec = build_shared_prefix_attention_spec( + batch_spec = build_prefix_tree_attention_spec( group_ids=group_ids_cpu, parent_ids=parent_ids_cpu, ) - row_spec = batch_spec.rows[0] seq_len = int(group_ids_cpu.shape[1]) - slices = _full_row_slices_with_padding( - row_slices=row_spec.slices, - valid_tokens=int(row_spec.valid_tokens), - seq_len=seq_len, - ) - if not slices: - return _empty_block_mask(seq_len=seq_len, block_size=block_size, device=device) - return build_block_mask( - FlexMaskSpec( - q_len=seq_len, - k_len=seq_len, - block_size=block_size, - slices=slices, - exact_mask=ExactMaskMetadata( - q_token_indices=torch.arange(seq_len, dtype=torch.int64), - k_token_indices=torch.arange(seq_len, dtype=torch.int64), - cache_key=( - f"identity:{seq_len}" - if sliding_window is None - else f"identity:{seq_len}:sliding:{int(sliding_window)}" + row_masks = [] + token_indices = torch.arange(seq_len, dtype=torch.int64) + for row_spec in batch_spec.rows: + row_index = int(row_spec.row_index) + slices = tuple(replace(slice_, row_index=0) for slice_ in row_spec.slices) + if int(row_spec.valid_tokens) < seq_len: + padding_range = TokenRange(start=int(row_spec.valid_tokens), end=seq_len) + slices = ( + *slices, + AttnSlice( + q_range=padding_range, + k_range=padding_range, + mask_kind=AttnMaskKind.CAUSAL, + row_index=0, + family_index=None, ), - ), - ), - group_ids=group_ids_cpu[0], - parent_ids=parent_ids_cpu[0], - input_pos=None if input_pos_cpu is None else input_pos_cpu[0], - sliding_window=sliding_window, - device=device, + ) + if not slices: + row_masks.append( + _empty_block_mask(seq_len=seq_len, block_size=block_size, device=device) + ) + continue + row_masks.append( + build_block_mask_from_context( + FlexMaskSpec( + q_len=seq_len, + k_len=seq_len, + block_size=block_size, + slices=slices, + exact_mask=ExactMaskMetadata( + q_token_indices=token_indices, + k_token_indices=token_indices, + cache_key=( + f"identity:{seq_len}" + if sliding_window is None + else f"identity:{seq_len}:sliding:{int(sliding_window)}" + ), + ), + ), + context=prepare_block_mask_context( + group_ids=group_ids_cpu[row_index], + parent_ids=parent_ids_cpu[row_index], + input_pos=None + if input_pos_cpu is None + else input_pos_cpu[row_index], + ), + sliding_window=sliding_window, + device=device, + ) + ) + if not row_masks: + return _empty_block_mask(seq_len=seq_len, block_size=block_size, device=device) + return _stack_row_block_masks( + row_masks, + seq_len=seq_len, + block_size=block_size, ) -def _full_row_slices_with_padding( +def _stack_optional_block_tensors( + masks: list[BlockMask], + name: str, +) -> Tensor | None: + tensors = [getattr(mask, name) for mask in masks] + if any(tensor is None for tensor in tensors): + return None + return torch.cat(tensors, dim=0) + + +def _stack_row_block_masks( + masks: list[BlockMask], *, - row_slices: tuple[AttnSlice, ...], - valid_tokens: int, seq_len: int, -) -> tuple[AttnSlice, ...]: - if valid_tokens >= seq_len: - return row_slices - padding_range = TokenRange(start=int(valid_tokens), end=int(seq_len)) - if padding_range.is_empty(): - return row_slices - return ( - *row_slices, - AttnSlice( - q_range=padding_range, - k_range=padding_range, - mask_kind=AttnMaskKind.CAUSAL, - row_index=0, - family_index=None, - ), + block_size: tuple[int, int], +) -> BlockMask: + if len(masks) == 1: + return masks[0] + row_mask_mods = tuple(mask.mask_mod for mask in masks) + + def mask_mod( + batch_idx: Tensor, + head_idx: Tensor, + query_idx: Tensor, + kv_idx: Tensor, + ) -> Tensor: + result = torch.zeros_like(query_idx, dtype=torch.bool) + for row_index, row_mask_mod in enumerate(row_mask_mods): + result = torch.where( + batch_idx == row_index, + row_mask_mod(batch_idx, head_idx, query_idx, kv_idx), + result, + ) + return result + + return BlockMask( + seq_lengths=(int(seq_len), int(seq_len)), + kv_num_blocks=torch.cat([mask.kv_num_blocks for mask in masks], dim=0), + kv_indices=torch.cat([mask.kv_indices for mask in masks], dim=0), + full_kv_num_blocks=_stack_optional_block_tensors(masks, "full_kv_num_blocks"), + full_kv_indices=_stack_optional_block_tensors(masks, "full_kv_indices"), + q_num_blocks=_stack_optional_block_tensors(masks, "q_num_blocks"), + q_indices=_stack_optional_block_tensors(masks, "q_indices"), + full_q_num_blocks=_stack_optional_block_tensors(masks, "full_q_num_blocks"), + full_q_indices=_stack_optional_block_tensors(masks, "full_q_indices"), + BLOCK_SIZE=block_size, + mask_mod=mask_mod, ) @@ -271,7 +331,7 @@ def _false_mask( return torch.zeros_like(query_idx, dtype=torch.bool) -def _shared_prefix_block_size( +def _prefix_tree_block_size( device: torch.device, *, attention_head_dim: int | None, @@ -290,45 +350,18 @@ def _shared_prefix_block_size( ) -def _build_gdn_execution_spec_once( - group_ids: Tensor, - parent_ids: Tensor, - *, - build: bool, - cp_rank: int, - cp_size: int, - cp_group: Any | None, -) -> GdnPackedExecutionSpec | None: - if not build: - return None - if cp_size == 1: - return parse_gdn_shared_prefix_segments( - group_ids, parent_ids, min_completions_per_family=0 - ) - if ( - not torch.distributed.is_available() or not torch.distributed.is_initialized() # ty: ignore[possibly-missing-attribute] - ): - return parse_gdn_shared_prefix_segments( - group_ids, parent_ids, min_completions_per_family=0 - ) - return parse_gdn_shared_prefix_segments( - group_ids, parent_ids, min_completions_per_family=0 - ) - - def _build_gdn_execution_plan_once( spec: GdnPackedExecutionSpec | None, *, device: torch.device, cp_rank: int, cp_size: int, - cp_group: Any | None, attention_token_layout_index: TokenLayoutIndex | None, + planner_config: GdnPlannerConfig | None, ) -> GdnRankExecutionPlan | None: if spec is None: return None planner_device = torch.device("cpu") if device.type == "cuda" else device - del cp_group gc_was_enabled = gc.isenabled() if gc_was_enabled: gc.disable() @@ -339,6 +372,7 @@ def _build_gdn_execution_plan_once( cp_rank=cp_rank, cp_size=cp_size, attention_token_layout_index=attention_token_layout_index, + planner_config=planner_config, ) finally: if gc_was_enabled: @@ -346,7 +380,7 @@ def _build_gdn_execution_plan_once( return move_gdn_rank_execution_plan_to_device(plan, device) -def _gdn_cp_rank_size_group() -> tuple[int, int, Any | None]: +def _gdn_cp_rank_size() -> tuple[int, int]: try: from megatron.core import parallel_state as ps @@ -354,8 +388,7 @@ def _gdn_cp_rank_size_group() -> tuple[int, int, Any | None]: return ( int(ps.get_context_parallel_rank()), int(ps.get_context_parallel_world_size()), - ps.get_context_parallel_group(), ) except Exception: pass - return 0, 1, None + return 0, 1 diff --git a/src/art/megatron/routing_replay.py b/src/art/megatron/routing_replay.py index 69ce534bc..a34a65766 100644 --- a/src/art/megatron/routing_replay.py +++ b/src/art/megatron/routing_replay.py @@ -44,7 +44,7 @@ def _branch_rows_without_future_scored_tokens( vLLM returns routes for forwards it actually ran. During generation it does not run a forward for the final generated token, because no later token is - sampled from that row. ART shared-prefix packing may still leave a masked + sampled from that row. ART prefix-tree packing may still leave a masked suffix token in the same branch, so a physical group boundary is not enough to identify rows that can safely use synthetic replay routes. """ diff --git a/src/art/megatron/setup.sh b/src/art/megatron/setup.sh index 1e9c60eb1..3a325dd3a 100755 --- a/src/art/megatron/setup.sh +++ b/src/art/megatron/setup.sh @@ -35,3 +35,7 @@ if [ -x "${HOME}/.local/bin/uv" ]; then uv_bin="${HOME}/.local/bin/uv" fi "${uv_bin}" sync --extra megatron --frozen --active + +if [ "${INSTALL_VLLM_RUNTIME:-true}" = "true" ]; then + "${uv_bin}" sync --project vllm_runtime --frozen --no-dev +fi diff --git a/src/art/megatron/training/finalize_grads.py b/src/art/megatron/training/finalize_grads.py index cde0e7b06..e00cd8218 100644 --- a/src/art/megatron/training/finalize_grads.py +++ b/src/art/megatron/training/finalize_grads.py @@ -1,4 +1,3 @@ -from collections import defaultdict from collections.abc import Iterable from typing import Any, Literal, cast @@ -16,7 +15,6 @@ GRAD_SYNC_OP_NONE: GradSyncOp = "none" GRAD_SYNC_OP_SUM: GradSyncOp = "sum" GRAD_SYNC_OP_AVG: GradSyncOp = "avg" -VALID_DOMAINS = (TP_DEFAULT_GRAD_SYNC_DOMAIN, EXPERT_TP_GRAD_SYNC_DOMAIN) VALID_SYNC_OPS = (GRAD_SYNC_OP_NONE, GRAD_SYNC_OP_SUM, GRAD_SYNC_OP_AVG) @@ -28,6 +26,8 @@ def _iter_named_trainable_parameters( for name, param in model_chunk.named_parameters(): if not param.requires_grad: continue + if getattr(param, "_art_dynamic_lora_slot", False): + continue param_id = id(param) if param_id in seen: continue @@ -60,6 +60,48 @@ def _resolve_reduce_op(op: GradSyncOp) -> Any: raise RuntimeError(f"Unknown grad sync op: {op}") +def tensor_parallel_grad_sync( + param: torch.nn.Parameter, + *, + name: str, +) -> tuple[Any, Any] | None: + domain: GradSyncDomain = getattr( + param, "grad_sync_domain", TP_DEFAULT_GRAD_SYNC_DOMAIN + ) + group = _resolve_domain_group(domain) + if group is None: + return None + op: GradSyncOp = getattr(param, "grad_sync_op", GRAD_SYNC_OP_NONE) + if op not in VALID_SYNC_OPS: + raise RuntimeError(f"{name}: unsupported grad_sync_op={op}") + if op == GRAD_SYNC_OP_NONE: + return None + return group, _resolve_reduce_op(op) + + +def coalesced_all_reduce( + grads: list[torch.Tensor], + *, + group: Any, + op: Any, +) -> None: + coalesced = _flatten_dense_tensors(grads) + reduced = ( + coalesced.float() + if torch.is_floating_point(coalesced) and coalesced.dtype != torch.float32 + else coalesced + ) + torch.distributed.all_reduce( # ty: ignore[possibly-missing-attribute] + reduced, + op=op, + group=group, + ) + if reduced is not coalesced: + reduced = reduced.to(dtype=coalesced.dtype) + for grad, synced in zip(grads, _unflatten_dense_tensors(reduced, grads)): + grad.copy_(synced) + + def flush_param_grads_to_main_grads(model_chunks: Iterable[torch.nn.Module]) -> None: """Fallback for direct jobs when DDP post-hooks leave grads in param.grad. @@ -100,57 +142,30 @@ def finalize_model_grads_extended( ) buckets: dict[ - tuple[GradSyncDomain, GradSyncOp, torch.dtype, torch.device], - list[tuple[str, torch.Tensor]], - ] = defaultdict(list) + tuple[int, str, torch.dtype, torch.device], + tuple[Any, Any, list[torch.Tensor]], + ] = {} for name, param in _iter_named_trainable_parameters(model): - domain: GradSyncDomain = getattr( - param, "grad_sync_domain", TP_DEFAULT_GRAD_SYNC_DOMAIN - ) - if _resolve_domain_group(domain) is None: - continue - - op: GradSyncOp = getattr(param, "grad_sync_op", GRAD_SYNC_OP_NONE) - if op not in VALID_SYNC_OPS: - raise RuntimeError(f"{name}: unsupported grad_sync_op={op}") - if op == GRAD_SYNC_OP_NONE: + sync = tensor_parallel_grad_sync(param, name=name) + if sync is None: continue if not hasattr(param, "main_grad"): raise RuntimeError( - f"{name}: expected main_grad for domain={domain} reduce_op={op}, but attribute is missing" + f"{name}: expected main_grad for tensor-parallel grad sync, but attribute is missing" ) grad = param.main_grad if grad is None: raise RuntimeError( - f"{name}: expected non-None main_grad for domain={domain} reduce_op={op}" + f"{name}: expected non-None main_grad for tensor-parallel grad sync" ) local_grad = cast( # local part of dtensor torch.Tensor, grad._local_tensor if hasattr(grad, "_local_tensor") else grad ) - buckets[(domain, op, local_grad.dtype, local_grad.device)].append( - (name, local_grad) - ) + group, reduce_op = sync + key = (id(group), str(reduce_op), local_grad.dtype, local_grad.device) + buckets.setdefault(key, (group, reduce_op, []))[2].append(local_grad) - for (domain, op, _dtype, _device), entries in buckets.items(): - group = _resolve_domain_group( - domain - ) # already checked if the domain is one we are handling - - grads = [grad for _name, grad in entries] - coalesced = _flatten_dense_tensors(grads) - reduced = ( - coalesced.float() - if torch.is_floating_point(coalesced) and coalesced.dtype != torch.float32 - else coalesced - ) - torch.distributed.all_reduce( # ty: ignore[possibly-missing-attribute] - reduced, - op=_resolve_reduce_op(op), - group=group, - ) - if reduced is not coalesced: - reduced = reduced.to(dtype=coalesced.dtype) - for grad, synced in zip(grads, _unflatten_dense_tensors(reduced, grads)): - grad.copy_(synced) + for group, op, grads in buckets.values(): + coalesced_all_reduce(grads, group=group, op=op) diff --git a/src/art/megatron/training/microbatches.py b/src/art/megatron/training/microbatches.py index de6650cba..c3b515c26 100644 --- a/src/art/megatron/training/microbatches.py +++ b/src/art/megatron/training/microbatches.py @@ -17,7 +17,7 @@ PreparedMegatronBatch, ) from art.megatron.flex_attn.compiled import flash_sparse_block_size_for_head_dim -from art.megatron.shared_prefix_state import create_shared_prefix_state +from art.megatron.prefix_tree_state import create_prefix_tree_state from art.megatron.training.trace import ( packed_sequence_token_uids, sft_sequence_token_uids, @@ -313,15 +313,16 @@ def _causal_attention_state( seq_len: int, device: torch.device, *, + provider: Any, + model_support_handler: Any, sliding_windows: tuple[int, ...] = (), build_gdn_execution_spec: bool, - model_support_handler: Any, attention_head_dim: int | None = None, attention_value_head_dim: int | None = None, ) -> Any: group_ids = torch.zeros((1, seq_len), dtype=torch.int64, device="cpu") parent_ids = torch.zeros_like(group_ids) - return create_shared_prefix_state( + return create_prefix_tree_state( group_ids=group_ids, parent_ids=parent_ids, target_device=device, @@ -331,9 +332,23 @@ def _causal_attention_state( model_support_handler=model_support_handler, attention_head_dim=attention_head_dim, attention_value_head_dim=attention_value_head_dim, + gdn_planner_config=_gdn_planner_config_for_provider( + provider, + model_support_handler, + ), ) +def _gdn_planner_config_for_provider( + provider: Any, model_support_handler: Any +) -> Any | None: + if not bool(getattr(model_support_handler, "build_gdn_execution_spec", False)): + return None + from art.megatron.gdn.gdn_prefix_tree import GdnPlannerConfig + + return GdnPlannerConfig.from_provider(provider) + + def _next_micro_lookahead( micro_inputs: list[Any], micro_order: int, @@ -353,7 +368,7 @@ def _prepare_dense_rl_micro( model_support_handler: Any, ref_logprobs: torch.Tensor | None, ) -> PreparedRLMicroInputs: - attention_state = create_shared_prefix_state( + attention_state = create_prefix_tree_state( group_ids=micro["group_ids"], parent_ids=micro["parent_ids"], target_device=device, @@ -365,6 +380,10 @@ def _prepare_dense_rl_micro( model_support_handler=model_support_handler, attention_head_dim=getattr(provider, "kv_channels", None), attention_value_head_dim=getattr(provider, "kv_channels", None), + gdn_planner_config=_gdn_planner_config_for_provider( + provider, + model_support_handler, + ), ) _move_inputs_to_device(micro, device) shifted_labels = shift_tensor(micro["tokens"], -100) @@ -410,6 +429,10 @@ def _prepare_rl_cp_micro_full( build_gdn_execution_spec=bool( getattr(model_support_handler, "build_gdn_execution_spec", False) ), + gdn_planner_config=_gdn_planner_config_for_provider( + provider, + model_support_handler, + ), trace_token_uids=trace_token_uids, block_mask_variants=_art_flex_cp_block_mask_variants(provider, device), target_device=device, @@ -564,6 +587,7 @@ def _prepare_dense_sft_micro( build_gdn_execution_spec=bool( getattr(model_support_handler, "build_gdn_execution_spec", False) ), + provider=provider, model_support_handler=model_support_handler, attention_head_dim=getattr(provider, "kv_channels", None), attention_value_head_dim=getattr(provider, "kv_channels", None), @@ -626,7 +650,7 @@ def _prepare_sft_cp_micro_full( The synthetic sparse-packed metadata is constructed on CPU and only the rank-local dispatched tensors are moved to `device`. Constructing it on CUDA - would make shared-prefix planning read metadata back from the GPU. + would make prefix-tree planning read metadata back from the GPU. """ sparse_micro = _sft_inputs_to_sparse_packed_tensors( micro, @@ -641,6 +665,10 @@ def _prepare_sft_cp_micro_full( build_gdn_execution_spec=bool( getattr(model_support_handler, "build_gdn_execution_spec", False) ), + gdn_planner_config=_gdn_planner_config_for_provider( + provider, + model_support_handler, + ), trace_token_uids=trace_token_uids, block_mask_variants=_art_flex_cp_block_mask_variants(provider, device), target_device=device, diff --git a/src/art/megatron/weights/adapter_export.py b/src/art/megatron/weights/adapter_export.py index 15f8062c2..664e30a02 100644 --- a/src/art/megatron/weights/adapter_export.py +++ b/src/art/megatron/weights/adapter_export.py @@ -1,3 +1,4 @@ +from collections.abc import Callable, Sequence import math from typing import Any @@ -9,31 +10,15 @@ from art.megatron.lora import ( GatedDeltaNetInProjLoRA, LoRA, - MLPExpertsLinearFC1FusedLoRA, MLPExpertsLinearFC1LoRA, MLPExpertsLinearFC2LoRA, SelfAttentionLinearProjLoRA, SelfAttentionLinearQKVLoRA, SharedExpertsLinearFC1LoRA, - SharedExpertsLinearFC2LoRA, ) from art.megatron.weights.param_name_canonicalization import canonical_art_param_name -def layer_base_prefix( - module: TransformerLayer, - *, - module_name: str | None = None, -) -> str: - if module_name is not None: - canonical_name = canonical_art_param_name(module_name) - if canonical_name.startswith( - ("decoder.layers.", "language_model.decoder.layers.") - ): - return canonical_name - return f"language_model.decoder.layers.{module.layer_number - 1}" - - def _adapter_alpha_dim(lora: LoRA) -> tuple[int, int]: dim = int(lora.A_T.shape[-1]) alpha = float(lora.scale) * dim @@ -51,12 +36,6 @@ def _adapter_tensors( return a_t.transpose(-1, -2).contiguous(), b_t.transpose(-1, -2).contiguous() -def _adapter_param_prefix(base_prefix: str, adapter_key: str | None) -> str: - if adapter_key is None: - return f"{base_prefix}.adapter" - return f"{base_prefix}.adapter.{adapter_key}" - - def _adapter_weight( *, base_prefix: str, @@ -66,7 +45,8 @@ def _adapter_weight( linear_in: torch.Tensor, linear_out: torch.Tensor, ) -> AdapterWeight: - param_prefix = _adapter_param_prefix(base_prefix, adapter_key) + adapter_suffix = "" if adapter_key is None else f".{adapter_key}" + param_prefix = f"{base_prefix}.adapter{adapter_suffix}" return AdapterWeight( global_base_prefix=base_prefix, adapter_key=adapter_key, @@ -162,96 +142,198 @@ def _fused_pair_adapter_weight( ) -def add_standard_self_attention_adapter_weights( - adapter_weights_by_base: dict[str, list[Any]], - *, - layer_prefix: str, - self_attention: Any, +def _set_adapter_weights( + out: dict[str, list[Any]], + base_prefix: str, + *weights: AdapterWeight, + weight_suffix: str = ".weight", ) -> None: - linear_proj = getattr(self_attention, "linear_proj", None) - if isinstance(linear_proj, SelfAttentionLinearProjLoRA): - base_prefix = f"{layer_prefix}.self_attention.linear_proj" - adapter_weights_by_base[f"{base_prefix}.weight"] = [ - _simple_adapter_weight(base_prefix, linear_proj.lora) - ] + out[f"{base_prefix}{weight_suffix}"] = list(weights) - linear_qkv = getattr(self_attention, "linear_qkv", None) - if isinstance(linear_qkv, SelfAttentionLinearQKVLoRA): - base_prefix = f"{layer_prefix}.self_attention.linear_qkv" - adapter_weights = [] - q_proj_lora = getattr(linear_qkv, "q_proj_lora", None) - k_proj_lora = getattr(linear_qkv, "k_proj_lora", None) - v_proj_lora = getattr(linear_qkv, "v_proj_lora", None) - if q_proj_lora is not None: - adapter_weights.append( - _simple_adapter_weight( - base_prefix, - q_proj_lora, - adapter_key="adapter_q", - ) - ) - if k_proj_lora is not None: - adapter_weights.append( - _simple_adapter_weight( - base_prefix, - k_proj_lora, - adapter_key="adapter_k", - ) - ) - if v_proj_lora is not None: - adapter_weights.append( - _simple_adapter_weight( - base_prefix, - v_proj_lora, - adapter_key="adapter_v", - ) + +def _set_expert_adapter_weights( + out: dict[str, list[Any]], + base_prefix: str, + lora: LoRA, + build_weight: Callable[[int], AdapterWeight], +) -> None: + for local_expert_idx in range(lora.num_local_experts): + global_expert_idx = local_expert_idx + lora._expert_offset + _set_adapter_weights( + out, + base_prefix, + build_weight(local_expert_idx), + weight_suffix=f".weight{global_expert_idx}", + ) + + +def _set_lora_weights( + out: dict[str, list[Any]], + base_prefix: str, + *items: tuple[LoRA | None, str | None], +) -> None: + weights = [ + _simple_adapter_weight(base_prefix, lora, adapter_key=adapter_key) + for lora, adapter_key in items + if lora is not None + ] + if not weights: + return + _set_adapter_weights( + out, + base_prefix, + *weights, + ) + + +def layer_base_prefix( + module: TransformerLayer, + *, + module_name: str | None = None, +) -> str: + if module_name is not None: + canonical_name = canonical_art_param_name(module_name) + if canonical_name.startswith( + ("decoder.layers.", "language_model.decoder.layers.") + ): + return canonical_name + return f"language_model.decoder.layers.{module.layer_number - 1}" + + +def build_transformer_layer_adapter_weights( + model_chunks: Sequence[Any], + grouped_moe: bool = False, + language_layers_only: bool = False, +) -> dict[str, list[Any]]: + layer_filter = None + if language_layers_only: + from art.megatron.lora import ( + _is_language_transformer_layer_name as layer_filter, + ) + + add_mlp_adapter_weights = ( + _add_moe_mlp_adapter_weights_for_layer + if grouped_moe + else _add_dense_mlp_adapter_weights_for_layer + ) + adapter_weights_by_base: dict[str, list[Any]] = {} + for chunk in model_chunks: + for module_name, module in chunk.named_modules(): + if not isinstance(module, TransformerLayer): + continue + if layer_filter is not None and not layer_filter(module_name): + continue + layer_prefix = layer_base_prefix(module, module_name=module_name) + add_self_attention_adapter_weights( + adapter_weights_by_base, + layer_prefix=layer_prefix, + self_attention=module.self_attention, ) - if adapter_weights: - adapter_weights_by_base[f"{base_prefix}.weight"] = adapter_weights + add_mlp_adapter_weights(adapter_weights_by_base, layer_prefix, module) + return adapter_weights_by_base -def add_gated_delta_net_adapter_weights( +def add_self_attention_adapter_weights( adapter_weights_by_base: dict[str, list[Any]], *, layer_prefix: str, self_attention: Any, ) -> None: - out_proj = getattr(self_attention, "out_proj", None) - if isinstance(out_proj, SelfAttentionLinearProjLoRA): - base_prefix = f"{layer_prefix}.self_attention.out_proj" - adapter_weights_by_base[f"{base_prefix}.weight"] = [ - _simple_adapter_weight(base_prefix, out_proj.lora) - ] + for attr in ("linear_proj", "out_proj"): + linear_proj = getattr(self_attention, attr, None) + if isinstance(linear_proj, SelfAttentionLinearProjLoRA): + base_prefix = f"{layer_prefix}.self_attention.{attr}" + _set_lora_weights( + adapter_weights_by_base, + base_prefix, + (linear_proj.lora, None), + ) + + linear_qkv = getattr(self_attention, "linear_qkv", None) + if isinstance(linear_qkv, SelfAttentionLinearQKVLoRA): + base_prefix = f"{layer_prefix}.self_attention.linear_qkv" + _set_lora_weights( + adapter_weights_by_base, + base_prefix, + (linear_qkv.q_proj_lora, "adapter_q"), + (linear_qkv.k_proj_lora, "adapter_k"), + (linear_qkv.v_proj_lora, "adapter_v"), + ) in_proj = getattr(self_attention, "in_proj", None) if isinstance(in_proj, GatedDeltaNetInProjLoRA): base_prefix = f"{layer_prefix}.self_attention.in_proj" - adapter_weights_by_base[f"{base_prefix}.weight"] = [ + input_dim = int(in_proj.qkv_lora.A_T.shape[-2]) + output_dim = int(in_proj.num_value_heads_per_partition) + _set_adapter_weights( + adapter_weights_by_base, + base_prefix, _simple_adapter_weight( - base_prefix, - in_proj.qkv_lora, - adapter_key="adapter_qkv", + base_prefix, in_proj.qkv_lora, adapter_key="adapter_qkv" ), _simple_adapter_weight( - base_prefix, - in_proj.z_lora, - adapter_key="adapter_z", - ), - _zero_adapter_weight( - base_prefix=base_prefix, - adapter_key="adapter_b", - input_dim=int(in_proj.qkv_lora.A_T.shape[-2]), - output_dim=int(in_proj.num_value_heads_per_partition), - like=in_proj.qkv_lora.B_T, + base_prefix, in_proj.z_lora, adapter_key="adapter_z" ), - _zero_adapter_weight( - base_prefix=base_prefix, - adapter_key="adapter_a", - input_dim=int(in_proj.qkv_lora.A_T.shape[-2]), - output_dim=int(in_proj.num_value_heads_per_partition), - like=in_proj.qkv_lora.B_T, + *( + _zero_adapter_weight( + base_prefix=base_prefix, + adapter_key=adapter_key, + input_dim=input_dim, + output_dim=output_dim, + like=in_proj.qkv_lora.B_T, + ) + for adapter_key in ("adapter_b", "adapter_a") ), - ] + ) + + +def add_standard_self_attention_adapter_weights( + adapter_weights_by_base: dict[str, list[Any]], + *, + layer_prefix: str, + self_attention: Any, +) -> None: + add_self_attention_adapter_weights( + adapter_weights_by_base, + layer_prefix=layer_prefix, + self_attention=self_attention, + ) + + +def _add_dense_mlp_adapter_weights_for_layer( + adapter_weights_by_base: dict[str, list[Any]], + layer_prefix: str, + module: Any, +) -> None: + from art.megatron.model_support.handlers.default_dense import _require_dense_mlp + + _require_dense_mlp(module) + add_split_mlp_adapter_weights( + adapter_weights_by_base, + f"{layer_prefix}.mlp", + module.mlp, + ) + + +def _add_moe_mlp_adapter_weights_for_layer( + adapter_weights_by_base: dict[str, list[Any]], + layer_prefix: str, + module: Any, +) -> None: + from art.megatron.model_support.handlers.default_dense import _require_moe_experts + + add_grouped_moe_adapter_weights( + adapter_weights_by_base, + layer_prefix=layer_prefix, + experts=_require_moe_experts(module), + ) + shared_experts = getattr(module.mlp, "shared_experts", None) + if shared_experts is not None: + add_split_mlp_adapter_weights( + adapter_weights_by_base, + f"{layer_prefix}.mlp.shared_experts", + shared_experts, + ) def add_grouped_moe_adapter_weights( @@ -261,43 +343,44 @@ def add_grouped_moe_adapter_weights( experts: Any, ) -> None: linear_fc1 = getattr(experts, "linear_fc1", None) - if isinstance(linear_fc1, MLPExpertsLinearFC1FusedLoRA): - base_prefix = f"{layer_prefix}.mlp.experts.linear_fc1" - for local_expert_idx in range(linear_fc1.lora.num_local_experts): - global_expert_idx = local_expert_idx + linear_fc1.lora._expert_offset - adapter_weights_by_base[f"{base_prefix}.weight{global_expert_idx}"] = [ - _simple_adapter_weight( - base_prefix, - linear_fc1.lora, - expert_idx=local_expert_idx, - ) - ] - elif isinstance(linear_fc1, MLPExpertsLinearFC1LoRA): - base_prefix = f"{layer_prefix}.mlp.experts.linear_fc1" - for local_expert_idx in range(linear_fc1.gate_lora.num_local_experts): - global_expert_idx = local_expert_idx + linear_fc1.gate_lora._expert_offset - adapter_weights_by_base[f"{base_prefix}.weight{global_expert_idx}"] = [ - _fused_pair_adapter_weight( - base_prefix, - linear_fc1.gate_lora, - linear_fc1.up_lora, - first_expert_idx=local_expert_idx, - second_expert_idx=local_expert_idx, - ) - ] + base_prefix = f"{layer_prefix}.mlp.experts.linear_fc1" + if isinstance(linear_fc1, MLPExpertsLinearFC1LoRA): + if linear_fc1.fused_gate_up: + lora = linear_fc1.lora + build_weight = lambda local_expert_idx: _simple_adapter_weight( + base_prefix, + linear_fc1.lora, + expert_idx=local_expert_idx, + ) + else: + lora = linear_fc1.gate_lora + build_weight = lambda local_expert_idx: _fused_pair_adapter_weight( + base_prefix, + linear_fc1.gate_lora, + linear_fc1.up_lora, + first_expert_idx=local_expert_idx, + second_expert_idx=local_expert_idx, + ) + _set_expert_adapter_weights( + adapter_weights_by_base, + base_prefix, + lora, + build_weight, + ) linear_fc2 = getattr(experts, "linear_fc2", None) if isinstance(linear_fc2, MLPExpertsLinearFC2LoRA): base_prefix = f"{layer_prefix}.mlp.experts.linear_fc2" - for local_expert_idx in range(linear_fc2.lora.num_local_experts): - global_expert_idx = local_expert_idx + linear_fc2.lora._expert_offset - adapter_weights_by_base[f"{base_prefix}.weight{global_expert_idx}"] = [ - _simple_adapter_weight( - base_prefix, - linear_fc2.lora, - expert_idx=local_expert_idx, - ) - ] + _set_expert_adapter_weights( + adapter_weights_by_base, + base_prefix, + linear_fc2.lora, + lambda local_expert_idx: _simple_adapter_weight( + base_prefix, + linear_fc2.lora, + expert_idx=local_expert_idx, + ), + ) def add_dense_mlp_adapter_weights( @@ -306,28 +389,11 @@ def add_dense_mlp_adapter_weights( layer_prefix: str, mlp: Any, ) -> None: - linear_fc1 = getattr(mlp, "linear_fc1", None) - if isinstance(linear_fc1, SharedExpertsLinearFC1LoRA): - base_prefix = f"{layer_prefix}.mlp.linear_fc1" - adapter_weights_by_base[f"{base_prefix}.weight"] = [ - _simple_adapter_weight( - base_prefix, - linear_fc1.gate_lora, - adapter_key="adapter_gate", - ), - _simple_adapter_weight( - base_prefix, - linear_fc1.up_lora, - adapter_key="adapter_up", - ), - ] - - linear_fc2 = getattr(mlp, "linear_fc2", None) - if isinstance(linear_fc2, SharedExpertsLinearFC2LoRA): - base_prefix = f"{layer_prefix}.mlp.linear_fc2" - adapter_weights_by_base[f"{base_prefix}.weight"] = [ - _simple_adapter_weight(base_prefix, linear_fc2.row_parallel_lora.lora) - ] + add_split_mlp_adapter_weights( + adapter_weights_by_base, + f"{layer_prefix}.mlp", + mlp, + ) def add_shared_experts_adapter_weights( @@ -336,25 +402,33 @@ def add_shared_experts_adapter_weights( layer_prefix: str, shared_experts: Any, ) -> None: - linear_fc1 = getattr(shared_experts, "linear_fc1", None) + add_split_mlp_adapter_weights( + adapter_weights_by_base, + f"{layer_prefix}.mlp.shared_experts", + shared_experts, + ) + + +def add_split_mlp_adapter_weights( + adapter_weights_by_base: dict[str, list[Any]], + base_prefix: str, + mlp: Any, +) -> None: + linear_fc1 = getattr(mlp, "linear_fc1", None) if isinstance(linear_fc1, SharedExpertsLinearFC1LoRA): - base_prefix = f"{layer_prefix}.mlp.shared_experts.linear_fc1" - adapter_weights_by_base[f"{base_prefix}.weight"] = [ - _simple_adapter_weight( - base_prefix, - linear_fc1.gate_lora, - adapter_key="adapter_gate", - ), - _simple_adapter_weight( - base_prefix, - linear_fc1.up_lora, - adapter_key="adapter_up", - ), - ] - - linear_fc2 = getattr(shared_experts, "linear_fc2", None) - if isinstance(linear_fc2, SharedExpertsLinearFC2LoRA): - base_prefix = f"{layer_prefix}.mlp.shared_experts.linear_fc2" - adapter_weights_by_base[f"{base_prefix}.weight"] = [ - _simple_adapter_weight(base_prefix, linear_fc2.row_parallel_lora.lora) - ] + fc1_prefix = f"{base_prefix}.linear_fc1" + _set_lora_weights( + adapter_weights_by_base, + fc1_prefix, + (linear_fc1.gate_lora, "adapter_gate"), + (linear_fc1.up_lora, "adapter_up"), + ) + + linear_fc2 = getattr(mlp, "linear_fc2", None) + if isinstance(linear_fc2, SelfAttentionLinearProjLoRA): + fc2_prefix = f"{base_prefix}.linear_fc2" + _set_adapter_weights( + adapter_weights_by_base, + fc2_prefix, + _simple_adapter_weight(fc2_prefix, linear_fc2.lora), + ) diff --git a/src/art/megatron/weights/lora_publish.py b/src/art/megatron/weights/lora_publish.py index 1b3c4c5b0..f7b9c5427 100644 --- a/src/art/megatron/weights/lora_publish.py +++ b/src/art/megatron/weights/lora_publish.py @@ -1,16 +1,22 @@ from collections.abc import Iterable, Sequence -import re from typing import Any, NamedTuple import torch -from art.megatron.lora import LoRAPublishPlanner, LoraShardMeta +from art.megatron.lora import ( + LoRA, + LoRAPublishPlanner, + LoraShardMeta, + _block_for_key, + _dtype_name, +) +from art.megatron.lora import ( + _distributed_initialized as _distributed_ready, +) from art.megatron.model_support.lora_disk import save_vllm_lora_tensors from art.megatron.model_support.spec import ExpertPackedLoraGroup, ExpertPackedLoraSlot from art.megatron.training.model_chunks import ModelChunks -_LAYER_BLOCK_RE = re.compile(r"^(?P.*\.layers\.\d+)\.") - class PackedExpertShardMeta(NamedTuple): key: str @@ -58,35 +64,17 @@ def finish(self) -> None: self._events.clear() -def iter_lora_modules(model_chunks: ModelChunks) -> Iterable[Any]: +def iter_lora_modules(model_chunks: ModelChunks) -> Iterable[LoRA]: for chunk in model_chunks: for module in chunk.modules(): - yield module - - -def _dtype_name(dtype: torch.dtype) -> str: - return str(dtype).removeprefix("torch.") + if isinstance(module, LoRA): + yield module def _dtype_from_name(name: str) -> torch.dtype: - dtype = getattr(torch, name, None) - if not isinstance(dtype, torch.dtype): - raise RuntimeError(f"Unsupported LoRA tensor dtype={name!r}") - return dtype - - -def _block_for_key(key: str) -> str: - match = _LAYER_BLOCK_RE.match(key) - if match is not None: - return match.group("block") - return "__global__" - - -def _expert_prefix_projection(adapter_model_prefix: str) -> tuple[str, str] | None: - group_prefix, separator, projection = adapter_model_prefix.partition(".{expert}.") - if not separator: - return None - return group_prefix, projection + if isinstance(dtype := getattr(torch, name, None), torch.dtype): + return dtype + raise RuntimeError(f"Unsupported LoRA tensor dtype={name!r}") def _packed_expert_slot( @@ -94,10 +82,9 @@ def _packed_expert_slot( suffix: str, groups: Sequence[ExpertPackedLoraGroup], ) -> tuple[str, ExpertPackedLoraSlot] | None: - parts = _expert_prefix_projection(adapter_model_prefix) - if parts is None: + group_prefix, separator, projection = adapter_model_prefix.partition(".{expert}.") + if not separator: return None - group_prefix, projection = parts lora_name = suffix.removesuffix(".weight") for group in groups: if not group_prefix.endswith(group.art_group_suffix): @@ -109,23 +96,15 @@ def _packed_expert_slot( def _uses_packed_expert_publish( - module: Any, + module: LoRA, groups: Sequence[ExpertPackedLoraGroup], ) -> bool: - if int(getattr(module, "num_local_experts", 1)) <= 1: - return False - if not hasattr(module, "_lora_params"): + if module.num_local_experts <= 1: return False - adapter_model_prefix = getattr(module, "adapter_model_prefix", "") - if not isinstance(adapter_model_prefix, str): - return False - lora_suffixes = [ - suffix - for suffix, _param in module._lora_params() # type: ignore[attr-defined] - ] - return bool(lora_suffixes) and all( - _packed_expert_slot(adapter_model_prefix, suffix, groups) is not None - for suffix in lora_suffixes + params = tuple(module._lora_params()) + return bool(params) and all( + _packed_expert_slot(module.adapter_model_prefix, suffix, groups) is not None + for suffix, _param in params ) @@ -141,15 +120,12 @@ def collect_local_lora_entries( for module in iter_lora_modules(model_chunks): if _uses_packed_expert_publish(module, packed_expert_groups): continue - if hasattr(module, "sharded_lora_state_dict"): - module_state: dict[str, torch.Tensor] = module.sharded_lora_state_dict() # type: ignore[attr-defined] - for key, value in module_state.items(): - target_dtype = ( - adapter_model[key].dtype if key in adapter_model else value.dtype - ) - local_tensors[key] = value.to(target_dtype).contiguous() - if hasattr(module, "sharded_lora_manifest"): - local_manifest.update(module.sharded_lora_manifest()) # type: ignore[attr-defined] + for key, value in module.sharded_lora_state_dict().items(): + target_dtype = ( + adapter_model[key].dtype if key in adapter_model else value.dtype + ) + local_tensors[key] = value.to(target_dtype).contiguous() + local_manifest.update(module.sharded_lora_manifest()) if set(local_tensors) != set(local_manifest): raise RuntimeError( @@ -171,18 +147,6 @@ def collect_local_lora_entries( return local_tensors, metadata -def _target_dtype_for_lora_param( - module: Any, - adapter_model: dict[str, torch.Tensor], - suffix: str, - fallback: torch.dtype, -) -> torch.dtype: - keys = module._expected_weight_keys(suffix.removesuffix(".weight")) # type: ignore[attr-defined] - return ( - adapter_model[keys[0]].dtype if keys and keys[0] in adapter_model else fallback - ) - - def collect_local_packed_expert_entries( model_chunks: ModelChunks, adapter_model: dict[str, torch.Tensor], @@ -195,25 +159,24 @@ def collect_local_packed_expert_entries( for module in iter_lora_modules(model_chunks): if not _uses_packed_expert_publish(module, packed_expert_groups): continue - adapter_model_prefix = module.adapter_model_prefix # type: ignore[attr-defined] - expert_start = int(module._expert_offset) # type: ignore[attr-defined] - expert_count = int(module.num_local_experts) # type: ignore[attr-defined] - for suffix, param in module._lora_params(): # type: ignore[attr-defined] + expert_start = int(module._expert_offset) + expert_count = int(module.num_local_experts) + for suffix, param in module._lora_params(): slot_match = _packed_expert_slot( - adapter_model_prefix, + module.adapter_model_prefix, suffix, packed_expert_groups, ) - if slot_match is None or not module._should_export_parameter(param): # type: ignore[attr-defined] + if slot_match is None or not module._should_export_parameter(param): continue group_prefix, slot = slot_match key = f"{group_prefix}.{slot.output_suffix}" tensor = param.data.transpose(1, 2).contiguous() - target_dtype = _target_dtype_for_lora_param( - module, - adapter_model, - suffix, - tensor.dtype, + source_keys = module._expected_weight_keys(suffix.removesuffix(".weight")) + target_dtype = ( + adapter_model[source_keys[0]].dtype + if source_keys and source_keys[0] in adapter_model + else tensor.dtype ) tensor = tensor.to(target_dtype).contiguous() if key in local_tensors: @@ -225,7 +188,7 @@ def collect_local_packed_expert_entries( owner_rank=owner_rank, shape=tuple(int(dim) for dim in tensor.shape), dtype_name=_dtype_name(tensor.dtype), - manifest=module._manifest_for_param(param), # type: ignore[attr-defined] + manifest=module._manifest_for_param(param), expert_start=expert_start, expert_count=expert_count, pack_layout=slot.pack_layout, @@ -252,7 +215,12 @@ def _global_packed_expert_metadata( continue group_prefix, slot = slot_match shard_ranks = range(template.shard_world_size) if template.sharded else (0,) - for ep_rank in range(planner._expert_model_world_size()): + ep_world_size = 1 + if _distributed_ready(): + from megatron.core import parallel_state as ps + + ep_world_size = ps.get_expert_model_parallel_world_size() + for ep_rank in range(ep_world_size): expert_start = ep_rank * template.num_local_experts expert_key = ( f"{template.adapter_model_prefix.format(expert=expert_start)}." @@ -350,70 +318,66 @@ def _merge_sharded_tensor( return torch.cat(tuple(ordered_shards), dim=axis).contiguous() -def merge_sharded_adapter_entries( - entries_by_key: dict[str, list[tuple[dict[str, Any], torch.Tensor]]], -) -> dict[str, torch.Tensor]: - adapter_model: dict[str, torch.Tensor] = {} - for key, key_entries in entries_by_key.items(): - first_manifest = key_entries[0][0] - sharded = bool(first_manifest["sharded"]) - shard_world_size = int(first_manifest["shard_world_size"]) - for manifest_entry, _tensor in key_entries: - if bool(manifest_entry["sharded"]) != sharded: - raise RuntimeError(f"Inconsistent sharded flag for key={key}") - if int(manifest_entry["shard_world_size"]) != shard_world_size: - raise RuntimeError(f"Inconsistent shard world size for key={key}") - - if not sharded: - if len(key_entries) != 1: - raise RuntimeError( - f"Replicated key={key} expected 1 shard, got {len(key_entries)}" - ) - adapter_model[key] = key_entries[0][1] - continue - - shard_rank_to_tensor: dict[int, torch.Tensor] = {} - for manifest_entry, shard_tensor in key_entries: - shard_rank = int(manifest_entry["shard_rank"]) - if shard_rank in shard_rank_to_tensor: - raise RuntimeError(f"Duplicate shard_rank={shard_rank} for key={key}") - shard_rank_to_tensor[shard_rank] = shard_tensor +def _merge_manifest_entries( + key: str, + key_entries: Sequence[tuple[dict[str, Any], torch.Tensor]], + *, + manifest: dict[str, Any] | None = None, +) -> torch.Tensor: + first_manifest = key_entries[0][0] + sharded = bool(first_manifest["sharded"]) + shard_world_size = int(first_manifest["shard_world_size"]) + for entry_manifest, _tensor in key_entries: + if bool(entry_manifest["sharded"]) != sharded: + raise RuntimeError(f"Inconsistent sharded flag for key={key}") + if int(entry_manifest["shard_world_size"]) != shard_world_size: + raise RuntimeError(f"Inconsistent shard world size for key={key}") - expected_shard_ranks = set(range(shard_world_size)) - if set(shard_rank_to_tensor) != expected_shard_ranks: + if not sharded: + if len(key_entries) != 1: raise RuntimeError( - f"Shard rank coverage mismatch for key={key}: " - f"expected {sorted(expected_shard_ranks)}, got {sorted(shard_rank_to_tensor)}" + f"Replicated key={key} expected 1 shard, got {len(key_entries)}" ) + return key_entries[0][1] - ordered_shards = [ - shard_rank_to_tensor[shard_rank] for shard_rank in range(shard_world_size) - ] - adapter_model[key] = _merge_sharded_tensor( - key, - ordered_shards=ordered_shards, - manifest=first_manifest, + shard_rank_to_tensor: dict[int, torch.Tensor] = {} + for entry_manifest, shard_tensor in key_entries: + shard_rank = int(entry_manifest["shard_rank"]) + if shard_rank in shard_rank_to_tensor: + raise RuntimeError(f"Duplicate shard_rank={shard_rank} for key={key}") + shard_rank_to_tensor[shard_rank] = shard_tensor + + expected_shard_ranks = set(range(shard_world_size)) + if set(shard_rank_to_tensor) != expected_shard_ranks: + raise RuntimeError( + f"Shard rank coverage mismatch for key={key}: " + f"expected {sorted(expected_shard_ranks)}, got {sorted(shard_rank_to_tensor)}" ) - return adapter_model + return _merge_sharded_tensor( + key, + ordered_shards=[ + shard_rank_to_tensor[shard_rank] for shard_rank in range(shard_world_size) + ], + manifest=first_manifest if manifest is None else manifest, + ) -def _distributed_ready() -> bool: - is_initialized = getattr(torch.distributed, "is_initialized", None) - return ( - torch.distributed.is_available() - and callable(is_initialized) - and bool(is_initialized()) - ) +def merge_sharded_adapter_entries( + entries_by_key: dict[str, list[tuple[dict[str, Any], torch.Tensor]]], +) -> dict[str, torch.Tensor]: + return { + key: _merge_manifest_entries(key, key_entries) + for key, key_entries in entries_by_key.items() + } def _rank_and_device() -> tuple[int, torch.device]: - if _distributed_ready(): - rank = torch.distributed.get_rank() # type: ignore[possibly-missing-attribute] - else: - rank = 0 - if torch.cuda.is_available(): - return rank, torch.device("cuda", torch.cuda.current_device()) - return rank, torch.device("cpu") + return ( + torch.distributed.get_rank() if _distributed_ready() else 0, # type: ignore[possibly-missing-attribute] + torch.device("cuda", torch.cuda.current_device()) + if torch.cuda.is_available() + else torch.device("cpu"), + ) def _metadata_by_owner_dtype( @@ -514,45 +478,10 @@ def _merge_packed_expert_block( key: str, key_entries: list[tuple[dict[str, Any], torch.Tensor]], ) -> torch.Tensor: - first_manifest = key_entries[0][0] - sharded = bool(first_manifest["sharded"]) - shard_world_size = int(first_manifest["shard_world_size"]) - if not sharded: - if len(key_entries) != 1: - raise RuntimeError( - f"Replicated packed key={key} expected 1 shard, got {len(key_entries)}" - ) - return key_entries[0][1] - - shard_rank_to_tensor: dict[int, torch.Tensor] = {} - for manifest_entry, shard_tensor in key_entries: - if bool(manifest_entry["sharded"]) != sharded: - raise RuntimeError(f"Inconsistent sharded flag for packed key={key}") - if int(manifest_entry["shard_world_size"]) != shard_world_size: - raise RuntimeError(f"Inconsistent shard world size for packed key={key}") - shard_rank = int(manifest_entry["shard_rank"]) - if shard_rank in shard_rank_to_tensor: - raise RuntimeError( - f"Duplicate shard_rank={shard_rank} for packed key={key}" - ) - shard_rank_to_tensor[shard_rank] = shard_tensor - - expected_shard_ranks = set(range(shard_world_size)) - if set(shard_rank_to_tensor) != expected_shard_ranks: - raise RuntimeError( - f"Shard rank coverage mismatch for packed key={key}: " - f"expected {sorted(expected_shard_ranks)}, got {sorted(shard_rank_to_tensor)}" - ) - - manifest = dict(first_manifest) - manifest["export_shard_dim"] = int(manifest["export_shard_dim"]) + 1 - return _merge_sharded_tensor( - key, - ordered_shards=[ - shard_rank_to_tensor[shard_rank] for shard_rank in range(shard_world_size) - ], - manifest=manifest, - ) + manifest = dict(key_entries[0][0]) + if bool(manifest["sharded"]): + manifest["export_shard_dim"] = int(manifest["export_shard_dim"]) + 1 + return _merge_manifest_entries(key, key_entries, manifest=manifest) def _pack_merged_expert_blocks( diff --git a/src/art/model.py b/src/art/model.py index 597aa30a5..c92f14f36 100644 --- a/src/art/model.py +++ b/src/art/model.py @@ -29,6 +29,7 @@ from .preprocessing.vllm_tokens import attach_vllm_token_metadata_to_choice from .trajectories import Trajectory, TrajectoryGroup from .types import SFTMetricLoggingConfig, TrainSFTConfig +from .utils import wandb_sdk from .utils.trajectory_logging import write_trajectory_groups_parquet if TYPE_CHECKING: @@ -617,28 +618,31 @@ def _sync_wandb_config( def _get_wandb_run(self) -> Optional["Run"]: """Get or create the wandb run for this model.""" - import wandb - if "WANDB_API_KEY" not in os.environ: return None if self._wandb_run is None or self._wandb_run._is_finished: - run = wandb.init( - project=self.project, - name=self.name, - id=self.name, - config=self._wandb_config or None, - resume="allow", - reinit="create_new", - settings=wandb.Settings( - x_stats_open_metrics_endpoints={ - "vllm": "http://localhost:8000/metrics", - }, - x_stats_open_metrics_filters=( - "vllm.vllm:num_requests_waiting", - "vllm.vllm:num_requests_running", + try: + run = wandb_sdk.init( + project=self.project, + name=self.name, + id=self.name, + config=self._wandb_config or None, + resume="allow", + reinit="create_new", + settings=wandb_sdk.settings( + x_stats_open_metrics_endpoints={ + "vllm": "http://localhost:8000/metrics", + }, + x_stats_open_metrics_filters=( + "vllm.vllm:num_requests_waiting", + "vllm.vllm:num_requests_running", + ), ), - ), - ) + ) + except ModuleNotFoundError as e: + if e.name == "wandb" or (e.name or "").startswith("wandb."): + return None + raise self._wandb_run = run object.__setattr__( self, diff --git a/src/art/preprocessing/moe_routing.py b/src/art/preprocessing/moe_routing.py index e4a31b307..91cc053f1 100644 --- a/src/art/preprocessing/moe_routing.py +++ b/src/art/preprocessing/moe_routing.py @@ -47,15 +47,15 @@ class MoeRoutingAlignmentStats(BaseModel): class MoeRoutingPackStats(BaseModel): packed_tokens: int = 0 - shared_prefix_rows: int = 0 - shared_prefix_conflict_rows: int = 0 - shared_prefix_conflict_slots: int = 0 - shared_prefix_compared_slots: int = 0 + prefix_tree_rows: int = 0 + prefix_tree_conflict_rows: int = 0 + prefix_tree_conflict_slots: int = 0 + prefix_tree_compared_slots: int = 0 def add_alignment(self, stats: MoeRoutingAlignmentStats) -> None: - self.shared_prefix_conflict_rows += stats.overlap_conflict_rows - self.shared_prefix_conflict_slots += stats.overlap_conflict_slots - self.shared_prefix_compared_slots += stats.overlap_compared_slots + self.prefix_tree_conflict_rows += stats.overlap_conflict_rows + self.prefix_tree_conflict_slots += stats.overlap_conflict_slots + self.prefix_tree_compared_slots += stats.overlap_compared_slots class PackedMoeRoutingReplay(BaseModel): diff --git a/src/art/preprocessing/pack.py b/src/art/preprocessing/pack.py index c84fd9699..e42558947 100644 --- a/src/art/preprocessing/pack.py +++ b/src/art/preprocessing/pack.py @@ -1,11 +1,18 @@ +import math import os import random import time -from typing import Any, cast +from typing import Any, Literal, NamedTuple, cast import torch from typing_extensions import NotRequired, TypedDict, Unpack +from ..megatron.prefix_tree_packing import ( + estimate_prefix_tree_packed_tokens, +) +from ..megatron.prefix_tree_packing import ( + prefix_tree_pack as _prefix_tree_pack_sequences, +) from ..types import Verbosity from .moe_routing import ( MoeRoutingPackStats, @@ -38,6 +45,34 @@ class DiskPackedTensors(TypedDict): image_grid_thw: NotRequired[tuple[int, list[int]]] +class _PackedPrefixTreeRow(TypedDict): + token_ids: list[int] + group_ids: list[int] + parent_ids: list[int] + input_pos: list[int] + assistant_mask: list[int] + logprobs: list[float] + advantages: list[float] + weights: list[float] + pixel_values: torch.Tensor | None + image_grid_thw: torch.Tensor | None + moe_routes: list[TokenRoute | None] + + +class _PrefixTreePackItem(NamedTuple): + token_ids: tuple[int, ...] + input_pos: tuple[int, ...] + assistant_mask: tuple[int, ...] + logprobs: tuple[float, ...] + advantage: float + weight: float + prompt_id: int + shareable_length: int + pixel_values: torch.Tensor | None + image_grid_thw: torch.Tensor | None + moe_routes: tuple[TokenRoute | None, ...] | None + + def packed_tensors_from_tokenized_results( tokenized_results: list[TokenizedResult], seq_len: int, @@ -48,18 +83,30 @@ def packed_tensors_from_tokenized_results( pack_results: bool = True, include_moe_routing: bool = False, ) -> PackedTensors: - # TODO: This function could potentially be optimized with vectorized operations - token_ids: list[list[int]] = [[]] - group_ids: list[list[int]] = [[]] - parent_ids: list[list[int]] = [[]] - input_pos: list[list[int]] = [[]] - assistant_mask: list[list[int]] = [[]] - logprobs: list[list[float]] = [[]] - advantages: list[list[float]] = [[]] - weights: list[list[float]] = [[]] - pixel_values: list[list[torch.Tensor]] = [[]] - image_grid_thw: list[list[torch.Tensor]] = [[]] - moe_routes: list[list[TokenRoute | None]] = [[]] + return prefix_tree_pack( + tokenized_results=tokenized_results, + seq_len=seq_len, + pad_token_id=pad_token_id, + truncate_long_results=truncate_long_results, + advantage_balance=advantage_balance, + verbosity=verbosity, + pack_results=pack_results, + include_moe_routing=include_moe_routing, + ) + + +def prefix_tree_pack( + *, + tokenized_results: list[TokenizedResult], + seq_len: int, + pad_token_id: int = -100, + truncate_long_results: bool = True, + advantage_balance: float = 0.0, + verbosity: Verbosity = 1, + pack_results: bool = True, + include_moe_routing: bool = False, +) -> PackedTensors: + rows: list[list[_PrefixTreePackItem]] = [[]] moe_routing_pack_stats = MoeRoutingPackStats() for result in tokenized_results: @@ -72,85 +119,31 @@ def packed_tensors_from_tokenized_results( "MoE routing replay from trajectories was requested, but a " "tokenized result has no aligned routed experts" ) - result_without_prompt = result.without_prompt() - if sum(result_without_prompt.assistant_mask) == 0: + if sum(result.assistant_mask[result.prompt_length :]) == 0: if verbosity > 1: print("Result has no unique completion tokens, skipping") continue - if token_ids[-1] and ( + item = _prefix_tree_pack_item(result, seq_len=seq_len) + if rows[-1] and ( not pack_results - or len(token_ids[-1]) - + ( - len(result_without_prompt.token_ids) - if result.prompt_id in group_ids[-1] - else len(result.token_ids) - ) - > seq_len + or _packed_row_token_count([*rows[-1], item], seq_len=seq_len) > seq_len ): - token_ids.append([]) - group_ids.append([]) - parent_ids.append([]) - input_pos.append([]) - assistant_mask.append([]) - logprobs.append([]) - advantages.append([]) - weights.append([]) - pixel_values.append([]) - image_grid_thw.append([]) - moe_routes.append([]) - group_id = random.randint(-(2**63), 2**63 - 1) - if result.prompt_id in group_ids[-1]: - if include_moe_routing: - _record_shared_prefix_route_conflicts( - existing_group_ids=group_ids[-1], - existing_routes=moe_routes[-1], - result=result, - stats=moe_routing_pack_stats, - ) - result = result_without_prompt - token_ids[-1].extend(result.token_ids) - group_ids[-1].extend( - [result.prompt_id] * result.prompt_length - + [group_id] * (len(result.token_ids) - result.prompt_length) - ) - parent_ids[-1].extend([result.prompt_id] * len(result.token_ids)) - input_pos[-1].extend(result.input_pos) - assistant_mask[-1].extend(result.assistant_mask) - logprobs[-1].extend(result.logprobs) - advantages[-1].extend([result.advantage] * len(result.token_ids)) - weights[-1].extend([result.weight] * len(result.token_ids)) - if result.pixel_values is not None: - pixel_values[-1].append(result.pixel_values) - if result.image_grid_thw is not None: - image_grid_thw[-1].append(result.image_grid_thw) - if include_moe_routing: - assert result.moe_routed_experts is not None - moe_routes[-1].extend(result.moe_routed_experts) + rows.append([]) + rows[-1].append(item) if truncate_long_results: - token_ids[-1] = token_ids[-1][:seq_len] - group_ids[-1] = group_ids[-1][:seq_len] - parent_ids[-1] = parent_ids[-1][:seq_len] - input_pos[-1] = input_pos[-1][:seq_len] - assistant_mask[-1] = assistant_mask[-1][:seq_len] - logprobs[-1] = logprobs[-1][:seq_len] - advantages[-1] = advantages[-1][:seq_len] - weights[-1] = weights[-1][:seq_len] - if include_moe_routing: - moe_routes[-1] = moe_routes[-1][:seq_len] - - permutation = list(range(len(token_ids))) - random.shuffle(permutation) - token_ids = [token_ids[i] for i in permutation] - group_ids = [group_ids[i] for i in permutation] - parent_ids = [parent_ids[i] for i in permutation] - input_pos = [input_pos[i] for i in permutation] - assistant_mask = [assistant_mask[i] for i in permutation] - logprobs = [logprobs[i] for i in permutation] - advantages = [advantages[i] for i in permutation] - weights = [weights[i] for i in permutation] - pixel_values = [pixel_values[i] for i in permutation] - image_grid_thw = [image_grid_thw[i] for i in permutation] - moe_routes = [moe_routes[i] for i in permutation] + rows[-1][-1] = _truncate_prefix_tree_pack_item(rows[-1][-1], seq_len) + + random.shuffle(rows) + packed_rows = [ + _pack_prefix_tree_row( + row, + seq_len=seq_len, + pack_results=pack_results, + include_moe_routing=include_moe_routing, + moe_routing_pack_stats=moe_routing_pack_stats, + ) + for row in rows + ] def pad(values: list[list], pad_value) -> list[list]: max_len = seq_len @@ -158,6 +151,18 @@ def pad(values: list[list], pad_value) -> list[list]: value.extend([pad_value] * (max_len - len(value))) return values + token_ids = [row["token_ids"] for row in packed_rows] + group_ids = [row["group_ids"] for row in packed_rows] + parent_ids = [row["parent_ids"] for row in packed_rows] + input_pos = [row["input_pos"] for row in packed_rows] + assistant_mask = [row["assistant_mask"] for row in packed_rows] + logprobs = [row["logprobs"] for row in packed_rows] + advantages = [row["advantages"] for row in packed_rows] + weights = [row["weights"] for row in packed_rows] + pixel_values = [row["pixel_values"] for row in packed_rows] + image_grid_thw = [row["image_grid_thw"] for row in packed_rows] + moe_routes = [row["moe_routes"] for row in packed_rows] + assistant_mask_tensor = torch.tensor(pad(assistant_mask, 0), dtype=torch.bool) weights_tensor = torch.tensor(pad(weights, 0.0)) weights_tensor = torch.where( @@ -196,12 +201,8 @@ def pad(values: list[list], pad_value) -> list[list]: "logprobs": torch.tensor(pad(logprobs, float("nan"))), "advantages": advantages_tensor, "weights": weights_tensor, - "pixel_values": [ - torch.concat(tensors) if tensors else None for tensors in pixel_values - ], - "image_grid_thw": [ - torch.concat(tensors) if tensors else None for tensors in image_grid_thw - ], + "pixel_values": pixel_values, + "image_grid_thw": image_grid_thw, "moe_routing_replay": None, } if include_moe_routing: @@ -224,35 +225,219 @@ def pad(values: list[list], pad_value) -> list[list]: return packed_tensors -def _record_shared_prefix_route_conflicts( - *, - existing_group_ids: list[int], - existing_routes: list[TokenRoute | None], +def _prefix_tree_pack_item( result: TokenizedResult, + *, + seq_len: int, +) -> _PrefixTreePackItem: + shareable_length = min( + int(result.prompt_length), + max(_first_trainable_token_index(result) - 1, 0), + ) + item = _PrefixTreePackItem( + token_ids=tuple(int(value) for value in result.token_ids), + input_pos=tuple(int(value) for value in result.input_pos), + assistant_mask=tuple(int(value) for value in result.assistant_mask), + logprobs=tuple(float(value) for value in result.logprobs), + advantage=float(result.advantage), + weight=float(result.weight), + prompt_id=int(result.prompt_id), + shareable_length=shareable_length, + pixel_values=result.pixel_values, + image_grid_thw=result.image_grid_thw, + moe_routes=( + tuple(result.moe_routed_experts) + if result.moe_routed_experts is not None + else None + ), + ) + return _truncate_prefix_tree_pack_item(item, seq_len) + + +def _truncate_prefix_tree_pack_item( + item: _PrefixTreePackItem, + seq_len: int, +) -> _PrefixTreePackItem: + if len(item.token_ids) <= seq_len: + return item + return _PrefixTreePackItem( + token_ids=item.token_ids[:seq_len], + input_pos=item.input_pos[:seq_len], + assistant_mask=item.assistant_mask[:seq_len], + logprobs=item.logprobs[:seq_len], + advantage=item.advantage, + weight=item.weight, + prompt_id=item.prompt_id, + shareable_length=min(item.shareable_length, seq_len), + pixel_values=item.pixel_values, + image_grid_thw=item.image_grid_thw, + moe_routes=item.moe_routes[:seq_len] if item.moe_routes is not None else None, + ) + + +def _first_trainable_token_index(result: TokenizedResult) -> int: + return next( + ( + index + for index, (is_assistant, logprob) in enumerate( + zip(result.assistant_mask, result.logprobs, strict=True) + ) + if bool(is_assistant) or not math.isnan(float(logprob)) + ), + len(result.token_ids), + ) + + +def _packed_row_token_count( + row: list[_PrefixTreePackItem], + *, + seq_len: int, +) -> int: + if not row: + return 0 + count = estimate_prefix_tree_packed_tokens( + (torch.tensor(item.token_ids, dtype=torch.long) for item in row), + max_depth=seq_len, + shareable_lengths=(item.shareable_length for item in row), + ) + if count is None: + raise RuntimeError("CPU prefix-tree token estimate unexpectedly failed") + return count + + +def _pack_prefix_tree_row( + row: list[_PrefixTreePackItem], + *, + seq_len: int, + pack_results: bool, + include_moe_routing: bool, + moe_routing_pack_stats: MoeRoutingPackStats, +) -> _PackedPrefixTreeRow: + if not row: + return { + "token_ids": [], + "group_ids": [], + "parent_ids": [], + "input_pos": [], + "assistant_mask": [], + "logprobs": [], + "advantages": [], + "weights": [], + "pixel_values": None, + "image_grid_thw": None, + "moe_routes": [], + } + tree = _prefix_tree_pack_sequences( + (torch.tensor(item.token_ids, dtype=torch.long) for item in row), + max_depth=seq_len if pack_results else 0, + shareable_lengths=( + item.shareable_length if pack_results else 0 for item in row + ), + ) + token_ids = tree.tokens.reshape(-1).tolist() + assigned = [False] * len(token_ids) + input_pos = [0] * len(token_ids) + assistant_mask = [0] * len(token_ids) + logprobs = [float("nan")] * len(token_ids) + advantages = [0.0] * len(token_ids) + weights = [0.0] * len(token_ids) + moe_routes: list[TokenRoute | None] = [None] * len(token_ids) + for item_index, item in enumerate(row): + packed_positions = tree.positions_by_sequence[item_index].tolist() + for source_index, packed_index in enumerate(packed_positions): + _validate_prefix_tree_assignment( + item, + source_index=source_index, + packed_index=packed_index, + token_ids=token_ids, + input_pos=input_pos, + assigned=assigned, + ) + route = ( + item.moe_routes[source_index] + if item.moe_routes is not None and source_index < len(item.moe_routes) + else None + ) + if assigned[packed_index]: + if include_moe_routing: + _record_shared_route_conflict( + existing=moe_routes[packed_index], + candidate=route, + stats=moe_routing_pack_stats, + ) + continue + assigned[packed_index] = True + input_pos[packed_index] = item.input_pos[source_index] + assistant_mask[packed_index] = item.assistant_mask[source_index] + logprobs[packed_index] = item.logprobs[source_index] + advantages[packed_index] = item.advantage + weights[packed_index] = item.weight + moe_routes[packed_index] = route + return { + "token_ids": token_ids[:seq_len], + "group_ids": tree.group_ids.reshape(-1).tolist()[:seq_len], + "parent_ids": tree.parent_ids.reshape(-1).tolist()[:seq_len], + "input_pos": input_pos[:seq_len], + "assistant_mask": assistant_mask[:seq_len], + "logprobs": logprobs[:seq_len], + "advantages": advantages[:seq_len], + "weights": weights[:seq_len], + "pixel_values": _packed_row_tensor_list(row, "pixel_values"), + "image_grid_thw": _packed_row_tensor_list(row, "image_grid_thw"), + "moe_routes": moe_routes[:seq_len] if include_moe_routing else [], + } + + +def _validate_prefix_tree_assignment( + item: _PrefixTreePackItem, + *, + source_index: int, + packed_index: int, + token_ids: list[int], + input_pos: list[int], + assigned: list[bool], +) -> None: + if token_ids[packed_index] != item.token_ids[source_index]: + raise RuntimeError("Prefix-tree pack token assignment mismatch") + if not assigned[packed_index]: + return + if input_pos[packed_index] != item.input_pos[source_index]: + raise RuntimeError("Prefix-tree pack cannot share mismatched input positions") + if item.assistant_mask[source_index] or not math.isnan(item.logprobs[source_index]): + raise RuntimeError("Prefix-tree pack attempted to share a trainable token") + + +def _record_shared_route_conflict( + *, + existing: TokenRoute | None, + candidate: TokenRoute | None, stats: MoeRoutingPackStats, ) -> None: - assert result.moe_routed_experts is not None - prefix_positions = [ - index - for index, group_id in enumerate(existing_group_ids) - if group_id == result.prompt_id - ] - if len(prefix_positions) != result.prompt_length: - raise RuntimeError( - "Shared-prefix route comparison could not find the existing packed " - f"prefix rows: prompt_length={result.prompt_length}, " - f"existing_rows={len(prefix_positions)}" - ) - for prefix_offset, packed_index in enumerate(prefix_positions): - route = result.moe_routed_experts[prefix_offset] - existing = existing_routes[packed_index] - if route is None or existing is None: - raise RuntimeError("Shared-prefix MoE route is missing") - compared, conflicts = count_route_slot_conflicts(existing, route) - stats.shared_prefix_rows += 1 - stats.shared_prefix_compared_slots += compared - stats.shared_prefix_conflict_slots += conflicts - stats.shared_prefix_conflict_rows += int(conflicts > 0) + if existing is None or candidate is None: + raise RuntimeError("Prefix-tree MoE route is missing") + compared, conflicts = count_route_slot_conflicts(existing, candidate) + stats.prefix_tree_rows += 1 + stats.prefix_tree_compared_slots += compared + stats.prefix_tree_conflict_slots += conflicts + stats.prefix_tree_conflict_rows += int(conflicts > 0) + + +def _packed_row_tensor_list( + row: list[_PrefixTreePackItem], + attr: Literal["pixel_values", "image_grid_thw"], +) -> torch.Tensor | None: + tensors: list[torch.Tensor] = [] + seen_shared_prompts: set[int] = set() + for item in row: + tensor = getattr(item, attr) + if tensor is None: + continue + if item.shareable_length > 0: + if item.prompt_id in seen_shared_prompts: + continue + seen_shared_prompts.add(item.prompt_id) + tensors.append(tensor) + return torch.concat(tensors) if tensors else None def _tensorize_moe_routes( diff --git a/src/art/serverless/backend.py b/src/art/serverless/backend.py index f6a797a87..67659c90f 100644 --- a/src/art/serverless/backend.py +++ b/src/art/serverless/backend.py @@ -28,15 +28,16 @@ TrainConfig, TrainSFTConfig, ) +from ..utils import wandb_sdk from ..utils.record_provenance import record_provenance if TYPE_CHECKING: - import wandb + from wandb.sdk.artifacts.artifact import Artifact from ..model import Model, TrainableModel -def _extract_step_from_wandb_artifact(artifact: "wandb.Artifact") -> int | None: +def _extract_step_from_wandb_artifact(artifact: "Artifact") -> int | None: """Extract step number from a W&B artifact's aliases.""" for alias in artifact.aliases: if alias.startswith("step"): @@ -541,15 +542,13 @@ async def _train_sft( import tempfile import uuid - import wandb - from ..utils.sft import resolve_sft_batch_size assert model.id is not None, "Model ID is required" # Get the user's default entity from W&B if not set if model.entity is None: - api = wandb.Api(api_key=self._client.api_key) + api = wandb_sdk.api(api_key=self._client.api_key) model.entity = api.default_entity # Generate unique artifact name to avoid race conditions in distributed systems @@ -592,17 +591,17 @@ async def _train_sft( # Upload the file to W&B as a dataset artifact # Use the model's canonical run_id from database, or fall back to model name - run = wandb.init( + run = wandb_sdk.init( name=model.name, id=model.run_id or model.name, # Use stored run_id to match the canonical wandb run entity=model.entity, project=model.project, resume="allow", # Resume if this run already exists - settings=wandb.Settings(api_key=self._client.api_key), + settings=wandb_sdk.settings(api_key=self._client.api_key), ) try: - artifact = wandb.Artifact( + artifact = wandb_sdk.artifact( artifact_name, type="dataset", metadata={ @@ -735,12 +734,10 @@ async def _experimental_pull_model_checkpoint( import os import tempfile - import wandb - assert model.id is not None, "Model ID is required" # If entity is not set, use the user's default entity from W&B - api = wandb.Api(api_key=self._client.api_key) # ty:ignore[possibly-missing-attribute] + api = wandb_sdk.api(api_key=self._client.api_key) if model.entity is None: model.entity = api.default_entity if verbose: @@ -905,8 +902,6 @@ async def _experimental_fork_checkpoint( import os import tempfile - import wandb - from_project = from_project or model.project if from_s3_bucket is not None: @@ -962,7 +957,7 @@ async def _experimental_fork_checkpoint( selected_step = target_step else: # Pull from W&B artifacts - api = wandb.Api(api_key=self._client.api_key) # ty:ignore[possibly-missing-attribute] + api = wandb_sdk.api(api_key=self._client.api_key) from_entity = model.entity or api.default_entity # Iterate all artifact versions to find the best step. @@ -1012,17 +1007,17 @@ async def _experimental_fork_checkpoint( if verbose: print(f"Uploading forked checkpoint as W&B artifact for {model.name}...") - wandb.login(key=self._client.api_key) # ty:ignore[possibly-missing-attribute] - run = wandb.init( + wandb_sdk.login(key=self._client.api_key) + run = wandb_sdk.init( project=model.project, entity=model.entity, job_type="checkpoint-fork", name=f"fork-{from_model}-to-{model.name}", - settings=wandb.Settings(silent=True), + settings=wandb_sdk.settings(silent=True), ) assert run is not None - dest_artifact = wandb.Artifact(name=model.name, type="lora") + dest_artifact = wandb_sdk.artifact(name=model.name, type="lora") dest_artifact.add_dir(checkpoint_dir) aliases = ["latest"] if selected_step is not None: @@ -1031,7 +1026,7 @@ async def _experimental_fork_checkpoint( run.finish() # Copy provenance from the source model's W&B run to the destination model - api = wandb.Api(api_key=self._client.api_key) # ty:ignore[possibly-missing-attribute] + api = wandb_sdk.api(api_key=self._client.api_key) try: source_run = api.run(f"{model.entity}/{from_project}/{from_model}") source_provenance = source_run.config.get("wandb.provenance") diff --git a/src/art/utils/benchmarking/log_constant_metrics_wandb.py b/src/art/utils/benchmarking/log_constant_metrics_wandb.py index ada248105..6da9b07a8 100644 --- a/src/art/utils/benchmarking/log_constant_metrics_wandb.py +++ b/src/art/utils/benchmarking/log_constant_metrics_wandb.py @@ -1,8 +1,7 @@ """Utilities for logging constant baseline metrics to Weights & Biases.""" -import wandb - import art +from art.utils import wandb_sdk async def log_constant_metrics_wandb( @@ -30,7 +29,7 @@ async def log_constant_metrics_wandb( Example: `{"train": {"loss": 0.5}, "val": {"loss": 0.4, "accuracy": 0.8}}` """ - run = wandb.init( + run = wandb_sdk.init( project=model.project, name=logged_run_name if logged_run_name else model.name, reinit="create_new", diff --git a/src/art/utils/deployment/wandb.py b/src/art/utils/deployment/wandb.py index 9ddf778e8..8add1dd4f 100644 --- a/src/art/utils/deployment/wandb.py +++ b/src/art/utils/deployment/wandb.py @@ -5,6 +5,7 @@ from art.errors import UnsupportedBaseModelDeploymentError +from .. import wandb_sdk from .common import DeploymentConfig if TYPE_CHECKING: @@ -52,8 +53,6 @@ def deploy_wandb( Returns: The model name for inference: wandb-artifact:///{entity}/{project}/{name}:step{step} """ - import wandb - if model.base_model not in WANDB_SUPPORTED_BASE_MODELS: raise UnsupportedBaseModelDeploymentError( message=f"Base model {model.base_model} is not supported for serverless LoRA deployment by W&B. Supported models: {WANDB_SUPPORTED_BASE_MODELS}" @@ -64,23 +63,23 @@ def deploy_wandb( # Get the user's default entity from W&B if not set if model.entity is None: - api = wandb.Api() + api = wandb_sdk.api() model.entity = api.default_entity if verbose: print(f"Uploading checkpoint from {checkpoint_path} to W&B...") - run = wandb.init( + run = wandb_sdk.init( name=model.name + " (deployment)", entity=model.entity, project=model.project, - settings=wandb.Settings(api_key=os.environ["WANDB_API_KEY"]), + settings=wandb_sdk.settings(api_key=os.environ["WANDB_API_KEY"]), ) try: metadata: dict[str, object] = {"wandb.base_model": model.base_model} if config is not None: metadata["wandb.provenance"] = config.provenance - artifact = wandb.Artifact( + artifact = wandb_sdk.artifact( model.name, type="lora", metadata=metadata, diff --git a/src/art/utils/record_provenance.py b/src/art/utils/record_provenance.py index 84a8bce7f..9b202be35 100644 --- a/src/art/utils/record_provenance.py +++ b/src/art/utils/record_provenance.py @@ -3,18 +3,18 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - import wandb + from wandb.sdk.wandb_run import Run +from . import wandb_sdk -def record_provenance(run: wandb.Run, provenance: str) -> None: - """Record provenance on the latest artifact version's metadata.""" - import wandb as wandb_module - api = wandb_module.Api() +def record_provenance(run: Run, provenance: str) -> None: + """Record provenance on the latest artifact version's metadata.""" + api = wandb_sdk.api() artifact_path = f"{run.entity}/{run.project}/{run.name}:latest" try: artifact = api.artifact(artifact_path, type="lora") - except wandb_module.errors.CommError: + except wandb_sdk.comm_error_type(): return # No artifact exists yet existing = artifact.metadata.get("wandb.provenance") diff --git a/src/art/utils/wandb_sdk.py b/src/art/utils/wandb_sdk.py new file mode 100644 index 000000000..0c015c42a --- /dev/null +++ b/src/art/utils/wandb_sdk.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from wandb.apis.public import Api + from wandb.sdk.artifacts.artifact import Artifact + from wandb.sdk.wandb_run import Run + from wandb.sdk.wandb_settings import Settings + + +def api(*args: Any, **kwargs: Any) -> Api: + from wandb.apis.public import Api + + return Api(*args, **kwargs) + + +def artifact(*args: Any, **kwargs: Any) -> Artifact: + from wandb.sdk.artifacts.artifact import Artifact + + return Artifact(*args, **kwargs) + + +def init(*args: Any, **kwargs: Any) -> Run: + from wandb.sdk.wandb_init import init as wandb_init + + return wandb_init(*args, **kwargs) + + +def login(*args: Any, **kwargs: Any) -> Any: + from wandb.sdk.wandb_login import login as wandb_login + + return wandb_login(*args, **kwargs) + + +def settings(*args: Any, **kwargs: Any) -> Settings: + from wandb.sdk.wandb_settings import Settings + + return Settings(*args, **kwargs) + + +def comm_error_type() -> type[Exception]: + from wandb.errors import CommError + + return CommError diff --git a/tests/integration/megatron/cp_attn/test_attention_packed_vs_flattened.py b/tests/integration/megatron/cp_attn/test_attention_packed_vs_flattened.py index 3d3d51d4c..bed82490c 100644 --- a/tests/integration/megatron/cp_attn/test_attention_packed_vs_flattened.py +++ b/tests/integration/megatron/cp_attn/test_attention_packed_vs_flattened.py @@ -9,7 +9,7 @@ torch = pytest.importorskip("torch") from art.megatron.flex_attn.attention import FlexAttentionWrapper -from art.megatron.shared_prefix_state import create_shared_prefix_state +from art.megatron.prefix_tree_state import create_prefix_tree_state from tests.integration.megatron.gdn_shared_prefix.cases import default_phase0_cases from tests.integration.megatron.gdn_shared_prefix.metrics import ( GDN_CORRECTNESS_DTYPE, @@ -22,7 +22,7 @@ build_phase0_packed_tensors, ) from tests.integration.megatron.gdn_shared_prefix.parser_import import ( - parse_gdn_shared_prefix_segments, + parse_gdn_prefix_tree_segments, ) from tests.integration.megatron.model_support.oracle_harness import ( TEST_DEFAULT_FLEX_BACKEND, @@ -51,25 +51,23 @@ def _fp32_test_flex_backend(): @pytest.mark.skipif( not torch.cuda.is_available(), - reason="CUDA is required for compiled flex-attention shared-prefix coverage.", + reason="CUDA is required for compiled flex-attention prefix-tree coverage.", ) -def test_shared_prefix_attention_matches_flattened_grad_accumulation() -> None: +def test_prefix_tree_attention_matches_flattened_grad_accumulation() -> None: case = next( item for item in default_phase0_cases() if item.name == "multi_family_repeated" ) tensors = build_phase0_packed_tensors(case) group_ids = tensors["group_ids"].cuda() parent_ids = tensors["parent_ids"].cuda() - spec = parse_gdn_shared_prefix_segments( - group_ids.cpu(), parent_ids.cpu(), min_completions_per_family=1 - ) + spec = parse_gdn_prefix_tree_segments(group_ids.cpu(), parent_ids.cpu()) q, k, v = _attention_inputs(group_ids.shape, seed=20260425) q_ref = q.detach().clone().requires_grad_(True) k_ref = k.detach().clone().requires_grad_(True) v_ref = v.detach().clone().requires_grad_(True) output_grad = _packed_output_grad(spec, q.shape, seed=20260426) - attention_state = create_shared_prefix_state(group_ids, parent_ids) + attention_state = create_prefix_tree_state(group_ids, parent_ids) packed_out = FlexAttentionWrapper()( q, k, @@ -82,40 +80,19 @@ def test_shared_prefix_attention_matches_flattened_grad_accumulation() -> None: ref_out = torch.zeros_like(packed_out) ref_loss = q_ref.new_zeros(()) - for family in spec.families: - prefix = family.prefix - prefix_grad_used = False - for completion in family.completions: - indices = torch.tensor( - [ - *range(prefix.start, prefix.end), - *range(completion.start, completion.end), - ], - device=q.device, - dtype=torch.long, - ) - row = family.row_index - q_slice = q_ref[row : row + 1].index_select(2, indices) - k_slice = k_ref[row : row + 1].index_select(2, indices) - v_slice = v_ref[row : row + 1].index_select(2, indices) - flat_out = _dense_causal_attention(q_slice, k_slice, v_slice) - - ref_out[row, :, completion.start : completion.end] = flat_out[ - 0, :, prefix.length : - ] - flat_grad = torch.zeros_like(flat_out) - flat_grad[0, :, prefix.length :] = output_grad[ - row, :, completion.start : completion.end - ] - if not prefix_grad_used: - ref_out[row, :, prefix.start : prefix.end] = flat_out[ - 0, :, : prefix.length - ] - flat_grad[0, :, : prefix.length] = output_grad[ - row, :, prefix.start : prefix.end - ] - prefix_grad_used = True - ref_loss = ref_loss + (flat_out * flat_grad).sum() + for segment_index, segment in enumerate(spec.tree_segments): + indices, output_slice = _segment_context_positions(spec, segment_index) + index_tensor = torch.tensor(indices, device=q.device, dtype=torch.long) + row = segment.row_index + q_slice = q_ref[row : row + 1].index_select(2, index_tensor) + k_slice = k_ref[row : row + 1].index_select(2, index_tensor) + v_slice = v_ref[row : row + 1].index_select(2, index_tensor) + flat_out = _dense_causal_attention(q_slice, k_slice, v_slice) + + ref_out[row, :, segment.start : segment.end] = flat_out[0, :, output_slice] + flat_grad = torch.zeros_like(flat_out) + flat_grad[0, :, output_slice] = output_grad[row, :, segment.start : segment.end] + ref_loss = ref_loss + (flat_out * flat_grad).sum() ref_loss.backward() real_mask = _real_token_mask(spec, q.shape, device=q.device) @@ -133,7 +110,7 @@ def test_shared_prefix_attention_matches_flattened_grad_accumulation() -> None: @pytest.mark.skipif( not torch.cuda.is_available(), - reason="CUDA is required for compiled flex-attention shared-prefix coverage.", + reason="CUDA is required for compiled flex-attention prefix-tree coverage.", ) def test_physical_causal_attention_leaks_across_siblings() -> None: case = next( @@ -142,11 +119,9 @@ def test_physical_causal_attention_leaks_across_siblings() -> None: tensors = build_phase0_packed_tensors(case) group_ids = tensors["group_ids"].cuda() parent_ids = tensors["parent_ids"].cuda() - spec = parse_gdn_shared_prefix_segments( - group_ids.cpu(), parent_ids.cpu(), min_completions_per_family=1 - ) + spec = parse_gdn_prefix_tree_segments(group_ids.cpu(), parent_ids.cpu()) q, k, v = _attention_inputs(group_ids.shape, seed=20260427) - attention_state = create_shared_prefix_state(group_ids, parent_ids) + attention_state = create_prefix_tree_state(group_ids, parent_ids) packed_out = FlexAttentionWrapper()( q, k, @@ -225,11 +200,27 @@ def _completion_token_mask( spec: Any, shape: torch.Size, *, device: torch.device ) -> torch.Tensor: mask = torch.zeros(shape, device=device, dtype=torch.bool) - for family in spec.families: - for completion in family.completions: - mask[ - family.row_index, - :, - completion.start : completion.end, - ] = True + for index, segment in enumerate(spec.tree_segments): + if spec.tree_parent_indices[index] >= 0: + mask[segment.row_index, :, segment.start : segment.end] = True return mask + + +def _segment_context_positions( + spec: Any, segment_index: int +) -> tuple[list[int], slice]: + path = [] + cursor = segment_index + while cursor >= 0: + path.append(cursor) + cursor = spec.tree_parent_indices[cursor] + path.reverse() + positions = [ + position + for index in path + for position in range( + spec.tree_segments[index].start, spec.tree_segments[index].end + ) + ] + segment_length = spec.tree_segments[segment_index].length + return positions, slice(len(positions) - segment_length, len(positions)) diff --git a/tests/integration/megatron/gdn_shared_prefix/distributed_init.py b/tests/integration/megatron/gdn_shared_prefix/distributed_init.py new file mode 100644 index 000000000..b9b4075c8 --- /dev/null +++ b/tests/integration/megatron/gdn_shared_prefix/distributed_init.py @@ -0,0 +1,7 @@ +from pathlib import Path + + +def file_init_method(tmp_path: Path, name: str) -> str: + path = tmp_path / f"{name}.dist" + path.unlink(missing_ok=True) + return f"file://{path}" diff --git a/tests/integration/megatron/gdn_shared_prefix/layout_reference.py b/tests/integration/megatron/gdn_shared_prefix/layout_reference.py index 7369eaef7..ce3745354 100644 --- a/tests/integration/megatron/gdn_shared_prefix/layout_reference.py +++ b/tests/integration/megatron/gdn_shared_prefix/layout_reference.py @@ -7,9 +7,9 @@ from torch import Tensor from art.megatron.context_parallel.layout_index import TokenLayoutIndex -from art.megatron.gdn.gdn_shared_prefix import ( +from art.megatron.gdn.gdn_prefix_tree import ( GdnPackedExecutionSpec, - parse_gdn_shared_prefix_segments, + parse_gdn_prefix_tree_segments, ) from art.megatron.gdn.layout import ( GdnCpExchangePlan, @@ -19,7 +19,7 @@ class TestGdnCpLayoutPlan(BaseModel): - model_config = ConfigDict(frozen=True) + model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True) batch_size: int = Field(ge=1) sequence_length: int = Field(ge=1) @@ -39,9 +39,7 @@ def build_test_gdn_cp_layout_plan( gdn_token_ranges_by_rank: Sequence[Sequence[tuple[int, int, int]]] | None = None, device: torch.device | str | None = None, ) -> TestGdnCpLayoutPlan: - spec = parse_gdn_shared_prefix_segments( - group_ids, parent_ids, min_completions_per_family=0 - ) + spec = parse_gdn_prefix_tree_segments(group_ids, parent_ids) gdn_ranges = ( _normalize_rank_ranges(gdn_token_ranges_by_rank, cp_size=cp_size) if gdn_token_ranges_by_rank is not None @@ -89,7 +87,7 @@ def _build_full_exchange_plan( ) for transfer in local_plan.transfers: transfers.setdefault((transfer.source_rank, transfer.dest_rank), transfer) - return GdnCpExchangePlan.model_construct( + return GdnCpExchangePlan( cp_size=len(source_layout.token_counts_by_rank), source_token_counts_by_rank=source_layout.token_counts_by_rank, dest_token_counts_by_rank=tuple( @@ -133,7 +131,7 @@ def _split_gdn_token_ranges_by_rank( _segment_token_start(segment, spec.sequence_length), _segment_token_start(segment, spec.sequence_length) + segment.length, ) - for segment in spec.segments() + for segment in spec.tree_segments ), cp_size=cp_size, ) diff --git a/tests/integration/megatron/gdn_shared_prefix/oracles.py b/tests/integration/megatron/gdn_shared_prefix/oracles.py index 3d3f9ae12..e50887805 100644 --- a/tests/integration/megatron/gdn_shared_prefix/oracles.py +++ b/tests/integration/megatron/gdn_shared_prefix/oracles.py @@ -7,12 +7,14 @@ from torch import Tensor import torch.nn.functional as F +from art.megatron.gdn.gdn_prefix_tree import GdnPackedExecutionSpec, GdnSegmentSpec + from .metrics import ( mean_abs_pct, parameter_grad_mean_abs_pct_with_name, stable_output_mse_loss, ) -from .parser_import import parse_gdn_shared_prefix_segments +from .parser_import import parse_gdn_prefix_tree_segments class ToyGdnConfig(BaseModel): @@ -35,7 +37,7 @@ class ToyStatefulGdn(torch.nn.Module): """Small stateful block used to validate oracle mechanics on CPU. This is not a GDN approximation. It deliberately has the two state classes - that make GDN shared-prefix execution non-trivial: a finite conv tail and a + that make GDN prefix-tree execution non-trivial: a finite conv tail and a recurrent state. That is enough to prove parser routing, flattened accumulation, and known-bad physical-stream sensitivity before the real FLA kernels are invoked. @@ -107,27 +109,27 @@ def run_toy_packed( group_ids: Tensor, parent_ids: Tensor, ) -> Tensor: - spec = parse_gdn_shared_prefix_segments( - group_ids, parent_ids, min_completions_per_family=1 - ) + spec = parse_gdn_prefix_tree_segments(group_ids, parent_ids) output = torch.zeros_like(hidden) - for family in spec.families: - row = family.row_index - prefix_hidden = hidden[row, family.prefix.start : family.prefix.end] - prefix_out, prefix_conv, prefix_rec = module.forward_segment( - prefix_hidden, - conv_initial=module.zero_conv_state(hidden), - recurrent_initial=module.zero_recurrent_state(hidden), + conv_states: list[Tensor] = [] + rec_states: list[Tensor] = [] + for segment_index, segment in enumerate(spec.tree_segments): + row = segment.row_index + parent_index = spec.tree_parent_indices[segment_index] + if parent_index < 0: + conv_initial = module.zero_conv_state(hidden) + rec_initial = module.zero_recurrent_state(hidden) + else: + conv_initial = conv_states[parent_index] + rec_initial = rec_states[parent_index] + segment_out, conv_final, rec_final = module.forward_segment( + hidden[row, segment.start : segment.end], + conv_initial=conv_initial, + recurrent_initial=rec_initial, ) - output[row, family.prefix.start : family.prefix.end] = prefix_out - for completion in family.completions: - suffix_hidden = hidden[row, completion.start : completion.end] - suffix_out, _, _ = module.forward_segment( - suffix_hidden, - conv_initial=prefix_conv, - recurrent_initial=prefix_rec, - ) - output[row, completion.start : completion.end] = suffix_out + output[row, segment.start : segment.end] = segment_out + conv_states.append(conv_final) + rec_states.append(rec_final) return output @@ -138,30 +140,36 @@ def run_toy_flattened_reference( group_ids: Tensor, parent_ids: Tensor, ) -> Tensor: - spec = parse_gdn_shared_prefix_segments( - group_ids, parent_ids, min_completions_per_family=1 - ) + spec = parse_gdn_prefix_tree_segments(group_ids, parent_ids) output = torch.zeros_like(hidden) - for family in spec.families: - row = family.row_index - prefix_hidden = hidden[row, family.prefix.start : family.prefix.end] - prefix_len = family.prefix.length - for child_index, completion in enumerate(family.completions): - suffix_hidden = hidden[row, completion.start : completion.end] - flattened = torch.cat([prefix_hidden, suffix_hidden], dim=0) - flat_out, _, _ = module.forward_segment( - flattened, - conv_initial=module.zero_conv_state(hidden), - recurrent_initial=module.zero_recurrent_state(hidden), - ) - if child_index == 0: - output[row, family.prefix.start : family.prefix.end] = flat_out[ - :prefix_len - ] - output[row, completion.start : completion.end] = flat_out[prefix_len:] + for segment_index, segment in enumerate(spec.tree_segments): + path = _segment_path(spec, segment_index) + flattened = torch.cat( + [hidden[node.row_index, node.start : node.end] for node in path], + dim=0, + ) + flat_out, _, _ = module.forward_segment( + flattened, + conv_initial=module.zero_conv_state(hidden), + recurrent_initial=module.zero_recurrent_state(hidden), + ) + segment_len = segment.length + output[segment.row_index, segment.start : segment.end] = flat_out[-segment_len:] return output +def _segment_path( + spec: GdnPackedExecutionSpec, + segment_index: int, +) -> tuple[GdnSegmentSpec, ...]: + indices = [] + cursor = segment_index + while cursor >= 0: + indices.append(cursor) + cursor = spec.tree_parent_indices[cursor] + return tuple(spec.tree_segments[index] for index in reversed(indices)) + + def run_toy_physical_stream( module: ToyStatefulGdn, hidden: Tensor, diff --git a/tests/integration/megatron/gdn_shared_prefix/packed_layout.py b/tests/integration/megatron/gdn_shared_prefix/packed_layout.py index 45a41ff58..6175116a3 100644 --- a/tests/integration/megatron/gdn_shared_prefix/packed_layout.py +++ b/tests/integration/megatron/gdn_shared_prefix/packed_layout.py @@ -6,7 +6,7 @@ import torch from .cases import GdnPhase0Case -from .parser_import import GdnPackedExecutionSpec, parse_gdn_shared_prefix_segments +from .parser_import import GdnPackedExecutionSpec, parse_gdn_prefix_tree_segments class GdnCaseSummary(BaseModel): @@ -137,19 +137,23 @@ def summarize_case( conv_width: int, cp_sizes: tuple[int, ...] = (2, 4, 8), ) -> GdnCaseSummary: - spec = parse_gdn_shared_prefix_segments( - tensors["group_ids"], tensors["parent_ids"], min_completions_per_family=1 - ) + spec = parse_gdn_prefix_tree_segments(tensors["group_ids"], tensors["parent_ids"]) suffix_lengths = [ - segment.length for family in spec.families for segment in family.completions + segment.length + for index, segment in enumerate(spec.tree_segments) + if spec.tree_parent_indices[index] >= 0 ] boundary = _boundary_flags(spec, cp_sizes) return GdnCaseSummary( name=case.name, total_tokens=spec.real_token_count, family_count=spec.family_count, - completion_count=spec.completion_count, - max_segment_length=spec.max_segment_length, + completion_count=sum( + 1 for parent_index in spec.tree_parent_indices if parent_index >= 0 + ), + max_segment_length=max( + (segment.length for segment in spec.tree_segments), default=0 + ), suffix_shorter_than_conv=any(length < conv_width for length in suffix_lengths), suffix_equal_to_conv=any(length == conv_width for length in suffix_lengths), suffix_longer_than_conv=any(length > conv_width for length in suffix_lengths), @@ -227,19 +231,49 @@ def _boundary_flags( boundaries = {shard * rank for rank in range(1, cp_size)} if shard * (cp_size - 1) >= spec.real_token_count: flags["empty_trailing_rank"] = True - for family in spec.families: - family_start = _segment_real_start(family.prefix, spec, real_index) - family_end = _segment_real_end(family.completions[-1], spec, real_index) + for root in _root_segments(spec): + descendants = _descendant_segments(spec, root.family_index) + family_segments = (root, *descendants) + family_start = min( + _segment_real_start(segment, spec, real_index) + for segment in family_segments + ) + family_end = max( + _segment_real_end(segment, spec, real_index) + for segment in family_segments + ) if family_start in boundaries or family_end in boundaries: flags["family_boundary_at_partition"] = True - if _crosses_boundary(family.prefix, spec, real_index, boundaries): + if _crosses_boundary(root, spec, real_index, boundaries): flags["cp_boundary_prefix"] = True - for completion in family.completions: + for completion in descendants: if _crosses_boundary(completion, spec, real_index, boundaries): flags["cp_boundary_suffix"] = True return flags +def _root_segments(spec: GdnPackedExecutionSpec) -> tuple[Any, ...]: + return tuple( + segment + for index, segment in enumerate(spec.tree_segments) + if spec.tree_parent_indices[index] < 0 + ) + + +def _descendant_segments( + spec: GdnPackedExecutionSpec, root_index: int +) -> tuple[Any, ...]: + descendants = [] + for index, segment in enumerate(spec.tree_segments): + parent = spec.tree_parent_indices[index] + while parent >= 0: + if parent == root_index: + descendants.append(segment) + break + parent = spec.tree_parent_indices[parent] + return tuple(descendants) + + def _segment_real_start( segment: Any, spec: GdnPackedExecutionSpec, real_index: dict[int, int] ) -> int: diff --git a/tests/integration/megatron/gdn_shared_prefix/parser_import.py b/tests/integration/megatron/gdn_shared_prefix/parser_import.py index ce184d96e..333487471 100644 --- a/tests/integration/megatron/gdn_shared_prefix/parser_import.py +++ b/tests/integration/megatron/gdn_shared_prefix/parser_import.py @@ -9,9 +9,9 @@ def _load_parser_module() -> ModuleType: repo_root = Path(__file__).resolve().parents[4] - module_path = repo_root / "src/art/megatron/gdn/gdn_shared_prefix.py" + module_path = repo_root / "src/art/megatron/gdn/gdn_prefix_tree.py" spec = importlib.util.spec_from_file_location( - "_art_gdn_shared_prefix_for_tests", module_path + "_art_gdn_prefix_tree_for_tests", module_path ) if spec is None or spec.loader is None: raise RuntimeError(f"Failed to load parser module from {module_path}") @@ -24,6 +24,5 @@ def _load_parser_module() -> ModuleType: _MODULE = _load_parser_module() GdnPackedExecutionSpec: Any = _MODULE.GdnPackedExecutionSpec -build_gdn_cp_segment_schedule: Any = _MODULE.build_gdn_cp_segment_schedule build_gdn_rank_execution_plan: Any = _MODULE.build_gdn_rank_execution_plan -parse_gdn_shared_prefix_segments: Any = _MODULE.parse_gdn_shared_prefix_segments +parse_gdn_prefix_tree_segments: Any = _MODULE.parse_gdn_prefix_tree_segments diff --git a/tests/integration/megatron/gdn_shared_prefix/real_gdn_oracle.py b/tests/integration/megatron/gdn_shared_prefix/real_gdn_oracle.py index e69fef22b..78628bf17 100644 --- a/tests/integration/megatron/gdn_shared_prefix/real_gdn_oracle.py +++ b/tests/integration/megatron/gdn_shared_prefix/real_gdn_oracle.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, Literal +from typing import Any, Literal, NamedTuple from pydantic import BaseModel, ConfigDict import torch @@ -8,7 +8,7 @@ import torch.nn.functional as F from art.megatron.context_parallel.layout_index import TokenLayoutIndex -from art.megatron.gdn.gdn_shared_prefix import FLA_CHUNK_SIZE +from art.megatron.gdn.gdn_prefix_tree import FLA_CHUNK_SIZE from art.megatron.gdn.operator import ( _apply_gated_rms_norm, _chunk_gated_delta_rule, @@ -21,7 +21,7 @@ _out_proj, _zero_conv_state, _zero_recurrent_state, - gdn_shared_prefix_forward, + gdn_prefix_tree_forward, ) from .layout_reference import build_test_gdn_cp_layout_plan @@ -30,7 +30,7 @@ parameter_grad_mean_abs_pct_with_name, stable_output_mse_loss, ) -from .parser_import import parse_gdn_shared_prefix_segments +from .parser_import import parse_gdn_prefix_tree_segments class RealGdnOracleMetrics(BaseModel): @@ -61,6 +61,57 @@ class GdnChainBoundaryDebug(BaseModel): ] +class _TreeFamily(NamedTuple): + row_index: int + family_index: int + prefix: Any + completions: tuple[Any, ...] + segment_indices: tuple[int, ...] + parent_indices: tuple[int, ...] + + @property + def token_count(self) -> int: + return self.prefix.length + sum(segment.length for segment in self.completions) + + +def _segment_path(spec: Any, segment_index: int) -> tuple[Any, ...]: + path = [] + cursor = segment_index + while cursor >= 0: + path.append(cursor) + cursor = spec.tree_parent_indices[cursor] + return tuple(spec.tree_segments[index] for index in reversed(path)) + + +def _tree_families(spec: Any) -> tuple[_TreeFamily, ...]: + families = [] + for root_index, root in enumerate(spec.tree_segments): + if spec.tree_parent_indices[root_index] >= 0: + continue + segment_indices = [root_index] + for index in range(root_index + 1, len(spec.tree_segments)): + parent = spec.tree_parent_indices[index] + while parent >= 0: + if parent == root_index: + segment_indices.append(index) + break + parent = spec.tree_parent_indices[parent] + segments = tuple(spec.tree_segments[index] for index in segment_indices) + families.append( + _TreeFamily( + row_index=root.row_index, + family_index=root_index, + prefix=root, + completions=segments[1:], + segment_indices=tuple(segment_indices), + parent_indices=tuple( + spec.tree_parent_indices[index] for index in segment_indices + ), + ) + ) + return tuple(families) + + def compare_real_gdn_cp1_to_flattened( *, packed_gdn: Any, @@ -73,7 +124,7 @@ def compare_real_gdn_cp1_to_flattened( packed_hidden = hidden_states.clone().detach().requires_grad_(True) flat_hidden = hidden_states.clone().detach().requires_grad_(True) - packed_out, _ = gdn_shared_prefix_forward( + packed_out, _ = gdn_prefix_tree_forward( packed_gdn, packed_hidden, group_ids=group_ids, @@ -118,7 +169,7 @@ def compare_real_gdn_cp1_to_flattened_with_output_grad( packed_hidden = hidden_states.clone().detach().requires_grad_(True) flat_hidden = hidden_states.clone().detach().requires_grad_(True) - packed_out, _ = gdn_shared_prefix_forward( + packed_out, _ = gdn_prefix_tree_forward( packed_gdn, packed_hidden, group_ids=group_ids, @@ -296,35 +347,34 @@ def run_real_gdn_flattened_reference( parent_ids: Tensor, execution_spec: Any | None = None, ) -> Tensor: - spec = execution_spec or parse_gdn_shared_prefix_segments( - group_ids, parent_ids, min_completions_per_family=1 - ) + spec = execution_spec or parse_gdn_prefix_tree_segments(group_ids, parent_ids) output = torch.zeros_like(hidden_states) - for family in spec.families: - row = family.row_index - prefix_hidden = hidden_states[ - family.prefix.start : family.prefix.end, row : row + 1, : - ] - prefix_len = family.prefix.length - for child_index, completion in enumerate(family.completions): - suffix_hidden = hidden_states[ - completion.start : completion.end, row : row + 1, : - ] - flat_hidden = torch.cat([prefix_hidden, suffix_hidden], dim=0) - flat_out, _, _, _ = _run_gdn_segment( - gdn, - flat_hidden, - conv_initial=_zero_conv_state(gdn, hidden_states, row), - recurrent_initial=_zero_recurrent_state(gdn, hidden_states, row), - output_final_state=False, - ) - if child_index == 0: - output[family.prefix.start : family.prefix.end, row : row + 1, :] = ( - flat_out[:prefix_len] - ) - output[completion.start : completion.end, row : row + 1, :] = flat_out[ - prefix_len: - ] + for segment_index, segment in enumerate(spec.tree_segments): + flat_hidden = torch.cat( + [ + hidden_states[ + node.start : node.end, + node.row_index : node.row_index + 1, + :, + ] + for node in _segment_path(spec, segment_index) + ], + dim=0, + ) + flat_out, _, _, _ = _run_gdn_segment( + gdn, + flat_hidden, + conv_initial=_zero_conv_state(gdn, hidden_states, segment.row_index), + recurrent_initial=_zero_recurrent_state( + gdn, hidden_states, segment.row_index + ), + output_final_state=False, + ) + output[ + segment.start : segment.end, + segment.row_index : segment.row_index + 1, + :, + ] = flat_out[-segment.length :] return output @@ -359,9 +409,7 @@ def run_real_gdn_local_fork_reference( cp_size: int, attention_token_layout_index: TokenLayoutIndex | None = None, ) -> Tensor: - spec = parse_gdn_shared_prefix_segments( - group_ids, parent_ids, min_completions_per_family=0 - ) + spec = parse_gdn_prefix_tree_segments(group_ids, parent_ids) gdn_token_indices_by_rank = _split_gdn_families_by_rank(spec, cp_size=cp_size) gdn_token_ranges_by_rank = _rank_ranges_from_tokens_by_rank( gdn_token_indices_by_rank @@ -414,12 +462,12 @@ def _split_gdn_families_by_rank( raise ValueError(f"cp_size must be >= 1, got {cp_size}") ranks: list[list[int]] = [[] for _ in range(cp_size)] loads = [0] * cp_size - for family in spec.families: + for family in _tree_families(spec): rank = min(range(cp_size), key=lambda index: (loads[index], index)) family_tokens = tuple( token for segment in (family.prefix, *family.completions) - for token in segment.linear_indices(spec.sequence_length) + for token in _segment_linear_indices(segment, spec.sequence_length) ) ranks[rank].extend(family_tokens) loads[rank] += len(family_tokens) @@ -471,6 +519,11 @@ def _simulate_all_to_all_single( return tuple(outputs) +def _segment_linear_indices(segment: Any, sequence_length: int) -> range: + base = int(segment.row_index) * int(sequence_length) + return range(base + int(segment.start), base + int(segment.end)) + + def _transfer_positions(tensor: Tensor | None, *, count: int) -> tuple[int, ...]: if tensor is None: return tuple(range(count)) @@ -523,11 +576,9 @@ def run_real_gdn_suffix_only_chain_reference( mutation: GdnChainMutation | None = None, boundary_debug: list[GdnChainBoundaryDebug] | None = None, ) -> Tensor: - spec = parse_gdn_shared_prefix_segments( - group_ids, parent_ids, min_completions_per_family=0 - ) + spec = parse_gdn_prefix_tree_segments(group_ids, parent_ids) output = torch.zeros_like(hidden_states) - for family in spec.families: + for family in _tree_families(spec): row = family.row_index zero_conv = _zero_conv_state(gdn, hidden_states, batch_size=1) zero_rec = _zero_recurrent_state(gdn, hidden_states, batch_size=1) @@ -575,11 +626,9 @@ def run_real_gdn_chunk_native_reference( group_ids: Tensor, parent_ids: Tensor, ) -> Tensor: - spec = parse_gdn_shared_prefix_segments( - group_ids, parent_ids, min_completions_per_family=0 - ) + spec = parse_gdn_prefix_tree_segments(group_ids, parent_ids) output = torch.zeros_like(hidden_states) - for family in spec.families: + for family in _tree_families(spec): _scatter_family_output( output, family, @@ -597,13 +646,11 @@ def run_real_gdn_mixed_cp_reference( cp_size: int, local_fork_max_tokens: int, ) -> Tensor: - spec = parse_gdn_shared_prefix_segments( - group_ids, parent_ids, min_completions_per_family=0 - ) + spec = parse_gdn_prefix_tree_segments(group_ids, parent_ids) output = torch.zeros_like(hidden_states) local_count = 0 chain_count = 0 - for family in spec.families: + for family in _tree_families(spec): if family.token_count <= local_fork_max_tokens: local_count += 1 _scatter_family_output( @@ -700,7 +747,7 @@ def _run_local_fork_rank( local_group_ids, local_parent_ids = _local_fork_group_tensors( spec, local_token_indices, device=rank_hidden.device ) - local_output, _ = gdn_shared_prefix_forward( + local_output, _ = gdn_prefix_tree_forward( gdn, rank_hidden.unsqueeze(1).contiguous(), group_ids=local_group_ids, @@ -726,7 +773,7 @@ def _run_gdn_family_local_fork( local_group_ids, local_parent_ids = _family_group_tensors( family, device=hidden_states.device ) - local_output, _ = gdn_shared_prefix_forward( + local_output, _ = gdn_prefix_tree_forward( gdn, local_hidden, group_ids=local_group_ids, @@ -753,14 +800,21 @@ def _family_group_tensors( ) -> tuple[Tensor, Tensor]: group_ids = [] parent_ids = [] - prefix_group_id = 0 - group_ids.extend([prefix_group_id] * family.prefix.length) - parent_ids.extend([prefix_group_id] * family.prefix.length) - next_group_id = 1 - for completion in family.completions: - group_ids.extend([next_group_id] * completion.length) - parent_ids.extend([prefix_group_id] * completion.length) - next_group_id += 1 + local_group_by_global: dict[int, int] = {} + for local_group_id, (segment, global_index, parent_index) in enumerate( + zip( + (family.prefix, *family.completions), + family.segment_indices, + family.parent_indices, + strict=True, + ) + ): + local_group_by_global[global_index] = local_group_id + local_parent_id = ( + local_group_id if parent_index < 0 else local_group_by_global[parent_index] + ) + group_ids.extend([local_group_id] * segment.length) + parent_ids.extend([local_parent_id] * segment.length) return ( torch.tensor([group_ids], device=device, dtype=torch.long), torch.tensor([parent_ids], device=device, dtype=torch.long), @@ -883,12 +937,12 @@ def _local_fork_group_tensors( ) parent_ids = torch.full_like(group_ids, -1) next_group_id = 0 - for family in spec.families: + for family in _tree_families(spec): family_segments = (family.prefix, *family.completions) family_tokens = tuple( token_index for segment in family_segments - for token_index in segment.linear_indices(spec.sequence_length) + for token_index in _segment_linear_indices(segment, spec.sequence_length) ) token_is_local = tuple( token_index in local_position for token_index in family_tokens @@ -898,19 +952,23 @@ def _local_fork_group_tensors( if not all(token_is_local): raise ValueError("local-fork execution requires whole prompt families") - prefix_group_id = next_group_id - next_group_id += 1 - for token_index in family.prefix.linear_indices(spec.sequence_length): - position = local_position[token_index] - group_ids[position] = prefix_group_id - parent_ids[position] = prefix_group_id - for completion in family.completions: - child_group_id = next_group_id + group_by_segment_index: dict[int, int] = {} + for segment, global_index, parent_index in zip( + family_segments, + family.segment_indices, + family.parent_indices, + strict=True, + ): + group_id = next_group_id next_group_id += 1 - for token_index in completion.linear_indices(spec.sequence_length): + group_by_segment_index[global_index] = group_id + parent_group_id = ( + group_id if parent_index < 0 else group_by_segment_index[parent_index] + ) + for token_index in _segment_linear_indices(segment, spec.sequence_length): position = local_position[token_index] - group_ids[position] = child_group_id - parent_ids[position] = prefix_group_id + group_ids[position] = group_id + parent_ids[position] = parent_group_id if torch.any(group_ids == -1): raise RuntimeError("local-fork metadata left unassigned token rows") return group_ids.unsqueeze(0), parent_ids.unsqueeze(0) diff --git a/tests/integration/megatron/gdn_shared_prefix/test_fla_cp_native_recurrent.py b/tests/integration/megatron/gdn_shared_prefix/test_fla_cp_native_recurrent.py index bcf3a0cfb..6f5eefc17 100644 --- a/tests/integration/megatron/gdn_shared_prefix/test_fla_cp_native_recurrent.py +++ b/tests/integration/megatron/gdn_shared_prefix/test_fla_cp_native_recurrent.py @@ -1,7 +1,6 @@ from __future__ import annotations from pathlib import Path -import socket from typing import Any, cast import pytest @@ -20,6 +19,7 @@ chunk_gated_delta_rule_native_cp, ) +from .distributed_init import file_init_method # noqa: E402 from .metrics import GDN_CORRECTNESS_DTYPE, assert_mean_abs_pct # noqa: E402 _CP_SIZES = ( @@ -43,10 +43,10 @@ def test_native_fla_cp_recurrent_matches_single_rank( cp_size: int, tmp_path: Path ) -> None: - port = _find_free_port() + init_method = file_init_method(tmp_path, f"native_fla_recurrent_cp{cp_size}") mp.spawn( _native_fla_cp_worker, - args=(cp_size, port, str(tmp_path)), + args=(cp_size, init_method, str(tmp_path)), nprocs=cp_size, join=True, ) @@ -62,10 +62,10 @@ def test_native_fla_cp_recurrent_matches_single_rank( def test_native_fla_cp_recurrent_varlen_multichain_matches_single_rank( cp_size: int, tmp_path: Path ) -> None: - port = _find_free_port() + init_method = file_init_method(tmp_path, f"native_fla_varlen_cp{cp_size}") mp.spawn( _native_fla_cp_varlen_multichain_worker, - args=(cp_size, port, str(tmp_path)), + args=(cp_size, init_method, str(tmp_path)), nprocs=cp_size, join=True, ) @@ -119,13 +119,13 @@ def test_native_fla_summary_affine_debug_matches_final_state() -> None: def _native_fla_cp_worker( rank: int, cp_size: int, - port: int, + init_method: str, output_dir: str, ) -> None: torch.cuda.set_device(rank) init_process_group( backend="nccl", - init_method=f"tcp://127.0.0.1:{port}", + init_method=init_method, rank=rank, world_size=cp_size, ) @@ -201,13 +201,13 @@ def _native_fla_cp_worker( def _native_fla_cp_varlen_multichain_worker( rank: int, cp_size: int, - port: int, + init_method: str, output_dir: str, ) -> None: torch.cuda.set_device(rank) init_process_group( backend="nccl", - init_method=f"tcp://127.0.0.1:{port}", + init_method=init_method, rank=rank, world_size=cp_size, ) @@ -519,9 +519,3 @@ def _cat_varlen_slices( def _assert_grad_close(left: torch.Tensor, right_grad: torch.Tensor, name: str) -> None: assert left.grad is not None, name assert_mean_abs_pct(right_grad, left.grad, name) - - -def _find_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/integration/megatron/gdn_shared_prefix/test_gdn_cp_packed_correctness.py b/tests/integration/megatron/gdn_shared_prefix/test_gdn_cp_packed_correctness.py index 2151b41e1..875756162 100644 --- a/tests/integration/megatron/gdn_shared_prefix/test_gdn_cp_packed_correctness.py +++ b/tests/integration/megatron/gdn_shared_prefix/test_gdn_cp_packed_correctness.py @@ -2,7 +2,6 @@ from collections.abc import Callable from pathlib import Path -import socket from typing import Any import pytest @@ -15,12 +14,14 @@ from torch.distributed import destroy_process_group, init_process_group # noqa: E402 import torch.multiprocessing as mp # noqa: E402 -from art.megatron.gdn.gdn_shared_prefix import ( # noqa: E402 +from art.megatron.context_parallel.layout_index import TokenLayoutIndex # noqa: E402 +from art.megatron.gdn.gdn_prefix_tree import ( # noqa: E402 GdnPlannerConfig, build_gdn_rank_execution_plan, - parse_gdn_shared_prefix_segments, + parse_gdn_prefix_tree_segments, ) from art.megatron.gdn.operator import run_gdn_layer # noqa: E402 +from art.megatron.prefix_tree_packing import prefix_tree_pack # noqa: E402 from .cases import ( # noqa: E402 GdnFamilyShape, @@ -29,6 +30,7 @@ default_phase0_cases, ) from .distributed_grad import all_reduce_parameter_grads_coalesced # noqa: E402 +from .distributed_init import file_init_method # noqa: E402 from .metrics import ( # noqa: E402 GDN_CORRECTNESS_DTYPE, REAL_GDN_GRAD_MEAN_ABS_PCT_THRESHOLD, @@ -50,10 +52,10 @@ def test_gdn_cp_packed_matches_cp1_oracle_all_edge_cases( cp_size: int, tmp_path: Path ) -> None: _skip_without_gpus(cp_size) - port = _find_free_port() + init_method = file_init_method(tmp_path, f"cp1_oracle_cp{cp_size}") mp.spawn( _cp1_oracle_worker, - args=(cp_size, port, str(tmp_path), False), + args=(cp_size, init_method, str(tmp_path), False), nprocs=cp_size, join=True, ) @@ -66,10 +68,10 @@ def test_gdn_cp_packed_sibling_order_matches_cp1_oracle( cp_size: int, tmp_path: Path ) -> None: _skip_without_gpus(cp_size) - port = _find_free_port() + init_method = file_init_method(tmp_path, f"cp1_oracle_sibling_cp{cp_size}") mp.spawn( _cp1_oracle_worker, - args=(cp_size, port, str(tmp_path), True), + args=(cp_size, init_method, str(tmp_path), True), nprocs=cp_size, join=True, ) @@ -77,17 +79,61 @@ def test_gdn_cp_packed_sibling_order_matches_cp1_oracle( assert (tmp_path / f"cp1_oracle_sibling_rank_{rank}.ok").read_text() == "ok\n" +@pytest.mark.parametrize("cp_size", (2, 4)) +def test_gdn_cp_tree_chain_matches_cp1_oracle(cp_size: int, tmp_path: Path) -> None: + _skip_without_gpus(cp_size) + init_method = file_init_method(tmp_path, f"tree_chain_cp{cp_size}") + mp.spawn( + _tree_chain_oracle_worker, + args=(cp_size, init_method, str(tmp_path)), + nprocs=cp_size, + join=True, + ) + for rank in range(cp_size): + assert (tmp_path / f"tree_chain_rank_{rank}.ok").read_text() == "ok\n" + + +def test_gdn_cp_tree_fuzz_matches_cp1_oracle(tmp_path: Path) -> None: + cp_size = 4 + _skip_without_gpus(cp_size) + init_method = file_init_method(tmp_path, "tree_fuzz_cp4") + mp.spawn( + _tree_fuzz_oracle_worker, + args=(cp_size, init_method, str(tmp_path)), + nprocs=cp_size, + join=True, + ) + for rank in range(cp_size): + assert (tmp_path / f"tree_fuzz_rank_{rank}.ok").read_text() == "ok\n" + + +@pytest.mark.parametrize("cp_size", (2, 4)) +def test_gdn_cp_tree_trainability_updates_parameters( + cp_size: int, tmp_path: Path +) -> None: + _skip_without_gpus(cp_size) + init_method = file_init_method(tmp_path, f"tree_trainability_cp{cp_size}") + mp.spawn( + _tree_trainability_worker, + args=(cp_size, init_method, str(tmp_path)), + nprocs=cp_size, + join=True, + ) + for rank in range(cp_size): + assert (tmp_path / f"tree_trainability_rank_{rank}.ok").read_text() == "ok\n" + + def _cp1_oracle_worker( rank: int, cp_size: int, - port: int, + init_method: str, output_dir: str, sibling_only: bool, ) -> None: torch.cuda.set_device(rank) init_process_group( backend="nccl", - init_method=f"tcp://127.0.0.1:{port}", + init_method=init_method, rank=rank, world_size=cp_size, ) @@ -126,6 +172,167 @@ def _cp1_oracle_worker( destroy_process_group() +def _tree_chain_oracle_worker( + rank: int, + cp_size: int, + init_method: str, + output_dir: str, +) -> None: + torch.cuda.set_device(rank) + init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=cp_size, + ) + try: + ps.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=cp_size, + expert_model_parallel_size=1, + ) + ref_gdn, cp_gdn = _make_matching_gdn_pair(cp_size=cp_size) + _assert_tree_pack_matches_cp1( + "tree_chain", + ref_gdn, + cp_gdn, + _tree_chain_pack(), + rank=rank, + cp_size=cp_size, + seed=9090, + planner_config=_tree_chain_planner_config(), + require_chain=True, + ) + Path(output_dir, f"tree_chain_rank_{rank}.ok").write_text("ok\n") + finally: + if getattr(ps, "model_parallel_is_initialized", lambda: False)(): + ps.destroy_model_parallel() + destroy_process_group() + + +def _tree_fuzz_oracle_worker( + rank: int, + cp_size: int, + init_method: str, + output_dir: str, +) -> None: + torch.cuda.set_device(rank) + init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=cp_size, + ) + try: + ps.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=cp_size, + expert_model_parallel_size=1, + ) + ref_gdn, cp_gdn = _make_matching_gdn_pair(cp_size=cp_size) + for case_index, (name, pack) in enumerate(_tree_fuzz_packs()): + _assert_tree_pack_matches_cp1( + name, + ref_gdn, + cp_gdn, + pack, + rank=rank, + cp_size=cp_size, + seed=9190 + case_index, + planner_config=_tree_fuzz_planner_config(), + require_chain=False, + ) + torch.distributed.barrier() + Path(output_dir, f"tree_fuzz_rank_{rank}.ok").write_text("ok\n") + finally: + if getattr(ps, "model_parallel_is_initialized", lambda: False)(): + ps.destroy_model_parallel() + destroy_process_group() + + +def _tree_trainability_worker( + rank: int, + cp_size: int, + init_method: str, + output_dir: str, +) -> None: + torch.cuda.set_device(rank) + init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=cp_size, + ) + try: + ps.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=cp_size, + expert_model_parallel_size=1, + ) + _, cp_gdn = _make_matching_gdn_pair( + cp_size=cp_size, + params_dtype=GDN_CORRECTNESS_DTYPE, + ) + pack = _tree_trainability_pack() + group_ids = pack.group_ids.cuda() + parent_ids = pack.parent_ids.cuda() + spec = parse_gdn_prefix_tree_segments(group_ids, parent_ids) + plan = build_gdn_rank_execution_plan( + spec, + device=group_ids.device, + cp_rank=rank, + cp_size=cp_size, + attention_token_layout_index=_uniform_attention_layout( + spec.real_token_count, cp_size + ), + planner_config=_tree_chain_planner_config(), + ) + hidden = _tree_trainability_hidden(spec.real_token_count, cp_size=cp_size) + flat_hidden = hidden.transpose(0, 1).reshape(-1, hidden.shape[-1]) + local_index = torch.tensor( + plan.attention_token_indices, device=hidden.device, dtype=torch.long + ) + local_hidden = flat_hidden.index_select(0, local_index).unsqueeze(1) + optimizer = torch.optim.SGD(cp_gdn.parameters(), lr=5e-2) + + optimizer.zero_grad(set_to_none=True) + loss = _tree_training_local_loss( + cp_gdn, + local_hidden, + group_ids, + parent_ids, + spec, + plan, + ) + loss.backward() + all_reduce_parameter_grads_coalesced(cp_gdn) + grad_norm = _parameter_l1_norm( + parameter.grad + for parameter in cp_gdn.parameters() + if parameter.grad is not None + ) + assert torch.isfinite(grad_norm) + assert float(grad_norm.item()) > 0.0 + before_step = [ + parameter.detach().float().clone() for parameter in cp_gdn.parameters() + ] + optimizer.step() + update_norm = _parameter_l1_norm( + parameter.detach().float() - before + for parameter, before in zip(cp_gdn.parameters(), before_step, strict=True) + ) + assert torch.isfinite(update_norm) + assert float(update_norm.item()) > 0.0 + Path(output_dir, f"tree_trainability_rank_{rank}.ok").write_text("ok\n") + finally: + if getattr(ps, "model_parallel_is_initialized", lambda: False)(): + ps.destroy_model_parallel() + destroy_process_group() + + def _assert_case_matches_cp1( ref_gdn: torch.nn.Module, cp_gdn: torch.nn.Module, @@ -141,14 +348,15 @@ def _assert_case_matches_cp1( tensors = build_phase0_packed_tensors(case) group_ids = tensors["group_ids"].cuda() parent_ids = tensors["parent_ids"].cuda() - spec = parse_gdn_shared_prefix_segments( - group_ids, parent_ids, min_completions_per_family=0 - ) + spec = parse_gdn_prefix_tree_segments(group_ids, parent_ids) plan = build_gdn_rank_execution_plan( spec, device=group_ids.device, cp_rank=rank, cp_size=cp_size, + attention_token_layout_index=_uniform_attention_layout( + spec.real_token_count, cp_size + ), planner_config=planner_config or GdnPlannerConfig(), ) hidden, output_grad = _hidden_and_grad(case, seed=seed) @@ -212,6 +420,84 @@ def _assert_case_matches_cp1( ) +def _assert_tree_pack_matches_cp1( + name: str, + ref_gdn: torch.nn.Module, + cp_gdn: torch.nn.Module, + pack: Any, + *, + rank: int, + cp_size: int, + seed: int, + planner_config: GdnPlannerConfig, + require_chain: bool, +) -> None: + zero_parameter_grads(ref_gdn) + zero_parameter_grads(cp_gdn) + group_ids = pack.group_ids.cuda() + parent_ids = pack.parent_ids.cuda() + spec = parse_gdn_prefix_tree_segments(group_ids, parent_ids) + plan = build_gdn_rank_execution_plan( + spec, + device=group_ids.device, + cp_rank=rank, + cp_size=cp_size, + attention_token_layout_index=_uniform_attention_layout( + spec.real_token_count, cp_size + ), + planner_config=planner_config, + ) + if require_chain: + assert any(plan.tree_chain_buckets_by_depth) + hidden, output_grad = _tree_hidden_and_grad(spec.real_token_count, seed=seed) + ref_hidden = hidden.clone().detach().requires_grad_(True) + ref_out, _ = run_gdn_layer( + ref_gdn, + ref_hidden, + group_ids=group_ids, + parent_ids=parent_ids, + ) + ref_loss = (ref_out * output_grad).sum() + ref_loss.backward() + + flat_hidden = hidden.transpose(0, 1).reshape(-1, hidden.shape[-1]) + flat_grad = output_grad.transpose(0, 1).reshape(-1, output_grad.shape[-1]) + local_index = torch.tensor( + plan.attention_token_indices, device=hidden.device, dtype=torch.long + ) + local_hidden = ( + flat_hidden.index_select(0, local_index) + .unsqueeze(1) + .contiguous() + .detach() + .requires_grad_(True) + ) + local_output_grad = flat_grad.index_select(0, local_index).unsqueeze(1).contiguous() + cp_out, _ = run_gdn_layer( + cp_gdn, + local_hidden, + group_ids=group_ids, + parent_ids=parent_ids, + execution_spec=spec, + execution_plan=plan, + cp_group=torch.distributed.group.WORLD, + ) + cp_loss = (cp_out * local_output_grad).sum() + cp_loss.backward() + _assert_cp_matches_reference( + name, + ref_gdn, + cp_gdn, + ref_hidden, + ref_out, + ref_loss.detach(), + local_hidden, + cp_out, + cp_loss.detach(), + local_index, + ) + + def _assert_sibling_order_matches_cp1( ref_gdn: torch.nn.Module, cp_gdn: torch.nn.Module, @@ -233,14 +519,15 @@ def _assert_sibling_order_matches_cp1( swapped_parent_ids[0, 5:9] = 0 swapped_group_ids[0, 9:12] = 2 swapped_parent_ids[0, 9:12] = 0 - spec = parse_gdn_shared_prefix_segments( - swapped_group_ids, swapped_parent_ids, min_completions_per_family=0 - ) + spec = parse_gdn_prefix_tree_segments(swapped_group_ids, swapped_parent_ids) plan = build_gdn_rank_execution_plan( spec, device=group_ids.device, cp_rank=rank, cp_size=cp_size, + attention_token_layout_index=_uniform_attention_layout( + spec.real_token_count, cp_size + ), planner_config=GdnPlannerConfig(), ) hidden, output_grad = _hidden_and_grad(case, seed=20520426 + cp_size) @@ -349,6 +636,36 @@ def _assert_cp_matches_reference( torch.cuda.synchronize() +def _tree_training_local_loss( + gdn: torch.nn.Module, + local_hidden: torch.Tensor, + group_ids: torch.Tensor, + parent_ids: torch.Tensor, + spec: Any, + plan: Any, +) -> torch.Tensor: + out, _ = run_gdn_layer( + gdn, + local_hidden, + group_ids=group_ids, + parent_ids=parent_ids, + execution_spec=spec, + execution_plan=plan, + cp_group=torch.distributed.group.WORLD, + ) + return out.float().square().mean() + + +def _parameter_l1_norm(tensors: Any) -> torch.Tensor: + total: torch.Tensor | None = None + for tensor in tensors: + contribution = tensor.detach().float().abs().sum() + total = contribution if total is None else total + contribution + if total is None: + return torch.tensor(0.0, device="cuda") + return total + + class _TensorGradView: def __init__(self, grad: torch.Tensor) -> None: self.grad = grad @@ -377,6 +694,164 @@ def _hidden_and_grad( return hidden, grad +def _tree_hidden_and_grad( + sequence_length: int, *, seed: int +) -> tuple[torch.Tensor, torch.Tensor]: + generator = torch.Generator(device="cuda").manual_seed(seed) + hidden = torch.randn( + sequence_length, + 1, + 64, + device="cuda", + dtype=GDN_CORRECTNESS_DTYPE, + generator=generator, + ) + grad = torch.randn( + hidden.shape, + device="cuda", + dtype=GDN_CORRECTNESS_DTYPE, + generator=generator, + ) + torch.distributed.broadcast(hidden, src=0) + torch.distributed.broadcast(grad, src=0) + return hidden, grad + + +def _tree_trainability_hidden(sequence_length: int, *, cp_size: int) -> torch.Tensor: + generator = torch.Generator(device="cuda").manual_seed(6262026 + cp_size) + hidden = torch.randn( + sequence_length, + 1, + 64, + device="cuda", + dtype=GDN_CORRECTNESS_DTYPE, + generator=generator, + ) + torch.distributed.broadcast(hidden, src=0) + return hidden + + +def _tree_trainability_pack(): + root = torch.arange(11, 107) + mid = torch.arange(1001, 1161) + sibling = torch.arange(2001, 2081) + return prefix_tree_pack( + ( + torch.cat((root, mid, torch.tensor([301, 302]))), + torch.cat((root, mid, torch.tensor([303, 304, 305]))), + torch.cat((root, sibling, torch.tensor([401]))), + ), + max_depth=2, + ) + + +def _tree_chain_pack(): + long_root = torch.arange(11, 267) + short_root = torch.arange(1001, 1097) + long_mid = torch.arange(2001, 2641) + other_mid = torch.arange(3001, 3065) + return prefix_tree_pack( + ( + torch.cat((long_root, torch.tensor([301]))), + torch.cat((long_root, torch.tensor([302]))), + torch.cat((short_root, long_mid, torch.tensor([401]))), + torch.cat((short_root, long_mid, torch.tensor([402]))), + torch.cat((short_root, other_mid, torch.tensor([403]))), + ), + max_depth=2, + ) + + +def _tree_chain_planner_config() -> GdnPlannerConfig: + return GdnPlannerConfig( + runtime_local_recurrent_tokens_per_ms=10.0, + runtime_chain_recurrent_tokens_per_ms=1_000_000.0, + runtime_cp_summary_bandwidth_bytes_per_ms=1e18, + runtime_cp_summary_compute_segments_per_ms=1e18, + runtime_cp_suffix_scan_latency_ms=0.0, + runtime_cp_suffix_scan_segments_per_ms=1e18, + ) + + +def _tree_fuzz_planner_config() -> GdnPlannerConfig: + return GdnPlannerConfig() + + +def _tree_fuzz_packs() -> tuple[tuple[str, Any], ...]: + return ( + ( + "tree_fuzz_duplicates", + prefix_tree_pack(_duplicate_tree_sequences(), max_depth=4), + ), + ( + "tree_fuzz_ragged_depth4", + prefix_tree_pack(_random_tree_sequences(13, max_depth=4), max_depth=4), + ), + ( + "tree_fuzz_mixed_tiny_long", + prefix_tree_pack(_random_tree_sequences(29, max_depth=5), max_depth=5), + ), + ) + + +def _uniform_attention_layout(real_token_count: int, cp_size: int) -> TokenLayoutIndex: + ranges_by_rank: list[tuple[tuple[int, int, int], ...]] = [] + for rank in range(cp_size): + start = (int(real_token_count) * rank) // cp_size + end = (int(real_token_count) * (rank + 1)) // cp_size + ranges_by_rank.append(((start, end, 0),) if end > start else ()) + return TokenLayoutIndex( + ownership_ranges_by_rank=tuple(ranges_by_rank), + token_counts_by_rank=tuple( + sum(end - start for start, end, _position in ranges) + for ranges in ranges_by_rank + ), + ) + + +def _duplicate_tree_sequences() -> tuple[torch.Tensor, ...]: + root = torch.arange(11, 331) + mid_a = torch.arange(1001, 1261) + mid_b = torch.arange(2001, 2065) + leaf_a = torch.arange(3001, 3013) + leaf_b = torch.arange(4001, 4017) + first = torch.cat((root, mid_a, leaf_a)) + second = torch.cat((root, mid_a, leaf_b)) + third = torch.cat((root, mid_b, torch.tensor([91, 92, 93]))) + return (first, first, second, third, third) + + +def _random_tree_sequences(seed: int, *, max_depth: int) -> tuple[torch.Tensor, ...]: + generator = torch.Generator().manual_seed(seed) + next_token = 1 + + def randint(low: int, high: int) -> int: + return int(torch.randint(low, high + 1, (), generator=generator).item()) + + def tokens(length: int) -> torch.Tensor: + nonlocal next_token + out = torch.arange(next_token, next_token + length) + next_token += length + 997 + return out + + def segment_length(depth: int) -> int: + choices = (1, 3, 17, 64, 129, 257, 384 if depth == 0 else 96) + return choices[randint(0, len(choices) - 1)] + + def walk(prefix: torch.Tensor, depth: int) -> list[torch.Tensor]: + here = torch.cat((prefix, tokens(segment_length(depth)))) + if depth + 1 >= max_depth: + return [ + torch.cat((here, tokens(randint(1, 17)))) for _ in range(randint(2, 4)) + ] + leaves: list[torch.Tensor] = [] + for _ in range(randint(2, 3)): + leaves.extend(walk(here, depth + 1)) + return leaves + + return tuple(walk(torch.empty(0, dtype=torch.long), 0)) + + def _packed_correctness_cases() -> tuple[GdnPhase0Case, ...]: return ( *default_phase0_cases(conv_width=2), @@ -388,11 +863,7 @@ def _packed_correctness_cases() -> tuple[GdnPhase0Case, ...]: def _planner_config_for_case(case: GdnPhase0Case) -> GdnPlannerConfig | None: if case.name != "mixed_local_chain_edge": return None - return GdnPlannerConfig( - cp_chain_min_tokens_per_rank=16, - cp_chain_min_total_tokens=128, - cp_chain_min_prefix_only_tokens=128, - ) + return GdnPlannerConfig() def _mixed_local_chain_case() -> GdnPhase0Case: @@ -450,9 +921,3 @@ def _swap_siblings(tensor: torch.Tensor) -> torch.Tensor: def _skip_without_gpus(cp_size: int) -> None: if not torch.cuda.is_available() or torch.cuda.device_count() < cp_size: pytest.skip(f"Need {cp_size} CUDA devices for CP{cp_size} packed GDN.") - - -def _find_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/integration/megatron/gdn_shared_prefix/test_gdn_cp_train_prepare.py b/tests/integration/megatron/gdn_shared_prefix/test_gdn_cp_train_prepare.py index e0d2e831f..6ef5a8890 100644 --- a/tests/integration/megatron/gdn_shared_prefix/test_gdn_cp_train_prepare.py +++ b/tests/integration/megatron/gdn_shared_prefix/test_gdn_cp_train_prepare.py @@ -1,7 +1,6 @@ from __future__ import annotations from pathlib import Path -import socket from typing import Any, cast import pytest @@ -24,6 +23,7 @@ from art.preprocessing.pack import PackedTensors # noqa: E402 from .cases import default_phase0_cases # noqa: E402 +from .distributed_init import file_init_method # noqa: E402 from .packed_layout import build_phase0_packed_tensors # noqa: E402 @@ -31,10 +31,10 @@ def test_gdn_cp_training_batch_carries_prebuilt_rank_plan(tmp_path: Path) -> Non cp_size = 2 if not torch.cuda.is_available() or torch.cuda.device_count() < cp_size: pytest.skip(f"requires {cp_size} CUDA devices") - port = _find_free_port() + init_method = file_init_method(tmp_path, "gdn_cp_train_prepare") mp.spawn( _worker, - args=(cp_size, port, str(tmp_path)), + args=(cp_size, init_method, str(tmp_path)), nprocs=cp_size, join=True, ) @@ -42,11 +42,11 @@ def test_gdn_cp_training_batch_carries_prebuilt_rank_plan(tmp_path: Path) -> Non assert (tmp_path / f"rank_{rank}.ok").read_text() == "ok\n" -def _worker(rank: int, cp_size: int, port: int, output_dir: str) -> None: +def _worker(rank: int, cp_size: int, init_method: str, output_dir: str) -> None: torch.cuda.set_device(rank) init_process_group( backend="nccl", - init_method=f"tcp://127.0.0.1:{port}", + init_method=init_method, rank=rank, world_size=cp_size, ) @@ -101,12 +101,6 @@ def _worker(rank: int, cp_size: int, port: int, output_dir: str) -> None: destroy_process_group() -def _find_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]) - - def test_main_loss_matches_shifted_dispatched_loss_inputs() -> None: packed = cast( Any, diff --git a/tests/integration/megatron/gdn_shared_prefix/test_gdn_planner_runtime_model.py b/tests/integration/megatron/gdn_shared_prefix/test_gdn_planner_runtime_model.py new file mode 100644 index 000000000..82f20e44a --- /dev/null +++ b/tests/integration/megatron/gdn_shared_prefix/test_gdn_planner_runtime_model.py @@ -0,0 +1,135 @@ +from __future__ import annotations + +import pytest + +from art.megatron.gdn.gdn_prefix_tree import GdnPlannerConfig + + +def test_gdn_planner_runtime_model_preserves_qwen35_reference_shape() -> None: + reference = GdnPlannerConfig.from_model_shape( + hidden_size=2048, + tensor_model_parallel_size=1, + linear_num_key_heads=16, + linear_num_value_heads=32, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_conv_kernel_dim=4, + dtype_bytes=2, + ) + default = GdnPlannerConfig() + assert ( + reference.runtime_hidden_bytes_per_token + == default.runtime_hidden_bytes_per_token + ) + assert ( + reference.runtime_cp_summary_bytes_per_segment + == default.runtime_cp_summary_bytes_per_segment + ) + assert reference.runtime_local_recurrent_tokens_per_ms == pytest.approx( + default.runtime_local_recurrent_tokens_per_ms + ) + assert reference.runtime_chain_recurrent_tokens_per_ms == pytest.approx( + default.runtime_chain_recurrent_tokens_per_ms + ) + assert ( + reference.runtime_parent_state_bytes_per_exchange + / reference.runtime_parent_state_bandwidth_bytes_per_ms + ) == pytest.approx( + default.runtime_parent_state_bytes_per_exchange + / default.runtime_parent_state_bandwidth_bytes_per_ms + ) + + +def test_gdn_planner_runtime_model_scales_with_state_shape() -> None: + reference = GdnPlannerConfig.from_model_shape( + hidden_size=2048, + tensor_model_parallel_size=1, + linear_num_key_heads=16, + linear_num_value_heads=32, + linear_key_head_dim=128, + linear_value_head_dim=128, + ) + small = GdnPlannerConfig.from_model_shape( + hidden_size=1024, + tensor_model_parallel_size=1, + linear_num_key_heads=8, + linear_num_value_heads=16, + linear_key_head_dim=64, + linear_value_head_dim=64, + ) + large = GdnPlannerConfig.from_model_shape( + hidden_size=7168, + tensor_model_parallel_size=1, + linear_num_key_heads=56, + linear_num_value_heads=112, + linear_key_head_dim=128, + linear_value_head_dim=128, + ) + assert ( + small.runtime_hidden_bytes_per_token < reference.runtime_hidden_bytes_per_token + ) + assert small.runtime_cp_summary_bytes_per_segment < ( + reference.runtime_cp_summary_bytes_per_segment + ) + assert small.runtime_local_recurrent_tokens_per_ms > ( + reference.runtime_local_recurrent_tokens_per_ms + ) + assert ( + large.runtime_hidden_bytes_per_token > reference.runtime_hidden_bytes_per_token + ) + assert large.runtime_cp_summary_bytes_per_segment > ( + reference.runtime_cp_summary_bytes_per_segment + ) + assert large.runtime_local_recurrent_tokens_per_ms < ( + reference.runtime_local_recurrent_tokens_per_ms + ) + + +def test_gdn_planner_runtime_model_tracks_qwen35_shape_axes() -> None: + qwen35_35b = GdnPlannerConfig.from_model_shape( + hidden_size=2048, + tensor_model_parallel_size=1, + linear_num_key_heads=16, + linear_num_value_heads=32, + linear_key_head_dim=128, + linear_value_head_dim=128, + ) + qwen35_9b = GdnPlannerConfig.from_model_shape( + hidden_size=4096, + tensor_model_parallel_size=1, + linear_num_key_heads=16, + linear_num_value_heads=32, + linear_key_head_dim=128, + linear_value_head_dim=128, + ) + qwen35_397b = GdnPlannerConfig.from_model_shape( + hidden_size=4096, + tensor_model_parallel_size=1, + linear_num_key_heads=16, + linear_num_value_heads=64, + linear_key_head_dim=128, + linear_value_head_dim=128, + ) + + assert qwen35_9b.runtime_hidden_bytes_per_token == 2 * ( + qwen35_35b.runtime_hidden_bytes_per_token + ) + assert qwen35_9b.runtime_local_recurrent_tokens_per_ms == pytest.approx( + qwen35_35b.runtime_local_recurrent_tokens_per_ms + ) + assert qwen35_9b.runtime_cp_summary_bytes_per_segment == ( + qwen35_35b.runtime_cp_summary_bytes_per_segment + ) + + assert qwen35_397b.runtime_hidden_bytes_per_token == 2 * ( + qwen35_35b.runtime_hidden_bytes_per_token + ) + assert qwen35_397b.runtime_cp_summary_bytes_per_segment == 2 * ( + qwen35_35b.runtime_cp_summary_bytes_per_segment + ) + assert qwen35_397b.runtime_local_recurrent_tokens_per_ms == pytest.approx( + (2**-0.75) * qwen35_35b.runtime_local_recurrent_tokens_per_ms + ) + assert qwen35_397b.runtime_chain_recurrent_tokens_per_ms == pytest.approx( + (2**-0.75) * qwen35_35b.runtime_chain_recurrent_tokens_per_ms + ) diff --git a/tests/integration/megatron/gdn_shared_prefix/test_qwen35_full_model_cp1_packed_vs_flattened.py b/tests/integration/megatron/gdn_shared_prefix/test_qwen35_full_model_cp1_packed_vs_flattened.py index 19f33970c..abb7c8f6d 100644 --- a/tests/integration/megatron/gdn_shared_prefix/test_qwen35_full_model_cp1_packed_vs_flattened.py +++ b/tests/integration/megatron/gdn_shared_prefix/test_qwen35_full_model_cp1_packed_vs_flattened.py @@ -2,7 +2,8 @@ from collections.abc import Iterator from contextlib import ExitStack, contextmanager -import socket +from pathlib import Path +import tempfile from typing import Any import pytest @@ -22,7 +23,7 @@ from art.loss import shift_tensor from art.megatron.model_support.handlers.qwen3_5 import QWEN3_5_MOE_HANDLER -from art.megatron.shared_prefix_state import create_shared_prefix_state +from art.megatron.prefix_tree_state import create_prefix_tree_state from ..model_support.oracle_harness import TEST_DEFAULT_FLEX_BACKEND from ..model_support.oracle_worker import ( @@ -31,6 +32,7 @@ _apply_test_flex_inner_fp32_patch, ) from .cases import default_phase0_cases +from .distributed_init import file_init_method from .metrics import ( GDN_CORRECTNESS_DTYPE, MEAN_ABS_PCT_THRESHOLD, @@ -40,7 +42,7 @@ stable_output_mse_loss, ) from .packed_layout import build_phase0_packed_tensors -from .parser_import import parse_gdn_shared_prefix_segments +from .parser_import import parse_gdn_prefix_tree_segments from .real_gdn_oracle import ( attach_main_grads, zero_parameter_grads, @@ -64,7 +66,7 @@ def _fp32_test_flex_backend() -> Iterator[None]: @pytest.mark.skipif( not torch.cuda.is_available(), - reason="CUDA is required for Qwen3.5 full-model shared-prefix oracle coverage.", + reason="CUDA is required for Qwen3.5 full-model prefix-tree oracle coverage.", ) def test_qwen35_full_model_cp1_matches_flattened_grad_accumulation() -> None: with _single_rank_model_parallel(): @@ -96,66 +98,66 @@ def test_qwen35_full_model_cp1_matches_flattened_grad_accumulation() -> None: flat_loss_sum: torch.Tensor | None = None logits_mean_abs_pct = 0.0 - spec = parse_gdn_shared_prefix_segments( - group_ids.cpu(), parent_ids.cpu(), min_completions_per_family=1 - ) - for family in spec.families: - row = family.row_index - prefix = family.prefix - for completion in family.completions: - ref_tokens = torch.cat( - [ - tokens[row : row + 1, prefix.start : prefix.end], - tokens[row : row + 1, completion.start : completion.end], - ], - dim=1, - ) - ref_pos = torch.cat( - [ - input_pos[row : row + 1, prefix.start : prefix.end], - input_pos[row : row + 1, completion.start : completion.end], - ], - dim=1, - ) - ref_assistant_mask = torch.cat( - [ - torch.zeros( - (1, prefix.length), dtype=torch.bool, device=device - ), - assistant_mask[ - row : row + 1, completion.start : completion.end - ], - ], - dim=1, - ) - ref_group_ids = torch.zeros_like(ref_tokens) - ref_parent_ids = torch.zeros_like(ref_tokens) - ref_logits, ref_loss = _run_model_loss( - flat_model, - tokens=ref_tokens, - input_pos=ref_pos, - group_ids=ref_group_ids, - parent_ids=ref_parent_ids, - assistant_mask=ref_assistant_mask, - ) - ref_loss.backward() - flat_loss_sum = ( - ref_loss.detach() - if flat_loss_sum is None - else flat_loss_sum + ref_loss.detach() - ) + spec = parse_gdn_prefix_tree_segments(group_ids.cpu(), parent_ids.cpu()) + for segment_index, completion in enumerate(spec.tree_segments): + if spec.tree_parent_indices[segment_index] < 0: + continue + row = completion.row_index + path = _segment_path(spec, segment_index) + completion_offset = sum(segment.length for segment in path[:-1]) + ref_tokens = torch.cat( + [ + tokens[row : row + 1, segment.start : segment.end] + for segment in path + ], + dim=1, + ) + ref_pos = torch.cat( + [ + input_pos[row : row + 1, segment.start : segment.end] + for segment in path + ], + dim=1, + ) + ref_assistant_mask = torch.cat( + [ + torch.zeros( + (1, completion_offset), + dtype=torch.bool, + device=device, + ), + assistant_mask[row : row + 1, completion.start : completion.end], + ], + dim=1, + ) + ref_group_ids = torch.zeros_like(ref_tokens) + ref_parent_ids = torch.zeros_like(ref_tokens) + ref_logits, ref_loss = _run_model_loss( + flat_model, + tokens=ref_tokens, + input_pos=ref_pos, + group_ids=ref_group_ids, + parent_ids=ref_parent_ids, + assistant_mask=ref_assistant_mask, + ) + ref_loss.backward() + flat_loss_sum = ( + ref_loss.detach() + if flat_loss_sum is None + else flat_loss_sum + ref_loss.detach() + ) - if completion.length > 1: - packed_slice = packed_logits[ - row : row + 1, completion.start : completion.end - 1 - ] - ref_slice = ref_logits[ - :, prefix.length : prefix.length + completion.length - 1 - ] - logits_mean_abs_pct = max( - logits_mean_abs_pct, - mean_abs_pct(ref_slice, packed_slice), - ) + if completion.length > 1: + packed_slice = packed_logits[ + row : row + 1, completion.start : completion.end - 1 + ] + ref_slice = ref_logits[ + :, completion_offset : completion_offset + completion.length - 1 + ] + logits_mean_abs_pct = max( + logits_mean_abs_pct, + mean_abs_pct(ref_slice, packed_slice), + ) assert flat_loss_sum is not None grad_name, grad_pct = parameter_grad_mean_abs_pct_with_name( @@ -214,70 +216,64 @@ def _assert_logits_vjp_equivalence( flat_loss_sum: torch.Tensor | None = None logits_mean_abs_pct = 0.0 - spec = parse_gdn_shared_prefix_segments( - group_ids.cpu(), parent_ids.cpu(), min_completions_per_family=1 - ) - for family in spec.families: - row = family.row_index - prefix = family.prefix - for completion in family.completions: - ref_tokens = torch.cat( - [ - tokens[row : row + 1, prefix.start : prefix.end], - tokens[row : row + 1, completion.start : completion.end], - ], - dim=1, - ) - ref_pos = torch.cat( - [ - input_pos[row : row + 1, prefix.start : prefix.end], - input_pos[row : row + 1, completion.start : completion.end], - ], - dim=1, - ) - ref_logits = _run_model_logits( - flat_model, - tokens=ref_tokens, - input_pos=ref_pos, - group_ids=torch.zeros_like(ref_tokens), - parent_ids=torch.zeros_like(ref_tokens), - ) - ref_output_grad = torch.zeros_like(ref_logits) - ref_output_mask = torch.zeros( - ref_logits.shape[:2], - device=ref_logits.device, - dtype=torch.bool, - ) - if completion.length > 1: - ref_output_grad[ - :, prefix.length : prefix.length + completion.length - 1 - ] = output_grad[row : row + 1, completion.start : completion.end - 1] - ref_output_mask[ - :, prefix.length : prefix.length + completion.length - 1 - ] = True - ref_loss = stable_output_mse_loss( - ref_logits, - ref_output_grad, - mask=ref_output_mask.unsqueeze(-1), - denominator=loss_denominator, - ) - ref_loss.backward() - flat_loss_sum = ( - ref_loss.detach() - if flat_loss_sum is None - else flat_loss_sum + ref_loss.detach() + spec = parse_gdn_prefix_tree_segments(group_ids.cpu(), parent_ids.cpu()) + for segment_index, completion in enumerate(spec.tree_segments): + if spec.tree_parent_indices[segment_index] < 0: + continue + row = completion.row_index + path = _segment_path(spec, segment_index) + completion_offset = sum(segment.length for segment in path[:-1]) + ref_tokens = torch.cat( + [tokens[row : row + 1, segment.start : segment.end] for segment in path], + dim=1, + ) + ref_pos = torch.cat( + [input_pos[row : row + 1, segment.start : segment.end] for segment in path], + dim=1, + ) + ref_logits = _run_model_logits( + flat_model, + tokens=ref_tokens, + input_pos=ref_pos, + group_ids=torch.zeros_like(ref_tokens), + parent_ids=torch.zeros_like(ref_tokens), + ) + ref_output_grad = torch.zeros_like(ref_logits) + ref_output_mask = torch.zeros( + ref_logits.shape[:2], + device=ref_logits.device, + dtype=torch.bool, + ) + if completion.length > 1: + ref_output_grad[ + :, completion_offset : completion_offset + completion.length - 1 + ] = output_grad[row : row + 1, completion.start : completion.end - 1] + ref_output_mask[ + :, completion_offset : completion_offset + completion.length - 1 + ] = True + ref_loss = stable_output_mse_loss( + ref_logits, + ref_output_grad, + mask=ref_output_mask.unsqueeze(-1), + denominator=loss_denominator, + ) + ref_loss.backward() + flat_loss_sum = ( + ref_loss.detach() + if flat_loss_sum is None + else flat_loss_sum + ref_loss.detach() + ) + if completion.length > 1: + packed_slice = packed_logits[ + row : row + 1, completion.start : completion.end - 1 + ] + ref_slice = ref_logits[ + :, completion_offset : completion_offset + completion.length - 1 + ] + logits_mean_abs_pct = max( + logits_mean_abs_pct, + mean_abs_pct(ref_slice, packed_slice), ) - if completion.length > 1: - packed_slice = packed_logits[ - row : row + 1, completion.start : completion.end - 1 - ] - ref_slice = ref_logits[ - :, prefix.length : prefix.length + completion.length - 1 - ] - logits_mean_abs_pct = max( - logits_mean_abs_pct, - mean_abs_pct(ref_slice, packed_slice), - ) assert flat_loss_sum is not None grad_name, grad_pct = parameter_grad_mean_abs_pct_with_name( @@ -304,7 +300,7 @@ def _run_model_loss( group_ids=group_ids, parent_ids=parent_ids, ) - attention_state = create_shared_prefix_state( + attention_state = create_prefix_tree_state( group_ids=group_ids, parent_ids=parent_ids, build_gdn_execution_spec=True, @@ -339,7 +335,7 @@ def _run_model_logits( group_ids: torch.Tensor, parent_ids: torch.Tensor, ) -> torch.Tensor: - attention_state = create_shared_prefix_state( + attention_state = create_prefix_tree_state( group_ids=group_ids, parent_ids=parent_ids, build_gdn_execution_spec=True, @@ -359,6 +355,15 @@ def _run_model_logits( return logits +def _segment_path(spec: Any, segment_index: int) -> tuple[Any, ...]: + indices = [] + cursor = segment_index + while cursor >= 0: + indices.append(cursor) + cursor = spec.tree_parent_indices[cursor] + return tuple(spec.tree_segments[index] for index in reversed(indices)) + + def _make_matching_models() -> tuple[torch.nn.Module, torch.nn.Module]: model_parallel_cuda_manual_seed(1234) packed = _make_model() @@ -424,28 +429,23 @@ def _single_rank_model_parallel() -> Iterator[None]: if is_initialized(): pytest.skip("torch.distributed is already initialized in this process.") torch.cuda.set_device(0) - init_process_group( - backend="nccl", - init_method=f"tcp://127.0.0.1:{_find_free_port()}", - 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, + with tempfile.TemporaryDirectory(prefix="art_dist_") as tmp: + init_process_group( + backend="nccl", + init_method=file_init_method(Path(tmp), "qwen35_full_model_cp1"), + rank=0, + world_size=1, ) - yield - finally: - if getattr(ps, "model_parallel_is_initialized", lambda: False)(): - ps.destroy_model_parallel() - if is_initialized(): - destroy_process_group() - - -def _find_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]) + 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() + if is_initialized(): + destroy_process_group() diff --git a/tests/integration/megatron/gdn_shared_prefix/test_qwen35_gdn_topology_oracle.py b/tests/integration/megatron/gdn_shared_prefix/test_qwen35_gdn_topology_oracle.py index 2eff1264f..5cfded08c 100644 --- a/tests/integration/megatron/gdn_shared_prefix/test_qwen35_gdn_topology_oracle.py +++ b/tests/integration/megatron/gdn_shared_prefix/test_qwen35_gdn_topology_oracle.py @@ -26,7 +26,7 @@ @pytest.mark.parametrize("cp_size", _CP_SIZES) -def test_qwen35_gdn_shared_prefix_cp_topology_oracle( +def test_qwen35_gdn_prefix_tree_cp_topology_oracle( cp_size: int, capsys: pytest.CaptureFixture[str], monkeypatch: pytest.MonkeyPatch, @@ -63,7 +63,7 @@ def test_qwen35_gdn_shared_prefix_cp_topology_oracle( } ) variant = VariantSpec( - name=f"qwen35_gdn_shared_prefix_cp{cp_size}", + name=f"qwen35_gdn_prefix_tree_cp{cp_size}", objective="rl", topology=topology, ) diff --git a/tests/integration/megatron/gdn_shared_prefix/test_real_gdn_cp1_packed_vs_flattened.py b/tests/integration/megatron/gdn_shared_prefix/test_real_gdn_cp1_packed_vs_flattened.py index de6933582..f026b90c9 100644 --- a/tests/integration/megatron/gdn_shared_prefix/test_real_gdn_cp1_packed_vs_flattened.py +++ b/tests/integration/megatron/gdn_shared_prefix/test_real_gdn_cp1_packed_vs_flattened.py @@ -2,7 +2,8 @@ from collections.abc import Iterator from contextlib import contextmanager -import socket +from pathlib import Path +import tempfile import pytest @@ -18,13 +19,13 @@ from megatron.core.ssm.gated_delta_net import GatedDeltaNet from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from torch.distributed import ( - DistNetworkError, destroy_process_group, init_process_group, is_initialized, ) from .cases import default_phase0_cases +from .distributed_init import file_init_method from .metrics import ( GDN_CORRECTNESS_DTYPE, MEAN_ABS_PCT_MISMATCH_THRESHOLD, @@ -232,44 +233,23 @@ def _single_rank_model_parallel() -> Iterator[None]: if is_initialized(): pytest.skip("torch.distributed is already initialized in this process.") torch.cuda.set_device(0) - _init_single_rank_process_group() - try: - ps.initialize_model_parallel( - tensor_model_parallel_size=1, - pipeline_model_parallel_size=1, - context_parallel_size=1, - expert_model_parallel_size=1, + with tempfile.TemporaryDirectory(prefix="art_dist_") as tmp: + init_process_group( + backend="nccl", + init_method=file_init_method(Path(tmp), "single_rank"), + rank=0, + world_size=1, ) - yield - finally: - if getattr(ps, "model_parallel_is_initialized", lambda: False)(): - ps.destroy_model_parallel() - if is_initialized(): - destroy_process_group() - - -def _init_single_rank_process_group() -> None: - last_error: DistNetworkError | None = None - for _ in range(16): try: - init_process_group( - backend="nccl", - init_method=f"tcp://127.0.0.1:{_find_free_port()}", - rank=0, - world_size=1, + ps.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=1, + expert_model_parallel_size=1, ) - return - except DistNetworkError as error: - if "EADDRINUSE" not in str(error): - raise - last_error = error + yield + finally: + if getattr(ps, "model_parallel_is_initialized", lambda: False)(): + ps.destroy_model_parallel() if is_initialized(): destroy_process_group() - if last_error is not None: - raise last_error - - -def _find_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/integration/megatron/gdn_shared_prefix/test_real_gdn_native_fla_cp.py b/tests/integration/megatron/gdn_shared_prefix/test_real_gdn_native_fla_cp.py index e0d164c56..63c3ce43e 100644 --- a/tests/integration/megatron/gdn_shared_prefix/test_real_gdn_native_fla_cp.py +++ b/tests/integration/megatron/gdn_shared_prefix/test_real_gdn_native_fla_cp.py @@ -1,7 +1,6 @@ from __future__ import annotations from pathlib import Path -import socket from typing import cast import pytest @@ -22,11 +21,12 @@ from torch.distributed import destroy_process_group, init_process_group # noqa: E402 import torch.multiprocessing as mp # noqa: E402 -from art.megatron.gdn.gdn_shared_prefix import ( # noqa: E402 +from art.megatron.context_parallel.layout_index import TokenLayoutIndex # noqa: E402 +from art.megatron.gdn.gdn_prefix_tree import ( # noqa: E402 GdnPlannerConfig, GdnSegmentBucketPlan, build_gdn_rank_execution_plan, - parse_gdn_shared_prefix_segments, + parse_gdn_prefix_tree_segments, ) from art.megatron.gdn.operator import ( # noqa: E402 _project_gdn_inputs, @@ -37,6 +37,7 @@ ) from .cases import GdnFamilyShape, GdnPackedRowShape, GdnPhase0Case # noqa: E402 +from .distributed_init import file_init_method # noqa: E402 from .metrics import ( # noqa: E402 GDN_CORRECTNESS_DTYPE, MEAN_ABS_PCT_THRESHOLD, @@ -70,10 +71,10 @@ def test_real_qwen35_gdn_native_fla_cp_prepared_varlen_batch_matches_single_rank( cp_size: int, tmp_path: Path ) -> None: - port = _find_free_port() + init_method = file_init_method(tmp_path, f"native_gdn_prepared_cp{cp_size}") mp.spawn( _native_gdn_cp_prepared_varlen_worker, - args=(cp_size, port, str(tmp_path)), + args=(cp_size, init_method, str(tmp_path)), nprocs=cp_size, join=True, ) @@ -89,10 +90,10 @@ def test_real_qwen35_gdn_native_fla_cp_prepared_varlen_batch_matches_single_rank def test_real_qwen35_gdn_native_cp_packed_layer_matches_cp1( cp_size: int, tmp_path: Path ) -> None: - port = _find_free_port() + init_method = file_init_method(tmp_path, f"native_gdn_packed_cp{cp_size}") mp.spawn( _native_gdn_cp_packed_layer_worker, - args=(cp_size, port, str(tmp_path)), + args=(cp_size, init_method, str(tmp_path)), nprocs=cp_size, join=True, ) @@ -103,13 +104,13 @@ def test_real_qwen35_gdn_native_cp_packed_layer_matches_cp1( def _native_gdn_cp_packed_layer_worker( rank: int, cp_size: int, - port: int, + init_method: str, output_dir: str, ) -> None: torch.cuda.set_device(rank) init_process_group( backend="nccl", - init_method=f"tcp://127.0.0.1:{port}", + init_method=init_method, rank=rank, world_size=cp_size, ) @@ -127,29 +128,18 @@ def _native_gdn_cp_packed_layer_worker( tensors = build_phase0_packed_tensors(case) group_ids = tensors["group_ids"].cuda() parent_ids = tensors["parent_ids"].cuda() - spec = parse_gdn_shared_prefix_segments( - group_ids, parent_ids, min_completions_per_family=0 - ) + spec = parse_gdn_prefix_tree_segments(group_ids, parent_ids) plan = build_gdn_rank_execution_plan( spec, device=group_ids.device, cp_rank=rank, cp_size=cp_size, - planner_config=GdnPlannerConfig( - cp_chain_min_tokens_per_rank=16, - cp_chain_min_total_tokens=128, - cp_chain_min_prefix_only_tokens=128, - # This test is the native chain correctness guard, so force the - # planner onto chain prefix and completion buckets. - planner_chain_bucket_ms=0.0, - planner_chain_token_ms=0.0, - planner_local_bucket_ms=1.0, - planner_local_token_ms=1.0, - cp_chain_min_score_delta_ms=0.0, + attention_token_layout_index=_uniform_attention_layout( + spec.real_token_count, cp_size ), + planner_config=GdnPlannerConfig(), ) - assert plan.chain_prefix_buckets - assert plan.chain_completion_buckets + assert any(plan.tree_chain_buckets_by_depth) hidden, output_grad = _packed_hidden_and_grad(case, cp_size) ref_hidden = hidden.clone().detach().requires_grad_(True) ref_out, _ = run_gdn_layer( @@ -215,13 +205,13 @@ def _native_gdn_cp_packed_layer_worker( def _native_gdn_cp_prepared_varlen_worker( rank: int, cp_size: int, - port: int, + init_method: str, output_dir: str, ) -> None: torch.cuda.set_device(rank) init_process_group( backend="nccl", - init_method=f"tcp://127.0.0.1:{port}", + init_method=init_method, rank=rank, world_size=cp_size, ) @@ -465,6 +455,21 @@ def _packed_hidden_and_grad( return hidden, output_grad +def _uniform_attention_layout(real_token_count: int, cp_size: int) -> TokenLayoutIndex: + ranges_by_rank: list[tuple[tuple[int, int, int], ...]] = [] + for rank in range(cp_size): + start = (int(real_token_count) * rank) // cp_size + end = (int(real_token_count) * (rank + 1)) // cp_size + ranges_by_rank.append(((start, end, 0),) if end > start else ()) + return TokenLayoutIndex( + ownership_ranges_by_rank=tuple(ranges_by_rank), + token_counts_by_rank=tuple( + sum(end - start for start, end, _position in ranges) + for ranges in ranges_by_rank + ), + ) + + def _varlen_hidden_and_lengths(cp_size: int) -> tuple[torch.Tensor, torch.Tensor]: device = torch.device("cuda") lengths = torch.tensor((512, 1024, 1536), device=device, dtype=torch.long) @@ -492,8 +497,25 @@ def _varlen_bucket( cu_seqlens_cpu = torch.cat( [lengths_cpu.new_zeros(1), torch.cumsum(lengths_cpu, dim=0)] ) - offsets = torch.arange(max_len, device=device, dtype=torch.long).unsqueeze(1) - real_mask = offsets < lengths.unsqueeze(0) + row_indices = torch.cat( + [ + torch.full((int(length),), row, device=device, dtype=torch.long) + for row, length in enumerate(lengths_cpu.tolist()) + if int(length) > 0 + ], + dim=0, + ) + position_indices = torch.cat( + [ + torch.arange(int(length), device=device, dtype=torch.long) + for length in lengths_cpu.tolist() + if int(length) > 0 + ], + dim=0, + ) + real_mask = torch.ones( + int(lengths_cpu.sum().item()), device=device, dtype=torch.bool + ) return GdnSegmentBucketPlan( length=max_len, lengths=lengths, @@ -502,11 +524,8 @@ def _varlen_bucket( real_mask=real_mask, cu_seqlens=cu_seqlens_cpu.to(device=device), cu_seqlens_cpu=cu_seqlens_cpu, - row_indices=torch.arange(int(lengths.numel()), device=device, dtype=torch.long) - .unsqueeze(0) - .expand(max_len, -1) - .contiguous(), - position_indices=offsets.expand(-1, int(lengths.numel())).contiguous(), + row_indices=row_indices.contiguous(), + position_indices=position_indices.contiguous(), family_indices=torch.arange( int(lengths.numel()), device=device, dtype=torch.long ), @@ -599,9 +618,3 @@ def _all_reduce_parameter_grads(module: torch.nn.Module) -> None: main_grad = getattr(parameter, "main_grad", None) if main_grad is not None: torch.distributed.all_reduce(main_grad) - - -def _find_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/integration/megatron/gdn_shared_prefix/test_real_gdn_tp_lora.py b/tests/integration/megatron/gdn_shared_prefix/test_real_gdn_tp_lora.py index c4bd99abc..62e217daf 100644 --- a/tests/integration/megatron/gdn_shared_prefix/test_real_gdn_tp_lora.py +++ b/tests/integration/megatron/gdn_shared_prefix/test_real_gdn_tp_lora.py @@ -1,7 +1,6 @@ from __future__ import annotations from pathlib import Path -import socket import pytest @@ -26,6 +25,7 @@ from art.megatron.model_support.handlers import QWEN3_5_MOE_HANDLER # noqa: E402 from .cases import GdnPhase0Case, default_phase0_cases # noqa: E402 +from .distributed_init import file_init_method # noqa: E402 from .metrics import GDN_CORRECTNESS_DTYPE, assert_real_gdn_metrics # noqa: E402 from .packed_layout import build_phase0_packed_tensors # noqa: E402 from .real_gdn_oracle import ( # noqa: E402 @@ -68,10 +68,10 @@ def test_real_qwen35_gdn_lora_gradients_match_flattened() -> None: reason="At least two CUDA devices are required for TP2 GDN coverage.", ) def test_real_qwen35_gdn_tp2_gradients_match_flattened(tmp_path: Path) -> None: - port = _find_free_port() + init_method = file_init_method(tmp_path, "real_gdn_tp2_lora") mp.spawn( _tp2_worker, - args=(port, str(tmp_path)), + args=(init_method, str(tmp_path)), nprocs=2, join=True, ) @@ -79,11 +79,11 @@ def test_real_qwen35_gdn_tp2_gradients_match_flattened(tmp_path: Path) -> None: assert (tmp_path / f"rank_{rank}.ok").read_text() == "ok\n" -def _tp2_worker(rank: int, port: int, output_dir: str) -> None: +def _tp2_worker(rank: int, init_method: str, output_dir: str) -> None: torch.cuda.set_device(rank) init_process_group( backend="nccl", - init_method=f"tcp://127.0.0.1:{port}", + init_method=init_method, rank=rank, world_size=2, ) @@ -229,9 +229,3 @@ def _gdn_lora_grad_names(gdn: torch.nn.Module) -> tuple[str, ...]: and parameter.grad is not None and bool(parameter.grad.abs().max().item() > 0) ) - - -def _find_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/integration/megatron/lora/test_lora_disk_codecs.py b/tests/integration/megatron/lora/test_lora_disk_codecs.py index 6fbc32b70..c317c52b0 100644 --- a/tests/integration/megatron/lora/test_lora_disk_codecs.py +++ b/tests/integration/megatron/lora/test_lora_disk_codecs.py @@ -1,13 +1,18 @@ import importlib.util import json +import os from pathlib import Path +import shutil import subprocess import sys from typing import Any, cast +import pytest from safetensors.torch import load_file, save_file import torch +pytest.importorskip("megatron.bridge.models.gpt_provider") + from art.megatron import lora as lora_module from art.megatron.lora import LoRA, LoRAParallelSpec, LoRAPublishPlanner from art.megatron.model_support.handlers import ( @@ -16,6 +21,7 @@ QWEN3_5_MOE_HANDLER, QWEN3_MOE_HANDLER, ) +from art.megatron.model_support.handlers.gemma4 import GEMMA4_MOE_HANDLER from art.megatron.model_support.lora_disk import ( load_lora_tensors_for_megatron, normalize_lora_checkpoint_to_vllm, @@ -32,6 +38,66 @@ REPO_ROOT = Path(__file__).parents[4] VLLM_PYTHON = REPO_ROOT / "vllm_runtime/.venv/bin/python" VLLM_RUNTIME_SRC = REPO_ROOT / "vllm_runtime/src" +_VLLM_RUNTIME_UNAVAILABLE_REASON: str | None | object = object() + + +def _vllm_python_cmd() -> list[str]: + override = os.environ.get("ART_TEST_VLLM_PYTHON") + if override: + return [override] + if VLLM_PYTHON.exists(): + return [str(VLLM_PYTHON)] + uv = shutil.which("uv") + if uv is None: + raise RuntimeError( + f"{VLLM_PYTHON} does not exist and uv is not available to run " + "the locked vLLM runtime project" + ) + return [ + uv, + "run", + "--project", + str(REPO_ROOT / "vllm_runtime"), + "--frozen", + "--no-dev", + "python", + ] + + +def _vllm_runtime_unavailable_reason() -> str | None: + global _VLLM_RUNTIME_UNAVAILABLE_REASON + if isinstance(_VLLM_RUNTIME_UNAVAILABLE_REASON, str): + return _VLLM_RUNTIME_UNAVAILABLE_REASON + if _VLLM_RUNTIME_UNAVAILABLE_REASON is None: + return None + try: + subprocess.run( + [ + *_vllm_python_cmd(), + "-c", + "import vllm; from vllm.lora.lora_model import LoRAModel", + ], + check=True, + text=True, + capture_output=True, + timeout=120, + ) + except Exception as exc: + _VLLM_RUNTIME_UNAVAILABLE_REASON = ( + "Stock vLLM loader runtime is unavailable. Run " + "`uv sync --project vllm_runtime --frozen --no-dev`, or set " + "`ART_TEST_VLLM_PYTHON` to a Python environment with vLLM installed. " + f"Original error: {exc}" + ) + return _VLLM_RUNTIME_UNAVAILABLE_REASON + _VLLM_RUNTIME_UNAVAILABLE_REASON = None + return None + + +def test_stock_vllm_loader_runtime_is_available() -> None: + reason = _vllm_runtime_unavailable_reason() + if reason is not None: + pytest.fail(reason) def _config(base_model: str, rank: int = 2, alpha: int = 4) -> dict: @@ -119,6 +185,8 @@ def _assert_stock_vllm_loads( expected_modules: set[str], mapper: str = "none", ) -> list[str]: + if reason := _vllm_runtime_unavailable_reason(): + pytest.skip(reason) script = r""" import json import sys @@ -145,7 +213,7 @@ def _assert_stock_vllm_loads( """ result = subprocess.run( [ - str(VLLM_PYTHON), + *_vllm_python_cmd(), "-c", script, str(path), @@ -728,6 +796,84 @@ def test_qwen35_vllm_config_preserves_shared_expert_targets_when_present(): _assert_tensors_equal(roundtrip, original) +def test_gemma4_shared_experts_plural_keys_map_to_vllm_dense_mlp(tmp_path: Path): + art_prefix = "base_model.model.model.layers.0" + hidden_size = 3 + model_dir = tmp_path / "gemma4" + model_dir.mkdir() + (model_dir / "config.json").write_text( + json.dumps({"num_hidden_layers": 1}), + encoding="utf-8", + ) + save_file( + { + "model.layers.0.pre_feedforward_layernorm.weight": torch.tensor( + [2.0, 4.0, 8.0] + ), + "model.layers.0.pre_feedforward_layernorm_2.weight": torch.tensor( + [1.0, 2.0, 4.0] + ), + }, + model_dir / "model-00001-of-00001.safetensors", + ) + (model_dir / "model.safetensors.index.json").write_text( + json.dumps( + { + "weight_map": { + "model.layers.0.pre_feedforward_layernorm.weight": ( + "model-00001-of-00001.safetensors" + ), + "model.layers.0.pre_feedforward_layernorm_2.weight": ( + "model-00001-of-00001.safetensors" + ), + } + } + ), + encoding="utf-8", + ) + original = { + f"{art_prefix}.mlp.shared_experts.gate_proj.lora_A.weight": torch.ones( + 2, + hidden_size, + ), + f"{art_prefix}.mlp.shared_experts.gate_proj.lora_B.weight": torch.ones(4, 2), + f"{art_prefix}.mlp.shared_experts.up_proj.lora_A.weight": torch.ones( + 2, + hidden_size, + ), + f"{art_prefix}.mlp.shared_experts.up_proj.lora_B.weight": torch.ones(4, 2), + f"{art_prefix}.mlp.shared_experts.down_proj.lora_A.weight": torch.ones(2, 4), + f"{art_prefix}.mlp.shared_experts.down_proj.lora_B.weight": torch.ones( + hidden_size, + 2, + ), + } + adapter_config = _config(str(model_dir)) + vllm_tensors, _ = GEMMA4_MOE_HANDLER.to_vllm_lora_tensors( + original, + adapter_config=adapter_config, + ) + + assert set(vllm_tensors) == { + f"{art_prefix}.mlp.gate_proj.lora_A.weight", + f"{art_prefix}.mlp.gate_proj.lora_B.weight", + f"{art_prefix}.mlp.up_proj.lora_A.weight", + f"{art_prefix}.mlp.up_proj.lora_B.weight", + f"{art_prefix}.mlp.down_proj.lora_A.weight", + f"{art_prefix}.mlp.down_proj.lora_B.weight", + } + assert not any("shared_expert" in key for key in vllm_tensors) + assert torch.equal( + vllm_tensors[f"{art_prefix}.mlp.gate_proj.lora_A.weight"], + torch.full((2, hidden_size), 0.5), + ) + roundtrip = GEMMA4_MOE_HANDLER.from_vllm_lora_tensors( + vllm_tensors, + adapter_config=adapter_config, + ) + _assert_tensors_equal(roundtrip, original) + + def test_gpt_oss_vllm_canonical_roundtrip_and_stock_loader(tmp_path: Path): art_prefix = "base_model.model.model.layers.0" original = _gpt_oss_moe_art_tensors(art_prefix) diff --git a/tests/integration/megatron/lora/test_lora_module_loading.py b/tests/integration/megatron/lora/test_lora_module_loading.py new file mode 100644 index 000000000..2c0c1f1d4 --- /dev/null +++ b/tests/integration/megatron/lora/test_lora_module_loading.py @@ -0,0 +1,43 @@ +import pytest +import torch + +from art.megatron.lora import LoRA + + +def test_load_lora_treats_absent_site_as_identity() -> None: + module = LoRA( + "base_model.model.foo", + in_features=3, + out_features=5, + rank=2, + alpha=32, + dtype=torch.float32, + device=torch.device("cpu"), + ) + adapter = { + "base_model.model.foo.lora_A.weight": torch.ones(2, 3), + "base_model.model.foo.lora_B.weight": torch.ones(5, 2), + } + x = torch.ones(4, 3) + + module.load_lora(adapter) + assert module(x).abs().sum() > 0 + + module.load_lora({}) + assert torch.count_nonzero(module.B_T) == 0 + assert torch.allclose(module(x), torch.zeros(4, 5)) + + +def test_load_lora_rejects_partially_present_site() -> None: + module = LoRA( + "base_model.model.foo", + in_features=3, + out_features=5, + rank=2, + alpha=32, + dtype=torch.float32, + device=torch.device("cpu"), + ) + + with pytest.raises(KeyError, match="Incomplete LoRA adapter keys"): + module.load_lora({"base_model.model.foo.lora_A.weight": torch.ones(2, 3)}) diff --git a/tests/integration/megatron/model_support/chat_template_rollout.py b/tests/integration/megatron/model_support/chat_template_rollout.py index f126fb3ee..beb950513 100644 --- a/tests/integration/megatron/model_support/chat_template_rollout.py +++ b/tests/integration/megatron/model_support/chat_template_rollout.py @@ -14,6 +14,7 @@ _messages_for_chat_template, tokenize_trajectory, tokenize_trajectory_groups, + tokenize_vllm_trajectory_histories, ) from art.trajectories import History from tests.support.chat_template_conformance_cases import ( @@ -131,27 +132,31 @@ def run_chat_template_rollout(base_model: str) -> ChatTemplateRolloutReport: ) ) - non_final_tool_call_base = tokenize_trajectory( + non_final_tool_call_base_results = tokenize_vllm_trajectory_histories( tokenizer=tokenizer, - image_processor=None, - history=_history(inputs.non_final_tool_call_base), + histories=[_history(inputs.non_final_tool_call_base)], advantage=1.0, allow_training_without_logprobs=False, trajectory=inputs.non_final_tool_call_base, ) - non_final_tool_call_mutated = tokenize_trajectory( + non_final_tool_call_mutated_results = tokenize_vllm_trajectory_histories( tokenizer=tokenizer, - image_processor=None, - history=_history(inputs.non_final_tool_call_mutated), + histories=[_history(inputs.non_final_tool_call_mutated)], advantage=1.0, allow_training_without_logprobs=False, trajectory=inputs.non_final_tool_call_mutated, ) - if non_final_tool_call_base is None or non_final_tool_call_mutated is None: + if not non_final_tool_call_base_results or not non_final_tool_call_mutated_results: raise RuntimeError("tool-call tokenization produced no trainable tokens") + non_final_tool_call_base = non_final_tool_call_base_results[-1] + non_final_tool_call_mutated = non_final_tool_call_mutated_results[-1] if ( - len(non_final_tool_call_base.choice_offsets) < 2 - or len(non_final_tool_call_mutated.choice_offsets) < 2 + sum(len(result.choice_offsets) for result in non_final_tool_call_base_results) + < 2 + or sum( + len(result.choice_offsets) for result in non_final_tool_call_mutated_results + ) + < 2 ): raise RuntimeError("expected non-final tool call and final assistant answer") non_final_tool_call_prefix_changed = _assistant_prefix_tokens( @@ -164,10 +169,17 @@ def run_chat_template_rollout(base_model: str) -> ChatTemplateRolloutReport: scenarios.append( ChatTemplateScenarioReport( name="rl_non_final_tool_call_prefill_mutation", - entrypoint="tokenize_trajectory", + entrypoint="tokenize_vllm_trajectory_histories", passed=non_final_tool_call_prefix_changed - and int(sum(non_final_tool_call_base.assistant_mask)) > 0, - assistant_token_count=int(sum(non_final_tool_call_base.assistant_mask)), + and sum( + int(sum(result.assistant_mask)) + for result in non_final_tool_call_base_results + ) + > 0, + assistant_token_count=sum( + int(sum(result.assistant_mask)) + for result in non_final_tool_call_base_results + ), mutation_changed_prompt=non_final_tool_call_prefix_changed, ) ) diff --git a/tests/integration/megatron/model_support/forward_trace.py b/tests/integration/megatron/model_support/forward_trace.py index 4a0410c6f..a075b0098 100644 --- a/tests/integration/megatron/model_support/forward_trace.py +++ b/tests/integration/megatron/model_support/forward_trace.py @@ -1243,6 +1243,7 @@ def _canonicalize_row_aligned_value( def _canonicalize_call_row_token_order(cls, call: dict[str, Any]) -> None: """Canonicalizes all row-aligned call tensors to global token order.""" cls._align_exact_zero_padding_row_token_uids(call) + cls._drop_exact_zero_padding_rows(call) row_token_uids = call.get("row_token_uids") if not isinstance(row_token_uids, torch.Tensor) or row_token_uids.ndim != 1: return @@ -1263,6 +1264,38 @@ def _canonicalize_call_row_token_order(cls, call: dict[str, Any]) -> None: ) call["row_token_uids"] = row_token_uids.index_select(0, order).contiguous() + @classmethod + def _drop_exact_zero_padding_rows(cls, call: dict[str, Any]) -> None: + """Removes traced sequence-padding rows before comparing compact CP traces.""" + row_token_uids = call.get("row_token_uids") + tensor = call.get("primary_output") + if ( + not isinstance(row_token_uids, torch.Tensor) + or row_token_uids.ndim != 1 + or not isinstance(tensor, torch.Tensor) + or tensor.ndim == 0 + or int(tensor.shape[0]) != int(row_token_uids.numel()) + ): + return + row_count = int(row_token_uids.numel()) + padding_rows = row_token_uids < 0 + if row_count == 0 or not bool(padding_rows.any().item()): + return + flat = tensor.detach().reshape(row_count, -1) + if not bool((flat[padding_rows] == 0).all().item()): + return + valid_rows = torch.nonzero(~padding_rows, as_tuple=False).reshape(-1) + original_call = dict(call) + for key, value in original_call.items(): + if key == "row_token_uids": + continue + call[key] = cls._slice_row_aligned_value( + value, + row_indices=valid_rows, + total_rows=row_count, + ) + call["row_token_uids"] = row_token_uids.index_select(0, valid_rows).contiguous() + @staticmethod def _align_exact_zero_padding_row_token_uids(call: dict[str, Any]) -> None: """Moves padding UID markers onto exact-zero sequence-parallel pad rows.""" diff --git a/tests/integration/megatron/model_support/oracle_harness.py b/tests/integration/megatron/model_support/oracle_harness.py index 2ea15447e..db2ad58a3 100644 --- a/tests/integration/megatron/model_support/oracle_harness.py +++ b/tests/integration/megatron/model_support/oracle_harness.py @@ -256,6 +256,9 @@ def _without_context_parallel(topology: Topology) -> Topology: mutation: CP_ATTENTION_SENSITIVITY_TOPOLOGY for mutation in CP_ATTENTION_SENSITIVITY_MUTATIONS } +SENSITIVITY_TOPOLOGY_BY_MUTATION["attn_skip_flash_lse_normalize"] = Topology( + tp=1, ep=2, etp=1, dp=1, cp=4, sp=False +) SENSITIVITY_TOPOLOGY_BY_MUTATION["bwd_skip_sync_fc1_a"] = Topology( tp=2, ep=1, etp=2, dp=1, sp=True ) @@ -712,6 +715,8 @@ def sensitivity_topology_for_mutation( }: return DENSE_DP_SENSITIVITY_TOPOLOGY if mutation in CP_ATTENTION_SENSITIVITY_MUTATIONS: + if mutation == "attn_skip_flash_lse_normalize": + return Topology(tp=1, ep=1, etp=1, dp=1, cp=4, sp=False) return DENSE_CP_ATTENTION_SENSITIVITY_TOPOLOGY return DENSE_SENSITIVITY_TOPOLOGY return SENSITIVITY_TOPOLOGY_BY_MUTATION[mutation] @@ -837,210 +842,10 @@ def _build_packed_tensors( config: PackedTensorConfig, seed: int, ) -> dict[str, Any]: - """Generates deterministic synthetic packed tensors used in integration runs.""" - import torch - - if config.num_sequences <= 1: - raise ValueError("num_sequences must be greater than 1") - shape = (config.num_sequences, config.sequence_length) - generator = torch.Generator().manual_seed(seed) - tokens = torch.zeros(shape, dtype=torch.long) - token_low = 10 - token_span = max(1, config.vocab_high - token_low) - group_ids = torch.full(shape, -1, dtype=torch.long) - parent_ids = torch.full(shape, -1, dtype=torch.long) - assistant_mask = torch.zeros(shape, dtype=torch.bool) - input_pos = torch.zeros(shape, dtype=torch.long) - logprobs = torch.full(shape, float("nan"), dtype=torch.float32) - advantages = torch.zeros(shape, dtype=torch.float32) - weights = torch.zeros(shape, dtype=torch.float32) - - prefix_length = max(1, min(config.sequence_length - 1, config.prefill_tokens)) - max_completion_tokens = max(1, config.sequence_length - prefix_length) - base_completion_tokens = max(1, min(config.decode_tokens, max_completion_tokens)) - jitter_width = min(config.decode_tokens_jitter, max_completion_tokens - 1) - - def _sample_completion_length() -> int: - if jitter_width > 0: - jitter = int( - torch.randint( - low=-jitter_width, - high=jitter_width + 1, - size=(1,), - generator=generator, - dtype=torch.long, - ).item() - ) - else: - jitter = 0 - return max( - 1, - min(max_completion_tokens, base_completion_tokens + jitter), - ) - - def _sample_token_block(length: int) -> torch.Tensor: - return torch.randint( - low=token_low, - high=config.vocab_high, - size=(length,), - dtype=torch.long, - generator=generator, - ) - - def _sample_logprob_block(length: int) -> torch.Tensor: - return ( - torch.randn( - (length,), - generator=generator, - dtype=torch.float32, - ) - * 0.25 - - 1.75 - ) - - def _sample_advantage_value() -> float: - return float( - ( - torch.randn( - (1,), - generator=generator, - dtype=torch.float32, - ) - * 0.5 - ).item() - ) - - for sequence_index in range(config.num_sequences): - cursor = 0 - next_group_id = 0 - while cursor < config.sequence_length: - prompt_group_id = next_group_id - next_group_id += 1 - completion_lengths = [ - _sample_completion_length() - for _ in range(config.completion_branches_per_prefix) - ] - remaining = config.sequence_length - cursor + """Generates deterministic nested prefix-tree tensors used in integration runs.""" + from .prefix_tree_workloads import build_complex_prefix_tree_packed_tensors - if config.packing_mode == "stop_early": - included_completion_lengths = list(completion_lengths) - while ( - included_completion_lengths - and (prefix_length + sum(included_completion_lengths)) > remaining - ): - included_completion_lengths.pop() - if not included_completion_lengths: - break - - prompt_end = cursor + prefix_length - tokens[sequence_index, cursor:prompt_end] = _sample_token_block( - prefix_length - ) - group_ids[sequence_index, cursor:prompt_end] = prompt_group_id - parent_ids[sequence_index, cursor:prompt_end] = prompt_group_id - input_pos[sequence_index, cursor:prompt_end] = torch.arange( - prefix_length, dtype=torch.long - ) - cursor = prompt_end - - for completion_length in included_completion_lengths: - completion_group_id = next_group_id - next_group_id += 1 - completion_end = cursor + completion_length - tokens[sequence_index, cursor:completion_end] = _sample_token_block( - completion_length - ) - group_ids[sequence_index, cursor:completion_end] = ( - completion_group_id - ) - parent_ids[sequence_index, cursor:completion_end] = prompt_group_id - input_pos[sequence_index, cursor:completion_end] = torch.arange( - prefix_length, - prefix_length + completion_length, - dtype=torch.long, - ) - assistant_mask[sequence_index, cursor:completion_end] = True - logprobs[sequence_index, cursor:completion_end] = ( - _sample_logprob_block(completion_length) - ) - advantages[sequence_index, cursor:completion_end] = ( - _sample_advantage_value() - ) - weights[sequence_index, cursor:completion_end] = 1.0 - cursor = completion_end - continue - - prompt_take = min(prefix_length, remaining) - prompt_end = cursor + prompt_take - tokens[sequence_index, cursor:prompt_end] = _sample_token_block(prompt_take) - group_ids[sequence_index, cursor:prompt_end] = prompt_group_id - parent_ids[sequence_index, cursor:prompt_end] = prompt_group_id - input_pos[sequence_index, cursor:prompt_end] = torch.arange( - prompt_take, dtype=torch.long - ) - cursor = prompt_end - if cursor >= config.sequence_length: - break - - for completion_length in completion_lengths: - if cursor >= config.sequence_length: - break - completion_group_id = next_group_id - next_group_id += 1 - remaining = config.sequence_length - cursor - completion_take = min(completion_length, remaining) - completion_end = cursor + completion_take - tokens[sequence_index, cursor:completion_end] = _sample_token_block( - completion_take - ) - group_ids[sequence_index, cursor:completion_end] = completion_group_id - parent_ids[sequence_index, cursor:completion_end] = prompt_group_id - input_pos[sequence_index, cursor:completion_end] = torch.arange( - prefix_length, - prefix_length + completion_take, - dtype=torch.long, - ) - assistant_mask[sequence_index, cursor:completion_end] = True - logprobs[sequence_index, cursor:completion_end] = _sample_logprob_block( - completion_take - ) - advantages[sequence_index, cursor:completion_end] = ( - _sample_advantage_value() - ) - weights[sequence_index, cursor:completion_end] = 1.0 - cursor = completion_end - if completion_take < completion_length: - break - - # Ensure paired cross-DP rows are never token-identical across valid tokens. - half = config.num_sequences // 2 - if half > 0 and config.num_sequences % 2 == 0: - valid_lengths = (group_ids != -1).sum(dim=1) - for pair_index in range(half): - left_index = pair_index - right_index = pair_index + half - left_valid = int(valid_lengths[left_index].item()) - right_valid = int(valid_lengths[right_index].item()) - if left_valid != right_valid or left_valid == 0: - continue - if torch.equal( - tokens[left_index, :left_valid], tokens[right_index, :right_valid] - ): - tokens[right_index, 0] = ( - (tokens[right_index, 0] - token_low + 1) % token_span - ) + token_low - return { - "tokens": tokens, - "group_ids": group_ids, - "parent_ids": parent_ids, - "input_pos": input_pos, - "assistant_mask": assistant_mask, - "logprobs": logprobs, - "advantages": advantages, - "weights": weights, - "pixel_values": [None] * config.num_sequences, - "image_grid_thw": [None] * config.num_sequences, - } + return build_complex_prefix_tree_packed_tensors(config, seed) def _create_packed_tensors( @@ -1306,7 +1111,7 @@ def _sample_valid_lengths(self) -> tuple[int, ...]: if self._sample_valid_lengths_cache is not None: return self._sample_valid_lengths_cache from art.megatron.context_parallel.builder import ( - build_shared_prefix_attention_spec, + build_prefix_tree_attention_spec, ) from art.preprocessing.pack import packed_tensors_from_dir @@ -1317,7 +1122,7 @@ def _sample_valid_lengths(self) -> tuple[int, ...]: parent_ids = packed_tensors["parent_ids"] self._sample_valid_lengths_cache = tuple( int( - build_shared_prefix_attention_spec( + build_prefix_tree_attention_spec( group_ids=group_ids[row_index : row_index + 1], parent_ids=parent_ids[row_index : row_index + 1], ) @@ -1397,8 +1202,6 @@ def _trim_trace_padding( if not isinstance(sample_index, int): continue valid_length = int(valid_lengths[sample_index]) - if valid_length >= sequence_length: - continue row_token_uids = call.get("row_token_uids") if ( isinstance(row_token_uids, torch.Tensor) @@ -1409,9 +1212,7 @@ def _trim_trace_padding( (row_token_uids >= 0) & (local_token_uids < valid_length), as_tuple=False, ).reshape(-1) - if int(keep_rows.numel()) > 0 and int(keep_rows.numel()) < int( - row_token_uids.numel() - ): + if int(keep_rows.numel()) < int(row_token_uids.numel()): call["row_token_uids"] = row_token_uids.index_select( 0, keep_rows ).contiguous() @@ -1430,6 +1231,8 @@ def _trim_trace_padding( 0, keep_rows ).contiguous() continue + if valid_length >= sequence_length: + continue for key in ("primary_output", "router_topk_scores", "router_topk_ids"): tensor = call.get(key) if not isinstance(tensor, torch.Tensor) or tensor.ndim == 0: @@ -2250,12 +2053,6 @@ def run_sensitivity_suite( reports: list[VariantReport] = [] ran_any_variants = False for objective in selected_oracle_objectives(): - runner = VariantRunner( - objective=objective, - case_config=case_config, - oracle_flex_backend=oracle_flex_backend, - variant_flex_backend=variant_flex_backend, - ) objective_mutations = selected_sensitivity_mutations_for_objective( objective, mutations, @@ -2263,29 +2060,91 @@ def run_sensitivity_suite( ) if not objective_mutations: continue - variants = [] - for mutation in objective_mutations: - topology = sensitivity_topology_for_mutation( - mutation, - is_moe=case_config.is_moe, - ) - if max_world_size is not None and topology.world_size() > max_world_size: + for flex_backend, flex_mutations in ( + ( + None, + [ + mutation + for mutation in objective_mutations + if mutation != "attn_skip_flash_lse_normalize" + ], + ), + ( + "FLASH", + [ + mutation + for mutation in objective_mutations + if mutation == "attn_skip_flash_lse_normalize" + ], + ), + ): + if not flex_mutations: continue - variants.append( - VariantSpec( - name=f"{objective}_sensitivity_{mutation}", - objective=objective, - topology=topology, - mutation=mutation, - expected_signal="fail", - pass_fn_by_phase=phase_pass, - flex_backend=variant_flex_backend, + oracle_slug = ( + None + if flex_backend is None + else oracle_output_slug( + objective, + oracle_topology(is_moe=case_config.is_moe), + "flash", ) ) - if not variants: - continue - ran_any_variants = True - reports.extend(runner.run_suite(variants)) + runner_case_config = ( + case_config + if flex_backend is None or case_config.precision == "bf16" + else case_config.model_copy(update={"precision": "bf16"}) + ) + runner = VariantRunner( + objective=objective, + case_config=runner_case_config, + oracle_flex_backend=( + oracle_flex_backend if flex_backend is None else flex_backend + ), + variant_flex_backend=( + variant_flex_backend if flex_backend is None else flex_backend + ), + oracle_slug_override=oracle_slug, + ) + variants = [] + for mutation in flex_mutations: + topology = sensitivity_topology_for_mutation( + mutation, + is_moe=runner_case_config.is_moe, + ) + if ( + max_world_size is not None + and topology.world_size() > max_world_size + ): + continue + variants.append( + VariantSpec( + name=f"{objective}_sensitivity_{mutation}", + objective=objective, + topology=topology, + output_slug=( + None + if flex_backend is None + else oracle_output_slug( + objective, + topology, + f"{mutation}_flash", + ) + ), + reference_slug=oracle_slug, + mutation=mutation, + expected_signal="fail", + pass_fn_by_phase=phase_pass, + flex_backend=( + variant_flex_backend + if flex_backend is None + else flex_backend + ), + ) + ) + if not variants: + continue + ran_any_variants = True + reports.extend(runner.run_suite(variants)) if ran_any_variants: return reports requested = ", ".join(mutations) diff --git a/tests/integration/megatron/model_support/oracle_worker.py b/tests/integration/megatron/model_support/oracle_worker.py index 42da1eab5..85fc0f596 100644 --- a/tests/integration/megatron/model_support/oracle_worker.py +++ b/tests/integration/megatron/model_support/oracle_worker.py @@ -788,7 +788,8 @@ def _matches_grad_sync_skip_mutation( return ".self_attention.linear_proj.lora.B_T" in param_name if mutation == "bwd_skip_sync_fc1_a": return ( - ".mlp.experts.linear_fc1.gate_lora.A_T" in param_name + ".mlp.experts.linear_fc1.lora.A_T" in param_name + or ".mlp.experts.linear_fc1.gate_lora.A_T" in param_name or ".mlp.experts.linear_fc1.up_lora.A_T" in param_name or ".mlp.linear_fc1.gate_lora.A_T" in param_name or ".mlp.linear_fc1.up_lora.A_T" in param_name @@ -1136,12 +1137,16 @@ def _reference_forward( def _reference_fc1_forward(self: Any, x: torch.Tensor, tokens_per_expert: Any): base_out, bias_out = self.linear_fc1(x, tokens_per_expert) - adapter_out = torch.cat( - ( - self.gate_lora(x, tokens_per_expert), - self.up_lora(x, tokens_per_expert), - ), - dim=1, + adapter_out = ( + self.lora(x, tokens_per_expert) + if self.fused_gate_up + else torch.cat( + ( + self.gate_lora(x, tokens_per_expert), + self.up_lora(x, tokens_per_expert), + ), + dim=1, + ) ) return base_out + adapter_out, bias_out diff --git a/tests/integration/megatron/model_support/packed_position_ids.py b/tests/integration/megatron/model_support/packed_position_ids.py index 397e7ce12..ba8567923 100644 --- a/tests/integration/megatron/model_support/packed_position_ids.py +++ b/tests/integration/megatron/model_support/packed_position_ids.py @@ -16,7 +16,8 @@ from art.megatron import train as megatron_train from art.megatron.model_support.discovery import inspect_architecture -from art.megatron.shared_prefix_state import create_shared_prefix_state +from art.megatron.prefix_tree import parse_prefix_tree_row +from art.megatron.prefix_tree_state import create_prefix_tree_state from ..artifacts import GitRepoState, pinned_git_state from .oracle_harness import ( @@ -34,10 +35,11 @@ _configure_provider, provider_topology_env, ) +from .prefix_tree_workloads import build_complex_prefix_tree_packed_tensors # Qwen3.5/3.6 hybrid MoE runs show small shape-dependent logit drift between # the single packed forward and many shorter reference forwards, even when the -# rotary grouping and shared-prefix semantics are correct. Keep the bound tight, +# rotary grouping and prefix-tree semantics are correct. Keep the bound tight, # but above the observed ~0.13% truncate-case jitter. _LOGITS_MEAN_ABS_PCT_LIMIT = 0.2 _DEBUG_ENV = "ART_PACKED_POSITION_IDS_DEBUG" @@ -285,263 +287,38 @@ def _build_art_realistic_packed_tensors( config: PackedTensorConfig, seed: int, ) -> dict[str, Any]: - if config.num_sequences <= 1: - raise ValueError("num_sequences must be greater than 1") - if config.prefill_tokens < 2: - raise ValueError( - "prefill_tokens must be at least 2 to build ART-style branch context" - ) - if config.sequence_length < 3: - raise ValueError( - "sequence_length must leave room for shared prompt, branch context, " - "and at least one trainable token" - ) - - shape = (config.num_sequences, config.sequence_length) - generator = torch.Generator().manual_seed(seed) - tokens = torch.zeros(shape, dtype=torch.long) - group_ids = torch.full(shape, -1, dtype=torch.long) - parent_ids = torch.full(shape, -1, dtype=torch.long) - input_pos = torch.zeros(shape, dtype=torch.long) - assistant_mask = torch.zeros(shape, dtype=torch.bool) - logprobs = torch.full(shape, float("nan"), dtype=torch.float32) - advantages = torch.zeros(shape, dtype=torch.float32) - weights = torch.zeros(shape, dtype=torch.float32) - - first_trainable_pos = max(2, min(config.sequence_length - 1, config.prefill_tokens)) - shared_prompt_length = first_trainable_pos - 1 - max_completion_tokens = max(1, config.sequence_length - first_trainable_pos) - base_completion_tokens = max(1, min(config.decode_tokens, max_completion_tokens)) - jitter_width = min(config.decode_tokens_jitter, max_completion_tokens - 1) - token_low = 10 - token_span = max(1, config.vocab_high - token_low) - - def _sample_completion_length() -> int: - if jitter_width > 0: - jitter = int( - torch.randint( - low=-jitter_width, - high=jitter_width + 1, - size=(1,), - generator=generator, - dtype=torch.long, - ).item() - ) - else: - jitter = 0 - return max(1, min(max_completion_tokens, base_completion_tokens + jitter)) - - def _sample_token_block(length: int) -> torch.Tensor: - return torch.randint( - low=token_low, - high=config.vocab_high, - size=(length,), - dtype=torch.long, - generator=generator, - ) - - def _sample_logprob_block(length: int) -> torch.Tensor: - return ( - torch.randn((length,), generator=generator, dtype=torch.float32) * 0.25 - - 1.75 - ) - - def _sample_advantage_value() -> float: - return float( - (torch.randn((1,), generator=generator, dtype=torch.float32) * 0.5).item() - ) - - def _write_prompt( - sequence_index: int, - cursor: int, - prompt_group_id: int, - ) -> tuple[int, int]: - prompt_tokens = _sample_token_block(first_trainable_pos) - prompt_end = cursor + shared_prompt_length - tokens[sequence_index, cursor:prompt_end] = prompt_tokens[:shared_prompt_length] - group_ids[sequence_index, cursor:prompt_end] = prompt_group_id - parent_ids[sequence_index, cursor:prompt_end] = prompt_group_id - input_pos[sequence_index, cursor:prompt_end] = torch.arange( - shared_prompt_length, - dtype=torch.long, - ) - return prompt_end, int(prompt_tokens[shared_prompt_length].item()) - - def _write_branch( - sequence_index: int, - cursor: int, - completion_group_id: int, - prompt_group_id: int, - context_token: int, - completion_length: int, - ) -> int: - branch_end = cursor + 1 + completion_length - tokens[sequence_index, cursor] = context_token - tokens[sequence_index, cursor + 1 : branch_end] = _sample_token_block( - completion_length - ) - group_ids[sequence_index, cursor:branch_end] = completion_group_id - parent_ids[sequence_index, cursor:branch_end] = prompt_group_id - input_pos[sequence_index, cursor:branch_end] = torch.arange( - shared_prompt_length, - shared_prompt_length + 1 + completion_length, - dtype=torch.long, - ) - trainable_start = cursor + 1 - assistant_mask[sequence_index, trainable_start:branch_end] = True - logprobs[sequence_index, trainable_start:branch_end] = _sample_logprob_block( - completion_length - ) - advantages[sequence_index, trainable_start:branch_end] = ( - _sample_advantage_value() - ) - weights[sequence_index, trainable_start:branch_end] = 1.0 / completion_length - return branch_end + return build_complex_prefix_tree_packed_tensors(config, seed) - for sequence_index in range(config.num_sequences): - cursor = 0 - next_group_id = 0 - while cursor < config.sequence_length: - prompt_group_id = next_group_id - next_group_id += 1 - completion_lengths = [ - _sample_completion_length() - for _ in range(config.completion_branches_per_prefix) - ] - remaining = config.sequence_length - cursor - if remaining <= shared_prompt_length + 1: - break - - if config.packing_mode == "stop_early": - included_completion_lengths = list(completion_lengths) - while included_completion_lengths and ( - shared_prompt_length - + sum(1 + length for length in included_completion_lengths) - > remaining - ): - included_completion_lengths.pop() - if not included_completion_lengths: - break - - cursor, context_token = _write_prompt( - sequence_index, - cursor, - prompt_group_id, - ) - for completion_length in included_completion_lengths: - completion_group_id = next_group_id - next_group_id += 1 - cursor = _write_branch( - sequence_index, - cursor, - completion_group_id, - prompt_group_id, - context_token, - completion_length, - ) - continue - cursor, context_token = _write_prompt( - sequence_index, - cursor, - prompt_group_id, - ) - for completion_length in completion_lengths: - remaining = config.sequence_length - cursor - if remaining <= 1: - break - completion_take = min(completion_length, remaining - 1) - completion_group_id = next_group_id - next_group_id += 1 - cursor = _write_branch( - sequence_index, - cursor, - completion_group_id, - prompt_group_id, - context_token, - completion_take, - ) - - half = config.num_sequences // 2 - if half > 0 and config.num_sequences % 2 == 0: - valid_lengths = (group_ids != -1).sum(dim=1) - for pair_index in range(half): - left_index = pair_index - right_index = pair_index + half - left_valid = int(valid_lengths[left_index].item()) - right_valid = int(valid_lengths[right_index].item()) - if left_valid != right_valid or left_valid == 0: - continue - if torch.equal( - tokens[left_index, :left_valid], - tokens[right_index, :right_valid], - ): - tokens[right_index, 0] = ( - (tokens[right_index, 0] - token_low + 1) % token_span - ) + token_low - - weights = torch.where(assistant_mask, weights, torch.zeros_like(weights)) - if bool(assistant_mask.any().item()): - weights[assistant_mask] /= weights[assistant_mask].mean() - advantages = torch.where( - assistant_mask, - advantages, - torch.zeros_like(advantages), - ) - advantage_scale = ( - advantages[assistant_mask].abs() * weights[assistant_mask] - ).mean() - if float(advantage_scale.item()) > 0.0: - advantages[assistant_mask] /= advantage_scale - - return { - "tokens": tokens, - "group_ids": group_ids, - "parent_ids": parent_ids, - "input_pos": input_pos, - "assistant_mask": assistant_mask, - "logprobs": logprobs, - "advantages": advantages, - "weights": weights, - "pixel_values": [None] * config.num_sequences, - "image_grid_thw": [None] * config.num_sequences, - } - - -def _prompt_family_segments( +def _prefix_tree_leaf_paths( group_ids: torch.Tensor, parent_ids: torch.Tensor, *, - required_completion_count: int = 2, -) -> list[tuple[tuple[int, int], list[tuple[int, int]]]]: - families: list[tuple[tuple[int, int], list[tuple[int, int]]]] = [] - valid_tokens = int((group_ids != -1).sum().item()) - cursor = 0 - while cursor < valid_tokens: - group_id = int(group_ids[cursor].item()) - parent_id = int(parent_ids[cursor].item()) - prompt_start = cursor - while cursor < valid_tokens and int(group_ids[cursor].item()) == group_id: - cursor += 1 - prompt_end = cursor - if group_id != parent_id: + required_leaf_count: int = 2, +) -> list[tuple[tuple[tuple[int, int], ...], tuple[int, int]]]: + tree = parse_prefix_tree_row(group_ids=group_ids, parent_ids=parent_ids) + segment_by_group = {segment.group_id: segment for segment in tree.segments} + child_count_by_group: dict[int, int] = {} + for segment in tree.segments: + if segment.group_id == segment.parent_id: continue - completions: list[tuple[int, int]] = [] - while cursor < valid_tokens: - completion_group_id = int(group_ids[cursor].item()) - completion_parent_id = int(parent_ids[cursor].item()) - if completion_parent_id != group_id or completion_group_id == group_id: - break - completion_start = cursor - while ( - cursor < valid_tokens - and int(group_ids[cursor].item()) == completion_group_id - ): - cursor += 1 - completions.append((completion_start, cursor)) - if len(completions) >= required_completion_count: - families.append(((prompt_start, prompt_end), completions)) - return families + child_count_by_group[segment.parent_id] = ( + child_count_by_group.get(segment.parent_id, 0) + 1 + ) + paths = [ + ( + tuple( + (segment_by_group[group_id].start, segment_by_group[group_id].end) + for group_id in leaf.ancestors + ), + (leaf.start, leaf.end), + ) + for leaf in tree.segments + if leaf.group_id != leaf.parent_id + and child_count_by_group.get(leaf.group_id, 0) == 0 + and leaf.end - leaf.start >= 2 + ] + return paths if len(paths) >= required_leaf_count else [] def _run_logits( @@ -600,13 +377,13 @@ def _logits_equivalence_check( for row_index in range(int(input_ids.shape[0])): row_group_ids = group_ids[row_index : row_index + 1] row_parent_ids = parent_ids[row_index : row_index + 1] - families = _prompt_family_segments(row_group_ids[0], row_parent_ids[0]) - if not families: - _debug_log(f"logits_check row={row_index} skipped no prompt family") + leaf_paths = _prefix_tree_leaf_paths(row_group_ids[0], row_parent_ids[0]) + if not leaf_paths: + _debug_log(f"logits_check row={row_index} skipped no prefix-tree leaves") continue row_input_ids = input_ids[row_index : row_index + 1] row_position_ids = position_ids[row_index : row_index + 1] - packed_bias = create_shared_prefix_state( + packed_bias = create_prefix_tree_state( group_ids=row_group_ids, parent_ids=row_parent_ids, input_pos=row_position_ids, @@ -618,7 +395,7 @@ def _logits_equivalence_check( attention_head_dim=getattr(provider, "kv_channels", None), attention_value_head_dim=getattr(provider, "kv_channels", None), ) - _debug_log(f"logits_check row={row_index} families={len(families)}") + _debug_log(f"logits_check row={row_index} leaves={len(leaf_paths)}") packed_logits = _time_block( f"logits_check row={row_index} packed_forward", lambda: _run_logits( @@ -630,96 +407,63 @@ def _logits_equivalence_check( ), device=row_input_ids.device, ) - for family_index, (prompt_segment, completion_segments) in enumerate(families): - prompt_start, prompt_end = prompt_segment + for leaf_index, (ancestor_segments, leaf_segment) in enumerate(leaf_paths): + leaf_start, leaf_end = leaf_segment + prompt_len = sum(end - start for start, end in ancestor_segments) + reference_segments = (*ancestor_segments, leaf_segment) _debug_log( "logits_check row=" - f"{row_index} family={family_index} " - f"prompt=({prompt_start},{prompt_end}) " - f"completions={completion_segments}" + f"{row_index} leaf={leaf_index} " + f"ancestors={ancestor_segments} leaf={leaf_segment}" + ) + reference_input_ids = torch.cat( + tuple(row_input_ids[:, start:end] for start, end in reference_segments), + dim=1, + ) + reference_position_ids = torch.cat( + tuple( + row_position_ids[:, start:end] for start, end in reference_segments + ), + dim=1, + ) + reference_group_ids = torch.zeros_like(reference_input_ids) + reference_parent_ids = torch.zeros_like(reference_input_ids) + reference_bias = create_prefix_tree_state( + group_ids=reference_group_ids, + parent_ids=reference_parent_ids, + input_pos=reference_position_ids, + sliding_windows=sliding_windows, + build_gdn_execution_spec=bool( + getattr(handler, "build_gdn_execution_spec", False) + ), + model_support_handler=handler, + attention_head_dim=getattr(provider, "kv_channels", None), + attention_value_head_dim=getattr(provider, "kv_channels", None), + ) + reference_logits = _time_block( + f"logits_check row={row_index} leaf={leaf_index} reference_forward", + lambda: _run_logits( + model=model, + handler=handler, + input_ids=reference_input_ids, + position_ids=reference_position_ids, + attention_bias=reference_bias, + ), + device=reference_input_ids.device, + ) + packed_completion_logits = packed_logits[:, leaf_start : leaf_end - 1, :] + reference_completion_logits = reference_logits[:, prompt_len:-1, :] + diff = (packed_completion_logits - reference_completion_logits).abs() + logits_abs_sum += float(diff.sum().item()) + logits_ref_abs_sum += float(reference_completion_logits.abs().sum().item()) + logits_numel += int(diff.numel()) + logits_max_abs_diff = max(logits_max_abs_diff, float(diff.max().item())) + completion_pair_count += 1 + _debug_log( + "logits_check row=" + f"{row_index} leaf={leaf_index} " + f"max_abs_diff={float(diff.max().item()):.6f}" ) - for completion_index, (completion_start, completion_end) in enumerate( - completion_segments - ): - reference_input_ids = torch.cat( - ( - row_input_ids[:, prompt_start:prompt_end], - row_input_ids[:, completion_start:completion_end], - ), - dim=1, - ) - reference_position_ids = torch.cat( - ( - row_position_ids[:, prompt_start:prompt_end], - row_position_ids[:, completion_start:completion_end], - ), - dim=1, - ) - reference_group_ids = torch.zeros_like(reference_input_ids) - reference_parent_ids = torch.zeros_like(reference_input_ids) - reference_bias = create_shared_prefix_state( - group_ids=reference_group_ids, - parent_ids=reference_parent_ids, - input_pos=reference_position_ids, - sliding_windows=sliding_windows, - build_gdn_execution_spec=bool( - getattr(handler, "build_gdn_execution_spec", False) - ), - model_support_handler=handler, - attention_head_dim=getattr(provider, "kv_channels", None), - attention_value_head_dim=getattr(provider, "kv_channels", None), - ) - _debug_log( - "logits_check row=" - f"{row_index} family={family_index} " - f"completion={completion_index} " - f"segment=({completion_start},{completion_end}) " - f"reference_seq={int(reference_input_ids.shape[1])}" - ) - reference_logits = _time_block( - ( - f"logits_check row={row_index} " - f"family={family_index} " - f"completion={completion_index} reference_forward" - ), - lambda: _run_logits( - model=model, - handler=handler, - input_ids=reference_input_ids, - position_ids=reference_position_ids, - attention_bias=reference_bias, - ), - device=reference_input_ids.device, - ) - if completion_end - completion_start < 2: - continue - packed_completion_logits = packed_logits[ - :, - completion_start : completion_end - 1, - :, - ] - reference_completion_logits = reference_logits[ - :, - prompt_end - prompt_start : -1, - :, - ] - diff = (packed_completion_logits - reference_completion_logits).abs() - logits_abs_sum += float(diff.sum().item()) - logits_ref_abs_sum += float( - reference_completion_logits.abs().sum().item() - ) - logits_numel += int(diff.numel()) - logits_max_abs_diff = max( - logits_max_abs_diff, - float(diff.max().item()), - ) - completion_pair_count += 1 - _debug_log( - "logits_check row=" - f"{row_index} family={family_index} " - f"completion={completion_index} " - f"max_abs_diff={float(diff.max().item()):.6f}" - ) if completion_pair_count > 0: mean_abs = logits_abs_sum / max(logits_numel, 1) typical_abs = logits_ref_abs_sum / max(logits_numel, 1) diff --git a/tests/integration/megatron/model_support/prefix_tree_workloads.py b/tests/integration/megatron/model_support/prefix_tree_workloads.py new file mode 100644 index 000000000..a24fa128f --- /dev/null +++ b/tests/integration/megatron/model_support/prefix_tree_workloads.py @@ -0,0 +1,285 @@ +from __future__ import annotations + +from typing import Any + +import torch + + +def build_complex_prefix_tree_packed_tensors( + config: Any, + seed: int, +) -> dict[str, Any]: + """Build a deterministic nested prefix-tree packed workload. + + Each packed row repeats this tree while it fits: + + root + -> mid_a -> leaf_a_short, leaf_a_long + -> mid_b -> leaf_b + -> direct_leaf + + Internal nodes are non-trainable. Each leaf starts with one non-trainable + context token followed by trainable completion tokens, matching production + prefix-tree packing's trainable-token boundary. + """ + + num_sequences = int(config.num_sequences) + sequence_length = int(config.sequence_length) + if num_sequences <= 1: + raise ValueError("num_sequences must be greater than 1") + if sequence_length < 16: + raise ValueError("sequence_length must leave room for a nested prefix tree") + + generator = torch.Generator().manual_seed(seed) + shape = (num_sequences, sequence_length) + tokens = torch.zeros(shape, dtype=torch.long) + group_ids = torch.full(shape, -1, dtype=torch.long) + parent_ids = torch.full(shape, -1, dtype=torch.long) + input_pos = torch.zeros(shape, dtype=torch.long) + assistant_mask = torch.zeros(shape, dtype=torch.bool) + logprobs = torch.full(shape, float("nan"), dtype=torch.float32) + advantages = torch.zeros(shape, dtype=torch.float32) + weights = torch.zeros(shape, dtype=torch.float32) + + token_low = 10 + vocab_high = max(token_low + 1, int(config.vocab_high)) + token_span = vocab_high - token_low + prefill_tokens = max(4, min(sequence_length - 8, int(config.prefill_tokens))) + root_len = max(2, prefill_tokens // 2) + mid_budget = max(2, prefill_tokens - root_len) + mid_a_len = max(1, mid_budget // 2) + mid_b_len = max(1, mid_budget - mid_a_len) + max_completion_tokens = max(1, sequence_length - root_len - mid_a_len - 2) + base_completion_tokens = max( + 1, + min(int(config.decode_tokens), max_completion_tokens), + ) + jitter_width = min(int(config.decode_tokens_jitter), max_completion_tokens - 1) + + def sample_completion_length() -> int: + jitter = ( + int( + torch.randint( + low=-jitter_width, + high=jitter_width + 1, + size=(1,), + generator=generator, + dtype=torch.long, + ).item() + ) + if jitter_width > 0 + else 0 + ) + return max(1, min(max_completion_tokens, base_completion_tokens + jitter)) + + def sample_tokens(length: int) -> torch.Tensor: + return torch.randint( + low=token_low, + high=vocab_high, + size=(length,), + dtype=torch.long, + generator=generator, + ) + + def sample_logprobs(length: int) -> torch.Tensor: + return ( + torch.randn((length,), generator=generator, dtype=torch.float32) * 0.25 + - 1.75 + ) + + def sample_advantage() -> float: + return float( + (torch.randn((1,), generator=generator, dtype=torch.float32) * 0.5).item() + ) + + def write_segment( + *, + sequence_index: int, + cursor: int, + group_id: int, + parent_id: int, + length: int, + input_start: int, + trainable_offset: int | None = None, + ) -> int: + end = min(sequence_length, cursor + length) + take = end - cursor + if take <= 0: + return cursor + tokens[sequence_index, cursor:end] = sample_tokens(take) + group_ids[sequence_index, cursor:end] = group_id + parent_ids[sequence_index, cursor:end] = parent_id + input_pos[sequence_index, cursor:end] = torch.arange( + input_start, + input_start + take, + dtype=torch.long, + ) + if trainable_offset is not None and take > trainable_offset: + train_start = cursor + trainable_offset + train_len = end - train_start + assistant_mask[sequence_index, train_start:end] = True + logprobs[sequence_index, train_start:end] = sample_logprobs(train_len) + advantages[sequence_index, train_start:end] = sample_advantage() + weights[sequence_index, train_start:end] = 1.0 + return end + + def tree_token_budget(leaf_lengths: tuple[int, int, int, int]) -> int: + return ( + root_len + + mid_a_len + + mid_b_len + + sum(1 + length for length in leaf_lengths) + ) + + def fit_leaf_lengths( + leaf_lengths: tuple[int, int, int, int], + *, + remaining: int, + ) -> tuple[int, int, int, int] | None: + completion_budget = ( + remaining - root_len - mid_a_len - mid_b_len - len(leaf_lengths) + ) + if completion_budget < len(leaf_lengths): + return None + if sum(leaf_lengths) <= completion_budget: + return leaf_lengths + fitted = list(leaf_lengths) + overflow = sum(fitted) - completion_budget + while overflow > 0: + index = max(range(len(fitted)), key=lambda candidate: fitted[candidate]) + reducible = fitted[index] - 1 + if reducible <= 0: + return None + reduction = min(reducible, overflow) + fitted[index] -= reduction + overflow -= reduction + return (fitted[0], fitted[1], fitted[2], fitted[3]) + + for sequence_index in range(num_sequences): + cursor = 0 + next_group_id = 0 + while cursor < sequence_length: + leaf_lengths = ( + sample_completion_length(), + sample_completion_length(), + sample_completion_length(), + sample_completion_length(), + ) + remaining = sequence_length - cursor + if tree_token_budget(leaf_lengths) > remaining: + fitted_leaf_lengths = fit_leaf_lengths( + leaf_lengths, + remaining=remaining, + ) + if fitted_leaf_lengths is None: + break + leaf_lengths = fitted_leaf_lengths + root_group = next_group_id + next_group_id += 1 + cursor = write_segment( + sequence_index=sequence_index, + cursor=cursor, + group_id=root_group, + parent_id=root_group, + length=root_len, + input_start=0, + ) + mid_a_group = next_group_id + next_group_id += 1 + cursor = write_segment( + sequence_index=sequence_index, + cursor=cursor, + group_id=mid_a_group, + parent_id=root_group, + length=mid_a_len, + input_start=root_len, + ) + for leaf_length in leaf_lengths[:2]: + leaf_group = next_group_id + next_group_id += 1 + cursor = write_segment( + sequence_index=sequence_index, + cursor=cursor, + group_id=leaf_group, + parent_id=mid_a_group, + length=1 + leaf_length, + input_start=root_len + mid_a_len, + trainable_offset=1, + ) + mid_b_group = next_group_id + next_group_id += 1 + cursor = write_segment( + sequence_index=sequence_index, + cursor=cursor, + group_id=mid_b_group, + parent_id=root_group, + length=mid_b_len, + input_start=root_len, + ) + leaf_b_group = next_group_id + next_group_id += 1 + cursor = write_segment( + sequence_index=sequence_index, + cursor=cursor, + group_id=leaf_b_group, + parent_id=mid_b_group, + length=1 + leaf_lengths[2], + input_start=root_len + mid_b_len, + trainable_offset=1, + ) + direct_group = next_group_id + next_group_id += 1 + cursor = write_segment( + sequence_index=sequence_index, + cursor=cursor, + group_id=direct_group, + parent_id=root_group, + length=1 + leaf_lengths[3], + input_start=root_len, + trainable_offset=1, + ) + + half = num_sequences // 2 + if half > 0 and num_sequences % 2 == 0: + valid_lengths = (group_ids != -1).sum(dim=1) + for pair_index in range(half): + left_index = pair_index + right_index = pair_index + half + left_valid = int(valid_lengths[left_index].item()) + right_valid = int(valid_lengths[right_index].item()) + if left_valid != right_valid or left_valid == 0: + continue + if torch.equal( + tokens[left_index, :left_valid], + tokens[right_index, :right_valid], + ): + tokens[right_index, 0] = ( + (tokens[right_index, 0] - token_low + 1) % token_span + ) + token_low + + weights = torch.where(assistant_mask, weights, torch.zeros_like(weights)) + if bool(assistant_mask.any().item()): + weights[assistant_mask] /= weights[assistant_mask].mean() + advantages = torch.where( + assistant_mask, + advantages, + torch.zeros_like(advantages), + ) + advantage_scale = ( + advantages[assistant_mask].abs() * weights[assistant_mask] + ).mean() + if float(advantage_scale.item()) > 0.0: + advantages[assistant_mask] /= advantage_scale + + return { + "tokens": tokens, + "group_ids": group_ids, + "parent_ids": parent_ids, + "input_pos": input_pos, + "assistant_mask": assistant_mask, + "logprobs": logprobs, + "advantages": advantages, + "weights": weights, + "pixel_values": [None] * num_sequences, + "image_grid_thw": [None] * num_sequences, + } diff --git a/tests/integration/megatron/model_support/test_compile_flags.py b/tests/integration/megatron/model_support/test_compile_flags.py index 15654fc09..70641d6c7 100644 --- a/tests/integration/megatron/model_support/test_compile_flags.py +++ b/tests/integration/megatron/model_support/test_compile_flags.py @@ -1,3 +1,7 @@ +from art.megatron.model_support.handlers.gemma4 import ( + GEMMA4_DENSE_HANDLER, + GEMMA4_MOE_HANDLER, +) from art.megatron.model_support.handlers.qwen3_5 import QWEN3_5_MOE_HANDLER from art.megatron.model_support.handlers.qwen3_moe import QWEN3_MOE_HANDLER @@ -31,3 +35,31 @@ def test_qwen35_moe_compile_workarounds_cover_deepep_permute_restore() -> None: config = QWEN3_5_MOE_HANDLER.compile_workaround_config(provider) assert config.flags == _QWEN35_MOE_COMPILE_FLAGS assert config.unconditional_flags == () + + +def test_gemma4_wide_global_attention_uses_lower_triton_stage_count() -> None: + provider = type("Provider", (), {"global_head_dim": 512})() + + assert GEMMA4_DENSE_HANDLER.flex_attention_compile_crash_config( + provider + ).triton_num_stages_2_head_dims == (512,) + assert GEMMA4_MOE_HANDLER.flex_attention_compile_crash_config( + provider + ).triton_num_stages_2_head_dims == (512,) + + +def test_gemma4_standard_global_attention_keeps_default_triton_stage_count() -> None: + provider = type("Provider", (), {"global_head_dim": 256})() + + assert ( + GEMMA4_DENSE_HANDLER.flex_attention_compile_crash_config( + provider + ).triton_num_stages_2_head_dims + == () + ) + assert ( + GEMMA4_MOE_HANDLER.flex_attention_compile_crash_config( + provider + ).triton_num_stages_2_head_dims + == () + ) diff --git a/tests/integration/megatron/model_support/test_oracle_harness_invariants.py b/tests/integration/megatron/model_support/test_oracle_harness_invariants.py index 991e925ec..c0e702eca 100644 --- a/tests/integration/megatron/model_support/test_oracle_harness_invariants.py +++ b/tests/integration/megatron/model_support/test_oracle_harness_invariants.py @@ -30,6 +30,8 @@ selected_sensitivity_mutations_for_objective, sensitivity_topology_for_mutation, ) +from .oracle_worker import _matches_grad_sync_skip_mutation +from .prefix_tree_workloads import build_complex_prefix_tree_packed_tensors def _metric_row( @@ -98,6 +100,25 @@ def _expert_trace_call( } +def test_fc1_grad_sync_sensitivity_matches_split_and_fused_lora_names() -> None: + assert _matches_grad_sync_skip_mutation( + "chunk0.module.decoder.layers.0.mlp.experts.linear_fc1.lora.A_T", + "bwd_skip_sync_fc1_a", + ) + assert _matches_grad_sync_skip_mutation( + "chunk0.module.decoder.layers.0.mlp.experts.linear_fc1.gate_lora.A_T", + "bwd_skip_sync_fc1_a", + ) + assert _matches_grad_sync_skip_mutation( + "chunk0.module.decoder.layers.0.mlp.experts.linear_fc1.up_lora.A_T", + "bwd_skip_sync_fc1_a", + ) + assert not _matches_grad_sync_skip_mutation( + "chunk0.module.decoder.layers.0.mlp.experts.linear_fc2.lora.A_T", + "bwd_skip_sync_fc1_a", + ) + + def test_metric_threshold_rule_can_require_strictly_positive_values() -> None: rule = MetricThresholdRule(minimums={"candidate_abs_scale": 0.0}) @@ -412,6 +433,37 @@ def test_forward_trace_canonicalizes_row_outputs_by_token_uid() -> None: ) +def test_forward_trace_drops_exact_zero_padding_rows() -> None: + trace: dict[str, list[dict[str, Any]]] = { + "chunk0.module.decoder.layers.0.self_attention.out_proj": [ + { + "primary_output": torch.tensor( + [[0.0, 0.0], [30.0, 31.0], [10.0, 11.0], [20.0, 21.0]] + ), + "output": { + "hidden": torch.tensor( + [[0.0, 0.0], [3.0, 3.1], [1.0, 1.1], [2.0, 2.1]] + ) + }, + "row_token_uids": torch.tensor([-1, 3, 1, 2]), + } + ] + } + + ForwardTraceCapture.canonicalize_trace(trace) + + call = trace["chunk0.module.decoder.layers.0.self_attention.out_proj"][0] + assert torch.equal(call["row_token_uids"], torch.tensor([1, 2, 3])) + assert torch.equal( + call["primary_output"], + torch.tensor([[10.0, 11.0], [20.0, 21.0], [30.0, 31.0]]), + ) + assert torch.equal( + call["output"]["hidden"], + torch.tensor([[1.0, 1.1], [2.0, 2.1], [3.0, 3.1]]), + ) + + def test_forward_trace_canonicalizes_router_outputs_without_output_attrs() -> None: module = type("RouterWithDenseTraceUids", (), {})() probs = torch.tensor([[3.0], [1.0], [2.0]]) @@ -982,6 +1034,14 @@ def test_dense_sensitivity_keeps_dp_and_cp_attention_cases() -> None: ) == DENSE_CP_ATTENTION_SENSITIVITY_TOPOLOGY ) + assert sensitivity_topology_for_mutation( + "attn_skip_flash_lse_normalize", + is_moe=False, + ) == Topology(tp=1, ep=1, etp=1, dp=1, cp=4, sp=False) + assert sensitivity_topology_for_mutation( + "attn_skip_flash_lse_normalize", + is_moe=True, + ) == Topology(tp=1, ep=2, etp=1, dp=1, cp=4, sp=False) def test_case_config_base_model_can_be_overridden_by_env( @@ -1004,3 +1064,22 @@ def test_packed_tensor_defaults_match_main_rebase_oracle_tokens() -> None: assert config.decode_tokens_jitter == 32 assert config.packing_mode == "stop_early" assert config.vocab_high == 8192 + + +def test_prefix_tree_workload_fits_hf_parity_packed_size() -> None: + packed_tensors = build_complex_prefix_tree_packed_tensors( + PackedTensorConfig( + num_sequences=4, + sequence_length=256, + prefill_tokens=64, + completion_branches_per_prefix=2, + decode_tokens=64, + decode_tokens_jitter=32, + packing_mode="stop_early", + ), + seed=20260304, + ) + + assert int((packed_tensors["group_ids"] != -1).sum().item()) > 0 + assert int(packed_tensors["assistant_mask"].sum().item()) > 0 + assert int((packed_tensors["weights"] != 0).sum().item()) > 0 diff --git a/tests/integration/megatron/model_support/test_workflow.py b/tests/integration/megatron/model_support/test_workflow.py index 118fef764..1cef0ad01 100644 --- a/tests/integration/megatron/model_support/test_workflow.py +++ b/tests/integration/megatron/model_support/test_workflow.py @@ -11,6 +11,8 @@ from .validation_spec import ValidationReport, ValidationStageResult from .workflow import ( + INCLUDE_FLASH_SENSITIVITY_ENV, + KEEP_TOPOLOGY_ARTIFACTS_ENV, MANDATORY_VALIDATION_STAGES, NATIVE_VLLM_LORA_STAGE, SKIP_SENSITIVITY_ENV, @@ -39,6 +41,7 @@ @pytest.fixture(autouse=True) def _stub_pinned_git_state(monkeypatch) -> None: + monkeypatch.delenv(INCLUDE_FLASH_SENSITIVITY_ENV, raising=False) monkeypatch.setattr( "tests.integration.megatron.model_support.workflow.pinned_git_state", lambda suite_name: SimpleNamespace( @@ -457,6 +460,50 @@ def test_build_validation_report_populates_architecture_stage( assert native_vllm_lora_stage.artifact_dir == "/tmp/native-vllm-lora" +def test_build_validation_report_preserves_traces_when_sensitivity_runs( + monkeypatch, +) -> None: + seen_keep_env: list[str | None] = [] + + monkeypatch.delenv(KEEP_TOPOLOGY_ARTIFACTS_ENV, raising=False) + + monkeypatch.setattr( + "tests.integration.megatron.model_support.workflow.inspect_architecture", + lambda base_model: ArchitectureReport( + base_model=base_model, + model_key="qwen3_5_moe", + handler_key="qwen3_5_moe", + layer_families=[LayerFamilyInstance(key="standard_attention", count=1)], + recommended_min_layers=1, + ), + ) + + def _run_stage_in_subprocess( + *, + stage_name, + base_model, + architecture, + allow_unvalidated_arch=False, + ) -> ValidationStageResult: + del base_model, architecture, allow_unvalidated_arch + if stage_name == "correctness_sensitivity": + seen_keep_env.append(os.environ.get(KEEP_TOPOLOGY_ARTIFACTS_ENV)) + return ValidationStageResult(name=stage_name, passed=True, metrics={}) + + monkeypatch.setattr( + "tests.integration.megatron.model_support.workflow._run_stage_in_subprocess", + _run_stage_in_subprocess, + ) + + build_validation_report( + base_model="Qwen/Qwen3.5-35B-A3B", + include_sensitivity=True, + ) + + assert seen_keep_env == ["1"] + assert os.environ.get(KEEP_TOPOLOGY_ARTIFACTS_ENV) is None + + def test_build_validation_report_only_stage_skips_other_stages(monkeypatch) -> None: calls: list[str] = [] monkeypatch.setattr( @@ -781,6 +828,9 @@ def test_run_correctness_sensitivity_stage_runs_dense_models(monkeypatch) -> Non assert result.metrics["correctness_variant_count"] == 1 assert result.metrics["correctness_excluded_topologies"] == [] assert result.metrics["sensitivity_mutations"] == ["skip_finalize"] + assert result.metrics["default_excluded_sensitivity_mutations"] == [ + "attn_skip_flash_lse_normalize" + ] assert case_configs[0].is_moe is False @@ -804,7 +854,10 @@ def test_run_yes_no_trainability_stage(monkeypatch) -> None: "saturated_step": 2, }, ) - ) + ), + yes_no_trainability_passed=lambda report: ( + report.final_eval_reward >= report.reward_threshold + ), ), ) @@ -1126,6 +1179,9 @@ def test_run_correctness_sensitivity_stage_summarizes_reports(monkeypatch) -> No assert stage.metrics["is_moe"] is True assert stage.metrics["objectives"] == ["sft"] assert stage.metrics["sensitivity_mutations"] == ["skip_finalize"] + assert stage.metrics["default_excluded_sensitivity_mutations"] == [ + "attn_skip_flash_lse_normalize" + ] assert stage.metrics["available_gpu_count"] == 2 assert stage.metrics["required_gpu_count"] == 1 assert stage.metrics["correctness_variant_count"] == 1 @@ -1256,6 +1312,7 @@ def test_run_correctness_sensitivity_stage_can_skip_sensitivity_only( assert stage.metrics["required_gpu_count"] == 1 assert stage.metrics["correctness_variant_count"] == 1 assert stage.metrics["sensitivity_mutations"] == [] + assert stage.metrics["default_excluded_sensitivity_mutations"] == [] assert stage.metrics["sensitivity_skipped"] is True assert stage.metrics["sensitivity_skip_reason"] == f"{SKIP_SENSITIVITY_ENV}=1" assert stage.metrics["sensitivity_variant_count"] == 0 diff --git a/tests/integration/megatron/model_support/workflow.py b/tests/integration/megatron/model_support/workflow.py index 0550d92e3..99723acc7 100644 --- a/tests/integration/megatron/model_support/workflow.py +++ b/tests/integration/megatron/model_support/workflow.py @@ -37,7 +37,10 @@ LIVE_TRAINING_LOG_PATH = LOCAL_LOG_DIR / "live_training.log" ORACLE_LIVE_TRAINING_LOG_ENV = "ART_ORACLE_LIVE_TRAINING_LOG" SKIP_SENSITIVITY_ENV = "ART_MODEL_SUPPORT_SKIP_SENSITIVITY" +INCLUDE_FLASH_SENSITIVITY_ENV = "ART_MODEL_SUPPORT_INCLUDE_FLASH_SENSITIVITY" +KEEP_TOPOLOGY_ARTIFACTS_ENV = "ART_ORACLE_KEEP_TOPOLOGY_ARTIFACTS" WORKFLOW_ARTIFACT_SUITE_NAME = "Megatron model-support validation workflow" +FLASH_SENSITIVITY_MUTATION = "attn_skip_flash_lse_normalize" MANDATORY_VALIDATION_STAGES = ( "dependency_resolution", @@ -515,6 +518,7 @@ def run_correctness_sensitivity_stage( if topology.world_size() > max_world_size ] mutations: list[str] = [] + default_excluded_sensitivity_mutations: list[str] = [] excluded_sensitivity_mutations: list[str] = [] if not skip_sensitivity: for objective in objectives: @@ -543,10 +547,16 @@ def run_correctness_sensitivity_stage( > 1 ) ] + if not _truthy_env(INCLUDE_FLASH_SENSITIVITY_ENV): + default_excluded_sensitivity_mutations.append(FLASH_SENSITIVITY_MUTATION) mutations = [ mutation for mutation in mutations - if mutation not in excluded_sensitivity_mutations + if mutation + not in { + *excluded_sensitivity_mutations, + *default_excluded_sensitivity_mutations, + } ] LIVE_TRAINING_LOG_PATH.parent.mkdir(parents=True, exist_ok=True) LIVE_TRAINING_LOG_PATH.write_text("", encoding="utf-8") @@ -601,6 +611,9 @@ def run_correctness_sensitivity_stage( "objectives": objectives, "sensitivity_mutations": mutations, "excluded_sensitivity_mutations": excluded_sensitivity_mutations, + "default_excluded_sensitivity_mutations": ( + default_excluded_sensitivity_mutations + ), "available_gpu_count": available_gpu_count, "max_world_size": max_world_size, "required_gpu_count": oracle_world_size, @@ -738,14 +751,7 @@ def run_yes_no_trainability_stage( base_model=base_model, allow_unvalidated_arch=allow_unvalidated_arch, ) - passed = ( - report.saturated_step is not None - and report.saturated_step > 0 - and report.initial_eval_reward < report.reward_threshold - and report.final_eval_reward is not None - and report.final_eval_reward >= report.reward_threshold - and report.final_eval_reward > report.initial_eval_reward - ) + passed = yes_no_trainability.yes_no_trainability_passed(report) return ValidationStageResult( name=YES_NO_TRAINABILITY_STAGE, passed=passed, @@ -883,11 +889,11 @@ def build_validation_report( YES_NO_TRAINABILITY_STAGE: run_yes_no_trainability_stage, NATIVE_VLLM_LORA_STAGE: run_native_vllm_lora_stage, } - env = ( - {SKIP_SENSITIVITY_ENV: "0" if include_sensitivity else "1"} - if include_sensitivity is not None - else {} - ) + env = {} + if include_sensitivity is not None: + env[SKIP_SENSITIVITY_ENV] = "0" if include_sensitivity else "1" + if include_sensitivity: + env[KEEP_TOPOLOGY_ARTIFACTS_ENV] = "1" skip_stages = skip_stages or set() architecture: ArchitectureReport | None = None context = _temporary_env(**env) if env else nullcontext() diff --git a/tests/integration/megatron/train_inf_mismatch/output_parity.py b/tests/integration/megatron/train_inf_mismatch/output_parity.py index 839404894..34bb4413d 100644 --- a/tests/integration/megatron/train_inf_mismatch/output_parity.py +++ b/tests/integration/megatron/train_inf_mismatch/output_parity.py @@ -11,6 +11,8 @@ from pydantic import BaseModel, ConfigDict, Field, model_validator +from art.megatron.prefix_tree import parse_prefix_tree_row + from ..model_support.workflow_resources import ( handler_workflow_resources_for_base_model, resolve_stage_resources_for_current_host, @@ -139,9 +141,8 @@ class LogicalPrompt(BaseModel): sample_id: int family_id: int completion_id: int - # Packed prompt rows are the shared prefix segment exactly: prompt_end-start. - # ART stores the final context token at the start of each completion segment, - # so vLLM's generated-token logprobs start one token after this boundary. + # ART stores the final context token at the start of each leaf segment, so + # vLLM's generated-token logprobs start one token after the ancestor path. packed_prompt_length: int scored_token_start_index: int token_ids: list[int] @@ -433,6 +434,17 @@ def config_from_env() -> TrainInfOutputParityConfig: config.topology = config.topology.model_copy(update=updates) if raw_targets := os.environ.get("ART_TRAIN_INF_MISMATCH_LORA_TARGET_MODULES"): config.lora_target_modules = _parse_str_list(raw_targets) + if raw_vllm_memory := os.environ.get( + "ART_TRAIN_INF_MISMATCH_VLLM_GPU_MEMORY_UTILIZATION" + ): + config.engine_args["gpu_memory_utilization"] = float(raw_vllm_memory) + if raw_gdn_backend := os.environ.get("ART_TRAIN_INF_MISMATCH_GDN_PREFILL_BACKEND"): + raw_additional_config = config.engine_args.get("additional_config") + additional_config: dict[str, Any] = {} + if isinstance(raw_additional_config, dict): + additional_config.update(cast(dict[str, Any], raw_additional_config)) + additional_config["gdn_prefill_backend"] = raw_gdn_backend + config.engine_args["additional_config"] = additional_config raw_url = os.environ.get(train_inf_external_url_env) if raw_url is None and stage_resources is not None: raw_url = os.environ.get(model_support_external_url_env) @@ -457,40 +469,36 @@ def config_from_env() -> TrainInfOutputParityConfig: return config -def _prompt_family_segments( +def _prefix_tree_leaf_paths( group_ids: Any, parent_ids: Any, *, - required_completion_count: int = 1, -) -> list[tuple[tuple[int, int], list[tuple[int, int]]]]: - valid_tokens = int((group_ids != -1).sum().item()) - families: list[tuple[tuple[int, int], list[tuple[int, int]]]] = [] - cursor = 0 - while cursor < valid_tokens: - group_id = int(group_ids[cursor].item()) - parent_id = int(parent_ids[cursor].item()) - prompt_start = cursor - while cursor < valid_tokens and int(group_ids[cursor].item()) == group_id: - cursor += 1 - prompt_end = cursor - if group_id != parent_id: + required_leaf_count: int = 1, +) -> list[tuple[int, tuple[tuple[int, int], ...], tuple[int, int]]]: + tree = parse_prefix_tree_row(group_ids=group_ids, parent_ids=parent_ids) + segment_by_group = {segment.group_id: segment for segment in tree.segments} + child_count_by_group: dict[int, int] = {} + for segment in tree.segments: + if segment.group_id == segment.parent_id: continue - completions: list[tuple[int, int]] = [] - while cursor < valid_tokens: - completion_group_id = int(group_ids[cursor].item()) - completion_parent_id = int(parent_ids[cursor].item()) - if completion_parent_id != group_id or completion_group_id == group_id: - break - completion_start = cursor - while ( - cursor < valid_tokens - and int(group_ids[cursor].item()) == completion_group_id - ): - cursor += 1 - completions.append((completion_start, cursor)) - if len(completions) >= required_completion_count: - families.append(((prompt_start, prompt_end), completions)) - return families + child_count_by_group[segment.parent_id] = ( + child_count_by_group.get(segment.parent_id, 0) + 1 + ) + paths = [ + ( + leaf.family_index, + tuple( + (segment_by_group[group_id].start, segment_by_group[group_id].end) + for group_id in leaf.ancestors + ), + (leaf.start, leaf.end), + ) + for leaf in tree.segments + if leaf.group_id != leaf.parent_id + and child_count_by_group.get(leaf.group_id, 0) == 0 + and leaf.end - leaf.start >= 2 + ] + return paths if len(paths) >= required_leaf_count else [] def build_logical_token_map(packed_tensors: dict[str, Any]) -> LogicalTokenMap: @@ -513,64 +521,61 @@ def scored_token(sample_id: int, packed_i: int) -> bool: return True for sample_id in range(int(tokens.shape[0])): - families = _prompt_family_segments(group_ids[sample_id], parent_ids[sample_id]) - for family_id, (prompt_segment, completion_segments) in enumerate(families): - prompt_start, prompt_end = prompt_segment - prompt_len = prompt_end - prompt_start - for completion_id, (completion_start, completion_end) in enumerate( - completion_segments - ): - last_scored_i = None - for packed_i in range(completion_start + 1, completion_end): - if scored_token(sample_id, packed_i): - last_scored_i = packed_i - if last_scored_i is None: - continue - effective_completion_end = last_scored_i + 1 - flat = [ - int(value) - for value in tokens[sample_id, prompt_start:prompt_end].tolist() - ] + [ - int(value) - for value in tokens[ - sample_id, completion_start:effective_completion_end - ].tolist() - ] - flat_key = tuple(flat) - prompt_id = prompt_id_by_tokens.get(flat_key) - if prompt_id is None: - prompt_id = len(prompts) - prompt_id_by_tokens[flat_key] = prompt_id - prompts.append( - LogicalPrompt( - prompt_id=prompt_id, - sample_id=sample_id, - family_id=family_id, - completion_id=completion_id, - packed_prompt_length=prompt_len, - scored_token_start_index=prompt_len + 1, - token_ids=flat, - ) + leaf_paths = _prefix_tree_leaf_paths( + group_ids[sample_id], parent_ids[sample_id] + ) + for completion_id, (family_id, ancestor_segments, leaf_segment) in enumerate( + leaf_paths + ): + leaf_start, leaf_end = leaf_segment + last_scored_i = None + for packed_i in range(leaf_start + 1, leaf_end): + if scored_token(sample_id, packed_i): + last_scored_i = packed_i + if last_scored_i is None: + continue + effective_leaf_end = last_scored_i + 1 + prompt_len = sum(end - start for start, end in ancestor_segments) + reference_segments = (*ancestor_segments, (leaf_start, effective_leaf_end)) + flat = [ + int(value) + for start, end in reference_segments + for value in tokens[sample_id, start:end].tolist() + ] + flat_key = tuple(flat) + prompt_id = prompt_id_by_tokens.get(flat_key) + if prompt_id is None: + prompt_id = len(prompts) + prompt_id_by_tokens[flat_key] = prompt_id + prompts.append( + LogicalPrompt( + prompt_id=prompt_id, + sample_id=sample_id, + family_id=family_id, + completion_id=completion_id, + packed_prompt_length=prompt_len, + scored_token_start_index=prompt_len + 1, + token_ids=flat, ) - for packed_i in range(completion_start + 1, effective_completion_end): - if not scored_token(sample_id, packed_i): - continue - logical_tokens.append( - LogicalToken( - token_id=int(tokens[sample_id, packed_i].item()), - sample_id=sample_id, - family_id=family_id, - completion_id=completion_id, - prompt_id=prompt_id, - art_packed_token_index=packed_i, - art_logit_index=packed_i - 1, - vllm_prompt_token_index=prompt_len - + (packed_i - completion_start), - ) + ) + for packed_i in range(leaf_start + 1, effective_leaf_end): + if not scored_token(sample_id, packed_i): + continue + logical_tokens.append( + LogicalToken( + token_id=int(tokens[sample_id, packed_i].item()), + sample_id=sample_id, + family_id=family_id, + completion_id=completion_id, + prompt_id=prompt_id, + art_packed_token_index=packed_i, + art_logit_index=packed_i - 1, + vllm_prompt_token_index=prompt_len + (packed_i - leaf_start), ) + ) if not prompts or not logical_tokens: - raise RuntimeError("Shared-prefix probe produced no comparable logical tokens") + raise RuntimeError("Prefix-tree probe produced no comparable logical tokens") return LogicalTokenMap(prompts=prompts, tokens=logical_tokens) @@ -988,7 +993,7 @@ def _run_logits( ) -> Any: import torch - from art.megatron.shared_prefix_state import create_shared_prefix_state + from art.megatron.prefix_tree_state import create_prefix_tree_state from art.megatron.training.trace import ( packed_sequence_token_uids, prepare_replay_local_input_token_uids, @@ -1000,7 +1005,7 @@ def _run_logits( position_ids = packed_tensors["input_pos"].to(device=device) group_ids = packed_tensors["group_ids"].to(device=device) parent_ids = packed_tensors["parent_ids"].to(device=device) - attention_state = create_shared_prefix_state( + attention_state = create_prefix_tree_state( group_ids=group_ids, parent_ids=parent_ids, input_pos=position_ids, diff --git a/tests/integration/megatron/train_inf_mismatch/real_path.py b/tests/integration/megatron/train_inf_mismatch/real_path.py index 18f5ebb05..cd996cdb0 100644 --- a/tests/integration/megatron/train_inf_mismatch/real_path.py +++ b/tests/integration/megatron/train_inf_mismatch/real_path.py @@ -64,7 +64,7 @@ class RealPathConfig(BaseModel): output_parity: TrainInfOutputParityConfig = Field( default_factory=TrainInfOutputParityConfig ) - prompt_count: int = 2 + prompt_count: int = 4 rollouts_per_prompt: int = 2 max_completion_tokens: int = 16 prompt_sentence_count: int = 28 @@ -101,10 +101,10 @@ class RealPathBaseDiagnosticBundle(BaseModel): logical_prompt_count: int logical_token_count: int moe_routing_packed_tokens: int - moe_routing_shared_prefix_rows: int - moe_routing_shared_prefix_conflict_rows: int - moe_routing_shared_prefix_conflict_slots: int - moe_routing_shared_prefix_compared_slots: int + moe_routing_prefix_tree_rows: int + moe_routing_prefix_tree_conflict_rows: int + moe_routing_prefix_tree_conflict_slots: int + moe_routing_prefix_tree_compared_slots: int vllm_forward_trace_dir: str | None = None megatron_forward_trace_dir: str | None = None @@ -134,8 +134,8 @@ class RealPathTrainInfReport(BaseModel): base_logical_prompt_count: int | None = None base_logical_token_count: int | None = None base_moe_routing_packed_tokens: int | None = None - base_moe_routing_shared_prefix_conflict_rows: int | None = None - base_moe_routing_shared_prefix_conflict_slots: int | None = None + base_moe_routing_prefix_tree_conflict_rows: int | None = None + base_moe_routing_prefix_tree_conflict_slots: int | None = None adapter_path: str adapter_cache_key: str adapter_cache_hit: bool @@ -148,10 +148,12 @@ class RealPathTrainInfReport(BaseModel): lora: PairComparison lora_topk: TopKComparison moe_routing_packed_tokens: int - moe_routing_shared_prefix_rows: int - moe_routing_shared_prefix_conflict_rows: int - moe_routing_shared_prefix_conflict_slots: int - moe_routing_shared_prefix_compared_slots: int + moe_routing_prefix_tree_rows: int + moe_routing_prefix_tree_conflict_rows: int + moe_routing_prefix_tree_conflict_slots: int + moe_routing_prefix_tree_compared_slots: int + prompt_tree_depth: int = 0 + prompt_tree_branch_count: int = 0 mean_abs_pct_limit: float top20_kl_candidate_to_target_limit: float passed: bool @@ -175,13 +177,13 @@ def _real_path_rollout_weights_mode( _PROMPT_SENTENCES = [ "A careful systems engineer checks assumptions before changing thresholds.", - "The training batch contains shared prefixes and divergent completions.", + "The training batch contains prefix treees and divergent completions.", "Numerical parity should be measured on the exact tokens used by the policy.", "Sparse expert routing can create discontinuous output differences.", "A reproducible test writes enough artifacts to explain every comparison.", "LoRA adapters must be active and nonzero during both inference and training.", "The prompt should be realistic enough to exercise ordinary tokenizer paths.", - "Packed Megatron inputs use shared prefixes while vLLM receives flat requests.", + "Packed Megatron inputs use prefix treees while vLLM receives flat requests.", "If tokenization diverges, the mismatch should fail as early as possible.", "The report includes target logprobs and top token overlap for diagnosis.", "Routing replay should use vLLM expert ids captured from real rollouts.", @@ -200,6 +202,20 @@ def _real_path_rollout_weights_mode( "Validation code belongs in tests unless production needs the behavior.", ] _PROMPT_TOKENS_PER_SENTENCE_ESTIMATE = 12 +_PROMPT_TREE_ROOT = ( + "Write a concise continuation for a validation note. Preserve the technical " + "tone and keep the answer concrete." +) +_PROMPT_TREE_MIDS = ( + "Branch alpha: emphasize runtime validation and packed training inputs.", + "Branch beta: emphasize serving parity and reproducible comparison artifacts.", +) +_PROMPT_TREE_LEAVES = ( + "Case one: describe a route where a prefix tree splits into two completions.", + "Case two: describe a route where the same branch has a longer continuation.", + "Case three: describe a route where another branch reaches a different expert.", + "Case four: describe a direct leaf that skips the intermediate branch.", +) def config_from_env() -> RealPathConfig: @@ -276,10 +292,23 @@ def _apply_sliding_window_prompt_defaults(config: RealPathConfig) -> None: def _build_prompt_from_sentences(index: int, sentences: list[str]) -> str: - return ( - "Write a concise continuation for probe " - f"{index}. Preserve the technical tone.\n\n" + " ".join(sentences) + mid = _PROMPT_TREE_MIDS[(index // 2) % len(_PROMPT_TREE_MIDS)] + leaf = _PROMPT_TREE_LEAVES[index % len(_PROMPT_TREE_LEAVES)] + return f"{_PROMPT_TREE_ROOT}\n\n{mid}\n\n{leaf}\n\nNotes: " + " ".join(sentences) + + +def _prompt_tree_shape(prompts: list[str]) -> tuple[int, int]: + mid_count = len( + {mid for mid in _PROMPT_TREE_MIDS if any(mid in prompt for prompt in prompts)} ) + leaf_count = len( + { + leaf + for leaf in _PROMPT_TREE_LEAVES + if any(leaf in prompt for prompt in prompts) + } + ) + return (3 if mid_count and leaf_count else 1, mid_count + leaf_count) def _build_prompts(config: RealPathConfig) -> list[str]: @@ -756,14 +785,10 @@ async def _score_base_real_generation_path( logical_prompt_count=len(logical_map.prompts), logical_token_count=len(logical_map.tokens), moe_routing_packed_tokens=int(stats.packed_tokens), - moe_routing_shared_prefix_rows=int(stats.shared_prefix_rows), - moe_routing_shared_prefix_conflict_rows=int(stats.shared_prefix_conflict_rows), - moe_routing_shared_prefix_conflict_slots=int( - stats.shared_prefix_conflict_slots - ), - moe_routing_shared_prefix_compared_slots=int( - stats.shared_prefix_compared_slots - ), + moe_routing_prefix_tree_rows=int(stats.prefix_tree_rows), + moe_routing_prefix_tree_conflict_rows=int(stats.prefix_tree_conflict_rows), + moe_routing_prefix_tree_conflict_slots=int(stats.prefix_tree_conflict_slots), + moe_routing_prefix_tree_compared_slots=int(stats.prefix_tree_compared_slots), vllm_forward_trace_dir=( str(vllm_forward_trace_dir) if vllm_forward_trace_dir is not None else None ), @@ -1504,6 +1529,9 @@ async def run_real_path_train_inf_mismatch( parity_config.base_model, allow_unvalidated_arch=parity_config.allow_unvalidated_arch, ) + prompt_tree_depth, prompt_tree_branch_count = _prompt_tree_shape( + _build_prompts(config) + ) passed = ( comparison.mean_abs_pct <= mean_abs_pct_limit and topk_comparison.top20_intersection_kl_candidate_to_target @@ -1529,13 +1557,13 @@ async def run_real_path_train_inf_mismatch( if base_diagnostic is not None else None ), - base_moe_routing_shared_prefix_conflict_rows=( - base_diagnostic.moe_routing_shared_prefix_conflict_rows + base_moe_routing_prefix_tree_conflict_rows=( + base_diagnostic.moe_routing_prefix_tree_conflict_rows if base_diagnostic is not None else None ), - base_moe_routing_shared_prefix_conflict_slots=( - base_diagnostic.moe_routing_shared_prefix_conflict_slots + base_moe_routing_prefix_tree_conflict_slots=( + base_diagnostic.moe_routing_prefix_tree_conflict_slots if base_diagnostic is not None else None ), @@ -1557,16 +1585,16 @@ async def run_real_path_train_inf_mismatch( lora=comparison, lora_topk=topk_comparison, moe_routing_packed_tokens=int(stats.packed_tokens), - moe_routing_shared_prefix_rows=int(stats.shared_prefix_rows), - moe_routing_shared_prefix_conflict_rows=int( - stats.shared_prefix_conflict_rows - ), - moe_routing_shared_prefix_conflict_slots=int( - stats.shared_prefix_conflict_slots + moe_routing_prefix_tree_rows=int(stats.prefix_tree_rows), + moe_routing_prefix_tree_conflict_rows=int(stats.prefix_tree_conflict_rows), + moe_routing_prefix_tree_conflict_slots=int( + stats.prefix_tree_conflict_slots ), - moe_routing_shared_prefix_compared_slots=int( - stats.shared_prefix_compared_slots + moe_routing_prefix_tree_compared_slots=int( + stats.prefix_tree_compared_slots ), + prompt_tree_depth=prompt_tree_depth, + prompt_tree_branch_count=prompt_tree_branch_count, mean_abs_pct_limit=mean_abs_pct_limit, top20_kl_candidate_to_target_limit=top20_kl_limit, passed=passed, diff --git a/tests/integration/megatron/train_inf_mismatch/test_output_parity_invariants.py b/tests/integration/megatron/train_inf_mismatch/test_output_parity_invariants.py index 39c8ede0a..09ab421aa 100644 --- a/tests/integration/megatron/train_inf_mismatch/test_output_parity_invariants.py +++ b/tests/integration/megatron/train_inf_mismatch/test_output_parity_invariants.py @@ -31,7 +31,7 @@ ) -def test_logical_map_flattens_shared_prefix_branches() -> None: +def test_logical_map_flattens_prefix_tree_branches() -> None: packed = { "tokens": torch.tensor([[10, 11, 12, 13, 14, 12, 15, 16]]), "group_ids": torch.tensor([[0, 0, 1, 1, 1, 2, 2, 2]]), @@ -59,6 +59,61 @@ def test_logical_map_flattens_shared_prefix_branches() -> None: ] +def test_logical_map_flattens_nested_prefix_tree_leaves() -> None: + packed = { + "tokens": torch.tensor( + [[10, 11, 20, 30, 31, 32, 33, 34, 35, 40, 50, 51, 52, 60, 61, 62]] + ), + "group_ids": torch.tensor([[0, 0, 1, 2, 2, 2, 3, 3, 3, 4, 5, 5, 5, 6, 6, 6]]), + "parent_ids": torch.tensor([[0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 4, 4, 4, 0, 0, 0]]), + } + + logical_map = build_logical_token_map(packed) + + assert [prompt.token_ids for prompt in logical_map.prompts] == [ + [10, 11, 20, 30, 31, 32], + [10, 11, 20, 33, 34, 35], + [10, 11, 40, 50, 51, 52], + [10, 11, 60, 61, 62], + ] + assert [prompt.packed_prompt_length for prompt in logical_map.prompts] == [ + 3, + 3, + 3, + 2, + ] + assert [token.token_id for token in logical_map.tokens] == [ + 31, + 32, + 34, + 35, + 51, + 52, + 61, + 62, + ] + assert [token.art_logit_index for token in logical_map.tokens] == [ + 3, + 4, + 6, + 7, + 10, + 11, + 13, + 14, + ] + assert [token.vllm_prompt_token_index for token in logical_map.tokens] == [ + 4, + 5, + 4, + 5, + 4, + 5, + 3, + 4, + ] + + def test_aggregate_mean_abs_pct_uses_vllm_merge_formula() -> None: summary = aggregate_mean_abs_pct( candidate=torch.tensor([2.0, 4.0]), @@ -260,6 +315,26 @@ def test_config_from_env_accepts_lora_target_module_override( assert config.lora_target_modules == ["experts", "in_proj_qkv", "in_proj_z"] +def test_config_from_env_accepts_vllm_memory_utilization_override( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("ART_TRAIN_INF_MISMATCH_VLLM_GPU_MEMORY_UTILIZATION", "0.5") + + config = config_from_env() + + assert config.engine_args["gpu_memory_utilization"] == 0.5 + + +def test_config_from_env_accepts_gdn_prefill_backend_override( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("ART_TRAIN_INF_MISMATCH_GDN_PREFILL_BACKEND", "triton") + + config = config_from_env() + + assert config.engine_args["additional_config"] == {"gdn_prefill_backend": "triton"} + + def test_default_rollout_modes_follow_model_support_native_lora_status() -> None: assert TrainInfOutputParityConfig( base_model="Qwen/Qwen3.5-35B-A3B" @@ -317,3 +392,4 @@ def fake_run(*args, **kwargs): assert captured_env["ART_RUN_TRAIN_INF_MISMATCH_LIVE"] == "1" assert captured_env["ART_TRAIN_INF_MISMATCH_ALLOW_UNVALIDATED_ARCH"] == "1" assert captured_env["ART_REAL_PATH_MAX_COMPLETION_TOKENS"] == "16" + assert captured_env["ART_TRAIN_INF_MISMATCH_VLLM_GPU_MEMORY_UTILIZATION"] == "0.50" diff --git a/tests/integration/megatron/train_inf_mismatch/workflow_stage.py b/tests/integration/megatron/train_inf_mismatch/workflow_stage.py index ae0a7cef2..12e449887 100644 --- a/tests/integration/megatron/train_inf_mismatch/workflow_stage.py +++ b/tests/integration/megatron/train_inf_mismatch/workflow_stage.py @@ -76,6 +76,7 @@ def run_train_inf_mismatch( "1" if allow_unvalidated_arch else "0" ) env["ART_REAL_PATH_MAX_COMPLETION_TOKENS"] = "16" + env.setdefault("ART_TRAIN_INF_MISMATCH_VLLM_GPU_MEMORY_UTILIZATION", "0.50") existing_pythonpath = env.get("PYTHONPATH") tests_dir = str(REPO_ROOT / "tests") env["PYTHONPATH"] = ( diff --git a/tests/integration/megatron/trainability/test_config.py b/tests/integration/megatron/trainability/test_config.py index 75e5cc37c..dae5bda21 100644 --- a/tests/integration/megatron/trainability/test_config.py +++ b/tests/integration/megatron/trainability/test_config.py @@ -9,8 +9,21 @@ import art -from .test_live_length_trainability import _use_default_moe_dedicated_placement +from .test_live_length_trainability import ( + LengthSampleReport, + LengthTrainabilityReport, + _default_learning_rate, + _length_trainability_thresholds, + _prompt_for_index, + _use_default_moe_dedicated_placement, + length_trainability_passed, +) +from .test_live_length_trainability import ( + _prompt_tree_shape as _length_prompt_tree_shape, +) from .yes_no_trainability import ( + TrainabilityStepReport, + YesNoTrainabilityReport, _build_internal_config, _build_variant, _default_variant_name, @@ -21,6 +34,11 @@ _variant_packed_sequence_length, _variant_rollouts_per_prompt, _variant_train_kwargs, + build_prompts, + yes_no_trainability_passed, +) +from .yes_no_trainability import ( + _prompt_tree_shape as _yes_no_prompt_tree_shape, ) @@ -159,6 +177,153 @@ def test_qwen3_5_defaults_to_shared_lora_rollout() -> None: assert "inference_gpu_ids" not in config +def test_dense_yes_no_default_uses_dedicated_placement(monkeypatch) -> None: + monkeypatch.delenv("ART_MODEL_SUPPORT_YES_NO_VARIANT", raising=False) + + assert _default_variant_name("Qwen/Qwen3-32B") == "megatron_dedicated" + + +def test_yes_no_default_variant_env_override(monkeypatch) -> None: + monkeypatch.setenv("ART_MODEL_SUPPORT_YES_NO_VARIANT", "megatron_shared") + + assert _default_variant_name("Qwen/Qwen3-32B") == "megatron_shared" + + +def test_yes_no_trainability_passes_initially_saturated_stable_report() -> None: + report = YesNoTrainabilityReport( + variant="megatron_shared", + backend_name="megatron", + placement_mode="shared", + base_model="google/gemma-4-31B-it", + output_dir="/tmp/report", + trainer_gpu_ids=[0, 1], + inference_gpu_ids=[0, 1], + rollout_weights_mode="lora", + reward_threshold=0.9, + max_steps=4, + prompt_count=8, + eval_prompt_count=8, + rollouts_per_prompt=4, + latest_step=1, + initial_eval_reward=0.9375, + final_eval_reward=0.9375, + saturated_step=1, + step0_name="model@0", + latest_name="model@1", + steps=[ + TrainabilityStepReport( + step=1, + eval_reward=0.9375, + train_reward=0.875, + train_metrics={"grad_norm": 54.0}, + ) + ], + ) + + assert yes_no_trainability_passed(report) is True + + +def test_yes_no_prompts_form_prefix_tree_by_default(monkeypatch) -> None: + monkeypatch.delenv("ART_MODEL_SUPPORT_YES_NO_PROMPT", raising=False) + monkeypatch.setenv("ART_MODEL_SUPPORT_YES_NO_PROMPT_COUNT", "8") + + prompts = build_prompts() + + assert _yes_no_prompt_tree_shape(prompts) == (3, 6) + + +def test_qwen3_5_length_trainability_uses_stable_learning_rate() -> None: + assert _default_learning_rate("Qwen/Qwen3.5-35B-A3B") == 7e-5 + assert _default_learning_rate("Qwen/Qwen3-30B-A3B-Instruct-2507") == 1e-4 + + +def test_length_prompts_form_prefix_tree_by_default() -> None: + prompts = [_prompt_for_index(index)[0] for index in range(4)] + + assert _length_prompt_tree_shape(prompts) == (3, 6) + + +def test_length_trainability_accepts_near_baseline_learning_signal() -> None: + report = LengthTrainabilityReport( + base_model="google/gemma-4-31B-it", + max_steps=10, + max_steps_off_policy=0, + latest_step=3, + variant_name="megatron_dedicated", + trainer_gpu_ids=[0], + inference_gpu_ids=[1], + training_topology={"tp": 1, "cp": 1, "ep": 1, "etp": 1, "dp": 1, "sp": False}, + rollout_weights_mode="lora", + rollouts_per_prompt=4, + normalize_advantages=True, + summary_log_path="/tmp/length_trainability.log", + latest_summary_log_path="/tmp/latest_length_trainability.log", + thresholds=_length_trainability_thresholds("google/gemma-4-31B-it"), + initial_train_abs_error=3.875, + best_train_abs_error=0.5, + success_step=3, + final_train_reward=-0.05, + final_train_abs_error=0.5, + model_ids_after=["length@0", "length@3"], + samples=[ + LengthSampleReport( + split="train", + step=0, + scenario_index=0, + target_step=0, + target_tokens=10, + max_tokens=142, + prompt_word_count=300, + generated_tokens=14, + abs_error=4, + reward=-0.4, + text="a short answer", + ), + LengthSampleReport( + split="train", + step=0, + scenario_index=1, + target_step=0, + target_tokens=10, + max_tokens=142, + prompt_word_count=300, + generated_tokens=6, + abs_error=4, + reward=-0.4, + text="brief", + ), + LengthSampleReport( + split="train", + step=3, + scenario_index=2, + target_step=3, + target_tokens=10, + max_tokens=142, + prompt_word_count=300, + generated_tokens=10, + abs_error=0, + reward=0.0, + text="a target length answer", + ), + LengthSampleReport( + split="train", + step=3, + scenario_index=3, + target_step=3, + target_tokens=10, + max_tokens=142, + prompt_word_count=300, + generated_tokens=11, + abs_error=1, + reward=-0.1, + text="a slightly long answer", + ), + ], + ) + + assert length_trainability_passed(report) is True + + def test_validated_dense_model_uses_dense_shared_topology( monkeypatch, ) -> None: diff --git a/tests/integration/megatron/trainability/test_live_length_trainability.py b/tests/integration/megatron/trainability/test_live_length_trainability.py index 2921c2667..2a2221c8a 100644 --- a/tests/integration/megatron/trainability/test_live_length_trainability.py +++ b/tests/integration/megatron/trainability/test_live_length_trainability.py @@ -37,6 +37,8 @@ torch = pytest.importorskip("torch") DEFAULT_BASE_MODEL = "Qwen/Qwen3.5-35B-A3B" +DEFAULT_LENGTH_LEARNING_RATE = 1e-4 +LARGE_MOE_LENGTH_LEARNING_RATE = 7e-5 LIVE_ENV = "ART_RUN_LIVE_LENGTH_TRAINABILITY" TRAINER_GPU_IDS_ENV = "ART_MODEL_SUPPORT_TRAINER_GPU_IDS" INFERENCE_GPU_IDS_ENV = "ART_MODEL_SUPPORT_INFERENCE_GPU_IDS" @@ -65,6 +67,16 @@ "only as background texture. Use one sentence. Do not use bullets, numbering, " "code, or a preface." ) +LENGTH_PROMPT_MIDS = ( + "Branch alpha: the harbor office is preparing a routine status note.", + "Branch beta: the harbor office is summarizing a quiet maintenance record.", +) +LENGTH_PROMPT_LEAVES = ( + "Case one: mention calm water without adding drama.", + "Case two: mention ordinary work near the pier.", + "Case three: mention a simple observation from the office.", + "Case four: mention a reserved conclusion about the day.", +) FILLER_SENTENCES = ( "The morning ledger mentioned a bicycle bell near the old customs window.", "A folded receipt waited beside three dull pencils and a chipped mug.", @@ -159,6 +171,8 @@ class LengthTrainabilityReport(BaseModel): training_topology: dict[str, int | bool] rollout_weights_mode: str rollouts_per_prompt: int + prompt_tree_depth: int = 0 + prompt_tree_branch_count: int = 0 normalize_advantages: bool summary_log_path: str latest_summary_log_path: str @@ -208,6 +222,12 @@ def _target_tokens() -> int: return _get_env_int("ART_MODEL_SUPPORT_LENGTH_TARGET_TOKENS", 10) +def _default_learning_rate(base_model: str) -> float: + if base_model == DEFAULT_BASE_MODEL: + return LARGE_MOE_LENGTH_LEARNING_RATE + return DEFAULT_LENGTH_LEARNING_RATE + + def _use_default_moe_dedicated_placement(variant: Any, *, base_model: str) -> None: if not model_uses_expert_parallel(base_model, allow_unvalidated_arch=True): return @@ -298,16 +318,33 @@ def _prompt_for_index(index: int) -> tuple[str, int]: sentences = list(FILLER_SENTENCES) rng.shuffle(sentences) selected: list[str] = [] - prompt = BASE_PROMPT + mid = LENGTH_PROMPT_MIDS[(index // 2) % len(LENGTH_PROMPT_MIDS)] + leaf = LENGTH_PROMPT_LEAVES[index % len(LENGTH_PROMPT_LEAVES)] + prefix = f"{BASE_PROMPT}\n\n{mid}\n\n{leaf}" + prompt = prefix for sentence in sentences: if _word_count(prompt) >= target_words: break selected.append(sentence) - prompt = f"{BASE_PROMPT}\n\nNotes: {' '.join(selected)}" + prompt = f"{prefix}\n\nNotes: {' '.join(selected)}" _check_prompt_hides_target(prompt) return prompt, _word_count(prompt) +def _prompt_tree_shape(prompts: list[str]) -> tuple[int, int]: + mid_count = len( + {mid for mid in LENGTH_PROMPT_MIDS if any(mid in prompt for prompt in prompts)} + ) + leaf_count = len( + { + leaf + for leaf in LENGTH_PROMPT_LEAVES + if any(leaf in prompt for prompt in prompts) + } + ) + return (3 if mid_count and leaf_count else 1, mid_count + leaf_count) + + def _scenario( index: int, *, @@ -772,7 +809,7 @@ async def rollout_fn( max_steps_off_policy=max_steps_off_policy, learning_rate=_get_env_float( "ART_MODEL_SUPPORT_LENGTH_LEARNING_RATE", - 1e-4, + _default_learning_rate(base_model), ), loss_fn="cispo", normalize_advantages=normalize_advantages, @@ -803,7 +840,7 @@ async def rollout_fn( success_step = next( ( step - for step, abs_error in train_abs_error_by_step.items() + for step, abs_error in sorted(train_abs_error_by_step.items()) if _success_abs_error_passed(abs_error, thresholds) ), None, @@ -819,6 +856,13 @@ async def rollout_fn( if final_train_samples else None ) + prompt_tree_sample_count = 4 if scenario_limit is None else min(4, scenario_limit) + prompt_tree_depth, prompt_tree_branch_count = _prompt_tree_shape( + [ + _scenario(index, base_model=base_model).prompt + for index in range(prompt_tree_sample_count) + ] + ) topology = cast(Topology, variant.topology) report = LengthTrainabilityReport( base_model=base_model, @@ -831,6 +875,8 @@ async def rollout_fn( training_topology=cast(dict[str, int | bool], topology.model_dump()), rollout_weights_mode=rollout_weights_mode, rollouts_per_prompt=rollouts_per_prompt, + prompt_tree_depth=prompt_tree_depth, + prompt_tree_branch_count=prompt_tree_branch_count, normalize_advantages=normalize_advantages, summary_log_path=str(summary_log_path), latest_summary_log_path=str(LATEST_SUMMARY_LOG_PATH), diff --git a/tests/integration/megatron/trainability/test_live_yes_no_trainability.py b/tests/integration/megatron/trainability/test_live_yes_no_trainability.py index 119d3b74a..a12353752 100644 --- a/tests/integration/megatron/trainability/test_live_yes_no_trainability.py +++ b/tests/integration/megatron/trainability/test_live_yes_no_trainability.py @@ -4,7 +4,10 @@ import pytest -from .yes_no_trainability import run_yes_no_trainability_async +from .yes_no_trainability import ( + run_yes_no_trainability_async, + yes_no_trainability_passed, +) torch = pytest.importorskip("torch") @@ -29,12 +32,7 @@ def _unsloth_base_model() -> str: def _assert_passed(report) -> None: - assert report.saturated_step is not None - assert report.saturated_step > 0 - assert report.initial_eval_reward < report.reward_threshold - assert report.final_eval_reward is not None - assert report.final_eval_reward >= report.reward_threshold - assert report.final_eval_reward > report.initial_eval_reward + assert yes_no_trainability_passed(report) assert report.latest_step > 0 assert report.step0_name in report.model_ids_before assert report.latest_name in report.model_ids_after diff --git a/tests/integration/megatron/trainability/yes_no_trainability.py b/tests/integration/megatron/trainability/yes_no_trainability.py index 67aba7078..f9a213148 100644 --- a/tests/integration/megatron/trainability/yes_no_trainability.py +++ b/tests/integration/megatron/trainability/yes_no_trainability.py @@ -3,7 +3,6 @@ import asyncio from contextlib import asynccontextmanager, contextmanager, nullcontext import gc -from itertools import permutations import os from pathlib import Path import re @@ -35,6 +34,7 @@ _TRAINER_GPU_IDS_ENV = "ART_MODEL_SUPPORT_TRAINER_GPU_IDS" _INFERENCE_GPU_IDS_ENV = "ART_MODEL_SUPPORT_INFERENCE_GPU_IDS" _SHARED_GPU_IDS_ENV = "ART_MODEL_SUPPORT_SHARED_GPU_IDS" +_VARIANT_ENV = "ART_MODEL_SUPPORT_YES_NO_VARIANT" _EXTERNAL_VLLM_URL_ENV = "ART_MODEL_SUPPORT_EXTERNAL_VLLM_URL" _EXTERNAL_VLLM_API_KEY_ENV = "ART_MODEL_SUPPORT_EXTERNAL_VLLM_API_KEY" _TRAINABILITY_ROOT = ( @@ -75,6 +75,8 @@ class YesNoTrainabilityReport(BaseModel): prompt_count: int eval_prompt_count: int rollouts_per_prompt: int + prompt_tree_depth: int = 0 + prompt_tree_branch_count: int = 0 latest_step: int initial_eval_reward: float final_eval_reward: float | None = None @@ -96,28 +98,52 @@ class _TrainabilityVariant(BaseModel): inference_gpu_ids: list[int] = Field(default_factory=list) +_YES_NO_PROMPT_ROOT = ( + "Read the validation card and answer with one word from yes, no, or maybe." +) +_YES_NO_PROMPT_MIDS = ( + "Branch alpha: the card is about deployment readiness.", + "Branch beta: the card is about metric interpretation.", +) +_YES_NO_PROMPT_LEAVES = ( + "Case one: the safest answer is uncertain.", + "Case two: the report contains a contradiction.", + "Case three: the check has partial evidence.", + "Case four: the reviewer needs a cautious final word.", +) + + def build_prompts() -> list[str]: prompt = os.environ.get("ART_MODEL_SUPPORT_YES_NO_PROMPT", "").strip() prompt_count = _get_env_int("ART_MODEL_SUPPORT_YES_NO_PROMPT_COUNT", 8) if prompt: return [prompt] * max(1, prompt_count) prompts = [ - f"{prefix} exactly one of {body}" - for prefix in ("respond with", "just respond with") - for use_quotes in (True, False) - for length in (3, 2) - for words in permutations(("yes", "no", "maybe"), length) - for body in [ + "\n\n".join( ( - ", ".join(f"'{word}'" if use_quotes else word for word in words) - if length == 3 - else " or ".join(f"'{word}'" if use_quotes else word for word in words) + _YES_NO_PROMPT_ROOT, + _YES_NO_PROMPT_MIDS[(index // 2) % len(_YES_NO_PROMPT_MIDS)], + _YES_NO_PROMPT_LEAVES[index % len(_YES_NO_PROMPT_LEAVES)], + "Return only yes, no, or maybe.", ) - ] + ) + for index in range(max(1, prompt_count)) ] - if prompt_count <= len(prompts): - return prompts[: max(1, prompt_count)] - return [prompts[index % len(prompts)] for index in range(prompt_count)] + return prompts + + +def _prompt_tree_shape(prompts: list[str]) -> tuple[int, int]: + mid_count = len( + {mid for mid in _YES_NO_PROMPT_MIDS if any(mid in prompt for prompt in prompts)} + ) + leaf_count = len( + { + leaf + for leaf in _YES_NO_PROMPT_LEAVES + if any(leaf in prompt for prompt in prompts) + } + ) + return (3 if mid_count and leaf_count else 1, mid_count + leaf_count) def _slugify(value: str) -> str: @@ -574,6 +600,13 @@ def _default_variant_name( *, allow_unvalidated_arch: bool = False, ) -> _VARIANT_NAME: + if override := os.environ.get(_VARIANT_ENV, "").strip(): + if override not in {"megatron_shared", "megatron_dedicated"}: + raise ValueError( + f"Unsupported {_VARIANT_ENV}={override!r}. " + "Expected 'megatron_shared' or 'megatron_dedicated'." + ) + return cast(_VARIANT_NAME, override) workflow_resources = handler_workflow_resources_for_base_model( base_model, allow_unvalidated_arch=allow_unvalidated_arch, @@ -583,13 +616,15 @@ def _default_variant_name( and workflow_resources.yes_no_trainability_variant is not None ): return workflow_resources.yes_no_trainability_variant - if ( - _rollout_weights_mode( - base_model, - allow_unvalidated_arch=allow_unvalidated_arch, - ) - == "merged" - ): + is_moe = model_uses_expert_parallel( + base_model, + allow_unvalidated_arch=allow_unvalidated_arch, + ) + rollout_weights_mode = _rollout_weights_mode( + base_model, + allow_unvalidated_arch=allow_unvalidated_arch, + ) + if rollout_weights_mode == "merged" or not is_moe: return "megatron_dedicated" return "megatron_shared" @@ -896,6 +931,7 @@ async def run_yes_no_trainability_async( eval_prompt_count = _get_env_int("ART_MODEL_SUPPORT_YES_NO_EVAL_PROMPTS", 8) prompts = build_prompts() eval_prompts = prompts[:eval_prompt_count] + prompt_tree_depth, prompt_tree_branch_count = _prompt_tree_shape(prompts) internal_config = _build_internal_config( variant, base_model=base_model, @@ -961,6 +997,8 @@ async def run_yes_no_trainability_async( prompt_count=len(prompts), eval_prompt_count=len(eval_prompts), rollouts_per_prompt=rollouts_per_prompt, + prompt_tree_depth=prompt_tree_depth, + prompt_tree_branch_count=prompt_tree_branch_count, latest_step=0, initial_eval_reward=initial_eval_reward, step0_name=step0_name, @@ -1052,6 +1090,26 @@ def run_yes_no_trainability( ) +def yes_no_trainability_passed(report: YesNoTrainabilityReport) -> bool: + learned_from_below_threshold = ( + report.saturated_step is not None + and report.saturated_step > 0 + and report.initial_eval_reward < report.reward_threshold + and report.final_eval_reward is not None + and report.final_eval_reward >= report.reward_threshold + and report.final_eval_reward > report.initial_eval_reward + ) + already_saturated_and_stable = ( + report.initial_eval_reward >= report.reward_threshold + and report.latest_step > 0 + and report.final_eval_reward is not None + and report.final_eval_reward >= report.reward_threshold + and bool(report.steps) + and any(step.train_metrics.get("grad_norm", 0.0) > 0.0 for step in report.steps) + ) + return learned_from_below_threshold or already_saturated_and_stable + + def run_megatron_dedicated_yes_no_trainability( base_model: str, *, diff --git a/tests/integration/test_provenance.py b/tests/integration/test_provenance.py index 187fcad88..459e2297d 100644 --- a/tests/integration/test_provenance.py +++ b/tests/integration/test_provenance.py @@ -4,10 +4,10 @@ from datetime import datetime from dotenv import load_dotenv -import wandb import art from art.serverless.backend import ServerlessBackend +from art.utils import wandb_sdk load_dotenv() @@ -41,7 +41,7 @@ def get_latest_artifact_provenance( entity: str, project: str, name: str ) -> list[str] | None: """Fetch provenance from the latest W&B artifact's metadata.""" - api = wandb.Api() + api = wandb_sdk.api() artifact = api.artifact(f"{entity}/{project}/{name}:latest", type="lora") return artifact.metadata.get("wandb.provenance") diff --git a/tests/integration/test_push_and_fork.py b/tests/integration/test_push_and_fork.py index 1f590f92d..401e281d8 100644 --- a/tests/integration/test_push_and_fork.py +++ b/tests/integration/test_push_and_fork.py @@ -154,9 +154,9 @@ async def test_fork_checkpoint_from_wandb(): # Verify the forked checkpoint matches model A's checkpoint. # Pull both via W&B directly (the fork uploaded the artifact # with a step{N} alias matching the source step). - import wandb + from art.utils import wandb_sdk - api = wandb.Api(api_key=backend._client.api_key) # ty:ignore[possibly-missing-attribute] + api = wandb_sdk.api(api_key=backend._client.api_key) with tempfile.TemporaryDirectory() as tmpdir: dir_a = os.path.join(tmpdir, "a") dir_b = os.path.join(tmpdir, "b") diff --git a/tests/support/chat_template_conformance_cases.py b/tests/support/chat_template_conformance_cases.py index 5912c5784..960bc6599 100644 --- a/tests/support/chat_template_conformance_cases.py +++ b/tests/support/chat_template_conformance_cases.py @@ -7,8 +7,11 @@ from pydantic import BaseModel from transformers.tokenization_utils_base import PreTrainedTokenizerBase -from art.preprocessing.tokenize import _apply_chat_template_token_ids -from art.trajectories import History, Trajectory, TrajectoryGroup, get_messages +from art.preprocessing.tokenize import ( + _apply_chat_template_token_ids, + _messages_for_chat_template, +) +from art.trajectories import History, Trajectory, TrajectoryGroup from art.types import MessagesAndChoices, Tools @@ -121,7 +124,7 @@ def _rendered_ids( ) -> list[int]: return _apply_chat_template_token_ids( tokenizer, - cast(list[dict[str, Any]], get_messages(messages_and_choices)), + _messages_for_chat_template(tokenizer, messages_and_choices), tools=tools, tokenize=True, add_generation_prompt=False, diff --git a/tests/unit/test_megatron_reference_logprobs.py b/tests/unit/test_megatron_reference_logprobs.py index 262136709..5612b522b 100644 --- a/tests/unit/test_megatron_reference_logprobs.py +++ b/tests/unit/test_megatron_reference_logprobs.py @@ -164,7 +164,7 @@ def test_calculate_megatron_logprobs_replays_routes(monkeypatch) -> None: chunk = _Chunk(controller) monkeypatch.setattr( megatron_microbatches, - "create_shared_prefix_state", + "create_prefix_tree_state", lambda **kwargs: (kwargs["group_ids"], kwargs["parent_ids"]), ) monkeypatch.setattr( diff --git a/tests/unit/test_metric_routing.py b/tests/unit/test_metric_routing.py index 529cdf14a..6be608d4f 100644 --- a/tests/unit/test_metric_routing.py +++ b/tests/unit/test_metric_routing.py @@ -1,7 +1,6 @@ import json import os from pathlib import Path -import types from unittest.mock import MagicMock, patch import pytest @@ -46,13 +45,13 @@ def test_get_wandb_run_registers_taxonomy_sections(self, tmp_path: Path) -> None fake_run = MagicMock() fake_run._is_finished = False - fake_wandb = types.SimpleNamespace() - fake_wandb.init = MagicMock(return_value=fake_run) - fake_wandb.define_metric = MagicMock() - fake_wandb.Settings = lambda **kwargs: kwargs + fake_init = MagicMock(return_value=fake_run) with patch.dict(os.environ, {"WANDB_API_KEY": "test-key"}, clear=False): - with patch.dict("sys.modules", {"wandb": fake_wandb}): + with ( + patch("art.model.wandb_sdk.init", fake_init), + patch("art.model.wandb_sdk.settings", lambda **kwargs: kwargs), + ): model = Model( name="test-model", project="test-project", @@ -88,13 +87,13 @@ def test_log_metrics_defines_nested_cost_keys_with_training_step( fake_run._is_finished = False fake_run.config = MagicMock() - fake_wandb = types.SimpleNamespace() - fake_wandb.init = MagicMock(return_value=fake_run) - fake_wandb.define_metric = MagicMock() - fake_wandb.Settings = lambda **kwargs: kwargs + fake_init = MagicMock(return_value=fake_run) with patch.dict(os.environ, {"WANDB_API_KEY": "test-key"}, clear=False): - with patch.dict("sys.modules", {"wandb": fake_wandb}): + with ( + patch("art.model.wandb_sdk.init", fake_init), + patch("art.model.wandb_sdk.settings", lambda **kwargs: kwargs), + ): model = Model( name="test-model", project="test-project", @@ -134,10 +133,7 @@ def test_update_wandb_config_seeds_wandb_init(self, tmp_path: Path) -> None: fake_run._is_finished = False fake_run.config = MagicMock() - fake_wandb = types.SimpleNamespace() - fake_wandb.init = MagicMock(return_value=fake_run) - fake_wandb.define_metric = MagicMock() - fake_wandb.Settings = lambda **kwargs: kwargs + fake_init = MagicMock(return_value=fake_run) payload = { "experiment": {"learning_rate": 1e-5, "batch_size": 4}, @@ -145,7 +141,10 @@ def test_update_wandb_config_seeds_wandb_init(self, tmp_path: Path) -> None: } with patch.dict(os.environ, {"WANDB_API_KEY": "test-key"}, clear=False): - with patch.dict("sys.modules", {"wandb": fake_wandb}): + with ( + patch("art.model.wandb_sdk.init", fake_init), + patch("art.model.wandb_sdk.settings", lambda **kwargs: kwargs), + ): model = Model( name="test-model", project="test-project", @@ -155,7 +154,7 @@ def test_update_wandb_config_seeds_wandb_init(self, tmp_path: Path) -> None: run = model._get_wandb_run() assert run is fake_run - init_kwargs = fake_wandb.init.call_args.kwargs + init_kwargs = fake_init.call_args.kwargs assert init_kwargs["config"] == payload assert "allow_val_change" not in init_kwargs fake_run.config.update.assert_called_once_with(payload) @@ -165,13 +164,13 @@ def test_update_wandb_config_updates_active_run(self, tmp_path: Path) -> None: fake_run._is_finished = False fake_run.config = MagicMock() - fake_wandb = types.SimpleNamespace() - fake_wandb.init = MagicMock(return_value=fake_run) - fake_wandb.define_metric = MagicMock() - fake_wandb.Settings = lambda **kwargs: kwargs + fake_init = MagicMock(return_value=fake_run) with patch.dict(os.environ, {"WANDB_API_KEY": "test-key"}, clear=False): - with patch.dict("sys.modules", {"wandb": fake_wandb}): + with ( + patch("art.model.wandb_sdk.init", fake_init), + patch("art.model.wandb_sdk.settings", lambda **kwargs: kwargs), + ): model = Model( name="test-model", project="test-project", diff --git a/tests/unit/test_moe_routing_real_path.py b/tests/unit/test_moe_routing_real_path.py index dbd22dea6..569824db5 100644 --- a/tests/unit/test_moe_routing_real_path.py +++ b/tests/unit/test_moe_routing_real_path.py @@ -5,7 +5,9 @@ from openai.types.chat.chat_completion import Choice import pytest +import torch +from art.megatron.prefix_tree import parse_prefix_tree_row from art.megatron.routing_replay import ( build_moe_routing_replay_bundle_from_packed_tensors, ) @@ -119,38 +121,48 @@ def _tokenized( *, prompt_id: int, prompt_length: int, + trainable_start: int | None = None, + advantage: float = 1.0, + weight: float = 1.0, + pixel_values: torch.Tensor | None = None, + image_grid_thw: torch.Tensor | None = None, ) -> TokenizedResult: + trainable_start = prompt_length if trainable_start is None else trainable_start return TokenizedResult( - advantage=1.0, + advantage=advantage, chat="", token_ids=token_ids, input_pos=list(range(len(token_ids))), - assistant_mask=[0] * prompt_length + [1] * (len(token_ids) - prompt_length), - logprobs=[math.nan] * prompt_length + [-1.0] * (len(token_ids) - prompt_length), - pixel_values=None, - image_grid_thw=None, + assistant_mask=[0] * trainable_start + [1] * (len(token_ids) - trainable_start), + logprobs=[math.nan] * trainable_start + + [-1.0] * (len(token_ids) - trainable_start), + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, trajectory=Trajectory(), - choice_offsets=[prompt_length], + choice_offsets=[trainable_start], extra_logprobs={}, _tokenizer=_FakeTokenizer(), # type: ignore[arg-type] moe_routed_experts=cast(list[list[list[int]] | None], routes), prompt_id=prompt_id, prompt_length=prompt_length, + weight=weight, ) -def test_pack_carries_routes_through_shared_prefix_splicing() -> None: +def test_pack_carries_routes_through_prefix_tree_splicing() -> None: first = _tokenized( [10, 11, 20, 21], [_route(0), _route(10), _route(20), _route(30)], prompt_id=123, - prompt_length=2, + prompt_length=1, + trainable_start=2, ) second = _tokenized( [10, 11, 22, 23], - [_route(0), _route(99), _route(40), _route(50)], + [_route(99), _route(10), _route(40), _route(50)], prompt_id=123, - prompt_length=2, + prompt_length=1, + trainable_start=2, ) packed = packed_tensors_from_tokenized_results( @@ -161,21 +173,135 @@ def test_pack_carries_routes_through_shared_prefix_splicing() -> None: include_moe_routing=True, ) - assert packed["tokens"].tolist()[0][:6] == [10, 11, 20, 21, 22, 23] + assert packed["tokens"].tolist()[0][:7] == [10, 11, 20, 21, 11, 22, 23] routing_replay = packed["moe_routing_replay"] assert routing_replay is not None - assert routing_replay.expert_indices.tolist()[0][:6] == [ + assert routing_replay.expert_indices.tolist()[0][:7] == [ _route(0), _route(10), _route(20), _route(30), + _route(10), _route(40), _route(50), ] stats = routing_replay.pack_stats - assert stats.shared_prefix_rows == 2 - assert stats.shared_prefix_conflict_rows == 1 - assert stats.shared_prefix_conflict_slots == 4 + assert stats.prefix_tree_rows == 1 + assert stats.prefix_tree_conflict_rows == 1 + assert stats.prefix_tree_conflict_slots == 4 + + +def test_prefix_tree_pack_keeps_trainable_duplicates_in_leaf_metadata() -> None: + first = _tokenized( + [10, 11, 20, 21], + [_route(0), _route(10), _route(20), _route(30)], + prompt_id=123, + prompt_length=1, + trainable_start=2, + advantage=2.0, + weight=0.5, + pixel_values=torch.ones(1, 2), + image_grid_thw=torch.tensor([[1, 2, 3]]), + ) + second = _tokenized( + [10, 11, 20, 22], + [_route(0), _route(10), _route(40), _route(50)], + prompt_id=123, + prompt_length=1, + trainable_start=2, + advantage=4.0, + weight=0.25, + pixel_values=torch.full((1, 2), 2.0), + image_grid_thw=torch.tensor([[4, 5, 6]]), + ) + + packed = packed_tensors_from_tokenized_results( + [first, second], + seq_len=8, + pad_token_id=0, + truncate_long_results=False, + ) + + assert packed["tokens"].tolist()[0][:7] == [10, 11, 20, 21, 11, 20, 22] + assert packed["input_pos"].tolist()[0][:7] == [0, 1, 2, 3, 1, 2, 3] + assert packed["assistant_mask"].tolist()[0][:7] == [ + False, + False, + True, + True, + False, + True, + True, + ] + assert math.isnan(float(packed["logprobs"][0, 1])) + assert float(packed["logprobs"][0, 2]) == -1.0 + assert float(packed["logprobs"][0, 5]) == -1.0 + assert float(packed["advantages"][0, 2]) != float(packed["advantages"][0, 5]) + assert float(packed["weights"][0, 2]) != float(packed["weights"][0, 5]) + assert int(packed["group_ids"][0, 2]) != int(packed["group_ids"][0, 5]) + pixel_values = packed["pixel_values"][0] + image_grid_thw = packed["image_grid_thw"][0] + assert pixel_values is not None + assert image_grid_thw is not None + assert torch.equal(pixel_values, torch.ones(1, 2)) + assert torch.equal(image_grid_thw, torch.tensor([[1, 2, 3]])) + + +def test_prefix_tree_pack_public_api_emits_nested_metadata() -> None: + results = [ + _tokenized( + [10, 11, 20, 101, 201], + [_route(0), _route(10), _route(20), _route(30), _route(40)], + prompt_id=1, + prompt_length=4, + trainable_start=4, + ), + _tokenized( + [10, 11, 20, 102, 202], + [_route(0), _route(10), _route(20), _route(50), _route(60)], + prompt_id=2, + prompt_length=4, + trainable_start=4, + ), + _tokenized( + [10, 12, 30, 103, 203], + [_route(0), _route(70), _route(80), _route(90), _route(100)], + prompt_id=3, + prompt_length=4, + trainable_start=4, + ), + ] + + packed = packed_tensors_from_tokenized_results( + results, + seq_len=16, + pad_token_id=0, + truncate_long_results=False, + ) + tree = parse_prefix_tree_row( + group_ids=packed["group_ids"][0], + parent_ids=packed["parent_ids"][0], + ) + + assert packed["tokens"].tolist()[0][: tree.valid_tokens] == [ + 10, + 11, + 20, + 101, + 201, + 102, + 202, + 12, + 30, + 103, + 203, + ] + assert max(segment.depth for segment in tree.segments) == 2 + assert not packed["assistant_mask"][0, 2] + assert not packed["assistant_mask"][0, 3] + assert packed["assistant_mask"][0, 4] + assert packed["assistant_mask"][0, 6] + assert int(packed["group_ids"][0, 4]) != int(packed["group_ids"][0, 6]) def test_pack_infers_at_least_topk_experts_from_sparse_routes() -> None: diff --git a/tests/unit/test_prefix_tree.py b/tests/unit/test_prefix_tree.py new file mode 100644 index 000000000..f47e7bddc --- /dev/null +++ b/tests/unit/test_prefix_tree.py @@ -0,0 +1,544 @@ +from __future__ import annotations + +import pytest +import torch + +from art.megatron.context_parallel.layout_index import TokenLayoutIndex +from art.megatron.prefix_tree import parse_prefix_tree_row +from art.megatron.prefix_tree_packing import prefix_tree_pack + + +def test_parse_prefix_tree_row_tracks_ancestors_and_depth() -> None: + pack = prefix_tree_pack( + ( + torch.tensor([1, 2, 3, 4, 8]), + torch.tensor([1, 2, 3, 4, 9]), + torch.tensor([1, 2, 3, 5]), + torch.tensor([1, 6]), + ), + max_depth=3, + ) + + tree = parse_prefix_tree_row( + group_ids=pack.group_ids[0], + parent_ids=pack.parent_ids[0], + ) + + assert tree.valid_tokens == int(pack.tokens.numel()) + assert max(segment.depth for segment in tree.segments) == 3 + assert [(segment.group_id, segment.ancestors) for segment in tree.segments] == [ + (1, ()), + (2, (1,)), + (3, (1, 2)), + (4, (1, 2, 3)), + (5, (1, 2, 3)), + (6, (1, 2)), + (7, (1,)), + ] + + +def test_parse_prefix_tree_row_rejects_missing_parent() -> None: + with pytest.raises(RuntimeError, match="missing parent"): + parse_prefix_tree_row( + group_ids=torch.tensor([1, 2]), + parent_ids=torch.tensor([1, 3]), + ) + + +def test_parse_prefix_tree_row_rejects_non_contiguous_group() -> None: + with pytest.raises(RuntimeError, match="contiguous group runs"): + parse_prefix_tree_row( + group_ids=torch.tensor([1, 2, 1]), + parent_ids=torch.tensor([1, 1, 1]), + ) + + +def test_gdn_tree_parser_accepts_nested_tree() -> None: + pytest.importorskip("megatron.core.packed_seq_params") + from art.megatron.gdn.gdn_prefix_tree import ( + GdnPlannerConfig, + build_gdn_rank_execution_plan, + parse_gdn_prefix_tree_segments, + ) + + pack = prefix_tree_pack( + ( + torch.tensor([1, 2, 3, 4]), + torch.tensor([1, 2, 3, 5]), + torch.tensor([1, 6]), + ), + max_depth=2, + ) + + spec = parse_gdn_prefix_tree_segments( + group_ids=pack.group_ids, + parent_ids=pack.parent_ids, + ) + plan = build_gdn_rank_execution_plan(spec, device="cpu") + + assert spec.tree_parent_indices == (-1, 0, 1, 1, 0) + assert spec.tree_depths == (0, 1, 2, 2, 1) + assert [ + sum(bucket.segment_count for bucket in buckets) + for buckets in plan.tree_segment_buckets_by_depth + ] == [0, 2, 2] + child_bucket = plan.tree_segment_buckets_by_depth[1][0] + assert child_bucket.parent_indices is not None + assert child_bucket.parent_indices.tolist() == [-1, -1] + assert child_bucket.family_indices.tolist() == [1, 4] + assert child_bucket.position_indices.tolist() == [0, 1, 2, 0, 5] + assert child_bucket.output_mask is not None + assert child_bucket.output_mask.tolist() == [True, True, True, False, True] + + +def test_gdn_tree_parser_accepts_zero_depth_roots() -> None: + pytest.importorskip("megatron.core.packed_seq_params") + from art.megatron.gdn.gdn_prefix_tree import ( + build_gdn_rank_execution_plan, + parse_gdn_prefix_tree_segments, + ) + + pack = prefix_tree_pack( + ( + torch.tensor([1, 2]), + torch.tensor([1, 3]), + torch.tensor([4]), + ), + max_depth=0, + ) + + spec = parse_gdn_prefix_tree_segments( + group_ids=pack.group_ids, + parent_ids=pack.parent_ids, + ) + plan = build_gdn_rank_execution_plan(spec, device="cpu") + + assert spec.tree_parent_indices == (-1, -1, -1) + assert spec.tree_depths == (0, 0, 0) + assert [bucket.segment_count for bucket in plan.tree_segment_buckets_by_depth[0]] + assert not hasattr(plan, "local_prefix_buckets") + assert not hasattr(plan, "chain_completion_buckets") + assert not hasattr(plan, "prefix_boundary_buckets") + assert all( + not bucket.needs_final_state for bucket in plan.tree_segment_buckets_by_depth[0] + ) + + +def test_gdn_tree_planner_splits_leaf_and_internal_final_state_buckets() -> None: + pytest.importorskip("megatron.core.packed_seq_params") + from art.megatron.gdn.gdn_prefix_tree import ( + GdnPlannerConfig, + build_gdn_rank_execution_plan, + parse_gdn_prefix_tree_segments, + ) + + pack = prefix_tree_pack( + ( + torch.tensor([1, 2, 3, 4, 7]), + torch.tensor([1, 2, 3, 4, 8]), + torch.tensor([1, 2, 5, 6]), + ), + max_depth=2, + ) + + spec = parse_gdn_prefix_tree_segments( + group_ids=pack.group_ids, + parent_ids=pack.parent_ids, + ) + plan = build_gdn_rank_execution_plan( + spec, + device="cpu", + planner_config=GdnPlannerConfig(), + ) + tree_has_children = _tree_has_children(spec) + + depth_one_bucket = plan.tree_segment_buckets_by_depth[1][0] + bucket_state_flags = { + tree_has_children[family_index] + for family_index in depth_one_bucket.family_indices.tolist() + } + assert bucket_state_flags == {False, True} + assert depth_one_bucket.needs_final_state + + +def test_gdn_tree_cp_plan_chains_long_nodes() -> None: + pytest.importorskip("megatron.core.packed_seq_params") + from art.megatron.gdn.gdn_prefix_tree import ( + GdnPlannerConfig, + build_gdn_rank_execution_plan, + parse_gdn_prefix_tree_segments, + ) + + root = torch.arange(1, 321) + mid = torch.arange(1001, 1321) + other = torch.arange(2001, 2321) + pack = prefix_tree_pack( + ( + torch.cat((root, mid, torch.tensor([11]))), + torch.cat((root, mid, torch.tensor([12]))), + torch.cat((root, other, torch.tensor([13]))), + ), + max_depth=3, + ) + spec = parse_gdn_prefix_tree_segments( + group_ids=pack.group_ids, + parent_ids=pack.parent_ids, + ) + config = _chain_every_legal_segment_config() + plans = tuple( + build_gdn_rank_execution_plan( + spec, + device="cpu", + cp_rank=rank, + cp_size=4, + attention_token_layout_index=_uniform_attention_layout( + spec.real_token_count, 4 + ), + planner_config=config, + ) + for rank in range(4) + ) + + assert _covered_token_indices(plans) == set(range(spec.real_token_count)) + assert any(plans[0].tree_chain_buckets_by_depth[0]) + assert any( + bucket + for plan in plans + for depth_buckets in plan.tree_chain_buckets_by_depth[1:] + for bucket in depth_buckets + ) + _assert_remote_parent_state_transfers_cover(spec, plans) + for plan in plans: + assert sum(plan.gdn_token_count for plan in plans) == spec.real_token_count + for depth_buckets in plan.tree_chain_buckets_by_depth: + for bucket in depth_buckets: + assert bucket.lengths_by_rank_cpu is not None + assert tuple(bucket.lengths_by_rank_cpu.shape)[0] == 4 + assert bucket.parent_indices is not None + + +def test_gdn_tree_cp_plan_exchanges_remote_parent_states() -> None: + pytest.importorskip("megatron.core.packed_seq_params") + from art.megatron.gdn.gdn_prefix_tree import ( + build_gdn_rank_execution_plan, + parse_gdn_prefix_tree_segments, + ) + + root = torch.arange(1, 17) + mid = torch.arange(1001, 1321) + pack = prefix_tree_pack( + ( + torch.cat((root, mid, torch.tensor([11]))), + torch.cat((root, mid, torch.tensor([12]))), + torch.cat((root, torch.tensor([99]))), + ), + max_depth=2, + ) + spec = parse_gdn_prefix_tree_segments( + group_ids=pack.group_ids, + parent_ids=pack.parent_ids, + ) + plans = tuple( + build_gdn_rank_execution_plan( + spec, + device="cpu", + cp_rank=rank, + cp_size=4, + attention_token_layout_index=_uniform_attention_layout( + spec.real_token_count, 4 + ), + planner_config=_chain_every_legal_segment_config(), + ) + for rank in range(4) + ) + assert _covered_token_indices(plans) == set(range(spec.real_token_count)) + assert any( + bucket + for plan in plans + for depth_buckets in plan.tree_chain_buckets_by_depth[1:] + for bucket in depth_buckets + ) + assert _remote_parent_state_transfer_count(plans) > 0 + _assert_remote_parent_state_transfers_cover(spec, plans) + + +def test_gdn_tree_cp_randomized_plans_cover_each_token_once() -> None: + pytest.importorskip("megatron.core.packed_seq_params") + from art.megatron.gdn.gdn_prefix_tree import ( + build_gdn_rank_execution_plan, + parse_gdn_prefix_tree_segments, + ) + + config = _chain_every_legal_segment_config() + for seed in range(8): + pack = prefix_tree_pack( + _random_tree_sequences(seed), + max_depth=4, + ) + spec = parse_gdn_prefix_tree_segments( + group_ids=pack.group_ids, + parent_ids=pack.parent_ids, + ) + plans = tuple( + build_gdn_rank_execution_plan( + spec, + device="cpu", + cp_rank=rank, + cp_size=4, + attention_token_layout_index=_uniform_attention_layout( + spec.real_token_count, 4 + ), + planner_config=config, + ) + for rank in range(4) + ) + + assert _covered_token_indices(plans) == set(range(spec.real_token_count)) + assert sum(plan.gdn_token_count for plan in plans) == spec.real_token_count + for plan in plans: + for depth_buckets in ( + *plan.tree_segment_buckets_by_depth, + *plan.tree_chain_buckets_by_depth, + ): + for bucket in depth_buckets: + assert bucket.parent_indices is not None + assert int(bucket.real_token_count) > 0 + + +def test_gdn_tree_cp_randomized_plans_pass_health_checks() -> None: + pytest.importorskip("megatron.core.packed_seq_params") + from art.megatron.gdn.gdn_prefix_tree import ( + GdnPlannerConfig, + build_gdn_rank_execution_plan, + parse_gdn_prefix_tree_segments, + ) + + config = GdnPlannerConfig() + for seed in range(16): + pack = prefix_tree_pack( + _random_tree_sequences(seed + 100, max_depth=5), + max_depth=5, + ) + spec = parse_gdn_prefix_tree_segments( + group_ids=pack.group_ids, + parent_ids=pack.parent_ids, + ) + plans = tuple( + build_gdn_rank_execution_plan( + spec, + device="cpu", + cp_rank=rank, + cp_size=4, + attention_token_layout_index=_uniform_attention_layout( + spec.real_token_count, 4 + ), + planner_config=config, + ) + for rank in range(4) + ) + + _assert_tree_plan_health(spec, plans) + + +def _chain_every_legal_segment_config(): + from art.megatron.gdn.gdn_prefix_tree import GdnPlannerConfig + + return GdnPlannerConfig( + cp_chain_min_runtime_delta_ms=0.0, + runtime_local_recurrent_tokens_per_ms=1.0, + runtime_chain_recurrent_tokens_per_ms=1_000_000_000.0, + runtime_cp_summary_bytes_per_segment=0, + runtime_cp_summary_exchange_count_per_bucket=0, + runtime_cp_summary_compute_segments_per_ms=1_000_000_000.0, + runtime_cp_suffix_scan_latency_ms=0.0, + runtime_parent_state_bytes_per_exchange=0, + ) + + +def _covered_token_indices(plans) -> set[int]: + return { + token + for plan in plans + for start, end, _position in plan.gdn_token_ranges + for token in range(start, end) + } + + +def _local_owner_by_family(plans) -> dict[int, int]: + owner_by_family = {} + for rank, plan in enumerate(plans): + for depth_buckets in plan.tree_segment_buckets_by_depth: + for bucket in depth_buckets: + for family_index in bucket.family_indices.tolist(): + previous = owner_by_family.setdefault(int(family_index), rank) + assert previous == rank + return owner_by_family + + +def _chained_families(plans) -> set[int]: + return { + int(family_index) + for plan in plans + for depth_buckets in plan.tree_chain_buckets_by_depth + for bucket in depth_buckets + for family_index in bucket.family_indices.tolist() + } + + +def _assert_remote_parent_state_transfers_cover(spec, plans) -> None: + owner_by_family = _local_owner_by_family(plans) + chained_families = _chained_families(plans) + + def assert_transfer( + parent_index: int, source_rank: int, dest_rank: int, depth: int + ) -> None: + source_exchange = plans[source_rank].tree_state_exchanges_by_depth[depth] + dest_exchange = plans[dest_rank].tree_state_exchanges_by_depth[depth] + assert source_exchange is not None + assert dest_exchange is not None + assert parent_index in source_exchange.source_family_indices + assert parent_index in dest_exchange.dest_family_indices + matching = [ + transfer + for transfer in dest_exchange.exchange.transfers + if transfer.source_rank == source_rank and transfer.dest_rank == dest_rank + ] + assert matching + + for family_index, parent_index in enumerate(spec.tree_parent_indices): + if parent_index < 0 or parent_index not in owner_by_family: + continue + source_rank = owner_by_family[parent_index] + depth = spec.tree_depths[family_index] + if family_index in chained_families: + for dest_rank in range(len(plans)): + if source_rank != dest_rank: + assert_transfer(parent_index, source_rank, dest_rank, depth) + continue + assert family_index in owner_by_family + dest_rank = owner_by_family[family_index] + if source_rank != dest_rank: + assert_transfer(parent_index, source_rank, dest_rank, depth) + + +def _remote_parent_state_transfer_count(plans) -> int: + return sum( + exchange.exchange.cross_rank_token_count + for plan in plans + for exchange in plan.tree_state_exchanges_by_depth + if exchange is not None + ) // len(plans) + + +def _tree_has_children(spec) -> list[bool]: + has_children = [False] * spec.family_count + for parent_index in spec.tree_parent_indices: + if parent_index >= 0: + has_children[parent_index] = True + return has_children + + +def _assert_tree_plan_health(spec, plans) -> None: + tree_has_children = _tree_has_children(spec) + token_counts = [0] * int(spec.real_token_count) + for plan in plans: + range_tokens = sum( + end - start for start, end, _position in plan.gdn_token_ranges + ) + assert range_tokens == int(plan.gdn_token_count) + assert len(plan.attention_token_indices) == int(plan.attention_token_count) + + bucket_tokens = 0 + for depth_buckets in plan.tree_segment_buckets_by_depth: + for bucket in depth_buckets: + bucket_tokens += int(bucket.real_token_count) + assert bucket.parent_indices is not None + assert int(bucket.parent_indices.numel()) == int(bucket.segment_count) + assert int(bucket.real_token_count) > 0 + bucket_state_flags = { + tree_has_children[family_index] + for family_index in bucket.family_indices.tolist() + } + assert bucket_state_flags == {bucket.needs_final_state} + for family_index, parent_index in zip( + bucket.family_indices.tolist(), + bucket.parent_indices.tolist(), + strict=True, + ): + assert spec.tree_parent_indices[family_index] == parent_index + + for depth_buckets in plan.tree_chain_buckets_by_depth: + for bucket in depth_buckets: + bucket_tokens += int(bucket.real_token_count) + assert bucket.parent_indices is not None + assert int(bucket.parent_indices.numel()) == int(bucket.segment_count) + assert int(bucket.real_token_count) > 0 + bucket_state_flags = { + tree_has_children[family_index] + for family_index in bucket.family_indices.tolist() + } + if bucket.needs_final_state: + assert any(bucket_state_flags) + else: + assert bucket_state_flags == {False} + for family_index, parent_index in zip( + bucket.family_indices.tolist(), + bucket.parent_indices.tolist(), + strict=True, + ): + assert spec.tree_parent_indices[family_index] == parent_index + assert bucket_tokens == int(plan.gdn_token_count) + + for start, end, _position in plan.gdn_token_ranges: + for token_index in range(start, end): + token_counts[token_index] += 1 + + _assert_remote_parent_state_transfers_cover(spec, plans) + assert token_counts == [1] * int(spec.real_token_count) + rank_tokens = [int(plan.gdn_token_count) for plan in plans] + assert max(rank_tokens) - min(rank_tokens) <= max(256, spec.real_token_count // 3) + + +def _uniform_attention_layout(real_token_count: int, cp_size: int) -> TokenLayoutIndex: + ranges_by_rank: list[tuple[tuple[int, int, int], ...]] = [] + for rank in range(cp_size): + start = (int(real_token_count) * rank) // cp_size + end = (int(real_token_count) * (rank + 1)) // cp_size + ranges_by_rank.append(((start, end, 0),) if end > start else ()) + return TokenLayoutIndex( + ownership_ranges_by_rank=tuple(ranges_by_rank), + token_counts_by_rank=tuple( + sum(end - start for start, end, _position in ranges) + for ranges in ranges_by_rank + ), + ) + + +def _random_tree_sequences( + seed: int, *, max_depth: int = 4 +) -> tuple[torch.Tensor, ...]: + generator = torch.Generator().manual_seed(seed) + next_token = 1 + + def tokens(length: int) -> torch.Tensor: + nonlocal next_token + out = torch.arange(next_token, next_token + length) + next_token += length + return out + + def randint(low: int, high: int) -> int: + return int(torch.randint(low, high + 1, (), generator=generator).item()) + + def walk(prefix: torch.Tensor, depth: int) -> list[torch.Tensor]: + segment_length = [1, 3, 17, 64, 129, 257][randint(0, 5)] + here = torch.cat((prefix, tokens(segment_length))) + if depth + 1 >= max_depth: + return [ + torch.cat((here, tokens(randint(1, 9)))) for _ in range(randint(2, 4)) + ] + leaves: list[torch.Tensor] = [] + for _ in range(randint(2, 3)): + leaves.extend(walk(here, depth + 1)) + return leaves + + return tuple(walk(torch.empty(0, dtype=torch.long), 0)) diff --git a/tests/unit/test_prefix_tree_attention_builder.py b/tests/unit/test_prefix_tree_attention_builder.py new file mode 100644 index 000000000..78455e7ee --- /dev/null +++ b/tests/unit/test_prefix_tree_attention_builder.py @@ -0,0 +1,633 @@ +from __future__ import annotations + +import pytest +import torch +from torch.nn.attention.flex_attention import BlockMask +from torch.nn.attention.flex_attention import create_block_mask as torch_block_mask + +pytest.importorskip("megatron.core.packed_seq_params") + +from art.megatron.context_parallel.block_mask import ( + build_block_mask_from_context, + prepare_block_mask_context, +) +from art.megatron.context_parallel.builder import ( + build_dense_reference_mask, + build_prefix_tree_attention_spec, +) +from art.megatron.context_parallel.runtime import get_or_build_runtime_plan +from art.megatron.context_parallel.types import ( + AttnMaskKind, + AttnSlice, + ContextParallelConfig, + ExactMaskMetadata, + FlexMaskSpec, + ParallelTopology, + TokenRange, +) +from art.megatron.prefix_tree_packing import PrefixTreePack, prefix_tree_pack +from art.megatron.prefix_tree_state import create_prefix_tree_state + + +def build_block_mask( + spec: FlexMaskSpec, + *, + group_ids: torch.Tensor, + parent_ids: torch.Tensor, + input_pos: torch.Tensor | None = None, + sliding_window: int | None = None, + device: torch.device, +) -> BlockMask | None: + return build_block_mask_from_context( + spec, + context=prepare_block_mask_context( + group_ids=group_ids, + parent_ids=parent_ids, + input_pos=input_pos, + ), + sliding_window=sliding_window, + device=device, + ) + + +def test_prefix_tree_attention_spec_supports_branching_completions() -> None: + group_ids, parent_ids = _branching_prefix_inputs() + + spec = build_prefix_tree_attention_spec( + group_ids=group_ids, + parent_ids=parent_ids, + ) + dense = build_dense_reference_mask(row_spec=spec.rows[0]) + + assert dense.int().tolist() == [ + [1, 0, 0, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0, 0], + [1, 1, 1, 0, 0, 0, 0], + [1, 1, 1, 1, 0, 0, 0], + [1, 1, 1, 0, 1, 0, 0], + [1, 1, 1, 0, 0, 1, 0], + [1, 1, 1, 0, 0, 1, 1], + ] + + +def test_prefix_tree_attention_spec_matches_tree_reference() -> None: + group_ids, parent_ids = _branching_prefix_inputs() + + spec = build_prefix_tree_attention_spec( + group_ids=group_ids, + parent_ids=parent_ids, + ) + dense = build_dense_reference_mask(row_spec=spec.rows[0]) + + assert dense.equal(_reference_tree_mask(group_ids[0], parent_ids[0])) + + +def test_prefix_tree_can_build_context_parallel_layout() -> None: + group_ids, parent_ids = _branching_prefix_inputs() + spec = build_prefix_tree_attention_spec( + group_ids=group_ids, + parent_ids=parent_ids, + ) + + plan = get_or_build_runtime_plan( + spec, + topology=ParallelTopology(cp=2), + config=ContextParallelConfig(planner_chunk_size=2, planner_max_search_steps=1), + original_seq_len=int(group_ids.numel()), + ) + + assert sum(plan[rank].local_valid_lengths[0] for rank in range(2)) == int( + group_ids.numel() + ) + + +def test_sparse_block_mask_exact_predicate_matches_dense_reference() -> None: + group_ids, parent_ids = _branching_prefix_inputs() + spec = build_prefix_tree_attention_spec( + group_ids=group_ids, + parent_ids=parent_ids, + ) + row = spec.rows[0] + token_indices = torch.arange(row.valid_tokens, dtype=torch.long) + block_mask = build_block_mask( + FlexMaskSpec( + q_len=row.valid_tokens, + k_len=row.valid_tokens, + block_size=(2, 2), + slices=row.slices, + exact_mask=ExactMaskMetadata( + q_token_indices=token_indices, + k_token_indices=token_indices, + cache_key="depth-two", + ), + ), + group_ids=group_ids[0], + parent_ids=parent_ids[0], + device=torch.device("cpu"), + ) + + assert block_mask is not None + q_indices = torch.arange(row.valid_tokens)[:, None] + k_indices = torch.arange(row.valid_tokens)[None, :] + actual = block_mask.mask_mod( + torch.zeros_like(q_indices), + torch.zeros_like(q_indices), + q_indices, + k_indices, + ) + + assert actual.equal(build_dense_reference_mask(row_spec=row)) + + +@pytest.mark.parametrize( + ("name", "pack"), + ( + ( + "no-sharing", + prefix_tree_pack( + ( + torch.tensor([1, 2, 3]), + torch.tensor([4, 5]), + torch.tensor([6, 7, 8, 9]), + ), + max_depth=0, + ), + ), + ( + "depth-one", + prefix_tree_pack( + ( + torch.tensor([1, 2, 3, 4]), + torch.tensor([1, 2, 3, 5]), + torch.tensor([1, 2, 6]), + ), + max_depth=1, + ), + ), + ( + "depth-three", + prefix_tree_pack( + ( + torch.tensor([1, 2, 3, 4, 8]), + torch.tensor([1, 2, 3, 4, 9]), + torch.tensor([1, 2, 3, 5]), + torch.tensor([1, 6]), + ), + max_depth=3, + ), + ), + ), +) +def test_sparse_block_mask_matches_torch_block_metadata( + name: str, + pack: PrefixTreePack, +) -> None: + del name + spec = build_prefix_tree_attention_spec( + group_ids=pack.group_ids, + parent_ids=pack.parent_ids, + ) + row = spec.rows[0] + token_indices = torch.arange(row.valid_tokens, dtype=torch.long) + block_mask = build_block_mask( + FlexMaskSpec( + q_len=row.valid_tokens, + k_len=row.valid_tokens, + block_size=(2, 2), + slices=row.slices, + exact_mask=ExactMaskMetadata( + q_token_indices=token_indices, + k_token_indices=token_indices, + cache_key="torch-parity", + ), + ), + group_ids=pack.group_ids[0], + parent_ids=pack.parent_ids[0], + device=torch.device("cpu"), + ) + + assert block_mask is not None + _assert_matches_torch_block_mask(block_mask) + + +def test_sparse_block_mask_prunes_exact_blocks_rejected_by_group_tree() -> None: + group_ids = torch.tensor([1, 1, 1, 1, 2, 2, 2, 2], dtype=torch.long) + parent_ids = torch.tensor([1, 1, 1, 1, 2, 2, 2, 2], dtype=torch.long) + block_mask = build_block_mask( + FlexMaskSpec( + q_len=4, + k_len=4, + block_size=(2, 2), + slices=( + AttnSlice( + q_range=TokenRange(start=0, end=4), + k_range=TokenRange(start=0, end=4), + mask_kind=AttnMaskKind.CAUSAL, + row_index=0, + ), + ), + exact_mask=ExactMaskMetadata( + q_token_indices=torch.tensor([4, 5, 6, 7], dtype=torch.long), + k_token_indices=torch.tensor([0, 1, 2, 3], dtype=torch.long), + cache_key="all-false-cross-family", + ), + ), + group_ids=group_ids, + parent_ids=parent_ids, + device=torch.device("cpu"), + ) + + assert block_mask is not None + assert int(block_mask.kv_num_blocks.sum().item()) == 0 + assert block_mask.full_kv_num_blocks is not None + assert int(block_mask.full_kv_num_blocks.sum().item()) == 0 + _assert_matches_torch_block_mask(block_mask) + + +def test_prefix_tree_state_builds_batched_block_mask() -> None: + group_ids = torch.tensor( + [ + [1, 1, 2, 2, -1], + [10, 11, 11, -1, -1], + ], + dtype=torch.long, + ) + parent_ids = torch.tensor( + [ + [1, 1, 1, 1, -1], + [10, 10, 10, -1, -1], + ], + dtype=torch.long, + ) + + state = create_prefix_tree_state( + group_ids=group_ids, + parent_ids=parent_ids, + target_device=torch.device("cpu"), + ) + seq_len = int(group_ids.shape[1]) + batch_idx = torch.arange(2)[:, None, None].expand(2, seq_len, seq_len) + query_idx = torch.arange(seq_len)[None, :, None].expand(2, seq_len, seq_len) + kv_idx = torch.arange(seq_len)[None, None, :].expand(2, seq_len, seq_len) + actual = state.block_mask.mask_mod( + batch_idx, + torch.zeros_like(batch_idx), + query_idx, + kv_idx, + ) + spec = build_prefix_tree_attention_spec( + group_ids=group_ids, + parent_ids=parent_ids, + ) + assert int(state.block_mask.kv_num_blocks.shape[0]) == 2 + for row_index, row_spec in enumerate(spec.rows): + valid_tokens = int(row_spec.valid_tokens) + expected = torch.zeros((seq_len, seq_len), dtype=torch.bool) + expected[:valid_tokens, :valid_tokens] = build_dense_reference_mask( + row_spec=row_spec + ) + if valid_tokens < seq_len: + padding_len = seq_len - valid_tokens + expected[valid_tokens:, valid_tokens:] = torch.tril( + torch.ones((padding_len, padding_len), dtype=torch.bool) + ) + assert actual[row_index].equal(expected) + _assert_matches_torch_block_mask(state.block_mask, batch_size=2) + + +def test_sparse_block_mask_matches_torch_for_sliding_window_metadata() -> None: + seq_len = 128 + group_ids = torch.ones(seq_len, dtype=torch.long) + parent_ids = torch.ones(seq_len, dtype=torch.long) + input_pos = torch.arange(seq_len, dtype=torch.long) + block_mask = build_block_mask( + FlexMaskSpec( + q_len=seq_len, + k_len=seq_len, + block_size=(8, 8), + slices=( + AttnSlice( + q_range=TokenRange(start=0, end=seq_len), + k_range=TokenRange(start=0, end=seq_len), + mask_kind=AttnMaskKind.CAUSAL, + row_index=0, + ), + ), + exact_mask=ExactMaskMetadata( + q_token_indices=torch.arange(seq_len, dtype=torch.long), + k_token_indices=torch.arange(seq_len, dtype=torch.long), + cache_key="sliding-window", + ), + ), + group_ids=group_ids, + parent_ids=parent_ids, + input_pos=input_pos, + sliding_window=8, + device=torch.device("cpu"), + ) + + assert block_mask is not None + _assert_matches_torch_block_mask(block_mask) + full_causal_blocks = (seq_len // 8) * (seq_len // 8 + 1) // 2 + assert int(block_mask.kv_num_blocks.sum().item()) < full_causal_blocks + + +def test_context_parallel_stage_masks_match_dense_nested_tree() -> None: + _assert_context_parallel_stage_masks_match_dense( + prefix_tree_pack( + ( + torch.tensor([1, 2, 3, 4, 8]), + torch.tensor([1, 2, 3, 4, 9]), + torch.tensor([1, 2, 3, 5]), + torch.tensor([1, 6]), + ), + max_depth=3, + ), + require_remote_stage=True, + ) + _assert_context_parallel_stage_masks_match_dense( + prefix_tree_pack( + ( + torch.tensor([1, 2, 3]), + torch.tensor([4, 5, 6]), + torch.tensor([7, 8]), + torch.tensor([9, 10, 11, 12]), + ), + max_depth=3, + ), + require_remote_stage=False, + ) + + +def _assert_context_parallel_stage_masks_match_dense( + pack: PrefixTreePack, + *, + require_remote_stage: bool, +) -> None: + spec = build_prefix_tree_attention_spec( + group_ids=pack.group_ids, + parent_ids=pack.parent_ids, + ) + row = spec.rows[0] + dense = build_dense_reference_mask(row_spec=row) + topology = ParallelTopology(cp=2) + config = ContextParallelConfig( + block_size=2, + planner_chunk_size=2, + planner_max_search_steps=1, + planner_remote_stage_token_floor=1, + planner_remote_stage_pair_floor=1, + ) + plan = get_or_build_runtime_plan( + spec, + topology=topology, + config=config, + original_seq_len=int(pack.tokens.numel()), + ) + + checked_stages = 0 + checked_remote_stages = 0 + for rank_plan in plan: + for stage in rank_plan.stage_plans: + if stage.mask_metadata is None: + continue + block_mask = build_block_mask( + FlexMaskSpec( + q_len=stage.q_len, + k_len=stage.k_len, + block_size=(2, 2), + slices=stage.slices, + exact_mask=stage.mask_metadata, + ), + group_ids=pack.group_ids[0], + parent_ids=pack.parent_ids[0], + device=torch.device("cpu"), + ) + assert block_mask is not None + q_offsets = torch.arange(stage.q_len)[:, None] + k_offsets = torch.arange(stage.k_len)[None, :] + actual = block_mask.mask_mod( + torch.zeros_like(q_offsets), + torch.zeros_like(q_offsets), + q_offsets, + k_offsets, + ) + q_tokens = stage.mask_metadata.q_token_indices + k_tokens = stage.mask_metadata.k_token_indices + expected = ( + dense[q_tokens.clamp_min(0)[:, None], k_tokens.clamp_min(0)[None, :]] + & (q_tokens[:, None] >= 0) + & (k_tokens[None, :] >= 0) + ) + + assert actual.equal(expected) + assert _effective_block_mask(block_mask).equal(expected) + _assert_matches_torch_block_mask(block_mask) + checked_stages += 1 + checked_remote_stages += int(not stage.is_local_stage) + + assert checked_stages + if require_remote_stage: + assert checked_remote_stages + + +def _effective_block_mask(block_mask: BlockMask) -> torch.Tensor: + q_len, k_len = block_mask.seq_lengths + q_block, k_block = block_mask.BLOCK_SIZE + effective = torch.zeros((q_len, k_len), dtype=torch.bool) + _fill_full_blocks(effective, block_mask, q_block=q_block, k_block=k_block) + _fill_partial_blocks(effective, block_mask, q_block=q_block, k_block=k_block) + return effective + + +def _fill_full_blocks( + effective: torch.Tensor, + block_mask: BlockMask, + *, + q_block: int, + k_block: int, +) -> None: + if block_mask.full_kv_num_blocks is None or block_mask.full_kv_indices is None: + return + for q_block_index in range(int(block_mask.full_kv_num_blocks.shape[-1])): + q_slice = slice(q_block_index * q_block, (q_block_index + 1) * q_block) + block_count = int(block_mask.full_kv_num_blocks[0, 0, q_block_index]) + for k_block_index in block_mask.full_kv_indices[ + 0, 0, q_block_index, :block_count + ].tolist(): + k_slice = slice( + int(k_block_index) * k_block, + (int(k_block_index) + 1) * k_block, + ) + effective[q_slice, k_slice] = True + + +def _fill_partial_blocks( + effective: torch.Tensor, + block_mask: BlockMask, + *, + q_block: int, + k_block: int, +) -> None: + for q_block_index in range(int(block_mask.kv_num_blocks.shape[-1])): + q_offsets = torch.arange( + q_block_index * q_block, + min((q_block_index + 1) * q_block, effective.shape[0]), + )[:, None] + block_count = int(block_mask.kv_num_blocks[0, 0, q_block_index]) + for k_block_index in block_mask.kv_indices[ + 0, 0, q_block_index, :block_count + ].tolist(): + k_offsets = torch.arange( + int(k_block_index) * k_block, + min((int(k_block_index) + 1) * k_block, effective.shape[1]), + )[None, :] + effective[q_offsets, k_offsets] |= block_mask.mask_mod( + torch.zeros_like(q_offsets), + torch.zeros_like(q_offsets), + q_offsets, + k_offsets, + ) + + +def test_sparse_block_mask_supports_non_monotonic_remote_k_indices() -> None: + q_token_indices = torch.tensor([4, 5, 6, 7], dtype=torch.long) + k_token_indices = torch.tensor([0, 1, 6, 2, 3, 4], dtype=torch.long) + block_mask = build_block_mask( + FlexMaskSpec( + q_len=int(q_token_indices.numel()), + k_len=int(k_token_indices.numel()), + block_size=(2, 2), + slices=( + AttnSlice( + q_range=TokenRange(start=0, end=int(q_token_indices.numel())), + k_range=TokenRange(start=0, end=int(k_token_indices.numel())), + mask_kind=AttnMaskKind.CAUSAL, + row_index=0, + ), + ), + exact_mask=ExactMaskMetadata( + q_token_indices=q_token_indices, + k_token_indices=k_token_indices, + cache_key="non-monotonic-k", + ), + ), + group_ids=torch.ones(8, dtype=torch.long), + parent_ids=torch.ones(8, dtype=torch.long), + device=torch.device("cpu"), + ) + + assert block_mask is not None + q_indices = torch.arange(q_token_indices.numel())[:, None] + k_indices = torch.arange(k_token_indices.numel())[None, :] + + actual = block_mask.mask_mod( + torch.zeros_like(q_indices), + torch.zeros_like(q_indices), + q_indices, + k_indices, + ) + + assert actual.equal(q_token_indices[:, None] >= k_token_indices[None, :]) + _assert_matches_torch_block_mask(block_mask) + + +def _assert_matches_torch_block_mask( + block_mask: BlockMask, + *, + batch_size: int = 1, +) -> None: + q_len, k_len = block_mask.seq_lengths + reference = torch_block_mask( + block_mask.mask_mod, + B=batch_size, + H=1, + Q_LEN=q_len, + KV_LEN=k_len, + device="cpu", + BLOCK_SIZE=block_mask.BLOCK_SIZE, + ) + assert _effective_block_mask(block_mask).equal(_effective_block_mask(reference)) + for counts_name, indices_name in ( + ("kv_num_blocks", "kv_indices"), + ("full_kv_num_blocks", "full_kv_indices"), + ("q_num_blocks", "q_indices"), + ("full_q_num_blocks", "full_q_indices"), + ): + assert _block_entries(block_mask, counts_name, indices_name) == _block_entries( + reference, + counts_name, + indices_name, + ) + + +def _block_entries( + block_mask: BlockMask, + counts_name: str, + indices_name: str, +) -> set[tuple[int, int, int, int]]: + counts = getattr(block_mask, counts_name) + indices = getattr(block_mask, indices_name) + if counts is None or indices is None: + return set() + entries = set() + for batch_index in range(int(counts.shape[0])): + for head_index in range(int(counts.shape[1])): + for block_index in range(int(counts.shape[2])): + block_count = int(counts[batch_index, head_index, block_index]) + for other_block in indices[ + batch_index, + head_index, + block_index, + :block_count, + ].tolist(): + entries.add( + ( + batch_index, + head_index, + block_index, + int(other_block), + ) + ) + return entries + + +def _branching_prefix_inputs() -> tuple[torch.Tensor, torch.Tensor]: + return ( + torch.tensor([[1, 1, 1, 2, 3, 4, 4]], dtype=torch.long), + torch.tensor([[1, 1, 1, 1, 1, 1, 1]], dtype=torch.long), + ) + + +def _reference_tree_mask( + group_ids: torch.Tensor, parent_ids: torch.Tensor +) -> torch.Tensor: + group_list = [int(value) for value in group_ids.tolist()] + parent_by_group: dict[int, int | None] = {} + for group_id, parent_id in zip(group_list, parent_ids.tolist(), strict=True): + group_id = int(group_id) + parent_id = int(parent_id) + if group_id not in parent_by_group: + parent_by_group[group_id] = None if parent_id == group_id else parent_id + + ancestors_by_group = { + group_id: _ancestors(group_id, parent_by_group) for group_id in parent_by_group + } + dense = torch.zeros((len(group_list), len(group_list)), dtype=torch.bool) + for q_pos, q_group in enumerate(group_list): + allowed_groups = ancestors_by_group[q_group] | {q_group} + for k_pos, k_group in enumerate(group_list): + dense[q_pos, k_pos] = k_pos <= q_pos and k_group in allowed_groups + return dense + + +def _ancestors( + group_id: int, + parent_by_group: dict[int, int | None], +) -> set[int]: + ancestors: set[int] = set() + cursor = parent_by_group[group_id] + while cursor is not None and cursor not in ancestors: + ancestors.add(cursor) + cursor = parent_by_group.get(cursor) + return ancestors diff --git a/tests/unit/test_prefix_tree_grad_parity.py b/tests/unit/test_prefix_tree_grad_parity.py new file mode 100644 index 000000000..0be02730e --- /dev/null +++ b/tests/unit/test_prefix_tree_grad_parity.py @@ -0,0 +1,279 @@ +from __future__ import annotations + +from copy import deepcopy + +import pytest +import torch +from torch import nn +import torch.nn.functional as F + +from art.megatron.prefix_tree_packing import PrefixTreePack, prefix_tree_pack + + +class _ToyCausalLM(nn.Module): + def __init__(self) -> None: + super().__init__() + self.token_embedding = nn.Embedding(32, 8, dtype=torch.float64) + self.position_embedding = nn.Embedding(8, 8, dtype=torch.float64) + self.mix = nn.Linear(8, 8, bias=False, dtype=torch.float64) + self.output = nn.Linear(8, 32, bias=False, dtype=torch.float64) + + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + causal_mask: torch.Tensor, + ) -> torch.Tensor: + states = self.token_embedding(input_ids) + self.position_embedding(position_ids) + context = causal_mask.to(states.dtype) @ states + return self.output(torch.tanh(self.mix(context))) + + +@pytest.mark.parametrize("max_depth", (1, 2, 3)) +@pytest.mark.parametrize("multi_target", (False, True)) +def test_prefix_tree_ce_parameter_grads_match_independent_sequences( + *, + max_depth: int, + multi_target: bool, +) -> None: + input_ids = _input_ids() + target_ids = tuple( + _targets(tokens, multi_target=multi_target) for tokens in input_ids + ) + pack = prefix_tree_pack(input_ids, max_depth=max_depth) + + assert int(pack.tokens.numel()) < sum(len(row) for row in input_ids) + + torch.manual_seed(20260518) + naive_model = _ToyCausalLM() + packed_model = deepcopy(naive_model) + + naive_loss = torch.stack( + [ + _sequence_ce_loss(naive_model, tokens, labels) + for tokens, labels in zip(input_ids, target_ids, strict=True) + ] + ).sum() + packed_loss = _packed_ce_loss(packed_model, pack, target_ids) + + torch.testing.assert_close(packed_loss, naive_loss, rtol=1e-12, atol=1e-12) + naive_loss.backward() + packed_loss.backward() + + for (name, naive_param), packed_param in zip( + naive_model.named_parameters(), + packed_model.parameters(), + strict=True, + ): + assert naive_param.grad is not None, name + assert packed_param.grad is not None, name + torch.testing.assert_close( + packed_param.grad, + naive_param.grad, + rtol=1e-10, + atol=1e-10, + msg=lambda msg, name=name: f"{name} grad mismatch:\n{msg}", + ) + + +@pytest.mark.parametrize("max_depth", (1, 2, 3)) +def test_same_layout_mutation_preserves_forward_outputs(max_depth: int) -> None: + pack = prefix_tree_pack(_input_ids(), max_depth=max_depth) + torch.manual_seed(20260518) + model = _ToyCausalLM() + logits = _packed_logits(model, pack) + + for positions in pack.positions_by_sequence: + mutated_logits = _packed_logits(model, _mutated_pack(pack, keep=positions)) + torch.testing.assert_close( + mutated_logits.index_select(0, positions), + logits.index_select(0, positions), + rtol=0.0, + atol=0.0, + ) + + +@pytest.mark.parametrize("max_depth", (1, 2, 3)) +@pytest.mark.parametrize("sequence_index", (0, 2, 4)) +def test_same_layout_mutation_preserves_target_loss_grads( + max_depth: int, + sequence_index: int, +) -> None: + input_ids = _input_ids() + target_ids = tuple(_targets(tokens, multi_target=True) for tokens in input_ids) + pack = prefix_tree_pack(input_ids, max_depth=max_depth) + mutated = _mutated_pack(pack, keep=pack.positions_by_sequence[sequence_index]) + + torch.manual_seed(20260518) + base_model = _ToyCausalLM() + mutated_model = deepcopy(base_model) + + base_loss = _packed_sequence_ce_loss(base_model, pack, target_ids, sequence_index) + mutated_loss = _packed_sequence_ce_loss( + mutated_model, + mutated, + target_ids, + sequence_index, + ) + + torch.testing.assert_close(mutated_loss, base_loss, rtol=0.0, atol=0.0) + base_loss.backward() + mutated_loss.backward() + _assert_matching_grads(mutated_model, base_model) + + +def _input_ids() -> tuple[torch.Tensor, ...]: + return ( + torch.tensor([1, 2, 3, 4, 5]), + torch.tensor([1, 2, 3, 4, 6]), + torch.tensor([1, 2, 3, 7]), + torch.tensor([1, 2, 8]), + torch.tensor([9, 10, 11]), + ) + + +def _targets(tokens: torch.Tensor, *, multi_target: bool) -> torch.Tensor: + labels = (tokens * 3 + 5) % 31 + if not multi_target: + return labels + alternate = (tokens * 5 + 7) % 31 + stacked = torch.stack((labels, alternate), dim=1) + if int(stacked.numel()) > 2: + stacked[1, 1] = -100 + return stacked + + +def _sequence_ce_loss( + model: _ToyCausalLM, + input_ids: torch.Tensor, + target_ids: torch.Tensor, +) -> torch.Tensor: + seq_len = int(input_ids.numel()) + logits = model( + input_ids, + torch.arange(seq_len), + torch.ones((seq_len, seq_len), dtype=torch.bool).tril(), + ) + return _target_ce_loss(logits, target_ids) + + +def _packed_ce_loss( + model: _ToyCausalLM, + pack: PrefixTreePack, + target_ids: tuple[torch.Tensor, ...], +) -> torch.Tensor: + logits = _packed_logits(model, pack) + losses = [ + _target_ce_loss(logits.index_select(0, positions), labels) + for positions, labels in zip( + pack.positions_by_sequence, + target_ids, + strict=True, + ) + ] + return torch.stack(losses).sum() + + +def _packed_sequence_ce_loss( + model: _ToyCausalLM, + pack: PrefixTreePack, + target_ids: tuple[torch.Tensor, ...], + sequence_index: int, +) -> torch.Tensor: + return _target_ce_loss( + _packed_logits(model, pack).index_select( + 0, + pack.positions_by_sequence[sequence_index], + ), + target_ids[sequence_index], + ) + + +def _packed_logits(model: _ToyCausalLM, pack: PrefixTreePack) -> torch.Tensor: + return model( + pack.tokens.reshape(-1), + pack.position_ids.reshape(-1), + _prefix_tree_causal_mask(pack), + ) + + +def _target_ce_loss(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + if labels.ndim == 1: + return F.cross_entropy(logits, labels, ignore_index=-100, reduction="sum") + expanded = logits.unsqueeze(1).expand(-1, int(labels.shape[1]), -1) + return F.cross_entropy( + expanded.reshape(-1, int(logits.shape[-1])), + labels.reshape(-1), + ignore_index=-100, + reduction="sum", + ) + + +def _mutated_pack(pack: PrefixTreePack, *, keep: torch.Tensor) -> PrefixTreePack: + tokens = pack.tokens.clone() + mutate = torch.ones(int(tokens.shape[1]), dtype=torch.bool) + mutate[keep] = False + replacement = torch.arange(int(tokens.shape[1]), dtype=tokens.dtype) + 17 + tokens[0, mutate] = replacement[mutate] % 31 + return PrefixTreePack( + tokens=tokens, + group_ids=pack.group_ids, + parent_ids=pack.parent_ids, + position_ids=pack.position_ids, + positions_by_sequence=pack.positions_by_sequence, + ) + + +def _assert_matching_grads(actual_model: nn.Module, expected_model: nn.Module) -> None: + for (name, expected_param), actual_param in zip( + expected_model.named_parameters(), + actual_model.parameters(), + strict=True, + ): + assert expected_param.grad is not None, name + assert actual_param.grad is not None, name + torch.testing.assert_close( + actual_param.grad, + expected_param.grad, + rtol=1e-10, + atol=1e-10, + msg=lambda msg, name=name: f"{name} grad mismatch:\n{msg}", + ) + + +def _prefix_tree_causal_mask(pack: PrefixTreePack) -> torch.Tensor: + group_ids = pack.group_ids.reshape(-1).tolist() + parent_ids = pack.parent_ids.reshape(-1).tolist() + position_ids = pack.position_ids.reshape(-1).tolist() + parent_by_group: dict[int, int] = {} + for group_id, parent_id in zip(group_ids, parent_ids, strict=True): + previous = parent_by_group.setdefault(group_id, parent_id) + assert previous == parent_id + + ancestors = { + group_id: _ancestor_groups(group_id, parent_by_group) + for group_id in parent_by_group + } + mask = torch.zeros((len(group_ids), len(group_ids)), dtype=torch.bool) + for query_index, query_group in enumerate(group_ids): + query_ancestors = ancestors[query_group] + query_position = position_ids[query_index] + for key_index, key_group in enumerate(group_ids): + if ( + key_group in query_ancestors + and position_ids[key_index] <= query_position + ): + mask[query_index, key_index] = True + return mask + + +def _ancestor_groups(group_id: int, parent_by_group: dict[int, int]) -> set[int]: + ancestors = {group_id} + parent_id = parent_by_group[group_id] + while parent_id != group_id: + if parent_id in ancestors: + raise AssertionError("prefix-tree group parents contain a cycle") + ancestors.add(parent_id) + group_id = parent_id + parent_id = parent_by_group[group_id] + return ancestors diff --git a/tests/unit/test_prefix_tree_packing.py b/tests/unit/test_prefix_tree_packing.py new file mode 100644 index 000000000..3f27eac9c --- /dev/null +++ b/tests/unit/test_prefix_tree_packing.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +import pytest +import torch + +from art.megatron.prefix_tree_packing import ( + _local_position_pairs, + estimate_prefix_tree_packed_tokens, + prefix_tree_pack, +) + + +def test_prefix_tree_pack_support_depth_one() -> None: + inputs = ( + torch.tensor([1, 2, 3, 4]), + torch.tensor([1, 2, 5]), + torch.tensor([9]), + ) + + pack = prefix_tree_pack(inputs, max_depth=1) + + assert pack.tokens.tolist() == [[1, 2, 3, 4, 5, 9]] + assert pack.group_ids.tolist() == [[1, 1, 2, 2, 3, 4]] + assert pack.parent_ids.tolist() == [[1, 1, 1, 1, 1, 4]] + assert pack.position_ids.tolist() == [[0, 1, 2, 3, 2, 0]] + assert [positions.tolist() for positions in pack.positions_by_sequence] == [ + [0, 1, 2, 3], + [0, 1, 4], + [5], + ] + + +def test_prefix_tree_pack_support_zero_depth_without_sharing() -> None: + pack = prefix_tree_pack( + ( + torch.tensor([1, 2]), + torch.tensor([1, 3]), + torch.tensor([4]), + ), + max_depth=0, + ) + + assert pack.tokens.tolist() == [[1, 2, 1, 3, 4]] + assert pack.group_ids.tolist() == [[1, 1, 2, 2, 3]] + assert pack.parent_ids.tolist() == [[1, 1, 2, 2, 3]] + assert pack.position_ids.tolist() == [[0, 1, 0, 1, 0]] + assert [positions.tolist() for positions in pack.positions_by_sequence] == [ + [0, 1], + [2, 3], + [4], + ] + + +def test_prefix_tree_pack_support_deeper_trees() -> None: + pack = prefix_tree_pack( + ( + torch.tensor([1, 2, 3, 4]), + torch.tensor([1, 2, 3, 5]), + torch.tensor([1, 6, 7]), + ), + max_depth=2, + ) + + assert pack.tokens.tolist() == [[1, 2, 3, 4, 5, 6, 7]] + assert pack.group_ids.tolist() == [[1, 2, 2, 3, 4, 5, 5]] + assert pack.parent_ids.tolist() == [[1, 1, 1, 2, 2, 1, 1]] + assert pack.position_ids.tolist() == [[0, 1, 2, 3, 3, 1, 2]] + assert [positions.tolist() for positions in pack.positions_by_sequence] == [ + [0, 1, 2, 3], + [0, 1, 2, 4], + [0, 5, 6], + ] + + +def test_packing_preserves_first_seen_branch_order() -> None: + pack = prefix_tree_pack( + (torch.tensor([9]), torch.tensor([1])), + max_depth=1, + ) + + assert pack.tokens.tolist() == [[9, 1]] + assert [positions.tolist() for positions in pack.positions_by_sequence] == [ + [0], + [1], + ] + + +def test_packing_handles_empty_sequences() -> None: + pack = prefix_tree_pack( + (torch.empty(0, dtype=torch.long), torch.empty(0, dtype=torch.long)), + max_depth=1, + ) + + assert pack.tokens.tolist() == [[]] + assert pack.group_ids.tolist() == [[]] + assert pack.parent_ids.tolist() == [[]] + assert [positions.tolist() for positions in pack.positions_by_sequence] == [[], []] + + +def test_prefix_tree_pack_respects_shareable_lengths() -> None: + inputs = ( + torch.tensor([1, 2, 3]), + torch.tensor([1, 2, 4]), + ) + + pack = prefix_tree_pack(inputs, max_depth=4, shareable_lengths=(1, 1)) + + assert pack.tokens.tolist() == [[1, 2, 3, 2, 4]] + assert [positions.tolist() for positions in pack.positions_by_sequence] == [ + [0, 1, 2], + [0, 3, 4], + ] + assert estimate_prefix_tree_packed_tokens( + inputs, + max_depth=4, + shareable_lengths=(1, 1), + ) == int(pack.tokens.numel()) + + +def test_packed_token_estimator_matches_real_packing() -> None: + cases = [ + (torch.tensor([1, 2, 3]), torch.tensor([1, 2, 4]), torch.tensor([5])), + ( + torch.tensor([1, 2, 3, 4]), + torch.tensor([1, 2, 3, 5]), + torch.tensor([1, 2, 6, 7]), + torch.tensor([1, 8]), + ), + ( + torch.tensor([9, 1, 2]), + torch.tensor([9, 1, 3]), + torch.tensor([9, 4, 5]), + torch.tensor([6, 7]), + torch.tensor([], dtype=torch.long), + ), + ] + + for inputs in cases: + for depth in range(5): + pack = prefix_tree_pack(inputs, max_depth=depth) + + assert estimate_prefix_tree_packed_tokens(inputs, max_depth=depth) == int( + pack.tokens.numel() + ) + + +def test_packed_token_estimator_matches_randomized_packing() -> None: + generator = torch.Generator().manual_seed(123) + inputs = [] + for family in range(5): + prefix = torch.randint(1, 100, (4,), generator=generator) + for branch in range(4): + middle = torch.tensor([family, branch]) + suffix = torch.randint(1, 100, (3,), generator=generator) + inputs.append(torch.cat((prefix, middle, suffix))) + + for depth in range(5): + pack = prefix_tree_pack(inputs, max_depth=depth) + + assert estimate_prefix_tree_packed_tokens(inputs, max_depth=depth) == int( + pack.tokens.numel() + ) + + +def test_packing_rejects_non_1d_sequences() -> None: + with pytest.raises(ValueError, match="expects 1-D tensors"): + prefix_tree_pack((torch.tensor([[1, 2], [3, 4]]),), max_depth=1) + + +def test_local_position_pairs_preserve_requested_order_without_dense_match() -> None: + local_global_positions = torch.tensor([[2, -1, 0, 4, 1]]) + item_positions = torch.tensor([0, 1, 2, 3, 4]) + + local_positions, source_positions = _local_position_pairs( + local_global_positions, + item_positions, + ) + + assert local_positions.tolist() == [2, 4, 0, 3] + assert source_positions.tolist() == [0, 1, 2, 4] diff --git a/tests/unit/test_serverless_pipeline_trainer_compat.py b/tests/unit/test_serverless_pipeline_trainer_compat.py index 25dd1a415..66a63f848 100644 --- a/tests/unit/test_serverless_pipeline_trainer_compat.py +++ b/tests/unit/test_serverless_pipeline_trainer_compat.py @@ -1,4 +1,3 @@ -import sys from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -250,12 +249,6 @@ def log_artifact(self, artifact): def finish(self) -> None: pass - fake_wandb = SimpleNamespace( - Artifact=FakeArtifact, - init=MagicMock(return_value=FakeRun()), - Settings=lambda **kwargs: kwargs, - ) - trajectory = Trajectory( messages_and_choices=[ {"role": "user", "content": "prompt"}, @@ -264,7 +257,11 @@ def finish(self) -> None: ) with patch.object(model, "_get_wandb_run", return_value=None): - with patch.dict(sys.modules, {"wandb": fake_wandb}): + with ( + patch("art.serverless.backend.wandb_sdk.artifact", FakeArtifact), + patch("art.serverless.backend.wandb_sdk.init", return_value=FakeRun()), + patch("art.serverless.backend.wandb_sdk.settings", lambda **kwargs: kwargs), + ): with patch("art.serverless.backend.asyncio.sleep", no_sleep): async for _ in backend._train_sft( model,