torchrl 0.11.0__cp314-cp314-win_amd64.whl

torchrl/modules/distributions/truncated_normal.py

@@ -0,0 +1,187 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# from https://github.com/toshas/torch_truncnorm
+from __future__ import annotations
+
+import math
+from numbers import Number
+
+import torch
+from torch.distributions import constraints, Distribution
+from torch.distributions.utils import broadcast_all
+
+CONST_SQRT_2 = math.sqrt(2)
+CONST_INV_SQRT_2PI = 1 / math.sqrt(2 * math.pi)
+CONST_INV_SQRT_2 = 1 / math.sqrt(2)
+CONST_LOG_INV_SQRT_2PI = math.log(CONST_INV_SQRT_2PI)
+CONST_LOG_SQRT_2PI_E = 0.5 * math.log(2 * math.pi * math.e)
+
+
+class TruncatedStandardNormal(Distribution):
+    """Truncated Standard Normal distribution.
+
+    Source: https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    """
+
+    arg_constraints = {
+        "a": constraints.real,
+        "b": constraints.real,
+    }
+    has_rsample = True
+    eps = 1e-6
+
+    def __init__(self, a, b, validate_args=None, device=None):
+        self.a, self.b = broadcast_all(a, b)
+        _non_blocking = device is not None and torch.device(device).type == "cuda"
+        self.a = self.a.to(device, non_blocking=_non_blocking)
+        self.b = self.b.to(device, non_blocking=_non_blocking)
+        if isinstance(a, Number) and isinstance(b, Number):
+            batch_shape = torch.Size()
+        else:
+            batch_shape = self.a.size()
+        super().__init__(batch_shape, validate_args=validate_args)
+        if self.a.dtype != self.b.dtype:
+            raise ValueError("Truncation bounds types are different")
+        if any(
+            (self.a >= self.b)
+            .view(
+                -1,
+            )
+            .tolist()
+        ):
+            raise ValueError("Incorrect truncation range")
+        eps = self.eps
+        self._dtype_min_gt_0 = eps
+        self._dtype_max_lt_1 = 1 - eps
+        self._little_phi_a = self._little_phi(self.a)
+        self._little_phi_b = self._little_phi(self.b)
+        self._big_phi_a = self._big_phi(self.a)
+        self._big_phi_b = self._big_phi(self.b)
+        self._Z = (self._big_phi_b - self._big_phi_a).clamp(eps, 1 - eps)
+        self._log_Z = self._Z.log()
+        little_phi_coeff_a = torch.nan_to_num(self.a, nan=math.nan)
+        little_phi_coeff_b = torch.nan_to_num(self.b, nan=math.nan)
+        self._lpbb_m_lpaa_d_Z = (
+            self._little_phi_b * little_phi_coeff_b
+            - self._little_phi_a * little_phi_coeff_a
+        ) / self._Z
+        self._mean = -(self._little_phi_b - self._little_phi_a) / self._Z
+        self._variance = (
+            1
+            - self._lpbb_m_lpaa_d_Z
+            - ((self._little_phi_b - self._little_phi_a) / self._Z) ** 2
+        )
+        self._entropy = CONST_LOG_SQRT_2PI_E + self._log_Z - 0.5 * self._lpbb_m_lpaa_d_Z
+
+    @constraints.dependent_property
+    def support(self):
+        return constraints.interval(self.a, self.b)
+
+    @property
+    def mean(self):
+        return self._mean
+
+    @property
+    def deterministic_sample(self):
+        return self.mean
+
+    @property
+    def variance(self):
+        return self._variance
+
+    def entropy(self):
+        return self._entropy
+
+    @property
+    def auc(self):
+        return self._Z
+
+    @staticmethod
+    def _little_phi(x):
+        return (-(x**2) * 0.5).exp() * CONST_INV_SQRT_2PI
+
+    def _big_phi(self, x):
+        phi = 0.5 * (1 + (x * CONST_INV_SQRT_2).erf())
+        return phi.clamp(self.eps, 1 - self.eps)
+
+    @staticmethod
+    def _inv_big_phi(x):
+        return CONST_SQRT_2 * (2 * x - 1).erfinv()
+
+    def cdf(self, value):
+        if self._validate_args:
+            self._validate_sample(value)
+        return ((self._big_phi(value) - self._big_phi_a) / self._Z).clamp(0, 1)
+
+    def icdf(self, value):
+        y = self._big_phi_a + value * self._Z
+        y = y.clamp(self.eps, 1 - self.eps)
+        return self._inv_big_phi(y)
+
+    def log_prob(self, value):
+        if self._validate_args:
+            self._validate_sample(value)
+        return CONST_LOG_INV_SQRT_2PI - self._log_Z - (value**2) * 0.5
+
+    def rsample(self, sample_shape=None):
+        if sample_shape is None:
+            sample_shape = torch.Size([])
+        shape = self._extended_shape(sample_shape)
+        p = torch.empty(shape, device=self.a.device).uniform_(
+            self._dtype_min_gt_0, self._dtype_max_lt_1
+        )
+        return self.icdf(p)
+
+
+class TruncatedNormal(TruncatedStandardNormal):
+    """Truncated Normal distribution.
+
+    https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    """
+
+    has_rsample = True
+
+    def __init__(self, loc, scale, a, b, validate_args=None, device=None):
+        scale = scale.clamp_min(self.eps)
+        self.loc, self.scale, a, b = broadcast_all(loc, scale, a, b)
+        _non_blocking = device is not None and torch.device(device).type == "cuda"
+        a = a.to(device, non_blocking=_non_blocking)
+        b = b.to(device, non_blocking=_non_blocking)
+        self._non_std_a = a
+        self._non_std_b = b
+        a = (a - self.loc) / self.scale
+        b = (b - self.loc) / self.scale
+        super().__init__(a, b, validate_args=validate_args)
+        self._log_scale = self.scale.log()
+        self._mean = self._mean * self.scale + self.loc
+        self._variance = self._variance * self.scale**2
+        self._entropy += self._log_scale
+
+    def _to_std_rv(self, value):
+        return (value - self.loc) / self.scale
+
+    def _from_std_rv(self, value):
+        return value * self.scale + self.loc
+
+    def cdf(self, value):
+        return super().cdf(self._to_std_rv(value))
+
+    def icdf(self, value):
+        sample = self._from_std_rv(super().icdf(value))
+
+        # clamp data but keep gradients
+        sample_clip = torch.stack(
+            [sample.detach(), self._non_std_a.detach().expand_as(sample)], 0
+        ).max(0)[0]
+        sample_clip = torch.stack(
+            [sample_clip, self._non_std_b.detach().expand_as(sample)], 0
+        ).min(0)[0]
+        sample.data.copy_(sample_clip)
+        return sample
+
+    def log_prob(self, value):
+        value = self._to_std_rv(value)
+        return super().log_prob(value) - self._log_scale

torchrl/modules/distributions/utils.py

@@ -0,0 +1,233 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import torch
+from torch import autograd, distributions as d
+from torch.distributions import Independent, Transform, TransformedDistribution
+
+try:
+    from torch.compiler import is_dynamo_compiling
+except ImportError:
+    from torch._dynamo import is_compiling as is_dynamo_compiling
+
+
+def _cast_device(elt: torch.Tensor | float, device) -> torch.Tensor | float:
+    if isinstance(elt, torch.Tensor):
+        _non_blocking = device is not None and torch.device(device).type == "cuda"
+        return elt.to(device, non_blocking=_non_blocking)
+    return elt
+
+
+def _cast_transform_device(transform, device):
+    if transform is None:
+        return transform
+    _non_blocking = device is not None and torch.device(device).type == "cuda"
+    if isinstance(transform, d.ComposeTransform):
+        for i, t in enumerate(transform.parts):
+            transform.parts[i] = _cast_transform_device(t, device)
+    elif isinstance(transform, d.Transform):
+        for attribute in dir(transform):
+            value = getattr(transform, attribute)
+            if isinstance(value, torch.Tensor):
+                setattr(
+                    transform, attribute, value.to(device, non_blocking=_non_blocking)
+                )
+        return transform
+    else:
+        raise TypeError(
+            f"Cannot perform device casting for transform of type {type(transform)}"
+        )
+
+
+class FasterTransformedDistribution(TransformedDistribution):
+    """A faster implementation of TransformedDistribution."""
+
+    __doc__ = __doc__ + TransformedDistribution.__doc__
+
+    def __init__(self, base_distribution, transforms, validate_args=None):
+        if is_dynamo_compiling():
+            return super().__init__(
+                base_distribution, transforms, validate_args=validate_args
+            )
+        if isinstance(transforms, Transform):
+            self.transforms = [transforms]
+        elif isinstance(transforms, list):
+            raise ValueError("Make a ComposeTransform first.")
+        else:
+            raise ValueError(
+                f"transforms must be a Transform or list, but was {transforms}"
+            )
+        transform = self.transforms[0]
+        # Reshape base_distribution according to transforms.
+        base_shape = base_distribution.batch_shape + base_distribution.event_shape
+        base_event_dim = len(base_distribution.event_shape)
+        # transform = ComposeTransform(self.transforms)
+        # if len(base_shape) < transform.domain.event_dim:
+        #     raise ValueError("base_distribution needs to have shape with size at least {}, but got {}."
+        #                      .format(transform.domain.event_dim, base_shape))
+        transform_codomain_event_dim = transform.codomain.event_dim
+        transform_domain_event_dim = transform.domain.event_dim
+
+        forward_shape = transform.forward_shape(base_shape)
+        expanded_base_shape = transform.inverse_shape(forward_shape)
+        if base_shape != expanded_base_shape:
+            base_batch_shape = expanded_base_shape[
+                : len(expanded_base_shape) - base_event_dim
+            ]
+            base_distribution = base_distribution.expand(base_batch_shape)
+        reinterpreted_batch_ndims = transform_domain_event_dim - base_event_dim
+        if reinterpreted_batch_ndims > 0:
+            base_distribution = Independent(
+                base_distribution, reinterpreted_batch_ndims
+            )
+        self.base_dist = base_distribution
+
+        # Compute shapes.
+        transform_change_in_event_dim = (
+            transform_codomain_event_dim - transform_domain_event_dim
+        )
+        event_dim = max(
+            transform_codomain_event_dim,  # the transform is coupled
+            base_event_dim + transform_change_in_event_dim,  # the base dist is coupled
+        )
+        cut = len(forward_shape) - event_dim
+        batch_shape = forward_shape[:cut]
+        event_shape = forward_shape[cut:]
+        super(TransformedDistribution, self).__init__(
+            batch_shape, event_shape, validate_args=validate_args
+        )
+
+
+def _safetanh(x, eps):  # noqa: D103
+    lim = 1.0 - eps
+    y = x.tanh()
+    return y.clamp(-lim, lim)
+
+
+def _safeatanh(y, eps):  # noqa: D103
+    lim = 1.0 - eps
+    return y.clamp(-lim, lim).atanh()
+
+
+class _SafeTanh(autograd.Function):
+    generate_vmap_rule = True
+
+    @staticmethod
+    def forward(input, eps):
+        output = input.tanh()
+        lim = 1.0 - eps
+        output = output.clamp(-lim, lim)
+        # ctx.save_for_backward(output)
+        return output
+
+    @staticmethod
+    def setup_context(ctx, inputs, output):
+        # input, eps = inputs
+        # ctx.mark_non_differentiable(ind, ind_inv)
+        # # Tensors must be saved via ctx.save_for_backward. Please do not
+        # # assign them directly onto the ctx object.
+        ctx.save_for_backward(output)
+
+    @staticmethod
+    def backward(ctx, *grad):
+        grad = grad[0]
+        (output,) = ctx.saved_tensors
+        return (grad * (1 - output.pow(2)), None)
+
+
+class _SafeTanhNoEps(autograd.Function):
+    generate_vmap_rule = True
+
+    @staticmethod
+    def forward(input):
+        output = input.tanh()
+        eps = torch.finfo(input.dtype).resolution
+        lim = 1.0 - eps
+        output = output.clamp(-lim, lim)
+        return output
+
+    @staticmethod
+    def setup_context(ctx, inputs, output):
+        ctx.save_for_backward(output)
+
+    @staticmethod
+    def backward(ctx, *grad):
+        grad = grad[0]
+        (output,) = ctx.saved_tensors
+        return (grad * (1 - output.pow(2)),)
+
+
+class _SafeaTanh(autograd.Function):
+    generate_vmap_rule = True
+
+    @staticmethod
+    def forward(tanh_val, eps):
+        if eps is None:
+            eps = torch.finfo(tanh_val.dtype).resolution
+        lim = 1.0 - eps
+        output = tanh_val.clamp(-lim, lim)
+        # ctx.save_for_backward(output)
+        output = output.atanh()
+        return output
+
+    @staticmethod
+    def setup_context(ctx, inputs, output):
+        tanh_val, eps = inputs
+
+        # ctx.mark_non_differentiable(ind, ind_inv)
+        # # Tensors must be saved via ctx.save_for_backward. Please do not
+        # # assign them directly onto the ctx object.
+        ctx.save_for_backward(tanh_val)
+        ctx.eps = eps
+
+    @staticmethod
+    def backward(ctx, *grad):
+        grad = grad[0]
+        (tanh_val,) = ctx.saved_tensors
+        eps = ctx.eps
+        lim = 1.0 - eps
+        output = tanh_val.clamp(-lim, lim)
+        return (grad / (1 - output.pow(2)), None)
+
+
+class _SafeaTanhNoEps(autograd.Function):
+    generate_vmap_rule = True
+
+    @staticmethod
+    def forward(tanh_val):
+        eps = torch.finfo(tanh_val.dtype).resolution
+        lim = 1.0 - eps
+        output = tanh_val.clamp(-lim, lim)
+        # ctx.save_for_backward(output)
+        output = output.atanh()
+        return output
+
+    @staticmethod
+    def setup_context(ctx, inputs, output):
+        tanh_val = inputs[0]
+        eps = torch.finfo(tanh_val.dtype).resolution
+
+        # ctx.mark_non_differentiable(ind, ind_inv)
+        # # Tensors must be saved via ctx.save_for_backward. Please do not
+        # # assign them directly onto the ctx object.
+        ctx.save_for_backward(tanh_val)
+        ctx.eps = eps
+
+    @staticmethod
+    def backward(ctx, *grad):
+        grad = grad[0]
+        (tanh_val,) = ctx.saved_tensors
+        eps = ctx.eps
+        lim = 1.0 - eps
+        output = tanh_val.clamp(-lim, lim)
+        return (grad / (1 - output.pow(2)),)
+
+
+safetanh = _SafeTanh.apply
+safeatanh = _SafeaTanh.apply
+
+safetanh_noeps = _SafeTanhNoEps.apply
+safeatanh_noeps = _SafeaTanhNoEps.apply

torchrl/modules/llm/__init__.py

@@ -0,0 +1,62 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+"""LLM utilities for TorchRL.
+
+Note:
+    This package contains optional integrations (e.g. vLLM) that may rely on native
+    extensions. To keep `import torchrl` / `import torchrl.envs` lightweight and
+    robust, we **avoid importing optional backends at module import time** and
+    instead only import those backends on demand.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .policies.common import ChatHistory, LLMWrapperBase, LogProbs, Masks, Text, Tokens
+from .policies.transformers_wrapper import (
+    RemoteTransformersWrapper,
+    TransformersWrapper,
+)
+from .policies.vllm_wrapper import vLLMWrapper
+
+__all__ = [
+    # Data structures
+    "ChatHistory",
+    "LogProbs",
+    "Masks",
+    "Text",
+    "Tokens",
+    # Wrapper base class
+    "LLMWrapperBase",
+    # Local wrappers
+    "TransformersWrapper",
+    "vLLMWrapper",
+    # Remote wrappers
+    "RemoteTransformersWrapper",
+    # Async vLLM (recommended)
+    "AsyncVLLM",
+    "make_async_vllm_engine",
+    "stateless_init_process_group_async",
+    # Sync vLLM utilities
+    "make_vllm_worker",
+    "stateless_init_process_group",
+]
+
+
+def __getattr__(name: str) -> Any:  # noqa: ANN401
+    # Keep backends optional and on-demand to avoid importing vLLM native extensions
+    # as a side-effect of importing torchrl.
+    if name in {
+        "AsyncVLLM",
+        "make_async_vllm_engine",
+        "make_vllm_worker",
+        "stateless_init_process_group",
+        "stateless_init_process_group_async",
+    }:
+        from . import backends  # local import is intentional / required
+
+        return getattr(backends, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

torchrl/modules/llm/backends/__init__.py

@@ -0,0 +1,65 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+"""LLM backends.
+
+These backends can be optional and may rely on native extensions. We avoid
+importing them at module import time and lazily load on attribute access.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+__all__ = [
+    # Base classes
+    "RLvLLMEngine",
+    # Sync vLLM
+    "make_vllm_worker",
+    "RayLLMWorker",
+    "LocalLLMWrapper",
+    # Async vLLM
+    "_AsyncvLLMWorker",
+    "_AsyncLLMEngine",
+    "AsyncVLLM",
+    "make_async_vllm_engine",
+    # Utilities
+    "stateless_init_process_group",
+    "stateless_init_process_group_async",
+]
+
+_LAZY_ATTRS: dict[str, tuple[str, str]] = {
+    # Base classes and interfaces
+    "RLvLLMEngine": ("torchrl.modules.llm.backends.vllm", "RLvLLMEngine"),
+    # Sync vLLM
+    "make_vllm_worker": ("torchrl.modules.llm.backends.vllm", "make_vllm_worker"),
+    "RayLLMWorker": ("torchrl.modules.llm.backends.vllm", "RayLLMWorker"),
+    "LocalLLMWrapper": ("torchrl.modules.llm.backends.vllm", "LocalLLMWrapper"),
+    # Async vLLM
+    "_AsyncvLLMWorker": ("torchrl.modules.llm.backends.vllm", "_AsyncvLLMWorker"),
+    "_AsyncLLMEngine": ("torchrl.modules.llm.backends.vllm", "_AsyncLLMEngine"),
+    "AsyncVLLM": ("torchrl.modules.llm.backends.vllm", "AsyncVLLM"),
+    "make_async_vllm_engine": (
+        "torchrl.modules.llm.backends.vllm",
+        "make_async_vllm_engine",
+    ),
+    # Utilities
+    "stateless_init_process_group": (
+        "torchrl.modules.llm.backends.vllm",
+        "stateless_init_process_group",
+    ),
+    "stateless_init_process_group_async": (
+        "torchrl.modules.llm.backends.vllm",
+        "stateless_init_process_group_async",
+    ),
+}
+
+
+def __getattr__(name: str) -> Any:  # noqa: ANN401
+    target = _LAZY_ATTRS.get(name)
+    if target is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    module_name, attr_name = target
+    module = __import__(module_name, fromlist=[attr_name])
+    return getattr(module, attr_name)

torchrl/modules/llm/backends/vllm/__init__.py

@@ -0,0 +1,94 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""vLLM backends for TorchRL.
+
+This module provides comprehensive vLLM integration including:
+- Base classes and interfaces
+- Synchronous vLLM workers
+- Asynchronous vLLM services
+- Shared utilities
+
+Examples:
+    >>> # Create an async vLLM service (recommended)
+    >>> from torchrl.modules.llm.backends.vllm import AsyncVLLM
+    >>> service = AsyncVLLM.from_pretrained("Qwen/Qwen2.5-3B")
+
+    >>> # Create a sync Ray worker
+    >>> from torchrl.modules.llm.backends.vllm import make_vllm_worker
+    >>> worker = make_vllm_worker("Qwen/Qwen2.5-3B", make_ray_worker=True)
+
+    >>> # All engines implement the same interface
+    >>> from torchrl.modules.llm.backends.vllm import RLvLLMEngine
+    >>> updater = vLLMUpdaterV2(any_engine)  # Works with any RLvLLMEngine
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+__all__ = [
+    # Base classes and interfaces
+    "RLvLLMEngine",
+    # Synchronous vLLM
+    "make_vllm_worker",
+    "RayLLMWorker",
+    "LocalLLMWrapper",
+    # Asynchronous vLLM
+    "AsyncVLLM",
+    "make_async_vllm_engine",
+    "_AsyncLLMEngine",
+    "_AsyncvLLMWorker",
+    # Utilities
+    "stateless_init_process_group",
+    "stateless_init_process_group_async",
+]
+
+_LAZY_ATTRS: dict[str, tuple[str, str]] = {
+    # Base
+    "RLvLLMEngine": ("torchrl.modules.llm.backends.vllm.base", "RLvLLMEngine"),
+    # Sync
+    "make_vllm_worker": (
+        "torchrl.modules.llm.backends.vllm.vllm_sync",
+        "make_vllm_worker",
+    ),
+    "RayLLMWorker": ("torchrl.modules.llm.backends.vllm.vllm_sync", "RayLLMWorker"),
+    "LocalLLMWrapper": (
+        "torchrl.modules.llm.backends.vllm.vllm_sync",
+        "LocalLLMWrapper",
+    ),
+    # Async
+    "_AsyncLLMEngine": (
+        "torchrl.modules.llm.backends.vllm.vllm_async",
+        "_AsyncLLMEngine",
+    ),
+    "_AsyncvLLMWorker": (
+        "torchrl.modules.llm.backends.vllm.vllm_async",
+        "_AsyncvLLMWorker",
+    ),
+    "AsyncVLLM": ("torchrl.modules.llm.backends.vllm.vllm_async", "AsyncVLLM"),
+    "make_async_vllm_engine": (
+        "torchrl.modules.llm.backends.vllm.vllm_async",
+        "make_async_vllm_engine",
+    ),
+    # Utils
+    "stateless_init_process_group": (
+        "torchrl.modules.llm.backends.vllm.vllm_utils",
+        "stateless_init_process_group",
+    ),
+    "stateless_init_process_group_async": (
+        "torchrl.modules.llm.backends.vllm.vllm_utils",
+        "stateless_init_process_group_async",
+    ),
+}
+
+
+def __getattr__(name: str) -> Any:  # noqa: ANN401
+    target = _LAZY_ATTRS.get(name)
+    if target is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    module_name, attr_name = target
+    module = __import__(module_name, fromlist=[attr_name])
+    return getattr(module, attr_name)

torchrl/modules/llm/backends/vllm/_models.py

@@ -0,0 +1,46 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Override the last layers of your models here."""
+
+from __future__ import annotations
+
+import os
+
+import torch
+
+try:
+    from vllm.config import VllmConfig
+    from vllm.model_executor.models.qwen3 import Qwen3ForCausalLM
+except ImportError:
+
+    class VllmConfig:
+        """Placeholder for VllmConfig class when vLLM is not installed."""
+
+    class Qwen3ForCausalLM:
+        """Placeholder for Qwen3ForCausalLM class when vLLM is not installed."""
+
+
+def is_fp32_output_enabled() -> bool:
+    """Check if FP32 output is enabled."""
+    return os.getenv("VLLM_ENABLE_FP32_OUTPUT", "0") == "1"
+
+
+class Qwen3ForCausalLMFP32(Qwen3ForCausalLM):
+    """Qwen3ForCausalLM with FP32 output."""
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__(vllm_config=vllm_config, prefix=prefix)
+        if is_fp32_output_enabled():
+            self.lm_head.float()
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor | None:
+        if is_fp32_output_enabled():
+            hidden_states = hidden_states.float()
+        logits = self.logits_processor(self.lm_head, hidden_states)
+        return logits