diff --git a/vime/backends/megatron_utils/linear_logp_provider.py b/vime/backends/megatron_utils/linear_logp_provider.py index 35e8c828..a81e3837 100644 --- a/vime/backends/megatron_utils/linear_logp_provider.py +++ b/vime/backends/megatron_utils/linear_logp_provider.py @@ -73,6 +73,8 @@ class LinearLogpContext: hidden: torch.Tensor projection: LinearProjection vocab_partition: VocabPartition + reuse_local_logits: bool = False + local_logits: torch.Tensor | None = None def __post_init__(self) -> None: if not isinstance(self.hidden, torch.Tensor) or self.hidden.ndim != 2: @@ -83,6 +85,17 @@ def __post_init__(self) -> None: raise ValueError("linear_logp projection and vocabulary shard widths do not match") if self.hidden.device != self.projection.weight.device: raise ValueError("linear_logp context tensors must share a device") + if not isinstance(self.reuse_local_logits, bool): + raise TypeError("linear_logp reuse_local_logits must be a bool") + if self.local_logits is not None: + if self.local_logits.ndim != 2: + raise ValueError("linear_logp context.local_logits must be [T, V_local]") + if self.local_logits.size(0) != self.hidden.size(0): + raise ValueError("linear_logp context local-logits rows must match hidden rows") + if self.local_logits.size(1) != self.vocab_partition.local_size: + raise ValueError("linear_logp context local-logits width must match vocabulary shard") + if self.local_logits.device != self.hidden.device: + raise ValueError("linear_logp context local logits must share the hidden device") @dataclass(frozen=True) diff --git a/vime/backends/megatron_utils/model.py b/vime/backends/megatron_utils/model.py index 2a1a6f19..6f71e9f3 100644 --- a/vime/backends/megatron_utils/model.py +++ b/vime/backends/megatron_utils/model.py @@ -90,6 +90,7 @@ def capture(module, inputs, kwargs, *, context_owner=owner): real_size=real_vocab_size, padded_size=padded_vocab_size, ), + reuse_local_logits=bool(getattr(module, "__rl_kernel_reusable_local_logits__", False)), ) output_layer._vime_linear_logp_capture_handle = output_layer.register_forward_pre_hook( @@ -101,6 +102,18 @@ def _take_linear_logp_context(model_chunk): owner = _unwrap_model_chunk(model_chunk) context = getattr(owner, "_vime_linear_logp_context", None) owner._vime_linear_logp_context = None + output_layer = getattr(owner, "output_layer", None) + local_logits = getattr(output_layer, "_rl_kernel_local_logits", None) + if output_layer is not None: + output_layer._rl_kernel_local_logits = None + if context is not None and context.reuse_local_logits: + if not isinstance(local_logits, torch.Tensor): + raise RuntimeError("strict reusable LM head did not publish local logits") + if local_logits.ndim == 3: + local_logits = local_logits.transpose(0, 1).contiguous().reshape(-1, local_logits.size(-1)) + elif local_logits.ndim != 2: + raise RuntimeError(f"unsupported strict reusable LM-head logits shape: {tuple(local_logits.shape)}") + context = dataclasses.replace(context, local_logits=local_logits) return context