Deduplicate the GLM-5.2 and DSV4 exact LM heads onto shared machinery

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

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",
 ]

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",

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,14 +86,15 @@ 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,
             local_lora_B,
             local_token_ids,
             local_temperature,
             (None, None, None),
+            _distributed_row_plan(local_hidden, group),
             _FakeDistributedComponent(),
         )
         torch.testing.assert_close(

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,14 +175,15 @@ 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,
         lora_b,
         token_ids,
         temperature,
         (None, None, None),
+        _distributed_row_plan(hidden, group, FakeComponent.source_ordinal),
         FakeComponent(),
     )
     logprob.sum().backward()

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,14 +286,15 @@ 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,
         lora_B,
         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)

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"),
     ]