diff --git a/src/xorl/models/transformers/deepseek_v4/exact_lm_head.py b/src/xorl/models/transformers/deepseek_v4/exact_lm_head.py index ecfe7064..7b2af0ca 100644 --- a/src/xorl/models/transformers/deepseek_v4/exact_lm_head.py +++ b/src/xorl/models/transformers/deepseek_v4/exact_lm_head.py @@ -10,8 +10,6 @@ from __future__ import annotations from dataclasses import dataclass -from functools import lru_cache -from typing import Any import torch import torch.distributed as dist @@ -20,19 +18,21 @@ from xorl.distributed.canonical_moe import LogicalRowOwnership from xorl.lora.modules.linear import LoraLinear +from xorl.models.transformers.exact_lm_head_shared import ( + ExactHeadRowPlan, + ExactLmHeadFunction, + check_exact_head_tp_group, + exact_lora_local_logits, + filtered_surrogate_local_grad_logits, + rank_order_row_counts, + rank_order_variable_row_all_gather, + rank_order_vocab_all_gather, + surrogate_local_grad_logits, +) from xorl.ops.bi_families_v2 import exact_temperature_scale_bf16_logits from xorl.ops.exact_sampling_transforms import ( - EXACT_FILTER_ROW_CHUNK, - exact_sampling_identity_rows, - exact_sampling_support, - exact_selected_logprob_partitioned_from_support, normalize_temperature_rows, -) -from xorl.ops.exact_sampling_transforms import ( - selected_logprob_reference_grad as _selected_logprob_reference_grad, -) -from xorl.ops.exact_sampling_transforms import ( - selected_logprob_reference_grad_partitioned as _selected_logprob_reference_grad_partitioned, + score_with_sampling_transforms, ) from xorl.ops.exact_sampling_transforms import ( validate_temperature_rows as _validate_temperature_rows, @@ -69,35 +69,14 @@ def dsv4_lm_head_shard(tp_rank: int) -> Dsv4LmHeadShard: return Dsv4LmHeadShard(tp_rank, start, start + DSV4_LM_HEAD_LOCAL_VOCAB_SIZE) -@lru_cache(maxsize=64) -def _single_adapter_batch_info(device_index: int, rows: int) -> Any: - from sglang.srt.lora.utils import LoRABatchInfo # noqa: PLC0415 - - device = torch.device("cuda", device_index) - return LoRABatchInfo( - use_cuda_graph=False, - bs=1, - num_segments=1, - seg_indptr=torch.tensor([0, rows], dtype=torch.int32, device=device), - weight_indices=torch.zeros(1, dtype=torch.int32, device=device), - lora_ranks=torch.ones(1, dtype=torch.int32, device=device), - scalings=torch.ones(1, dtype=torch.float32, device=device), - max_len=rows, - seg_lens=torch.tensor([rows], dtype=torch.int32, device=device), - permutation=None, - expected_tokens=rows, - has_active_lora=True, - ) - - def _rank_order_row_counts(local_rows: int, device: torch.device, group: dist.ProcessGroup) -> tuple[int, ...]: - local = torch.tensor([local_rows], dtype=torch.int64, device=device) - gathered = torch.empty(DSV4_LM_HEAD_TP_SIZE, dtype=torch.int64, device=device) - dist.all_gather_into_tensor(gathered, local, group=group) - counts = tuple(int(value) for value in gathered.cpu().tolist()) - if any(value < 0 for value in counts): - raise RuntimeError(f"DSV4 exact lm head received negative row counts: {counts}") - return counts + return rank_order_row_counts( + local_rows, + device, + group, + world_size=DSV4_LM_HEAD_TP_SIZE, + program="DSV4 exact lm head", + ) def _rank_order_variable_row_all_gather( @@ -107,39 +86,42 @@ def _rank_order_variable_row_all_gather( row_counts: tuple[int, ...], padded_rows: int, ) -> Tensor: - if value.ndim == 0 or not value.is_contiguous(): - raise ValueError("DSV4 row all-gather requires a contiguous non-scalar tensor") - if value.shape[0] > padded_rows or len(row_counts) != DSV4_LM_HEAD_TP_SIZE: - raise ValueError(f"Invalid DSV4 variable-row geometry: local={value.shape[0]}, counts={row_counts}") - if value.shape[0] < padded_rows: - padding = value.new_zeros((padded_rows - value.shape[0], *value.shape[1:])) - value = torch.cat((value, padding), dim=0) - gathered = torch.empty( - (DSV4_LM_HEAD_TP_SIZE * padded_rows, *value.shape[1:]), - dtype=value.dtype, - device=value.device, + return rank_order_variable_row_all_gather( + value, + group, + world_size=DSV4_LM_HEAD_TP_SIZE, + row_counts=row_counts, + padded_rows=padded_rows, ) - dist.all_gather_into_tensor(gathered, value, group=group) - pieces = [ - gathered[rank * padded_rows : rank * padded_rows + count] for rank, count in enumerate(row_counts) if count - ] - return torch.cat(pieces, dim=0) if pieces else gathered[:0] def _rank_order_vocab_all_gather(local_logits: Tensor, group: dist.ProcessGroup) -> Tensor: - if local_logits.dtype is not torch.bfloat16 or not local_logits.is_contiguous(): - raise ValueError("DSV4 local logits must be contiguous BF16") - rows = local_logits.shape[0] - gathered = torch.empty( - (DSV4_LM_HEAD_TP_SIZE * rows, DSV4_LM_HEAD_LOCAL_VOCAB_SIZE), - dtype=local_logits.dtype, - device=local_logits.device, + return rank_order_vocab_all_gather( + local_logits, + group, + expected_world_size=DSV4_LM_HEAD_TP_SIZE, + expected_local_vocab_size=DSV4_LM_HEAD_LOCAL_VOCAB_SIZE, + expected_dtype=torch.bfloat16, ) - dist.all_gather_into_tensor(gathered, local_logits, group=group) - return ( - gathered.view(DSV4_LM_HEAD_TP_SIZE, rows, DSV4_LM_HEAD_LOCAL_VOCAB_SIZE) - .permute(1, 0, 2) - .reshape(rows, DSV4_LM_HEAD_VOCAB_SIZE) + + +def _distributed_row_plan(local_hidden: Tensor, group: dist.ProcessGroup, source_ordinal: int) -> ExactHeadRowPlan: + """Ragged rank-order row blocks: TP8 ranks may own different row counts.""" + + row_counts = _rank_order_row_counts(local_hidden.shape[0], local_hidden.device, group) + if sum(row_counts) <= 0: + raise ValueError("DSV4 exact lm head requires at least one live row across TP8") + padded_rows = max(row_counts) + source_offset = sum(row_counts[:source_ordinal]) + source_rows = row_counts[source_ordinal] + return ExactHeadRowPlan( + lambda value: _rank_order_variable_row_all_gather( + value, + group, + row_counts=row_counts, + padded_rows=padded_rows, + ), + lambda value: value.narrow(0, source_offset, source_rows), ) @@ -200,153 +182,6 @@ def _local_surrogate_vjp( return by_label.get("hidden"), by_label.get("a"), by_label.get("b") -class _Dsv4ExactDistributedHeadFunction(torch.autograd.Function): - @staticmethod - def forward( - ctx, - local_hidden: Tensor, - local_weight: Tensor, - lora_a: Tensor, - local_lora_b: Tensor, - local_token_ids: Tensor, - local_temperature: Tensor | None, - local_sampling_transforms: tuple[Tensor | None, Tensor | None, Tensor | None], - component: "Dsv4ExactTP8LmHeadSelectedLogprob", - ) -> Tensor: - group = component._validate_tp_group() - row_counts = _rank_order_row_counts(local_hidden.shape[0], local_hidden.device, group) - if sum(row_counts) <= 0: - raise ValueError("DSV4 exact lm head requires at least one live row across TP8") - padded_rows = max(row_counts) - - effective_a = lora_a.to(torch.bfloat16).contiguous() - effective_b = local_lora_b.to(torch.bfloat16).contiguous() - gathered_hidden = _rank_order_variable_row_all_gather( - local_hidden, - group, - row_counts=row_counts, - padded_rows=padded_rows, - ) - gathered_ids = _rank_order_variable_row_all_gather( - local_token_ids, - group, - row_counts=row_counts, - padded_rows=padded_rows, - ) - gathered_temperature = ( - None - if local_temperature is None - else _rank_order_variable_row_all_gather( - local_temperature, - group, - row_counts=row_counts, - padded_rows=padded_rows, - ) - ) - gathered_sampling_transforms = tuple( - None - if value is None - else _rank_order_variable_row_all_gather( - value, - group, - row_counts=row_counts, - padded_rows=padded_rows, - ) - for value in local_sampling_transforms - ) - has_sampling_filter = gathered_sampling_transforms[0] is not None - if not has_sampling_filter: - gathered_logprob = component._exact_forward_value( - gathered_hidden, - local_weight, - effective_a, - effective_b, - gathered_ids, - gathered_temperature, - ) - else: - gathered_logprob = component._exact_forward_value_filtered( - gathered_hidden, - local_weight, - effective_a, - effective_b, - gathered_ids, - gathered_temperature, - gathered_sampling_transforms, - ) - source_ordinal = component.source_ordinal - rows = row_counts[source_ordinal] - source_offset = sum(row_counts[:source_ordinal]) - local_logprob = gathered_logprob.narrow(0, source_offset, rows).contiguous() - ctx.set_materialize_grads(False) - ctx.component = component - ctx.local_rows = rows - ctx.source_ordinal = source_ordinal - ctx.source_offset = source_offset - ctx.row_counts = row_counts - ctx.padded_rows = padded_rows - ctx.has_sampling_filter = has_sampling_filter - ctx.save_for_backward( - gathered_hidden, - local_weight, - effective_a, - effective_b, - gathered_ids, - gathered_temperature - if gathered_temperature is not None - else torch.empty((0,), dtype=torch.float32, device=local_hidden.device), - lora_a, - local_lora_b, - *gathered_sampling_transforms, - ) - return local_logprob - - @staticmethod - def backward(ctx, grad_local_logprob: Tensor | None): - if grad_local_logprob is None: - return (None, None, None, None, None, None, None, None) - ( - gathered_hidden, - local_weight, - effective_a, - effective_b, - gathered_ids, - stored_temperature, - _a_master, - _b_master, - top_ks, - top_ps, - min_ps, - ) = ctx.saved_tensors - temperature = None if stored_temperature.numel() == 0 else stored_temperature - group = ctx.component._validate_tp_group() - gathered_grad = _rank_order_variable_row_all_gather( - grad_local_logprob.contiguous(), - group, - row_counts=ctx.row_counts, - padded_rows=ctx.padded_rows, - ) - vjp = ctx.component._surrogate_vjp_filtered if ctx.has_sampling_filter else ctx.component._surrogate_vjp - args = ( - gathered_hidden, - local_weight, - effective_a, - effective_b, - gathered_ids, - gathered_grad, - temperature, - ) - if ctx.has_sampling_filter: - args = (*args, (top_ks, top_ps, min_ps)) - grad_hidden, grad_a, grad_b = vjp( - *args, - needs_input_grad=(ctx.needs_input_grad[0], ctx.needs_input_grad[2], ctx.needs_input_grad[3]), - ) - if grad_hidden is not None: - grad_hidden = grad_hidden.narrow(0, ctx.source_offset, ctx.local_rows).contiguous() - return (grad_hidden, None, grad_a, grad_b, None, None, None, None) - - class Dsv4ExactTP8LmHeadLoraLinear(LoraLinear): _dsv4_exact_tp8_lm_head = True @@ -390,17 +225,18 @@ def _validate_tp_group(self) -> dist.ProcessGroup: if group is None or not dist.is_initialized(): raise RuntimeError("DSV4 exact lm head requires an initialized TP8 group") ranks = tuple(dist.get_process_group_ranks(group)) - expected_ranks = ranks if self.physical_ranks is None else self.physical_ranks - if ( - dist.get_world_size(group) != DSV4_LM_HEAD_TP_SIZE - or ranks != expected_ranks - or dist.get_rank(group) != self.shard.tp_rank - or dist.get_rank(group) != self.source_ordinal - or dist.get_rank() != ranks[self.shard.tp_rank] - ): - raise RuntimeError("DSV4 exact lm-head TP8 rank/order mismatch") - if str(dist.get_backend(group)).lower() != "nccl": - raise RuntimeError("DSV4 exact lm head requires NCCL") + check_exact_head_tp_group( + program="DSV4 exact lm-head", + world_size=dist.get_world_size(group), + group_rank=dist.get_rank(group), + global_rank=dist.get_rank(), + group_ranks=ranks, + backend=str(dist.get_backend(group)).lower(), + expected_world_size=DSV4_LM_HEAD_TP_SIZE, + expected_ranks=ranks if self.physical_ranks is None else self.physical_ranks, + shard_rank=self.shard.tp_rank, + source_ordinal=self.source_ordinal, + ) return group def _validate_operands( @@ -438,23 +274,16 @@ def _validate_operands( @staticmethod def _exact_local_logits(hidden: Tensor, weight: Tensor, effective_a: Tensor, effective_b: Tensor) -> Tensor: - from sglang.kernels.ops.gemm.sgemm_lora_a import sgemm_lora_a_fwd # noqa: PLC0415 - from sglang.kernels.ops.gemm.sgemm_lora_b import sgemm_lora_b_fwd # noqa: PLC0415 - base_logits = F.linear(hidden, weight) if base_logits.dtype is not torch.bfloat16: raise RuntimeError("DSV4 sampler-aligned base head must store BF16 logits") - batch_info = _single_adapter_batch_info(hidden.device.index, hidden.shape[0]) - lora_a_output = sgemm_lora_a_fwd(hidden, effective_a.unsqueeze(0), batch_info) - output = sgemm_lora_b_fwd( - lora_a_output, - effective_b.unsqueeze(0), - batch_info, - base_output=base_logits, - ) - if output.data_ptr() != base_logits.data_ptr() or output.dtype is not torch.bfloat16: - raise RuntimeError("DSV4 sampler-aligned LoRA B must update the BF16 base-logit buffer in place") - return output.contiguous() + return exact_lora_local_logits( + hidden, + effective_a, + effective_b, + base_logits=base_logits, + rank=1, + ).contiguous() def _exact_forward_value( self, @@ -497,18 +326,6 @@ def _exact_forward_value_filtered( top_ks, top_ps, min_ps = sampling_transforms if top_ks is None or top_ps is None or min_ps is None: raise ValueError("filtered DSV4 exact scoring requires complete row metadata") - support = exact_sampling_support( - score_logits, - top_ks, - top_ps, - min_ps, - ) - identity_rows = exact_sampling_identity_rows( - top_ks, - top_ps, - min_ps, - vocab_size=score_logits.shape[1], - ) from sglang.srt.batch_invariant_ops.batch_invariant_ops import ( # noqa: PLC0415 log_softmax as _bi_log_softmax, ) @@ -519,15 +336,35 @@ def _native_score(native_logits: Tensor, native_ids: Tensor): native_logprob = native_logprobs.gather(1, native_ids.unsqueeze(1)).squeeze(1) return native_logprob, native_selected - native_logprob, native_selected - logprob, _, _ = exact_selected_logprob_partitioned_from_support( + logprob, _, _ = score_with_sampling_transforms( score_logits, token_ids, - support, - identity_rows, + top_ks, + top_ps, + min_ps, _native_score, ) return logprob.contiguous() + def _reference_full_logits_fn(self, local_weight, effective_a, effective_b, group): + """The differentiable rank-1 reference, gathered in serving byte order.""" + + def reference_full_logits(hidden_chunk: Tensor) -> Tensor: + local_reference = ( + ( + F.linear(hidden_chunk, local_weight) + + F.linear( + F.linear(hidden_chunk.float(), effective_a.float()), + effective_b.float(), + ).to(torch.bfloat16) + ) + .float() + .contiguous() + ) + return _rank_order_vocab_all_gather(local_reference.to(torch.bfloat16), group).float() + + return reference_full_logits + def _surrogate_vjp( self, hidden: Tensor, @@ -546,23 +383,14 @@ def _surrogate_vjp( rows=hidden.shape[0], device=hidden.device, ) - with torch.no_grad(): - local_reference = ( - ( - F.linear(hidden, local_weight) - + F.linear(F.linear(hidden.float(), effective_a.float()), effective_b.float()).to(torch.bfloat16) - ) - .float() - .contiguous() - ) - full_reference = _rank_order_vocab_all_gather(local_reference.to(torch.bfloat16), group).float() - full_grad = _selected_logprob_reference_grad( - full_reference, + local_grad = surrogate_local_grad_logits( + hidden, token_ids, grad_logprob, temperature, + reference_full_logits_fn=self._reference_full_logits_fn(local_weight, effective_a, effective_b, group), + local_vocab_slice=slice(self.shard.vocab_start, self.shard.vocab_end), ) - local_grad = full_grad[:, self.shard.vocab_start : self.shard.vocab_end].contiguous() grad_hidden, grad_a, grad_b = _local_surrogate_vjp( hidden, local_weight, @@ -592,68 +420,31 @@ def _surrogate_vjp_filtered( ) -> tuple[Tensor | None, Tensor | None, Tensor | None]: group = self._validate_tp_group() temperature = _validate_temperature_rows(temperature, rows=hidden.shape[0], device=hidden.device) - top_ks, top_ps, min_ps = sampling_transforms - if top_ks is None or top_ps is None or min_ps is None: - raise ValueError("filtered DSV4 exact scoring requires complete row metadata") - local_grad = torch.empty( - (hidden.shape[0], DSV4_LM_HEAD_LOCAL_VOCAB_SIZE), - dtype=torch.float32, - device=hidden.device, - ) - with torch.no_grad(): - for row_start in range(0, hidden.shape[0], EXACT_FILTER_ROW_CHUNK): - row_end = min(row_start + EXACT_FILTER_ROW_CHUNK, hidden.shape[0]) - row_slice = slice(row_start, row_end) - hidden_chunk = hidden[row_slice].contiguous() - - # Recreate the literal value path only for this bounded row - # chunk. The dense support mask is transient and never saved - # on the autograd context. - exact_local = self._exact_local_logits( - hidden_chunk, - local_weight, - effective_a, - effective_b, - ) - exact_full = _rank_order_vocab_all_gather(exact_local, group) - score_logits = _temperature_scale_bf16_logits( - exact_full, - None if temperature is None else temperature[row_slice].contiguous(), - ) - support = exact_sampling_support( - score_logits, - top_ks[row_slice], - top_ps[row_slice], - min_ps[row_slice], - ) - identity_rows = exact_sampling_identity_rows( - top_ks[row_slice], - top_ps[row_slice], - min_ps[row_slice], - vocab_size=score_logits.shape[1], - ) - local_reference = ( - ( - F.linear(hidden_chunk, local_weight) - + F.linear( - F.linear(hidden_chunk.float(), effective_a.float()), - effective_b.float(), - ).to(torch.bfloat16) - ) - .float() - .contiguous() - ) - full_reference = _rank_order_vocab_all_gather(local_reference.to(torch.bfloat16), group).float() - full_grad = _selected_logprob_reference_grad_partitioned( - full_reference, - token_ids[row_slice], - grad_logprob[row_slice], - None if temperature is None else temperature[row_slice], - support, - identity_rows, - ) - local_grad[row_slice] = full_grad[:, self.shard.vocab_start : self.shard.vocab_end] + def _exact_score_logits(hidden_chunk: Tensor, temperature_chunk: Tensor | None) -> Tensor: + exact_local = self._exact_local_logits( + hidden_chunk.contiguous(), + local_weight, + effective_a, + effective_b, + ) + exact_full = _rank_order_vocab_all_gather(exact_local, group) + return _temperature_scale_bf16_logits( + exact_full, + None if temperature_chunk is None else temperature_chunk.contiguous(), + ) + + local_grad = filtered_surrogate_local_grad_logits( + hidden, + token_ids, + grad_logprob, + temperature, + sampling_transforms, + exact_score_logits_fn=_exact_score_logits, + reference_full_logits_fn=self._reference_full_logits_fn(local_weight, effective_a, effective_b, group), + local_vocab_slice=slice(self.shard.vocab_start, self.shard.vocab_end), + local_vocab_size=DSV4_LM_HEAD_LOCAL_VOCAB_SIZE, + ) grad_hidden, grad_a, grad_b = _local_surrogate_vjp( hidden, local_weight, @@ -686,8 +477,8 @@ def distributed_selected_logprob( local_token_ids, local_temperature, ) - self._validate_tp_group() - return _Dsv4ExactDistributedHeadFunction.apply( + row_plan = _distributed_row_plan(local_hidden, self._validate_tp_group(), self.source_ordinal) + return ExactLmHeadFunction.apply( local_hidden, local_weight, lora_a, @@ -695,6 +486,7 @@ def distributed_selected_logprob( local_token_ids, local_temperature, local_sampling_transforms, + row_plan, self, ) diff --git a/src/xorl/models/transformers/exact_lm_head_shared.py b/src/xorl/models/transformers/exact_lm_head_shared.py new file mode 100644 index 00000000..de6bdbd3 --- /dev/null +++ b/src/xorl/models/transformers/exact_lm_head_shared.py @@ -0,0 +1,550 @@ +"""Shared machinery for the exact TP-sharded LM-head value programs. + +The GLM-5.2 and DSV4 exact heads run the same program skeleton — local +base+LoRA logits through the literal serving SGEMMs, a rank-order vocabulary +all-gather, native or filtered selected scoring, and a straight-through +surrogate VJP over rank-local training rows. They pin different byte +contracts (FP32 vs BF16 logit buffers, ``head_v2`` vs plain base GEMM, rank-r +vs rank-1 factors, equal vs ragged row ownership). This module owns the +skeleton; each family injects its contract-bearing pieces as closures, +constants, and an :class:`ExactHeadRowPlan`. + +Nothing here computes a family's value bytes on its own: the base projection, +the temperature dtype boundary, and the native selected-score kernel always +come from the family module. +""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from functools import lru_cache +from typing import Any + +import torch +import torch.distributed as dist +from torch import Tensor + +from xorl.ops.exact_sampling_transforms import ( + EXACT_FILTER_ROW_CHUNK, + exact_sampling_identity_rows, + exact_sampling_support, + selected_logprob_reference_grad, + selected_logprob_reference_grad_partitioned, +) + + +SamplingTransforms = tuple[Tensor | None, Tensor | None, Tensor | None] + + +# --------------------------------------------------------------------------- +# Serving-kernel plumbing +# --------------------------------------------------------------------------- + + +@lru_cache(maxsize=64) +def single_adapter_lora_batch_info(device_index: int, rows: int, rank: int = 1, scaling: float = 1.0) -> Any: + """One active exact adapter, in the metadata format used by serving.""" + + from sglang.srt.lora.utils import LoRABatchInfo # noqa: PLC0415 + + device = torch.device("cuda", device_index) + return LoRABatchInfo( + use_cuda_graph=False, + bs=1, + num_segments=1, + seg_indptr=torch.tensor([0, rows], dtype=torch.int32, device=device), + weight_indices=torch.zeros(1, dtype=torch.int32, device=device), + lora_ranks=torch.full((1,), rank, dtype=torch.int32, device=device), + scalings=torch.full((1,), scaling, dtype=torch.float32, device=device), + max_len=rows, + seg_lens=torch.tensor([rows], dtype=torch.int32, device=device), + permutation=None, + expected_tokens=rows, + has_active_lora=True, + ) + + +def exact_lora_local_logits( + hidden_2d: Tensor, + effective_a: Tensor, + effective_b: Tensor, + *, + base_logits: Tensor, + rank: int, + scaling: float = 1.0, +) -> Tensor: + """Run the literal serving A/B SGEMMs with the fused add into ``base_logits``. + + The caller computes ``base_logits`` through its family's pinned base + projection (FP32 ``head_v2`` for GLM-5.2, BF16 ``F.linear`` for DSV4); this + helper owns only the shared LoRA-kernel choreography and its invariants. + """ + + try: + from sglang.kernels.ops.gemm.sgemm_lora_a import sgemm_lora_a_fwd # noqa: PLC0415 + from sglang.kernels.ops.gemm.sgemm_lora_b import sgemm_lora_b_fwd # noqa: PLC0415 + except Exception as exc: + raise RuntimeError("Pinned exact LoRA SGEMM kernels are required") from exc + + rows = hidden_2d.shape[0] + batch_info = single_adapter_lora_batch_info(hidden_2d.device.index, rows, rank, scaling) + lora_a_output = sgemm_lora_a_fwd(hidden_2d, effective_a.unsqueeze(0), batch_info) + if lora_a_output.dtype is not torch.bfloat16 or tuple(lora_a_output.shape) != (rows, rank): + raise RuntimeError("The pinned exact A SGEMM did not produce the required BF16 rank store") + output = sgemm_lora_b_fwd( + lora_a_output, + effective_b.unsqueeze(0), + batch_info, + base_output=base_logits, + ) + if output.data_ptr() != base_logits.data_ptr(): + raise RuntimeError("The pinned exact B SGEMM did not perform the required in-place base+delta store") + if output.dtype is not base_logits.dtype or not output.is_contiguous(): + raise RuntimeError("The pinned exact B SGEMM did not preserve the contiguous base-logit buffer") + return output + + +# --------------------------------------------------------------------------- +# Rank-order collectives +# --------------------------------------------------------------------------- + + +def rank_order_vocab_from_stacked( + stacked_logits: Tensor, + *, + expected_world_size: int, + expected_local_vocab_size: int, + expected_dtype: torch.dtype = torch.float32, +) -> Tensor: + """Apply serving's ``[rank,row,vocab] -> [row,rank*vocab]`` order.""" + + if stacked_logits.ndim != 3 or stacked_logits.shape[0] != expected_world_size: + raise ValueError( + "Rank-stacked LM-head logits must be [world, rows, local_vocab], got " + f"{tuple(stacked_logits.shape)} for world={expected_world_size}" + ) + if stacked_logits.shape[-1] != expected_local_vocab_size: + raise ValueError( + f"Rank-stacked local vocabulary width must be {expected_local_vocab_size}, got {stacked_logits.shape[-1]}" + ) + if stacked_logits.dtype is not expected_dtype: + raise TypeError(f"Rank-stacked LM-head logits must be {expected_dtype}, got {stacked_logits.dtype}") + if not stacked_logits.is_contiguous(): + raise ValueError("Rank-stacked LM-head logits must be contiguous in collective rank order") + rows = stacked_logits.shape[1] + return stacked_logits.permute(1, 0, 2).reshape(rows, expected_world_size * expected_local_vocab_size) + + +def rank_order_vocab_all_gather( + local_logits: Tensor, + group: dist.ProcessGroup, + *, + expected_world_size: int, + expected_local_vocab_size: int, + expected_dtype: torch.dtype = torch.float32, +) -> Tensor: + """Match serving's concat-style ``all_gather(..., dim=-1)`` byte order.""" + + if not dist.is_initialized(): + raise RuntimeError("Rank-order LM-head all-gather requires initialized torch.distributed") + if dist.get_world_size(group) != expected_world_size: + raise RuntimeError( + f"LM-head all-gather expected world size {expected_world_size}, got {dist.get_world_size(group)}" + ) + if local_logits.ndim != 2 or local_logits.shape[1] != expected_local_vocab_size: + raise ValueError( + f"Local LM-head logits must be [rows, {expected_local_vocab_size}], got {tuple(local_logits.shape)}" + ) + if local_logits.dtype is not expected_dtype or not local_logits.is_contiguous(): + raise ValueError(f"Local LM-head logits must be contiguous {expected_dtype} before the rank-order gather") + + rows = local_logits.shape[0] + gathered = torch.empty( + (expected_world_size * rows, expected_local_vocab_size), + dtype=local_logits.dtype, + device=local_logits.device, + ) + dist.all_gather_into_tensor(gathered, local_logits, group=group) + return rank_order_vocab_from_stacked( + gathered.view(expected_world_size, rows, expected_local_vocab_size), + expected_world_size=expected_world_size, + expected_local_vocab_size=expected_local_vocab_size, + expected_dtype=expected_dtype, + ) + + +def rank_order_row_all_gather(value: Tensor, group: dist.ProcessGroup) -> Tensor: + """Gather equal-shaped local row blocks in process-group rank order.""" + + if not dist.is_initialized(): + raise RuntimeError("Rank-order LM-head row gathering requires initialized torch.distributed") + if value.ndim == 0: + raise ValueError("Rank-order LM-head row gathering requires at least one dimension") + if not value.is_contiguous(): + raise ValueError("Rank-order LM-head row gathering requires a contiguous tensor") + world_size = dist.get_world_size(group) + gathered = torch.empty( + (world_size * value.shape[0], *value.shape[1:]), + dtype=value.dtype, + device=value.device, + ) + dist.all_gather_into_tensor(gathered, value, group=group) + return gathered + + +def rank_order_row_counts( + local_rows: int, + device: torch.device, + group: dist.ProcessGroup, + *, + world_size: int, + program: str, +) -> tuple[int, ...]: + """Exchange per-rank row counts in process-group rank order.""" + + local = torch.tensor([local_rows], dtype=torch.int64, device=device) + gathered = torch.empty(world_size, dtype=torch.int64, device=device) + dist.all_gather_into_tensor(gathered, local, group=group) + counts = tuple(int(value) for value in gathered.cpu().tolist()) + if any(value < 0 for value in counts): + raise RuntimeError(f"{program} received negative row counts: {counts}") + return counts + + +def rank_order_variable_row_all_gather( + value: Tensor, + group: dist.ProcessGroup, + *, + world_size: int, + row_counts: tuple[int, ...], + padded_rows: int, +) -> Tensor: + """Gather ragged local row blocks (padded to ``padded_rows``) in rank order.""" + + if value.ndim == 0 or not value.is_contiguous(): + raise ValueError("Variable-row all-gather requires a contiguous non-scalar tensor") + if value.shape[0] > padded_rows or len(row_counts) != world_size: + raise ValueError(f"Invalid variable-row geometry: local={value.shape[0]}, counts={row_counts}") + if value.shape[0] < padded_rows: + padding = value.new_zeros((padded_rows - value.shape[0], *value.shape[1:])) + value = torch.cat((value, padding), dim=0) + gathered = torch.empty( + (world_size * padded_rows, *value.shape[1:]), + dtype=value.dtype, + device=value.device, + ) + dist.all_gather_into_tensor(gathered, value, group=group) + pieces = [ + gathered[rank * padded_rows : rank * padded_rows + count] for rank, count in enumerate(row_counts) if count + ] + return torch.cat(pieces, dim=0) if pieces else gathered[:0] + + +def require_equal_nonzero_row_count(value: Tensor, group: dist.ProcessGroup, *, program: str) -> None: + """Fail before payload collectives when source-row shapes diverge.""" + + local_rows = torch.tensor([value.shape[0]], dtype=torch.int64, device=value.device) + world_size = dist.get_world_size(group) + gathered_rows = torch.empty(world_size, dtype=torch.int64, device=value.device) + dist.all_gather_into_tensor(gathered_rows, local_rows, group=group) + if bool((gathered_rows <= 0).any().item()): + raise ValueError(f"{program} requires at least one source row on every rank") + if bool((gathered_rows != gathered_rows[0]).any().item()): + raise ValueError( + f"{program} requires equal source-row counts across the group, got {gathered_rows.cpu().tolist()}" + ) + + +def all_reduce_sum_fp32(value: Tensor, group: dist.ProcessGroup) -> Tensor: + """In-place logical-owner sum used by the hidden/factor surrogate gradients.""" + + if value.dtype is not torch.float32 or not value.is_contiguous(): + raise ValueError("Exact LM-head surrogate reductions require contiguous FP32 tensors") + dist.all_reduce(value, op=dist.ReduceOp.SUM, group=group) + return value + + +def check_exact_head_tp_group( + *, + program: str, + world_size: int, + group_rank: int, + global_rank: int, + group_ranks: tuple[int, ...], + backend: str, + expected_world_size: int, + expected_ranks: Sequence[int], + shard_rank: int, + source_ordinal: int | None = None, +) -> None: + """Validate one family's TP group geometry against its declared contract. + + Pure checks over already-queried values, so family modules keep owning + their ``torch.distributed`` calls (and their tests' patch points). + """ + + expected_ranks = tuple(int(rank) for rank in expected_ranks) + if world_size != expected_world_size: + raise RuntimeError(f"{program} requires TP{expected_world_size}, got TP{world_size}") + if group_ranks != expected_ranks: + raise RuntimeError( + f"{program} gather order must match its expected TP{expected_world_size} group; " + f"expected {expected_ranks}, got {group_ranks}" + ) + if ( + group_rank != shard_rank + or (source_ordinal is not None and group_rank != source_ordinal) + or global_rank != group_ranks[shard_rank] + ): + raise RuntimeError( + f"{program} shard/group rank mismatch: " + f"shard_rank={shard_rank}, group_rank={group_rank}, global_rank={global_rank}" + ) + if backend != "nccl": + raise RuntimeError(f"{program} production group must use NCCL, got {backend}") + + +# --------------------------------------------------------------------------- +# The one exact-head autograd boundary +# --------------------------------------------------------------------------- + + +class ExactHeadRowPlan: + """How this rank's training rows map onto the head's replicated row set. + + ``gather`` assembles every contributor's rows in rank order (an identity + for already-replicated rows); ``narrow_local`` returns this rank's slice + of a row-aligned result. + """ + + __slots__ = ("gather", "narrow_local") + + def __init__( + self, + gather: Callable[[Tensor], Tensor], + narrow_local: Callable[[Tensor], Tensor], + ) -> None: + self.gather = gather + self.narrow_local = narrow_local + + +REPLICATED_ROW_PLAN = ExactHeadRowPlan(lambda value: value, lambda value: value) + + +class ExactLmHeadFunction(torch.autograd.Function): + """Own the literal exact-head forward; delegate only the VJP to the family. + + The forward gathers this rank's rows per the row plan, runs the family + component's exact value program (native or sampling-filtered), and returns + the caller's local rows. The backward gathers the downstream rank-local + gradients, applies the family's declared straight-through surrogate VJP to + the identical global row block on every vocabulary rank, and returns only + the caller's hidden-state slice. + """ + + @staticmethod + def forward( + ctx, + hidden_states: Tensor, + local_weight: Tensor, + lora_a: Tensor, + local_lora_b: Tensor, + token_ids: Tensor, + temperature: Tensor | None, + sampling_transforms: SamplingTransforms, + row_plan: ExactHeadRowPlan, + component, + ) -> Tensor: + effective_a = lora_a.to(torch.bfloat16).contiguous() + effective_b = local_lora_b.to(torch.bfloat16).contiguous() + gathered_hidden = row_plan.gather(hidden_states) + gathered_token_ids = row_plan.gather(token_ids) + gathered_temperature = None if temperature is None else row_plan.gather(temperature) + gathered_transforms = tuple(None if value is None else row_plan.gather(value) for value in sampling_transforms) + has_sampling_filter = gathered_transforms[0] is not None + if has_sampling_filter: + gathered_logprob = component._exact_forward_value_filtered( + gathered_hidden, + local_weight, + effective_a, + effective_b, + gathered_token_ids, + gathered_temperature, + gathered_transforms, + ) + else: + gathered_logprob = component._exact_forward_value( + gathered_hidden, + local_weight, + effective_a, + effective_b, + gathered_token_ids, + gathered_temperature, + ) + local_logprob = row_plan.narrow_local(gathered_logprob).contiguous() + if local_logprob.dtype is not torch.float32 or tuple(local_logprob.shape) != tuple(token_ids.shape): + raise RuntimeError( + f"{type(component).__name__} returned an invalid selected-logprob tensor: " + f"dtype={local_logprob.dtype}, shape={tuple(local_logprob.shape)}, " + f"expected FP32 {tuple(token_ids.shape)}" + ) + ctx.set_materialize_grads(False) + ctx.component = component + ctx.row_plan = row_plan + ctx.has_sampling_filter = has_sampling_filter + # local_weight and the masters are also version-counter sentinels. + ctx.save_for_backward( + gathered_hidden.detach(), + local_weight, + effective_a, + effective_b, + gathered_token_ids, + gathered_temperature + if gathered_temperature is not None + else torch.empty((0,), dtype=torch.float32, device=hidden_states.device), + lora_a, + local_lora_b, + *gathered_transforms, + ) + return local_logprob + + @staticmethod + def backward(ctx, grad_local_logprob: Tensor | None): + ( + gathered_hidden, + local_weight, + effective_a, + effective_b, + gathered_token_ids, + stored_temperature, + _lora_a_master, + _local_lora_b_master, + top_ks, + top_ps, + min_ps, + ) = ctx.saved_tensors + if grad_local_logprob is None: + return (None,) * 9 + temperature = None if stored_temperature.numel() == 0 else stored_temperature + gathered_grad_logprob = ctx.row_plan.gather(grad_local_logprob.float().contiguous()) + vjp = ctx.component._surrogate_vjp_filtered if ctx.has_sampling_filter else ctx.component._surrogate_vjp + args = ( + gathered_hidden, + local_weight, + effective_a, + effective_b, + gathered_token_ids, + gathered_grad_logprob, + temperature, + ) + if ctx.has_sampling_filter: + args = (*args, (top_ks, top_ps, min_ps)) + grad_hidden, grad_a, grad_b = vjp( + *args, + needs_input_grad=(ctx.needs_input_grad[0], ctx.needs_input_grad[2], ctx.needs_input_grad[3]), + ) + if grad_hidden is not None: + grad_hidden = ctx.row_plan.narrow_local(grad_hidden).contiguous() + return (grad_hidden, None, grad_a, grad_b, None, None, None, None, None) + + +# --------------------------------------------------------------------------- +# The shared surrogate-VJP plumbing +# --------------------------------------------------------------------------- + + +def surrogate_local_grad_logits( + hidden_2d: Tensor, + token_ids_1d: Tensor, + grad_logprob_1d: Tensor, + temperature_1d: Tensor | None, + *, + reference_full_logits_fn: Callable[[Tensor], Tensor], + local_vocab_slice: slice, +) -> Tensor: + """Local-shard reference dlogits for the unfiltered surrogate VJP.""" + + with torch.no_grad(), torch.autocast(device_type=hidden_2d.device.type, enabled=False): + full_reference_logits = reference_full_logits_fn(hidden_2d) + full_grad_logits = selected_logprob_reference_grad( + full_reference_logits, + token_ids_1d, + grad_logprob_1d, + temperature_1d, + ) + return full_grad_logits[:, local_vocab_slice].contiguous() + + +def filtered_surrogate_local_grad_logits( + hidden_2d: Tensor, + token_ids_1d: Tensor, + grad_logprob_1d: Tensor, + temperature_1d: Tensor | None, + sampling_transforms: SamplingTransforms, + *, + exact_score_logits_fn: Callable[[Tensor, Tensor | None], Tensor], + reference_full_logits_fn: Callable[[Tensor], Tensor], + local_vocab_slice: slice, + local_vocab_size: int, + row_chunk: int = EXACT_FILTER_ROW_CHUNK, +) -> Tensor: + """Local-shard reference dlogits over the exact support, in bounded chunks. + + Per chunk: the family recreates its literal value logits (post-gather, + post-temperature) to derive the support, recreates its differentiable + reference logits, and the shared partitioned reference grad is sliced to + the local vocabulary columns. The dense support mask is transient and + never saved on an autograd context. + """ + + top_ks, top_ps, min_ps = sampling_transforms + if top_ks is None or top_ps is None or min_ps is None: + raise ValueError("filtered exact scoring requires complete row metadata") + rows = hidden_2d.shape[0] + local_grad_logits = torch.empty( + (rows, local_vocab_size), + dtype=torch.float32, + device=hidden_2d.device, + ) + for start in range(0, rows, row_chunk): + end = min(start + row_chunk, rows) + hidden_chunk = hidden_2d[start:end] + temperature_chunk = None if temperature_1d is None else temperature_1d[start:end] + chunk_transforms = (top_ks[start:end], top_ps[start:end], min_ps[start:end]) + with torch.no_grad(), torch.autocast(device_type=hidden_2d.device.type, enabled=False): + exact_score_logits = exact_score_logits_fn(hidden_chunk, temperature_chunk) + support = exact_sampling_support(exact_score_logits, *chunk_transforms) + identity_rows = exact_sampling_identity_rows( + *chunk_transforms, + vocab_size=exact_score_logits.shape[1], + ) + full_reference_logits = reference_full_logits_fn(hidden_chunk) + full_grad_logits = selected_logprob_reference_grad_partitioned( + full_reference_logits, + token_ids_1d[start:end], + grad_logprob_1d[start:end], + temperature_chunk, + support, + identity_rows, + ) + local_grad_logits[start:end] = full_grad_logits[:, local_vocab_slice] + return local_grad_logits + + +__all__ = [ + "REPLICATED_ROW_PLAN", + "ExactHeadRowPlan", + "ExactLmHeadFunction", + "all_reduce_sum_fp32", + "check_exact_head_tp_group", + "exact_lora_local_logits", + "filtered_surrogate_local_grad_logits", + "rank_order_row_all_gather", + "rank_order_row_counts", + "rank_order_variable_row_all_gather", + "rank_order_vocab_all_gather", + "rank_order_vocab_from_stacked", + "require_equal_nonzero_row_count", + "single_adapter_lora_batch_info", + "surrogate_local_grad_logits", +] diff --git a/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py b/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py index 53c4694d..74381d8d 100644 --- a/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py +++ b/src/xorl/models/transformers/glm5/exact_lm_head_qlora.py @@ -22,8 +22,6 @@ from __future__ import annotations from dataclasses import dataclass -from functools import lru_cache -from typing import Any import torch import torch.distributed as dist @@ -31,20 +29,30 @@ from torch import Tensor, nn from xorl.lora.modules.linear import LoraLinear +from xorl.models.transformers.exact_lm_head_shared import ( + REPLICATED_ROW_PLAN, + ExactHeadRowPlan, + ExactLmHeadFunction, + check_exact_head_tp_group, + exact_lora_local_logits, + filtered_surrogate_local_grad_logits, + require_equal_nonzero_row_count, + surrogate_local_grad_logits, +) +from xorl.models.transformers.exact_lm_head_shared import ( + all_reduce_sum_fp32 as _all_reduce_sum_fp32, +) +from xorl.models.transformers.exact_lm_head_shared import ( + rank_order_row_all_gather as _rank_order_row_all_gather, +) +from xorl.models.transformers.exact_lm_head_shared import ( + rank_order_vocab_all_gather as _rank_order_vocab_all_gather, +) from xorl.models.transformers.glm5.exact_lora_contract import glm52_exact_lora_scaling from xorl.ops.bi_families_v2 import exact_temperature_scale_fp32_logits from xorl.ops.exact_sampling_transforms import ( - EXACT_FILTER_ROW_CHUNK, - exact_sampling_identity_rows, - exact_sampling_support, - exact_selected_logprob_partitioned_from_support, normalize_temperature_rows, -) -from xorl.ops.exact_sampling_transforms import ( - selected_logprob_reference_grad as _selected_logprob_reference_grad, -) -from xorl.ops.exact_sampling_transforms import ( - selected_logprob_reference_grad_partitioned as _selected_logprob_reference_grad_partitioned, + score_with_sampling_transforms, ) from xorl.ops.exact_sampling_transforms import ( validate_temperature_rows as _validate_temperature_rows, @@ -131,133 +139,22 @@ def glm52_lm_head_shard(tp_rank: int) -> Glm52LmHeadShard: return shard -@lru_cache(maxsize=64) -def _single_adapter_lm_head_batch_info(device_index: int, rows: int, rank: int = 1, scaling: float = 1.0) -> Any: - """One active exact adapter, in the metadata format used by S4.""" - - from sglang.srt.lora.utils import LoRABatchInfo # noqa: PLC0415 - - device = torch.device("cuda", device_index) - return LoRABatchInfo( - use_cuda_graph=False, - bs=1, - num_segments=1, - seg_indptr=torch.tensor([0, rows], dtype=torch.int32, device=device), - weight_indices=torch.zeros(1, dtype=torch.int32, device=device), - lora_ranks=torch.full((1,), rank, dtype=torch.int32, device=device), - scalings=torch.full((1,), scaling, dtype=torch.float32, device=device), - max_len=rows, - seg_lens=torch.tensor([rows], dtype=torch.int32, device=device), - permutation=None, - expected_tokens=rows, - has_active_lora=True, - ) - - -def _rank_order_vocab_from_stacked( - stacked_logits: Tensor, - *, - expected_world_size: int, - expected_local_vocab_size: int, -) -> Tensor: - """Apply SGLang's ``[rank,row,vocab] -> [row,rank*vocab]`` order.""" - - expected_shape_tail = (expected_local_vocab_size,) - if stacked_logits.ndim != 3 or stacked_logits.shape[0] != expected_world_size: - raise ValueError( - "Rank-stacked LM-head logits must be [world, rows, local_vocab], got " - f"{tuple(stacked_logits.shape)} for world={expected_world_size}" - ) - if tuple(stacked_logits.shape[-1:]) != expected_shape_tail: - raise ValueError( - f"Rank-stacked local vocabulary width must be {expected_local_vocab_size}, got {stacked_logits.shape[-1]}" - ) - if stacked_logits.dtype is not torch.float32: - raise TypeError(f"Rank-stacked LM-head logits must be FP32, got {stacked_logits.dtype}") - if not stacked_logits.is_contiguous(): - raise ValueError("Rank-stacked LM-head logits must be contiguous in collective rank order") - rows = stacked_logits.shape[1] - return stacked_logits.permute(1, 0, 2).reshape(rows, expected_world_size * expected_local_vocab_size) - - -def _rank_order_vocab_all_gather( - local_logits: Tensor, - tp_group: dist.ProcessGroup, - *, - expected_world_size: int, - expected_local_vocab_size: int, -) -> Tensor: - """Match S4's concat-style ``all_gather(..., dim=-1)`` byte order.""" - - if not dist.is_initialized(): - raise RuntimeError("Rank-order LM-head all-gather requires initialized torch.distributed") - if dist.get_world_size(tp_group) != expected_world_size: - raise RuntimeError( - f"LM-head all-gather expected world size {expected_world_size}, got {dist.get_world_size(tp_group)}" - ) - if local_logits.ndim != 2 or tuple(local_logits.shape[1:]) != (expected_local_vocab_size,): - raise ValueError( - f"Local LM-head logits must be [rows, {expected_local_vocab_size}], got {tuple(local_logits.shape)}" - ) - if local_logits.dtype is not torch.float32 or not local_logits.is_contiguous(): - raise ValueError("Local LM-head logits must be contiguous FP32 before the rank-order gather") - - rows = local_logits.shape[0] - gathered = torch.empty( - (expected_world_size * rows, expected_local_vocab_size), - dtype=torch.float32, - device=local_logits.device, - ) - dist.all_gather_into_tensor(gathered, local_logits, group=tp_group) - return _rank_order_vocab_from_stacked( - gathered.view(expected_world_size, rows, expected_local_vocab_size), - expected_world_size=expected_world_size, - expected_local_vocab_size=expected_local_vocab_size, - ) - - -def _all_reduce_sum_fp32(value: Tensor, group: dist.ProcessGroup) -> Tensor: - """In-place logical-owner sum used by the hidden/A surrogate gradients.""" +def _require_equal_nonzero_row_count(value: Tensor, group: dist.ProcessGroup) -> None: + """Fail before payload collectives when TP16 source-row shapes diverge.""" - if value.dtype is not torch.float32 or not value.is_contiguous(): - raise ValueError("GLM-5.2 LM-head surrogate reductions require contiguous FP32 tensors") - dist.all_reduce(value, op=dist.ReduceOp.SUM, group=group) - return value + require_equal_nonzero_row_count(value, group, program="The exact GLM-5.2 lm head") -def _rank_order_row_all_gather(value: Tensor, group: dist.ProcessGroup) -> Tensor: - """Gather equal-shaped local row blocks in process-group rank order.""" +def _distributed_row_plan(local_hidden: Tensor, group: dist.ProcessGroup) -> ExactHeadRowPlan: + """Equal rank-order row blocks: every TP16 rank contributes the same count.""" - if not dist.is_initialized(): - raise RuntimeError("Rank-order LM-head row gathering requires initialized torch.distributed") - if value.ndim == 0: - raise ValueError("Rank-order LM-head row gathering requires at least one dimension") - if not value.is_contiguous(): - raise ValueError("Rank-order LM-head row gathering requires a contiguous tensor") - world_size = dist.get_world_size(group) - gathered = torch.empty( - (world_size * value.shape[0], *value.shape[1:]), - dtype=value.dtype, - device=value.device, + _require_equal_nonzero_row_count(local_hidden, group) + local_rows = local_hidden.shape[0] + source_rank = dist.get_rank(group) + return ExactHeadRowPlan( + lambda value: _rank_order_row_all_gather(value, group), + lambda value: value.narrow(0, source_rank * local_rows, local_rows), ) - dist.all_gather_into_tensor(gathered, value, group=group) - return gathered - - -def _require_equal_nonzero_row_count(value: Tensor, group: dist.ProcessGroup) -> None: - """Fail before payload collectives when TP16 source-row shapes diverge.""" - - local_rows = torch.tensor([value.shape[0]], dtype=torch.int64, device=value.device) - world_size = dist.get_world_size(group) - gathered_rows = torch.empty(world_size, dtype=torch.int64, device=value.device) - dist.all_gather_into_tensor(gathered_rows, local_rows, group=group) - if bool((gathered_rows <= 0).any().item()): - raise ValueError("The exact GLM-5.2 lm head requires at least one source row on every TP16 rank") - if bool((gathered_rows != gathered_rows[0]).any().item()): - raise ValueError( - "The exact GLM-5.2 lm head requires equal source-row counts across TP16, got " - f"{gathered_rows.cpu().tolist()}" - ) def _local_qlora_surrogate_logits( @@ -338,229 +235,6 @@ def _local_qlora_surrogate_vjp( return grad_hidden, by_label.get("A"), by_label.get("B") -class _Glm52ExactTP16LmHeadFunction(torch.autograd.Function): - """Own the literal forward and delegate only the VJP to the QLoRA oracle.""" - - @staticmethod - def forward( - ctx, - hidden_states: Tensor, - local_weight: Tensor, - lora_A: Tensor, - local_lora_B: Tensor, - token_ids: Tensor, - temperature: Tensor | None, - sampling_transforms: tuple[Tensor | None, Tensor | None, Tensor | None], - component: Glm52ExactTP16LmHeadSelectedLogprob, - ) -> Tensor: - effective_A = lora_A.to(torch.bfloat16).contiguous() - effective_B = local_lora_B.to(torch.bfloat16).contiguous() - has_sampling_filter = sampling_transforms[0] is not None - if not has_sampling_filter: - logprob = component._exact_forward_value( - hidden_states, local_weight, effective_A, effective_B, token_ids, temperature - ) - else: - logprob = component._exact_forward_value_filtered( - hidden_states, - local_weight, - effective_A, - effective_B, - token_ids, - temperature, - sampling_transforms, - ) - if logprob.dtype is not torch.float32 or tuple(logprob.shape) != tuple(token_ids.shape): - raise RuntimeError( - "GLM-5.2 exact TP16 LM head returned an invalid selected-logprob tensor: " - f"dtype={logprob.dtype}, shape={tuple(logprob.shape)}, expected FP32 {tuple(token_ids.shape)}" - ) - ctx.set_materialize_grads(False) - ctx.component = component - ctx.has_sampling_filter = has_sampling_filter - # local_weight and the masters are also version-counter sentinels. - ctx.save_for_backward( - hidden_states.detach(), - local_weight, - effective_A, - effective_B, - token_ids, - temperature - if temperature is not None - else torch.empty((0,), dtype=torch.float32, device=hidden_states.device), - lora_A, - local_lora_B, - *sampling_transforms, - ) - return logprob - - @staticmethod - def backward(ctx, grad_logprob: Tensor | None): - ( - hidden_states, - local_weight, - effective_A, - effective_B, - token_ids, - stored_temperature, - _lora_A_master, - _local_lora_B_master, - top_ks, - top_ps, - min_ps, - ) = ctx.saved_tensors - if grad_logprob is None: - return (None, None, None, None, None, None, None, None) - temperature = None if stored_temperature.numel() == 0 else stored_temperature - vjp = ctx.component._surrogate_vjp_filtered if ctx.has_sampling_filter else ctx.component._surrogate_vjp - args = ( - hidden_states, - local_weight, - effective_A, - effective_B, - token_ids, - grad_logprob, - temperature, - ) - if ctx.has_sampling_filter: - args = (*args, (top_ks, top_ps, min_ps)) - grad_hidden, grad_A, grad_B = vjp( - *args, - needs_input_grad=(ctx.needs_input_grad[0], ctx.needs_input_grad[2], ctx.needs_input_grad[3]), - ) - return (grad_hidden, None, grad_A, grad_B, None, None, None, None) - - -class _Glm52ExactDistributedTP16LmHeadFunction(torch.autograd.Function): - """Bridge rank-local training rows to the replicated TP16 head program. - - The lm-head-only TP group is carved from CP or DP ranks, so its members may - own different token rows. Forward gathers those row blocks before running - the literal TP16 selected-logprob program. The function returns only the - caller's rank-local rows. Backward gathers the downstream rank-local RL - gradients, applies the component's declared hybrid VJP to the identical - global row block on every vocabulary rank, and returns only the caller's - hidden-state slice. - """ - - @staticmethod - def forward( - ctx, - local_hidden_states: Tensor, - local_weight: Tensor, - lora_A: Tensor, - local_lora_B: Tensor, - local_token_ids: Tensor, - local_temperature: Tensor | None, - local_sampling_transforms: tuple[Tensor | None, Tensor | None, Tensor | None], - component: Glm52ExactTP16LmHeadSelectedLogprob, - ) -> Tensor: - group = component._validate_tp_group() - _require_equal_nonzero_row_count(local_hidden_states, group) - effective_A = lora_A.to(torch.bfloat16).contiguous() - effective_B = local_lora_B.to(torch.bfloat16).contiguous() - gathered_hidden = _rank_order_row_all_gather(local_hidden_states, group) - gathered_token_ids = _rank_order_row_all_gather(local_token_ids, group) - gathered_temperature = ( - None if local_temperature is None else _rank_order_row_all_gather(local_temperature, group) - ) - gathered_sampling_transforms = tuple( - None if value is None else _rank_order_row_all_gather(value, group) for value in local_sampling_transforms - ) - has_sampling_filter = gathered_sampling_transforms[0] is not None - if not has_sampling_filter: - gathered_logprob = component._exact_forward_value( - gathered_hidden, - local_weight, - effective_A, - effective_B, - gathered_token_ids, - gathered_temperature, - ) - else: - gathered_logprob = component._exact_forward_value_filtered( - gathered_hidden, - local_weight, - effective_A, - effective_B, - gathered_token_ids, - gathered_temperature, - gathered_sampling_transforms, - ) - local_rows = local_hidden_states.shape[0] - source_rank = dist.get_rank(group) - local_logprob = gathered_logprob.narrow(0, source_rank * local_rows, local_rows).contiguous() - if local_logprob.dtype is not torch.float32 or tuple(local_logprob.shape) != tuple(local_token_ids.shape): - raise RuntimeError( - "GLM-5.2 distributed exact LM head returned an invalid local logprob tensor: " - f"dtype={local_logprob.dtype}, shape={tuple(local_logprob.shape)}" - ) - - ctx.set_materialize_grads(False) - ctx.component = component - ctx.local_rows = local_rows - ctx.source_rank = source_rank - ctx.has_sampling_filter = has_sampling_filter - ctx.save_for_backward( - gathered_hidden, - local_weight, - effective_A, - effective_B, - gathered_token_ids, - gathered_temperature - if gathered_temperature is not None - else torch.empty((0,), dtype=torch.float32, device=local_hidden_states.device), - lora_A, - local_lora_B, - *gathered_sampling_transforms, - ) - return local_logprob - - @staticmethod - def backward(ctx, grad_local_logprob: Tensor | None): - ( - gathered_hidden, - local_weight, - effective_A, - effective_B, - gathered_token_ids, - stored_temperature, - _lora_A_master, - _local_lora_B_master, - top_ks, - top_ps, - min_ps, - ) = ctx.saved_tensors - if grad_local_logprob is None: - return (None, None, None, None, None, None, None, None) - temperature = None if stored_temperature.numel() == 0 else stored_temperature - group = ctx.component._validate_tp_group() - gathered_grad_logprob = _rank_order_row_all_gather(grad_local_logprob.float().contiguous(), group) - vjp = ctx.component._surrogate_vjp_filtered if ctx.has_sampling_filter else ctx.component._surrogate_vjp - args = ( - gathered_hidden, - local_weight, - effective_A, - effective_B, - gathered_token_ids, - gathered_grad_logprob, - temperature, - ) - if ctx.has_sampling_filter: - args = (*args, (top_ks, top_ps, min_ps)) - grad_hidden, grad_A, grad_B = vjp( - *args, - needs_input_grad=(ctx.needs_input_grad[0], ctx.needs_input_grad[2], ctx.needs_input_grad[3]), - ) - if grad_hidden is not None: - grad_hidden = grad_hidden.narrow( - 0, - ctx.source_rank * ctx.local_rows, - ctx.local_rows, - ).contiguous() - return (grad_hidden, None, grad_A, grad_B, None, None, None, None) - - class Glm52ExactTP16LmHeadLoraLinear(LoraLinear): """Logical full head whose only admitted value path is selected logprob.""" @@ -652,24 +326,17 @@ def _validate_tp_group(self) -> dist.ProcessGroup: raise RuntimeError("GLM-5.2 exact LM head has no bound TP process group") if not dist.is_initialized(): raise RuntimeError("GLM-5.2 exact LM head requires initialized torch.distributed") - world_size = dist.get_world_size(group) - group_rank = dist.get_rank(group) - group_ranks = tuple(dist.get_process_group_ranks(group)) - if world_size != GLM52_LM_HEAD_TP_SIZE: - raise RuntimeError(f"GLM-5.2 exact LM head requires TP16, got TP{world_size}") - if group_ranks != self.expected_group_ranks: - raise RuntimeError( - "GLM-5.2 exact LM-head gather order must match its stage-local TP16 group; " - f"expected {self.expected_group_ranks}, got {group_ranks}" - ) - if group_rank != self.shard.tp_rank or dist.get_rank() != group_ranks[self.shard.tp_rank]: - raise RuntimeError( - "GLM-5.2 exact LM-head shard/group rank mismatch: " - f"shard_rank={self.shard.tp_rank}, group_rank={group_rank}, global_rank={dist.get_rank()}" - ) - backend = str(dist.get_backend(group)).lower() - if backend != "nccl": - raise RuntimeError(f"GLM-5.2 exact LM-head production group must use NCCL, got {backend}") + check_exact_head_tp_group( + program="GLM-5.2 exact LM-head", + world_size=dist.get_world_size(group), + group_rank=dist.get_rank(group), + global_rank=dist.get_rank(), + group_ranks=tuple(dist.get_process_group_ranks(group)), + backend=str(dist.get_backend(group)).lower(), + expected_world_size=GLM52_LM_HEAD_TP_SIZE, + expected_ranks=self.expected_group_ranks, + shard_rank=self.shard.tp_rank, + ) return group @staticmethod @@ -813,8 +480,6 @@ def _exact_local_logits( raise ValueError("Exact local LM-head operands must be contiguous") try: - from sglang.kernels.ops.gemm.sgemm_lora_a import sgemm_lora_a_fwd # noqa: PLC0415 - from sglang.kernels.ops.gemm.sgemm_lora_b import sgemm_lora_b_fwd # noqa: PLC0415 from sglang.srt.batch_invariant_ops import head_v2_full_logits_with_lse # noqa: PLC0415 except Exception as exc: raise RuntimeError("Pinned exact v2 batch-invariant and LoRA kernels are required") from exc @@ -826,24 +491,14 @@ def _exact_local_logits( GLM52_LM_HEAD_LOCAL_VOCAB_SIZE, ): raise RuntimeError("Pinned exact v2 base LM-head kernel returned an invalid local-logit buffer") - batch_info = _single_adapter_lm_head_batch_info(hidden_2d.device.index, rows, self.max_lora_rank, self.scaling) - lora_a_output = sgemm_lora_a_fwd(hidden_2d, effective_A.unsqueeze(0), batch_info) - if lora_a_output.dtype is not torch.bfloat16 or tuple(lora_a_output.shape) != ( - rows, - self.max_lora_rank, - ): - raise RuntimeError("Pinned exact A SGEMM did not produce the required BF16 rank store") - output = sgemm_lora_b_fwd( - lora_a_output, - effective_B.unsqueeze(0), - batch_info, - base_output=base_logits, + return exact_lora_local_logits( + hidden_2d, + effective_A, + effective_B, + base_logits=base_logits, + rank=self.max_lora_rank, + scaling=self.scaling, ) - if output.data_ptr() != base_logits.data_ptr(): - raise RuntimeError("Pinned exact B SGEMM did not perform the required in-place base+delta store") - if output.dtype is not torch.float32 or not output.is_contiguous(): - raise RuntimeError("Pinned exact B SGEMM did not preserve the contiguous FP32 local-logit buffer") - return output @staticmethod def _selected_logprob_from_gathered( @@ -885,29 +540,18 @@ def _selected_logprob_from_gathered_filtered( top_ks, top_ps, min_ps = sampling_transforms if top_ks is None or top_ps is None or min_ps is None: raise ValueError("filtered GLM exact scoring requires complete row metadata") - support = exact_sampling_support( - score_logits, - top_ks, - top_ps, - min_ps, - ) - identity_rows = exact_sampling_identity_rows( - top_ks, - top_ps, - min_ps, - vocab_size=score_logits.shape[1], - ) try: from sglang.srt.batch_invariant_ops import ( # noqa: PLC0415 head_v2_selected_logprob_from_logits, ) except Exception as exc: raise RuntimeError("Pinned exact v2 selected-logprob tail is required") from exc - logprob, _, _ = exact_selected_logprob_partitioned_from_support( + logprob, _, _ = score_with_sampling_transforms( score_logits, token_ids, - support, - identity_rows, + top_ks, + top_ps, + min_ps, lambda native_logits, native_ids: head_v2_selected_logprob_from_logits( native_logits, native_ids, @@ -975,6 +619,26 @@ def _exact_forward_value_filtered( ) return logprob.view_as(token_ids) + def _reference_full_logits_fn(self, local_weight, effective_A, effective_B, group): + """The differentiable QLoRA reference, gathered in serving byte order.""" + + def reference_full_logits(hidden_chunk: Tensor) -> Tensor: + local_reference_logits = _local_qlora_surrogate_logits( + hidden_chunk, + local_weight, + effective_A, + effective_B, + self.scaling, + ).contiguous() + return _rank_order_vocab_all_gather( + local_reference_logits, + group, + expected_world_size=GLM52_LM_HEAD_TP_SIZE, + expected_local_vocab_size=GLM52_LM_HEAD_LOCAL_VOCAB_SIZE, + ) + + return reference_full_logits + def _surrogate_vjp( self, hidden_states: Tensor, @@ -996,27 +660,14 @@ def _surrogate_vjp( rows=rows, device=hidden_states.device, ) - with torch.no_grad(), torch.autocast(device_type=hidden_states.device.type, enabled=False): - local_reference_logits = _local_qlora_surrogate_logits( - hidden_2d, - local_weight, - effective_A, - effective_B, - self.scaling, - ).contiguous() - full_reference_logits = _rank_order_vocab_all_gather( - local_reference_logits, - group, - expected_world_size=GLM52_LM_HEAD_TP_SIZE, - expected_local_vocab_size=GLM52_LM_HEAD_LOCAL_VOCAB_SIZE, - ) - full_grad_logits = _selected_logprob_reference_grad( - full_reference_logits, + local_grad_logits = surrogate_local_grad_logits( + hidden_2d, token_ids_1d, grad_logprob.reshape(rows), temperature_1d, + reference_full_logits_fn=self._reference_full_logits_fn(local_weight, effective_A, effective_B, group), + local_vocab_slice=slice(self.shard.vocab_start, self.shard.vocab_end), ) - local_grad_logits = full_grad_logits[:, self.shard.vocab_start : self.shard.vocab_end].contiguous() grad_hidden, grad_A, grad_B = _local_qlora_surrogate_vjp( hidden_2d, local_weight, @@ -1050,65 +701,30 @@ def _surrogate_vjp_filtered( hidden_2d = hidden_states.view(rows, GLM52_LM_HEAD_HIDDEN_SIZE) token_ids_1d = token_ids.view(rows) temperature_1d = _validate_temperature_rows(temperature, rows=rows, device=hidden_states.device) - top_ks, top_ps, min_ps = sampling_transforms - if top_ks is None or top_ps is None or min_ps is None: - raise ValueError("filtered GLM exact VJP requires complete row metadata") - grad_logprob_1d = grad_logprob.reshape(rows) - local_grad_logits = torch.empty( - (rows, GLM52_LM_HEAD_LOCAL_VOCAB_SIZE), - dtype=torch.float32, - device=hidden_states.device, - ) - for start in range(0, rows, EXACT_FILTER_ROW_CHUNK): - end = min(start + EXACT_FILTER_ROW_CHUNK, rows) - hidden_chunk = hidden_2d[start:end] - temperature_chunk = None if temperature_1d is None else temperature_1d[start:end] - chunk_transforms = (top_ks[start:end], top_ps[start:end], min_ps[start:end]) - with torch.no_grad(), torch.autocast(device_type=hidden_states.device.type, enabled=False): - exact_local_logits = self._exact_local_logits( - hidden_chunk, - local_weight, - effective_A, - effective_B, - ) - exact_full_logits = _rank_order_vocab_all_gather( - exact_local_logits, - group, - expected_world_size=GLM52_LM_HEAD_TP_SIZE, - expected_local_vocab_size=GLM52_LM_HEAD_LOCAL_VOCAB_SIZE, - ) - exact_score_logits = ( - exact_full_logits - if temperature_chunk is None - else exact_temperature_scale_fp32_logits(exact_full_logits, temperature_chunk) - ) - support = exact_sampling_support(exact_score_logits, *chunk_transforms) - identity_rows = exact_sampling_identity_rows( - *chunk_transforms, - vocab_size=exact_score_logits.shape[1], - ) - local_reference_logits = _local_qlora_surrogate_logits( - hidden_chunk, - local_weight, - effective_A, - effective_B, - self.scaling, - ).contiguous() - full_reference_logits = _rank_order_vocab_all_gather( - local_reference_logits, - group, - expected_world_size=GLM52_LM_HEAD_TP_SIZE, - expected_local_vocab_size=GLM52_LM_HEAD_LOCAL_VOCAB_SIZE, - ) - full_grad_logits = _selected_logprob_reference_grad_partitioned( - full_reference_logits, - token_ids_1d[start:end], - grad_logprob_1d[start:end], - temperature_chunk, - support, - identity_rows, + + def _exact_score_logits(hidden_chunk: Tensor, temperature_chunk: Tensor | None) -> Tensor: + exact_local_logits = self._exact_local_logits(hidden_chunk, local_weight, effective_A, effective_B) + exact_full_logits = _rank_order_vocab_all_gather( + exact_local_logits, + group, + expected_world_size=GLM52_LM_HEAD_TP_SIZE, + expected_local_vocab_size=GLM52_LM_HEAD_LOCAL_VOCAB_SIZE, ) - local_grad_logits[start:end] = full_grad_logits[:, self.shard.vocab_start : self.shard.vocab_end] + if temperature_chunk is None: + return exact_full_logits + return exact_temperature_scale_fp32_logits(exact_full_logits, temperature_chunk) + + local_grad_logits = filtered_surrogate_local_grad_logits( + hidden_2d, + token_ids_1d, + grad_logprob.reshape(rows), + temperature_1d, + sampling_transforms, + exact_score_logits_fn=_exact_score_logits, + reference_full_logits_fn=self._reference_full_logits_fn(local_weight, effective_A, effective_B, group), + local_vocab_slice=slice(self.shard.vocab_start, self.shard.vocab_end), + local_vocab_size=GLM52_LM_HEAD_LOCAL_VOCAB_SIZE, + ) grad_hidden, grad_A, grad_B = _local_qlora_surrogate_vjp( hidden_2d, local_weight, @@ -1144,7 +760,7 @@ def forward( require_cuda=True, ) self._validate_tp_group() - return _Glm52ExactTP16LmHeadFunction.apply( + return ExactLmHeadFunction.apply( hidden_states, local_weight, lora_A, @@ -1152,6 +768,7 @@ def forward( token_ids, temperature, sampling_transforms, + REPLICATED_ROW_PLAN, self, ) @@ -1176,8 +793,8 @@ def distributed_selected_logprob( local_temperature, require_cuda=True, ) - self._validate_tp_group() - return _Glm52ExactDistributedTP16LmHeadFunction.apply( + row_plan = _distributed_row_plan(local_hidden_states, self._validate_tp_group()) + return ExactLmHeadFunction.apply( local_hidden_states, local_weight, lora_A, @@ -1185,6 +802,7 @@ def distributed_selected_logprob( local_token_ids, local_temperature, local_sampling_transforms, + row_plan, self, ) diff --git a/src/xorl/ops/loss/bi_fused_lm_head.py b/src/xorl/ops/loss/bi_fused_lm_head.py index 8b90b54f..90583e85 100644 --- a/src/xorl/ops/loss/bi_fused_lm_head.py +++ b/src/xorl/ops/loss/bi_fused_lm_head.py @@ -45,8 +45,8 @@ from xorl.ops.loss.sampling_transform_ce import ( ChunkedScoringPolicy, chunked_transform_scored_ce, - gather_vocab_shards, ) +from xorl.utils.dist_utils import gather_vocab_shards _TEMPERATURE_MATERIALIZE_ROW_CHUNK = EXACT_FILTER_ROW_CHUNK diff --git a/src/xorl/ops/loss/sampling_transform_ce.py b/src/xorl/ops/loss/sampling_transform_ce.py index 91021e7c..3e31d6c4 100644 --- a/src/xorl/ops/loss/sampling_transform_ce.py +++ b/src/xorl/ops/loss/sampling_transform_ce.py @@ -37,6 +37,7 @@ validate_sampling_transform_rows, validate_temperature_rows, ) +from xorl.utils.dist_utils import gather_vocab_shards LogitsFn = Callable[[torch.Tensor, torch.Tensor], torch.Tensor] @@ -128,27 +129,6 @@ def _resolve_tp_layout( return _TpVocabLayout(tp_group, vocab_sizes, dist.get_rank(tp_group)) -def gather_vocab_shards( - local_logits: torch.Tensor, - *, - vocab_sizes: tuple[int, ...], - group: dist.ProcessGroup, -) -> torch.Tensor: - """Gather possibly ragged vocabulary shards in process-group rank order.""" - - max_vocab = max(vocab_sizes) - padded = local_logits.new_zeros((local_logits.shape[0], max_vocab)) - padded[:, : local_logits.shape[1]].copy_(local_logits) - world_size = dist.get_world_size(group) - gathered = local_logits.new_empty((world_size * local_logits.shape[0], max_vocab)) - dist.all_gather_into_tensor(gathered, padded.contiguous(), group=group) - rank_major = gathered.view(world_size, local_logits.shape[0], max_vocab) - return torch.cat( - [rank_major[rank, :, :vocab_size] for rank, vocab_size in enumerate(vocab_sizes)], - dim=1, - ).contiguous() - - def _mm_accumulate_fp32(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: """Matrix multiply with an FP32-accumulated FP32 result.""" @@ -506,6 +486,5 @@ def sampling_transform_per_token_ce( __all__ = [ "ChunkedScoringPolicy", "chunked_transform_scored_ce", - "gather_vocab_shards", "sampling_transform_per_token_ce", ] diff --git a/src/xorl/utils/dist_utils.py b/src/xorl/utils/dist_utils.py index 77dedcee..b62a99c8 100644 --- a/src/xorl/utils/dist_utils.py +++ b/src/xorl/utils/dist_utils.py @@ -24,6 +24,31 @@ def all_gather(tensor: "torch.Tensor", world_size: int) -> "torch.Tensor": return output_tensor.view(-1, *tensor.size()[1:]) +def gather_vocab_shards( + local_logits: "torch.Tensor", + *, + vocab_sizes: tuple[int, ...], + group: "ProcessGroup", +) -> "torch.Tensor": + """Gather possibly ragged vocabulary shards in process-group rank order. + + Shards are padded to the widest shard for one ``all_gather_into_tensor`` + and reassembled as ``[rows, sum(vocab_sizes)]`` in rank order. + """ + + max_vocab = max(vocab_sizes) + padded = local_logits.new_zeros((local_logits.shape[0], max_vocab)) + padded[:, : local_logits.shape[1]].copy_(local_logits) + world_size = dist.get_world_size(group) + gathered = local_logits.new_empty((world_size * local_logits.shape[0], max_vocab)) + dist.all_gather_into_tensor(gathered, padded.contiguous(), group=group) + rank_major = gathered.view(world_size, local_logits.shape[0], max_vocab) + return torch.cat( + [rank_major[rank, :, :vocab_size] for rank, vocab_size in enumerate(vocab_sizes)], + dim=1, + ).contiguous() + + def all_reduce( data: Union[int, float, List[Union[int, float]], "torch.Tensor"], op: Literal["mean", "sum", "max", "min"] = "mean", diff --git a/tests/distributed/test_glm52_exact_lm_head_qlora_collectives.py b/tests/distributed/test_glm52_exact_lm_head_qlora_collectives.py index a899159c..bd99c811 100644 --- a/tests/distributed/test_glm52_exact_lm_head_qlora_collectives.py +++ b/tests/distributed/test_glm52_exact_lm_head_qlora_collectives.py @@ -6,9 +6,10 @@ import torch import torch.distributed as dist +from xorl.models.transformers.exact_lm_head_shared import ExactLmHeadFunction from xorl.models.transformers.glm5.exact_lm_head_qlora import ( _all_reduce_sum_fp32, - _Glm52ExactDistributedTP16LmHeadFunction, + _distributed_row_plan, _rank_order_vocab_all_gather, _require_equal_nonzero_row_count, ) @@ -85,7 +86,7 @@ def _surrogate_vjp( local_lora_B = torch.tensor([[0.75]], dtype=torch.float32, requires_grad=True) local_token_ids = torch.tensor([rank + 10], dtype=torch.int64) local_temperature = torch.tensor([0.7 + rank * 0.6], dtype=torch.float32) - local_logprob = _Glm52ExactDistributedTP16LmHeadFunction.apply( + local_logprob = ExactLmHeadFunction.apply( local_hidden, local_weight, lora_A, @@ -93,6 +94,7 @@ def _surrogate_vjp( local_token_ids, local_temperature, (None, None, None), + _distributed_row_plan(local_hidden, group), _FakeDistributedComponent(), ) torch.testing.assert_close( diff --git a/tests/models/test_dsv4_exact_lm_head_temperature.py b/tests/models/test_dsv4_exact_lm_head_temperature.py index 35e189dd..0511a595 100644 --- a/tests/models/test_dsv4_exact_lm_head_temperature.py +++ b/tests/models/test_dsv4_exact_lm_head_temperature.py @@ -7,10 +7,11 @@ import xorl.models.transformers.deepseek_v4.exact_lm_head as exact_head from xorl.models.transformers.deepseek_v4.exact_lm_head import ( - _Dsv4ExactDistributedHeadFunction, + _distributed_row_plan, _rank_order_variable_row_all_gather, _temperature_scale_bf16_logits, ) +from xorl.models.transformers.exact_lm_head_shared import ExactLmHeadFunction from xorl.ops.exact_sampling_transforms import ( selected_logprob_reference_grad as _selected_logprob_reference_grad, ) @@ -166,7 +167,6 @@ def _surrogate_vjp( monkeypatch.setattr(exact_head, "_rank_order_row_counts", lambda *_args: (2, 0, 0, 0, 0, 0, 0, 0)) monkeypatch.setattr(exact_head, "_rank_order_variable_row_all_gather", lambda value, *_args, **_kwargs: value) - monkeypatch.setattr(exact_head.dist, "get_rank", lambda _group: 0) hidden = torch.arange(6, dtype=torch.float32).reshape(2, 3).to(torch.bfloat16).requires_grad_(True) weight = torch.zeros(5, 3, dtype=torch.bfloat16) @@ -175,7 +175,7 @@ def _surrogate_vjp( token_ids = torch.tensor([0, 4], dtype=torch.int64) temperature = torch.tensor([0.7, 1.3], dtype=torch.float32) - logprob = _Dsv4ExactDistributedHeadFunction.apply( + logprob = ExactLmHeadFunction.apply( hidden, weight, lora_a, @@ -183,6 +183,7 @@ def _surrogate_vjp( token_ids, temperature, (None, None, None), + _distributed_row_plan(hidden, group, FakeComponent.source_ordinal), FakeComponent(), ) logprob.sum().backward() diff --git a/tests/models/test_glm52_exact_lm_head_qlora.py b/tests/models/test_glm52_exact_lm_head_qlora.py index a0875b37..94603aa1 100644 --- a/tests/models/test_glm52_exact_lm_head_qlora.py +++ b/tests/models/test_glm52_exact_lm_head_qlora.py @@ -7,6 +7,14 @@ import torch.nn.functional as F import xorl.models.transformers.glm5.exact_lm_head_qlora as lm_head_impl +from xorl.models.transformers.exact_lm_head_shared import ( + REPLICATED_ROW_PLAN, + ExactLmHeadFunction, + single_adapter_lora_batch_info, +) +from xorl.models.transformers.exact_lm_head_shared import ( + rank_order_vocab_from_stacked as _rank_order_vocab_from_stacked, +) from xorl.models.transformers.glm5.exact_lm_head_qlora import ( GLM52_EXACT_TP16_LM_HEAD_CONTRACT_VERSION, GLM52_LM_HEAD_HIDDEN_SIZE, @@ -15,9 +23,7 @@ GLM52_LM_HEAD_TP_SIZE, GLM52_LM_HEAD_VOCAB_SIZE, Glm52ExactTP16LmHeadSelectedLogprob, - _Glm52ExactTP16LmHeadFunction, _local_qlora_surrogate_vjp, - _rank_order_vocab_from_stacked, glm52_lm_head_shard, ) from xorl.ops.bi_families_v2 import exact_temperature_scale_fp32_logits @@ -280,7 +286,7 @@ def _surrogate_vjp( token_ids = torch.tensor([0, 4], dtype=torch.int64) temperature = torch.tensor([0.7, 1.3], dtype=torch.float32) - logprob = _Glm52ExactTP16LmHeadFunction.apply( + logprob = ExactLmHeadFunction.apply( hidden, weight, lora_A, @@ -288,6 +294,7 @@ def _surrogate_vjp( token_ids, temperature, (None, None, None), + REPLICATED_ROW_PLAN, FakeComponent(), ) assert logprob.requires_grad @@ -457,7 +464,7 @@ def test_official_local_shard_literal_v2_bytes_tail_and_surrogate_gradients() -> effective_A = lora_A.detach().to(torch.bfloat16).contiguous() effective_B = lora_B.detach().to(torch.bfloat16).contiguous() - batch_info = lm_head_impl._single_adapter_lm_head_batch_info(device.index, rows) + batch_info = single_adapter_lora_batch_info(device.index, rows) direct_base, _direct_lse = head_v2_full_logits_with_lse(hidden, local_weight) direct_a = sgemm_lora_a_fwd(hidden, effective_A.unsqueeze(0), batch_info) direct_delta = sgemm_lora_b_fwd(direct_a, effective_B.unsqueeze(0), batch_info) diff --git a/tests/ops/test_exact_sampling_transforms.py b/tests/ops/test_exact_sampling_transforms.py index 2b15a255..3211238f 100644 --- a/tests/ops/test_exact_sampling_transforms.py +++ b/tests/ops/test_exact_sampling_transforms.py @@ -198,6 +198,7 @@ def test_filtered_exact_heads_do_not_save_dense_support_on_autograd_contexts(): modules = [ importlib.import_module("xorl.ops.loss.bi_fused_lm_head"), importlib.import_module("xorl.ops.loss.sampling_transform_ce"), + importlib.import_module("xorl.models.transformers.exact_lm_head_shared"), importlib.import_module("xorl.models.transformers.glm5.exact_lm_head_qlora"), importlib.import_module("xorl.models.transformers.deepseek_v4.exact_lm_head"), ]